From 102eadc9c2cb1c19a1785394219188aef3c95175 Mon Sep 17 00:00:00 2001 From: Ruslan Rakhimov Date: Mon, 24 Aug 2026 17:51:44 +0300 Subject: [PATCH 1/2] fix(kernel): clear the latched CUDA error before raising in pinned_tensor TORCH_CHECK reports a failed CUDA call but does not read the status out of the runtime, so the error stays latched in the calling thread. The errors raised here are non-sticky and the context remains usable, but the next unrelated CUDA call picks up the stale status and fails as if it were its own. That turns a handled exception into a process-wide fault. host_banks.pin() catches the RuntimeError from host_register() and re-raises a friendlier message; any torch call made while handling that failure would report the stale cudaHostRegister error instead of its own. Route every CUDA call in the file through an FT_CUDA_CHECK macro that drains the status with cudaGetLastError() before raising. The two cudaDeviceGetAttribute calls in host_ptr_identity() were previously unchecked and are now checked as well. Check messages are unchanged. Add a test that drives the failure path directly and asserts the context survives it. --- .../freetoken/kernel/csrc/pinned_tensor.cpp | 57 +++++++++++-------- tests/kernels/test_pinned_tensor.py | 17 ++++++ 2 files changed, 50 insertions(+), 24 deletions(-) diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..e2678d85 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -4,6 +4,20 @@ namespace { +// A failed runtime call latches its status in the calling thread until someone +// reads it. TORCH_CHECK reports the error but does not clear it, so the next +// unrelated CUDA call -- torch's, in practice -- reports *this* failure instead +// of its own. Drain it before raising: every error below is non-sticky, so the +// context stays usable for a caller that handles the exception. +#define FT_CUDA_CHECK(expr, ...) \ + do { \ + const cudaError_t ft_err_ = (expr); \ + if (ft_err_ != cudaSuccess) { \ + cudaGetLastError(); \ + TORCH_CHECK(false, __VA_ARGS__, cudaGetErrorString(ft_err_)); \ + } \ + } while (0) + void free_pinned(void *ptr) { if (ptr != nullptr) { cudaFreeHost(ptr); @@ -34,9 +48,8 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) { const size_t alloc_nbytes = static_cast(nbytes == 0 ? 1 : nbytes); void *data_ptr = nullptr; - const cudaError_t alloc_err = cudaMallocHost(&data_ptr, alloc_nbytes); - TORCH_CHECK(alloc_err == cudaSuccess, - "cudaMallocHost failed: ", cudaGetErrorString(alloc_err)); + FT_CUDA_CHECK(cudaMallocHost(&data_ptr, alloc_nbytes), + "cudaMallocHost failed: "); auto options = input.options().device(torch::kCPU).pinned_memory(true); @@ -58,10 +71,9 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // Portable + mapped: the offload gather kernel reads these banks straight // from host memory (zero-copy), which requires device-mapped pinned pages. void *data_ptr = nullptr; - const cudaError_t alloc_err = cudaHostAlloc( - &data_ptr, alloc_nbytes, cudaHostAllocPortable | cudaHostAllocMapped); - TORCH_CHECK(alloc_err == cudaSuccess, - "cudaHostAlloc failed: ", cudaGetErrorString(alloc_err)); + FT_CUDA_CHECK(cudaHostAlloc(&data_ptr, alloc_nbytes, + cudaHostAllocPortable | cudaHostAllocMapped), + "cudaHostAlloc failed: "); auto options = torch::TensorOptions() .dtype(dtype) @@ -76,37 +88,34 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // device address). Zero-copy consumers resolve bank base addresses through these. bool host_ptr_identity() { int device = 0; - const cudaError_t err = cudaGetDevice(&device); - TORCH_CHECK(err == cudaSuccess, "cudaGetDevice failed: ", cudaGetErrorString(err)); + FT_CUDA_CHECK(cudaGetDevice(&device), "cudaGetDevice failed: "); int uva = 0, reg = 0; - cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); - cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); + FT_CUDA_CHECK(cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device), + "cudaDeviceGetAttribute(UnifiedAddressing) failed: "); + FT_CUDA_CHECK(cudaDeviceGetAttribute( + ®, cudaDevAttrCanUseHostPointerForRegisteredMem, device), + "cudaDeviceGetAttribute(CanUseHostPointerForRegisteredMem) failed: "); return uva == 1 && reg == 1; } int64_t host_device_ptr(int64_t host_ptr) { void *dev_ptr = nullptr; - const cudaError_t err = - cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast(host_ptr), 0); - TORCH_CHECK(err == cudaSuccess, - "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ", - cudaGetErrorString(err)); + FT_CUDA_CHECK( + cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast(host_ptr), 0), + "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): "); return reinterpret_cast(dev_ptr); } void host_register(int64_t addr, int64_t nbytes) { - const cudaError_t err = - cudaHostRegister(reinterpret_cast(addr), static_cast(nbytes), - cudaHostRegisterPortable | cudaHostRegisterMapped); - TORCH_CHECK(err == cudaSuccess, - "cudaHostRegister failed: ", cudaGetErrorString(err)); + FT_CUDA_CHECK(cudaHostRegister(reinterpret_cast(addr), + static_cast(nbytes), + cudaHostRegisterPortable | cudaHostRegisterMapped), + "cudaHostRegister failed: "); } int64_t driver_cuda_version() { int version = 0; // stays 0 when no driver is installed - const cudaError_t err = cudaDriverGetVersion(&version); - TORCH_CHECK(err == cudaSuccess, - "cudaDriverGetVersion failed: ", cudaGetErrorString(err)); + FT_CUDA_CHECK(cudaDriverGetVersion(&version), "cudaDriverGetVersion failed: "); return version; } diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd..589565bd 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -128,6 +128,23 @@ def test_host_device_ptr_is_identity_under_uva(): assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() +def test_failed_pinned_call_leaves_the_context_usable(): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + + from freetoken.kernel.pinned import _load_pinned_extension + + torch.cuda.init() + ext = _load_pinned_extension() + # A rejected registration is a normal outcome -- host_banks.pin() catches it. + # The runtime latches the status until someone reads it, so if the extension + # raises without draining it, the next unrelated CUDA call reports this + # failure as its own and a caller that handled the exception dies anyway. + with pytest.raises(RuntimeError, match="cudaHostRegister failed"): + ext.host_register(0, 64) + torch.randn(4, device="cuda").sum().item() + + def test_host_bank_pin_registers_and_translates(): if not torch.cuda.is_available(): pytest.skip("needs CUDA") From a1daa7aeea57935c566514a15eae03408831d3d8 Mon Sep 17 00:00:00 2001 From: Ruslan Rakhimov Date: Mon, 24 Aug 2026 20:54:39 +0300 Subject: [PATCH 2/2] fix(kernel): drain in the deleter too, and keep the attribute probe non-fatal Two corrections to the previous commit. cudaFreeHost was the one call left uncovered, and it is the worst place to leak a status: it runs as a from_blob deleter during GC, with no exception to attribute the failure to. It drains unconditionally and never throws. host_ptr_identity's two cudaDeviceGetAttribute calls previously discarded their status with uva/reg pre-initialised to 0 -- the same idiom as driver_cuda_version's "stays 0 when no driver is installed", i.e. a deliberate fallback rather than an oversight. Raising there turned an unqueryable attribute into a fatal error on the offload path, and _host_ptr_identity is lru_cached, which does not cache exceptions, so it would have re-raised on every call. Drain without raising and keep the 0 fallback: "no identity" routes device_ptr through host_device_ptr, the real translation, which is correct everywhere. Also soften the macro comment: these calls report whatever is latched on the thread, so "every error below is non-sticky" was a stronger claim than holds. Draining leaves the context usable for the failures these calls originate; a sticky error latched elsewhere re-latches on the next call. --- .../freetoken/kernel/csrc/pinned_tensor.cpp | 35 +++++++++++++------ 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index e2678d85..d35e9b98 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -7,20 +7,27 @@ namespace { // A failed runtime call latches its status in the calling thread until someone // reads it. TORCH_CHECK reports the error but does not clear it, so the next // unrelated CUDA call -- torch's, in practice -- reports *this* failure instead -// of its own. Drain it before raising: every error below is non-sticky, so the -// context stays usable for a caller that handles the exception. +// of its own. Drain it before raising. The failures these calls can originate +// are non-sticky, so draining leaves the context usable; a sticky error latched +// elsewhere will simply re-latch on the next call, which is correct. #define FT_CUDA_CHECK(expr, ...) \ do { \ const cudaError_t ft_err_ = (expr); \ if (ft_err_ != cudaSuccess) { \ - cudaGetLastError(); \ + [[maybe_unused]] const cudaError_t ft_drained_ = cudaGetLastError(); \ TORCH_CHECK(false, __VA_ARGS__, cudaGetErrorString(ft_err_)); \ } \ } while (0) void free_pinned(void *ptr) { - if (ptr != nullptr) { - cudaFreeHost(ptr); + if (ptr == nullptr) { + return; + } + // A from_blob deleter runs during GC, with no exception to attribute a failure + // to -- so this drains and never throws. Left latched it would be the worst + // case of the bug above: a stale status with no visible origin at all. + if (cudaFreeHost(ptr) != cudaSuccess) { + [[maybe_unused]] const cudaError_t drained = cudaGetLastError(); } } @@ -89,12 +96,20 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, bool host_ptr_identity() { int device = 0; FT_CUDA_CHECK(cudaGetDevice(&device), "cudaGetDevice failed: "); + // An unqueryable attribute stays 0 and answers "no identity", which is the safe + // answer: device_ptr() then goes through host_device_ptr(), the real translation, + // correct on every platform. Only the latch is new here -- do not raise, or an + // attribute query that used to degrade gracefully becomes fatal on the offload path. int uva = 0, reg = 0; - FT_CUDA_CHECK(cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device), - "cudaDeviceGetAttribute(UnifiedAddressing) failed: "); - FT_CUDA_CHECK(cudaDeviceGetAttribute( - ®, cudaDevAttrCanUseHostPointerForRegisteredMem, device), - "cudaDeviceGetAttribute(CanUseHostPointerForRegisteredMem) failed: "); + if (cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device) != cudaSuccess) { + [[maybe_unused]] const cudaError_t drained = cudaGetLastError(); + uva = 0; + } + if (cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, + device) != cudaSuccess) { + [[maybe_unused]] const cudaError_t drained = cudaGetLastError(); + reg = 0; + } return uva == 1 && reg == 1; }