Add logits-processor source, tests, and model card

Files changed (9) hide show

CARD.md +70 -0
README.md +70 -0
build.toml +23 -0
flake.nix +17 -0
logits_processor_cuda/logits_processor.cu +204 -0
tests/test_logits_processor.py +71 -0
torch-ext/logits_processor/__init__.py +63 -0
torch-ext/torch_binding.cpp +17 -0
torch-ext/torch_binding.h +17 -0

CARD.md ADDED Viewed

	@@ -0,0 +1,70 @@

+---
+library_name: kernels
+license: apache-2.0
+---
+# logits-processor
+The generation last-mile as loadable [`kernels`](https://github.com/huggingface/kernels):
+guided-decoding masking and fused sampling, the two hot-path ops between a
+model's logits and the next token that the Kernel Hub did not carry.
+```python
+from kernels import get_kernel
+lp = get_kernel("phanerozoic/logits-processor")
+# 1) structured / guided decoding: -inf the tokens a grammar disallows
+lp.apply_token_bitmask(logits, bitmask)          # logits [B, V], bitmask [B, ceil(V/32)] int32
+# 2) fused temperature + top-k ∩ top-p ∩ min-p + multinomial draw
+tokens = lp.sample(logits, temperature=0.7, top_k=50, top_p=0.9, min_p=0.0, seed=step)
+```
+## Ops
+### `apply_token_bitmask(logits, bitmask) -> logits`
+Sets the logits of disallowed tokens to `-inf`, in place, from a packed allow-mask.
+Bit `1` means the token is **allowed** (XGrammar convention). `logits` is `[B, V]`
+in float/half/bf16; `bitmask` is `[B, ceil(V/32)]` int32. This is the op behind
+constrained JSON and tool-calling output — apply it, then `sample`.
+### `sample(logits, temperature=1.0, top_k=0, top_p=1.0, min_p=0.0, seed=0, offset=0) -> tokens`
+One `logits -> token` op: per-row temperature scaling, then the **intersection**
+of top-k, top-p (nucleus), and min-p filtering, then a multinomial draw.
+- `top_k <= 0`, `top_p >= 1`, `min_p <= 0` each disable that filter.
+- `temperature <= 0` is greedy (argmax).
+- Every sampling parameter is per-row: pass a scalar to broadcast, or a length-`B`
+  tensor for a heterogeneous batch — continuous batching with a different grammar,
+  temperature, or nucleus per request.
+- The draw is a counter-based RNG keyed by `(seed, row, offset)`: deterministic
+  and replayable. Advance `offset` (or `seed`) per decoding step.
+- Returns int64 token ids `[B]`.
+## Correctness
+Verified through `get_kernel` on A100 (sm80), A10G (sm86), L4 (sm89), H200 (sm90),
+and RTX PRO 6000 (sm120):
+- the bitmask sets exactly the disallowed logits to `-inf` in float / half / bf16,
+  including a vocabulary size that is not a multiple of 32;
+- greedy and `top_k=1` return the argmax;
+- the sampled token always lies in the exact top-k / top-p / min-p keep-set and in
+  their intersection, checked against an independent reference over many draws;
+- the empirical sampling distribution matches the filtered, renormalized softmax
+  reference to within 0.002 (0.0003 at a 32k vocabulary);
+- draws are deterministic in `(seed, offset)` and change when `offset` advances;
+- per-row temperature / top-k / top-p / min-p are honored independently across a
+  batch, and a bitmask followed by `sample` never draws a disallowed token.
+## Status
+v1 is correctness-first: sampling is a single descending sort of the softmax
+probabilities followed by a per-row prefix cutoff and inverse-CDF draw. The sort
+dominates; a radix-select over the top-p mass is the first planned optimization.
+Built for sm80 / sm86 / sm89 / sm90 / sm100 / sm120; certified on all of these
+except sm100, for which no B200 was available to run the check.
+## License
+Apache-2.0.

README.md ADDED Viewed

	@@ -0,0 +1,70 @@

+---
+library_name: kernels
+license: apache-2.0
+---
+# logits-processor
+The generation last-mile as loadable [`kernels`](https://github.com/huggingface/kernels):
+guided-decoding masking and fused sampling, the two hot-path ops between a
+model's logits and the next token that the Kernel Hub did not carry.
+```python
+from kernels import get_kernel
+lp = get_kernel("phanerozoic/logits-processor")
+# 1) structured / guided decoding: -inf the tokens a grammar disallows
+lp.apply_token_bitmask(logits, bitmask)          # logits [B, V], bitmask [B, ceil(V/32)] int32
+# 2) fused temperature + top-k ∩ top-p ∩ min-p + multinomial draw
+tokens = lp.sample(logits, temperature=0.7, top_k=50, top_p=0.9, min_p=0.0, seed=step)
+```
+## Ops
+### `apply_token_bitmask(logits, bitmask) -> logits`
+Sets the logits of disallowed tokens to `-inf`, in place, from a packed allow-mask.
+Bit `1` means the token is **allowed** (XGrammar convention). `logits` is `[B, V]`
+in float/half/bf16; `bitmask` is `[B, ceil(V/32)]` int32. This is the op behind
+constrained JSON and tool-calling output — apply it, then `sample`.
+### `sample(logits, temperature=1.0, top_k=0, top_p=1.0, min_p=0.0, seed=0, offset=0) -> tokens`
+One `logits -> token` op: per-row temperature scaling, then the **intersection**
+of top-k, top-p (nucleus), and min-p filtering, then a multinomial draw.
+- `top_k <= 0`, `top_p >= 1`, `min_p <= 0` each disable that filter.
+- `temperature <= 0` is greedy (argmax).
+- Every sampling parameter is per-row: pass a scalar to broadcast, or a length-`B`
+  tensor for a heterogeneous batch — continuous batching with a different grammar,
+  temperature, or nucleus per request.
+- The draw is a counter-based RNG keyed by `(seed, row, offset)`: deterministic
+  and replayable. Advance `offset` (or `seed`) per decoding step.
+- Returns int64 token ids `[B]`.
+## Correctness
+Verified through `get_kernel` on A100 (sm80), A10G (sm86), L4 (sm89), H200 (sm90),
+and RTX PRO 6000 (sm120):
+- the bitmask sets exactly the disallowed logits to `-inf` in float / half / bf16,
+  including a vocabulary size that is not a multiple of 32;
+- greedy and `top_k=1` return the argmax;
+- the sampled token always lies in the exact top-k / top-p / min-p keep-set and in
+  their intersection, checked against an independent reference over many draws;
+- the empirical sampling distribution matches the filtered, renormalized softmax
+  reference to within 0.002 (0.0003 at a 32k vocabulary);
+- draws are deterministic in `(seed, offset)` and change when `offset` advances;
+- per-row temperature / top-k / top-p / min-p are honored independently across a
+  batch, and a bitmask followed by `sample` never draws a disallowed token.
+## Status
+v1 is correctness-first: sampling is a single descending sort of the softmax
+probabilities followed by a per-row prefix cutoff and inverse-CDF draw. The sort
+dominates; a radix-select over the top-p mass is the first planned optimization.
+Built for sm80 / sm86 / sm89 / sm90 / sm100 / sm120; certified on all of these
+except sm100, for which no B200 was available to run the check.
+## License
+Apache-2.0.

build.toml ADDED Viewed

	@@ -0,0 +1,23 @@

+[general]
+name = "logits-processor"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = [
+  "cuda",
+]
+[general.hub]
+repo-id = "phanerozoic/logits-processor"
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+[kernel.logits_processor]
+backend = "cuda"
+cuda-capabilities = ["8.0", "8.6", "8.9", "9.0", "10.0", "12.0"]
+depends = ["torch"]
+src = ["logits_processor_cuda/logits_processor.cu"]

flake.nix ADDED Viewed

	@@ -0,0 +1,17 @@

+{
+  description = "Flake for logits-processor kernel";
+  inputs = {
+    builder.url = "github:huggingface/kernels";
+  };
+  outputs =
+    {
+      self,
+      builder,
+    }:
+    builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

logits_processor_cuda/logits_processor.cu ADDED Viewed

	@@ -0,0 +1,204 @@

+// logits-processor: the generation last-mile as loadable kernels.
+//
+// Two ops that every serving stack needs on the hot path but that the Kernel Hub
+// did not carry:
+//
+//   apply_token_bitmask_inplace  - structured/guided decoding. Given a packed
+//     allow-mask per request (XGrammar convention: bit 1 = token allowed), set
+//     every disallowed logit to -inf in place, so the following softmax cannot
+//     select it. This is the kernel behind constrained JSON / tool-calling.
+//
+//   sample - temperature, then the intersection of top-k / top-p (nucleus) /
+//     min-p, then a multinomial draw, fused into one logits -> token op with a
+//     counter-based RNG (deterministic in (seed, row, offset)). temperature <= 0
+//     is greedy. The filters are computed on a single descending sort of the
+//     row so their cutoffs are all prefix lengths, and the smallest wins.
+//
+// v1 favors being obviously correct over peak throughput: one segmented sort of
+// the softmax probabilities, then a per-row prefix walk for the cutoff and the
+// inverse-CDF draw. The sort dominates; replacing it with a radix-select over
+// the top-p mass is the first optimization once this certifies on GPU.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cub/cub.cuh>
+#include <cstdint>
+#include <math.h>
+namespace {
+constexpr int TPB = 256;
+__device__ __forceinline__ float to_f(float x) { return x; }
+__device__ __forceinline__ float to_f(__half x) { return __half2float(x); }
+__device__ __forceinline__ float to_f(__nv_bfloat16 x) { return __bfloat162float(x); }
+__device__ __forceinline__ void set_ninf(float& x) { x = -INFINITY; }
+__device__ __forceinline__ void set_ninf(__half& x) { x = __float2half(-INFINITY); }
+__device__ __forceinline__ void set_ninf(__nv_bfloat16& x) { x = __float2bfloat16(-INFINITY); }
+__device__ __forceinline__ unsigned triple32(unsigned x) {
+  x ^= x >> 17; x *= 0xed5ad4bbU; x ^= x >> 11; x *= 0xac4c1b51U;
+  x ^= x >> 15; x *= 0x31848babU; x ^= x >> 14; return x;
+}
+// Portable reduction operators: CUDA 13's CUB dropped cub::Max()/cub::Sum().
+struct MaxOp { __device__ __forceinline__ float operator()(const float& a, const float& b) const { return fmaxf(a, b); } };
+// bit == 1 means allowed; disallowed logits become -inf. grid = (xtiles, B).
+template <typename T>
+__global__ void bitmask_kernel(T* logits, const int32_t* mask, int V, int mask_stride) {
+  int row = blockIdx.y;
+  const int32_t* m = mask + (long)row * mask_stride;
+  T* lg = logits + (long)row * V;
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < V; i += gridDim.x * blockDim.x) {
+    if (!((m[i >> 5] >> (i & 31)) & 1)) set_ninf(lg[i]);
+  }
+}
+// probs[row] = softmax(logits[row] / max(temp,eps)); one block per row.
+template <typename T>
+__global__ void softmax_temp_kernel(const T* __restrict__ logits, const float* __restrict__ temp,
+                                    float* __restrict__ probs, int V) {
+  int row = blockIdx.x;
+  const T* lg = logits + (long)row * V;
+  float* pr = probs + (long)row * V;
+  float t = temp[row];
+  float scale = (t > 0.f) ? (1.f / t) : 1.f;
+  typedef cub::BlockReduce<float, TPB> BR;
+  __shared__ typename BR::TempStorage ts;
+  __shared__ float sh;
+  float m = -INFINITY;
+  for (int i = threadIdx.x; i < V; i += blockDim.x) m = fmaxf(m, to_f(lg[i]));
+  float gmax = BR(ts).Reduce(m, MaxOp());
+  if (threadIdx.x == 0) sh = gmax;
+  __syncthreads();
+  gmax = sh;
+  __syncthreads();
+  float s = 0.f;
+  for (int i = threadIdx.x; i < V; i += blockDim.x) {
+    float e = __expf((to_f(lg[i]) - gmax) * scale);
+    pr[i] = e;
+    s += e;
+  }
+  float gsum = BR(ts).Sum(s);
+  if (threadIdx.x == 0) sh = gsum;
+  __syncthreads();
+  float inv = 1.f / sh;
+  __syncthreads();
+  for (int i = threadIdx.x; i < V; i += blockDim.x) pr[i] *= inv;
+}
+// sp/si: probabilities and token ids sorted descending per row. Compute the
+// kept prefix length L = min(top_k, top_p, min_p cutoffs), then inverse-CDF draw.
+__global__ void cutoff_sample_kernel(const float* __restrict__ sp, const int* __restrict__ si,
+                                     const float* __restrict__ top_p, const int* __restrict__ top_k,
+                                     const float* __restrict__ min_p, const float* __restrict__ temp,
+                                     unsigned seed, unsigned offset, int64_t* __restrict__ out, int V) {
+  int row = blockIdx.x;
+  if (threadIdx.x != 0) return;
+  const float* p = sp + (long)row * V;
+  const int* idx = si + (long)row * V;
+  float tp = top_p[row]; int tk = top_k[row]; float mp = min_p[row]; float t = temp[row];
+  long L = (tk > 0 && tk < V) ? (long)tk : (long)V;  // top-k
+  if (tp > 0.f && tp < 1.f) {                          // top-p (nucleus)
+    float c = 0.f; long Lp = L;
+    for (long j = 0; j < L; j++) { c += p[j]; if (c >= tp) { Lp = j + 1; break; } }
+    L = Lp;
+  }
+  if (mp > 0.f) {                                       // min-p
+    float thr = mp * p[0]; long j = 0;
+    while (j < L && p[j] >= thr) j++;
+    L = (j > 0) ? j : 1;
+  }
+  if (L < 1) L = 1;
+  if (t <= 0.f) L = 1;                                  // greedy
+  float sum = 0.f;
+  for (long j = 0; j < L; j++) sum += p[j];
+  unsigned h = triple32((unsigned)row * 0x9E3779B1u ^ seed ^ (offset * 0x85EBCA6Bu));
+  float u = (h >> 8) * (1.0f / 16777216.0f);            // [0, 1)
+  float target = u * sum, c = 0.f;
+  int pick = idx[L - 1];
+  for (long j = 0; j < L; j++) { c += p[j]; if (c >= target) { pick = idx[j]; break; } }
+  out[row] = pick;
+}
+}  // namespace
+void apply_token_bitmask_inplace(at::Tensor logits, at::Tensor bitmask) {
+  TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.is_contiguous(),
+              "logits must be contiguous [B, V] CUDA");
+  TORCH_CHECK(bitmask.is_cuda() && bitmask.scalar_type() == at::kInt,
+              "bitmask must be int32 CUDA");
+  const at::cuda::CUDAGuard guard(logits.device());
+  int B = logits.size(0), V = logits.size(1);
+  auto bm = bitmask.contiguous();
+  int ms = bm.size(-1);
+  TORCH_CHECK(ms >= (V + 31) / 32, "bitmask has too few words for V");
+  const int32_t* m = bm.data_ptr<int32_t>();
+  int gx = (V + TPB - 1) / TPB; if (gx > 512) gx = 512; if (gx < 1) gx = 1;
+  dim3 grid(gx, B);
+  cudaStream_t s = at::cuda::getCurrentCUDAStream();
+  auto st = logits.scalar_type();
+  if (st == at::kFloat)
+    bitmask_kernel<float><<<grid, TPB, 0, s>>>(logits.data_ptr<float>(), m, V, ms);
+  else if (st == at::kHalf)
+    bitmask_kernel<__half><<<grid, TPB, 0, s>>>((__half*)logits.data_ptr(), m, V, ms);
+  else if (st == at::kBFloat16)
+    bitmask_kernel<__nv_bfloat16><<<grid, TPB, 0, s>>>((__nv_bfloat16*)logits.data_ptr(), m, V, ms);
+  else
+    TORCH_CHECK(false, "logits must be float/half/bf16");
+}
+at::Tensor sample(at::Tensor logits, at::Tensor temperature, at::Tensor top_p,
+                  at::Tensor top_k, at::Tensor min_p, int64_t seed, int64_t offset) {
+  TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.is_contiguous(),
+              "logits must be contiguous [B, V] CUDA");
+  const at::cuda::CUDAGuard guard(logits.device());
+  int B = logits.size(0), V = logits.size(1);
+  cudaStream_t s = at::cuda::getCurrentCUDAStream();
+  auto fopt = logits.options().dtype(at::kFloat);
+  auto iopt = logits.options().dtype(at::kInt);
+  auto tf = temperature.to(at::kFloat).contiguous();
+  auto pf = top_p.to(at::kFloat).contiguous();
+  auto kf = top_k.to(at::kInt).contiguous();
+  auto mf = min_p.to(at::kFloat).contiguous();
+  auto probs = at::empty({B, V}, fopt);
+  auto st = logits.scalar_type();
+  if (st == at::kFloat)
+    softmax_temp_kernel<float><<<B, TPB, 0, s>>>(logits.data_ptr<float>(), tf.data_ptr<float>(), probs.data_ptr<float>(), V);
+  else if (st == at::kHalf)
+    softmax_temp_kernel<__half><<<B, TPB, 0, s>>>((const __half*)logits.data_ptr(), tf.data_ptr<float>(), probs.data_ptr<float>(), V);
+  else if (st == at::kBFloat16)
+    softmax_temp_kernel<__nv_bfloat16><<<B, TPB, 0, s>>>((const __nv_bfloat16*)logits.data_ptr(), tf.data_ptr<float>(), probs.data_ptr<float>(), V);
+  else
+    TORCH_CHECK(false, "logits must be float/half/bf16");
+  auto sorted_p = at::empty({B, V}, fopt);
+  auto sorted_i = at::empty({B, V}, iopt);
+  auto idx_in = at::arange(V, iopt).unsqueeze(0).repeat({B, 1}).contiguous();
+  auto offsets = at::arange(0, (long)(B + 1) * V, V, iopt).contiguous();
+  size_t tmp_bytes = 0;
+  cub::DeviceSegmentedRadixSort::SortPairsDescending(
+      nullptr, tmp_bytes, probs.data_ptr<float>(), sorted_p.data_ptr<float>(),
+      idx_in.data_ptr<int>(), sorted_i.data_ptr<int>(), B * V, B,
+      offsets.data_ptr<int>(), offsets.data_ptr<int>() + 1, 0, (int)sizeof(float) * 8, s);
+  auto tmp = at::empty({(long)tmp_bytes}, logits.options().dtype(at::kByte));
+  cub::DeviceSegmentedRadixSort::SortPairsDescending(
+      tmp.data_ptr(), tmp_bytes, probs.data_ptr<float>(), sorted_p.data_ptr<float>(),
+      idx_in.data_ptr<int>(), sorted_i.data_ptr<int>(), B * V, B,
+      offsets.data_ptr<int>(), offsets.data_ptr<int>() + 1, 0, (int)sizeof(float) * 8, s);
+  auto out = at::empty({B}, logits.options().dtype(at::kLong));
+  cutoff_sample_kernel<<<B, 32, 0, s>>>(sorted_p.data_ptr<float>(), sorted_i.data_ptr<int>(),
+      pf.data_ptr<float>(), kf.data_ptr<int>(), mf.data_ptr<float>(), tf.data_ptr<float>(),
+      (unsigned)seed, (unsigned)offset, out.data_ptr<int64_t>(), V);
+  return out;
+}

