chore: import upstream snapshot with attribution

python/tokenspeed/runtime/sampling/backends/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+
+# Concrete backend modules are imported lazily in sampling.registry to avoid
+# circular imports (default.py imports register_backend from registry).

python/tokenspeed/runtime/sampling/backends/base.py

@@ -0,0 +1,285 @@
+# 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.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput
+    from tokenspeed.runtime.sampling.dp_sampling_config import DpSamplingRuntimeConfig
+    from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo
+    from tokenspeed.runtime.sampling.sampling_params import SamplingParams
+    from tokenspeed.runtime.utils.server_args import ServerArgs
+
+
+DEFAULT_RANDOM_SEED = 48
+CUDA_GRAPH_VARIANT_DEFAULT = "default"
+SPECULATIVE_ACCEPT_THRESHOLD_SINGLE = 1.0
+SPECULATIVE_ACCEPT_THRESHOLD_ACC = 1.0
+
+
+@dataclass
+class SamplingBackendConfig:
+
+    enable_nan_detection: bool = False
+
+    # Optional logprob features — OFF by default. These are checked at server
+    # start / graph capture time so the fast path has zero extra compute.
+    # Enabling any of these enlarges the captured graph footprint.
+    enable_output_logprobs: bool = False
+
+    # Sizing for pre-allocated per-backend buffers (e.g. coin buffers for
+    # rejection sampling). Required to keep RNG out of the CUDA graph.
+    max_bs: int = 1
+    max_draft_tokens_per_req: int = 1
+
+    # Sizing for backend-owned per-request state (e.g. token-count buffers
+    # for penalties in FlashInferFullSamplingBackend). Indexed by req_pool_idx, not
+    # batch row, so the data survives batch membership changes.
+    max_req_pool_size: int = 0
+    vocab_size: int = 0
+
+    device: torch.device | None = None
+    random_seed: int = DEFAULT_RANDOM_SEED
+
+    # Attention TP group for sampler-output broadcast (rank 0 wins).
+    tp_group: tuple[int, ...] | None = None
+    enable_tp_sync: bool = True
+
+    @classmethod
+    def from_server_args(
+        cls,
+        server_args: ServerArgs,
+        *,
+        max_bs: int,
+        max_draft_tokens_per_req: int,
+        device: str,
+        random_seed: int = DEFAULT_RANDOM_SEED,
+        max_req_pool_size: int = 0,
+        vocab_size: int = 0,
+        tp_group: tuple[int, ...] | None = None,
+    ) -> SamplingBackendConfig:
+
+        return cls(
+            enable_nan_detection=server_args.enable_nan_detection,
+            enable_output_logprobs=server_args.enable_output_logprobs,
+            max_bs=max_bs,
+            max_draft_tokens_per_req=max(max_draft_tokens_per_req, 1),
+            max_req_pool_size=max_req_pool_size,
+            vocab_size=vocab_size,
+            device=device,
+            random_seed=random_seed,
+            tp_group=tp_group,
+            enable_tp_sync=not server_args.disable_sampling_tp_sync,
+        )
+
+
+class SamplingBackend(ABC):
+    """Shared contract for single-step sampling and multi-step spec-decode verification.
+
+    Both methods return (output_tokens, accept_lengths). For sample(),
+    accept_lengths is all-ones so the downstream contract matches verify().
+
+    Backends that need random state override prepare() to refill per-request
+    buffers outside of any CUDA graph capture.
+
+    Requests asking for params a backend doesn't implement are NOT rejected;
+    the backend silently applies only what it supports, so all requests go
+    through the same captured graph.
+    """
+
+    # Subclasses that hold per-pool-idx state (scalars like temperature /
+    # top_k, plus large rows like _counts / _logit_bias) flip this to True
+    # so prepare_step() performs flip detection + _reset_slot. Stateless
+    # backends (greedy) leave it False and the whole prepare_step call is
+    # a no-op.
+    _HAS_POOL_STATE: bool = False
+    _SUPPORTS_DP_VERIFY: bool = False
+
+    def __init__(self, config: SamplingBackendConfig) -> None:
+
+        self.config = config
+
+        # Sentinel of "which rid currently owns each slot from this backend's
+        # point of view". rid is just a comparison value here, not a lookup
+        # key, so this is pool-keyed state (size O(pool_rows) strings), not
+        # rid-keyed state. A mismatch against the incoming rid is a flip.
+        if self._HAS_POOL_STATE:
+            pool_rows = config.max_req_pool_size + 1
+            self._last_rid_per_slot: list[str | None] = [None] * pool_rows
+
+        # Resolved once; None means maybe_broadcast is a no-op.
+        self._tp_pg = None
+        self._tp_src_global_rank: int | None = None
+        if (
+            config.enable_tp_sync
+            and config.tp_group is not None
+            and len(config.tp_group) > 1
+        ):
+            from tokenspeed.runtime.distributed.process_group_manager import (
+                process_group_manager as pg_manager,
+            )
+
+            self._tp_pg = pg_manager.get_process_group("nccl", config.tp_group)
+            self._tp_src_global_rank = config.tp_group[0]
+
+    def configure_dp_sampling(self, runtime: DpSamplingRuntimeConfig) -> None:
+        """Configure optional DP sampling state.
+
+        Stateless or unsupported backends ignore this; DP-capable backends
+        override it to initialize backend-local communication buffers.
+        """
+
+    def maybe_broadcast(self, *tensors: torch.Tensor) -> None:
+        """Broadcast each tensor from tp_group[0] so all attention-TP ranks
+        agree. No-op when sync is off or tp_size <= 1. Graph-safe."""
+        if self._tp_pg is None:
+            return
+        for t in tensors:
+            dist.broadcast(t, src=self._tp_src_global_rank, group=self._tp_pg)
+
+    def prepare_step(
+        self,
+        request_ids: list[str],
+        request_pool_indices: list[int],
+        sampling_params_list: list[SamplingParams],
+        num_tokens_per_req: int = 1,
+    ) -> None:
+        """Called once per step, outside the CUDA graph. Two jobs:
+
+        1. Flip detection: a slot's owning rid changed since last step
+           (first-use and rid-recycling look the same). Delegates to
+           _reset_slot which scatters all per-slot persistent state
+           (scalars, counts, bias, generators).
+        2. Per-step dynamic refill: coin buffers, etc. Delegated to the
+           subclass via _prepare_step_hook.
+
+        Stateless backends (greedy) short-circuit both phases.
+        """
+
+        if not self._HAS_POOL_STATE:
+            return
+
+        assert (
+            len(request_ids) == len(request_pool_indices) == len(sampling_params_list)
+        ), (
+            f"prepare_step expects aligned per-request lists; got "
+            f"rids={len(request_ids)}, pool_indices={len(request_pool_indices)}, "
+            f"sp_list={len(sampling_params_list)}"
+        )
+
+        pool_rows = len(self._last_rid_per_slot)
+        for rid, pool_idx, sp in zip(
+            request_ids, request_pool_indices, sampling_params_list
+        ):
+            assert (
+                0 <= pool_idx < pool_rows
+            ), f"pool_idx {pool_idx} out of range [0, {pool_rows}) for rid={rid}"
+            if self._last_rid_per_slot[pool_idx] != rid:
+                self._reset_slot(pool_idx, sp)
+                self._last_rid_per_slot[pool_idx] = rid
+
+        self._prepare_step_hook(
+            num_tokens_per_req=num_tokens_per_req,
+            bs=len(request_pool_indices),
+            request_pool_indices=request_pool_indices,
+        )
+
+    def prepare_capture(self, bs: int, num_tokens_per_req: int = 1) -> None:
+        """Per-step refill for the capture/warm-up path. No flip detection;
+        the backend uses its stub generator for any RNG-fed buffers so the
+        captured graph sees a fully-written state.
+        Default: no-op.
+        """
+        self._prepare_step_hook(
+            num_tokens_per_req=num_tokens_per_req,
+            bs=bs,
+            request_pool_indices=None,
+        )
+
+    def cuda_graph_capture_variants(self, num_tokens_per_req: int) -> tuple[str, ...]:
+        """Return sampler-specific CUDA graph variants to capture."""
+        return (CUDA_GRAPH_VARIANT_DEFAULT,)
+
+    def prepare_capture_variant(
+        self,
+        bs: int,
+        num_tokens_per_req: int,
+        variant: str,
+    ) -> None:
+        if variant != CUDA_GRAPH_VARIANT_DEFAULT:
+            raise ValueError(f"Unsupported CUDA graph variant: {variant}")
+        self.prepare_capture(bs=bs, num_tokens_per_req=num_tokens_per_req)
+
+    def cuda_graph_replay_variant(self, num_tokens_per_req: int) -> str:
+        return CUDA_GRAPH_VARIANT_DEFAULT
+
+    def _prepare_step_hook(
+        self,
+        num_tokens_per_req: int,
+        bs: int,
+        request_pool_indices: list[int] | None,
+    ) -> None:
+        """Subclass hook for per-step dynamic state (coin buffers, etc).
+        request_pool_indices=None is the capture path; otherwise the CPU
+        list from forward_op.request_pool_indices.
+        Default: no-op."""
+
+    def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None:
+        """Scatter all per-slot persistent state for a newly-assigned slot.
+        Called from prepare_step on flip. Stateful backends override."""
+        raise NotImplementedError
+
+    def reset_capture_state(self) -> None:
+        """Clear any per-pool state that warm-up iterations may have dirtied
+        before CUDA graph capture. Warm-up runs sample()/verify() against
+        pool row 0 (see CudaGraphWrapper capture path); stateful backends
+        override this to zero whatever row 0 accumulates. Default: no-op."""
+
+    def get_packed_output_d2h(
+        self,
+        output_tokens: torch.Tensor,
+        output_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        """If the backend wrote both outputs into a single contiguous GPU
+        buffer, return CPU views obtained from one D2H copy. Otherwise
+        return None and let the caller fall back to two separate D2Hs."""
+        return None
+
+    @abstractmethod
+    def sample(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+
+    @abstractmethod
+    def verify(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        candidates: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...

python/tokenspeed/runtime/sampling/backends/flashinfer.py

@@ -0,0 +1,621 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax
+from tokenspeed_kernel.ops.sampling.cuda import (
+    chain_speculative_sampling_target_only,
+    fused_topk_topp_prepare,
+    fused_topk_topp_renorm,
+    verify_chain_greedy,
+)
+from tokenspeed_kernel.ops.sampling.flashinfer import (
+    softmax,
+    top_k_renorm_prob,
+    top_k_top_p_sampling_from_probs,
+    top_p_renorm_prob,
+)
+from tokenspeed_kernel.ops.sampling.triton import gather_and_expand_scalars
+from tokenspeed_kernel.platform import current_platform
+from tokenspeed_kernel.torch_compile import get_compiler_backend
+
+# Resolved once at import: the fused top-k + top-p kernel is NVIDIA-only.
+# On non-NVIDIA platforms (e.g. ROCm) we fall back to the back-to-back
+# flashinfer renorm calls. Defining this at module scope keeps the hot path
+# branch-free in the captured graph.
+_FUSED_TOPK_TOPP_AVAILABLE = current_platform().is_nvidia
+
+from tokenspeed.runtime.distributed.dp_sampling_comm import DpSamplingComm
+from tokenspeed.runtime.sampling.backends.base import (
+    SPECULATIVE_ACCEPT_THRESHOLD_ACC,
+    SPECULATIVE_ACCEPT_THRESHOLD_SINGLE,
+    SamplingBackend,
+    SamplingBackendConfig,
+)
+from tokenspeed.runtime.sampling.dp_sampling_config import (
+    DpSamplingRuntimeConfig,
+    slice_dp_vocab_mask,
+)
+from tokenspeed.runtime.sampling.registry import register_backend
+from tokenspeed.runtime.sampling.utils import (
+    coin_eps,
+    gather_token_logprobs_torch,
+)
+from tokenspeed.runtime.utils.nvtx import nvtx_range
+from tokenspeed.runtime.utils.pdl import pdl_enabled
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput
+    from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo
+    from tokenspeed.runtime.sampling.sampling_params import SamplingParams
+
+
+class FlashInferSamplingBackend(SamplingBackend):
+    """Fast backend: fused softmax(temperature) + top_k_top_p_sampling_from_probs
+    for stochastic single-step sampling; cuda chain kernels (greedy +
+    rejection) for multi-step verification.
+
+    Scope is deliberately narrow — temperature / top_k / top_p only —
+    keeping the hot path to 2 kernels. Requests asking for min_p, penalties,
+    or logit_bias are silently ignored; use `flashinfer_full` if any of those
+    matter for the workload.
+    """
+
+    _HAS_POOL_STATE = True
+    _SUPPORTS_DP_VERIFY = True
+
+    def __init__(self, config: SamplingBackendConfig) -> None:
+
+        super().__init__(config)
+        self._init_dp_sampling(config)
+        self._init_shared_buffers(config)
+        self._init_pool_scalars(config)
+        # Pre-create the side stream used by fused_topk_topp_renorm. Must
+        # happen before any CUDA graph capture — cudaStreamCreate is illegal
+        # inside capture, and verify() runs from the captured graph.
+        fused_topk_topp_prepare(config.device)
+
+    def _init_dp_sampling(self, config: SamplingBackendConfig) -> None:
+        self._dp_tp_group = config.tp_group
+        self._dp_tp_size = (
+            len(self._dp_tp_group) if self._dp_tp_group is not None else 1
+        )
+        self._dp_rank = 0
+        self._dp_comm: DpSamplingComm | None = None
+        self._dp_comm_vocab_size = 0
+
+        if self._dp_tp_size <= 1:
+            self._dp_max_pad_bs = config.max_bs
+            self._dp_max_reqs_per_rank = config.max_bs
+            return
+
+        self._dp_max_pad_bs = (
+            (config.max_bs + self._dp_tp_size - 1) // self._dp_tp_size
+        ) * self._dp_tp_size
+        self._dp_max_reqs_per_rank = self._dp_max_pad_bs // self._dp_tp_size
+
+    def configure_dp_sampling(self, runtime: DpSamplingRuntimeConfig) -> None:
+        if not runtime.enabled:
+            return
+        if (
+            runtime.vocab_size is None
+            or runtime.max_bucket_bs is None
+            or runtime.topology is None
+            or runtime.device is None
+        ):
+            raise RuntimeError("enabled DP sampling runtime is incomplete")
+        topology = runtime.topology
+        if topology.tp_size != self._dp_tp_size:
+            raise RuntimeError(
+                f"DP sampling runtime tp_size={topology.tp_size} "
+                f"does not match backend tp_size={self._dp_tp_size}"
+            )
+        if topology.tp_group != self._dp_tp_group:
+            raise RuntimeError("DP sampling runtime tp_group does not match backend")
+        if self._dp_tp_group is None:
+            raise RuntimeError("dp_sampling requires a tp_group")
+        self._dp_rank = topology.tp_rank
+        if runtime.max_bucket_bs > self._dp_max_pad_bs:
+            raise RuntimeError(
+                f"DP sampling max_bucket_bs={runtime.max_bucket_bs} exceeds "
+                f"backend max_pad_bs={self._dp_max_pad_bs}"
+            )
+        if runtime.vocab_size % self._dp_tp_size != 0:
+            raise RuntimeError(
+                f"DP sampling vocab_size={runtime.vocab_size} must be divisible by "
+                f"tp_size={self._dp_tp_size}"
+            )
+        self._init_dp_verify_buffers(runtime.device)
+        if runtime.vocab_size == self._dp_comm_vocab_size:
+            return
+        if self._dp_comm is not None and self._dp_comm.is_initialized:
+            raise RuntimeError("Cannot resize DP sampling comm after use")
+        self._dp_comm_vocab_size = runtime.vocab_size
+        self._dp_comm = DpSamplingComm(
+            tp_size=self._dp_tp_size,
+            rank=self._dp_rank,
+            group=self._dp_tp_group,
+            max_pad_bs=self._dp_max_pad_bs,
+            num_tokens_per_req=runtime.num_tokens_per_req,
+            vocab_size=runtime.vocab_size,
+            logits_dtype=None,
+            device=runtime.device,
+        )
+
+    def _init_dp_verify_buffers(self, device: torch.device | str) -> None:
+        if self._predict_local_buf is not None:
+            return
+
+        max_n = self.config.max_draft_tokens_per_req
+        self._predict_local_buf = torch.zeros(
+            (self._dp_max_reqs_per_rank * max_n,), dtype=torch.int32, device=device
+        )
+        self._accept_index_local_buf = torch.zeros(
+            (self._dp_max_reqs_per_rank * max_n,), dtype=torch.int32, device=device
+        )
+        self._accept_length_local_buf = torch.zeros(
+            (self._dp_max_reqs_per_rank,), dtype=torch.int32, device=device
+        )
+
+    def _init_pool_scalars(self, config: SamplingBackendConfig) -> None:
+        # Capture warm-up reads row 0 with req_pool_indices zeroed, so row 0
+        # must carry neutral-sampling values that can't produce nan/inf.
+        pool_rows = config.max_req_pool_size + 1
+
+        self._temperature_pool = torch.ones(
+            (pool_rows,), dtype=torch.float32, device=config.device
+        )
+        self._top_k_pool = torch.ones(
+            (pool_rows,), dtype=torch.int32, device=config.device
+        )
+        self._top_p_pool = torch.ones(
+            (pool_rows,), dtype=torch.float32, device=config.device
+        )
+        self._seed_pool = torch.zeros(
+            (pool_rows,), dtype=torch.int64, device=config.device
+        )
+
+        # Per-slot CPU-side torch.Generators used to advance speculative
+        # coin buffers outside the CUDA graph. Seeded on flip from sp.seed.
+        # Slot 0 is pre-filled with _capture_gen so capture warm-up works
+        # without any real request having been registered.
+        #
+        # Retract-resume note: if a request is retracted and later takes a
+        # different pool slot on resume, _reset_slot re-seeds a fresh
+        # Generator from sp.seed. Sampling stays deterministic given the same
+        # seed, and flashinfer's Philox path (seed + seq_len offset) already
+        # gives per-step uniqueness independent of the torch.Generator.
+        self._cpu_generator_per_slot: list[torch.Generator | None] = [None] * pool_rows
+        self._cpu_generator_per_slot[0] = self._capture_gen
+
+    def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None:
+        self._temperature_pool[pool_idx].fill_(float(sp.temperature))
+        self._top_k_pool[pool_idx].fill_(int(sp.top_k))
+        self._top_p_pool[pool_idx].fill_(float(sp.top_p))
+        self._seed_pool[pool_idx].fill_(int(sp.seed))
+
+        cpu_gen = torch.Generator(device="cpu")
+        cpu_gen.manual_seed(int(sp.seed))
+        self._cpu_generator_per_slot[pool_idx] = cpu_gen
+
+    def _init_shared_buffers(self, config: SamplingBackendConfig) -> None:
+
+        max_pad_bs = self._dp_max_pad_bs
+        max_n = config.max_draft_tokens_per_req
+        # Persistent coin buffers. Filled per-request in prepare() outside the
+        # CUDA graph so verify() only reads from them.
+        self._coins_buf = torch.zeros(
+            (max_pad_bs, max_n),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._final_coins_buf = torch.zeros(
+            (max_pad_bs,), dtype=torch.float32, device=config.device
+        )
+
+        # Stub generator used during CUDA-graph capture/warm-up (no requests yet).
+        self._capture_gen = torch.Generator(device=config.device)
+        self._capture_gen.manual_seed(config.random_seed)
+
+        # Pre-allocated persistent buffers — no per-step alloc in the hot path.
+        self._ones_buf = torch.ones(
+            (max_pad_bs,), dtype=torch.int32, device=config.device
+        )
+        # predict + accept_length share one packed backing store.
+        # Layout: [0, max_bs * max_n) is predict, [max_bs * max_n, total)
+        # is accept_length.
+        self._predict_max = max_pad_bs * max_n
+        self._output_pack_buf = torch.zeros(
+            (self._predict_max + max_pad_bs,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._predict_buf = self._output_pack_buf[: self._predict_max]
+        self._accept_length_buf = self._output_pack_buf[self._predict_max :]
+        # Flat layout so [:bs * n].view(bs, n) is contiguous for any bs/n.
+        self._accept_index_buf = torch.zeros(
+            (max_pad_bs * max_n,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+
+        self._predict_local_buf: torch.Tensor | None = None
+        self._accept_index_local_buf: torch.Tensor | None = None
+        self._accept_length_local_buf: torch.Tensor | None = None
+
+    @torch.compile(dynamic=True, backend=get_compiler_backend())
+    def _prepare_step_hook(
+        self,
+        num_tokens_per_req: int,
+        bs: int,
+        request_pool_indices: list[int] | None = None,
+    ) -> None:
+        """Refill persistent coin buffers outside the captured graph.
+        request_pool_indices=None is the capture/warm-up path — uses
+        _capture_gen for all rows. Otherwise reads per-slot generators
+        populated via _reset_slot."""
+        if bs <= 0:
+            return
+
+        n = min(num_tokens_per_req, self.config.max_draft_tokens_per_req)
+        lo = coin_eps(self._coins_buf.dtype)
+        if request_pool_indices is None:
+            self._coins_buf[:bs, :n].uniform_(lo, 1.0, generator=self._capture_gen)
+            self._final_coins_buf[:bs].uniform_(lo, 1.0, generator=self._capture_gen)
+            return
+
+        cpu_coins = torch.empty((bs, n), dtype=torch.float32, pin_memory=True)
+        cpu_final = torch.empty((bs,), dtype=torch.float32, pin_memory=True)
+
+        for i, pool_idx in enumerate(request_pool_indices):
+            gen = self._cpu_generator_per_slot[pool_idx]
+            if gen is None:
+                raise RuntimeError(
+                    f"sampling slot {pool_idx} was not initialized before "
+                    "coin-buffer refill"
+                )
+            cpu_coins[i, :n].uniform_(lo, 1.0, generator=gen)
+            cpu_final[i].uniform_(lo, 1.0, generator=gen)
+
+        self._coins_buf[:bs, :n].copy_(cpu_coins, non_blocking=True)
+        self._final_coins_buf[:bs].copy_(cpu_final, non_blocking=True)
+
+    @nvtx_range("sampling:sample", color="yellow")
+    def sample(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        logits = logits_output.next_token_logits
+
+        # Grammar bitmask apply — captured inside the CUDA graph. Buffer is
+        # pre-bound by bind_grammar_mask_buf; non-grammar rows stay all-ones.
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits, vocab_mask=sampling_info.vocab_mask
+            )
+
+        if sampling_info.is_all_greedy:
+
+            batch_next_token_ids = sampling_argmax(logits)
+
+        else:
+
+            temperatures, top_ks, top_ps, _, seeds, offsets = gather_and_expand_scalars(
+                sampling_info.req_pool_indices,
+                temperature=self._temperature_pool,
+                top_k=self._top_k_pool,
+                top_p=self._top_p_pool,
+                seed=self._seed_pool,
+                offsets=sampling_info.valid_cache_lengths,
+                enable_pdl=pdl_enabled(),
+            )
+
+            probs = softmax(
+                logits,
+                temperature=temperatures.view(-1, 1),
+                enable_pdl=pdl_enabled(),
+            )
+            batch_next_token_ids = top_k_top_p_sampling_from_probs(
+                probs,
+                top_ks,
+                top_ps,
+                filter_apply_order="joint",
+                seed=seeds,
+                offset=offsets,
+                deterministic=True,
+            )
+
+        sampled = batch_next_token_ids.to(torch.int32)
+
+        # TP-rank sync: rank 0 wins.
+        self.maybe_broadcast(sampled)
+
+        if self.config.enable_output_logprobs:
+            logits_output.next_token_logprobs = gather_token_logprobs_torch(
+                logits, sampled
+            )
+
+        bs = logits.shape[0]
+
+        return sampled, self._ones_buf[:bs]
+
+    @nvtx_range("sampling:verify", color="yellow")
+    def verify(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        candidates: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        bs = candidates.shape[0]
+        num_tokens_per_req = candidates.shape[1]
+        vocab_mask = sampling_info.vocab_mask
+        logits_layout_plan = getattr(logits_output, "logits_layout_plan", None)
+        dp_sampling = logits_layout_plan is not None
+
+        if dp_sampling:
+            if self._dp_comm is None:
+                raise RuntimeError(
+                    "dp_sampling requires tp_size > 1, a resolved tp_group, "
+                    "and a configured DP comm"
+                )
+            dp_comm = self._dp_comm
+            tp_size = self._dp_tp_size
+            rank = self._dp_rank
+            effective_bs = logits_layout_plan.effective_bs
+            pad_bs = logits_layout_plan.bucket_bs
+            if effective_bs != bs:
+                raise RuntimeError(
+                    f"DP sampling effective_bs={effective_bs} must match "
+                    f"candidate batch size {bs}"
+                )
+            if (
+                pad_bs < effective_bs
+                or pad_bs > self._dp_max_pad_bs
+                or pad_bs % tp_size != 0
+            ):
+                raise RuntimeError(
+                    f"invalid DP sampling pad_bs={pad_bs} for effective_bs={effective_bs}, "
+                    f"max_pad_bs={self._dp_max_pad_bs}, tp_size={tp_size}"
+                )
+            bs = pad_bs // tp_size
+
+            # Shard by request so each request's draft chain stays on one rank.
+            shard = slice(rank * bs, (rank + 1) * bs)
+            if pad_bs > effective_bs:
+                candidates = torch.nn.functional.pad(
+                    candidates, (0, 0, 0, pad_bs - effective_bs)
+                )[shard]
+                pool_indices = torch.nn.functional.pad(
+                    sampling_info.req_pool_indices, (0, pad_bs - effective_bs)
+                )[shard]
+            else:
+                candidates = candidates[shard]
+                pool_indices = sampling_info.req_pool_indices[shard]
+            vocab_mask = slice_dp_vocab_mask(
+                vocab_mask,
+                full_bs=effective_bs,
+                pad_bs=pad_bs,
+                num_tokens_per_req=num_tokens_per_req,
+                shard=shard,
+            )
+            coins = self._coins_buf[shard]
+            final_coins = self._final_coins_buf[shard]
+            if (
+                self._predict_local_buf is None
+                or self._accept_index_local_buf is None
+                or self._accept_length_local_buf is None
+            ):
+                raise RuntimeError("DP sampling verify buffers are not initialized")
+            predict = self._predict_local_buf[: bs * num_tokens_per_req]
+            accept_index = (
+                self._accept_index_local_buf[: bs * num_tokens_per_req]
+                .view(bs, num_tokens_per_req)
+                .fill_(-1)
+            )
+            accept_length = self._accept_length_local_buf[:bs]
+        else:
+            pool_indices = sampling_info.req_pool_indices
+            coins = self._coins_buf
+            final_coins = self._final_coins_buf
+            predict = self._predict_buf[: bs * num_tokens_per_req]
+            accept_index = (
+                self._accept_index_buf[: bs * num_tokens_per_req]
+                .view(bs, num_tokens_per_req)
+                .fill_(-1)
+            )
+            accept_length = self._accept_length_buf[:bs]
+
+        logits = logits_output.next_token_logits
+        if dp_sampling:
+            expected_rows = bs * num_tokens_per_req
+            if logits.shape[0] != expected_rows:
+                raise RuntimeError(
+                    f"DP sampling logits rows {logits.shape[0]} != expected "
+                    f"{expected_rows}"
+                )
+
+        # Per-draft-position grammar bitmask: buffer shape
+        # [bs * num_tokens_per_req, V/32] matches the flat target logits.
+        if vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits,
+                vocab_mask=vocab_mask,
+            )
+
+        if sampling_info.is_all_greedy:
+
+            target_predict = sampling_argmax(logits).reshape(bs, num_tokens_per_req)
+
+            verify_chain_greedy(
+                predicts=predict,
+                accept_index=accept_index,
+                accept_token_num=accept_length,
+                candidates=candidates,
+                target_predict=target_predict,
+                batch_size=bs,
+                num_draft_tokens=num_tokens_per_req,
+                enable_pdl=pdl_enabled(),
+            )
+
+        else:
+
+            # Each request's N verified positions share one (temp, top_k, top_p)
+            # tuple; flat [bs*N] per-row knobs match the flat [bs*N, vocab] logits.
+            n = num_tokens_per_req
+            temperatures, top_ks, top_ps, _, _, _ = gather_and_expand_scalars(
+                pool_indices,
+                temperature=self._temperature_pool,
+                top_k=self._top_k_pool,
+                top_p=self._top_p_pool,
+                n=n,
+                enable_pdl=pdl_enabled(),
+            )
+
+            target_probs = softmax(
+                logits,
+                temperature=temperatures,
+                enable_pdl=pdl_enabled(),
+            )
+            if _FUSED_TOPK_TOPP_AVAILABLE:
+                # Fused replacement for the back-to-back top_k_renorm_prob +
+                # top_p_renorm_prob(is_deterministic=True) pair. Sentinel
+                # K = 1<<30 in top_ks routes per-row through the radix top-p
+                # only path.
+                target_probs = fused_topk_topp_renorm(
+                    target_probs,
+                    top_ks,
+                    top_ps,
+                    enable_pdl=pdl_enabled(),
+                )
+            else:
+                target_probs = top_k_renorm_prob(target_probs, top_ks)
+                target_probs = top_p_renorm_prob(
+                    target_probs, top_ps, is_deterministic=True
+                )
+            target_probs = target_probs.reshape(bs, n, -1)
+
+            chain_speculative_sampling_target_only(
+                predicts=predict,
+                accept_index=accept_index,
+                accept_token_num=accept_length,
+                candidates=candidates,
+                uniform_samples=coins[:bs, :n],
+                uniform_samples_for_final_sampling=final_coins[:bs],
+                target_probs=target_probs,
+                draft_probs=None,
+                threshold_single=SPECULATIVE_ACCEPT_THRESHOLD_SINGLE,
+                threshold_acc=SPECULATIVE_ACCEPT_THRESHOLD_ACC,
+                deterministic=not dp_sampling,
+                enable_pdl=pdl_enabled(),
+            )
+
+        accept_length += 1
+        logprobs_local = None
+        if self.config.enable_output_logprobs and dp_sampling:
+            # DP verify logits are still sharded by request at this point.
+            # Compute scalar logprobs for local predictions before gathering
+            # predictions to full-batch shape; the non-DP writer requires
+            # matching logits/token row counts.
+            logprobs_local = gather_token_logprobs_torch(logits, predict).view(
+                bs, num_tokens_per_req
+            )
+
+        if dp_sampling:
+            n = num_tokens_per_req
+            dp_comm.prepare_verify_outputs(logits_output.next_token_logits.dtype)
+            (
+                predict_full,
+                accept_index_full,
+                accept_length_full,
+            ) = dp_comm.gather_verify_outputs(
+                predict_local=predict.view(bs, n),
+                accept_index_local=accept_index,
+                accept_length_local=accept_length,
+                pad_bs=pad_bs,
+            )
+            predict = predict_full.view(-1)[: effective_bs * n]
+            accept_index = accept_index_full[:effective_bs]
+            accept_length = accept_length_full[:effective_bs]
+            if logprobs_local is not None:
+                logprobs_full = dp_comm.gather_verify_logprobs(
+                    logprobs_local,
+                    pad_bs=pad_bs,
+                )
+                logits_output.next_token_logprobs = logprobs_full.view(-1)[
+                    : effective_bs * n
+                ]
+        # TP-rank sync: rank 0 wins on the full verify-output triple.
+        # Load-bearing: flashinfer top_k_renorm_prob has no is_deterministic
+        # knob and produces non-bit-identical results across ranks (sub-ulp
+        # FP accumulation order).
+        # PDL still uses rank-0 outputs to keep ranks aligned. Without PDL,
+        # fused top-k + top-p is bit-identical across ranks and does not need
+        # a broadcast.
+        elif pdl_enabled():
+            self.maybe_broadcast(predict, accept_index, accept_length)
+        elif not _FUSED_TOPK_TOPP_AVAILABLE:
+            self.maybe_broadcast(predict, accept_index, accept_length)
+
+        if self.config.enable_output_logprobs and not dp_sampling:
+            logits_output.next_token_logprobs = gather_token_logprobs_torch(
+                logits, predict
+            )
+
+        return predict, accept_length
+
+    def get_packed_output_d2h(
+        self,
+        output_tokens: torch.Tensor,
+        output_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        """One D2H of the packed predict+accept_length region.
+
+        Only applies when both outputs alias into ``_output_pack_buf`` (the
+        verify() path). For ``sample()``, ``output_tokens`` is a fresh
+        argmax/top_k_top_p result and ``output_lengths`` is ``_ones_buf``,
+        neither of which lives in the pack. We fall back to two D2Hs.
+        """
+        if (
+            output_tokens.data_ptr() != self._output_pack_buf.data_ptr()
+            or output_lengths.data_ptr() != self._accept_length_buf.data_ptr()
+        ):
+            return None
+        n_t = output_tokens.numel()
+        n_l = output_lengths.numel()
+        # Copy the whole [0, predict_max + n_l). The gap [n_t, predict_max)
+        # is stale padding (max_bs * max_n  vs.  bs * n) — small enough that
+        # the saved launch beats the wasted bandwidth.
+        size = self._predict_max + n_l
+        cpu_pack = torch.empty(size, dtype=torch.int32, pin_memory=True)
+        cpu_pack.copy_(self._output_pack_buf[:size], non_blocking=True)
+        return (
+            cpu_pack[:n_t].view(output_tokens.shape),
+            cpu_pack[self._predict_max : self._predict_max + n_l],
+        )
+
+
+register_backend("flashinfer", FlashInferSamplingBackend)

python/tokenspeed/runtime/sampling/backends/flashinfer_full.py

@@ -0,0 +1,495 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from tokenspeed_kernel.ops.sampling.cuda import (
+    chain_speculative_sampling_target_only,
+    fused_topk_topp_renorm,
+)
+from tokenspeed_kernel.ops.sampling.flashinfer import (
+    min_p_sampling_from_probs,
+    softmax,
+    top_k_renorm_prob,
+    top_p_renorm_prob,
+)
+from tokenspeed_kernel.ops.sampling.triton import (
+    gather_and_expand_scalars,
+    min_p_renorm_prob,
+)
+from tokenspeed_kernel.torch_compile import get_compiler_backend
+
+from tokenspeed.runtime.sampling.backends.base import (
+    SPECULATIVE_ACCEPT_THRESHOLD_ACC,
+    SPECULATIVE_ACCEPT_THRESHOLD_SINGLE,
+    SamplingBackendConfig,
+)
+from tokenspeed.runtime.sampling.backends.flashinfer import (
+    _FUSED_TOPK_TOPP_AVAILABLE,
+    FlashInferSamplingBackend,
+)
+from tokenspeed.runtime.sampling.registry import register_backend
+from tokenspeed.runtime.utils.nvtx import nvtx_range
+from tokenspeed.runtime.utils.pdl import pdl_enabled
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput
+    from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo
+    from tokenspeed.runtime.sampling.sampling_params import SamplingParams
+
+
+class FlashInferFullSamplingBackend(FlashInferSamplingBackend):
+    """Superset of `flashinfer` adding min_p, frequency/presence/repetition
+    penalties, and per-token logit_bias, for both single-step sampling and
+    multi-step spec-decode verification.
+
+    Stochastic path runs the 4-kernel sequence softmax(temperature) →
+    top_k_renorm → top_p_renorm → min_p_sampling, unconditionally (requests
+    with min_p == 0 are a no-op through min_p_sampling_from_probs) so the
+    captured CUDA graph matches the runtime flow.
+
+    Layout:
+      * Per-pool-idx token counts (int32[max_req_pool_size, vocab]) —
+        accumulated after each sample/verify. Zeroed when a pool slot is
+        re-assigned to a new rid (see `on_pool_assignment`).
+      * Per-pool-idx logit bias (bf16[max_req_pool_size, vocab]) — zero by
+        default, scattered from SamplingParams.logit_bias on pool
+        assignment. Added to logits per step.
+      * Per-batch-row bf16 penalty scalars flowing through SamplingBatchInfo.
+
+    sample() / verify() apply (in order, BEFORE temperature/softmax):
+      1. repetition (multiplicative): logits = where(count>0,
+         where(logits>0, logits/rep, logits*rep), logits)
+      2. frequency + presence (additive):
+         logits -= freq_pen * count + pres_pen * (count>0)
+      3. logit_bias (additive): logits += logit_bias[req_pool_idx]
+
+    Post-sample/verify, accumulate accepted tokens into counts.
+
+    Out of scope in this iteration: min_new_tokens EOS mask, grammar vocab
+    mask. Both remain silently-ignored no-ops.
+    """
+
+    _SUPPORTS_DP_VERIFY = False
+
+    def __init__(self, config: SamplingBackendConfig) -> None:
+
+        super().__init__(config)
+
+        if config.max_req_pool_size <= 0 or config.vocab_size <= 0:
+
+            raise ValueError(
+                "FlashInferFullSamplingBackend requires max_req_pool_size > 0 and "
+                f"vocab_size > 0; got max_req_pool_size={config.max_req_pool_size}, "
+                f"vocab_size={config.vocab_size}"
+            )
+
+        # Valid pool indices run 0..max_req_pool_size inclusive.
+        pool_rows = config.max_req_pool_size + 1
+
+        self._counts = torch.zeros(
+            (pool_rows, config.vocab_size),
+            dtype=torch.int32,
+            device=config.device,
+        )
+
+        # bf16 is enough precision for typical client-supplied bias values
+        # (OpenAI caps |logit_bias| at 100).
+        self._logit_bias = torch.zeros(
+            (pool_rows, config.vocab_size),
+            dtype=torch.bfloat16,
+            device=config.device,
+        )
+
+        # Per-request penalty scalars + min_p. rep_pen starts at 1.0
+        # (multiplicative identity); others at 0.0 (additive identity).
+        self._min_p_pool = torch.zeros(
+            (pool_rows,), dtype=torch.float32, device=config.device
+        )
+        self._freq_pen_pool = torch.zeros(
+            (pool_rows,), dtype=torch.bfloat16, device=config.device
+        )
+        self._pres_pen_pool = torch.zeros(
+            (pool_rows,), dtype=torch.bfloat16, device=config.device
+        )
+        self._rep_pen_pool = torch.full(
+            (pool_rows,), 1.0, dtype=torch.bfloat16, device=config.device
+        )
+
+    # ------------------------------------------------------------------
+    # Lifecycle hooks
+    # ------------------------------------------------------------------
+
+    def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None:
+
+        # Scatter scalars inherited from the parent backend (temperature, top_k,
+        # top_p, seed).
+        super()._reset_slot(pool_idx, sp)
+
+        # Penalty + min_p scalars.
+        self._min_p_pool[pool_idx].fill_(float(sp.min_p))
+        self._freq_pen_pool[pool_idx].fill_(float(sp.frequency_penalty))
+        self._pres_pen_pool[pool_idx].fill_(float(sp.presence_penalty))
+        self._rep_pen_pool[pool_idx].fill_(float(sp.repetition_penalty))
+
+        # Zero the slot's count row (history from the previous occupant is
+        # no longer applicable).
+        self._counts[pool_idx].fill_(0)
+
+        # Zero + scatter logit_bias for the new rid. Zeroing the whole row
+        # first rather than diffing because the previous occupant's bias
+        # keys are unknown here.
+        self._logit_bias[pool_idx].fill_(0.0)
+
+        bias_map = getattr(sp, "logit_bias", None) if sp is not None else None
+        if bias_map:
+            vocab = self._logit_bias.shape[1]
+            raw_ids = [int(tid) for tid in bias_map.keys()]
+            assert all(0 <= tid < vocab for tid in raw_ids), (
+                f"logit_bias contains out-of-vocab token id(s); "
+                f"vocab_size={vocab}, offending={[t for t in raw_ids if not 0 <= t < vocab]}"
+            )
+            token_ids = torch.tensor(
+                raw_ids,
+                device=self._logit_bias.device,
+                dtype=torch.long,
+            )
+            bias_values = torch.tensor(
+                list(bias_map.values()),
+                device=self._logit_bias.device,
+                dtype=torch.bfloat16,
+            )
+            self._logit_bias[pool_idx, token_ids] = bias_values
+
+    def reset_capture_state(self) -> None:
+
+        # Warm-up iterations route all pool indices to row 0, which
+        # accumulates sampled tokens into _counts[0]. Zero it so the graph
+        # captures reads against a clean baseline. _logit_bias[0] is only
+        # written in on_pool_assignment, so it stays zero across warm-up.
+        self._counts[0].fill_(0)
+
+    # ------------------------------------------------------------------
+    # Penalty + bias application (shared by sample and verify)
+    # ------------------------------------------------------------------
+
+    @nvtx_range("sampling:penalties", color="yellow")
+    @torch.compile(dynamic=True, backend=get_compiler_backend())
+    def _apply_penalties_and_bias(
+        self,
+        logits: torch.Tensor,
+        sampling_info: SamplingBatchInfo,
+        num_tokens_per_req: int = 1,
+    ) -> torch.Tensor:
+        """logits is [bs * num_tokens_per_req, V]. Penalty scalars are gathered
+        from the pool-indexed buffers. num_tokens_per_req > 1 is the spec-decode
+        verify() path where per-request scalars are repeat_interleave'd to
+        align with flat logits.
+        """
+
+        pool_idx = sampling_info.req_pool_indices
+
+        if num_tokens_per_req > 1:
+
+            pool_idx = torch.repeat_interleave(pool_idx, num_tokens_per_req, dim=0)
+
+        counts = self._counts.index_select(0, pool_idx)  # [bs*N, V]
+        active = counts > 0
+        counts_f = counts.to(logits.dtype)
+        active_f = active.to(logits.dtype)
+
+        # Gather per-request penalty scalars from the pool. [bs*N] → [bs*N, 1]
+        # for broadcast against [bs*N, V] logits.
+        rep = (
+            self._rep_pen_pool.index_select(0, pool_idx).to(logits.dtype).unsqueeze(-1)
+        )
+        freq = (
+            self._freq_pen_pool.index_select(0, pool_idx).to(logits.dtype).unsqueeze(-1)
+        )
+        presence = (
+            self._pres_pen_pool.index_select(0, pool_idx).to(logits.dtype).unsqueeze(-1)
+        )
+
+        # 1. Repetition (multiplicative). scales is 1.0 where count==0, else
+        # rep_pen. Apply as logits/scales where logits>0, logits*scales else.
+        scales = torch.where(active, rep.expand_as(logits), torch.ones_like(logits))
+        logits = torch.where(logits > 0, logits / scales, logits * scales)
+
+        # 2. Frequency + presence (additive). Fused into a single subtract.
+        logits = logits - freq * counts_f - presence * active_f
+
+        # 3. Per-token logit_bias (additive). Rows without a logit_bias are
+        # all-zero, so the add is a no-op for them.
+        logits = logits + self._logit_bias.index_select(0, pool_idx)
+
+        return logits
+
+    @nvtx_range("sampling:accum_counts", color="yellow")
+    @torch.compile(dynamic=True, backend=get_compiler_backend())
+    def _accumulate_counts(
+        self,
+        pool_idx: torch.Tensor,
+        tokens: torch.Tensor,
+        weights: torch.Tensor,
+    ) -> None:
+        """Graph-safe in-place scatter: counts[pool_idx, tokens] += weights.
+        weights is int32; 0 masks invalid rows, 1 accumulates."""
+
+        self._counts.index_put_(
+            (pool_idx, tokens.long()),
+            weights.to(torch.int32),
+            accumulate=True,
+        )
+
+    # ------------------------------------------------------------------
+    # Sample / verify
+    # ------------------------------------------------------------------
+
+    @nvtx_range("sampling:sample", color="yellow")
+    def sample(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        logits = logits_output.next_token_logits.float()
+
+        # Grammar bitmask apply — captured inside the CUDA graph. Buffer is
+        # pre-bound by bind_grammar_mask_buf; non-grammar rows stay all-ones.
+        # Applied before raw_logprobs capture so constrained logprobs reflect
+        # the grammar-masked distribution.
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits, vocab_mask=sampling_info.vocab_mask
+            )
+
+        # Raw-distribution logprobs (pre-penalty, pre-temperature) when the
+        # server flag is on. Gather is done after we know the sampled id.
+        raw_logprobs = (
+            torch.log_softmax(logits, dim=-1)
+            if self.config.enable_output_logprobs
+            else None
+        )
+
+        logits = self._apply_penalties_and_bias(logits, sampling_info)
+
+        temperatures, top_ks, top_ps, min_ps, seeds, offsets = (
+            gather_and_expand_scalars(
+                sampling_info.req_pool_indices,
+                temperature=self._temperature_pool,
+                top_k=self._top_k_pool,
+                top_p=self._top_p_pool,
+                min_p=self._min_p_pool,
+                seed=self._seed_pool,
+                offsets=sampling_info.valid_cache_lengths,
+                enable_pdl=pdl_enabled(),
+            )
+        )
+
+        probs = softmax(
+            logits, temperature=temperatures.view(-1, 1), enable_pdl=pdl_enabled()
+        )
+
+        if _FUSED_TOPK_TOPP_AVAILABLE:
+            # Fused replacement for the back-to-back top_k_renorm_prob +
+            # top_p_renorm_prob(is_deterministic=True) pair. Sentinel
+            # K = 1<<30 in top_ks routes per-row through the radix top-p
+            # only path.
+            probs = fused_topk_topp_renorm(
+                probs,
+                top_ks,
+                top_ps,
+                enable_pdl=pdl_enabled(),
+            )
+        else:
+            probs = top_k_renorm_prob(probs, top_ks)
+            probs = top_p_renorm_prob(probs, top_ps, is_deterministic=True)
+
+        batch_next_token_ids = min_p_sampling_from_probs(
+            probs,
+            min_ps,
+            seed=seeds,
+            offset=offsets,
+            deterministic=True,
+        )
+
+        sampled = batch_next_token_ids.to(torch.int32)
+
+        # TP-rank sync BEFORE _accumulate_counts so per-rank counts stay aligned.
+        # For fused top-k + top-p, the results are bit-identical across ranks.
+        # So we don't need to broadcast the results.
+        if not _FUSED_TOPK_TOPP_AVAILABLE:
+            self.maybe_broadcast(sampled)
+
+        if raw_logprobs is not None:
+
+            logits_output.next_token_logprobs = raw_logprobs.gather(
+                -1, sampled.unsqueeze(-1)
+            ).squeeze(-1)
+
+        # Accumulate sampled tokens into counts (greedy path accumulates too
+        # so mixed later batches see the correct history).
+        self._accumulate_counts(
+            sampling_info.req_pool_indices,
+            sampled,
+            torch.ones_like(sampled, dtype=torch.int32),
+        )
+
+        bs = logits.shape[0]
+
+        return sampled, self._ones_buf[:bs]
+
+    @nvtx_range("sampling:verify", color="yellow")
+    def verify(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        candidates: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        bs = candidates.shape[0]
+        num_tokens_per_req = candidates.shape[1]
+
+        predict = self._predict_buf[: bs * num_tokens_per_req]
+        accept_index = (
+            self._accept_index_buf[: bs * num_tokens_per_req]
+            .view(bs, num_tokens_per_req)
+            .fill_(-1)
+        )
+        accept_length = self._accept_length_buf[:bs]
+
+        logits = logits_output.next_token_logits.float()
+
+        # Per-draft-position grammar bitmask: buffer shape
+        # [bs * num_tokens_per_req, V/32] matches the flat target logits.
+        # Applied before raw_logprobs capture so constrained logprobs reflect
+        # the grammar-masked distribution.
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits,
+                vocab_mask=sampling_info.vocab_mask,
+            )
+
+        # Raw (pre-penalty) logprobs captured before penalty application to
+        # match sample()'s semantics.
+        raw_logprobs = (
+            torch.log_softmax(logits, dim=-1)
+            if self.config.enable_output_logprobs
+            else None
+        )
+
+        logits = self._apply_penalties_and_bias(
+            logits,
+            sampling_info,
+            num_tokens_per_req=num_tokens_per_req,
+        )
+
+        temperatures, top_ks, top_ps, min_ps, _, _ = gather_and_expand_scalars(
+            sampling_info.req_pool_indices,
+            temperature=self._temperature_pool,
+            top_k=self._top_k_pool,
+            top_p=self._top_p_pool,
+            min_p=self._min_p_pool,
+            n=num_tokens_per_req,
+            enable_pdl=pdl_enabled(),
+        )
+
+        target_probs = softmax(
+            logits, temperature=temperatures.view(-1, 1), enable_pdl=pdl_enabled()
+        )
+        if _FUSED_TOPK_TOPP_AVAILABLE:
+            # Fused replacement for the back-to-back top_k_renorm_prob +
+            # top_p_renorm_prob(is_deterministic=True) pair. Sentinel
+            # K = 1<<30 in top_ks routes per-row through the radix top-p
+            # only path.
+            target_probs = fused_topk_topp_renorm(
+                target_probs,
+                top_ks,
+                top_ps,
+                enable_pdl=pdl_enabled(),
+            )
+        else:
+            target_probs = top_k_renorm_prob(target_probs, top_ks)
+            target_probs = top_p_renorm_prob(
+                target_probs, top_ps, is_deterministic=True
+            )
+
+        target_probs = min_p_renorm_prob(target_probs, min_ps, enable_pdl=pdl_enabled())
+
+        target_probs = target_probs.reshape(bs, num_tokens_per_req, -1)
+
+        coins = self._coins_buf[:bs, :num_tokens_per_req]
+        coins_for_final_sampling = self._final_coins_buf[:bs]
+
+        chain_speculative_sampling_target_only(
+            predicts=predict,
+            accept_index=accept_index,
+            accept_token_num=accept_length,
+            candidates=candidates.to(torch.int32),
+            uniform_samples=coins,
+            uniform_samples_for_final_sampling=coins_for_final_sampling,
+            target_probs=target_probs,
+            draft_probs=None,
+            threshold_single=SPECULATIVE_ACCEPT_THRESHOLD_SINGLE,
+            threshold_acc=SPECULATIVE_ACCEPT_THRESHOLD_ACC,
+            deterministic=True,
+            enable_pdl=pdl_enabled(),
+        )
+
+        accept_length += 1
+
+        # TP-rank sync BEFORE _accumulate_counts so per-rank counts stay aligned.
+        # For fused top-k + top-p, the results are bit-identical across ranks.
+        # So we don't need to broadcast the results.
+        if not _FUSED_TOPK_TOPP_AVAILABLE:
+            self.maybe_broadcast(predict, accept_index, accept_length)
+
+        # Accumulate accepted tokens into counts. accept_index is [bs, N]
+        # with -1 in unused slots; clamp to a safe index and mask with a
+        # weight of 0 so invalid slots are no-ops.
+        valid = accept_index >= 0  # [bs, N]
+        safe_positions = accept_index.clamp(min=0).long()  # [bs, N]
+        accepted_tokens = predict.long().gather(0, safe_positions.view(-1))
+
+        pool_idx_expanded = (
+            sampling_info.req_pool_indices.unsqueeze(-1)
+            .expand(-1, num_tokens_per_req)
+            .reshape(-1)
+        )
+
+        self._accumulate_counts(
+            pool_idx_expanded,
+            accepted_tokens,
+            valid.reshape(-1).to(torch.int32),
+        )
+
+        if raw_logprobs is not None:
+
+            logits_output.next_token_logprobs = raw_logprobs.gather(
+                -1, predict.unsqueeze(-1)
+            ).squeeze(-1)
+
+        return predict, accept_length
+
+
+register_backend("flashinfer_full", FlashInferFullSamplingBackend)

python/tokenspeed/runtime/sampling/backends/greedy.py

@@ -0,0 +1,253 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax
+from tokenspeed_kernel.ops.sampling.cuda import (
+    verify_chain_greedy as _verify_chain_greedy_cuda,
+)
+from tokenspeed_kernel.registry import error_fn
+
+from tokenspeed.runtime.sampling.backends.base import (
+    SamplingBackend,
+    SamplingBackendConfig,
+)
+from tokenspeed.runtime.sampling.registry import register_backend
+from tokenspeed.runtime.sampling.utils import gather_token_logprobs_torch
+from tokenspeed.runtime.utils.nvtx import nvtx_range
+from tokenspeed.runtime.utils.pdl import pdl_enabled
+
+if TYPE_CHECKING:
+
+    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput
+    from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo
+
+
+def _verify_chain_greedy_torch(
+    predicts: torch.Tensor,  # [bs * N] int32, in/out
+    accept_index: torch.Tensor,  # [bs, N] int32, in/out (-1-filled on entry)
+    accept_token_num: torch.Tensor,  # [bs] int32, out
+    candidates: torch.Tensor,  # [bs, N] int32
+    target_predict: torch.Tensor,  # [bs, N] int64 (argmax output)
+    batch_size: int,
+    num_draft_tokens: int,
+) -> None:
+    """Pure-torch equivalent of tokenspeed_kernel.verify_chain_greedy.
+
+    Used on non-CUDA devices and when the CUDA kernel is unavailable.
+    """
+
+    bs = batch_size
+    n = num_draft_tokens
+
+    # For i in 1..n-1: candidates[b, i] accepted iff it equals target_predict[b, i-1].
+    # Accepted prefix length per row = longest-leading-1s of the match array.
+    match = candidates[:, 1:] == target_predict[:, :-1].to(
+        candidates.dtype
+    )  # [bs, n-1]
+    leading = torch.cumprod(match.to(torch.int32), dim=1)  # [bs, n-1]
+    num_accepted = leading.sum(dim=1).to(torch.int32)  # [bs]
+
+    # Fill all of `predicts` with target_predict; slots outside the accepted
+    # prefix are harmless because accept_index keeps them at -1 and callers
+    # mask on that. Matches the CUDA kernel's observable state.
+    predicts.copy_(target_predict.reshape(-1).to(torch.int32))
+
+    device = candidates.device
+    pos = torch.arange(n, device=device).unsqueeze(0)  # [1, n]
+    batch_off = torch.arange(bs, device=device).unsqueeze(1) * n  # [bs, 1]
+    flat_idx = (batch_off + pos).to(torch.int32)  # [bs, n]
+    valid = pos <= num_accepted.unsqueeze(1)  # [bs, n]
+    accept_index.copy_(torch.where(valid, flat_idx, torch.full_like(accept_index, -1)))
+    accept_token_num.copy_(num_accepted)
+
+
+def _verify_chain_greedy(
+    predicts: torch.Tensor,
+    accept_index: torch.Tensor,
+    accept_token_num: torch.Tensor,
+    candidates: torch.Tensor,
+    target_predict: torch.Tensor,
+    batch_size: int,
+    num_draft_tokens: int,
+    enable_pdl: bool = False,
+) -> None:
+
+    # Prefer the CUDA kernel when available AND the tensors are on CUDA.
+    if _verify_chain_greedy_cuda is not error_fn and candidates.is_cuda:
+
+        _verify_chain_greedy_cuda(
+            predicts=predicts,
+            accept_index=accept_index,
+            accept_token_num=accept_token_num,
+            candidates=candidates,
+            target_predict=target_predict,
+            batch_size=batch_size,
+            num_draft_tokens=num_draft_tokens,
+            enable_pdl=enable_pdl,
+        )
+        return
+
+    _verify_chain_greedy_torch(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates,
+        target_predict=target_predict,
+        batch_size=batch_size,
+        num_draft_tokens=num_draft_tokens,
+    )
+
+
+class GreedySamplingBackend(SamplingBackend):
+    """Greedy-only backend: argmax for single-step, chain-greedy verify for
+    multi-step verification. No flashinfer / min_p / penalty machinery, no
+    coin buffers. Verify uses the fused CUDA kernel when available; falls
+    back to a pure-torch implementation otherwise (CPU, ROCm, etc.).
+
+    sampling_info is ignored for single-step (always argmax). verify() also
+    treats every request as greedy — stochastic verification is not
+    supported. Intended as the default backend and as a fallback when
+    flashinfer is unavailable."""
+
+    def __init__(self, config: SamplingBackendConfig) -> None:
+
+        super().__init__(config)
+
+        self._ones_buf = torch.ones(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+        # Pre-allocated int32 buffer for ``sample``'s argmax output: lets the
+        # cute_dsl kernel write int32 token ids directly, skipping the
+        # ``.to(torch.int32)`` cast and its elementwise launch in the
+        # CUDA-graph-captured hot path.
+        self._sample_token_buf = torch.empty(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+        self._predict_buf = torch.zeros(
+            (config.max_bs * config.max_draft_tokens_per_req,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        # Flat layout so [:bs * n].view(bs, n) is contiguous for any bs/n
+        # (required by maybe_broadcast / NCCL).
+        self._accept_index_buf = torch.zeros(
+            (config.max_bs * config.max_draft_tokens_per_req,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._accept_length_buf = torch.zeros(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+
+    @nvtx_range("sampling:sample", color="yellow")
+    def sample(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        logits = logits_output.next_token_logits
+        # Grammar bitmask apply — captured inside the CUDA graph. Buffer is
+        # pre-bound by bind_grammar_mask_buf; non-grammar rows stay all-ones
+        # so apply is a no-op.
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits, vocab_mask=sampling_info.vocab_mask
+            )
+        bs = logits.shape[0]
+        tokens = sampling_argmax(logits, out=self._sample_token_buf[:bs])
+
+        # TP-rank sync (rank 0 wins), mirrors FlashInferSamplingBackend.sample.
+        # All-gathered logits are not bit-identical across ranks, so per-rank
+        # argmax can diverge; an unsynced token id desyncs batch composition and
+        # deadlocks a downstream model all-reduce.
+        self.maybe_broadcast(tokens)
+
+        if self.config.enable_output_logprobs:
+            logits_output.next_token_logprobs = gather_token_logprobs_torch(
+                logits, tokens
+            )
+
+        return tokens, self._ones_buf[:bs]
+
+    @nvtx_range("sampling:verify", color="yellow")
+    def verify(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        candidates: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        bs = candidates.shape[0]
+        num_tokens_per_req = candidates.shape[1]
+
+        predict = self._predict_buf[: bs * num_tokens_per_req]
+        accept_index = (
+            self._accept_index_buf[: bs * num_tokens_per_req]
+            .view(bs, num_tokens_per_req)
+            .fill_(-1)
+        )
+        accept_length = self._accept_length_buf[:bs]
+
+        logits = logits_output.next_token_logits
+
+        # Per-draft-position grammar bitmask: buffer shape
+        # [bs * num_tokens_per_req, V/32] matches the flat target logits.
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits,
+                vocab_mask=sampling_info.vocab_mask,
+            )
+        target_predict = sampling_argmax(logits).reshape(bs, num_tokens_per_req)
+
+        _verify_chain_greedy(
+            predicts=predict,
+            accept_index=accept_index,
+            accept_token_num=accept_length,
+            candidates=candidates.to(torch.int32),
+            target_predict=target_predict,
+            batch_size=bs,
+            num_draft_tokens=num_tokens_per_req,
+            enable_pdl=pdl_enabled(),
+        )
+
+        accept_length += 1
+
+        # TP-rank sync on the full verify-output triple, mirrors
+        # FlashInferSamplingBackend.verify. Per-rank argmax / accept-length
+        # divergence (logits not bit-identical across ranks) desyncs batch
+        # composition and deadlocks the model all-reduce. Buffers are laid out
+        # flat so these views are NCCL-contiguous.
+        self.maybe_broadcast(predict, accept_index, accept_length)
+
+        if self.config.enable_output_logprobs:
+            logits_output.next_token_logprobs = gather_token_logprobs_torch(
+                logits, predict
+            )
+
+        return predict, accept_length
+
+
+register_backend("greedy", GreedySamplingBackend)

python/tokenspeed/runtime/sampling/backends/triton.py

@@ -0,0 +1,707 @@
+# 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 keeps its request-pool state and backend boundary.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from tokenspeed_kernel.ops.sampling.cute_dsl import argmax as cute_argmax
+from tokenspeed_kernel.ops.sampling.triton import (
+    _QRITA_PERCENTILE_TO_STD_TABLE,
+    gumbel_sample_from_pools,
+    gumbel_sample_from_pools_compact,
+    gumbel_sample_from_pools_generic,
+    gumbel_sample_top_k_top_p_from_pools,
+    gumbel_sample_top_k_top_p_qrita_from_pools,
+    gumbel_sample_top_p_parallel_from_pools,
+    selected_token_logprobs,
+    verify_chain_target_sampled,
+)
+
+from tokenspeed.runtime.sampling.backends.base import (
+    CUDA_GRAPH_VARIANT_DEFAULT,
+    SamplingBackend,
+    SamplingBackendConfig,
+)
+from tokenspeed.runtime.sampling.registry import register_backend
+from tokenspeed.runtime.sampling.sampling_params import _SAMPLING_EPS, _TOP_K_DISABLED
+from tokenspeed.runtime.sampling.utils import nan_guard_logits
+from tokenspeed.runtime.utils.nvtx import nvtx_range
+from tokenspeed.runtime.utils.pdl import pdl_enabled
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput
+    from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo
+    from tokenspeed.runtime.sampling.sampling_params import SamplingParams
+
+
+_GUMBEL_BLOCK_SIZE = 1024
+_COMPACT_GUMBEL_BLOCK_SIZE = 4096
+_COMPACT_GUMBEL_VOCAB_MAX = 32768
+_TOP_K_TOP_P_SMALL_BLOCK_SIZE = 1024
+_TOP_K_TOP_P_TUNED_VOCAB_MAX = 32768
+_TOP_K_TOP_P_PAD = 128
+_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE = 1024
+_TOP_P_PARALLEL_SAMPLE_ATTEMPTS = 3
+_TOP_P_PARALLEL_VERIFY_ATTEMPTS = 4
+_TOP_P_PARALLEL_MAX_ATTEMPTS = max(
+    _TOP_P_PARALLEL_SAMPLE_ATTEMPTS, _TOP_P_PARALLEL_VERIFY_ATTEMPTS
+)
+_QRITA_VERIFY_MIN_ROWS = 128
+
+_SAMPLE_ROUTE_GUMBEL_GENERIC = 0
+_SAMPLE_ROUTE_GUMBEL_NO_FILTER = 1
+_SAMPLE_ROUTE_GUMBEL_TOP_K = 2
+_SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P = 3
+_SAMPLE_ROUTE_GUMBEL_TOP_P = 4
+CUDA_GRAPH_VARIANT_TRITON_NO_FILTER = "triton_no_filter"
+CUDA_GRAPH_VARIANT_TRITON_TOP_K = "triton_top_k"
+CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P = "triton_top_k_top_p"
+CUDA_GRAPH_VARIANT_TRITON_TOP_P = "triton_top_p"
+CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER = "triton_verify_no_filter"
+
+_CUDA_GRAPH_VARIANT_SAMPLE_ROUTES = {
+    CUDA_GRAPH_VARIANT_TRITON_NO_FILTER: _SAMPLE_ROUTE_GUMBEL_NO_FILTER,
+    CUDA_GRAPH_VARIANT_TRITON_TOP_K: _SAMPLE_ROUTE_GUMBEL_TOP_K,
+    CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P: _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+    CUDA_GRAPH_VARIANT_TRITON_TOP_P: _SAMPLE_ROUTE_GUMBEL_TOP_P,
+    CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER: _SAMPLE_ROUTE_GUMBEL_NO_FILTER,
+}
+
+
+class TritonSamplingBackend(SamplingBackend):
+    """TokenSpeed pool-state backend using Triton Gumbel-Max kernels."""
+
+    _HAS_POOL_STATE = True
+
+    def __init__(self, config: SamplingBackendConfig) -> None:
+        super().__init__(config)
+        self._init_triton_pool_state(config)
+        self._init_triton_buffers(config)
+        self._sample_route = _SAMPLE_ROUTE_GUMBEL_GENERIC
+        self._top_k_top_p_pad = _TOP_K_TOP_P_PAD
+
+    def _init_triton_pool_state(self, config: SamplingBackendConfig) -> None:
+        pool_rows = config.max_req_pool_size + 1
+        self._temperature_pool = torch.ones(
+            (pool_rows,), dtype=torch.float32, device=config.device
+        )
+        self._top_k_pool = torch.ones(
+            (pool_rows,), dtype=torch.int32, device=config.device
+        )
+        self._top_p_pool = torch.ones(
+            (pool_rows,), dtype=torch.float32, device=config.device
+        )
+        self._seed_pool = torch.zeros(
+            (pool_rows,), dtype=torch.int64, device=config.device
+        )
+
+        self._ones_buf = torch.ones(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+        self._predict_buf = torch.zeros(
+            (config.max_bs * config.max_draft_tokens_per_req,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        # Flat layout so [:bs * n].view(bs, n) is contiguous for any bs/n
+        # (required by maybe_broadcast / NCCL).
+        self._accept_index_buf = torch.zeros(
+            (config.max_bs * config.max_draft_tokens_per_req,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._accept_length_buf = torch.zeros(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+
+    def _init_triton_buffers(self, config: SamplingBackendConfig) -> None:
+        pool_rows = config.max_req_pool_size + 1
+        self._zero_offsets_pool = torch.zeros(
+            (pool_rows,), dtype=torch.int64, device=config.device
+        )
+
+        vocab_size = max(int(config.vocab_size), 1)
+        gumbel_blocks = (vocab_size + _GUMBEL_BLOCK_SIZE - 1) // _GUMBEL_BLOCK_SIZE
+        self._gumbel_local_ids = torch.empty(
+            (config.max_bs, gumbel_blocks),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._gumbel_local_scores = torch.empty(
+            (config.max_bs, gumbel_blocks),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._gumbel_out = torch.empty(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+        self._req_pool_indices_i32 = torch.empty(
+            (config.max_bs,), dtype=torch.int32, device=config.device
+        )
+        self._gumbel_verify_out = torch.empty(
+            (config.max_bs * config.max_draft_tokens_per_req,),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._gumbel_verify_local_ids = torch.empty(
+            (config.max_bs * config.max_draft_tokens_per_req, gumbel_blocks),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._gumbel_verify_local_scores = torch.empty(
+            (config.max_bs * config.max_draft_tokens_per_req, gumbel_blocks),
+            dtype=torch.float32,
+            device=config.device,
+        )
+
+        topk_blocks = (vocab_size + _TOP_K_TOP_P_SMALL_BLOCK_SIZE - 1) // (
+            _TOP_K_TOP_P_SMALL_BLOCK_SIZE
+        )
+        topk_candidates = topk_blocks * _TOP_K_TOP_P_PAD
+        self._topk_candidate_ids = torch.empty(
+            (config.max_bs, topk_candidates),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._topk_candidate_logits = torch.empty(
+            (config.max_bs, topk_candidates),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._topk_verify_candidate_ids = torch.empty(
+            (config.max_bs * config.max_draft_tokens_per_req, topk_candidates),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._topk_verify_candidate_logits = torch.empty(
+            (config.max_bs * config.max_draft_tokens_per_req, topk_candidates),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        max_verify_rows = max(config.max_bs * config.max_draft_tokens_per_req, 1)
+        num_sms = torch.cuda.get_device_properties(config.device).multi_processor_count
+        self._qrita_verify_num_programs = min(num_sms, max_verify_rows)
+        self._qrita_verify_buffer = torch.empty(
+            (self._qrita_verify_num_programs, vocab_size),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._qrita_percentile_to_std_table = torch.tensor(
+            _QRITA_PERCENTILE_TO_STD_TABLE,
+            dtype=torch.float32,
+            device=config.device,
+        )
+
+        top_p_blocks = (vocab_size + _TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE - 1) // (
+            _TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE
+        )
+        top_p_rows = max(config.max_bs * config.max_draft_tokens_per_req, 1)
+        self._top_p_local_max = torch.empty(
+            (top_p_rows, top_p_blocks), dtype=torch.float32, device=config.device
+        )
+        self._top_p_local_sum = torch.empty(
+            (top_p_rows, top_p_blocks), dtype=torch.float32, device=config.device
+        )
+        self._top_p_local_argmax = torch.empty(
+            (top_p_rows, top_p_blocks), dtype=torch.int32, device=config.device
+        )
+        self._top_p_local_scores = torch.empty(
+            (top_p_rows, top_p_blocks, _TOP_P_PARALLEL_MAX_ATTEMPTS),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._top_p_local_logits = torch.empty(
+            (top_p_rows, top_p_blocks, _TOP_P_PARALLEL_MAX_ATTEMPTS),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._top_p_local_ids = torch.empty(
+            (top_p_rows, top_p_blocks, _TOP_P_PARALLEL_MAX_ATTEMPTS),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._top_p_row_max = torch.empty(
+            (top_p_rows,), dtype=torch.float32, device=config.device
+        )
+        self._top_p_row_total = torch.empty(
+            (top_p_rows,), dtype=torch.float32, device=config.device
+        )
+        self._top_p_row_argmax = torch.empty(
+            (top_p_rows,), dtype=torch.int32, device=config.device
+        )
+        self._top_p_row_candidate_logits = torch.empty(
+            (top_p_rows, _TOP_P_PARALLEL_MAX_ATTEMPTS),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._top_p_row_candidate_ids = torch.empty(
+            (top_p_rows, _TOP_P_PARALLEL_MAX_ATTEMPTS),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._top_p_accepted = torch.empty(
+            (top_p_rows,), dtype=torch.int32, device=config.device
+        )
+        self._selected_logprob_out = torch.empty(
+            (top_p_rows,), dtype=torch.float32, device=config.device
+        )
+
+    def _req_pool_indices_for_kernels(
+        self, req_pool_indices: torch.Tensor, rows: int
+    ) -> torch.Tensor:
+        req_pool_indices = req_pool_indices[:rows]
+        if req_pool_indices.dtype == torch.int32:
+            return req_pool_indices
+        if req_pool_indices.dtype != torch.int64:
+            raise ValueError(
+                "Triton sampling requires int32/int64 req_pool_indices, "
+                f"got {req_pool_indices.dtype}"
+            )
+        out = self._req_pool_indices_i32[:rows]
+        out.copy_(req_pool_indices, non_blocking=True)
+        return out
+
+    def _write_logprob_outputs(
+        self,
+        logits_output: LogitsProcessorOutput,
+        logits: torch.Tensor,
+        sampled: torch.Tensor,
+    ) -> None:
+        if not self.config.enable_output_logprobs:
+            return
+
+        rows = logits.shape[0]
+        selected_out = self._selected_logprob_out[:rows]
+        logits_output.next_token_logprobs = selected_token_logprobs(
+            logits, sampled, selected_out
+        )
+
+    @staticmethod
+    def _select_sample_route(
+        sampling_params_list: list[SamplingParams],
+    ) -> int:
+        if len(sampling_params_list) == 0:
+            return _SAMPLE_ROUTE_GUMBEL_GENERIC
+
+        top_ks = [int(sp.top_k) for sp in sampling_params_list]
+        top_ps = [float(sp.top_p) for sp in sampling_params_list]
+        all_top_p_one = all(abs(p - 1.0) <= _SAMPLING_EPS for p in top_ps)
+        all_top_k_disabled = all(k == _TOP_K_DISABLED for k in top_ks)
+        all_top_k_finite = all(k != _TOP_K_DISABLED for k in top_ks)
+
+        if all_top_k_disabled and all_top_p_one:
+            return _SAMPLE_ROUTE_GUMBEL_NO_FILTER
+        if all_top_k_disabled:
+            return _SAMPLE_ROUTE_GUMBEL_TOP_P
+        if all_top_k_finite:
+            if all_top_p_one:
+                return _SAMPLE_ROUTE_GUMBEL_TOP_K
+            return _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P
+        return _SAMPLE_ROUTE_GUMBEL_GENERIC
+
+    def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None:
+        self._temperature_pool[pool_idx].fill_(float(sp.temperature))
+        self._top_k_pool[pool_idx].fill_(int(sp.top_k))
+        self._top_p_pool[pool_idx].fill_(float(sp.top_p))
+        self._seed_pool[pool_idx].fill_(int(sp.seed))
+
+    def prepare_step(
+        self,
+        request_ids: list[str],
+        request_pool_indices: list[int],
+        sampling_params_list: list[SamplingParams],
+        num_tokens_per_req: int = 1,
+    ) -> None:
+        SamplingBackend.prepare_step(
+            self,
+            request_ids=request_ids,
+            request_pool_indices=request_pool_indices,
+            sampling_params_list=sampling_params_list,
+            num_tokens_per_req=num_tokens_per_req,
+        )
+        self._sample_route = self._select_sample_route(sampling_params_list)
+        self._top_k_top_p_pad = self._select_top_k_top_p_pad(sampling_params_list)
+
+    @staticmethod
+    def _select_top_k_top_p_pad(sampling_params_list: list[SamplingParams]) -> int:
+        finite_top_ks = [
+            int(sp.top_k)
+            for sp in sampling_params_list
+            if int(sp.top_k) != _TOP_K_DISABLED
+        ]
+        if finite_top_ks and max(finite_top_ks) <= 64:
+            return 64
+        return _TOP_K_TOP_P_PAD
+
+    def _use_qrita_verify_top_k_route(self, rows: int, vocab_size: int) -> bool:
+        return (
+            self._sample_route
+            in (_SAMPLE_ROUTE_GUMBEL_TOP_K, _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P)
+            and rows >= _QRITA_VERIFY_MIN_ROWS
+            and vocab_size >= _TOP_K_TOP_P_TUNED_VOCAB_MAX
+            and (
+                vocab_size > _TOP_K_TOP_P_TUNED_VOCAB_MAX or self._top_k_top_p_pad > 64
+            )
+        )
+
+    def prepare_capture(self, bs: int, num_tokens_per_req: int = 1) -> None:
+        self._sample_route = _SAMPLE_ROUTE_GUMBEL_GENERIC
+        self._top_k_top_p_pad = _TOP_K_TOP_P_PAD
+        SamplingBackend.prepare_capture(
+            self, bs=bs, num_tokens_per_req=num_tokens_per_req
+        )
+
+    def cuda_graph_capture_variants(self, num_tokens_per_req: int) -> tuple[str, ...]:
+        variants = (
+            CUDA_GRAPH_VARIANT_DEFAULT,
+            CUDA_GRAPH_VARIANT_TRITON_NO_FILTER,
+            CUDA_GRAPH_VARIANT_TRITON_TOP_P,
+            CUDA_GRAPH_VARIANT_TRITON_TOP_K,
+            CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P,
+        )
+        if num_tokens_per_req <= 1:
+            return variants
+        return (*variants, CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER)
+
+    def prepare_capture_variant(
+        self,
+        bs: int,
+        num_tokens_per_req: int,
+        variant: str,
+    ) -> None:
+        sample_route = _CUDA_GRAPH_VARIANT_SAMPLE_ROUTES.get(variant)
+        if sample_route is not None:
+            self._sample_route = sample_route
+            if sample_route in (
+                _SAMPLE_ROUTE_GUMBEL_TOP_K,
+                _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+            ):
+                self._top_k_top_p_pad = _TOP_K_TOP_P_PAD
+            SamplingBackend.prepare_capture(
+                self,
+                bs=bs,
+                num_tokens_per_req=num_tokens_per_req,
+            )
+            return
+        if variant == CUDA_GRAPH_VARIANT_DEFAULT:
+            self.prepare_capture(bs=bs, num_tokens_per_req=num_tokens_per_req)
+            return
+        raise ValueError(f"Unsupported CUDA graph variant: {variant}")
+
+    def cuda_graph_replay_variant(self, num_tokens_per_req: int) -> str:
+        if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER:
+            if num_tokens_per_req > 1:
+                return CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER
+            return CUDA_GRAPH_VARIANT_TRITON_NO_FILTER
+        if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_K:
+            return CUDA_GRAPH_VARIANT_TRITON_TOP_K
+        if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P:
+            return CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P
+        if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P:
+            return CUDA_GRAPH_VARIANT_TRITON_TOP_P
+        return CUDA_GRAPH_VARIANT_DEFAULT
+
+    def sample(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        logits = nan_guard_logits(
+            logits_output.next_token_logits, self.config.enable_nan_detection
+        )
+
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits, vocab_mask=sampling_info.vocab_mask
+            )
+
+        if sampling_info.is_all_greedy:
+            batch_next_token_ids = cute_argmax(logits)
+        else:
+            offsets_pool = (
+                sampling_info.valid_cache_lengths
+                if sampling_info.valid_cache_lengths is not None
+                else self._zero_offsets_pool
+            )
+            bs = logits.shape[0]
+            req_pool_indices = self._req_pool_indices_for_kernels(
+                sampling_info.req_pool_indices, bs
+            )
+            if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER:
+                if logits.shape[1] <= _COMPACT_GUMBEL_VOCAB_MAX:
+                    batch_next_token_ids = gumbel_sample_from_pools_compact(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        self._gumbel_out[:bs],
+                        block_size=_COMPACT_GUMBEL_BLOCK_SIZE,
+                    )
+                else:
+                    batch_next_token_ids = gumbel_sample_from_pools(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        self._gumbel_local_ids[:bs],
+                        self._gumbel_local_scores[:bs],
+                        self._gumbel_out[:bs],
+                    )
+            elif self._sample_route in (
+                _SAMPLE_ROUTE_GUMBEL_TOP_K,
+                _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+            ):
+                batch_next_token_ids = gumbel_sample_top_k_top_p_from_pools(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_k_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._topk_candidate_ids[:bs],
+                    self._topk_candidate_logits[:bs],
+                    self._gumbel_out[:bs],
+                    block_size=_TOP_K_TOP_P_SMALL_BLOCK_SIZE,
+                    top_k_pad=self._top_k_top_p_pad,
+                )
+            elif self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P:
+                batch_next_token_ids = gumbel_sample_top_p_parallel_from_pools(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._top_p_local_max[:bs],
+                    self._top_p_local_sum[:bs],
+                    self._top_p_local_argmax[:bs],
+                    self._top_p_local_scores[:bs],
+                    self._top_p_local_logits[:bs],
+                    self._top_p_local_ids[:bs],
+                    self._top_p_row_max[:bs],
+                    self._top_p_row_total[:bs],
+                    self._top_p_row_argmax[:bs],
+                    self._top_p_row_candidate_logits[:bs],
+                    self._top_p_row_candidate_ids[:bs],
+                    self._top_p_accepted[:bs],
+                    self._gumbel_out[:bs],
+                    block_size=_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE,
+                    num_attempts=_TOP_P_PARALLEL_SAMPLE_ATTEMPTS,
+                )
+            else:
+                batch_next_token_ids = gumbel_sample_from_pools_generic(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_k_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._gumbel_out[:bs],
+                )
+
+        sampled = batch_next_token_ids.to(torch.int32)
+        self.maybe_broadcast(sampled)
+
+        self._write_logprob_outputs(logits_output, logits, sampled)
+
+        return sampled, self._ones_buf[: logits.shape[0]]
+
+    @nvtx_range("sampling:verify", color="yellow")
+    def verify(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        candidates: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        bs = candidates.shape[0]
+        num_tokens_per_req = candidates.shape[1]
+
+        predict = self._predict_buf[: bs * num_tokens_per_req]
+        accept_index = (
+            self._accept_index_buf[: bs * num_tokens_per_req]
+            .view(bs, num_tokens_per_req)
+            .fill_(-1)
+        )
+        accept_length = self._accept_length_buf[:bs]
+
+        logits = nan_guard_logits(
+            logits_output.next_token_logits, self.config.enable_nan_detection
+        )
+
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits,
+                vocab_mask=sampling_info.vocab_mask,
+            )
+
+        if sampling_info.is_all_greedy:
+            target_sampled = cute_argmax(logits)
+            verify_chain_target_sampled(
+                predicts=predict,
+                accept_index=accept_index,
+                accept_token_num=accept_length,
+                candidates=candidates,
+                target_sampled=target_sampled,
+                enable_pdl=pdl_enabled(),
+            )
+        else:
+            offsets_pool = (
+                sampling_info.valid_cache_lengths
+                if sampling_info.valid_cache_lengths is not None
+                else self._zero_offsets_pool
+            )
+            req_pool_indices = self._req_pool_indices_for_kernels(
+                sampling_info.req_pool_indices, bs
+            )
+            if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER:
+                if logits.shape[1] <= _COMPACT_GUMBEL_VOCAB_MAX:
+                    target_sampled = gumbel_sample_from_pools_compact(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        self._gumbel_verify_out[: bs * num_tokens_per_req],
+                        block_size=_COMPACT_GUMBEL_BLOCK_SIZE,
+                        num_tokens_per_req=num_tokens_per_req,
+                    )
+                else:
+                    target_sampled = gumbel_sample_from_pools(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        self._gumbel_verify_local_ids[: bs * num_tokens_per_req],
+                        self._gumbel_verify_local_scores[: bs * num_tokens_per_req],
+                        self._gumbel_verify_out[: bs * num_tokens_per_req],
+                        num_tokens_per_req=num_tokens_per_req,
+                    )
+            elif self._sample_route in (
+                _SAMPLE_ROUTE_GUMBEL_TOP_K,
+                _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+            ):
+                rows = bs * num_tokens_per_req
+                if self._use_qrita_verify_top_k_route(rows, logits.shape[1]):
+                    target_sampled = gumbel_sample_top_k_top_p_qrita_from_pools(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._top_k_pool,
+                        self._top_p_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        self._qrita_verify_buffer,
+                        self._qrita_percentile_to_std_table,
+                        self._gumbel_verify_out[:rows],
+                        num_tokens_per_req=num_tokens_per_req,
+                        num_programs=min(self._qrita_verify_num_programs, rows),
+                    )
+                else:
+                    target_sampled = gumbel_sample_top_k_top_p_from_pools(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._top_k_pool,
+                        self._top_p_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        self._topk_verify_candidate_ids[:rows],
+                        self._topk_verify_candidate_logits[:rows],
+                        self._gumbel_verify_out[:rows],
+                        block_size=_TOP_K_TOP_P_SMALL_BLOCK_SIZE,
+                        top_k_pad=self._top_k_top_p_pad,
+                        num_tokens_per_req=num_tokens_per_req,
+                    )
+            elif self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P:
+                rows = bs * num_tokens_per_req
+                target_sampled = gumbel_sample_top_p_parallel_from_pools(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._top_p_local_max[:rows],
+                    self._top_p_local_sum[:rows],
+                    self._top_p_local_argmax[:rows],
+                    self._top_p_local_scores[:rows],
+                    self._top_p_local_logits[:rows],
+                    self._top_p_local_ids[:rows],
+                    self._top_p_row_max[:rows],
+                    self._top_p_row_total[:rows],
+                    self._top_p_row_argmax[:rows],
+                    self._top_p_row_candidate_logits[:rows],
+                    self._top_p_row_candidate_ids[:rows],
+                    self._top_p_accepted[:rows],
+                    self._gumbel_verify_out[:rows],
+                    block_size=_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE,
+                    num_attempts=_TOP_P_PARALLEL_VERIFY_ATTEMPTS,
+                    num_tokens_per_req=num_tokens_per_req,
+                )
+            else:
+                target_sampled = gumbel_sample_from_pools_generic(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_k_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._gumbel_verify_out[: bs * num_tokens_per_req],
+                    num_tokens_per_req=num_tokens_per_req,
+                )
+            verify_chain_target_sampled(
+                predicts=predict,
+                accept_index=accept_index,
+                accept_token_num=accept_length,
+                candidates=candidates,
+                target_sampled=target_sampled,
+                enable_pdl=pdl_enabled(),
+            )
+
+        accept_length += 1
+
+        # Rank 0 remains the source of truth for attention-TP agreement.
+        self.maybe_broadcast(predict, accept_index, accept_length)
+
+        if self.config.enable_output_logprobs:
+            self._write_logprob_outputs(
+                logits_output,
+                logits,
+                predict,
+            )
+
+        return predict, accept_length
+
+
+register_backend("triton", TritonSamplingBackend)

python/tokenspeed/runtime/sampling/backends/triton_full.py

@@ -0,0 +1,572 @@
+# 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 keeps pool-owned counts and logit-bias state.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from tokenspeed_kernel.ops.sampling.triton import (
+    accumulate_counts_inplace,
+    apply_penalties_logit_bias_inplace,
+    gumbel_sample_from_pools,
+    gumbel_sample_from_pools_compact,
+    gumbel_sample_from_pools_generic,
+    gumbel_sample_min_p_from_pools,
+    gumbel_sample_min_p_from_pools_parallel,
+    gumbel_sample_top_k_top_p_from_pools,
+    gumbel_sample_top_k_top_p_qrita_from_pools,
+    gumbel_sample_top_p_parallel_from_pools,
+    verify_chain_target_sampled,
+)
+
+from tokenspeed.runtime.sampling.backends.base import (
+    CUDA_GRAPH_VARIANT_DEFAULT,
+    SamplingBackend,
+    SamplingBackendConfig,
+)
+from tokenspeed.runtime.sampling.backends.triton import (
+    _COMPACT_GUMBEL_BLOCK_SIZE,
+    _COMPACT_GUMBEL_VOCAB_MAX,
+    _SAMPLE_ROUTE_GUMBEL_NO_FILTER,
+    _SAMPLE_ROUTE_GUMBEL_TOP_K,
+    _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+    _SAMPLE_ROUTE_GUMBEL_TOP_P,
+    _TOP_K_TOP_P_PAD,
+    _TOP_K_TOP_P_SMALL_BLOCK_SIZE,
+    _TOP_P_PARALLEL_SAMPLE_ATTEMPTS,
+    _TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE,
+    _TOP_P_PARALLEL_VERIFY_ATTEMPTS,
+    TritonSamplingBackend,
+)
+from tokenspeed.runtime.sampling.registry import register_backend
+from tokenspeed.runtime.sampling.utils import nan_guard_logits
+from tokenspeed.runtime.utils.nvtx import nvtx_range
+from tokenspeed.runtime.utils.pdl import pdl_enabled
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput
+    from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo
+    from tokenspeed.runtime.sampling.sampling_params import SamplingParams
+
+
+CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P = "triton_full_min_p"
+CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P = "triton_full_top_k_top_p_min_p"
+
+
+class TritonFullSamplingBackend(TritonSamplingBackend):
+    """Full sampling backend with TokenSpeed-owned state and Triton kernels."""
+
+    def __init__(self, config: SamplingBackendConfig) -> None:
+        super().__init__(config)
+        if config.max_req_pool_size <= 0 or config.vocab_size <= 0:
+            raise ValueError(
+                "TritonFullSamplingBackend requires max_req_pool_size > 0 and "
+                f"vocab_size > 0; got max_req_pool_size={config.max_req_pool_size}, "
+                f"vocab_size={config.vocab_size}"
+            )
+
+        pool_rows = config.max_req_pool_size + 1
+        self._counts = torch.zeros(
+            (pool_rows, config.vocab_size),
+            dtype=torch.int32,
+            device=config.device,
+        )
+        self._logit_bias = torch.zeros(
+            (pool_rows, config.vocab_size),
+            dtype=torch.bfloat16,
+            device=config.device,
+        )
+        self._min_p_pool = torch.zeros(
+            (pool_rows,), dtype=torch.float32, device=config.device
+        )
+        self._freq_pen_pool = torch.zeros(
+            (pool_rows,), dtype=torch.bfloat16, device=config.device
+        )
+        self._pres_pen_pool = torch.zeros(
+            (pool_rows,), dtype=torch.bfloat16, device=config.device
+        )
+        self._rep_pen_pool = torch.full(
+            (pool_rows,), 1.0, dtype=torch.bfloat16, device=config.device
+        )
+        self._min_p_row_max = torch.empty(
+            (config.max_bs * config.max_draft_tokens_per_req,),
+            dtype=torch.float32,
+            device=config.device,
+        )
+        self._full_has_min_p = True
+
+    def prepare_step(
+        self,
+        request_ids: list[str],
+        request_pool_indices: list[int],
+        sampling_params_list: list[SamplingParams],
+        num_tokens_per_req: int = 1,
+    ) -> None:
+        super().prepare_step(
+            request_ids=request_ids,
+            request_pool_indices=request_pool_indices,
+            sampling_params_list=sampling_params_list,
+            num_tokens_per_req=num_tokens_per_req,
+        )
+        self._full_has_min_p = any(float(sp.min_p) > 0.0 for sp in sampling_params_list)
+
+    def prepare_capture(self, bs: int, num_tokens_per_req: int = 1) -> None:
+        self._full_has_min_p = True
+        super().prepare_capture(bs=bs, num_tokens_per_req=num_tokens_per_req)
+
+    def prepare_capture_variant(
+        self,
+        bs: int,
+        num_tokens_per_req: int,
+        variant: str,
+    ) -> None:
+        if variant == CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P:
+            self._full_has_min_p = True
+            self._sample_route = _SAMPLE_ROUTE_GUMBEL_NO_FILTER
+            SamplingBackend.prepare_capture(
+                self,
+                bs=bs,
+                num_tokens_per_req=num_tokens_per_req,
+            )
+            return
+        if variant == CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P:
+            self._full_has_min_p = True
+            self._sample_route = _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P
+            self._top_k_top_p_pad = _TOP_K_TOP_P_PAD
+            SamplingBackend.prepare_capture(
+                self,
+                bs=bs,
+                num_tokens_per_req=num_tokens_per_req,
+            )
+            return
+        self._full_has_min_p = variant == CUDA_GRAPH_VARIANT_DEFAULT
+        super().prepare_capture_variant(
+            bs=bs,
+            num_tokens_per_req=num_tokens_per_req,
+            variant=variant,
+        )
+
+    def cuda_graph_capture_variants(self, num_tokens_per_req: int) -> tuple[str, ...]:
+        return (
+            *super().cuda_graph_capture_variants(num_tokens_per_req),
+            CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P,
+            CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P,
+        )
+
+    def cuda_graph_replay_variant(self, num_tokens_per_req: int) -> str:
+        if self._full_has_min_p:
+            if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER:
+                return CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P
+            if self._sample_route in (
+                _SAMPLE_ROUTE_GUMBEL_TOP_K,
+                _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+            ):
+                return CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P
+            return CUDA_GRAPH_VARIANT_DEFAULT
+        return super().cuda_graph_replay_variant(num_tokens_per_req)
+
+    def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None:
+        super()._reset_slot(pool_idx, sp)
+
+        self._min_p_pool[pool_idx].fill_(float(sp.min_p))
+        self._freq_pen_pool[pool_idx].fill_(float(sp.frequency_penalty))
+        self._pres_pen_pool[pool_idx].fill_(float(sp.presence_penalty))
+        self._rep_pen_pool[pool_idx].fill_(float(sp.repetition_penalty))
+
+        self._counts[pool_idx].fill_(0)
+        self._logit_bias[pool_idx].fill_(0.0)
+
+        bias_map = getattr(sp, "logit_bias", None) if sp is not None else None
+        if bias_map:
+            vocab = self._logit_bias.shape[1]
+            raw_ids = [int(tid) for tid in bias_map.keys()]
+            assert all(0 <= tid < vocab for tid in raw_ids), (
+                f"logit_bias contains out-of-vocab token id(s); "
+                f"vocab_size={vocab}, offending="
+                f"{[t for t in raw_ids if not 0 <= t < vocab]}"
+            )
+            token_ids = torch.tensor(
+                raw_ids,
+                device=self._logit_bias.device,
+                dtype=torch.long,
+            )
+            bias_values = torch.tensor(
+                list(bias_map.values()),
+                device=self._logit_bias.device,
+                dtype=torch.bfloat16,
+            )
+            self._logit_bias[pool_idx, token_ids] = bias_values
+
+    def reset_capture_state(self) -> None:
+        self._counts[0].fill_(0)
+
+    @nvtx_range("sampling:penalties", color="yellow")
+    def _apply_penalties_and_bias(
+        self,
+        logits: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        num_tokens_per_req: int = 1,
+    ) -> torch.Tensor:
+        return apply_penalties_logit_bias_inplace(
+            logits,
+            req_pool_indices,
+            self._counts,
+            self._logit_bias,
+            self._freq_pen_pool,
+            self._pres_pen_pool,
+            self._rep_pen_pool,
+            num_tokens_per_req=num_tokens_per_req,
+        )
+
+    @nvtx_range("sampling:accum_counts", color="yellow")
+    def _accumulate_counts(
+        self,
+        pool_idx: torch.Tensor,
+        tokens: torch.Tensor,
+        weights: torch.Tensor,
+    ) -> None:
+        accumulate_counts_inplace(self._counts, pool_idx, tokens, weights)
+
+    def _gumbel_sample_full_logits(
+        self,
+        logits: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        offsets_pool: torch.Tensor,
+        out: torch.Tensor,
+        *,
+        num_tokens_per_req: int = 1,
+    ) -> torch.Tensor:
+        rows = logits.shape[0]
+        if (
+            self._full_has_min_p
+            and self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER
+        ):
+            if logits.shape[1] > _COMPACT_GUMBEL_VOCAB_MAX:
+                local_ids = (
+                    self._gumbel_local_ids
+                    if num_tokens_per_req == 1
+                    else self._gumbel_verify_local_ids
+                )
+                local_scores = (
+                    self._gumbel_local_scores
+                    if num_tokens_per_req == 1
+                    else self._gumbel_verify_local_scores
+                )
+                return gumbel_sample_min_p_from_pools_parallel(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._min_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    local_ids[:rows],
+                    local_scores[:rows],
+                    self._min_p_row_max[:rows],
+                    out[:rows],
+                    num_tokens_per_req=num_tokens_per_req,
+                )
+            return gumbel_sample_min_p_from_pools(
+                logits,
+                req_pool_indices,
+                self._temperature_pool,
+                self._min_p_pool,
+                self._seed_pool,
+                offsets_pool,
+                out[:rows],
+                num_tokens_per_req=num_tokens_per_req,
+            )
+
+        if self._sample_route in (
+            _SAMPLE_ROUTE_GUMBEL_TOP_K,
+            _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P,
+        ):
+            if (
+                not self._full_has_min_p
+                and num_tokens_per_req > 1
+                and self._use_qrita_verify_top_k_route(rows, logits.shape[1])
+            ):
+                return gumbel_sample_top_k_top_p_qrita_from_pools(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_k_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._qrita_verify_buffer,
+                    self._qrita_percentile_to_std_table,
+                    out[:rows],
+                    num_tokens_per_req=num_tokens_per_req,
+                    num_programs=min(self._qrita_verify_num_programs, rows),
+                )
+            candidate_ids = (
+                self._topk_candidate_ids
+                if num_tokens_per_req == 1
+                else self._topk_verify_candidate_ids
+            )
+            candidate_logits = (
+                self._topk_candidate_logits
+                if num_tokens_per_req == 1
+                else self._topk_verify_candidate_logits
+            )
+            return gumbel_sample_top_k_top_p_from_pools(
+                logits,
+                req_pool_indices,
+                self._temperature_pool,
+                self._top_k_pool,
+                self._top_p_pool,
+                self._seed_pool,
+                offsets_pool,
+                candidate_ids[:rows],
+                candidate_logits[:rows],
+                out[:rows],
+                min_p_pool=self._min_p_pool if self._full_has_min_p else None,
+                block_size=_TOP_K_TOP_P_SMALL_BLOCK_SIZE,
+                top_k_pad=self._top_k_top_p_pad,
+                num_tokens_per_req=num_tokens_per_req,
+            )
+
+        if not self._full_has_min_p:
+            if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER:
+                if logits.shape[1] <= _COMPACT_GUMBEL_VOCAB_MAX:
+                    return gumbel_sample_from_pools_compact(
+                        logits,
+                        req_pool_indices,
+                        self._temperature_pool,
+                        self._seed_pool,
+                        offsets_pool,
+                        out[:rows],
+                        block_size=_COMPACT_GUMBEL_BLOCK_SIZE,
+                        num_tokens_per_req=num_tokens_per_req,
+                    )
+                local_ids = (
+                    self._gumbel_local_ids
+                    if num_tokens_per_req == 1
+                    else self._gumbel_verify_local_ids
+                )
+                local_scores = (
+                    self._gumbel_local_scores
+                    if num_tokens_per_req == 1
+                    else self._gumbel_verify_local_scores
+                )
+                return gumbel_sample_from_pools(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    local_ids[:rows],
+                    local_scores[:rows],
+                    out[:rows],
+                    num_tokens_per_req=num_tokens_per_req,
+                )
+            if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P:
+                return gumbel_sample_top_p_parallel_from_pools(
+                    logits,
+                    req_pool_indices,
+                    self._temperature_pool,
+                    self._top_p_pool,
+                    self._seed_pool,
+                    offsets_pool,
+                    self._top_p_local_max[:rows],
+                    self._top_p_local_sum[:rows],
+                    self._top_p_local_argmax[:rows],
+                    self._top_p_local_scores[:rows],
+                    self._top_p_local_logits[:rows],
+                    self._top_p_local_ids[:rows],
+                    self._top_p_row_max[:rows],
+                    self._top_p_row_total[:rows],
+                    self._top_p_row_argmax[:rows],
+                    self._top_p_row_candidate_logits[:rows],
+                    self._top_p_row_candidate_ids[:rows],
+                    self._top_p_accepted[:rows],
+                    out[:rows],
+                    block_size=_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE,
+                    num_attempts=(
+                        _TOP_P_PARALLEL_SAMPLE_ATTEMPTS
+                        if num_tokens_per_req == 1
+                        else _TOP_P_PARALLEL_VERIFY_ATTEMPTS
+                    ),
+                    num_tokens_per_req=num_tokens_per_req,
+                )
+
+        return gumbel_sample_from_pools_generic(
+            logits,
+            req_pool_indices,
+            self._temperature_pool,
+            self._top_k_pool,
+            self._top_p_pool,
+            self._seed_pool,
+            offsets_pool,
+            out[:rows],
+            min_p_pool=self._min_p_pool,
+            num_tokens_per_req=num_tokens_per_req,
+        )
+
+    @nvtx_range("sampling:sample", color="yellow")
+    def sample(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        logits = nan_guard_logits(
+            logits_output.next_token_logits, self.config.enable_nan_detection
+        ).float()
+
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits, vocab_mask=sampling_info.vocab_mask
+            )
+
+        logits_for_logprobs = (
+            logits.clone() if self.config.enable_output_logprobs else None
+        )
+
+        req_pool_indices = self._req_pool_indices_for_kernels(
+            sampling_info.req_pool_indices, logits.shape[0]
+        )
+        logits = self._apply_penalties_and_bias(logits, req_pool_indices)
+        offsets_pool = (
+            sampling_info.valid_cache_lengths
+            if sampling_info.valid_cache_lengths is not None
+            else self._zero_offsets_pool
+        )
+        sampled = self._gumbel_sample_full_logits(
+            logits,
+            req_pool_indices,
+            offsets_pool,
+            self._gumbel_out[: logits.shape[0]],
+        ).to(torch.int32)
+
+        self.maybe_broadcast(sampled)
+
+        if logits_for_logprobs is not None:
+            self._write_logprob_outputs(
+                logits_output,
+                logits_for_logprobs,
+                sampled,
+            )
+
+        self._accumulate_counts(
+            req_pool_indices,
+            sampled,
+            torch.ones_like(sampled, dtype=torch.int32),
+        )
+
+        return sampled, self._ones_buf[: logits.shape[0]]
+
+    @nvtx_range("sampling:verify", color="yellow")
+    def verify(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        candidates: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        bs = candidates.shape[0]
+        num_tokens_per_req = candidates.shape[1]
+
+        predict = self._predict_buf[: bs * num_tokens_per_req]
+        accept_index = (
+            self._accept_index_buf[: bs * num_tokens_per_req]
+            .view(bs, num_tokens_per_req)
+            .fill_(-1)
+        )
+        accept_length = self._accept_length_buf[:bs]
+
+        logits = nan_guard_logits(
+            logits_output.next_token_logits, self.config.enable_nan_detection
+        ).float()
+
+        if sampling_info.vocab_mask is not None:
+            sampling_info.apply_vocab_mask(
+                logits=logits,
+                vocab_mask=sampling_info.vocab_mask,
+            )
+
+        logits_for_logprobs = (
+            logits.clone() if self.config.enable_output_logprobs else None
+        )
+
+        req_pool_indices = self._req_pool_indices_for_kernels(
+            sampling_info.req_pool_indices, bs
+        )
+        logits = self._apply_penalties_and_bias(
+            logits,
+            req_pool_indices,
+            num_tokens_per_req=num_tokens_per_req,
+        )
+
+        offsets_pool = (
+            sampling_info.valid_cache_lengths
+            if sampling_info.valid_cache_lengths is not None
+            else self._zero_offsets_pool
+        )
+        target_sampled = self._gumbel_sample_full_logits(
+            logits,
+            req_pool_indices,
+            offsets_pool,
+            self._gumbel_verify_out[: bs * num_tokens_per_req],
+            num_tokens_per_req=num_tokens_per_req,
+        )
+        verify_chain_target_sampled(
+            predicts=predict,
+            accept_index=accept_index,
+            accept_token_num=accept_length,
+            candidates=candidates.to(torch.int32),
+            target_sampled=target_sampled,
+            enable_pdl=pdl_enabled(),
+        )
+
+        accept_length += 1
+
+        self.maybe_broadcast(predict, accept_index, accept_length)
+
+        valid = accept_index >= 0
+        safe_positions = accept_index.clamp(min=0).long()
+        accepted_tokens = predict.long().gather(0, safe_positions.view(-1))
+
+        pool_idx_expanded = (
+            req_pool_indices.unsqueeze(-1).expand(-1, num_tokens_per_req).reshape(-1)
+        )
+
+        self._accumulate_counts(
+            pool_idx_expanded,
+            accepted_tokens,
+            valid.reshape(-1).to(torch.int32),
+        )
+
+        if logits_for_logprobs is not None:
+            self._write_logprob_outputs(
+                logits_output,
+                logits_for_logprobs,
+                predict,
+            )
+
+        return predict, accept_length
+
+
+register_backend("triton_full", TritonFullSamplingBackend)

python/tokenspeed/runtime/sampling/custom_logit_processor.py

@@ -0,0 +1,58 @@
+# 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.
+
+import json
+from abc import ABC, abstractmethod
+from functools import cache
+from typing import Any, Self
+
+import dill
+import torch
+
+
+@cache
+def _cache_from_str(json_str: str) -> "CustomLogitProcessor":
+    """Deserialize a json string to a Callable object.
+    This function is cached to avoid redundant deserialization.
+    """
+    data = json.loads(json_str)
+    return dill.loads(bytes.fromhex(data["callable"]))
+
+
+class CustomLogitProcessor(ABC):
+    """Abstract base class for callable functions."""
+
+    @abstractmethod
+    def __call__(
+        self,
+        logits: torch.Tensor,
+        custom_param_list: list[dict[str, Any]] | None = None,
+    ) -> torch.Tensor:
+        """Define the callable behavior."""
+        raise NotImplementedError
+
+    def to_str(self) -> str:
+        """Serialize the callable function to a JSON-compatible string."""
+        return json.dumps({"callable": dill.dumps(self).hex()})
+
+    @classmethod
+    def from_str(cls, json_str: str) -> Self:
+        """Deserialize a callable function from a JSON string."""
+        return _cache_from_str(json_str)

python/tokenspeed/runtime/sampling/dp_sampling_config.py

@@ -0,0 +1,190 @@
+# 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.
+
+from __future__ import annotations
+
+import dataclasses
+
+import torch
+
+
+@dataclasses.dataclass(frozen=True)
+class DpSamplingSupport:
+    requested: bool
+    enabled: bool
+    infra_supports: bool
+    drafter_available: bool
+    backend_supports_verify: bool
+    tp_size: int
+    tp_group_set: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class DpSamplingTopology:
+    tp_rank: int
+    tp_size: int
+    tp_group: tuple[int, ...] | None
+    skip_all_gather: bool
+    tie_word_embeddings: bool = False
+
+    @property
+    def tp_group_set(self) -> bool:
+        return self.tp_group is not None
+
+
+@dataclasses.dataclass(frozen=True)
+class DpSamplingRuntimeConfig:
+    enabled: bool = False
+    vocab_size: int | None = None
+    max_bucket_bs: int | None = None
+    min_bs: int | None = None
+    num_tokens_per_req: int = 1
+    topology: DpSamplingTopology | None = None
+    device: torch.device | str | None = None
+
+
+@dataclasses.dataclass(frozen=True)
+class DpSamplingRuntimeLimits:
+    runtime_vocab_size: int
+    max_num_seqs: int
+    data_parallel_size: int
+    num_tokens_per_req: int
+    configured_min_bs: int | None
+    device: torch.device | str
+
+
+def resolve_dp_sampling_support(
+    *,
+    requested: bool,
+    drafter_available: bool,
+    backend_supports_verify: bool,
+    topology: DpSamplingTopology,
+) -> DpSamplingSupport:
+    tp_size = int(topology.tp_size)
+    tp_group_set = topology.tp_group_set
+    infra_supports = (
+        drafter_available and backend_supports_verify and tp_size > 1 and tp_group_set
+    )
+    support = DpSamplingSupport(
+        requested=bool(requested),
+        enabled=infra_supports and bool(requested),
+        infra_supports=infra_supports,
+        drafter_available=drafter_available,
+        backend_supports_verify=backend_supports_verify,
+        tp_size=tp_size,
+        tp_group_set=tp_group_set,
+    )
+    if support.requested and not support.infra_supports:
+        raise RuntimeError(
+            "--dp-sampling was set but Batch-DP spec-verify "
+            "preconditions are not met: "
+            f"drafter={support.drafter_available}, "
+            f"backend_supports_dp_verify={support.backend_supports_verify}, "
+            f"tp_size={support.tp_size}, "
+            f"tp_group_set={support.tp_group_set}"
+        )
+    return support
+
+
+def resolve_dp_sampling_runtime(
+    *,
+    support: DpSamplingSupport,
+    lm_head_rows: int,
+    topology: DpSamplingTopology,
+    limits: DpSamplingRuntimeLimits,
+) -> DpSamplingRuntimeConfig:
+    if not support.enabled:
+        return DpSamplingRuntimeConfig(
+            num_tokens_per_req=limits.num_tokens_per_req,
+            topology=topology,
+            device=limits.device,
+        )
+    validate_dp_sampling_lm_head_vocab(
+        lm_head_rows=lm_head_rows,
+        vocab_size=limits.runtime_vocab_size,
+        tp_size=topology.tp_size,
+        skip_all_gather=topology.skip_all_gather,
+        tie_word_embeddings=topology.tie_word_embeddings,
+    )
+    dp_vocab_size = int(lm_head_rows)
+    if not topology.skip_all_gather:
+        dp_vocab_size *= int(topology.tp_size)
+    dp_vocab_size = (
+        (dp_vocab_size + int(topology.tp_size) - 1) // int(topology.tp_size)
+    ) * int(topology.tp_size)
+    max_bs = limits.max_num_seqs // max(limits.data_parallel_size, 1)
+    max_bucket_bs = (
+        (max_bs + topology.tp_size - 1) // topology.tp_size
+    ) * topology.tp_size
+    min_bs = (
+        int(limits.configured_min_bs)
+        if limits.configured_min_bs is not None
+        else 2 * int(topology.tp_size)
+    )
+    if min_bs < 1:
+        raise ValueError("dp_sampling_min_bs must be >= 1")
+    return DpSamplingRuntimeConfig(
+        enabled=True,
+        vocab_size=dp_vocab_size,
+        max_bucket_bs=max_bucket_bs,
+        min_bs=min_bs,
+        num_tokens_per_req=limits.num_tokens_per_req,
+        topology=topology,
+        device=limits.device,
+    )
+
+
+def slice_dp_vocab_mask(
+    vocab_mask: torch.Tensor | None,
+    *,
+    full_bs: int,
+    pad_bs: int,
+    num_tokens_per_req: int,
+    shard: slice,
+) -> torch.Tensor | None:
+    if vocab_mask is None:
+        return None
+    n = num_tokens_per_req
+    if pad_bs > full_bs:
+        vocab_mask = torch.nn.functional.pad(
+            vocab_mask,
+            (0, 0, 0, (pad_bs - full_bs) * n),
+            value=-1,
+        )
+    return vocab_mask.view(pad_bs, n, -1)[shard].reshape(-1, vocab_mask.shape[-1])
+
+
+def validate_dp_sampling_lm_head_vocab(
+    *,
+    lm_head_rows: int,
+    vocab_size: int,
+    tp_size: int,
+    skip_all_gather: bool,
+    tie_word_embeddings: bool,
+) -> None:
+    if skip_all_gather and int(lm_head_rows) < int(vocab_size):
+        raise RuntimeError(
+            "Batch-DP sampling with skip_all_gather requires a replicated/"
+            "full-vocab LM head. Got a sharded LM head with "
+            f"lm_head_rows={lm_head_rows}, vocab_size={vocab_size}, "
+            f"tp_size={tp_size}, tie_word_embeddings={tie_word_embeddings}. "
+            "Disable --dp-sampling or use a model path that resolves a "
+            "replicated LM head."
+        )

python/tokenspeed/runtime/sampling/logits_layout.py

@@ -0,0 +1,119 @@
+# 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.
+
+"""Declarative logits layout plans used by sampling."""
+
+from __future__ import annotations
+
+import dataclasses
+
+import torch
+
+from tokenspeed.runtime.distributed.comm_backend import Group
+from tokenspeed.runtime.distributed.dp_sampling_comm import DpSamplingComm
+
+
+@dataclasses.dataclass(frozen=True)
+class LogitsLayoutPlan:
+    effective_bs: int
+    bucket_bs: int
+    tp_size: int
+    num_tokens_per_req: int
+
+
+class LogitsLayoutExecutor:
+    """Executes sampling-provided logits layout plans."""
+
+    def __init__(
+        self,
+        *,
+        tp_rank: int,
+        tp_size: int,
+        tp_group: Group,
+        max_bucket_bs: int,
+        num_tokens_per_req: int,
+        vocab_size: int,
+        device: torch.device | str,
+    ) -> None:
+        self._tp_rank = tp_rank
+        self._tp_size = tp_size
+        self._num_tokens_per_req = num_tokens_per_req
+        self._comm = DpSamplingComm(
+            tp_size=tp_size,
+            rank=tp_rank,
+            group=tp_group,
+            max_pad_bs=max_bucket_bs,
+            num_tokens_per_req=num_tokens_per_req,
+            vocab_size=vocab_size,
+            logits_dtype=None,
+            device=device,
+        )
+
+    def slice_hidden_states(
+        self,
+        hidden_states: torch.Tensor,
+        plan: LogitsLayoutPlan,
+    ) -> torch.Tensor:
+        n = self._tokens_per_req(plan)
+        rows = hidden_states.shape[0]
+        if rows % n != 0:
+            raise ValueError(f"hidden_states have {rows} rows, not divisible by N={n}")
+        bs = rows // n
+        if bs != plan.effective_bs:
+            raise ValueError(
+                f"hidden_states imply effective_bs={bs}, but logits layout plan has "
+                f"effective_bs={plan.effective_bs}"
+            )
+        pad_rows = (plan.bucket_bs - plan.effective_bs) * n
+        if pad_rows > 0:
+            hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, pad_rows))
+        reqs_per_rank = plan.bucket_bs // self._tp_size
+        start = self._tp_rank * reqs_per_rank * n
+        return hidden_states[start : start + reqs_per_rank * n]
+
+    def swap_batch_vocab(
+        self,
+        local_logits: torch.Tensor,
+        plan: LogitsLayoutPlan,
+    ) -> torch.Tensor:
+        n = self._tokens_per_req(plan)
+        rows = local_logits.shape[0]
+        if rows % n != 0:
+            raise ValueError(f"local logits have {rows} rows, not divisible by N={n}")
+        bs = rows // n
+        if bs != plan.effective_bs:
+            raise ValueError(
+                f"local logits imply effective_bs={bs}, but logits layout plan has "
+                f"effective_bs={plan.effective_bs}"
+            )
+        pad_rows = (plan.bucket_bs - plan.effective_bs) * n
+        if pad_rows > 0:
+            local_logits = torch.nn.functional.pad(local_logits, (0, 0, 0, pad_rows))
+        return self._comm.swap_batch_vocab(local_logits, pad_bs=plan.bucket_bs)
+
+    def _tokens_per_req(self, plan: LogitsLayoutPlan) -> int:
+        if (
+            plan.tp_size != self._tp_size
+            or plan.num_tokens_per_req != self._num_tokens_per_req
+            or plan.bucket_bs < plan.effective_bs
+            or plan.bucket_bs % self._tp_size != 0
+        ):
+            raise RuntimeError("invalid DP logits layout plan")
+        return plan.num_tokens_per_req

python/tokenspeed/runtime/sampling/registry.py

@@ -0,0 +1,93 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from tokenspeed_kernel.platform import current_platform
+
+from tokenspeed.runtime.sampling.backends.base import (
+    DEFAULT_RANDOM_SEED,
+    SamplingBackend,
+    SamplingBackendConfig,
+)
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.utils.server_args import ServerArgs
+
+
+_BACKEND_REGISTRY: dict[str, type[SamplingBackend]] = {}
+
+
+def _get_default_backend_name() -> str:
+    if current_platform().is_nvidia:
+        return "flashinfer"
+    return "greedy"
+
+
+def _resolve_backend_name(server_args: ServerArgs) -> str:
+    return server_args.sampling_backend or _get_default_backend_name()
+
+
+def register_backend(name: str, cls: type[SamplingBackend]) -> None:
+    _BACKEND_REGISTRY[name] = cls
+
+
+def create_sampling_backend(
+    server_args: ServerArgs,
+    *,
+    max_bs: int,
+    max_draft_tokens_per_req: int,
+    device: str,
+    random_seed: int = DEFAULT_RANDOM_SEED,
+    max_req_pool_size: int = 0,
+    vocab_size: int = 0,
+    tp_group: tuple[int, ...] | None = None,
+) -> SamplingBackend:
+    # Trigger concrete-backend registration on first use.
+    from tokenspeed.runtime.sampling.backends import flashinfer as _fi  # noqa: F401
+    from tokenspeed.runtime.sampling.backends import (  # noqa: F401
+        flashinfer_full as _ff,
+    )
+    from tokenspeed.runtime.sampling.backends import greedy as _g  # noqa: F401
+    from tokenspeed.runtime.sampling.backends import triton as _t  # noqa: F401
+    from tokenspeed.runtime.sampling.backends import triton_full as _tf  # noqa: F401
+
+    name = _resolve_backend_name(server_args)
+    if name not in _BACKEND_REGISTRY:
+        raise ValueError(
+            f"Unknown sampling backend: {name!r}. "
+            f"Available: {list(_BACKEND_REGISTRY)}"
+        )
+    cls = _BACKEND_REGISTRY[name]
+
+    return cls(
+        SamplingBackendConfig.from_server_args(
+            server_args,
+            max_bs=max_bs,
+            max_draft_tokens_per_req=max_draft_tokens_per_req,
+            device=device,
+            random_seed=random_seed,
+            max_req_pool_size=max_req_pool_size,
+            vocab_size=vocab_size,
+            tp_group=tp_group,
+        )
+    )

python/tokenspeed/runtime/sampling/sampling_batch_info.py

@@ -0,0 +1,140 @@
+# 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.
+
+from __future__ import annotations
+
+import dataclasses
+import threading
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import torch
+
+from tokenspeed.runtime.utils import get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.engine.schedule_batch import ScheduleBatch
+
+
+@dataclasses.dataclass
+class SamplingBatchInfo:
+    # Basic batched sampling params. Disaggregated decode populates these via
+    # from_schedule_batch. The standard hot path leaves them None; sampling
+    # backends gather params from their own pool-indexed buffers.
+    temperatures: torch.Tensor | None = None
+    top_ps: torch.Tensor | None = None
+    top_ks: torch.Tensor | None = None
+    min_ps: torch.Tensor | None = None
+
+    # Whether all requests use greedy sampling
+    is_all_greedy: bool = False
+
+    # Masking tensors for grammar-guided structured outputs
+    vocab_size: int = 0
+    grammars: list | None = None
+    vocab_mask: torch.Tensor | None = None
+    # Backend-specific in-place fn ``(logits, vocab_mask) -> None``,
+    # bound by ``capturable_grammar.bind_grammar_mask_buf`` so the
+    # captured sampler can apply the bitmask without branching on
+    # backend.
+    apply_vocab_mask: Callable[[torch.Tensor, torch.Tensor], None] | None = None
+
+    # An event used for overlap schedule
+    sampling_info_done: threading.Event | None = None
+
+    # int64[bs] — req_pool_idx per batch row. Sampling backends gather
+    # their pool-indexed scalar buffers (temperature / top_k / top_p /
+    # seeds / penalties / logit_bias / counts) against this index.
+    req_pool_indices: torch.Tensor | None = None
+
+    # int32[pool_rows] — RuntimeStates.valid_cache_lengths, read-only
+    # reference. Sampling backends derive the per-request Philox offset
+    # from `valid_cache_lengths.index_select(0, req_pool_indices)`;
+    # carrying the reference rather than the gathered view keeps the
+    # index_select inside the captured graph.
+    valid_cache_lengths: torch.Tensor | None = None
+
+    # Device
+    device: str = "cuda"
+
+    def __getitem__(self, s: slice) -> SamplingBatchInfo:
+        """Row-slice batch-indexed fields; pool/scalar fields pass through.
+
+        Used by hybrid-batch samplers (MIXED + spec-dec) that apply
+        different sampler ops to a prefix vs suffix of rows. Only ``slice``
+        is supported — int indexing would yield 0-dim tensors and break
+        downstream gathers.
+
+        ``is_all_greedy`` is inherited from the parent; when ``top_ks`` is
+        populated the slice refines it from the sliced tensor (one GPU
+        sync, only on the disagg slice path).
+        """
+        if not isinstance(s, slice):
+            raise TypeError(
+                f"SamplingBatchInfo only supports slice indexing, got {type(s).__name__}"
+            )
+
+        def _slice(t):
+            return t[s] if t is not None else None
+
+        return dataclasses.replace(
+            self,
+            temperatures=_slice(self.temperatures),
+            top_ps=_slice(self.top_ps),
+            top_ks=_slice(self.top_ks),
+            min_ps=_slice(self.min_ps),
+            is_all_greedy=self.is_all_greedy,
+            req_pool_indices=_slice(self.req_pool_indices),
+            vocab_mask=_slice(self.vocab_mask),
+            grammars=_slice(self.grammars),
+        )
+
+    @classmethod
+    def from_schedule_batch(
+        cls, batch: ScheduleBatch, vocab_size: int
+    ) -> SamplingBatchInfo:
+        reqs = batch.reqs
+        device = batch.device
+        temperatures = torch.tensor(
+            [r.sampling_params.temperature for r in reqs], dtype=torch.float
+        ).to(device, non_blocking=True)
+        top_ps = torch.tensor(
+            [r.sampling_params.top_p for r in reqs], dtype=torch.float
+        ).to(device, non_blocking=True)
+        top_ks = torch.tensor(
+            [r.sampling_params.top_k for r in reqs], dtype=torch.int32
+        ).to(device, non_blocking=True)
+        min_ps = torch.tensor(
+            [r.sampling_params.min_p for r in reqs], dtype=torch.float
+        ).to(device, non_blocking=True)
+
+        ret = cls(
+            temperatures=temperatures,
+            top_ps=top_ps,
+            top_ks=top_ks,
+            min_ps=min_ps,
+            is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
+            vocab_size=vocab_size,
+            device=device,
+        )
+        return ret

python/tokenspeed/runtime/sampling/sampling_params.py

@@ -0,0 +1,243 @@
+# 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 parameters for text generation."""
+
+import zlib
+from typing import Any
+
+_SAMPLING_EPS = 1e-6
+
+# Sentinel for "top_k is disabled" (sample from whole vocab). We rewrite
+# top_k=-1 (API convention) to this value so downstream code can pass it
+# unchanged to top_k kernels that expect a positive cutoff.
+_TOP_K_DISABLED = 1 << 30
+
+# Upper bound the fused top-k + top-p kernel sorts in its on-chip top-K
+# branch. Requests with a finite top_k above this would silently fall through
+# to the top-p-only branch, so reject them at request time. Must stay in sync
+# with K_TOPK_MAX in fused_topk_topp.h.
+_TOP_K_FUSED_MAX = 128
+
+
+class SamplingParams:
+    """
+    The sampling parameters.
+
+    See docs/backend/sampling_params.md or
+    https://docs.tokenspeed.ai/backend/sampling_params.html
+    for the documentation.
+    """
+
+    def __init__(
+        self,
+        max_new_tokens: int | None = None,
+        stop: str | list[str] | None = None,
+        stop_token_ids: list[int] | None = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        min_p: float = 0.0,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        repetition_penalty: float = 1.0,
+        min_new_tokens: int = 0,
+        json_schema: str | None = None,
+        regex: str | None = None,
+        ebnf: str | None = None,
+        structural_tag: str | None = None,
+        ignore_eos: bool = False,
+        skip_special_tokens: bool = True,
+        spaces_between_special_tokens: bool = True,
+        no_stop_trim: bool = False,
+        thinking_budget: int | None = None,
+        custom_params: dict[str, Any] | None = None,
+        stream_interval: int | None = None,
+        logit_bias: dict[str, float] | None = None,
+        seed: int | None = None,
+        # vLLM-style output logprobs. None = off; 0 = the sampled (generated)
+        # token's logprob at each output position. Other values are rejected by
+        # verify().
+        logprobs: int | None = None,
+        # OpenAI-compat: `n` is a request-level fanout (number of choices)
+        # that the serving layer forwards on every sampling_params dict.
+        # TokenSpeed does not multiplex a single request into n completions,
+        # so accept and ignore.
+        n: int = 1,
+    ) -> None:
+        self.max_new_tokens = max_new_tokens
+        self.stop_strs = stop
+        if stop_token_ids:
+            self.stop_token_ids = set(stop_token_ids)
+        else:
+            self.stop_token_ids = None
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+        self.min_p = min_p
+        self.frequency_penalty = frequency_penalty
+        self.presence_penalty = presence_penalty
+        self.repetition_penalty = repetition_penalty
+        self.min_new_tokens = min_new_tokens
+        self.regex = regex
+        self.json_schema = json_schema
+        self.ebnf = ebnf
+        self.structural_tag = structural_tag
+        self.ignore_eos = ignore_eos
+        self.skip_special_tokens = skip_special_tokens
+        self.spaces_between_special_tokens = spaces_between_special_tokens
+        self.no_stop_trim = no_stop_trim
+        self.custom_params = custom_params
+        self.thinking_budget = thinking_budget
+        self.stream_interval = stream_interval
+        self.logit_bias = logit_bias
+        self.seed = seed
+        self.logprobs = logprobs
+
+        # Process some special cases
+        if self.temperature < _SAMPLING_EPS:
+            # top_k = 1 means greedy sampling
+            self.temperature = 1.0
+            self.top_k = 1
+        if self.top_k == -1:
+            self.top_k = _TOP_K_DISABLED
+
+    def verify(self, vocab_size: int) -> None:
+        if self.temperature < 0.0:
+            raise ValueError(
+                f"temperature must be non-negative, got {self.temperature}."
+            )
+        if not 0.0 < self.top_p <= 1.0:
+            raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
+        if not 0.0 <= self.min_p <= 1.0:
+            raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.")
+        if self.top_k < -1 or self.top_k == 0:
+            raise ValueError(
+                f"top_k must be -1 (disable), or at least 1, " f"got {self.top_k}."
+            )
+        if self.top_k != _TOP_K_DISABLED and self.top_k >= _TOP_K_FUSED_MAX:
+            raise ValueError(
+                f"top_k must be < {_TOP_K_FUSED_MAX} (fused kernel limit) "
+                f"or -1 (disable), got {self.top_k}."
+            )
+        if not -2.0 <= self.frequency_penalty <= 2.0:
+            raise ValueError(
+                "frequency_penalty must be in [-2, 2], got "
+                f"{self.frequency_penalty}."
+            )
+        if not -2.0 <= self.presence_penalty <= 2.0:
+            raise ValueError(
+                "presence_penalty must be in [-2, 2], got " f"{self.presence_penalty}."
+            )
+        if not 0.0 <= self.repetition_penalty <= 2.0:
+            raise ValueError(
+                "repetition_penalty must be in (0, 2], got "
+                f"{self.repetition_penalty}."
+            )
+        if not 0 <= self.min_new_tokens:
+            raise ValueError(
+                f"min_new_tokens must be in (0, max_new_tokens], got "
+                f"{self.min_new_tokens}."
+            )
+        if self.max_new_tokens is not None:
+            if self.max_new_tokens < 0:
+                raise ValueError(
+                    f"max_new_tokens must be at least 0, got {self.max_new_tokens}."
+                )
+            if not self.min_new_tokens <= self.max_new_tokens:
+                raise ValueError(
+                    f"min_new_tokens must be in (0, max_new_tokens({self.max_new_tokens})], got "
+                    f"{self.min_new_tokens}."
+                )
+        if self.logit_bias is not None:
+            for token_id in self.logit_bias:
+                if not 0 <= int(token_id) < vocab_size:
+                    raise ValueError(
+                        f"logit_bias must has keys in [0, {vocab_size - 1}], got "
+                        f"{token_id}."
+                    )
+
+        if self.logprobs is not None and self.logprobs != 0:
+            # Only the sampled token's logprob (logprobs=0) is materialized;
+            # top-k (>0) and full-vocab (-1) output logprobs are not supported.
+            raise ValueError(
+                f"logprobs={self.logprobs} is not supported; use logprobs=0 "
+                "(the sampled token's logprob)."
+            )
+
+        grammars = [
+            self.json_schema,
+            self.regex,
+            self.ebnf,
+        ]  # since mutually exclusive, only one can be set
+        if sum(x is not None for x in grammars) > 1:
+            raise ValueError("Only one of regex, json_schema, or ebnf can be set.")
+
+    def requested_features(self) -> "set[str]":
+        """Return the set of backend-facing feature names this request needs.
+
+        `temperature`, `top_k`, `top_p`, `min_p` each appear only when the
+        corresponding field is not at its neutral default. Used by
+        SamplingBackend.register() to reject requests asking for features
+        the active backend does not implement."""
+        out: set[str] = set()
+        if abs(self.temperature - 1.0) > _SAMPLING_EPS:
+            out.add("temperature")
+        # top_k=_TOP_K_DISABLED and top_k=1 (greedy short-circuit from __init__) are neutral.
+        if self.top_k != _TOP_K_DISABLED and self.top_k != 1:
+            out.add("top_k")
+        if self.top_p < 1.0:
+            out.add("top_p")
+        if self.min_p > 0.0:
+            out.add("min_p")
+        if self.frequency_penalty != 0.0:
+            out.add("frequency_penalty")
+        if self.presence_penalty != 0.0:
+            out.add("presence_penalty")
+        if self.repetition_penalty != 1.0:
+            out.add("repetition_penalty")
+        if self.logit_bias:
+            out.add("logit_bias")
+        return out
+
+    def resolve_seed(self, rid: str) -> None:
+        """If the caller didn't supply a seed, derive one deterministically
+        from rid. Called at the single request-materialization point so all
+        TP/DP ranks agree on the seed."""
+        if self.seed is None:
+            self.seed = zlib.crc32(rid.encode("utf-8")) & 0xFFFFFFFF
+
+    def normalize(self, tokenizer) -> None:
+        # Process stop strings
+        if self.stop_strs is None:
+            self.stop_strs = []
+            self.stop_str_max_len = 0
+        else:
+            if isinstance(self.stop_strs, str):
+                self.stop_strs = [self.stop_strs]
+
+            stop_str_max_len = 0
+            for stop_str in self.stop_strs:
+                if tokenizer is not None:
+                    stop_str_ids = tokenizer.encode(stop_str, add_special_tokens=False)
+                    stop_str_max_len = max(stop_str_max_len, len(stop_str_ids))
+                else:
+                    stop_str_max_len = max(stop_str_max_len, len(stop_str))
+            self.stop_str_max_len = stop_str_max_len

python/tokenspeed/runtime/sampling/utils.py

@@ -0,0 +1,86 @@
+# 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.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel.torch_compile import get_compiler_backend
+
+from tokenspeed.runtime.utils import crash_on_warnings, get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+# Smallest positive value per dtype, used as the lower bound for `uniform_`
+# draws that feed rejection-sampling kernels. A coin of exact 0 silently
+# accepts a zero-probability draft in `chain_speculative_sampling_target_only`
+# (the kernel condition `coin <= target_prob / threshold_acc` reduces to
+# `0 <= 0`), so the coin must be strictly positive.
+COIN_EPS = {
+    torch.float32: torch.finfo(torch.float32).tiny,
+    torch.bfloat16: torch.finfo(torch.bfloat16).tiny,
+}
+
+
+def coin_eps(dtype: torch.dtype) -> float:
+    """Lower bound for uniform coin draws of the given dtype. See COIN_EPS."""
+    return COIN_EPS[dtype]
+
+
+def nan_guard_logits(
+    logits: torch.Tensor,
+    enable_nan_detection: bool,
+) -> torch.Tensor:
+    """Replace NaNs with -1e5 and optionally crash; no-op when detection is disabled."""
+    if not enable_nan_detection:
+        return logits
+
+    if not torch.any(torch.isnan(logits)):
+        return logits
+
+    logger.warning("Detected errors during sampling! NaN in the logits.")
+    logits = torch.where(torch.isnan(logits), torch.full_like(logits, -1e5), logits)
+    if crash_on_warnings():
+        raise ValueError("Detected errors during sampling! NaN in the logits.")
+    return logits
+
+
+@torch.compile(dynamic=True, backend=get_compiler_backend())
+def gather_token_logprobs_torch(
+    logits: torch.Tensor,
+    tokens: torch.Tensor,
+) -> torch.Tensor:
+    """Per-row log_softmax(logits)[tokens]. Fuses cast + online softmax + gather
+    into one Triton kernel sequence so the full [B, V] log_softmax matrix is
+    never materialized."""
+    raw_logprobs = torch.log_softmax(logits.float(), dim=-1)
+    return raw_logprobs.gather(-1, tokens.unsqueeze(-1)).squeeze(-1)
+
+
+@torch.compile(dynamic=True, backend=get_compiler_backend())
+def top_p_normalize_probs_torch(
+    probs: torch.Tensor,
+    top_ps: torch.Tensor,
+) -> torch.Tensor:
+    """Pure-torch nucleus renorm — used by the prefill-logprob path."""
+    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    return torch.zeros_like(probs_sort).scatter_(-1, probs_idx, probs_sort)