chore: import upstream snapshot with attribution

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/__init__.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Sampling kernel entry points."""
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope
+from tokenspeed_kernel.selection import NoKernelFoundError, select_kernel
+from tokenspeed_kernel.signature import dense_tensor_format, format_signature
+
+__all__ = ["argmax"]
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+_SUPPORTED_OUT_DTYPES = (torch.int32, torch.int64)
+
+
+def _validate_argmax_out(logits: torch.Tensor, out: torch.Tensor) -> None:
+    if out.shape != (logits.shape[0],):
+        raise ValueError(
+            f"out must have shape (M,)={(logits.shape[0],)}, got {tuple(out.shape)}"
+        )
+    if out.dtype not in _SUPPORTED_OUT_DTYPES:
+        raise ValueError(f"out must be int32 or int64; got {out.dtype}")
+    if out.device != logits.device:
+        raise ValueError("out must be on the same device as logits")
+
+
+def _argmax_torch_fallback(
+    logits: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is not None:
+        _validate_argmax_out(logits, out)
+    result = torch.argmax(logits, dim=-1)
+    if out is not None:
+        out.copy_(result)
+        return out
+    return result
+
+
+def argmax(
+    logits: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+    solution: str | None = None,
+    override: str | None = None,
+) -> torch.Tensor:
+    """Return row-wise argmax indices over the last logits dimension.
+
+    Args:
+        logits: Input logits with shape ``(M, N)``. The argmax is taken
+            over the last dimension.
+        out: Optional output buffer with shape ``(M,)`` and dtype int32 or
+            int64 on the same device as ``logits``.
+        solution: Optional kernel solution to force through normal selection.
+        override: Optional exact kernel-name or solution override.
+
+    NaN handling:
+        Kernel-backed sampling paths treat NaNs as invalid candidates.
+        Rows with at least one non-NaN value return the index of the
+        maximum non-NaN value, with ties broken toward the lowest index.
+        Rows with no valid non-NaN values return the ``-1`` sentinel.
+        Unsupported inputs fall back to ``torch.argmax`` semantics.
+
+    Returns:
+        A tensor containing argmax indices for each row of ``logits``.
+    """
+    if (
+        logits.dim() != 2
+        or logits.shape[0] == 0
+        or not logits.is_cuda
+        or logits.dtype not in _SUPPORTED_DTYPES
+    ):
+        return _argmax_torch_fallback(logits, out=out)
+
+    signature = format_signature(logits=dense_tensor_format(logits.dtype))
+    try:
+        kernel = select_kernel(
+            "sampling",
+            "argmax",
+            signature,
+            solution=solution,
+            override=override,
+        )
+    except NoKernelFoundError:
+        return _argmax_torch_fallback(logits, out=out)
+
+    shape_params = {
+        "M": logits.shape[0],
+        "N": logits.shape[1],
+        "has_out": out is not None,
+    }
+    ShapeCapture.get().record(
+        "sampling", "argmax", kernel.name, logits.dtype, shape_params
+    )
+    with kernel_scope(
+        "sampling", "argmax", logits.dtype, kernel_name=kernel.name, **shape_params
+    ):
+        return kernel(logits, out=out)
+
+
+# Backend registration (side-effect imports).
+import tokenspeed_kernel.ops.sampling.cute_dsl  # noqa: E402,F401
+import tokenspeed_kernel.ops.sampling.gluon  # noqa: E402,F401

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/cuda.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""CUDA sampling kernels."""
+
+from tokenspeed_kernel.platform import current_platform
+from tokenspeed_kernel.registry import error_fn
+
+chain_speculative_sampling_target_only = error_fn
+# Prepare is a side-stream registration helper. On non-NVIDIA the fused
+# kernel doesn't exist, so callers should never reach the renorm path —
+# but prepare gets called unconditionally from FlashInfer backends'
+# __init__, so we keep it as a silent no-op rather than error_fn.
+fused_topk_topp_prepare = lambda *_args, **_kwargs: None  # noqa: E731
+fused_topk_topp_renorm = error_fn
+fused_topk_topp_workspace_size = error_fn
+verify_chain_greedy = error_fn
+
+if current_platform().is_nvidia:
+    try:
+        from tokenspeed_kernel.thirdparty.cuda.sampling_chain import (
+            chain_speculative_sampling_target_only,
+            verify_chain_greedy,
+        )
+    except ImportError:
+        pass
+
+    try:
+        from tokenspeed_kernel.thirdparty.cuda.fused_topk_topp import (
+            fused_topk_topp_renorm,
+            fused_topk_topp_workspace_size,
+        )
+        from tokenspeed_kernel.thirdparty.cuda.fused_topk_topp import (
+            prepare_for_device as fused_topk_topp_prepare,
+        )
+    except ImportError:
+        pass
+
+__all__ = [
+    "chain_speculative_sampling_target_only",
+    "fused_topk_topp_prepare",
+    "fused_topk_topp_renorm",
+    "fused_topk_topp_workspace_size",
+    "verify_chain_greedy",
+]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/flashinfer.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""FlashInfer sampling kernels."""
+
+from tokenspeed_kernel.platform import current_platform
+from tokenspeed_kernel.registry import error_fn
+
+min_p_sampling_from_probs = error_fn
+softmax = error_fn
+top_k_renorm_prob = error_fn
+top_k_top_p_sampling_from_probs = error_fn
+top_k_top_p_sampling_from_logits = error_fn
+top_p_renorm_prob = error_fn
+top_p_renorm_probs = error_fn
+
+if current_platform().is_nvidia:
+    try:
+        from flashinfer.sampling import (
+            min_p_sampling_from_probs,
+            top_k_renorm_prob,
+            top_k_top_p_sampling_from_logits,
+            top_k_top_p_sampling_from_probs,
+            top_p_renorm_prob,
+            top_p_renorm_probs,
+        )
+    except ImportError:
+        pass
+
+    try:
+        from tokenspeed_kernel.thirdparty.cuda.flashinfer_softmax import softmax
+    except ImportError:
+        pass
+
+__all__ = [
+    "min_p_sampling_from_probs",
+    "softmax",
+    "top_k_renorm_prob",
+    "top_k_top_p_sampling_from_logits",
+    "top_k_top_p_sampling_from_probs",
+    "top_p_renorm_prob",
+    "top_p_renorm_probs",
+]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/gluon/__init__.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Registration shims for AMD Gluon sampling kernels."""
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel.platform import ArchVersion, CapabilityRequirement
+from tokenspeed_kernel.registry import Priority, register_kernel
+from tokenspeed_kernel.signature import format_signatures
+
+try:
+    from tokenspeed_kernel_amd.ops.sampling.gluon.argmax_gfx950 import (
+        gluon_argmax_gfx950 as _argmax_impl,
+    )
+except ImportError as exc:
+    _IMPORT_ERROR = exc
+    _argmax_impl = None
+else:
+    _IMPORT_ERROR = None
+
+
+if _argmax_impl is not None:
+
+    @register_kernel(
+        "sampling",
+        "argmax",
+        name="gluon_argmax_gfx950",
+        solution="gluon",
+        capability=CapabilityRequirement(
+            min_arch_version=ArchVersion(9, 5),
+            max_arch_version=ArchVersion(9, 5),
+            vendors=frozenset({"amd"}),
+        ),
+        signatures=format_signatures(
+            "logits", "dense", {torch.float16, torch.bfloat16, torch.float32}
+        ),
+        priority=Priority.SPECIALIZED,
+        tags={"latency", "throughput"},
+    )
+    def gluon_argmax_gfx950(
+        logits: torch.Tensor,
+        *,
+        out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return _argmax_impl(logits, out=out)
+
+else:
+
+    def gluon_argmax_gfx950(
+        logits: torch.Tensor,
+        *,
+        out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        raise ImportError(
+            "gluon_argmax_gfx950 requires tokenspeed-kernel-amd"
+        ) from _IMPORT_ERROR
+
+
+__all__ = ["gluon_argmax_gfx950"]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/__init__.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Triton sampling kernel entry points."""
+
+from .common import gather_and_expand_scalars
+from .generic import gumbel_sample_from_pools_generic
+from .gumbel import (
+    gumbel_sample_from_pools,
+    gumbel_sample_from_pools_compact,
+)
+from .logprobs import selected_token_logprobs
+from .min_p import (
+    gumbel_sample_min_p_from_pools,
+    gumbel_sample_min_p_from_pools_parallel,
+)
+from .penalties import accumulate_counts_inplace, apply_penalties_logit_bias_inplace
+from .probability import min_p_renorm_prob
+from .topk_topp import (
+    _QRITA_PERCENTILE_TO_STD_TABLE,
+    gumbel_sample_top_k_top_p_from_pools,
+    gumbel_sample_top_k_top_p_qrita_from_pools,
+)
+from .topp import gumbel_sample_top_p_parallel_from_pools
+from .verify import verify_chain_target_sampled
+
+__all__ = [
+    "_QRITA_PERCENTILE_TO_STD_TABLE",
+    "gather_and_expand_scalars",
+    "gumbel_sample_from_pools",
+    "gumbel_sample_from_pools_compact",
+    "gumbel_sample_min_p_from_pools",
+    "gumbel_sample_min_p_from_pools_parallel",
+    "gumbel_sample_from_pools_generic",
+    "gumbel_sample_top_p_parallel_from_pools",
+    "gumbel_sample_top_k_top_p_from_pools",
+    "gumbel_sample_top_k_top_p_qrita_from_pools",
+    "min_p_renorm_prob",
+    "apply_penalties_logit_bias_inplace",
+    "accumulate_counts_inplace",
+    "selected_token_logprobs",
+    "verify_chain_target_sampled",
+]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/common.py

@@ -0,0 +1,179 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Triton sampling helper kernels."""
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+
+@triton.jit
+def _gather_and_expand_scalars_kernel(
+    index_ptr,
+    temperature_ptr,
+    top_k_ptr,
+    top_p_ptr,
+    min_p_ptr,
+    seed_ptr,
+    offsets_ptr,
+    out_temperature_ptr,
+    out_top_k_ptr,
+    out_top_p_ptr,
+    out_min_p_ptr,
+    out_seed_ptr,
+    out_offsets_ptr,
+    n: tl.constexpr,
+    N_BLOCK: tl.constexpr,
+    ENABLE_PDL: tl.constexpr,
+):
+    # PDL: wait for producer (e.g., penalty kernel writing into pools) to drain.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    bi = tl.program_id(0)
+    idx = tl.load(index_ptr + bi)
+
+    t = tl.load(temperature_ptr + idx)
+    k = tl.load(top_k_ptr + idx)
+    p = tl.load(top_p_ptr + idx)
+    if min_p_ptr is not None:
+        mp = tl.load(min_p_ptr + idx)
+    if seed_ptr is not None:
+        s = tl.load(seed_ptr + idx)
+    if offsets_ptr is not None:
+        # Cast int32 valid_cache_lengths to int64 for flashinfer's offset arg.
+        o = tl.load(offsets_ptr + idx).to(tl.int64)
+
+    n_off = tl.arange(0, N_BLOCK)
+    mask = n_off < n
+    base = bi * n
+
+    tl.store(out_temperature_ptr + base + n_off, t, mask=mask)
+    tl.store(out_top_k_ptr + base + n_off, k, mask=mask)
+    tl.store(out_top_p_ptr + base + n_off, p, mask=mask)
+    if out_min_p_ptr is not None:
+        tl.store(out_min_p_ptr + base + n_off, mp, mask=mask)
+    if out_seed_ptr is not None:
+        tl.store(out_seed_ptr + base + n_off, s, mask=mask)
+    if out_offsets_ptr is not None:
+        tl.store(out_offsets_ptr + base + n_off, o, mask=mask)
+
+    # PDL: signal that dependents (e.g., flashinfer softmax) can begin preamble.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def gather_and_expand_scalars(
+    index: torch.Tensor,
+    *,
+    temperature: torch.Tensor,
+    top_k: torch.Tensor,
+    top_p: torch.Tensor,
+    min_p: torch.Tensor | None = None,
+    seed: torch.Tensor | None = None,
+    offsets: torch.Tensor | None = None,
+    n: int = 1,
+    enable_pdl: bool = False,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor | None,
+]:
+    """Fused gather-and-broadcast for per-request sampling scalars.
+
+    Replaces the pattern ``index_select(pool, index)`` followed by
+    ``repeat_interleave(..., n)`` across up to six streams with one Triton
+    launch. ``offsets`` (int32) is cast to int64 inside the kernel.
+
+    Optional streams (min_p, seed, offsets) pass through as ``None`` — Triton
+    specializes the kernel on pointer-None-ness at JIT time and the gated
+    load/store paths are dead-code-eliminated.
+
+    Args:
+        ...
+        enable_pdl: opt into Programmatic Dependent Launch (Hopper+). Lets the
+            downstream flashinfer softmax/renorm kernels start their preamble
+            while our writes drain.
+
+    Returns ``(temperatures, top_ks, top_ps, min_ps_or_None, seeds_or_None,
+    offsets_or_None)``, each shape ``[bs * n]`` (or ``None`` when the
+    corresponding pool was omitted).
+    """
+    bs = index.size(0)
+    total = bs * n
+    device = index.device
+
+    out_temperature = torch.empty(total, dtype=temperature.dtype, device=device)
+    out_top_k = torch.empty(total, dtype=top_k.dtype, device=device)
+    out_top_p = torch.empty(total, dtype=top_p.dtype, device=device)
+    out_min_p = (
+        torch.empty(total, dtype=min_p.dtype, device=device)
+        if min_p is not None
+        else None
+    )
+    out_seed = (
+        torch.empty(total, dtype=seed.dtype, device=device)
+        if seed is not None
+        else None
+    )
+    out_offsets = (
+        torch.empty(total, dtype=torch.int64, device=device)
+        if offsets is not None
+        else None
+    )
+
+    if bs == 0:
+        return (
+            out_temperature,
+            out_top_k,
+            out_top_p,
+            out_min_p,
+            out_seed,
+            out_offsets,
+        )
+
+    extra_kwargs = {"launch_pdl": True} if enable_pdl else {}
+    _gather_and_expand_scalars_kernel[(bs,)](
+        index,
+        temperature,
+        top_k,
+        top_p,
+        min_p,
+        seed,
+        offsets,
+        out_temperature,
+        out_top_k,
+        out_top_p,
+        out_min_p,
+        out_seed,
+        out_offsets,
+        n=n,
+        N_BLOCK=triton.next_power_of_2(max(n, 1)),
+        ENABLE_PDL=enable_pdl,
+        num_warps=1,
+        **extra_kwargs,
+    )
+
+    return out_temperature, out_top_k, out_top_p, out_min_p, out_seed, out_offsets

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/generic.py

@@ -0,0 +1,321 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# TokenSpeed-specific mixed-parameter route.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+_GENERIC_GUMBEL_BLOCK_SIZE = 2048
+_GENERIC_GUMBEL_TOP_K_PAD = 128
+_GENERIC_GUMBEL_NUM_ATTEMPTS = 8
+_TOP_K_DISABLED = 1 << 30
+
+
+@triton.jit
+def _gumbel_sample_generic_pool_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    top_k_pool_ptr,
+    top_p_pool_ptr,
+    min_p_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    out_ptr,
+    logits_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TOPK_PAD: tl.constexpr,
+    NUM_ATTEMPTS: tl.constexpr,
+    TOP_K_DISABLED: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    token_offsets = tl.arange(0, BLOCK_SIZE)
+    rank_offsets = tl.arange(0, TOPK_PAD)
+
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    top_k_raw = tl.load(top_k_pool_ptr + pool_idx)
+    top_p = tl.load(top_p_pool_ptr + pool_idx).to(tl.float32)
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+
+    top_k_disabled = top_k_raw == TOP_K_DISABLED
+    top_k = tl.minimum(tl.maximum(top_k_raw, 1), TOPK_PAD)
+    top_k = tl.minimum(top_k, vocab_size)
+    min_p = tl.full((), 0.0, tl.float32)
+    if min_p_pool_ptr is not None:
+        min_p = tl.load(min_p_pool_ptr + pool_idx).to(tl.float32)
+    min_p_log_threshold = tl.log(tl.maximum(min_p, 1.0e-20))
+
+    row_max = tl.full((), float("-inf"), tl.float32)
+    row_argmax = tl.full((), 2147483647, tl.int32)
+
+    top_vals = tl.full((TOPK_PAD,), float("-inf"), tl.float32)
+    top_ids = tl.full((TOPK_PAD,), 2147483647, tl.int32)
+
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + token_offsets
+        mask = cols < vocab_size
+        vals = tl.load(
+            logits_ptr + row * logits_row_stride + cols,
+            mask=mask,
+            other=float("-inf"),
+        ).to(tl.float32)
+        vals = vals / temperature
+
+        block_max = tl.max(vals, axis=0)
+        block_argmax = tl.min(tl.where(vals == block_max, cols, 2147483647), axis=0)
+        better_max = (block_max > row_max) | (
+            (block_max == row_max) & (block_argmax < row_argmax)
+        )
+        row_max = tl.where(better_max, block_max, row_max)
+        row_argmax = tl.where(better_max, block_argmax, row_argmax)
+
+        block_top_vals = tl.topk(vals, TOPK_PAD)
+        remaining = vals
+        for block_rank in tl.static_range(0, TOPK_PAD):
+            cand_val = tl.max(
+                tl.where(rank_offsets == block_rank, block_top_vals, float("-inf")),
+                axis=0,
+            )
+            cand_id = tl.min(
+                tl.where(mask & (remaining == cand_val), cols, 2147483647),
+                axis=0,
+            )
+
+            worst_val = tl.min(top_vals, axis=0)
+            worst_id = tl.max(
+                tl.where(top_vals == worst_val, top_ids, -1),
+                axis=0,
+            )
+            worst_pos = tl.min(
+                tl.where(
+                    (top_vals == worst_val) & (top_ids == worst_id),
+                    rank_offsets,
+                    TOPK_PAD,
+                ),
+                axis=0,
+            )
+            better = (cand_val > worst_val) | (
+                (cand_val == worst_val) & (cand_id < worst_id)
+            )
+            replace = (rank_offsets == worst_pos) & better
+            top_vals = tl.where(replace, cand_val, top_vals)
+            top_ids = tl.where(replace, cand_id, top_ids)
+            remaining = tl.where(cols == cand_id, float("-inf"), remaining)
+
+    sorted_vals = tl.full((TOPK_PAD,), float("-inf"), tl.float32)
+    sorted_ids = tl.full((TOPK_PAD,), 2147483647, tl.int32)
+    work_vals = top_vals
+    work_ids = top_ids
+    for rank in tl.static_range(0, TOPK_PAD):
+        best_val = tl.max(work_vals, axis=0)
+        best_id = tl.min(tl.where(work_vals == best_val, work_ids, 2147483647), axis=0)
+        active_rank = rank < top_k
+        sorted_vals = tl.where(
+            (rank_offsets == rank) & active_rank,
+            best_val,
+            sorted_vals,
+        )
+        sorted_ids = tl.where(
+            (rank_offsets == rank) & active_rank,
+            best_id,
+            sorted_ids,
+        )
+        work_vals = tl.where(work_ids == best_id, float("-inf"), work_vals)
+
+    top_max = tl.max(sorted_vals, axis=0)
+    top_weights = tl.exp(sorted_vals - top_max)
+    top_weights = tl.where(rank_offsets < top_k, top_weights, 0.0)
+    top_denom = tl.maximum(tl.sum(top_weights, axis=0), 1.0e-20)
+    top_probs = top_weights / top_denom
+    top_cumulative_before = tl.cumsum(top_probs) - top_probs
+    top_keep = (
+        (rank_offsets < top_k)
+        & (top_cumulative_before < top_p)
+        & (sorted_vals >= top_max + min_p_log_threshold)
+    )
+
+    finite_seed = tl.randint(seed, offset)
+    finite_rand_offsets = tl.where(top_keep, sorted_ids, 0)
+    finite_uniform = tl.maximum(tl.rand(finite_seed, finite_rand_offsets), 1.0e-7)
+    finite_gumbel = -tl.log(-tl.log(finite_uniform))
+    finite_scores = tl.where(top_keep, sorted_vals + finite_gumbel, float("-inf"))
+    finite_score = tl.max(finite_scores, axis=0)
+    finite_token = tl.min(
+        tl.where(finite_scores == finite_score, sorted_ids, 2147483647),
+        axis=0,
+    )
+
+    disabled_token = tl.full((), 2147483647, tl.int32)
+    disabled_found = tl.full((), 0, tl.int32)
+    for attempt in tl.static_range(0, NUM_ATTEMPTS):
+        attempt_seed = tl.randint(seed, offset + attempt)
+        best_score = tl.full((), float("-inf"), tl.float32)
+        best_id = tl.full((), 2147483647, tl.int32)
+        best_logit = tl.full((), float("-inf"), tl.float32)
+
+        for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+            cols = start + token_offsets
+            mask = cols < vocab_size
+            vals = tl.load(
+                logits_ptr + row * logits_row_stride + cols,
+                mask=mask,
+                other=float("-inf"),
+            ).to(tl.float32)
+            vals = vals / temperature
+            uniform = tl.maximum(tl.rand(attempt_seed, cols), 1.0e-7)
+            gumbel = -tl.log(-tl.log(uniform))
+            scores = tl.where(mask, vals + gumbel, float("-inf"))
+            block_score = tl.max(scores, axis=0)
+            block_id = tl.min(tl.where(scores == block_score, cols, 2147483647), axis=0)
+            block_logit = tl.max(
+                tl.where(cols == block_id, vals, float("-inf")), axis=0
+            )
+            better = (block_score > best_score) | (
+                (block_score == best_score) & (block_id < best_id)
+            )
+            best_score = tl.where(better, block_score, best_score)
+            best_id = tl.where(better, block_id, best_id)
+            best_logit = tl.where(better, block_logit, best_logit)
+
+        total = tl.full((), 0.0, tl.float32)
+        before = tl.full((), 0.0, tl.float32)
+        for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+            cols = start + token_offsets
+            mask = cols < vocab_size
+            vals = tl.load(
+                logits_ptr + row * logits_row_stride + cols,
+                mask=mask,
+                other=float("-inf"),
+            ).to(tl.float32)
+            vals = vals / temperature
+            weights = tl.exp(vals - row_max)
+            weights = tl.where(mask, weights, 0.0)
+            total += tl.sum(weights, axis=0)
+            before_mask = (vals > best_logit) | (
+                (vals == best_logit) & (cols < best_id)
+            )
+            before += tl.sum(tl.where(mask & before_mask, weights, 0.0), axis=0)
+
+        min_p_accept = best_logit >= row_max + min_p_log_threshold
+        accepted = (before < (top_p * total)) & min_p_accept
+        take = (disabled_found == 0) & accepted
+        disabled_token = tl.where(take, best_id, disabled_token)
+        disabled_found = tl.where(accepted, 1, disabled_found)
+
+    disabled_token = tl.where(disabled_found != 0, disabled_token, row_argmax)
+    token = tl.where(top_k_disabled, disabled_token, finite_token)
+    tl.store(out_ptr + row, token)
+
+
+def gumbel_sample_from_pools_generic(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    top_k_pool: torch.Tensor,
+    top_p_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    min_p_pool: torch.Tensor | None = None,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Graph-safe Gumbel sampler for mixed top-k/top-p rows."""
+    if logits.ndim != 2:
+        raise ValueError(f"gumbel_sample_from_pools_generic expects 2D logits")
+    if logits.device.type != "cuda":
+        raise ValueError("gumbel_sample_from_pools_generic requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            "gumbel_sample_from_pools_generic requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    rows, vocab_size = logits.shape
+    if vocab_size <= 0:
+        raise ValueError("gumbel_sample_from_pools_generic requires non-empty vocab")
+    if num_tokens_per_req <= 0:
+        raise ValueError("num_tokens_per_req must be positive")
+    if rows % num_tokens_per_req != 0:
+        raise ValueError(
+            "logits rows must be divisible by num_tokens_per_req, "
+            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
+        )
+    request_rows = rows // num_tokens_per_req
+    if req_pool_indices.shape[0] != request_rows:
+        raise ValueError(
+            "req_pool_indices length must match request rows, "
+            f"got {req_pool_indices.shape[0]} and {request_rows}"
+        )
+    if req_pool_indices.dtype != torch.int32:
+        raise ValueError(
+            f"req_pool_indices must be int32, got {req_pool_indices.dtype}"
+        )
+    if top_k_pool.dtype != torch.int32:
+        raise ValueError(f"top_k_pool must be int32, got {top_k_pool.dtype}")
+    if seed_pool.dtype != torch.int64:
+        raise ValueError(f"seed_pool must be int64, got {seed_pool.dtype}")
+    if out.dtype != torch.int32:
+        raise ValueError(f"out must be int32, got {out.dtype}")
+    if out.shape[0] < rows:
+        raise ValueError(f"out is too small: {out.shape[0]} < {rows}")
+    if min_p_pool is not None:
+        if min_p_pool.device.type != "cuda":
+            raise ValueError("min_p_pool must be CUDA")
+        if min_p_pool.ndim != 1:
+            raise ValueError(f"min_p_pool must be 1D, got {min_p_pool.ndim}D")
+    if rows == 0:
+        return out[:0]
+
+    _gumbel_sample_generic_pool_kernel[(rows,)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        top_k_pool,
+        top_p_pool,
+        min_p_pool,
+        seed_pool,
+        offsets_pool,
+        out,
+        logits_row_stride=logits.stride(0),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_GENERIC_GUMBEL_BLOCK_SIZE,
+        TOPK_PAD=_GENERIC_GUMBEL_TOP_K_PAD,
+        NUM_ATTEMPTS=_GENERIC_GUMBEL_NUM_ATTEMPTS,
+        TOP_K_DISABLED=_TOP_K_DISABLED,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=8,
+        num_stages=3,
+    )
+    return out[:rows]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/gumbel.py

@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# TokenSpeed-specific pool state, scratch ownership, and tie-breaking.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+_GUMBEL_BLOCK_SIZE = 1024
+_GUMBEL_COMPACT_BLOCK_SIZE = 2048
+
+
+@triton.jit
+def _gumbel_sample_pool_stage1_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    local_ids_ptr,
+    local_scores_ptr,
+    logits_row_stride: tl.constexpr,
+    local_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    block_idx = tl.program_id(1)
+    token_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = token_offsets < vocab_size
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+
+    logits = tl.load(
+        logits_ptr + row * logits_row_stride + token_offsets,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+    gumbel_seed = tl.randint(seed, offset)
+    uniform = tl.maximum(tl.rand(gumbel_seed, token_offsets), 1.0e-7)
+    gumbel = -tl.log(-tl.log(uniform))
+    scores = tl.where(mask, logits / temperature + gumbel, float("-inf"))
+
+    max_score = tl.max(scores, axis=0)
+    token_id = tl.min(
+        tl.where(scores == max_score, token_offsets, vocab_size + BLOCK_SIZE),
+        axis=0,
+    )
+
+    tl.store(local_ids_ptr + row * local_row_stride + block_idx, token_id)
+    tl.store(local_scores_ptr + row * local_row_stride + block_idx, max_score)
+
+
+@triton.jit
+def _gumbel_sample_stage2_kernel(
+    local_ids_ptr,
+    local_scores_ptr,
+    out_ptr,
+    local_row_stride: tl.constexpr,
+    num_blocks: tl.constexpr,
+    NUM_BLOCKS_PAD: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block_offsets = tl.arange(0, NUM_BLOCKS_PAD)
+    mask = block_offsets < num_blocks
+
+    scores = tl.load(
+        local_scores_ptr + row * local_row_stride + block_offsets,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    ids = tl.load(
+        local_ids_ptr + row * local_row_stride + block_offsets,
+        mask=mask,
+        other=2147483647,
+    )
+
+    max_score = tl.max(scores, axis=0)
+    token_id = tl.min(tl.where(scores == max_score, ids, 2147483647), axis=0)
+    tl.store(out_ptr + row, token_id)
+
+
+@triton.jit
+def _gumbel_sample_compact_pool_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    out_ptr,
+    logits_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    token_offsets = tl.arange(0, BLOCK_SIZE)
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+    gumbel_seed = tl.randint(seed, offset)
+
+    best_score = tl.full((), float("-inf"), tl.float32)
+    best_id = tl.full((), 2147483647, tl.int32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + token_offsets
+        mask = cols < vocab_size
+        logits = tl.load(
+            logits_ptr + row * logits_row_stride + cols,
+            mask=mask,
+            other=float("-inf"),
+        ).to(tl.float32)
+        uniform = tl.maximum(tl.rand(gumbel_seed, cols), 1.0e-7)
+        gumbel = -tl.log(-tl.log(uniform))
+        scores = tl.where(mask, logits / temperature + gumbel, float("-inf"))
+
+        block_score = tl.max(scores, axis=0)
+        block_id = tl.min(tl.where(scores == block_score, cols, 2147483647), axis=0)
+        better = (block_score > best_score) | (
+            (block_score == best_score) & (block_id < best_id)
+        )
+        best_score = tl.where(better, block_score, best_score)
+        best_id = tl.where(better, block_id, best_id)
+
+    tl.store(out_ptr + row, best_id)
+
+
+def _check_gumbel_pool_inputs(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    local_ids: torch.Tensor,
+    local_scores: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    fn_name: str,
+    block_size: int,
+    num_tokens_per_req: int,
+) -> tuple[int, int, int]:
+    if logits.ndim != 2:
+        raise ValueError(f"{fn_name} expects 2D logits, got {logits.ndim}D")
+    if logits.device.type != "cuda":
+        raise ValueError(f"{fn_name} requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            f"{fn_name} requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    rows, vocab_size = logits.shape
+    if vocab_size <= 0:
+        raise ValueError(f"{fn_name} requires non-empty vocab dimension")
+
+    for name, tensor, ndim in (
+        ("req_pool_indices", req_pool_indices, 1),
+        ("temperature_pool", temperature_pool, 1),
+        ("seed_pool", seed_pool, 1),
+        ("offsets_pool", offsets_pool, 1),
+        ("out", out, 1),
+    ):
+        if tensor.device.type != "cuda":
+            raise ValueError(f"{name} must be CUDA")
+        if tensor.ndim != ndim:
+            raise ValueError(f"{name} must be {ndim}D, got {tensor.ndim}D")
+    if num_tokens_per_req <= 0:
+        raise ValueError("num_tokens_per_req must be positive")
+    if rows % num_tokens_per_req != 0:
+        raise ValueError(
+            "logits rows must be divisible by num_tokens_per_req, "
+            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
+        )
+    request_rows = rows // num_tokens_per_req
+    if req_pool_indices.shape[0] != request_rows:
+        raise ValueError(
+            "req_pool_indices length must match request rows, "
+            f"got {req_pool_indices.shape[0]} and {request_rows}"
+        )
+    if out.shape[0] < rows:
+        raise ValueError(f"out is too small: {out.shape[0]} < {rows}")
+    if req_pool_indices.dtype != torch.int32:
+        raise ValueError(
+            f"req_pool_indices must be int32, got {req_pool_indices.dtype}"
+        )
+    if seed_pool.dtype != torch.int64:
+        raise ValueError(f"seed_pool must be int64, got {seed_pool.dtype}")
+
+    num_blocks = triton.cdiv(vocab_size, block_size)
+    if local_ids.device.type != "cuda" or local_scores.device.type != "cuda":
+        raise ValueError("gumbel pool scratch tensors must be CUDA")
+    if local_ids.ndim != 2 or local_scores.ndim != 2:
+        raise ValueError("gumbel pool scratch tensors must be 2D")
+    if local_ids.shape[0] < rows or local_scores.shape[0] < rows:
+        raise ValueError("gumbel pool scratch tensors have too few rows")
+    if local_ids.shape[1] < num_blocks or local_scores.shape[1] < num_blocks:
+        raise ValueError(
+            "gumbel pool scratch tensors have too few blocks: "
+            f"need {num_blocks}, got {local_ids.shape[1]} / {local_scores.shape[1]}"
+        )
+    if local_ids.dtype != torch.int32:
+        raise ValueError(f"local_ids must be int32, got {local_ids.dtype}")
+    if local_scores.dtype != torch.float32:
+        raise ValueError(f"local_scores must be float32, got {local_scores.dtype}")
+    if local_ids.stride(-1) != 1 or local_scores.stride(-1) != 1:
+        raise ValueError("gumbel pool scratch tensors require stride-1 block dimension")
+    return rows, vocab_size, num_blocks
+
+
+def gumbel_sample_from_pools(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    local_ids: torch.Tensor,
+    local_scores: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Sample token ids from logits with Gumbel-Max and pool-indexed scalars."""
+    rows, vocab_size, num_blocks = _check_gumbel_pool_inputs(
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        seed_pool,
+        offsets_pool,
+        local_ids,
+        local_scores,
+        out,
+        fn_name="gumbel_sample_from_pools",
+        block_size=_GUMBEL_BLOCK_SIZE,
+        num_tokens_per_req=num_tokens_per_req,
+    )
+    if rows == 0:
+        return out[:0]
+
+    _gumbel_sample_pool_stage1_kernel[(rows, num_blocks)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        seed_pool,
+        offsets_pool,
+        local_ids,
+        local_scores,
+        logits_row_stride=logits.stride(0),
+        local_row_stride=local_ids.stride(0),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_GUMBEL_BLOCK_SIZE,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+    )
+    _gumbel_sample_stage2_kernel[(rows,)](
+        local_ids,
+        local_scores,
+        out,
+        local_row_stride=local_ids.stride(0),
+        num_blocks=num_blocks,
+        NUM_BLOCKS_PAD=triton.next_power_of_2(num_blocks),
+        num_warps=1,
+    )
+    return out[:rows]
+
+
+def gumbel_sample_from_pools_compact(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    block_size: int = _GUMBEL_COMPACT_BLOCK_SIZE,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Single-kernel Gumbel-Max path for vocabularies that fit one scan loop."""
+    if logits.ndim != 2:
+        raise ValueError(
+            f"gumbel_sample_from_pools_compact expects 2D logits, got {logits.ndim}D"
+        )
+    if logits.device.type != "cuda":
+        raise ValueError("gumbel_sample_from_pools_compact requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            "gumbel_sample_from_pools_compact requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    rows, vocab_size = logits.shape
+    if vocab_size <= 0:
+        raise ValueError("gumbel_sample_from_pools_compact requires non-empty vocab")
+    if num_tokens_per_req <= 0:
+        raise ValueError("num_tokens_per_req must be positive")
+    if rows % num_tokens_per_req != 0:
+        raise ValueError(
+            "logits rows must be divisible by num_tokens_per_req, "
+            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
+        )
+    request_rows = rows // num_tokens_per_req
+    if req_pool_indices.shape[0] != request_rows:
+        raise ValueError(
+            "req_pool_indices length must match request rows, "
+            f"got {req_pool_indices.shape[0]} and {request_rows}"
+        )
+    if out.shape[0] < rows:
+        raise ValueError(f"out is too small: {out.shape[0]} < {rows}")
+    if rows == 0:
+        return out[:0]
+
+    _gumbel_sample_compact_pool_kernel[(rows,)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        seed_pool,
+        offsets_pool,
+        out,
+        logits_row_stride=logits.stride(0),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=block_size,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=8,
+        num_stages=3,
+    )
+    return out[:rows]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/logprobs.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Computes selected-token logprobs without full-vocabulary materialization.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+
+@triton.jit
+def _selected_token_logprobs_kernel(
+    logits_ptr,
+    tokens_ptr,
+    out_ptr,
+    vocab_size: tl.constexpr,
+    logits_row_stride: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    row_ptr = logits_ptr + row * logits_row_stride
+
+    row_max = tl.full((), float("-inf"), tl.float32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offsets
+        mask = cols < vocab_size
+        vals = tl.load(row_ptr + cols, mask=mask, other=float("-inf")).to(tl.float32)
+        row_max = tl.maximum(
+            row_max, tl.max(tl.where(mask, vals, float("-inf")), axis=0)
+        )
+
+    denom = tl.full((), 0.0, tl.float32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offsets
+        mask = cols < vocab_size
+        vals = tl.load(row_ptr + cols, mask=mask, other=float("-inf")).to(tl.float32)
+        weights = tl.exp(vals - row_max)
+        denom += tl.sum(tl.where(mask, weights, 0.0), axis=0)
+
+    token = tl.load(tokens_ptr + row).to(tl.int64)
+    selected = tl.load(row_ptr + token).to(tl.float32)
+    tl.store(out_ptr + row, selected - row_max - tl.log(tl.maximum(denom, 1.0e-20)))
+
+
+def selected_token_logprobs(
+    logits: torch.Tensor,
+    tokens: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Compute ``log_softmax(logits)[row, tokens[row]]`` without materializing it."""
+    if logits.ndim != 2:
+        raise ValueError(f"selected_token_logprobs expects 2D logits")
+    if logits.device.type != "cuda":
+        raise ValueError("selected_token_logprobs requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            "selected_token_logprobs requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    rows, vocab_size = logits.shape
+    if tokens.numel() != rows:
+        raise ValueError(
+            f"tokens length must match rows, got {tokens.numel()} and {rows}"
+        )
+    if tokens.dtype not in (torch.int32, torch.int64):
+        raise ValueError(f"tokens must be int32 or int64, got {tokens.dtype}")
+    if out is None:
+        out = torch.empty((rows,), dtype=torch.float32, device=logits.device)
+    if out.dtype != torch.float32:
+        raise ValueError(f"out must be float32, got {out.dtype}")
+    if out.shape[0] < rows:
+        raise ValueError(f"out is too small: {out.shape[0]} < {rows}")
+    if rows == 0:
+        return out[:0]
+
+    _selected_token_logprobs_kernel[(rows,)](
+        logits,
+        tokens.reshape(-1),
+        out,
+        vocab_size=vocab_size,
+        logits_row_stride=logits.stride(0),
+        BLOCK_SIZE=1024,
+        num_warps=4,
+        num_stages=3,
+    )
+    return out[:rows]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/min_p.py

@@ -0,0 +1,365 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# TokenSpeed-native min-p Gumbel kernels.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+from .gumbel import _check_gumbel_pool_inputs, _gumbel_sample_stage2_kernel
+
+_MIN_P_GUMBEL_BLOCK_SIZE = 1024
+_MIN_P_PARALLEL_BLOCK_SIZE = 1024
+
+
+@triton.jit
+def _gumbel_sample_min_p_pool_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    min_p_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    out_ptr,
+    logits_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    min_p = tl.load(min_p_pool_ptr + pool_idx).to(tl.float32)
+    min_p_log_threshold = tl.log(tl.maximum(min_p, 1.0e-20))
+
+    row_max = tl.full((), float("-inf"), tl.float32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offsets
+        mask = cols < vocab_size
+        vals = tl.load(
+            logits_ptr + row * logits_row_stride + cols,
+            mask=mask,
+            other=float("-inf"),
+        ).to(tl.float32)
+        vals = vals / temperature
+        row_max = tl.maximum(
+            row_max, tl.max(tl.where(mask, vals, float("-inf")), axis=0)
+        )
+
+    threshold = row_max + min_p_log_threshold
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    step_offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+    rng_seed = tl.randint(seed, step_offset)
+
+    best_score = tl.full((), float("-inf"), tl.float32)
+    best_id = tl.full((), 2147483647, tl.int32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offsets
+        mask = cols < vocab_size
+        vals = tl.load(
+            logits_ptr + row * logits_row_stride + cols,
+            mask=mask,
+            other=float("-inf"),
+        ).to(tl.float32)
+        vals = vals / temperature
+        keep = mask & (vals >= threshold)
+        uniform = tl.maximum(tl.rand(rng_seed, cols), 1.0e-7)
+        gumbel = -tl.log(-tl.log(uniform))
+        scores = tl.where(keep, vals + gumbel, float("-inf"))
+        block_score = tl.max(scores, axis=0)
+        block_id = tl.min(tl.where(scores == block_score, cols, 2147483647), axis=0)
+        better = (block_score > best_score) | (
+            (block_score == best_score) & (block_id < best_id)
+        )
+        best_score = tl.where(better, block_score, best_score)
+        best_id = tl.where(better, block_id, best_id)
+
+    tl.store(out_ptr + row, best_id)
+
+
+def gumbel_sample_min_p_from_pools(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    min_p_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Gumbel-Max sampler for no top-k/top-p rows with a min-p cutoff."""
+    if logits.ndim != 2:
+        raise ValueError(f"gumbel_sample_min_p_from_pools expects 2D logits")
+    if logits.device.type != "cuda":
+        raise ValueError("gumbel_sample_min_p_from_pools requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            "gumbel_sample_min_p_from_pools requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    if num_tokens_per_req <= 0:
+        raise ValueError("num_tokens_per_req must be positive")
+    rows, vocab_size = logits.shape
+    if rows % num_tokens_per_req != 0:
+        raise ValueError(
+            "logits rows must be divisible by num_tokens_per_req, "
+            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
+        )
+    request_rows = rows // num_tokens_per_req
+    if req_pool_indices.shape[0] != request_rows:
+        raise ValueError(
+            "req_pool_indices length must match request rows, "
+            f"got {req_pool_indices.shape[0]} and {request_rows}"
+        )
+    if req_pool_indices.dtype != torch.int32:
+        raise ValueError(
+            f"req_pool_indices must be int32, got {req_pool_indices.dtype}"
+        )
+    if min_p_pool.ndim != 1:
+        raise ValueError(f"min_p_pool must be 1D, got {min_p_pool.ndim}D")
+    if seed_pool.dtype != torch.int64:
+        raise ValueError(f"seed_pool must be int64, got {seed_pool.dtype}")
+    if out.dtype != torch.int32:
+        raise ValueError(f"out must be int32, got {out.dtype}")
+    if out.shape[0] < rows:
+        raise ValueError(f"out is too small: {out.shape[0]} < {rows}")
+    if rows == 0:
+        return out[:0]
+
+    _gumbel_sample_min_p_pool_kernel[(rows,)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        min_p_pool,
+        seed_pool,
+        offsets_pool,
+        out,
+        logits_row_stride=logits.stride(0),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_MIN_P_GUMBEL_BLOCK_SIZE,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+        num_stages=3,
+    )
+    return out[:rows]
+
+
+@triton.jit
+def _min_p_local_max_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    local_max_ptr,
+    logits_row_stride: tl.constexpr,
+    local_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block_idx = tl.program_id(1)
+    req_row = row // NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    cols = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = cols < vocab_size
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    vals = tl.load(
+        logits_ptr + row * logits_row_stride + cols,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    vals = vals / temperature
+    tl.store(local_max_ptr + row * local_row_stride + block_idx, tl.max(vals, axis=0))
+
+
+@triton.jit
+def _min_p_row_max_kernel(
+    local_max_ptr,
+    row_max_ptr,
+    local_row_stride: tl.constexpr,
+    num_blocks: tl.constexpr,
+    NUM_BLOCKS_PAD: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, NUM_BLOCKS_PAD)
+    mask = offsets < num_blocks
+    vals = tl.load(
+        local_max_ptr + row * local_row_stride + offsets,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    tl.store(row_max_ptr + row, tl.max(vals, axis=0))
+
+
+@triton.jit
+def _min_p_local_gumbel_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    min_p_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    row_max_ptr,
+    local_ids_ptr,
+    local_scores_ptr,
+    logits_row_stride: tl.constexpr,
+    local_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block_idx = tl.program_id(1)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    cols = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = cols < vocab_size
+
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    min_p = tl.load(min_p_pool_ptr + pool_idx).to(tl.float32)
+    threshold = tl.load(row_max_ptr + row).to(tl.float32) + tl.log(
+        tl.maximum(min_p, 1.0e-20)
+    )
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    step_offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+    rng_seed = tl.randint(seed, step_offset)
+
+    vals = tl.load(
+        logits_ptr + row * logits_row_stride + cols,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    vals = vals / temperature
+    keep = mask & (vals >= threshold)
+    uniform = tl.maximum(tl.rand(rng_seed, cols), 1.0e-7)
+    gumbel = -tl.log(-tl.log(uniform))
+    scores = tl.where(keep, vals + gumbel, float("-inf"))
+    block_score = tl.max(scores, axis=0)
+    block_id = tl.min(tl.where(scores == block_score, cols, 2147483647), axis=0)
+    tl.store(local_ids_ptr + row * local_row_stride + block_idx, block_id)
+    tl.store(local_scores_ptr + row * local_row_stride + block_idx, block_score)
+
+
+def gumbel_sample_min_p_from_pools_parallel(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    min_p_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    local_ids: torch.Tensor,
+    local_scores: torch.Tensor,
+    row_max: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Parallel large-vocab min-p Gumbel sampler."""
+    rows, vocab_size, num_blocks = _check_gumbel_pool_inputs(
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        seed_pool,
+        offsets_pool,
+        local_ids,
+        local_scores,
+        out,
+        fn_name="gumbel_sample_min_p_from_pools_parallel",
+        block_size=_MIN_P_PARALLEL_BLOCK_SIZE,
+        num_tokens_per_req=num_tokens_per_req,
+    )
+    if min_p_pool.device.type != "cuda":
+        raise ValueError("min_p_pool must be CUDA")
+    if min_p_pool.ndim != 1:
+        raise ValueError(f"min_p_pool must be 1D, got {min_p_pool.ndim}D")
+    if row_max.device.type != "cuda":
+        raise ValueError("row_max must be CUDA")
+    if row_max.dtype != torch.float32:
+        raise ValueError(f"row_max must be float32, got {row_max.dtype}")
+    if row_max.shape[0] < rows:
+        raise ValueError(f"row_max is too small: {row_max.shape[0]} < {rows}")
+    if rows == 0:
+        return out[:0]
+
+    _min_p_local_max_kernel[(rows, num_blocks)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        local_scores,
+        logits_row_stride=logits.stride(0),
+        local_row_stride=local_scores.stride(0),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_MIN_P_PARALLEL_BLOCK_SIZE,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+    )
+    _min_p_row_max_kernel[(rows,)](
+        local_scores,
+        row_max,
+        local_row_stride=local_scores.stride(0),
+        num_blocks=num_blocks,
+        NUM_BLOCKS_PAD=triton.next_power_of_2(num_blocks),
+        num_warps=1,
+    )
+    _min_p_local_gumbel_kernel[(rows, num_blocks)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        min_p_pool,
+        seed_pool,
+        offsets_pool,
+        row_max,
+        local_ids,
+        local_scores,
+        logits_row_stride=logits.stride(0),
+        local_row_stride=local_ids.stride(0),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_MIN_P_PARALLEL_BLOCK_SIZE,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+    )
+    _gumbel_sample_stage2_kernel[(rows,)](
+        local_ids,
+        local_scores,
+        out,
+        local_row_stride=local_ids.stride(0),
+        num_blocks=num_blocks,
+        NUM_BLOCKS_PAD=triton.next_power_of_2(num_blocks),
+        num_warps=1,
+    )
+    return out[:rows]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/penalties.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# TokenSpeed-owned penalty, logit-bias, and count kernels.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+
+@triton.jit
+def _apply_penalties_logit_bias_inplace_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    counts_ptr,
+    logit_bias_ptr,
+    freq_pen_pool_ptr,
+    pres_pen_pool_ptr,
+    rep_pen_pool_ptr,
+    vocab_size: tl.constexpr,
+    logits_row_stride: tl.constexpr,
+    counts_row_stride: tl.constexpr,
+    bias_row_stride: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block = tl.program_id(1)
+    req_row = row // NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    cols = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = cols < vocab_size
+
+    logits_offsets = row * logits_row_stride + cols
+    state_offsets = pool_idx * counts_row_stride + cols
+    bias_offsets = pool_idx * bias_row_stride + cols
+
+    vals = tl.load(logits_ptr + logits_offsets, mask=mask, other=0.0).to(tl.float32)
+    counts = tl.load(counts_ptr + state_offsets, mask=mask, other=0).to(tl.float32)
+    active = counts > 0.0
+
+    rep = tl.load(rep_pen_pool_ptr + pool_idx).to(tl.float32)
+    freq = tl.load(freq_pen_pool_ptr + pool_idx).to(tl.float32)
+    presence = tl.load(pres_pen_pool_ptr + pool_idx).to(tl.float32)
+
+    rep_vals = tl.where(vals > 0.0, vals / rep, vals * rep)
+    vals = tl.where(active, rep_vals, vals)
+    vals = vals - freq * counts - presence * active.to(tl.float32)
+    vals += tl.load(logit_bias_ptr + bias_offsets, mask=mask, other=0.0).to(tl.float32)
+
+    tl.store(logits_ptr + logits_offsets, vals, mask=mask)
+
+
+def apply_penalties_logit_bias_inplace(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    counts: torch.Tensor,
+    logit_bias: torch.Tensor,
+    freq_pen_pool: torch.Tensor,
+    pres_pen_pool: torch.Tensor,
+    rep_pen_pool: torch.Tensor,
+    *,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Apply repetition/frequency/presence penalties and logit_bias in-place."""
+    if logits.ndim != 2:
+        raise ValueError(f"logits must be 2D, got {logits.ndim}D")
+    if counts.ndim != 2 or logit_bias.ndim != 2:
+        raise ValueError("counts and logit_bias must be 2D")
+    if logits.device.type != "cuda":
+        raise ValueError("apply_penalties_logit_bias_inplace requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            "apply_penalties_logit_bias_inplace requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    if req_pool_indices.dtype != torch.int32:
+        raise ValueError(
+            f"req_pool_indices must be int32, got {req_pool_indices.dtype}"
+        )
+    if counts.dtype != torch.int32:
+        raise ValueError(f"counts must be int32, got {counts.dtype}")
+    if num_tokens_per_req <= 0:
+        raise ValueError("num_tokens_per_req must be positive")
+
+    rows, vocab_size = logits.shape
+    if rows % num_tokens_per_req != 0:
+        raise ValueError(
+            "logits rows must be divisible by num_tokens_per_req, "
+            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
+        )
+    request_rows = rows // num_tokens_per_req
+    if req_pool_indices.shape[0] != request_rows:
+        raise ValueError(
+            "req_pool_indices length must match request rows, "
+            f"got {req_pool_indices.shape[0]} and {request_rows}"
+        )
+    if counts.shape[1] < vocab_size or logit_bias.shape[1] < vocab_size:
+        raise ValueError(
+            "counts/logit_bias vocab dimension must cover logits vocab, "
+            f"got counts={counts.shape}, logit_bias={logit_bias.shape}, logits={logits.shape}"
+        )
+    if rows == 0:
+        return logits
+
+    num_blocks = triton.cdiv(vocab_size, 1024)
+    _apply_penalties_logit_bias_inplace_kernel[(rows, num_blocks)](
+        logits,
+        req_pool_indices,
+        counts,
+        logit_bias,
+        freq_pen_pool,
+        pres_pen_pool,
+        rep_pen_pool,
+        vocab_size=vocab_size,
+        logits_row_stride=logits.stride(0),
+        counts_row_stride=counts.stride(0),
+        bias_row_stride=logit_bias.stride(0),
+        BLOCK_SIZE=1024,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+        num_stages=3,
+    )
+    return logits
+
+
+@triton.jit
+def _accumulate_counts_inplace_kernel(
+    counts_ptr,
+    pool_idx_ptr,
+    tokens_ptr,
+    weights_ptr,
+    total: tl.constexpr,
+    counts_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < total
+    weights = tl.load(weights_ptr + offs, mask=mask, other=0).to(tl.int32)
+    pool_idx = tl.load(pool_idx_ptr + offs, mask=mask, other=0).to(tl.int64)
+    tokens = tl.load(tokens_ptr + offs, mask=mask, other=0).to(tl.int64)
+    valid = mask & (weights != 0) & (tokens >= 0) & (tokens < vocab_size)
+    tl.atomic_add(
+        counts_ptr + pool_idx * counts_row_stride + tokens,
+        weights,
+        sem="relaxed",
+        mask=valid,
+    )
+
+
+def accumulate_counts_inplace(
+    counts: torch.Tensor,
+    pool_idx: torch.Tensor,
+    tokens: torch.Tensor,
+    weights: torch.Tensor,
+) -> None:
+    """Graph-safe ``counts[pool_idx, tokens] += weights``."""
+    if counts.ndim != 2:
+        raise ValueError(f"counts must be 2D, got {counts.ndim}D")
+    if counts.device.type != "cuda":
+        raise ValueError("accumulate_counts_inplace requires CUDA counts")
+    if counts.dtype != torch.int32:
+        raise ValueError(f"counts must be int32, got {counts.dtype}")
+    if pool_idx.dtype != torch.int32:
+        raise ValueError(f"pool_idx must be int32, got {pool_idx.dtype}")
+    if weights.dtype != torch.int32:
+        raise ValueError(f"weights must be int32, got {weights.dtype}")
+    if tokens.dtype not in (torch.int32, torch.int64):
+        raise ValueError(f"tokens must be int32 or int64, got {tokens.dtype}")
+    total = int(tokens.numel())
+    if pool_idx.numel() != total or weights.numel() != total:
+        raise ValueError(
+            "pool_idx, tokens, and weights must have the same number of elements"
+        )
+    if total == 0:
+        return
+
+    _accumulate_counts_inplace_kernel[(triton.cdiv(total, 256),)](
+        counts,
+        pool_idx.reshape(-1),
+        tokens.reshape(-1),
+        weights.reshape(-1),
+        total=total,
+        counts_row_stride=counts.stride(0),
+        vocab_size=counts.shape[1],
+        BLOCK_SIZE=256,
+        num_warps=4,
+    )

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/probability.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Probability-route compatibility helper for the existing FlashInfer full backend.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+
+@triton.jit
+def _min_p_renorm_prob_kernel(
+    probs_ptr,
+    min_p_ptr,
+    out_ptr,
+    vocab_size: tl.constexpr,
+    probs_row_stride: tl.constexpr,
+    out_row_stride: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    ENABLE_PDL: tl.constexpr,
+):
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    row = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    probs_row = probs_ptr + row * probs_row_stride
+    out_row = out_ptr + row * out_row_stride
+
+    max_prob = tl.full((), 0.0, tl.float32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offs
+        mask = cols < vocab_size
+        vals = tl.load(probs_row + cols, mask=mask, other=0.0).to(tl.float32)
+        max_prob = tl.maximum(max_prob, tl.max(tl.where(mask, vals, 0.0), axis=0))
+
+    threshold = max_prob * tl.load(min_p_ptr + row).to(tl.float32)
+    denom = tl.full((), 0.0, tl.float32)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offs
+        mask = cols < vocab_size
+        vals = tl.load(probs_row + cols, mask=mask, other=0.0).to(tl.float32)
+        keep = mask & (vals >= threshold)
+        denom += tl.sum(tl.where(keep, vals, 0.0), axis=0)
+
+    inv_denom = 1.0 / tl.maximum(denom, 1.0e-20)
+    for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+        cols = start + offs
+        mask = cols < vocab_size
+        vals = tl.load(probs_row + cols, mask=mask, other=0.0).to(tl.float32)
+        keep = mask & (vals >= threshold)
+        out = tl.where(keep, vals * inv_denom, 0.0)
+        tl.store(out_row + cols, out, mask=mask)
+
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def min_p_renorm_prob(
+    probs: torch.Tensor,
+    min_p: torch.Tensor,
+    *,
+    enable_pdl: bool = False,
+) -> torch.Tensor:
+    """Renormalize probabilities after applying a per-row min-p cutoff.
+
+    For each row, this computes ``threshold = min_p[row] * max(probs[row])``,
+    zeros probabilities below the threshold, and renormalizes the surviving
+    probabilities so the row sums to one.
+    """
+    if probs.ndim != 2:
+        raise ValueError(f"min_p_renorm_prob expects 2D probs, got {probs.ndim}D")
+    if min_p.ndim != 1:
+        raise ValueError(f"min_p_renorm_prob expects 1D min_p, got {min_p.ndim}D")
+    if min_p.shape[0] != probs.shape[0]:
+        raise ValueError(
+            "min_p length must match probs rows, "
+            f"got {min_p.shape[0]} and {probs.shape[0]}"
+        )
+    if probs.device.type != "cuda" or min_p.device.type != "cuda":
+        raise ValueError("min_p_renorm_prob requires CUDA tensors")
+    if probs.stride(-1) != 1:
+        raise ValueError(
+            f"min_p_renorm_prob requires stride-1 vocab dimension, got stride={probs.stride()}"
+        )
+    if not min_p.is_contiguous():
+        min_p = min_p.contiguous()
+
+    out = torch.empty_like(probs)
+    rows, vocab_size = probs.shape
+    if rows == 0:
+        return out
+
+    block_size = min(4096, triton.next_power_of_2(vocab_size))
+    num_warps = 4 if block_size <= 1024 else 8
+    extra_kwargs = {"launch_pdl": True} if enable_pdl else {}
+    _min_p_renorm_prob_kernel[(rows,)](
+        probs,
+        min_p,
+        out,
+        vocab_size=vocab_size,
+        probs_row_stride=probs.stride(0),
+        out_row_stride=out.stride(0),
+        BLOCK_SIZE=block_size,
+        ENABLE_PDL=enable_pdl,
+        num_warps=num_warps,
+        num_stages=3,
+        **extra_kwargs,
+    )
+    return out

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/topp.py

@@ -0,0 +1,577 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# TokenSpeed-specific top-p-only rejection/repair layout.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+_TOP_P_PARALLEL_BLOCK_SIZE = 1024
+_TOP_P_PARALLEL_NUM_ATTEMPTS = 3
+_TOP_P_REPAIR_NUM_ATTEMPTS = 8
+
+
+@triton.jit
+def _top_p_parallel_stage1_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    local_max_ptr,
+    local_sum_ptr,
+    local_argmax_ptr,
+    local_scores_ptr,
+    local_logits_ptr,
+    local_ids_ptr,
+    logits_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    num_blocks: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_ATTEMPTS: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block = tl.program_id(1)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    token_offsets = tl.arange(0, BLOCK_SIZE)
+    cols = block * BLOCK_SIZE + token_offsets
+    mask = cols < vocab_size
+
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+
+    vals = tl.load(
+        logits_ptr + row * logits_row_stride + cols,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    vals = vals / temperature
+
+    block_max = tl.max(vals, axis=0)
+    safe_block_max = tl.where(block_max > -float("inf"), block_max, 0.0)
+    block_sum = tl.sum(
+        tl.where(mask & (vals > -float("inf")), tl.exp(vals - safe_block_max), 0.0),
+        axis=0,
+    )
+    block_argmax = tl.min(tl.where(vals == block_max, cols, 2147483647), axis=0)
+    block_base = row * num_blocks + block
+    tl.store(local_max_ptr + block_base, block_max)
+    tl.store(local_sum_ptr + block_base, block_sum)
+    tl.store(local_argmax_ptr + block_base, block_argmax)
+
+    for attempt in tl.static_range(0, NUM_ATTEMPTS):
+        attempt_seed = tl.randint(seed, offset + attempt)
+        uniform = tl.maximum(tl.rand(attempt_seed, cols), 1.0e-7)
+        gumbel = -tl.log(-tl.log(uniform))
+        scores = tl.where(mask, vals + gumbel, float("-inf"))
+        best_score = tl.max(scores, axis=0)
+        best_id = tl.min(tl.where(scores == best_score, cols, 2147483647), axis=0)
+        best_logit = tl.max(tl.where(cols == best_id, vals, float("-inf")), axis=0)
+        out_offset = block_base * NUM_ATTEMPTS + attempt
+        tl.store(local_scores_ptr + out_offset, best_score)
+        tl.store(local_logits_ptr + out_offset, best_logit)
+        tl.store(local_ids_ptr + out_offset, best_id)
+
+
+@triton.jit
+def _top_p_parallel_stage2_kernel(
+    local_max_ptr,
+    local_sum_ptr,
+    local_argmax_ptr,
+    local_scores_ptr,
+    local_logits_ptr,
+    local_ids_ptr,
+    row_max_ptr,
+    row_total_ptr,
+    row_argmax_ptr,
+    row_candidate_logits_ptr,
+    row_candidate_ids_ptr,
+    num_blocks: tl.constexpr,
+    NUM_BLOCKS_PAD: tl.constexpr,
+    NUM_ATTEMPTS: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block_offsets = tl.arange(0, NUM_BLOCKS_PAD)
+    block_mask = block_offsets < num_blocks
+    block_base = row * num_blocks + block_offsets
+
+    local_max = tl.load(
+        local_max_ptr + block_base, mask=block_mask, other=-float("inf")
+    )
+    local_sum = tl.load(local_sum_ptr + block_base, mask=block_mask, other=0.0)
+    row_max = tl.max(local_max, axis=0)
+    safe_row_max = tl.where(row_max > -float("inf"), row_max, 0.0)
+    total = tl.sum(
+        tl.where(block_mask, local_sum * tl.exp(local_max - safe_row_max), 0.0),
+        axis=0,
+    )
+    local_argmax = tl.load(
+        local_argmax_ptr + block_base, mask=block_mask, other=2147483647
+    )
+    row_argmax = tl.min(
+        tl.where((local_max == row_max) & block_mask, local_argmax, 2147483647),
+        axis=0,
+    )
+
+    tl.store(row_max_ptr + row, row_max)
+    tl.store(row_total_ptr + row, total)
+    tl.store(row_argmax_ptr + row, row_argmax)
+
+    for attempt in tl.static_range(0, NUM_ATTEMPTS):
+        candidate_base = (row * num_blocks + block_offsets) * NUM_ATTEMPTS + attempt
+        scores = tl.load(
+            local_scores_ptr + candidate_base, mask=block_mask, other=-float("inf")
+        )
+        ids = tl.load(local_ids_ptr + candidate_base, mask=block_mask, other=2147483647)
+        logits = tl.load(
+            local_logits_ptr + candidate_base, mask=block_mask, other=-float("inf")
+        )
+        best_score = tl.max(scores, axis=0)
+        best_id = tl.min(
+            tl.where((scores == best_score) & block_mask, ids, 2147483647),
+            axis=0,
+        )
+        best_logit = tl.max(tl.where(ids == best_id, logits, -float("inf")), axis=0)
+        row_candidate_offset = row * NUM_ATTEMPTS + attempt
+        tl.store(row_candidate_ids_ptr + row_candidate_offset, best_id)
+        tl.store(row_candidate_logits_ptr + row_candidate_offset, best_logit)
+
+
+@triton.jit
+def _top_p_parallel_stage3_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    row_max_ptr,
+    row_candidate_logits_ptr,
+    row_candidate_ids_ptr,
+    partial_before_ptr,
+    logits_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    num_blocks: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_ATTEMPTS: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block = tl.program_id(1)
+    req_row = row // NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    token_offsets = tl.arange(0, BLOCK_SIZE)
+    cols = block * BLOCK_SIZE + token_offsets
+    mask = cols < vocab_size
+    row_max = tl.load(row_max_ptr + row)
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    vals = tl.load(
+        logits_ptr + row * logits_row_stride + cols,
+        mask=mask,
+        other=float("-inf"),
+    ).to(tl.float32)
+    vals = vals / temperature
+    weights = tl.exp(vals - row_max)
+
+    for attempt in tl.static_range(0, NUM_ATTEMPTS):
+        row_candidate_offset = row * NUM_ATTEMPTS + attempt
+        candidate_logit = tl.load(row_candidate_logits_ptr + row_candidate_offset)
+        candidate_id = tl.load(row_candidate_ids_ptr + row_candidate_offset)
+        before_mask = (vals > candidate_logit) | (
+            (vals == candidate_logit) & (cols < candidate_id)
+        )
+        before = tl.sum(tl.where(mask & before_mask, weights, 0.0), axis=0)
+        out_offset = (row * num_blocks + block) * NUM_ATTEMPTS + attempt
+        tl.store(partial_before_ptr + out_offset, before)
+
+
+@triton.jit
+def _top_p_parallel_stage4_kernel(
+    top_p_pool_ptr,
+    req_pool_indices_ptr,
+    row_total_ptr,
+    row_argmax_ptr,
+    row_candidate_ids_ptr,
+    partial_before_ptr,
+    accepted_ptr,
+    out_ptr,
+    num_blocks: tl.constexpr,
+    NUM_BLOCKS_PAD: tl.constexpr,
+    NUM_ATTEMPTS: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    req_row = row // NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    top_p = tl.load(top_p_pool_ptr + pool_idx).to(tl.float32)
+    target_mass = top_p * tl.load(row_total_ptr + row)
+    block_offsets = tl.arange(0, NUM_BLOCKS_PAD)
+    block_mask = block_offsets < num_blocks
+    token = tl.load(row_argmax_ptr + row)
+    found = tl.full((), 0, tl.int32)
+
+    for attempt in tl.static_range(0, NUM_ATTEMPTS):
+        before_base = (row * num_blocks + block_offsets) * NUM_ATTEMPTS + attempt
+        before = tl.sum(
+            tl.load(partial_before_ptr + before_base, mask=block_mask, other=0.0),
+            axis=0,
+        )
+        accepted = before < target_mass
+        candidate_id = tl.load(row_candidate_ids_ptr + row * NUM_ATTEMPTS + attempt)
+        take = (found == 0) & accepted
+        token = tl.where(take, candidate_id, token)
+        found = tl.where(take, 1, found)
+
+    tl.store(accepted_ptr + row, found)
+    tl.store(out_ptr + row, token)
+
+
+@triton.jit
+def _top_p_parallel_repair_kernel(
+    logits_ptr,
+    req_pool_indices_ptr,
+    temperature_pool_ptr,
+    top_p_pool_ptr,
+    seed_pool_ptr,
+    offsets_pool_ptr,
+    row_max_ptr,
+    row_total_ptr,
+    row_argmax_ptr,
+    accepted_ptr,
+    out_ptr,
+    logits_row_stride: tl.constexpr,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    START_ATTEMPT: tl.constexpr,
+    NUM_ATTEMPTS_TOTAL: tl.constexpr,
+    NUM_TOKENS_PER_REQ: tl.constexpr,
+):
+    row = tl.program_id(0)
+    accepted_found = tl.load(accepted_ptr + row)
+    accepted_token = tl.load(out_ptr + row)
+    req_row = row // NUM_TOKENS_PER_REQ
+    spec_pos = row - req_row * NUM_TOKENS_PER_REQ
+    pool_idx = tl.load(req_pool_indices_ptr + req_row)
+    token_offsets = tl.arange(0, BLOCK_SIZE)
+
+    temperature = tl.maximum(
+        tl.load(temperature_pool_ptr + pool_idx).to(tl.float32), 1.0e-20
+    )
+    top_p = tl.load(top_p_pool_ptr + pool_idx).to(tl.float32)
+    seed = tl.load(seed_pool_ptr + pool_idx).to(tl.int64)
+    offset = tl.load(offsets_pool_ptr + pool_idx).to(tl.int64) + spec_pos
+    row_max = tl.load(row_max_ptr + row)
+    total = tl.load(row_total_ptr + row)
+    target_mass = top_p * total
+    row_argmax = tl.load(row_argmax_ptr + row)
+
+    attempt = tl.full((), START_ATTEMPT, tl.int32)
+    while (attempt < NUM_ATTEMPTS_TOTAL) & (accepted_found == 0):
+        attempt_seed = tl.randint(seed, offset + attempt)
+        best_score = tl.full((), float("-inf"), tl.float32)
+        best_id = tl.full((), 2147483647, tl.int32)
+        best_logit = tl.full((), float("-inf"), tl.float32)
+
+        for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+            cols = start + token_offsets
+            mask = cols < vocab_size
+            vals = tl.load(
+                logits_ptr + row * logits_row_stride + cols,
+                mask=mask,
+                other=float("-inf"),
+            ).to(tl.float32)
+            vals = vals / temperature
+            uniform = tl.maximum(tl.rand(attempt_seed, cols), 1.0e-7)
+            gumbel = -tl.log(-tl.log(uniform))
+            scores = tl.where(mask, vals + gumbel, float("-inf"))
+            block_score = tl.max(scores, axis=0)
+            block_id = tl.min(tl.where(scores == block_score, cols, 2147483647), axis=0)
+            block_logit = tl.max(
+                tl.where(cols == block_id, vals, float("-inf")), axis=0
+            )
+            better = (block_score > best_score) | (
+                (block_score == best_score) & (block_id < best_id)
+            )
+            best_score = tl.where(better, block_score, best_score)
+            best_id = tl.where(better, block_id, best_id)
+            best_logit = tl.where(better, block_logit, best_logit)
+
+        before = tl.full((), 0.0, tl.float32)
+        for start in tl.range(0, vocab_size, BLOCK_SIZE, num_stages=3):
+            cols = start + token_offsets
+            mask = cols < vocab_size
+            vals = tl.load(
+                logits_ptr + row * logits_row_stride + cols,
+                mask=mask,
+                other=float("-inf"),
+            ).to(tl.float32)
+            vals = vals / temperature
+            weights = tl.exp(vals - row_max)
+            before_mask = (vals > best_logit) | (
+                (vals == best_logit) & (cols < best_id)
+            )
+            before += tl.sum(tl.where(mask & before_mask, weights, 0.0), axis=0)
+
+        accepted = before < target_mass
+        accepted_token = tl.where(accepted, best_id, accepted_token)
+        accepted_found = tl.where(accepted, 1, accepted_found)
+        attempt += 1
+
+    token = tl.where(accepted_found != 0, accepted_token, row_argmax)
+    tl.store(out_ptr + row, token)
+    tl.store(accepted_ptr + row, accepted_found)
+
+
+def gumbel_sample_top_p_parallel_from_pools(
+    logits: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    temperature_pool: torch.Tensor,
+    top_p_pool: torch.Tensor,
+    seed_pool: torch.Tensor,
+    offsets_pool: torch.Tensor,
+    local_max: torch.Tensor,
+    local_sum: torch.Tensor,
+    local_argmax: torch.Tensor,
+    local_scores: torch.Tensor,
+    local_logits: torch.Tensor,
+    local_ids: torch.Tensor,
+    row_max: torch.Tensor,
+    row_total: torch.Tensor,
+    row_argmax: torch.Tensor,
+    row_candidate_logits: torch.Tensor,
+    row_candidate_ids: torch.Tensor,
+    accepted: torch.Tensor,
+    out: torch.Tensor,
+    *,
+    block_size: int = _TOP_P_PARALLEL_BLOCK_SIZE,
+    num_attempts: int = _TOP_P_PARALLEL_NUM_ATTEMPTS,
+    num_tokens_per_req: int = 1,
+) -> torch.Tensor:
+    """Block-parallel top-p-only Gumbel sampler."""
+    if logits.ndim != 2:
+        raise ValueError("gumbel_sample_top_p_parallel_from_pools expects 2D logits")
+    if logits.device.type != "cuda":
+        raise ValueError("gumbel_sample_top_p_parallel_from_pools requires CUDA logits")
+    if logits.stride(-1) != 1:
+        raise ValueError(
+            "gumbel_sample_top_p_parallel_from_pools requires stride-1 vocab dimension, "
+            f"got stride={logits.stride()}"
+        )
+    rows, vocab_size = logits.shape
+    if vocab_size <= 0:
+        raise ValueError(
+            "gumbel_sample_top_p_parallel_from_pools requires non-empty vocab"
+        )
+    if num_tokens_per_req <= 0:
+        raise ValueError("num_tokens_per_req must be positive")
+    if rows % num_tokens_per_req != 0:
+        raise ValueError(
+            "logits rows must be divisible by num_tokens_per_req, "
+            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
+        )
+    request_rows = rows // num_tokens_per_req
+    if req_pool_indices.shape[0] != request_rows:
+        raise ValueError(
+            "req_pool_indices length must match request rows for parallel top-p sample, "
+            f"got {req_pool_indices.shape[0]} and {request_rows}"
+        )
+    if num_attempts <= 0:
+        raise ValueError("num_attempts must be positive")
+
+    num_blocks = triton.cdiv(vocab_size, block_size)
+    num_blocks_pad = triton.next_power_of_2(num_blocks)
+    for name, tensor, dtype in (
+        ("req_pool_indices", req_pool_indices, torch.int32),
+        ("seed_pool", seed_pool, torch.int64),
+        ("local_argmax", local_argmax, torch.int32),
+        ("local_ids", local_ids, torch.int32),
+        ("row_argmax", row_argmax, torch.int32),
+        ("row_candidate_ids", row_candidate_ids, torch.int32),
+        ("accepted", accepted, torch.int32),
+        ("out", out, torch.int32),
+    ):
+        if tensor.device.type != "cuda":
+            raise ValueError(f"{name} must be CUDA")
+        if tensor.dtype != dtype:
+            raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}")
+    for name, tensor in (
+        ("temperature_pool", temperature_pool),
+        ("top_p_pool", top_p_pool),
+        ("offsets_pool", offsets_pool),
+        ("local_max", local_max),
+        ("local_sum", local_sum),
+        ("local_scores", local_scores),
+        ("local_logits", local_logits),
+        ("row_max", row_max),
+        ("row_total", row_total),
+        ("row_candidate_logits", row_candidate_logits),
+    ):
+        if tensor.device.type != "cuda":
+            raise ValueError(f"{name} must be CUDA")
+    if out.shape[0] < rows:
+        raise ValueError(f"out is too small: {out.shape[0]} < {rows}")
+    if rows == 0:
+        return out[:0]
+
+    local_shape = (rows, num_blocks)
+    candidate_shape = (rows, num_blocks, num_attempts)
+    row_candidate_shape = (rows, num_attempts)
+    if local_max.shape[0] < rows or local_max.shape[1] < num_blocks:
+        raise ValueError(f"local_max must cover {local_shape}, got {local_max.shape}")
+    if local_sum.shape[0] < rows or local_sum.shape[1] < num_blocks:
+        raise ValueError(f"local_sum must cover {local_shape}, got {local_sum.shape}")
+    if local_argmax.shape[0] < rows or local_argmax.shape[1] < num_blocks:
+        raise ValueError(
+            f"local_argmax must cover {local_shape}, got {local_argmax.shape}"
+        )
+    for name, tensor in (
+        ("local_scores", local_scores),
+        ("local_logits", local_logits),
+        ("local_ids", local_ids),
+    ):
+        if (
+            tensor.shape[0] < rows
+            or tensor.shape[1] < num_blocks
+            or tensor.shape[2] < num_attempts
+        ):
+            raise ValueError(f"{name} must cover {candidate_shape}, got {tensor.shape}")
+    for name, tensor in (
+        ("row_max", row_max),
+        ("row_total", row_total),
+        ("row_argmax", row_argmax),
+        ("accepted", accepted),
+    ):
+        if tensor.shape[0] < rows:
+            raise ValueError(f"{name} is too small: {tensor.shape[0]} < {rows}")
+    for name, tensor in (
+        ("row_candidate_logits", row_candidate_logits),
+        ("row_candidate_ids", row_candidate_ids),
+    ):
+        if tensor.shape[0] < rows or tensor.shape[1] < num_attempts:
+            raise ValueError(
+                f"{name} must cover {row_candidate_shape}, got {tensor.shape}"
+            )
+
+    _top_p_parallel_stage1_kernel[(rows, num_blocks)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        seed_pool,
+        offsets_pool,
+        local_max,
+        local_sum,
+        local_argmax,
+        local_scores,
+        local_logits,
+        local_ids,
+        logits_row_stride=logits.stride(0),
+        vocab_size=vocab_size,
+        num_blocks=num_blocks,
+        BLOCK_SIZE=block_size,
+        NUM_ATTEMPTS=num_attempts,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+        num_stages=3,
+    )
+    _top_p_parallel_stage2_kernel[(rows,)](
+        local_max,
+        local_sum,
+        local_argmax,
+        local_scores,
+        local_logits,
+        local_ids,
+        row_max,
+        row_total,
+        row_argmax,
+        row_candidate_logits,
+        row_candidate_ids,
+        num_blocks=num_blocks,
+        NUM_BLOCKS_PAD=num_blocks_pad,
+        NUM_ATTEMPTS=num_attempts,
+        num_warps=8,
+        num_stages=3,
+    )
+    _top_p_parallel_stage3_kernel[(rows, num_blocks)](
+        logits,
+        req_pool_indices,
+        temperature_pool,
+        row_max,
+        row_candidate_logits,
+        row_candidate_ids,
+        local_scores,
+        logits_row_stride=logits.stride(0),
+        vocab_size=vocab_size,
+        num_blocks=num_blocks,
+        BLOCK_SIZE=block_size,
+        NUM_ATTEMPTS=num_attempts,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=4,
+        num_stages=3,
+    )
+    _top_p_parallel_stage4_kernel[(rows,)](
+        top_p_pool,
+        req_pool_indices,
+        row_total,
+        row_argmax,
+        row_candidate_ids,
+        local_scores,
+        accepted,
+        out,
+        num_blocks=num_blocks,
+        NUM_BLOCKS_PAD=num_blocks_pad,
+        NUM_ATTEMPTS=num_attempts,
+        NUM_TOKENS_PER_REQ=num_tokens_per_req,
+        num_warps=8,
+        num_stages=3,
+    )
+    if num_attempts < _TOP_P_REPAIR_NUM_ATTEMPTS:
+        _top_p_parallel_repair_kernel[(rows,)](
+            logits,
+            req_pool_indices,
+            temperature_pool,
+            top_p_pool,
+            seed_pool,
+            offsets_pool,
+            row_max,
+            row_total,
+            row_argmax,
+            accepted,
+            out,
+            logits_row_stride=logits.stride(0),
+            vocab_size=vocab_size,
+            BLOCK_SIZE=block_size,
+            START_ATTEMPT=num_attempts,
+            NUM_ATTEMPTS_TOTAL=_TOP_P_REPAIR_NUM_ATTEMPTS,
+            NUM_TOKENS_PER_REQ=num_tokens_per_req,
+            num_warps=4,
+            num_stages=3,
+        )
+    return out[:rows]

tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/verify.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: MIT AND Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# TokenSpeed samples target rows first, then verifies by token-id comparison.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+
+
+@triton.jit
+def _verify_chain_target_sampled_kernel(
+    predicts_ptr,
+    accept_index_ptr,
+    accept_token_num_ptr,
+    candidates_ptr,
+    target_sampled_ptr,
+    NUM_DRAFT_TOKENS: tl.constexpr,
+    ENABLE_PDL: tl.constexpr,
+):
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    row = tl.program_id(0)
+    base = row * NUM_DRAFT_TOKENS
+
+    tl.store(accept_index_ptr + base, base)
+
+    active = tl.full((), 1, tl.int32)
+    num_accepted = tl.full((), 0, tl.int32)
+    for i in tl.range(1, NUM_DRAFT_TOKENS):
+        target_id = tl.load(target_sampled_ptr + base + i - 1)
+        draft_id = tl.load(candidates_ptr + base + i)
+        match = (active != 0) & (draft_id == target_id)
+
+        tl.store(predicts_ptr + base + i - 1, target_id, mask=match)
+        tl.store(accept_index_ptr + base + i, base + i, mask=match)
+
+        num_accepted = tl.where(match, i, num_accepted)
+        active = tl.where(match, 1, 0)
+
+    final_id = tl.load(target_sampled_ptr + base + num_accepted)
+    tl.store(accept_token_num_ptr + row, num_accepted)
+    tl.store(predicts_ptr + base + num_accepted, final_id)
+
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def verify_chain_target_sampled(
+    predicts: torch.Tensor,
+    accept_index: torch.Tensor,
+    accept_token_num: torch.Tensor,
+    candidates: torch.Tensor,
+    target_sampled: torch.Tensor,
+    *,
+    enable_pdl: bool = False,
+) -> None:
+    """Verify a speculative chain against already-sampled target tokens."""
+    if candidates.ndim != 2:
+        raise ValueError(f"candidates must be 2D, got {candidates.ndim}D")
+    if accept_index.shape != candidates.shape:
+        raise ValueError(
+            f"accept_index shape {accept_index.shape} must match candidates {candidates.shape}"
+        )
+    bs, num_draft_tokens = candidates.shape
+    total = bs * num_draft_tokens
+    if predicts.shape[0] < total:
+        raise ValueError(f"predicts is too small: {predicts.shape[0]} < {total}")
+    if accept_token_num.shape[0] < bs:
+        raise ValueError(
+            f"accept_token_num is too small: {accept_token_num.shape[0]} < {bs}"
+        )
+    if target_sampled.numel() < total:
+        raise ValueError(
+            f"target_sampled is too small: {target_sampled.numel()} < {total}"
+        )
+    if candidates.dtype != torch.int32:
+        raise ValueError(f"candidates must be int32, got {candidates.dtype}")
+    if predicts.dtype != torch.int32:
+        raise ValueError(f"predicts must be int32, got {predicts.dtype}")
+    if accept_index.dtype != torch.int32:
+        raise ValueError(f"accept_index must be int32, got {accept_index.dtype}")
+    if accept_token_num.dtype != torch.int32:
+        raise ValueError(
+            f"accept_token_num must be int32, got {accept_token_num.dtype}"
+        )
+    if target_sampled.dtype not in (torch.int32, torch.int64):
+        raise ValueError(
+            f"target_sampled must be int32 or int64, got {target_sampled.dtype}"
+        )
+    if candidates.device.type != "cuda":
+        raise ValueError("verify_chain_target_sampled requires CUDA tensors")
+    if bs == 0:
+        return
+
+    target_sampled = target_sampled.reshape(-1)
+    extra_kwargs = {"launch_pdl": True} if enable_pdl else {}
+    _verify_chain_target_sampled_kernel[(bs,)](
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        target_sampled,
+        NUM_DRAFT_TOKENS=num_draft_tokens,
+        ENABLE_PDL=enable_pdl,
+        num_warps=1,
+        **extra_kwargs,
+    )