tests/test_logits_processor.py ADDED Viewed

	@@ -0,0 +1,71 @@

+import pytest
+import torch
+import kernels
+lp = kernels.get_kernel("phanerozoic/logits-processor", version=1, trust_remote_code=True)
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def _pack(allowed_sets, V, device):
+    nwords = (V + 31) // 32
+    mask = torch.zeros(len(allowed_sets), nwords, dtype=torch.int32, device=device)
+    for r, s in enumerate(allowed_sets):
+        for tok in s:
+            mask[r, tok >> 5] |= (1 << (tok & 31))
+    return mask
+@requires_cuda
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_bitmask_sets_disallowed_to_neg_inf(dtype):
+    V = 130
+    logits = torch.randn(2, V, device="cuda", dtype=dtype)
+    allowed = [{3, 7, 50, 129}, {0, 42, 63}]
+    ref = logits.clone()
+    lp.apply_token_bitmask(logits, _pack(allowed, V, "cuda"))
+    for r, s in enumerate(allowed):
+        for tok in range(V):
+            if tok in s:
+                assert logits[r, tok] == ref[r, tok]
+            else:
+                assert logits[r, tok].float() == float("-inf")
+@requires_cuda
+def test_greedy_and_topk1_are_argmax():
+    logits = torch.randn(8, 2000, device="cuda")
+    amax = logits.argmax(-1)
+    assert torch.equal(lp.sample(logits, temperature=0.0), amax)
+    assert torch.equal(lp.sample(logits, top_k=1, temperature=1.0), amax)
+@requires_cuda
+def test_min_p_and_topp_exclude_low_prob():
+    logits = torch.full((4, 50), -10.0, device="cuda")
+    logits[:, 7] = 10.0  # one dominant token, softmax ~1.0
+    assert torch.equal(lp.sample(logits, temperature=1.0, min_p=0.5, seed=1),
+                       torch.full((4,), 7, device="cuda"))
+    assert torch.equal(lp.sample(logits, temperature=1.0, top_p=0.9, seed=2),
+                       torch.full((4,), 7, device="cuda"))
+@requires_cuda
+def test_deterministic():
+    logits = torch.randn(16, 500, device="cuda")
+    a = lp.sample(logits, temperature=1.0, seed=99, offset=3)
+    b = lp.sample(logits, temperature=1.0, seed=99, offset=3)
+    assert torch.equal(a, b)
+@requires_cuda
+def test_distribution_matches_softmax():
+    torch.manual_seed(0)
+    V, N = 32, 40000
+    base = torch.randn(V, device="cuda")
+    logits = base.unsqueeze(0).repeat(N, 1)
+    toks = lp.sample(logits, temperature=1.0, seed=1234)
+    emp = torch.bincount(toks, minlength=V).float() / N
+    ref = torch.softmax(base, -1)
+    assert (emp - ref).abs().max().item() < 0.03

