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

[rakhimovv]

Thanks for FreeToken — it set up and ran cleanly on a fresh H100 box, and the offload design was a pleasure to read through.

While running the suite I hit a failure that turned out to be worth a fix. Happy to adjust anything here, including dropping it if you would rather solve it differently.

What breaks

Most CUDA calls in python/freetoken/kernel/csrc/pinned_tensor.cpp are guarded with
TORCH_CHECK(err == cudaSuccess, ...); two attribute queries and the deleter were not
checked at all. TORCH_CHECK reports the error to the caller but never reads it out of the
runtime, so the failure stays latched in the calling thread. The failures these calls can
originate are non-sticky — the context is still perfectly usable — but the next unrelated
CUDA call picks up the stale status and fails as if it were its own.

The result is that a handled exception from this extension silently poisons the rest
of the process.

Revised after review: cudaFreeHost was not covered, and the attribute probe was
raising where it used to fall back safely. Both corrected; see "The fix".

Related: #125 fixes the test-side symptom of the same CUDA 13 behaviour; see
"Relationship to #125" below.

Observed failure chain

On this box the full suite fails twice, and the second failure is collateral from the
first:

2 failed, 1371 passed, 9 skipped

1. tests/kernels/test_pinned_tensor.py::test_host_device_ptr_is_identity_under_uva
   calls host_device_ptr() on unregistered pageable memory. On CUDA 13 that returns
   cudaErrorInvalidValue, so the extension raises:

   RuntimeError: cudaHostGetDevicePointer failed (host memory must be pinned+mapped): invalid argument
   

2. The error is never drained. Later in the same pytest process,
   tests/kernels/test_triton_attention.py::test_paged_triton_attention_matches_reference[None-256]
   dies on its very first line — q = torch.randn(4, num_q_heads, head_dim, device=device) —
   with:

   torch.AcceleratorError: CUDA error: invalid argument
   

That the second failure is collateral and not a Triton bug is easy to confirm: run that
file on its own and it is green.

$ .venv/bin/python -m pytest tests/kernels/test_triton_attention.py
33 passed in 10.76s

A direct probe shows the error is merely uncleared rather than fatal: calling
host_device_ptr on pageable memory, then cudaGetLastError() through ctypes, returns
1 (cudaErrorInvalidValue) on the first read and 0 on the second, after which
torch.randn(4, device="cuda") succeeds normally.

Probe output before the change:

host_ptr_identity: True
pageable -> RuntimeError: cudaHostGetDevicePointer failed (host memory must be pinned+mapped): invalid argument
post-error cuda alloc: POISONED -> AcceleratorError CUDA error: invalid argument

and after:

host_ptr_identity: True
pageable -> RuntimeError: cudaHostGetDevicePointer failed (host memory must be pinned+mapped): invalid argument
post-error cuda alloc: OK
pinned host_device_ptr identity: True
production device_ptr == data_ptr: True

Why this is a real bug, not a test artifact

The clearest live example is kernel/backend.py:driver_cuda_version(), which wraps the
extension call in except Exception: return None and carries on using CUDA afterwards —
it runs at config time on the main thread (via moe/nvfp4_backends.py), so a latched
status there lands on whatever CUDA call the engine makes next.

cudaFreeHost is the other one that matters, and for a different reason: it is a
from_blob deleter, so it runs during GC with no exception to attribute a failure to. A
status left latched there has no visible origin at all.

pinned.device_ptr() short-circuits on _host_ptr_identity(), so the failing pointer
translation is not itself on the hot path on Linux.

The test suite is where it showed up first because pytest keeps one process alive across
files; a long-running server has the same property.

The fix

An FT_CUDA_CHECK macro that calls cudaGetLastError() to drain the latched status
before raising, on the seven calls that already raised: cudaMallocHost, cudaHostAlloc,
cudaGetDevice, cudaHostGetDevicePointer, cudaHostRegister and cudaDriverGetVersion.
The message text of each check is unchanged — old and new both take TORCH_CHECK's
variadic path, so the strings are byte-identical and it is still a c10::Error.

The two remaining sites drain but deliberately do not raise:

