From 1b330573d763eb2de93c73be76851fd7e9da2f32 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 00:58:11 -0400 Subject: [PATCH 1/5] feat(kernel): build the JIT CUDA kernels on pre-sm_70 GPUs Two device-side constructs in the JIT kernels are sm_70+ and make nvcc/ptxas reject the whole translation unit on Pascal (sm_6x), so the offload gather, the index kernels and the KV store never build there: - `__grid_constant__` on kernel parameters (compute_70+) - the `L1::no_allocate` modifier on `ld.global` (sm_70+) Both are codegen hints with no semantic content: the parameter is passed identically without the annotation, and the load returns the same bytes without the cache hint. Gate each behind `__CUDA_ARCH__ < 700` so every sm_70+ device pass and the host pass emit exactly what they did before. `__nanosleep` in the same header was already guarded this way. Also stop a failed CUDA call in the pinned-tensor extension from poisoning the process: the runtime keeps a failure in the per-thread last-error slot, so the next unrelated `C10_CUDA_CHECK` reported it instead of its own result. This was visible as an unrelated `torch.empty` on CUDA raising "invalid argument" after `host_device_ptr` rejected unregistered memory. `cudaHostGetDevicePointer` on *unregistered* memory is unspecified -- newer arches let UVA degenerate it to identity, Pascal validates registration and returns cudaErrorInvalidValue. Widen the test to accept either, and keep asserting the invariant that actually matters: it must never return a different nonzero alias. The in-contract case stays covered by test_host_bank_pin_registers_and_translates. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01KQy3DziJN9peJ7nA8L59ns --- .../kernel/csrc/include/freetoken/utils.cuh | 10 ++++++ .../kernel/csrc/jit/fast_index_copy.cuh | 18 ++++++++--- python/freetoken/kernel/csrc/jit/index.cu | 4 +-- python/freetoken/kernel/csrc/jit/store.cu | 2 +- .../freetoken/kernel/csrc/pinned_tensor.cpp | 31 ++++++++++++------- tests/kernels/test_pinned_tensor.py | 13 +++++--- 6 files changed, 55 insertions(+), 23 deletions(-) diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832c..04c9d7b3c 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -10,6 +10,16 @@ #include #include +// __grid_constant__ is compute_70+; Pascal's ptxas rejects it outright. It is only ever +// a codegen hint -- the kernel parameter is passed identically without it -- so dropping +// it below sm_70 costs nothing but the constant-bank optimisation. The host pass (no +// __CUDA_ARCH__) and every sm_70+ device pass keep the annotation verbatim. +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700 +#define FT_GRID_CONSTANT +#else +#define FT_GRID_CONSTANT __grid_constant__ +#endif + namespace device { inline constexpr auto kWarpThreads = 32u; diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23ed..5a5c00384 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -33,21 +33,31 @@ inline constexpr auto get_mem_package() { } } +// The L1::no_allocate hint keeps this streaming host gather from evicting L1, but the +// modifier is sm_70+ and Pascal's ptxas rejects it outright. It is a cache hint only, so +// pre-sm_70 issues the plain load and loses nothing but the anti-pollution optimisation. +// (st.global.wt below assembles fine on sm_6x, so the stores are left alone.) +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700 +#define FT_LD_STREAM "ld.global." +#else +#define FT_LD_STREAM "ld.global.L1::no_allocate." +#endif + __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { uint32_t tmp; - asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); + asm volatile(FT_LD_STREAM "b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { uint32_t tmp0, tmp1; - asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); + asm volatile(FT_LD_STREAM "v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { uint32_t tmp0, tmp1, tmp2, tmp3; - asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); + asm volatile(FT_LD_STREAM "v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; } @@ -485,7 +495,7 @@ struct MultiIndexCopyParams { template __global__ __launch_bounds__(kNumThreads) void fast_index_copy_multi( - const __grid_constant__ MultiIndexCopyParams p + const FT_GRID_CONSTANT MultiIndexCopyParams p ) { const int b = static_cast(blockIdx.x / kBlocksPerBank); if (b >= p.num_banks) { diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index ca0e1db26..90908f297 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -31,7 +31,7 @@ struct MaskedKernelParams { template __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void // - index_kernel(const __grid_constant__ IndexKernelParams params) { + index_kernel(const FT_GRID_CONSTANT IndexKernelParams params) { using namespace device; constexpr auto kSize = kElementSize; constexpr auto kSizePerWarp = kSize / kNumSplits; @@ -62,7 +62,7 @@ template __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void // masked_index_kernel( - const __grid_constant__ MaskedKernelParams mask_params) { + const FT_GRID_CONSTANT MaskedKernelParams mask_params) { using namespace device; constexpr auto kSize = kElementSize; constexpr auto kSizePerWarp = kSize / kNumSplits; diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu index 8d84d76ef..93773cf64 100644 --- a/python/freetoken/kernel/csrc/jit/store.cu +++ b/python/freetoken/kernel/csrc/jit/store.cu @@ -25,7 +25,7 @@ struct StoreKernelParams { template __global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void // - store_kv_cache(const __grid_constant__ StoreKernelParams params) { + store_kv_cache(const FT_GRID_CONSTANT StoreKernelParams params) { using namespace device; constexpr auto kWarpPerBlock = diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adfa..cf2af25e7 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -10,6 +10,18 @@ void free_pinned(void *ptr) { } } +// A failed CUDA runtime call leaves its status in the per-thread "last error" slot, so +// the next unrelated C10_CUDA_CHECK anywhere in the process reports it instead of its +// own result. Consume it before throwing to keep the failure local to this call. It +// matters most for host_device_ptr, which drivers may legitimately reject (pre-sm_70 +// validates registration where newer arches let UVA degenerate the lookup to identity). +void check_cuda(cudaError_t err, const char *what) { + if (err != cudaSuccess) { + cudaGetLastError(); + TORCH_CHECK(false, what, ": ", cudaGetErrorString(err)); + } +} + torch::Tensor create_pinned_tensor_like(torch::Tensor input) { TORCH_CHECK(input.device().is_cpu(), "Input tensor must be on CPU"); TORCH_CHECK(input.layout() == torch::kStrided, @@ -35,8 +47,7 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) { void *data_ptr = nullptr; const cudaError_t alloc_err = cudaMallocHost(&data_ptr, alloc_nbytes); - TORCH_CHECK(alloc_err == cudaSuccess, - "cudaMallocHost failed: ", cudaGetErrorString(alloc_err)); + check_cuda(alloc_err, "cudaMallocHost failed"); auto options = input.options().device(torch::kCPU).pinned_memory(true); @@ -60,8 +71,7 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, 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)); + check_cuda(alloc_err, "cudaHostAlloc failed"); auto options = torch::TensorOptions() .dtype(dtype) @@ -77,7 +87,7 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, bool host_ptr_identity() { int device = 0; const cudaError_t err = cudaGetDevice(&device); - TORCH_CHECK(err == cudaSuccess, "cudaGetDevice failed: ", cudaGetErrorString(err)); + check_cuda(err, "cudaGetDevice failed"); int uva = 0, reg = 0; cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); @@ -88,9 +98,8 @@ 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)); + check_cuda(err, + "cudaHostGetDevicePointer failed (host memory must be pinned+mapped)"); return reinterpret_cast(dev_ptr); } @@ -98,15 +107,13 @@ 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)); + check_cuda(err, "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)); + check_cuda(err, "cudaDriverGetVersion failed"); return version; } diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd5..a8a3132bf 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -120,12 +120,17 @@ def test_host_device_ptr_is_identity_under_uva(): torch.cuda.init() if not _host_ptr_identity(): pytest.skip("non-UVA platform: host_device_ptr rejects unregistered memory instead") - # Under UVA cudaHostGetDevicePointer degenerates to identity for any host pointer - # (no registration validation); rejection of pageable memory only exists on - # non-identity platforms (Windows/WDDM), where the translation is real. + # cudaHostGetDevicePointer on *unregistered* memory is unspecified: recent arches let + # UVA degenerate it to identity, while Pascal validates registration and returns + # cudaErrorInvalidValue. Either is fine -- what must never happen is a different + # nonzero alias, which would silently misaddress the zero-copy gather. The in-contract + # case (registered memory -> identity) is test_host_bank_pin_registers_and_translates. pageable = torch.empty(64, dtype=torch.uint8) ext = _load_pinned_extension() - assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() + try: + assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() + except RuntimeError as exc: + assert "cudaHostGetDevicePointer failed" in str(exc) def test_host_bank_pin_registers_and_translates(): From b9a8edbd4dec1a8be132eb568e396e49116d0238 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 01:25:23 -0400 Subject: [PATCH 2/5] feat(kernel): fall back to libdevice tanh below sm_75 `_fast_tanh` emits `tanh.approx.f32` as inline PTX. That instruction is sm_75+, and ptxas rejects the whole module without it ("Feature 'tanh' requires .target sm_75 or higher"), so every GELU_TANH activation failed to compile on Pascal. Gate it on a `constexpr_function` reading the compilation target, the same idiom e4m3_compat.py uses to branch fp8e4nv below sm_89. Being compile-time, this adds no kernel parameter and leaves the cache key untouched: sm_75+ still takes the single-instruction path, folds the branch away, and emits the PTX it did before. Older cards compute the same tanh through libdevice, which this file already depends on for `libdevice.erf`. Adds tests/kernels/test_activation.py: the four *_and_mul kernels against torch references, plus a saturation case at +-1e4 that would catch a fallback that overflows to nan at the tails. gelu_tanh had no direct coverage before -- it was only reached transitively through the MoE tests. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01KQy3DziJN9peJ7nA8L59ns --- python/freetoken/kernel/triton/activation.py | 20 ++++- tests/kernels/test_activation.py | 81 ++++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) create mode 100644 tests/kernels/test_activation.py diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533e..2daaea3e9 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -21,6 +21,9 @@ from triton.language.extra import libdevice from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents +from triton.language import target_info +from triton.runtime.jit import constexpr_function + from freetoken.utils.arch import is_sm90_supported SILU = 0 @@ -46,6 +49,15 @@ def _pdl_supported() -> bool: return is_sm90_supported() +@constexpr_function +def _fast_tanh_cx(): + """Compile-time: can the target issue ``tanh.approx.f32``? It is an sm_75+ + instruction and pre-Turing ptxas rejects the whole module ("Feature 'tanh' requires + .target sm_75 or higher"), so older cards take the libdevice fallback below. Resolved + from the compilation target like :mod:`freetoken.kernel.triton.e4m3_compat` does.""" + return target_info.cuda_capability_geq(7, 5) + + @triton.jit def _fast_tanh(x): # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. @@ -96,9 +108,13 @@ def _act_and_mul_kernel( if ACT == 0: # SILU: x / (1 + exp(-x)) via ex2.approx act = gate / (1.0 + _fast_ex2(-gate * _LOG2E)) y = act * up - elif ACT == 2: # GELU_TANH via tanh.approx + elif ACT == 2: # GELU_TANH via tanh.approx (sm_75+), else libdevice inner = 0.7978845608028654 * (gate + 0.044715 * gate * gate * gate) - act = 0.5 * gate * (1.0 + _fast_tanh(inner)) + if _fast_tanh_cx(): + tanh_inner = _fast_tanh(inner) + else: + tanh_inner = libdevice.tanh(inner) + act = 0.5 * gate * (1.0 + tanh_inner) y = act * up elif ACT == 3: # SWIGLUOAI: clamped gate/up, sigmoid(alpha*gate), (up + 1) bias gate = tl.minimum(gate, limit) diff --git a/tests/kernels/test_activation.py b/tests/kernels/test_activation.py new file mode 100644 index 000000000..7bf870c7f --- /dev/null +++ b/tests/kernels/test_activation.py @@ -0,0 +1,81 @@ +"""Triton *_and_mul activation kernels vs torch references. + +Guards the GELU_TANH fallback in particular: `tanh.approx.f32` is an sm_75+ PTX +instruction, so pre-Turing takes a libdevice path that must agree with it. +""" + +import pytest +import torch + + +def _xy(rows, d, dtype): + torch.manual_seed(0) + x = torch.randn(rows, 2 * d, device="cuda", dtype=dtype) + return x, x[:, :d].float(), x[:, d:].float() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("d", [128, 1000]) +def test_silu_and_mul_matches_torch(dtype, d): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + from freetoken.kernel.triton.activation import silu_and_mul + + x, gate, up = _xy(17, d, dtype) + ref = torch.nn.functional.silu(gate) * up + torch.testing.assert_close(silu_and_mul(x).float(), ref, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("d", [128, 1000]) +def test_gelu_and_mul_matches_torch(dtype, d): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + from freetoken.kernel.triton.activation import gelu_and_mul + + x, gate, up = _xy(17, d, dtype) + ref = torch.nn.functional.gelu(gate, approximate="none") * up + torch.testing.assert_close(gelu_and_mul(x).float(), ref, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("d", [128, 1000]) +def test_gelu_tanh_and_mul_matches_torch(dtype, d): + """The sm_75+ `tanh.approx.f32` path and the libdevice fallback must agree here.""" + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + from freetoken.kernel.triton.activation import gelu_tanh_and_mul + + x, gate, up = _xy(17, d, dtype) + ref = torch.nn.functional.gelu(gate, approximate="tanh") * up + torch.testing.assert_close(gelu_tanh_and_mul(x).float(), ref, atol=2e-2, rtol=2e-2) + + +def test_gelu_tanh_saturates_at_both_tails(): + """tanh must saturate to +-1 rather than overflow: gelu_tanh(x) -> x for large +x + and -> 0 for large -x. An exp-based fallback that overflows would produce nan.""" + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + from freetoken.kernel.triton.activation import gelu_tanh_and_mul + + gate = torch.tensor([-1e4, -60.0, -8.0, 0.0, 8.0, 60.0, 1e4], device="cuda") + x = torch.cat([gate, torch.ones_like(gate)]).reshape(1, -1) + out = gelu_tanh_and_mul(x).float().flatten() + assert torch.isfinite(out).all(), out + ref = torch.nn.functional.gelu(gate, approximate="tanh") + torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_swigluoai_and_mul_matches_reference(dtype): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + from freetoken.kernel.triton.activation import swigluoai_and_mul + + alpha, limit = 1.702, 7.0 + x, gate, up = _xy(17, 256, dtype) + g = gate.clamp(max=limit) + u = up.clamp(-limit, limit) + ref = g * torch.sigmoid(alpha * g) * (u + 1.0) + got = swigluoai_and_mul(x, alpha=alpha, limit=limit).float() + torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2) From 745974b01420fb0ed170718bd453a3b98f3e293e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 01:25:52 -0400 Subject: [PATCH 3/5] feat(moe): route moe_align to the atomic-free staged path below sm_70 `moe_align.py` ranks its scatter with `tl.atomic_add`. Triton lowers every atomic to a scoped *and* ordered PTX encoding -- `atom.global.gpu..add` -- and both the `.gpu` scope and every memory order (`.relaxed` included, the weakest it can emit) arrived with sm_70. ptxas therefore rejects the entire module on Pascal, and no `sem=`/`scope=` argument avoids it. Verified by probing all four sems and all three scopes on an sm_61 device: every one fails to assemble. Rather than emulate atomics, route around them. `moe_impl.py` already carries a second implementation -- the staged `moe_align_block_size_triton`, a counts/cumsum/binary-search chain across five launches with no atomic anywhere. It produces the same three buffers, so pre-sm_70 dispatches there. Checked against a reference of the documented contract on sm_61 across decode and prefill shapes (1..1024 rows, 8..257 experts, block 16..64): identical padding, every token placed exactly once inside its own expert's region. sm_70+ is untouched -- same branch, same kernel, same launch as before. The staged path costs five launches instead of one, which is the right trade against not running. Not attempted: emulating the atomic with unqualified inline PTX (`atom.global.add.u32` is sm_20+). It assembles and gives correct ranks, but `tl.inline_asm_elementwise` is contracted for *pure* elementwise ops -- when the operand tensor is narrower than the thread block, the layout replicates elements across threads and a side-effecting instruction executes once per replica. A 128-thread block over an 8-wide tensor adds 16x. That silently miscounts at exactly the decode widths that matter (numel = batch x topk), so it is not a safe basis for expert dispatch. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01KQy3DziJN9peJ7nA8L59ns --- python/freetoken/moe/fused.py | 13 +++++++++++++ python/freetoken/utils/arch.py | 4 ++++ 2 files changed, 17 insertions(+) diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index fe7e417d7..2d129de3d 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -128,6 +128,19 @@ def moe_align_block_size( from freetoken.kernel.backend import is_sgl_kernel_installed if not is_sgl_kernel_installed(): + from freetoken.utils.arch import is_sm70_supported + + if not is_sm70_supported(): + # The fused triton path ranks the scatter with tl.atomic_add, and triton + # lowers every atomic to a scoped+ordered PTX encoding that only sm_70+ + # can assemble -- there is no sem=/scope= that avoids it. The staged + # implementation computes the same buffers with a counts/cumsum chain and + # no atomics, so pre-Volta routes there instead. Slower (5 launches vs 1), + # which is the right trade against not running at all. + from freetoken.kernel import moe_align_block_size_triton + + return moe_align_block_size_triton(topk_ids, block_size, num_experts) + from freetoken.kernel.triton.moe_align import ( moe_align_block_size as triton_moe_align_block_size, ) diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d5..7d90d3fb8 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -40,6 +40,10 @@ def is_sm100_family() -> bool: return _is_arch_family(10) +def is_sm70_supported() -> bool: + return is_arch_supported(7, 0) + + def is_sm90_supported() -> bool: return is_arch_supported(9, 0) From cbad20c4d82a4b9bf2b3e2c712cae1ce432c2fb5 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 01:34:24 -0400 Subject: [PATCH 4/5] feat(attention): size decode/extend tiles to the device shared-memory budget `_select_extend_tile` already shrank the prefill tile when a device's opt-in shared memory could not hold it, but two gaps let it overflow anyway: - head_dim <= 128 returned (128, 64) before any budget check, so the most common head sizes never consulted the device at all; - the ladder stopped one rung short of what a 48KB budget needs (pre-Volta gets 48KB per block with no opt-in at all). Carry the ladder down to (16, 16) -- the floor, since `tl.dot` needs N >= 16 -- and budget-check every rung including the largest. The byte estimate is also only a lower bound: what triton allocates depends on how it schedules the pipeline, and the split kernel wants about twice the q/k/v tile bytes. A head_dim 64 model passed the estimate and still needed 64KB. Rather than carry a per-kernel fudge factor, `_select_extend_tiles` now returns the descending ladder and `_launch_first_fitting_tile` walks it, letting the compiler have the final say: OutOfResources is raised during launch setup, after the compile and before any GPU work, so retrying smaller is side-effect free and triton caches each compile. The split-k decode kernel had no sizing at all: BLOCK_N=32 with num_stages=2 hardcoded, whose pipelined k/v tiles want 128KB at head_dim 512 -- fine on A100/H100, over budget on consumer Ampere/Ada, never mind Pascal. Add `_select_decode_tile`, stepping 32/2 -> 32/1 -> 16/1. Both selectors only ever shrink a config that does not fit, and the existing tile-selection test still pins the exact choice for every datacenter and consumer budget it covered before, so any device that fits the current default keeps it. `smem_optin == 0` (unavailable) keeps the prior choice. Devices that previously raised OutOfResources here may now run. head_dim 512 extend has no fitting configuration on a 48KB device rather than a merely slower one -- 16x16 q/k/v tiles alone want 48KB. Its test now skips with that reason, gated on the device's actual budget so it still runs everywhere it can. On sm_61: tests/kernels/test_triton_attention.py goes from 9 failed to 35 passed, 2 skipped. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01KQy3DziJN9peJ7nA8L59ns --- python/freetoken/kernel/triton/attention.py | 157 +++++++++++++------- tests/kernels/test_triton_attention.py | 40 ++++- 2 files changed, 143 insertions(+), 54 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84f..341d5a803 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -5,6 +5,7 @@ import torch import triton import triton.language as tl +from triton.runtime.errors import OutOfResources _MAX_KV_SPLITS = 8 @@ -18,28 +19,72 @@ def _optin_smem_bytes(device_index: int) -> int: return int(getattr(props, "shared_memory_per_block_optin", 0)) -def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]: - """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware. +def _select_extend_tiles(head_dim: int, block_d: int, smem_optin: int) -> list[tuple[int, int]]: + """Candidate ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, largest first. Larger tiles run materially faster (~2x for head_dim 512 on H100) but their bf16 q/k/v tiles need about ``(BLOCK_M + 2 * BLOCK_N) * BLOCK_D * 2`` bytes of shared memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once head_dim >= 256. Keep the fast tiles where the device's opt-in shared memory fits them (datacenter - A100/H100); shrink only where it does not. ``smem_optin == 0`` (unknown) conservatively - selects the small tiles, i.e. the prior consumer-safe behavior. + A100/H100); shrink only where it does not. + + That byte count is a *lower bound*: what triton actually allocates depends on how it + schedules the pipeline, and the split kernel wants roughly twice the tile bytes. So + this returns the whole descending ladder from the first plausible rung and the caller + lets the compiler have the final say -- see :func:`_launch_first_fitting_tile`. The + floor is (16, 16) because ``tl.dot`` requires N >= 16. ``smem_optin == 0`` (unknown) + conservatively starts at the small tiles, i.e. the prior consumer-safe behavior. """ budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch - def fits(block_m: int, block_n: int) -> bool: - return (block_m + 2 * block_n) * block_d * 2 <= budget - - if head_dim <= 128: - return 128, 64 if head_dim <= 256: - return (128, 64) if fits(128, 64) else (64, 32) - if head_dim <= 384: - return (32, 64) if fits(32, 64) else (32, 32) - return (32, 64) if fits(32, 64) else (16, 16) + ladder = [(128, 64), (64, 32), (32, 32), (16, 16)] + elif head_dim <= 384: + ladder = [(32, 64), (32, 32), (16, 16)] + else: + ladder = [(32, 64), (16, 16)] + if not smem_optin: + return ladder[1:] if len(ladder) > 1 else ladder + first = next( + (i for i, (m, n) in enumerate(ladder) if (m + 2 * n) * block_d * 2 <= budget), + len(ladder) - 1, + ) + return ladder[first:] + + +def _launch_first_fitting_tile(run, tiles: list[tuple[int, int]]): + """Call ``run(block_m, block_n)`` on the first tile triton can actually allocate. + + ``OutOfResources`` is raised while triton sets up the launch -- after compiling the + kernel, before any GPU work is issued -- so dropping to a smaller tile and retrying is + side-effect free. Triton caches each compile, so this costs one extra compile the + first time a shape is seen on an undersized device and nothing afterwards. + """ + for i, (block_m, block_n) in enumerate(tiles): + try: + return run(block_m, block_n) + except OutOfResources: + if i == len(tiles) - 1: + raise + + +def _select_decode_tile(block_d: int, block_dv: int, smem_optin: int) -> tuple[int, int]: + """Pick ``(BLOCK_N, num_stages)`` for the split-k decode kernel, shared-memory aware. + + The pipelined k/v tiles cost about ``num_stages * BLOCK_N * (BLOCK_D + BLOCK_DV) * 2`` + bytes. The default 32/2 needs 128KB at head_dim 512 -- fine on A100/H100, over budget + on consumer Ampere/Ada and far over Pascal's 48KB. Shrink only where it does not fit, + so every device that fits the default keeps it. ``smem_optin == 0`` (unknown) keeps + the default, matching the prior behavior. + """ + if not smem_optin: + return 32, 2 + budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch + + for block_n, stages in ((32, 2), (32, 1), (16, 1)): + if stages * block_n * (block_d + block_dv) * 2 <= budget: + return block_n, stages + return 16, 1 @triton.jit @@ -398,6 +443,9 @@ def decode_paged_attention( block_h = triton.next_power_of_2(valid_block_h) block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) + block_n, num_stages = _select_decode_tile( + block_d, block_dv, _optin_smem_bytes(q.device.index) + ) _decode_grouped_stage1_kernel[ (batch, triton.cdiv(num_q_heads, valid_block_h), max_kv_splits) @@ -428,7 +476,7 @@ def decode_paged_attention( NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, BLOCK_DV=block_dv, - BLOCK_N=32, + BLOCK_N=block_n, BLOCK_H=block_h, VALID_BLOCK_H=valid_block_h, MIN_BLOCK_KV=_MIN_BLOCK_KV, @@ -436,7 +484,7 @@ def decode_paged_attention( DV=head_dim, SLIDING_WINDOW=sliding_window or 0, num_warps=4, - num_stages=2, + num_stages=num_stages, ) _decode_stage2_kernel[(batch, num_q_heads)]( attn_logits, @@ -797,18 +845,20 @@ def extend_paged_attention( block_dv = triton.next_power_of_2(head_dim) # Tile size is shared-memory bound: keep the fast (large) tiles on GPUs whose opt-in # shared memory fits them, shrink on consumer GPUs (sm_89 ~99KB) where the default - # 128x64 overflows once head_dim >= 256 (e.g. gemma4: SWA 256, full-attention 512). - block_m, block_n = _select_extend_tile( - head_dim, block_d, _optin_smem_bytes(q.device.index) - ) - grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) - if k_extend is not None or v_extend is not None: + # 128x64 overflows once head_dim >= 256 (e.g. gemma4: SWA 256, full-attention 512), + # and further on pre-Volta, which caps a block at 48KB with no opt-in at all. + tiles = _select_extend_tiles(head_dim, block_d, _optin_smem_bytes(q.device.index)) + use_split = k_extend is not None or v_extend is not None + if use_split: assert k_extend is not None and v_extend is not None assert k_extend.is_cuda and v_extend.is_cuda assert k_extend.dim() == 3 and v_extend.dim() == 3 assert k_extend.shape[0] == num_q_tokens and v_extend.shape[0] == num_q_tokens assert k_extend.shape[1] == num_kv_heads and v_extend.shape[1] == num_kv_heads assert k_extend.shape[-1] == head_dim and v_extend.shape[-1] == head_dim + + def _run_split(block_m: int, block_n: int) -> None: + grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) _extend_attention_split_kernel[grid]( q, k_extend, @@ -845,38 +895,41 @@ def extend_paged_attention( num_warps=8, num_stages=1, ) - return o - _extend_attention_kernel[grid]( - q, - k_cache, - v_cache, - o, - qo_indptr, - kv_indptr, - kv_indices, - prefix_lens, - sm_scale, - sinks_arg, - q.stride(0), - q.stride(1), - k_cache.stride(0), - k_cache.stride(1), - v_cache.stride(0), - v_cache.stride(1), - o.stride(0), - o.stride(1), - GROUP=num_q_heads // num_kv_heads, - D=head_dim, - BLOCK_D=block_d, - BLOCK_DV=block_dv, - BLOCK_M=block_m, - BLOCK_N=block_n, - SLIDING_WINDOW=sliding_window or 0, - HAS_SINKS=sinks is not None, - num_warps=8, - num_stages=1, - ) + def _run_plain(block_m: int, block_n: int) -> None: + grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) + _extend_attention_kernel[grid]( + q, + k_cache, + v_cache, + o, + qo_indptr, + kv_indptr, + kv_indices, + prefix_lens, + sm_scale, + sinks_arg, + q.stride(0), + q.stride(1), + k_cache.stride(0), + k_cache.stride(1), + v_cache.stride(0), + v_cache.stride(1), + o.stride(0), + o.stride(1), + GROUP=num_q_heads // num_kv_heads, + D=head_dim, + BLOCK_D=block_d, + BLOCK_DV=block_dv, + BLOCK_M=block_m, + BLOCK_N=block_n, + SLIDING_WINDOW=sliding_window or 0, + HAS_SINKS=sinks is not None, + num_warps=8, + num_stages=1, + ) + + _launch_first_fitting_tile(_run_split if use_split else _run_plain, tiles) return o diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 6f4afca9e..e1df7a19b 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -6,6 +6,24 @@ import torch +def _skip_if_smem_too_small(head_dim: int) -> None: + """Skip when no valid extend tile fits this device's per-block shared memory. + + ``_select_extend_tiles`` bottoms out at (16, 16) -- ``tl.dot`` needs N >= 16 -- whose + q/k/v tiles alone want ``(16 + 2 * 16) * head_dim * 2`` bytes. Pre-Volta caps a block + at 48KB with no opt-in, so head_dim 512 has no fitting configuration on that hardware + rather than a merely slower one. + """ + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + # 0.8 mirrors the selector's own headroom for scores/acc/alignment/triton scratch. + if optin and (16 + 2 * 16) * head_dim * 2 > optin * 0.8: + pytest.skip( + f"head_dim {head_dim} needs more shared memory than this device offers " + f"per block ({optin} bytes); the smallest valid extend tile does not fit" + ) + + def _reference_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -441,6 +459,7 @@ def test_extend_triton_attention_matches_reference( ): from freetoken.kernel.triton.attention import extend_paged_attention + _skip_if_smem_too_small(head_dim) torch.manual_seed(2) device = torch.device("cuda") num_q_heads = 16 @@ -603,15 +622,32 @@ def test_extend_triton_attention_with_sinks_matches_reference(use_split_inputs: # unknown budget -> conservative small tiles (prior consumer-safe behavior) (256, 0, (64, 32)), (512, 0, (16, 16)), + # pre-Volta: 48KB per block, no opt-in + (64, 49152, (128, 64)), + (128, 49152, (64, 32)), + (256, 49152, (16, 16)), ], ) def test_select_extend_tile_is_shared_memory_aware(head_dim, smem_optin, expected): + """The first (largest) candidate is what a device actually runs, so it pins the + per-device choice; the rest of the ladder only comes into play when triton reports + the tile does not fit after all.""" import triton - from freetoken.kernel.triton.attention import _select_extend_tile + from freetoken.kernel.triton.attention import _select_extend_tiles block_d = triton.next_power_of_2(head_dim) - assert _select_extend_tile(head_dim, block_d, smem_optin) == expected + assert _select_extend_tiles(head_dim, block_d, smem_optin)[0] == expected + + +def test_select_extend_tiles_descends_to_the_dot_floor(): + import triton + + from freetoken.kernel.triton.attention import _select_extend_tiles + + tiles = _select_extend_tiles(128, triton.next_power_of_2(128), 49152) + assert tiles[-1] == (16, 16), "tl.dot needs N >= 16, so (16, 16) is the floor" + assert tiles == sorted(tiles, reverse=True), "candidates must descend" @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") From c8e4c08e6548eb185b8d278efee5a03e26522655 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 01:53:36 -0400 Subject: [PATCH 5/5] feat(sampling): torch fallback for top-k/top-p below sm_70 The triton top-p/top-k threshold search reduces through `tl.atomic_add` and `tl.atomic_max`. Triton lowers every atomic to a scoped, ordered PTX encoding, and both the scope and every memory order it can emit arrived with sm_70, so ptxas rejects those kernels on Pascal. No unit test covered them, so this surfaced only when serving: `top_k` and `top_p` are set by default for most models, and the first real request killed the scheduler with a raw ptxas dump. Sorting needs no atomics. `kernel/torch_sampling.py` reproduces the same selection with torch ops -- top-k by the k-th largest value, top-p by the shortest descending prefix that reaches p, both keeping the elements the kernel path keeps, including ties and the crossing element. `softmax` has no atomics and is re-exported from the triton module unchanged, so only the threshold search changes. The draw uses cumsum + searchsorted rather than torch.multinomial: it stays capturable in a CUDA graph and never syncs to the host, matching what the kernel path guarantees. Seeding mirrors `triton.sampling._gen_u`, including its plain-torch.rand path while a stream is capturing. Selected in `sample_impl` only when flashinfer is absent *and* the device is pre-sm_70, so every currently supported GPU keeps the kernel path untouched. Full-vocabulary sorting is slower than the bracketed histogram search it replaces. On a card that cannot run the kernels at all, sampling is not the bottleneck. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01KQy3DziJN9peJ7nA8L59ns --- python/freetoken/engine/sample.py | 5 + python/freetoken/kernel/torch_sampling.py | 137 ++++++++++++++++++++++ tests/kernels/test_torch_sampling.py | 90 ++++++++++++++ 3 files changed, 232 insertions(+) create mode 100644 python/freetoken/kernel/torch_sampling.py create mode 100644 tests/kernels/test_torch_sampling.py diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 01d14b1aa..50e5ce57d 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -28,9 +28,14 @@ def sample_impl( top_p: torch.Tensor | float | None, ) -> torch.Tensor: from freetoken.kernel.backend import is_flashinfer_installed + from freetoken.utils.arch import is_sm70_supported if is_flashinfer_installed(): import flashinfer.sampling as sampling + elif not is_sm70_supported(): + # The triton top-p/top-k threshold search accumulates through tl.atomic_*, which + # triton can only lower to sm_70+ encodings. Sort with torch instead. + import freetoken.kernel.torch_sampling as sampling else: import freetoken.kernel.triton.sampling as sampling diff --git a/python/freetoken/kernel/torch_sampling.py b/python/freetoken/kernel/torch_sampling.py new file mode 100644 index 000000000..11171b4e5 --- /dev/null +++ b/python/freetoken/kernel/torch_sampling.py @@ -0,0 +1,137 @@ +"""Torch sampling fallback for GPUs whose atomics triton cannot lower. + +``freetoken.kernel.triton.sampling`` finds its top-p / top-k thresholds with histogram +and reduction passes that accumulate through ``tl.atomic_add`` / ``tl.atomic_max``. +Triton lowers every atomic to a scoped, ordered PTX encoding (``atom.global.gpu.``) +whose scope and memory order both arrived with sm_70, so on Pascal ptxas rejects those +kernels outright and any request with ``top_k`` or ``top_p`` set -- the default for most +models -- kills the scheduler. + +The threshold search is what needs atomics; sorting does not. These wrappers reproduce +the same selection with plain torch ops, sharing the triton ``softmax`` (which has no +atomics and compiles everywhere). Sorting the whole vocabulary is slower than the +bracketed histogram search the triton path uses, but it is correct, it is CUDA-graph +capturable, and on a card in this class sampling is not the bottleneck. + +Selected by :func:`freetoken.engine.sample.sample_impl`; nothing else should import it. +""" + +from __future__ import annotations + +import torch + +from freetoken.kernel.triton.sampling import softmax # noqa: F401 (re-exported) + +__all__ = [ + "softmax", + "sampling_from_probs", + "top_k_renorm_probs", + "top_p_renorm_probs", + "top_k_sampling_from_probs", + "top_p_sampling_from_probs", + "top_k_top_p_sampling_from_probs", +] + + +def _per_row(value, rows: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """A scalar or per-row tensor as a ``(rows, 1)`` column, for broadcasting.""" + if isinstance(value, torch.Tensor): + return value.to(device=device, dtype=dtype).reshape(rows, 1) + return torch.full((rows, 1), value, device=device, dtype=dtype) + + +def _renormalize(kept: torch.Tensor) -> torch.Tensor: + total = kept.sum(-1, keepdim=True) + # A row whose whole mass was filtered out (all-zero probs) would divide by zero; + # leave it uniform rather than emitting nan, matching the kernel path's behaviour + # of always returning a drawable distribution. + return torch.where(total > 0, kept / total.clamp_min(torch.finfo(kept.dtype).tiny), + torch.full_like(kept, 1.0 / kept.shape[-1])) + + +def top_k_renorm_probs(probs: torch.Tensor, top_k) -> torch.Tensor: + """Zero everything below the k-th largest probability per row, then renormalize. + + Values tied with the k-th are kept, so a row can retain more than ``k`` entries -- + the same behaviour as the threshold-based kernel path, which cannot separate ties + either. + """ + probs = probs.float() + rows, vocab = probs.shape + k = _per_row(top_k, rows, probs.device, torch.long).clamp_(1, vocab) + kth = probs.sort(dim=-1, descending=True).values.gather(1, k - 1) + return _renormalize(probs.masked_fill(probs < kth, 0.0)) + + +def top_p_renorm_probs(probs: torch.Tensor, top_p) -> torch.Tensor: + """Keep the shortest descending prefix whose mass reaches ``top_p``, renormalize.""" + probs = probs.float() + rows, _ = probs.shape + p = _per_row(top_p, rows, probs.device, torch.float32) + ordered, order = probs.sort(dim=-1, descending=True) + # Exclusive cumulative mass: an entry is dropped only once everything *before* it + # already reached p, which keeps the element that crosses the threshold. + drop = (ordered.cumsum(-1) - ordered) >= p + kept = torch.zeros_like(probs).scatter_(1, order, ordered.masked_fill(drop, 0.0)) + return _renormalize(kept) + + +def _draw(probs: torch.Tensor, seed=None, offset=None) -> torch.Tensor: + """Inverse-CDF draw, one token per row. + + ``searchsorted`` over the cumulative distribution rather than ``torch.multinomial``: + it stays capturable in a CUDA graph and never syncs to the host. Seeding mirrors + ``triton.sampling._gen_u`` -- and like it, a capturing stream takes the plain + ``torch.rand`` path, since a graph replays whatever generator state it captured. + """ + rows, vocab = probs.shape + if seed is not None and not torch.cuda.is_current_stream_capturing(): + generator = torch.Generator(device=probs.device) + s = int(seed if not isinstance(seed, torch.Tensor) else seed.view(-1)[0]) + o = 0 if offset is None else int( + offset if not isinstance(offset, torch.Tensor) else offset.view(-1)[0] + ) + generator.manual_seed((s * 0x9E3779B97F4A7C15 + o) & 0x7FFFFFFFFFFFFFFF) + u = torch.rand(rows, 1, device=probs.device, dtype=torch.float32, generator=generator) + else: + u = torch.rand(rows, 1, device=probs.device, dtype=torch.float32) + cdf = probs.cumsum(-1) + idx = torch.searchsorted(cdf.contiguous(), (u * cdf[:, -1:]).contiguous(), right=True) + return idx.squeeze(-1).clamp_(max=vocab - 1).to(torch.int32) + + +def _finish(out: torch.Tensor, indices, return_valid: bool): + out = out.to(indices.dtype) if indices is not None else out + return (out, torch.ones_like(out, dtype=torch.bool)) if return_valid else out + + +def _source(probs: torch.Tensor, indices) -> torch.Tensor: + probs = probs.float() + return probs if indices is None else probs[indices].contiguous() + + +def sampling_from_probs(probs, indices=None, deterministic=True, generator=None, + check_nan=False, seed=None, offset=None, return_valid=False): + src = _source(probs, indices) + return _finish(_draw(src, seed, offset), indices, return_valid) + + +def top_k_sampling_from_probs(probs, top_k, indices=None, deterministic=True, generator=None, + check_nan=False, seed=None, offset=None, return_valid=False): + src = _source(probs, indices) + return _finish(_draw(top_k_renorm_probs(src, top_k), seed, offset), indices, return_valid) + + +def top_p_sampling_from_probs(probs, top_p, indices=None, deterministic=True, generator=None, + check_nan=False, seed=None, offset=None, return_valid=False): + src = _source(probs, indices) + return _finish(_draw(top_p_renorm_probs(src, top_p), seed, offset), indices, return_valid) + + +def top_k_top_p_sampling_from_probs(probs, top_k, top_p, indices=None, + filter_apply_order="top_k_first", deterministic=True, + generator=None, check_nan=False, seed=None, offset=None, + return_valid=False): + src = _source(probs, indices) + renormed = top_p_renorm_probs(top_k_renorm_probs(src, top_k), top_p) + return _finish(_draw(renormed, seed, offset), indices, return_valid) diff --git a/tests/kernels/test_torch_sampling.py b/tests/kernels/test_torch_sampling.py new file mode 100644 index 000000000..2234829fe --- /dev/null +++ b/tests/kernels/test_torch_sampling.py @@ -0,0 +1,90 @@ +"""The torch sampling fallback must select the same tokens as the kernel path. + +These assert the *renormalization* semantics (which tokens survive, with what mass) +rather than the draw, so they are deterministic and arch-independent. +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +@pytest.fixture +def probs(): + torch.manual_seed(0) + return torch.softmax(torch.randn(5, 512, device="cuda"), dim=-1) + + +def test_top_k_keeps_exactly_the_k_largest(probs): + from freetoken.kernel.torch_sampling import top_k_renorm_probs + + for k in (1, 3, 40, 512): + out = top_k_renorm_probs(probs, k) + assert torch.allclose(out.sum(-1), torch.ones(probs.shape[0], device="cuda"), atol=1e-5) + kept = out > 0 + assert (kept.sum(-1) == k).all(), f"k={k} kept {kept.sum(-1).tolist()}" + # the survivors are the k largest of the original row + expect = probs.topk(k, dim=-1).indices.sort(-1).values + assert torch.equal(kept.nonzero()[:, 1].reshape(probs.shape[0], k), expect) + + +def test_top_k_accepts_per_row_tensor(probs): + from freetoken.kernel.torch_sampling import top_k_renorm_probs + + k = torch.tensor([1, 2, 3, 4, 5], device="cuda", dtype=torch.int32) + kept = top_k_renorm_probs(probs, k) > 0 + assert kept.sum(-1).tolist() == [1, 2, 3, 4, 5] + + +def test_top_p_keeps_the_crossing_element(probs): + """The prefix must *reach* p, so the element that crosses it is retained.""" + from freetoken.kernel.torch_sampling import top_p_renorm_probs + + for p in (0.1, 0.5, 0.9, 1.0): + out = top_p_renorm_probs(probs, p) + assert torch.allclose(out.sum(-1), torch.ones(probs.shape[0], device="cuda"), atol=1e-5) + kept = out > 0 + for row in range(probs.shape[0]): + ordered = probs[row].sort(descending=True).values + n = int(kept[row].sum()) + assert ordered[:n].sum() >= p - 1e-5, f"p={p} row={row} prefix short" + if n > 1: + assert ordered[: n - 1].sum() < p, f"p={p} row={row} prefix longer than needed" + + +def test_top_p_1_keeps_everything_with_mass(probs): + from freetoken.kernel.torch_sampling import top_p_renorm_probs + + assert ((top_p_renorm_probs(probs, 1.0) > 0) == (probs > 0)).all() + + +def test_draw_only_returns_tokens_the_filter_kept(probs): + from freetoken.kernel.torch_sampling import top_k_sampling_from_probs + + k = 4 + allowed = probs.topk(k, dim=-1).indices + for seed in range(8): + out = top_k_sampling_from_probs(probs, k, seed=seed) + assert out.dtype == torch.int32 and out.shape == (probs.shape[0],) + assert (out[:, None] == allowed).any(-1).all(), f"drew a filtered-out token: {out}" + + +def test_draw_is_deterministic_for_a_seed(probs): + from freetoken.kernel.torch_sampling import top_p_sampling_from_probs + + a = top_p_sampling_from_probs(probs, 0.8, seed=1234) + b = top_p_sampling_from_probs(probs, 0.8, seed=1234) + assert torch.equal(a, b) + + +def test_degenerate_row_stays_drawable(): + """An all-zero row must not produce nan or an out-of-range token.""" + from freetoken.kernel.torch_sampling import top_k_sampling_from_probs + + p = torch.zeros(2, 32, device="cuda") + p[1, 7] = 1.0 + out = top_k_sampling_from_probs(p, 4, seed=0) + assert out.shape == (2,) + assert (out >= 0).all() and (out < 32).all() + assert out[1].item() == 7