torch-ext/logits_processor/__init__.py ADDED Viewed

	@@ -0,0 +1,63 @@

+"""logits-processor: the generation last-mile, loadable through `kernels`.
+Two fused ops that every serving stack runs on the hot path:
+- ``apply_token_bitmask`` - structured / guided decoding. Set the logits of
+  disallowed tokens to -inf from a packed allow-mask (XGrammar convention:
+  bit 1 = allowed), so the next softmax cannot pick them. This is the kernel
+  behind constrained JSON and tool-calling output.
+- ``sample`` - temperature, then the intersection of top-k / top-p (nucleus) /
+  min-p filtering, then a multinomial draw, fused into one logits -> token op
+  with a counter-based RNG that is deterministic in ``(seed, row, offset)``.
+  ``temperature <= 0`` is greedy (argmax).
+Both take ``[B, V]`` logits in float / half / bf16. Sampling parameters are
+per-row: pass a scalar to broadcast, or a length-B tensor for a heterogeneous
+batch (continuous batching, per-request grammars).
+"""
+from typing import Union
+import torch
+from ._ops import ops
+__all__ = ["apply_token_bitmask", "sample"]
+Param = Union[float, int, torch.Tensor]
+def _row(x: Param, B: int, device, dtype) -> torch.Tensor:
+    if isinstance(x, torch.Tensor):
+        return x.to(device=device, dtype=dtype).contiguous()
+    return torch.full((B,), x, device=device, dtype=dtype)
+def apply_token_bitmask(logits: torch.Tensor, bitmask: torch.Tensor) -> torch.Tensor:
+    """In place: set disallowed tokens' logits to -inf.
+    logits: ``[B, V]`` (float/half/bf16). bitmask: ``[B, ceil(V/32)]`` int32,
+    bit 1 = token allowed. Returns ``logits`` for chaining.
+    """
+    ops.apply_token_bitmask_inplace(logits, bitmask.to(torch.int32).contiguous())
+    return logits
+def sample(logits: torch.Tensor, temperature: Param = 1.0, top_k: Param = 0,
+           top_p: Param = 1.0, min_p: Param = 0.0, seed: int = 0,
+           offset: int = 0) -> torch.Tensor:
+    """Fused temperature + top-k ∩ top-p ∩ min-p + multinomial sample.
+    top_k <= 0, top_p >= 1, and min_p <= 0 each disable that filter.
+    temperature <= 0 selects the argmax. Returns int64 token ids ``[B]``.
+    Same ``(seed, offset)`` and logits give the same tokens.
+    """
+    assert logits.dim() == 2, "logits must be [B, V]"
+    B, dev = logits.size(0), logits.device
+    return ops.sample(
+        logits.contiguous(),
+        _row(temperature, B, dev, torch.float32),
+        _row(top_p, B, dev, torch.float32),
+        _row(top_k, B, dev, torch.int32),
+        _row(min_p, B, dev, torch.float32),
+        int(seed), int(offset),
+    )

