lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/logits_processor.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.

"""Logits processing."""

import dataclasses

import torch
import triton
import triton.language as tl
from tokenspeed_kernel.ops.communication.triton import all_gather_inner, create_state
from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax
from tokenspeed_kernel.ops.sampling.cute_dsl import (
    create_dist_argmax_state,
    distributed_argmax,
)
from tokenspeed_kernel.platform import current_platform
from torch import nn

from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor
from tokenspeed.runtime.distributed.process_group_manager import (
    process_group_manager as pg_manager,
)
from tokenspeed.runtime.execution.context import ForwardContext
from tokenspeed.runtime.execution.forward_batch_info import (
    CaptureHiddenMode,
    ForwardMode,
)
from tokenspeed.runtime.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding,
)
from tokenspeed.runtime.sampling.dp_sampling_config import (
    DpSamplingRuntimeConfig,
)
from tokenspeed.runtime.sampling.logits_layout import (
    LogitsLayoutExecutor,
    LogitsLayoutPlan,
)
from tokenspeed.runtime.utils import get_colorful_logger

logger = get_colorful_logger(__name__)


@dataclasses.dataclass
class LogitsProcessorOutput:
    ## Part 1: This part will be assigned in python/tokenspeed/runtime/layers/logits_processor.py::LogitsProcessor
    # The logits of the next tokens.       shape: [#seq, vocab_size]
    next_token_logits: torch.Tensor
    # Used when ``do_argmax=True``.   shape: [#seq]
    next_token_ids: torch.Tensor | None = None
    # Used by speculative decoding.
    # The last hidden layers
    hidden_states: torch.Tensor | None = None
    logits_layout_plan: LogitsLayoutPlan | None = None

    ## Part 2: Populated by the active SamplingBackend during sample()/verify().
    # The logprobs of the next tokens.                              shape: [#seq]
    next_token_logprobs: torch.Tensor | None = None
    # The logprobs and ids of the top-k tokens in output positions. shape: [#seq, k]
    next_token_top_logprobs_val: list | None = None
    next_token_top_logprobs_idx: list | None = None
    # The logprobs and ids of the requested token ids in output positions. shape: [#seq, n] (n is the number of requested token ids)
    next_token_token_ids_logprobs_val: list | None = None
    next_token_token_ids_logprobs_idx: list | None = None

    ## Part 3: Prefill-only. This part will be assigned in python/tokenspeed/runtime/layers/logits_processor.py::LogitsProcessor
    # The logprobs of input tokens.        shape: [#token]
    input_token_logprobs: torch.Tensor | None = None
    # The logprobs and ids of the top-k tokens in input positions.  shape: [#seq, #token, k]
    input_top_logprobs_val: list = None
    input_top_logprobs_idx: list = None
    # The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids)
    input_token_ids_logprobs_val: list | None = None
    input_token_ids_logprobs_idx: list | None = None


@dataclasses.dataclass
class LogitsMetadata:
    forward_mode: ForwardMode
    capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL
    gather_ids: torch.Tensor | None = None

    extend_return_logprob: bool = False
    extend_return_top_logprob: bool = False
    extend_token_ids_logprob: bool = False
    extend_seq_lens_cpu: list[int] | None = None
    extend_logprob_start_lens_cpu: list[int] | None = None
    extend_logprob_pruned_lens_cpu: list[int] | None = None
    top_logprobs_nums: list[int] | None = None
    extend_input_logprob_token_ids_gpu: torch.Tensor | None = None
    token_ids_logprobs: list[list[int]] | None = None

    # logits and logprobs post processing
    temp_scaled_logprobs: bool = False
    temperature: torch.Tensor = None
    top_p_normalized_logprobs: bool = False
    top_p: torch.Tensor = None

    # DP attention metadata. Not needed when DP attention is not used.
    # Number of tokens in the request.
    global_num_tokens_gpu: torch.Tensor | None = None
    # The start position of local hidden states.
    dp_local_start_pos: torch.Tensor | None = None
    dp_local_num_tokens: torch.Tensor | None = None
    gathered_buffer: torch.Tensor | None = None
    # Buffer to gather logits from all ranks.
    forward_batch_gathered_buffer: torch.Tensor | None = None

    @classmethod
    def from_forward_context(cls, ctx: ForwardContext):
        return cls(
            forward_mode=ctx.forward_mode,
            capture_hidden_mode=ctx.capture_hidden_mode,
            gather_ids=ctx.gather_ids,
        )


