diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..d35e9b98 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/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(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 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 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(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); + 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; } 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")