Support pre-Volta GPUs (sm_61 / Pascal)

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
 

python/freetoken/kernel/csrc/include/freetoken/utils.cuh

@@ -10,6 +10,16 @@
 #include <source_location>
 #include <type_traits>
 
+// __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;

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 <typename IdType, std::size_t kNumThreads, std::size_t kBlocksPerBank>
 __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<int>(blockIdx.x / kBlocksPerBank);
     if (b >= p.num_banks) {

python/freetoken/kernel/csrc/jit/index.cu

@@ -31,7 +31,7 @@ struct MaskedKernelParams {
 template <std::size_t kNumThreads, std::size_t kMaxOccupancy, bool kUsePDL,
           std::size_t kElementSize, std::size_t kNumSplits, std::integral T>
 __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 <std::size_t kNumThreads, std::size_t kMaxOccupancy, bool kUsePDL,
           std::size_t kElementSize, std::size_t kNumSplits, std::integral T>
 __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;

python/freetoken/kernel/csrc/jit/store.cu

@@ -25,7 +25,7 @@ struct StoreKernelParams {
 template <std::size_t kNumThreads, std::size_t kMaxOccupancy, bool kUsePDL,
           std::size_t kElementSize, std::integral T>
 __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 =

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<int64_t> 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<int64_t> 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(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device);
@@ -88,25 +98,22 @@ 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));
+  check_cuda(err,
+             "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));
+  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;
 }
 

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.<sem>``)
+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)

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)