_FUSED_LM_HEAD_GEMM = None


def _get_fused_lm_head_gemm():
    """Lazily import the fused lm_head GEMM kernel.

    The kernel is only present when tokenspeed-kernel was built with a
    compatible nvcc. Cache a sentinel when unavailable so we fall back
    to ``torch.matmul`` silently on subsequent calls.
    """
    global _FUSED_LM_HEAD_GEMM
    if _FUSED_LM_HEAD_GEMM is not None:
        return _FUSED_LM_HEAD_GEMM
    if not current_platform().is_nvidia:
        _FUSED_LM_HEAD_GEMM = (None, None)
        return _FUSED_LM_HEAD_GEMM
    try:
        from tokenspeed_kernel.thirdparty.cuda.lm_head_gemm import (
            lm_head_gemm,
            should_use_fused,
        )

        _FUSED_LM_HEAD_GEMM = (should_use_fused, lm_head_gemm)
    except Exception:
        _FUSED_LM_HEAD_GEMM = (None, None)
    return _FUSED_LM_HEAD_GEMM


def _lm_head_matmul(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """Compute ``hidden_states @ weight.T``.

    Routes to the fused ``lm_head_gemm`` when the shape matches a compiled
    template and the bench-driven perf gate accepts (``should_use_fused``).
    Otherwise falls back to ``torch.matmul``.

    Only enabled for Kimi (``model_type == "kimi_k2"``) at the call site —
    on DSv3 the fused kernel's PDL launch surface caused a downstream EAGLE3
    spec decode AR regression that we have not characterised end-to-end; on
    Kimi the perf win is the largest and the regression has not been
    reproduced, so we gate the fused path to Kimi only.
    """
    cast_hidden = hidden_states.to(weight.dtype)
    should_use_fused, lm_head_gemm = _get_fused_lm_head_gemm()
    if should_use_fused is not None and should_use_fused(cast_hidden, weight):
        return lm_head_gemm(cast_hidden, weight, enable_pdl=True)
    return torch.matmul(cast_hidden, weight.T)


class LogitsProcessor(nn.Module):

    _LOGITS_AG_MAX_TOKENS = 128
    _LOGITS_AG_STATE_UNINITIALIZED = object()
    _LOGITS_AG_STATES = {}

    _LOGITS_DIST_ARGMAX_MAX_TOKENS = 8192
    _LOGITS_DIST_ARGMAX_UNINITIALIZED = object()
    _LOGITS_DIST_ARGMAX_STATES = {}

    def __init__(
        self,
        config,
        skip_all_gather: bool = False,
        do_argmax: bool = False,
        logit_scale: float | None = None,
        tp_rank: int | None = None,
        tp_size: int | None = None,
        tp_group: tuple[int, ...] | None = None,
    ):
        super().__init__()
        self.config = config
        self.skip_all_gather = skip_all_gather
        self.do_argmax = do_argmax
        self.dp_sampling_enabled = False
        self.dp_num_tokens_per_req = 1
        self.dp_sampling_min_bs = 0
        self.logit_scale = logit_scale
        self._logits_layout_executor: LogitsLayoutExecutor | None = None

        if tp_rank is None:
            if tp_size is not None or tp_group is not None:
                raise ValueError("tp_size and tp_group require tp_rank.")
            tp_rank, tp_size = 0, 1
        elif tp_size is None:
            raise ValueError("tp_size is required when tp_rank is provided.")
        if not 0 <= tp_rank < tp_size:
            raise ValueError(f"Invalid tensor-parallel rank: {tp_rank}/{tp_size}.")
        if tp_size != 1 and tp_group is None:
            raise ValueError("tp_group is required when tp_size > 1.")
        self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group

        self._all_gather_state = self._LOGITS_AG_STATE_UNINITIALIZED
        self._dist_argmax_state = self._LOGITS_DIST_ARGMAX_UNINITIALIZED

        self.final_logit_softcapping = getattr(
            self.config, "final_logit_softcapping", None
        )
        if (
            self.final_logit_softcapping is not None
            and self.final_logit_softcapping < 0
        ):
            self.final_logit_softcapping = None

        # Gate the fused lm_head GEMM to Kimi only. See ``_lm_head_matmul``.
        self._use_fused_lm_head = getattr(self.config, "model_type", None) == "kimi_k2"

    def configure_dp_logits_layout(self, runtime: DpSamplingRuntimeConfig) -> None:
        if (
            not runtime.enabled
            or runtime.topology is None
            or runtime.min_bs is None
            or runtime.max_bucket_bs is None
            or runtime.vocab_size is None
            or runtime.device is None
        ):
            raise RuntimeError("enabled DP sampling runtime is incomplete")
        topology = runtime.topology
        self.dp_sampling_enabled = True
        self.dp_num_tokens_per_req = runtime.num_tokens_per_req
        self.dp_sampling_min_bs = runtime.min_bs
        self._logits_layout_executor = LogitsLayoutExecutor(
            tp_rank=topology.tp_rank,
            tp_size=topology.tp_size,
            tp_group=topology.tp_group,
            max_bucket_bs=runtime.max_bucket_bs,
            num_tokens_per_req=runtime.num_tokens_per_req,
            vocab_size=runtime.vocab_size,
            device=runtime.device,
        )

    def _resolve_logits_layout_plan(
        self,
        hidden_states: torch.Tensor,
        logits_metadata: LogitsMetadata,
    ) -> LogitsLayoutPlan | None:
        if not self.dp_sampling_enabled:
            return None
        if not logits_metadata.forward_mode.is_decode():
            return None
        n = self.dp_num_tokens_per_req
        rows = hidden_states.shape[0]
        if rows % n != 0:
            raise ValueError(f"hidden_states have {rows} rows, not divisible by N={n}")
        effective_bs = rows // n
        bucket_bs = ((effective_bs + self.tp_size - 1) // self.tp_size) * self.tp_size
        if effective_bs < self.dp_sampling_min_bs:
            return None
        return LogitsLayoutPlan(
            effective_bs=effective_bs,
            bucket_bs=bucket_bs,
            tp_size=self.tp_size,
            num_tokens_per_req=n,
        )

    def _init_all_gather_state(self, lm_head: VocabParallelEmbedding):
        if not current_platform().is_nvidia:
            return None

        if self.tp_size == 1 or self.skip_all_gather:
            return None

        vocab_padded = lm_head.weight.size(0) * self.tp_size
        if vocab_padded % (self.tp_size * 8) != 0:
            return None

        key = (self.tp_group, vocab_padded)
        if key not in self._LOGITS_AG_STATES:
            self._LOGITS_AG_STATES[key] = create_state(
                group=pg_manager.get_process_group("nccl", self.tp_group),
                rank_in_group=self.tp_rank,
                max_tokens=self._LOGITS_AG_MAX_TOKENS,
                hidden_size=vocab_padded,
            )
        return self._LOGITS_AG_STATES[key]

    def _init_dist_argmax_state(self, lm_head: VocabParallelEmbedding):
        if not current_platform().is_nvidia:
            return None

        if self.tp_size == 1 or self.skip_all_gather or self.dp_sampling_enabled:
            return None

        vocab_per_rank = lm_head.weight.size(0)
        if vocab_per_rank * self.tp_size != self.config.vocab_size:
            return None  # padded vocab: sharded argmax could pick a pad column
        if vocab_per_rank < 4096 or vocab_per_rank % 32 != 0:
            return None  # below the kernel's vocab floor / alignment

        key = (self.tp_group, vocab_per_rank)
        if key not in self._LOGITS_DIST_ARGMAX_STATES:
            self._LOGITS_DIST_ARGMAX_STATES[key] = create_dist_argmax_state(
                group=pg_manager.get_process_group("nccl", self.tp_group),
                rank_in_group=self.tp_rank,
                max_M=self._LOGITS_DIST_ARGMAX_MAX_TOKENS,
                dtype=lm_head.weight.dtype,
                device=lm_head.weight.device,
                skip_ping_pong=True,
            )

        return self._LOGITS_DIST_ARGMAX_STATES[key]

    def forward(
        self,
        input_ids,
        hidden_states,
        lm_head: VocabParallelEmbedding,
        logits_metadata: LogitsMetadata,
        aux_hidden_states: torch.Tensor | None = None,
    ) -> LogitsProcessorOutput:
        # Get the last hidden states and last logits for the next token prediction
        if not logits_metadata.extend_return_logprob:
            gather_ids = logits_metadata.gather_ids
            if gather_ids is not None:
                # Shapes align iff midlayer already pruned to one row per request
                # (draft first-step reduce). Other paths emit [N, H] with N > bs.
                if gather_ids.shape[0] == hidden_states.shape[0]:
                    pruned_states = hidden_states
                    if aux_hidden_states is not None:
                        aux_pruned_states = list(aux_hidden_states)
                else:
                    pruned_states = hidden_states[gather_ids]
                    if aux_hidden_states is not None:
                        aux_pruned_states = [h[gather_ids] for h in aux_hidden_states]
            else:
                if logits_metadata.forward_mode.is_extend_or_mixed():
                    raise RuntimeError(
                        "EXTEND/MIXED forward must set gather_ids on ForwardContext"
                    )
                pruned_states = hidden_states
                if aux_hidden_states is not None:
                    aux_pruned_states = list(aux_hidden_states)

            sample_indices = None
            input_logprob_indices = None
        else:
            # Input logprobs are required.
            # Find 3 different indices.
            # 1. pruned_states: hidden states that we want logprobs from.
            # 2. sample_indices: Indices that have sampled tokens.
            # 3. input_logprob_indices: Indices that have input logprob tokens.
            sample_index_pt = -1
            sample_indices = []
            input_logprob_indices_pt = 0
            input_logprob_indices = []
            pt, pruned_states = 0, []
            for extend_logprob_start_len, extend_len in zip(
                logits_metadata.extend_logprob_start_lens_cpu,
                logits_metadata.extend_seq_lens_cpu,
            ):
                # It can happen in chunked prefill. We still need to sample 1 token,
                # But we don't want to include it in input logprob.
                if extend_len == extend_logprob_start_len:
                    start_len = extend_logprob_start_len - 1
                else:
                    start_len = extend_logprob_start_len

                # We always need at least 1 token to sample because that's required
                # by a caller.
                if extend_len <= start_len:
                    raise RuntimeError("extend_len must be greater than start_len.")
                pruned_states.append(hidden_states[pt + start_len : pt + extend_len])
                pt += extend_len
                sample_index_pt += extend_len - start_len
                sample_indices.append(sample_index_pt)
                input_logprob_indices.extend(
                    [
                        input_logprob_indices_pt + i
                        for i in range(extend_len - extend_logprob_start_len)
                    ]
                )
                input_logprob_indices_pt += extend_len - start_len

            pruned_states = torch.cat(pruned_states)
            sample_indices = torch.tensor(
                sample_indices, device=pruned_states.device, dtype=torch.int64
            )
            input_logprob_indices = torch.tensor(
                input_logprob_indices, device=pruned_states.device, dtype=torch.int64
            )

        # Compute logits for both input and sampled tokens.
        logits_layout_plan = self._resolve_logits_layout_plan(
            pruned_states, logits_metadata
        )
        logits = self._get_logits(
            pruned_states, lm_head, logits_metadata, plan=logits_layout_plan
        )
        sampled_logits = (
            logits[sample_indices] if sample_indices is not None else logits
        )

        hidden_states_to_store: torch.Tensor | None = None
        if logits_metadata.capture_hidden_mode.need_capture():
            if logits_metadata.capture_hidden_mode.is_full():
                if aux_hidden_states is not None:
                    aux_hidden_states = (
                        aux_hidden_states[0]
                        if len(aux_hidden_states) == 1
                        else torch.cat(aux_hidden_states, dim=-1)
                    )
                    hidden_states_to_store = aux_hidden_states
                else:
                    hidden_states_to_store = hidden_states
            elif logits_metadata.capture_hidden_mode.is_last():
                # Get the last token hidden states. If sample_indices is None,
                # pruned states only contain the last tokens already.
                if aux_hidden_states is not None:
                    aux_pruned_states = (
                        aux_pruned_states[0]
                        if len(aux_pruned_states) == 1
                        else torch.cat(aux_pruned_states, dim=-1)
                    )
                    hidden_states_to_store = (
                        aux_pruned_states[sample_indices]
                        if sample_indices is not None
                        else aux_pruned_states
                    )
                else:
                    hidden_states_to_store = (
                        pruned_states[sample_indices]
                        if sample_indices is not None
                        else pruned_states
                    )
            else:
                raise RuntimeError("Should never reach")

        if not logits_metadata.extend_return_logprob:
            # Decode mode or extend mode without return_logprob.
            # Greedy draft path: emit token ids here, fusing the cross-rank
            # vocab reduction into the argmax when gated on.
            next_token_ids = self._argmax(sampled_logits) if self.do_argmax else None
            return LogitsProcessorOutput(
                next_token_logits=sampled_logits,
                next_token_ids=next_token_ids,
                hidden_states=hidden_states_to_store,
                logits_layout_plan=logits_layout_plan,
            )
        else:
            input_logprobs = logits[input_logprob_indices]
            del hidden_states, logits

            # Normalize the logprob w/o temperature, top-p
            pruned_lens = torch.tensor(
                logits_metadata.extend_logprob_pruned_lens_cpu,
                device=input_logprobs.device,
            )
            if logits_metadata.temp_scaled_logprobs:
                logits_metadata.temperature = torch.repeat_interleave(
                    logits_metadata.temperature.view(-1),
                    pruned_lens,
                ).view(-1, 1)
            if logits_metadata.top_p_normalized_logprobs:
                logits_metadata.top_p = torch.repeat_interleave(
                    logits_metadata.top_p,
                    pruned_lens,
                )
            input_logprobs = self.compute_temp_top_p_normalized_logprobs(
                input_logprobs, logits_metadata
            )

            # Get the logprob of top-k tokens
            if logits_metadata.extend_return_top_logprob:
                (
                    input_top_logprobs_val,
                    input_top_logprobs_idx,
                ) = self.get_top_logprobs(input_logprobs, logits_metadata)
            else:
                input_top_logprobs_val = input_top_logprobs_idx = None

            # Get the logprob of given token id
            if logits_metadata.extend_token_ids_logprob:
                (
                    input_token_ids_logprobs_val,
                    input_token_ids_logprobs_idx,
                ) = self.get_token_ids_logprobs(input_logprobs, logits_metadata)
            else:
                input_token_ids_logprobs_val = input_token_ids_logprobs_idx = None

            input_token_logprobs = input_logprobs[
                torch.arange(input_logprobs.shape[0], device=input_logprobs.device),
                logits_metadata.extend_input_logprob_token_ids_gpu,
            ]

            return LogitsProcessorOutput(
                next_token_logits=sampled_logits,
                input_token_logprobs=input_token_logprobs,
                input_top_logprobs_val=input_top_logprobs_val,
                input_top_logprobs_idx=input_top_logprobs_idx,
                hidden_states=hidden_states_to_store,
                input_token_ids_logprobs_val=input_token_ids_logprobs_val,
                input_token_ids_logprobs_idx=input_token_ids_logprobs_idx,
            )

    def _get_logits(
        self,
        hidden_states: torch.Tensor,
        lm_head: VocabParallelEmbedding,
        logits_metadata: LogitsMetadata,
        embedding_bias: torch.Tensor | None = None,
        plan: LogitsLayoutPlan | None = None,
    ) -> torch.Tensor:
        """Get logits from hidden_states.

        If sampled_logits_only is True, it means hidden_states only contain the
        last position (e.g., extend without input logprobs). The caller should
        guarantee the given hidden_states follow this constraint.
        """
        dp_sampling = plan is not None
        if dp_sampling and not self.dp_sampling_enabled:
            raise RuntimeError(
                "DP logits layout plan was provided but LogitsProcessor was not "
                "configured with dp_sampling"
            )

        if dp_sampling and self.skip_all_gather:
            if self._logits_layout_executor is None:
                raise RuntimeError(
                    "dp_sampling logits layout executor is not configured"
                )
            hidden_states = self._logits_layout_executor.slice_hidden_states(
                hidden_states, plan
            )

        if hasattr(lm_head, "weight"):
            if self._use_fused_lm_head:
                logits = _lm_head_matmul(hidden_states, lm_head.weight)
            else:
                logits = torch.matmul(
                    hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
                )
        else:
            # GGUF models
            logits = lm_head.linear_method.apply(lm_head, hidden_states, embedding_bias)

        if self.logit_scale is not None:
            logits.mul_(self.logit_scale)

        if dp_sampling and not self.skip_all_gather:
            if self._logits_layout_executor is None:
                raise RuntimeError(
                    "dp_sampling logits layout executor is not configured"
                )
            logits = self._logits_layout_executor.swap_batch_vocab(logits, plan)

        elif not dp_sampling and self.tp_size > 1 and not self.skip_all_gather:
            if self.do_argmax:
                if self._dist_argmax_state is self._LOGITS_DIST_ARGMAX_UNINITIALIZED:
                    self._dist_argmax_state = self._init_dist_argmax_state(lm_head)

                if (
                    self._dist_argmax_state is not None
                    and not self.final_logit_softcapping
                    and logits.size(0) <= self._LOGITS_DIST_ARGMAX_MAX_TOKENS
                ):
                    return logits

            if self._all_gather_state is self._LOGITS_AG_STATE_UNINITIALIZED:
                self._all_gather_state = self._init_all_gather_state(lm_head)

            if (
                self._all_gather_state is not None
                and logits.size(0) <= self._LOGITS_AG_MAX_TOKENS
            ):
                # skip_entry_sync=True assumes other sync points existing between two all_gather_inner calls.
                logits = all_gather_inner(
                    self._all_gather_state,
                    logits,
                    tp_hidden_dim=logits.size(-1) * self.tp_size,
                    skip_entry_sync=True,
                    safe=False,
                )
            else:
                num_rows = logits.size(0)
                local_vocab_size = logits.size(1)
                gathered_logits = torch.empty(
                    self.tp_size * num_rows,
                    local_vocab_size,
                    dtype=logits.dtype,
                    device=logits.device,
                )
                all_gather_into_tensor(gathered_logits, logits, self.tp_group)
                logits = (
                    gathered_logits.view(self.tp_size, num_rows, local_vocab_size)
                    .transpose(0, 1)
                    .contiguous()
                    .view(num_rows, local_vocab_size * self.tp_size)
                )

        logits = logits[:, : self.config.vocab_size].contiguous()

        if self.final_logit_softcapping:
            fused_softcap_generic(logits, self.final_logit_softcapping)

        return logits

    def _argmax(self, logits: torch.Tensor) -> torch.Tensor:
        if (
            self._dist_argmax_state
            not in (self._LOGITS_DIST_ARGMAX_UNINITIALIZED, None)
            and not self.final_logit_softcapping
            and logits.size(0) <= self._LOGITS_DIST_ARGMAX_MAX_TOKENS
        ):
            _, idx = distributed_argmax(self._dist_argmax_state, logits)
            return idx
        else:
            return sampling_argmax(logits)

    @staticmethod
    def get_top_logprobs(all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata):
        max_k = max(logits_metadata.top_logprobs_nums)
        ret = all_logprobs.topk(max_k, dim=1)
        values = ret.values.tolist()
        indices = ret.indices.tolist()

        input_top_logprobs_val, input_top_logprobs_idx = [], []

        pt = 0
        for k, pruned_len in zip(
            logits_metadata.top_logprobs_nums,
            logits_metadata.extend_logprob_pruned_lens_cpu,
        ):
            if pruned_len <= 0:
                input_top_logprobs_val.append([])
                input_top_logprobs_idx.append([])
                continue

            input_top_logprobs_val.append(
                [values[pt + j][:k] for j in range(pruned_len)]
            )
            input_top_logprobs_idx.append(
                [indices[pt + j][:k] for j in range(pruned_len)]
            )
            pt += pruned_len

        return input_top_logprobs_val, input_top_logprobs_idx

    @staticmethod
    def get_token_ids_logprobs(
        all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata
    ):
        input_token_ids_logprobs_val, input_token_ids_logprobs_idx = [], []
        pt = 0
        for token_ids, pruned_len in zip(
            logits_metadata.token_ids_logprobs,
            logits_metadata.extend_logprob_pruned_lens_cpu,
        ):
            if pruned_len <= 0:
                input_token_ids_logprobs_val.append([])
                input_token_ids_logprobs_idx.append([])
                continue

            input_token_ids_logprobs_val.append(
                [all_logprobs[pt + j, token_ids].tolist() for j in range(pruned_len)]
            )
            input_token_ids_logprobs_idx.append([token_ids for _ in range(pruned_len)])
            pt += pruned_len

        return input_token_ids_logprobs_val, input_token_ids_logprobs_idx

    @staticmethod
    def compute_temp_top_p_normalized_logprobs(
        last_logits: torch.Tensor, logits_metadata: LogitsMetadata
    ) -> torch.Tensor:
        """
        compute logprobs for the output token from the given logits.

        Returns:
            torch.Tensor: logprobs from logits
        """
        last_logits = last_logits.float()
        # Scale logits if temperature scaling is enabled
        if logits_metadata.temp_scaled_logprobs:
            last_logits = last_logits / logits_metadata.temperature

        # Normalize logprobs if top_p normalization is enabled
        #  only normalize logprobs when top_p is set and not equal to 1.0
        if (
            logits_metadata.top_p_normalized_logprobs
            and (logits_metadata.top_p != 1.0).any()
        ):
            from tokenspeed.runtime.sampling.utils import top_p_normalize_probs_torch

            probs = torch.softmax(last_logits, dim=-1)
            del last_logits
            probs = top_p_normalize_probs_torch(probs, logits_metadata.top_p)
            return torch.log(probs)
        else:
            return torch.nn.functional.log_softmax(last_logits, dim=-1)


@triton.jit
def fused_softcap_kernel(
    full_logits_ptr,
    softcapping_value,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0).to(tl.int64)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load values
    x = tl.load(full_logits_ptr + offsets, mask=mask).to(tl.float32)

    # Perform operations in-place
    x = x / softcapping_value

    # Stable tanh form; the exp ratio overflows to inf/inf for large logits.
    x = 2 * tl.sigmoid(2 * x) - 1

    x = x * softcapping_value

    # Store result
    tl.store(full_logits_ptr + offsets, x, mask=mask)


def fused_softcap(full_logits, final_logit_softcapping):
    n_elements = full_logits.numel()
    BLOCK_SIZE = 1024
    grid = ((n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE, 1, 1)

    fused_softcap_kernel[grid](
        full_logits_ptr=full_logits,
        softcapping_value=final_logit_softcapping,
        n_elements=n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    return full_logits


def fused_softcap_generic(full_logits, final_logit_softcapping):
    return fused_softcap(full_logits, final_logit_softcapping)