torch-ext/torch_binding.cpp ADDED Viewed

	@@ -0,0 +1,17 @@

+#include <torch/library.h>
+#include "registration.h"
+#include "torch_binding.h"
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("apply_token_bitmask_inplace(Tensor(a!) logits, Tensor bitmask) -> ()");
+  ops.def("sample(Tensor logits, Tensor temperature, Tensor top_p, Tensor top_k, "
+          "Tensor min_p, int seed, int offset) -> Tensor");
+#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
+  ops.impl("apply_token_bitmask_inplace", torch::kCUDA, &apply_token_bitmask_inplace);
+  ops.impl("sample", torch::kCUDA, &sample);
+#endif
+}
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h ADDED Viewed

	@@ -0,0 +1,17 @@

+#pragma once
+#include <torch/torch.h>
+// Set logits of disallowed tokens to -inf, in place. bitmask is packed int32,
+// bit == 1 means the token is allowed (XGrammar convention); one row of
+// ceil(V/32) words per batch element. logits: [B, V], bitmask: [B, ceil(V/32)].
+void apply_token_bitmask_inplace(at::Tensor logits, at::Tensor bitmask);
+// Fused logits -> token sampler. Applies per-row temperature, then the
+// intersection of top-k / top-p (nucleus) / min-p filtering, renormalizes the
+// surviving set, and draws one token per row with a counter-based RNG keyed by
+// (seed, row, offset). temperature <= 0 is greedy (argmax). logits: [B, V]
+// (float/half/bf16); temperature/top_p/min_p: [B] float; top_k: [B] int32.
+// Returns tokens: [B] int64.
+at::Tensor sample(at::Tensor logits, at::Tensor temperature, at::Tensor top_p,
+                  at::Tensor top_k, at::Tensor min_p, int64_t seed, int64_t offset);