• cudaFreeHost in the free_pinned deleter, because a deleter must not throw.
• the two cudaDeviceGetAttribute calls in host_ptr_identity(), which 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. That fallback is
  deliberate and it is the safe answer: "no identity" routes device_ptr() through
  host_device_ptr(), the real translation, which is correct on every platform. Raising
  there would turn 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
  re-raise on every call rather than once.

That covers all nine CUDA calls in the file.

Why a new macro rather than C10_CUDA_CHECK, which drains the same way: _pinned_tensor
is a CppExtension linking only cudart, not libc10_cuda, and C10_CUDA_CHECK would
also replace these specific messages with a generic "CUDA error: ...".

Relationship to #125

#125 fixes the first failure above from the test side: it stops
test_host_device_ptr_is_identity_under_uva from passing a pageable pointer at all, on
the correct grounds that unregistered memory was never inside host_device_ptr's
contract. That is a fair reading and I have no objection to it.

It does not address what this PR is about. #125 changes only
tests/kernels/test_pinned_tensor.py and states that runtime code is unchanged, so
after it lands the C++ still leaves its error latched — the suite just stops containing
a call that triggers it. Any real cudaHostRegister or cudaHostGetDevicePointer
failure, of the kind host_banks.pin() is written to catch, still poisons whatever CUDA
call comes next.

The two changes are complementary rather than competing, and both are needed for a green
suite on CUDA 13. This PR deliberately does not touch
test_host_device_ptr_is_identity_under_uva, to stay out of #125's way.

Test added

A new test, test_failed_pinned_call_leaves_the_context_usable, drives the failure path
directly and asserts the property production depends on:

with pytest.raises(RuntimeError, match="cudaHostRegister failed"):
    ext.host_register(0, 64)
torch.randn(4, device="cuda").sum().item()

It uses host_register, not host_device_ptr, so it does not overlap #125's subject
matter and does not depend on any pageable-pointer behaviour.

Before this change it fails on the line after the raises block:

>       torch.randn(4, device="cuda").sum().item()
E       torch.AcceleratorError: CUDA error: invalid argument
tests/kernels/test_pinned_tensor.py:145: AcceleratorError

After it, it passes.

Before / after

All runs on bd372b6, same box, same install.

Whole suite, unmodified tree:

$ .venv/bin/python -m pytest tests/ -v -rA
2 failed, 1371 passed, 9 skipped

failing on tests/kernels/test_pinned_tensor.py::test_host_device_ptr_is_identity_under_uva
and tests/kernels/test_triton_attention.py::test_paged_triton_attention_matches_reference[None-256].

Whole suite with this PR applied:

$ .venv/bin/python -m pytest tests/ -q
1 failed, 1373 passed, 9 skipped

(Wall-clock is omitted deliberately: the two runs differ mostly in JIT kernel-cache warmth,
not in anything this change does.)

The Triton collateral failure is gone and the new test passes. The one remaining failure
is test_host_device_ptr_is_identity_under_uva, which is #125's to fix and which this
branch leaves untouched.

Isolating the new test against the C++ change alone — same tree, extension rebuilt each
time:

$ .venv/bin/python -m pytest tests/kernels/test_pinned_tensor.py -v   # without the C++ fix
2 failed, 7 passed in 14.21s
$ .venv/bin/python -m pytest tests/kernels/test_pinned_tensor.py -v   # with it
1 failed, 8 passed in 14.98s

Tested on

• GPU: NVIDIA H100 80GB HBM3 (sm_90), one GPU of four in the box
• NVIDIA driver: 580.126.16
• CPU: Intel Xeon Platinum 8462Y+, 56 threads; 2015 GiB system RAM
• OS: Linux 5.15.0-157-generic
• CUDA toolkit: nvcc release 13.1, V13.1.115
• torch 2.11.0+cu130, Python 3.12.13
• FreeToken 0.1.2, at commit bd372b6
• Install: uv pip install -e ".[accel,dev]"

Commands:

uv pip install -e ".[accel,dev]"
.venv/bin/python -m pytest tests/ -v -rA
.venv/bin/python -m pytest tests/kernels/test_triton_attention.py