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

# 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)