fix(kernel): clear the latched CUDA error before raising in pinned_tensor

python/freetoken/kernel/csrc/pinned_tensor.cpp

@@ -4,9 +4,30 @@
 
 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. 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) {                                              \
+      [[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();
   }
 }
 
@@ -34,9 +55,8 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) {
   const size_t alloc_nbytes = static_cast<size_t>(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 +78,9 @@ torch::Tensor alloc_pinned_tensor(std::vector<int64_t> 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 +95,42 @@ torch::Tensor alloc_pinned_tensor(std::vector<int64_t> 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: ");
+  // 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;
-  cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device);
-  cudaDeviceGetAttribute(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device);
+  if (cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device) != cudaSuccess) {
+    [[maybe_unused]] const cudaError_t drained = cudaGetLastError();
+    uva = 0;
+  }
+  if (cudaDeviceGetAttribute(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem,
+                             device) != cudaSuccess) {
+    [[maybe_unused]] const cudaError_t drained = cudaGetLastError();
+    reg = 0;
+  }
   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<void *>(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<void *>(host_ptr), 0),
+      "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ");
   return reinterpret_cast<int64_t>(dev_ptr);
 }
 
 void host_register(int64_t addr, int64_t nbytes) {
-  const cudaError_t err =
-      cudaHostRegister(reinterpret_cast<void *>(addr), static_cast<size_t>(nbytes),
-                       cudaHostRegisterPortable | cudaHostRegisterMapped);
-  TORCH_CHECK(err == cudaSuccess,
-              "cudaHostRegister failed: ", cudaGetErrorString(err));
+  FT_CUDA_CHECK(cudaHostRegister(reinterpret_cast<void *>(addr),
+                                 static_cast<size_t>(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;
 }
 

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")