[perf] Compute input logprobs without materializing the full-vocab log-softmax (#31958)

python/sglang/srt/environ.py

@@ -1028,6 +1028,10 @@ class Envs:
     SGLANG_LOGPROB_CHUNK_SIZE = EnvIntWithAlias(
         2048, deprecated_name="SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"
     )
+    # Compute input logprobs from logits via per-row logsumexp instead of
+    # materializing the full-vocab log-softmax. Escape hatch only; the two
+    # paths are mathematically identical.
+    SGLANG_ENABLE_FAST_INPUT_LOGPROBS = EnvBool(True)
 
     # Tool-Call behavior
     SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF)

python/sglang/srt/layers/logprob_processor.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
 import torch
 
 from sglang.srt.environ import envs
+from sglang.srt.runtime_context import get_exec
 
 if TYPE_CHECKING:
     from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput
@@ -61,6 +62,26 @@ class LogprobResult:
             )
 
 
+def compute_row_log_normalizer(
+    logits: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-row ``(max, logsumexp - max)`` in fp32.
+
+    Consumers compute ``logprob[i] = (logit[i] - max) - log_sum``, the same
+    shift-invariant order as log-softmax; a single absolute normalizer would
+    round the log_sum term away for rows with a large common offset.
+    """
+    if logits.is_cuda:
+        from sglang.srt.layers.logsumexp import row_logsumexp
+
+        return row_logsumexp(logits)
+    x = logits.float()
+    row_max = x.amax(dim=-1)
+    row_log_sum = torch.logsumexp(x - row_max[:, None], dim=-1)
+    row_log_sum = torch.where(row_max.isinf(), 0.0, row_log_sum)
+    return row_max, row_log_sum
+
+
 def get_top_logprobs_raw(
     logprobs: torch.Tensor,
     top_logprobs_nums: List[int],
@@ -170,25 +191,42 @@ def get_top_logprobs_chunk(
     top_logprobs_val: List,
     top_logprobs_idx: List,
     split_pruned_len: int,
+    log_normalizer: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    precomputed_topk: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
 ) -> int:
     """Get top-k logprobs for each sequence in the chunk.
 
     Args:
-        logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
+        logprobs: Log probabilities tensor of shape [seq_len, vocab_size].
+            With ``log_normalizer`` set, raw logits instead; top-k runs on the
+            logits (same order) and values are normalized by subtraction.
         top_k_nums: List of top-k numbers for each sequence
         pruned_lens: List of pruned lengths for each sequence
         top_logprobs_val: List to store top-k logprob values
         top_logprobs_idx: List to store top-k token indices
         split_pruned_len: Length of pruned tokens from previous chunk
+        log_normalizer: Per-row (max, logsumexp - max) of the logits
+        precomputed_topk: (raw fp32 top values, indices) from the fused
+            logsumexp+top-k kernel, sorted with lowest-index tie-breaking
 
     Returns:
         int: Number of remaining tokens to process in next chunk
     """
     # Empty chunks still walk the slice to emit placeholder entries.
     max_k = max(top_k_nums)
-    ret = logprobs.topk(max_k, dim=1)
+    if log_normalizer is not None:
-    values = ret.values.tolist()
+        row_max, row_log_sum = log_normalizer
-    indices = ret.indices.tolist()
+        if precomputed_topk is not None:
+            values_tensor, indices_tensor = precomputed_topk
+        else:
+            values_tensor, indices_tensor = logprobs.topk(max_k, dim=1)
+        values_tensor = (values_tensor.float() - row_max[:, None]) - row_log_sum[
+            :, None
+        ]
+    else:
+        values_tensor, indices_tensor = logprobs.topk(max_k, dim=1)
+    values = values_tensor.tolist()
+    indices = indices_tensor.tolist()
 
     pt = 0
     next_split_pruned_len = 0
@@ -240,21 +278,27 @@ def get_token_ids_logprobs_chunk(
     token_ids_logprobs_val: List,
     token_ids_logprobs_idx: List,
     split_pruned_len: int = 0,
+    log_normalizer: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
 ):
     """Get token_ids logprobs for each sequence in the chunk.
 
     Args:
-        logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
+        logprobs: Log probabilities tensor of shape [seq_len, vocab_size].
+            With ``log_normalizer`` set, raw logits instead; gathered rows are
+            normalized by subtraction.
         token_ids_logprobs: List of token IDs for each sequence
         pruned_lens: List of pruned lengths for each sequence
         token_ids_logprobs_val: List to store token logprob values
         token_ids_logprobs_idx: List to store token indices
         split_pruned_len: Length of pruned tokens from previous chunk
+        log_normalizer: Per-row (max, logsumexp - max) of the logits
 
     Returns:
         int: Number of remaining tokens to process in next chunk
     """
     # Empty chunks still walk the slice to emit placeholder entries.
+    if log_normalizer is not None:
+        row_max, row_log_sum = log_normalizer
     pt = 0
     next_split_pruned_len = 0
     for n, (token_ids, pruned_len) in enumerate(
@@ -285,7 +329,10 @@ def get_token_ids_logprobs_chunk(
                 next_split_pruned_len = split_pruned_len + j
                 break
             if token_ids is not None:
-                val.append(logprobs[pt + j, token_ids].tolist())
+                row = logprobs[pt + j, token_ids]
+                if log_normalizer is not None:
+                    row = (row.float() - row_max[pt + j]) - row_log_sum[pt + j]
+                val.append(row.tolist())
                 idx.append(token_ids)
 
         # Split-sequence continuations extend; everyone else owns a fresh
@@ -366,6 +413,18 @@ def compute_spec_v2_logprobs(
         )
 
 
+def _deterministic_inference_enabled() -> bool:
+    """True when serving with --enable-deterministic-inference.
+
+    Fails open: bare constructions (unit tests) have no published config
+    namespaces, and plain serving is the not-deterministic case.
+    """
+    try:
+        return bool(get_exec().deterministic.enable_deterministic_inference)
+    except ValueError:
+        return False
+
+
 class InputLogprobProcessor:
     """Input (prefill) logprob processing: single-pass or chunked.
 
@@ -379,6 +438,15 @@ class InputLogprobProcessor:
         self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get()
         # chunk size for logprobs processing
         self.logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get()
+        # Compute input logprobs from logits + logsumexp, skipping the
+        # full-vocab log-softmax materialization. Deterministic inference
+        # keeps the exact log_softmax path: the fused logsumexp reduces in a
+        # different order, which breaks the prefill/decode logprob
+        # bit-identity that mode guarantees.
+        self.enable_fast_input_logprobs = (
+            envs.SGLANG_ENABLE_FAST_INPUT_LOGPROBS.get()
+            and not _deterministic_inference_enabled()
+        )
 
     def forward(
         self,
@@ -447,6 +515,15 @@ class InputLogprobProcessor:
         split_len_topk = 0
         split_len_token_ids = 0
 
+        fused_kernel, fused_max_k = None, 0
+        if self.enable_fast_input_logprobs and pruned_states.is_cuda:
+            from sglang.srt.layers.logsumexp import (
+                FUSED_TOPK_MAX_K,
+                row_logsumexp_topk,
+            )
+
+            fused_kernel, fused_max_k = row_logsumexp_topk, FUSED_TOPK_MAX_K
+
         for i in range(num_chunks):
             start_idx = i * chunk_size
             end_idx = min((i + 1) * chunk_size, total_size)
@@ -496,12 +573,8 @@ class InputLogprobProcessor:
                 sampled_logits[chunk_sample_mask] = chunk_logits[chunk_sample_indices]
 
             # Zero-logprob-row chunks still need the per-sequence bookkeeping below.
-            # Compute the logprobs of the chunk. Free the raw logits before the
-            # out-of-place log_softmax: keeping all three alive is a 3x peak,
-            # which OOMs when the single chunk covers a large batch.
             chunk_logprobs = chunk_logits[chunk_indices]
             del chunk_logits
-            chunk_logprobs = torch.nn.functional.log_softmax(chunk_logprobs, dim=-1)
 
             # End at the last row inside the chunk; token_to_seq_idx[end_idx]
             # belongs to the next chunk and would emit its sequence twice.
@@ -509,6 +582,33 @@ class InputLogprobProcessor:
                 token_to_seq_idx[start_idx], token_to_seq_idx[end_idx - 1] + 1
             )
 
+            chunk_precomputed_topk = None
+            if self.enable_fast_input_logprobs:
+                # Every consumer below needs only small gathers / top-k plus a
+                # per-row normalizer, so keep the raw logits and skip the
+                # full-vocab log-softmax materialization entirely. When top-k
+                # is requested, the fused kernel produces the normalizer and
+                # the top-k in the same single read of the logits.
+                max_k = (
+                    max(logits_metadata.top_logprobs_nums[chunk_slice])
+                    if logits_metadata.extend_return_top_logprob
+                    else 0
+                )
+                if 0 < max_k <= fused_max_k:
+                    row_max, row_log_sum, top_vals, top_idx = fused_kernel(
+                        chunk_logprobs, max_k
+                    )
+                    chunk_log_normalizer = (row_max, row_log_sum)
+                    chunk_precomputed_topk = (top_vals, top_idx)
+                else:
+                    chunk_log_normalizer = compute_row_log_normalizer(chunk_logprobs)
+            else:
+                # Free the raw logits before the out-of-place log_softmax:
+                # keeping all three alive is a 3x peak, which OOMs when the
+                # single chunk covers a large batch.
+                chunk_log_normalizer = None
+                chunk_logprobs = torch.nn.functional.log_softmax(chunk_logprobs, dim=-1)
+
             # Get the logprob of top-k tokens
             if logits_metadata.extend_return_top_logprob:
                 top_k_nums = logits_metadata.top_logprobs_nums[chunk_slice]
@@ -522,6 +622,8 @@ class InputLogprobProcessor:
                     top_logprobs_val,
                     top_logprobs_idx,
                     split_len_topk,
+                    log_normalizer=chunk_log_normalizer,
+                    precomputed_topk=chunk_precomputed_topk,
                 )
 
             # Get the logprob of given token id
@@ -537,6 +639,7 @@ class InputLogprobProcessor:
                     token_ids_logprobs_val,
                     token_ids_logprobs_idx,
                     split_len_token_ids,
+                    log_normalizer=chunk_log_normalizer,
                 )
 
             # Get the logprob of the requested token ids
@@ -544,6 +647,11 @@ class InputLogprobProcessor:
                 torch.arange(chunk_logprobs.shape[0], device=chunk_logprobs.device),
                 logits_metadata.extend_input_logprob_token_ids_gpu[mask_indices],
             ]
+            if chunk_log_normalizer is not None:
+                row_max, row_log_sum = chunk_log_normalizer
+                chunk_token_logprobs = (
+                    chunk_token_logprobs.float() - row_max
+                ) - row_log_sum
             token_logprobs.append(chunk_token_logprobs)
             # Free before the next chunk's logits (bf16 + fp32) materialize.
             del chunk_logprobs

python/sglang/srt/layers/logsumexp.py

@@ -0,0 +1,298 @@
+"""Single-pass online row logsumexp, optionally fused with a small top-k.
+
+Reads the input exactly once with fp32 accumulation and writes per-row
+``(max, log_sum_exp - max)`` fp32 pairs. Keeping the two terms separate lets
+callers compute ``logprob[i] = (logit[i] - max) - log_sum`` the same
+shift-invariant way log-softmax does: folding them into one absolute
+normalizer would round the small log-sum term away whenever rows carry a
+large common offset.
+
+``row_logsumexp_topk`` additionally maintains a running top-k
+(k <= FUSED_TOPK_MAX_K) in registers during the same pass, so callers that
+need both the normalizer and the top-k tokens pay for a single read of the
+input.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# Maximum k for the fused top-k kernel; larger k should fall back to
+# torch.topk. The per-block selection cost grows with k: measured on GB300
+# at vocab=151936 the fused kernel beats a separate logsumexp + top-k up to
+# k~6 and drops below plain torch.topk near k=32.
+FUSED_TOPK_MAX_K = 8
+
+
+@triton.jit
+def _accumulate_block_lse(x, m_i, l_i):
+    """Fold one block of fp32 values into the running (max, sum_exp).
+
+    Per-block max first: a single dominant logit stays exact instead of
+    being renormalized against a stale running max. A +/-inf block max
+    makes exp(x - m_blk) nan; the true normalized sum_exp there is 1.
+    Shared by both kernels so their (max, log_sum) outputs are bitwise
+    identical by construction.
+    """
+    m_blk = tl.max(x)
+    l_blk = tl.sum(tl.exp(x - m_blk))
+    l_blk = tl.where((m_blk == float("inf")) | (m_blk == float("-inf")), 1.0, l_blk)
+    m_i, l_i = _combine_lse(m_i, l_i, m_blk, l_blk)
+    return m_i, l_i, m_blk
+
+
+@triton.jit
+def _combine_lse(m_a, l_a, m_b, l_b):
+    """Merge two (max, sum_exp) accumulators.
+
+    The tl.where guards keep fully -inf inputs nan-free: exp(-inf - -inf)
+    would otherwise poison the sum.
+    """
+    m_new = tl.maximum(m_a, m_b)
+    l_new = l_a * tl.exp(tl.where(m_a == m_new, 0.0, m_a - m_new)) + l_b * tl.exp(
+        tl.where(m_b == m_new, 0.0, m_b - m_new)
+    )
+    return m_new, l_new
+
+
+@triton.jit
+def _row_logsumexp_kernel(
+    x_ptr,
+    out_max_ptr,
+    out_log_sum_ptr,
+    row_stride,
+    col_stride,
+    # int rather than tl.constexpr so triton does not unroll the loop
+    num_cols,
+    BLOCK_N: tl.constexpr,
+):
+    # One program per row.
+    row = tl.program_id(0).to(tl.int64)
+    row_ptr = x_ptr + row * row_stride
+
+    m_i = float("-inf")
+    l_i = 0.0
+    for start in range(0, num_cols, BLOCK_N):
+        offs = start + tl.arange(0, BLOCK_N)
+        x = tl.load(
+            row_ptr + offs * col_stride, mask=offs < num_cols, other=float("-inf")
+        ).to(tl.float32)
+        m_i, l_i, _ = _accumulate_block_lse(x, m_i, l_i)
+
+    tl.store(out_max_ptr + row, m_i)
+    tl.store(out_log_sum_ptr + row, tl.log(l_i))
+
+
+def row_logsumexp(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-row ``(max, logsumexp - max)`` of a 2D tensor, as fp32."""
+    assert x.ndim == 2
+    assert x.dtype in (torch.float16, torch.bfloat16, torch.float32)
+
+    num_rows, num_cols = x.shape
+    out_max = torch.empty(num_rows, device=x.device, dtype=torch.float32)
+    out_log_sum = torch.empty(num_rows, device=x.device, dtype=torch.float32)
+    if num_rows == 0:
+        return out_max, out_log_sum
+    if num_cols == 0:
+        return out_max.fill_(float("-inf")), out_log_sum.zero_()
+
+    BLOCK_N = triton.next_power_of_2(min(num_cols, 16384))
+    _row_logsumexp_kernel[(num_rows,)](
+        x,
+        out_max,
+        out_log_sum,
+        x.stride(0),
+        x.stride(1),
+        num_cols,
+        BLOCK_N=BLOCK_N,
+        num_warps=8,
+    )
+    return out_max, out_log_sum
+
+
+@triton.jit
+def _fpval_to_key(x):
+    """fp32 bits (as uint32) -> unsigned key with float order.
+
+    The standard sign-flip transform, specialized to 32-bit from
+    kernels.ops.moe.gate_topk's generic helpers.
+    """
+    return x ^ tl.where((x & 0x80000000) != 0, 0xFFFFFFFF, 0x80000000)
+
+
+@triton.jit
+def _key_to_fpval(x):
+    return x ^ tl.where((x & 0x80000000) == 0, 0xFFFFFFFF, 0x80000000)
+
+
+@triton.jit
+def _pack_key(vals, idxs):
+    """Pack fp32 value + int32 index into one order-encoding int64 key.
+
+    Key order == candidate order: higher value wins, value ties go to the
+    LOWER index. The fp32 bits are mapped through the sign-flip
+    transform so unsigned bit order matches float order, then placed above
+    the complemented index; all keys land in [0, 2**63), so signed int64
+    comparison is the candidate comparison. The 2147483647 (int32 max)
+    index sentinel yields the smallest key for a given value, so "empty"
+    (-inf, sentinel) candidates lose against everything, including genuine
+    -inf lanes.
+    """
+    sortable = _fpval_to_key(vals.to(tl.uint32, bitcast=True))
+    return (sortable.to(tl.int64) << 31) | (2147483647 - idxs).to(tl.int64)
+
+
+@triton.jit
+def _unpack_val(keys):
+    sortable = ((keys >> 31) & 0xFFFFFFFF).to(tl.uint32)
+    return _key_to_fpval(sortable).to(tl.float32, bitcast=True)
+
+
+@triton.jit
+def _unpack_idx(keys):
+    return 2147483647 - (keys & 0x7FFFFFFF).to(tl.int32)
+
+
+@triton.jit
+def _row_logsumexp_topk_kernel(
+    x_ptr,
+    out_max_ptr,
+    out_log_sum_ptr,
+    out_top_vals_ptr,
+    out_top_idx_ptr,
+    row_stride,
+    col_stride,
+    # int rather than tl.constexpr so triton does not unroll the loop
+    num_cols,
+    K: tl.constexpr,
+    # K rounded up to a power of 2: register tensors need pow2 shapes. The
+    # padding slots hold (-inf, sentinel) and lose every comparison.
+    K_PAD: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    # One program per row.
+    row = tl.program_id(0).to(tl.int64)
+    row_ptr = x_ptr + row * row_stride
+
+    m_i = float("-inf")
+    l_i = 0.0
+    slots = tl.arange(0, K_PAD)
+    # Running top-K as packed keys, sorted descending. Slots beyond the
+    # genuine entries hold the (-inf, sentinel) key, which loses everything.
+    run_keys = _pack_key(
+        tl.full((K_PAD,), float("-inf"), tl.float32),
+        tl.full((K_PAD,), 2147483647, tl.int32),
+    )
+    kth_key = tl.min(run_keys)
+
+    for start in range(0, num_cols, BLOCK_N):
+        offs = start + tl.arange(0, BLOCK_N)
+        x = tl.load(
+            row_ptr + offs * col_stride, mask=offs < num_cols, other=float("-inf")
+        ).to(tl.float32)
+
+        m_i, l_i, m_blk = _accumulate_block_lse(x, m_i, l_i)
+
+        # Top-k maintenance. Blocks are visited in ascending column order and
+        # the wrapper guarantees the first block holds >= K in-row lanes, so
+        # after the first block every running slot carries a genuine index
+        # smaller than any later block's. A later block whose max does not
+        # strictly beat the running k-th value therefore cannot contribute:
+        # on equal values the running entry's lower index wins. The first
+        # block must merge unconditionally to displace the sentinel slots.
+        # View the block as K_PAD contiguous segments of SEG lanes (a
+        # row-major reshape, so no data movement) and surface up to K_PAD
+        # candidates per pass: the best available lane of each segment, via
+        # minor-axis reductions. Random rows spread their top-K across
+        # segments, so a contributing block usually needs one merge pass;
+        # the worst case (all K in one segment) needs K.
+        SEG: tl.constexpr = BLOCK_N // K_PAD
+        x2 = tl.reshape(x, (K_PAD, SEG))
+        i2 = start + tl.arange(0, K_PAD)[:, None] * SEG + tl.arange(0, SEG)[None, :]
+        avail2 = i2 < num_cols
+        go = (start == 0) | (m_blk > _unpack_val(kth_key))
+        while go:
+            xa2 = tl.where(avail2, x2, float("-inf"))
+            cand_vals = tl.max(xa2, axis=1)
+            attain = avail2 & (xa2 == cand_vals[:, None])
+            cand_idxs = tl.min(tl.where(attain, i2, 2147483647), axis=1)
+            # Surfaced lanes are spent either way: winners enter the running
+            # top-K and losers lost against a set that only ever improves,
+            # so neither can matter again.
+            avail2 = avail2 & (i2 != cand_idxs[:, None])
+
+            cand_keys = _pack_key(cand_vals, cand_idxs)
+            go = tl.max(cand_keys) > kth_key
+            if go:
+                # Sorted-sequence merge: max(A, reverse(B)) holds exactly the
+                # top-K_PAD of the union (both inputs sorted descending, keys
+                # distinct), then re-sort. Compare-exchange networks only; no
+                # reductions.
+                cand_sorted = tl.sort(cand_keys, descending=True)
+                merged = tl.maximum(run_keys, tl.flip(cand_sorted, 0))
+                run_keys = tl.sort(merged, descending=True)
+                kth_key = tl.max(tl.where(slots == K - 1, run_keys, -1))
+
+    tl.store(out_max_ptr + row, m_i)
+    tl.store(out_log_sum_ptr + row, tl.log(l_i))
+    out_mask = slots < K
+    tl.store(out_top_vals_ptr + row * K + slots, _unpack_val(run_keys), mask=out_mask)
+    tl.store(
+        out_top_idx_ptr + row * K + slots,
+        _unpack_idx(run_keys).to(tl.int64),
+        mask=out_mask,
+    )
+
+
+def row_logsumexp_topk(
+    x: torch.Tensor, k: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Per-row ``(max, logsumexp - max, top_vals, top_idx)`` in one read.
+
+    ``top_vals`` are the RAW input values upcast to fp32 (normalize with
+    ``(v - max) - log_sum``; the subtraction is shift-invariant), sorted
+    descending with value ties broken by the lowest index. ``top_idx`` is
+    int64 like torch.topk's. The (max, log_sum) pair is bitwise identical to
+    ``row_logsumexp`` on the same input.
+
+    k is a tl.constexpr, so each distinct k compiles its own kernel: k is the
+    per-batch max of the requested top-logprob counts and rarely varies
+    within a deployment, while padding every launch to k=32 would make the
+    common k<=2 case do 16x the selection work.
+    """
+    assert x.ndim == 2
+    assert x.dtype in (torch.float16, torch.bfloat16, torch.float32)
+    num_rows, num_cols = x.shape
+    # Like torch.topk, selecting more entries than a row holds is an error.
+    assert 1 <= k <= min(FUSED_TOPK_MAX_K, num_cols), (k, num_cols)
+
+    out_max = torch.empty(num_rows, device=x.device, dtype=torch.float32)
+    out_log_sum = torch.empty(num_rows, device=x.device, dtype=torch.float32)
+    top_vals = torch.empty((num_rows, k), device=x.device, dtype=torch.float32)
+    top_idx = torch.empty((num_rows, k), device=x.device, dtype=torch.int64)
+    if num_rows == 0:
+        return out_max, out_log_sum, top_vals, top_idx
+
+    # Floor 2: K_PAD=1 would degenerate the kernel's [K_PAD, SEG] view
+    # (and flip/sort of single-lane tensors).
+    k_pad = max(2, triton.next_power_of_2(k))
+    # next_power_of_2(num_cols) >= num_cols >= k when num_cols < 16384, and
+    # 16384 > FUSED_TOPK_MAX_K otherwise: the kernel's first block always
+    # sees at least k in-row lanes. The K_PAD floor keeps SEG =
+    # BLOCK_N // K_PAD >= 1 for tiny vocabularies.
+    BLOCK_N = max(triton.next_power_of_2(min(num_cols, 16384)), k_pad)
+    _row_logsumexp_topk_kernel[(num_rows,)](
+        x,
+        out_max,
+        out_log_sum,
+        top_vals,
+        top_idx,
+        x.stride(0),
+        x.stride(1),
+        num_cols,
+        K=k,
+        K_PAD=k_pad,
+        BLOCK_N=BLOCK_N,
+        num_warps=8,
+    )
+    return out_max, out_log_sum, top_vals, top_idx

test/registered/unit/layers/test_logprob_fast_input.py

@@ -0,0 +1,404 @@
+"""Fast input-logprob path must match the log-softmax reference.
+
+The fast path (SGLANG_ENABLE_FAST_INPUT_LOGPROBS) computes token / top-k /
+token-ids logprobs directly from logits with a per-row logsumexp normalizer,
+never materializing the full-vocab log-softmax. Same math, so results must
+agree with the reference path to floating-point tolerance, with identical
+top-k indices, across chunk splits and heterogeneous per-sequence params.
+"""
+
+import itertools
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.layers.logprob_processor import (
+    InputLogprobProcessor,
+    compute_row_log_normalizer,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=30, suite="base-a-test-cpu")
+
+VOCAB = 11
+# Heterogeneous per-sequence parameters; uniform ones hide misalignment.
+TOPK_CYCLE = [2, 0, 3]
+# [] is a valid probe set distinct from None (opt-out).
+TOKEN_IDS_CYCLE = [[0, 3], None, [1], []]
+
+
+def _build_batch(seq_specs, dtype, vocab=VOCAB):
+    """seq_specs: list of (extend_len, logprob_start_len). Mirrors
+    LogitsProcessor._get_pruned_states for the extend-with-logprobs path."""
+    pruned_rows = []
+    token_to_seq_idx = []
+    sample_indices = []
+    input_logprob_indices = []
+    pruned_lens = []
+    sample_pt = -1
+    lp_pt = 0
+    for idx, (extend_len, start) in enumerate(seq_specs):
+        eff_start = start - 1 if extend_len == start else start
+        rows = extend_len - eff_start
+        pruned_rows.append(torch.randn(rows, vocab).to(dtype))
+        token_to_seq_idx.extend([idx] * rows)
+        sample_pt += rows
+        sample_indices.append(sample_pt)
+        n_lp = extend_len - start
+        input_logprob_indices.extend([lp_pt + i for i in range(n_lp)])
+        lp_pt += rows
+        pruned_lens.append(n_lp)
+    metadata = SimpleNamespace(
+        extend_return_top_logprob=True,
+        extend_token_ids_logprob=True,
+        top_logprobs_nums=[TOPK_CYCLE[i % 3] for i in range(len(seq_specs))],
+        extend_logprob_pruned_lens_cpu=pruned_lens,
+        extend_input_logprob_token_ids_gpu=torch.zeros(
+            len(input_logprob_indices), dtype=torch.int64
+        ),
+        token_ids_logprobs=[
+            TOKEN_IDS_CYCLE[i % len(TOKEN_IDS_CYCLE)] for i in range(len(seq_specs))
+        ],
+    )
+    return (
+        torch.cat(pruned_rows),
+        torch.tensor(sample_indices, dtype=torch.int64),
+        torch.tensor(input_logprob_indices, dtype=torch.int64),
+        token_to_seq_idx,
+        metadata,
+    )
+
+
+def _run(proc, batch, fast, chunk_size):
+    pruned_states, sample_indices, input_logprob_indices, t2s, metadata = batch
+    proc.enable_logprobs_chunk = chunk_size is not None
+    proc.logprobs_chunk_size = chunk_size if chunk_size is not None else 10**9
+    proc.enable_fast_input_logprobs = fast
+
+    def get_logits_fn(states, lm_head, logits_metadata, **kwargs):
+        return states
+
+    return proc.forward(
+        pruned_states=pruned_states,
+        sample_indices=sample_indices,
+        input_logprob_indices=input_logprob_indices,
+        token_to_seq_idx=t2s,
+        lm_head=None,
+        get_logits_fn=get_logits_fn,
+        logits_metadata=metadata,
+    )
+
+
+def _assert_nested_close(test, ref, got, label, rtol, atol):
+    test.assertEqual(_shape_of(ref), _shape_of(got), label)
+    ref_flat = _flatten(ref)
+    got_flat = _flatten(got)
+    if ref_flat:
+        torch.testing.assert_close(
+            torch.tensor(ref_flat, dtype=torch.float64),
+            torch.tensor(got_flat, dtype=torch.float64),
+            rtol=rtol,
+            atol=atol,
+            msg=label,
+        )
+
+
+def _flatten(nested):
+    if isinstance(nested, list):
+        return [x for item in nested for x in _flatten(item)]
+    return [nested]
+
+
+def _shape_of(nested):
+    if isinstance(nested, list):
+        return [_shape_of(item) for item in nested]
+    return None
+
+
+class TestFastInputLogprobs(CustomTestCase):
+    def _sweep(self, dtype, rtol, atol):
+        torch.manual_seed(0)
+        proc = InputLogprobProcessor()
+        # (extend_len, start); start == extend_len is the degenerate
+        # zero-logprob-row shape.
+        menu = [(1, 1), (3, 0), (4, 1), (5, 5), (6, 2)]
+        tried = 0
+        for n_seqs in (1, 2, 3):
+            for combo in itertools.product(menu, repeat=n_seqs):
+                batch = _build_batch(list(combo), dtype)
+                for chunk_size in (None, 1, 2, 3, 5):
+                    tried += 1
+                    ref, ref_sampled = _run(proc, batch, False, chunk_size)
+                    got, got_sampled = _run(proc, batch, True, chunk_size)
+                    label = f"specs={list(combo)} chunk={chunk_size} dtype={dtype}"
+                    # Top-k order comes from the same values shifted by a
+                    # per-row constant, so indices must match exactly.
+                    self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
+                    self.assertEqual(
+                        ref.token_ids_logprobs_idx, got.token_ids_logprobs_idx, label
+                    )
+                    _assert_nested_close(
+                        self,
+                        ref.top_logprobs_val,
+                        got.top_logprobs_val,
+                        label,
+                        rtol,
+                        atol,
+                    )
+                    _assert_nested_close(
+                        self,
+                        ref.token_ids_logprobs_val,
+                        got.token_ids_logprobs_val,
+                        label,
+                        rtol,
+                        atol,
+                    )
+                    torch.testing.assert_close(
+                        ref.token_logprobs.float(),
+                        got.token_logprobs.float(),
+                        rtol=rtol,
+                        atol=atol,
+                        msg=label,
+                    )
+                    torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
+        self.assertGreater(tried, 100)
+
+    def test_fast_matches_reference_fp32(self):
+        self._sweep(torch.float32, rtol=1e-5, atol=1e-5)
+
+    def test_fast_matches_float64_truth_bf16(self):
+        # bf16 log_softmax rounds near-ties together, so the reference path's
+        # top-k ORDER is not reproducible from raw logits; validate the fast
+        # path against float64 ground truth instead. The fast path only
+        # rounds at the bf16 logits themselves (normalizer is fp32), so it
+        # sits much closer to the truth than bf16 resolution.
+        torch.manual_seed(0)
+        proc = InputLogprobProcessor()
+        menu = [(1, 1), (3, 0), (4, 1), (5, 5), (6, 2)]
+        for n_seqs in (1, 2, 3):
+            for combo in itertools.product(menu, repeat=n_seqs):
+                batch = _build_batch(list(combo), torch.bfloat16)
+                pruned_states, _, input_logprob_indices, _, metadata = batch
+                truth = torch.log_softmax(pruned_states.double(), dim=-1)[
+                    input_logprob_indices
+                ]
+                for chunk_size in (None, 2, 5):
+                    got, _ = _run(proc, batch, True, chunk_size)
+                    label = f"specs={list(combo)} chunk={chunk_size}"
+                    self._assert_rows_match_truth(
+                        got, truth, metadata, label, atol=1e-4
+                    )
+
+    def _assert_rows_match_truth(self, got, truth, metadata, label, atol):
+        pt = 0
+        for s, pruned_len in enumerate(metadata.extend_logprob_pruned_lens_cpu):
+            if pruned_len <= 0:
+                self.assertEqual(got.top_logprobs_val[s], [], label)
+                continue
+            k = metadata.top_logprobs_nums[s]
+            probe_ids = metadata.token_ids_logprobs[s]
+            for j in range(pruned_len):
+                row_truth = truth[pt + j]
+                vals = got.top_logprobs_val[s][j]
+                idxs = got.top_logprobs_idx[s][j]
+                self.assertEqual(len(vals), k, label)
+                for v, i in zip(vals, idxs):
+                    self.assertAlmostEqual(
+                        v, row_truth[i].item(), delta=atol, msg=label
+                    )
+                if probe_ids is not None:
+                    probe_vals = got.token_ids_logprobs_val[s][j]
+                    for v, i in zip(probe_vals, probe_ids):
+                        self.assertAlmostEqual(
+                            v, row_truth[i].item(), delta=atol, msg=label
+                        )
+            pt += pruned_len
+
+    def test_shift_invariant_large_offset(self):
+        # Regression: a large common fp32 offset must not round the log-sum
+        # term away. Uniform logits at 1e8 have true logprob -log(vocab).
+        for device in ("cpu", "cuda") if torch.cuda.is_available() else ("cpu",):
+            logits = torch.full((4, 1000), 1e8, dtype=torch.float32, device=device)
+            row_max, row_log_sum = compute_row_log_normalizer(logits)
+            logprob = (logits[:, 0].float() - row_max) - row_log_sum
+            expected = -torch.log(torch.tensor(1000.0))
+            torch.testing.assert_close(
+                logprob.cpu(), expected.expand(4), rtol=1e-5, atol=1e-5
+            )
+
+    def test_fast_path_emits_fp32_logprobs(self):
+        # Chosen dtype policy: the fast path returns fp32 token logprobs
+        # regardless of logits dtype (the normalizer is fp32, so fp32 is the
+        # true precision of the result), while the log_softmax path keeps
+        # the logits dtype. Runs on CPU CI so the policy is pinned even
+        # where the CUDA kernels never execute.
+        proc = InputLogprobProcessor()
+        batch = _build_batch([(4, 1), (3, 0)], torch.bfloat16)
+        got, _ = _run(proc, batch, True, None)
+        self.assertEqual(got.token_logprobs.dtype, torch.float32)
+        ref, _ = _run(proc, batch, False, None)
+        self.assertEqual(ref.token_logprobs.dtype, torch.bfloat16)
+
+    def test_logsumexp_module_imports(self):
+        # Runs on CPU CI too: catches import rot in the CUDA-only kernel
+        # module, whose imports otherwise only execute on GPU machines.
+        import sglang.srt.layers.logsumexp  # noqa: F401
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_tiny_shapes(self):
+        from sglang.srt.layers.logsumexp import row_logsumexp_topk
+
+        for rows, cols, k in ((1, 1, 1), (3, 2, 2), (2, 3, 1), (2, 5, 5)):
+            logits = torch.randn(rows, cols, device="cuda")
+            got_m, got_ls, got_v, got_i = row_logsumexp_topk(logits, k)
+            ref_v, ref_i = torch.topk(logits, k, dim=-1, sorted=True)
+            self.assertTrue(torch.equal(got_v, ref_v.float()), (rows, cols, k))
+            self.assertTrue(torch.equal(got_i, ref_i), (rows, cols, k))
+            torch.testing.assert_close(
+                got_m + got_ls,
+                torch.logsumexp(logits.double(), dim=-1).float(),
+                rtol=1e-5,
+                atol=1e-5,
+            )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_run_to_run_deterministic(self):
+        from sglang.srt.layers.logsumexp import row_logsumexp_topk
+
+        torch.manual_seed(0)
+        logits = torch.randn(512, 151936, dtype=torch.bfloat16, device="cuda")
+        a = row_logsumexp_topk(logits, 5)
+        b = row_logsumexp_topk(logits, 5)
+        self.assertTrue(all(torch.equal(p, q) for p, q in zip(a, b)))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_logsumexp_topk_matches_torch(self):
+        from sglang.srt.layers.logsumexp import row_logsumexp, row_logsumexp_topk
+
+        torch.manual_seed(0)
+        for k in (1, 2, 3, 5, 8):
+            for dtype in (torch.float32, torch.bfloat16):
+                logits = torch.randn(64, 151936, dtype=dtype, device="cuda")
+                got_m, got_ls, got_v, got_i = row_logsumexp_topk(logits, k)
+                # The (max, log_sum) pair is bitwise the non-fused kernel's.
+                ref_m, ref_ls = row_logsumexp(logits)
+                self.assertTrue(torch.equal(got_m, ref_m), (k, dtype))
+                self.assertTrue(torch.equal(got_ls, ref_ls), (k, dtype))
+                ref_v, ref_i = torch.topk(logits, k, dim=-1, sorted=True)
+                self.assertTrue(torch.equal(got_v, ref_v.float()), (k, dtype))
+                if dtype == torch.float32:
+                    # fp32 randn is tie-free w.h.p.: indices match exactly.
+                    self.assertTrue(torch.equal(got_i, ref_i), (k, dtype))
+                else:
+                    # bf16 has value ties; indices must point at the values.
+                    self.assertTrue(
+                        torch.equal(logits.gather(-1, got_i).float(), got_v),
+                        (k, dtype),
+                    )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_topk_tie_break_is_lowest_index(self):
+        from sglang.srt.layers.logsumexp import row_logsumexp_topk
+
+        logits = torch.zeros(1, 1000, device="cuda")
+        logits[0, [7, 3, 500]] = 5.0
+        _, _, _, got_i = row_logsumexp_topk(logits, 3)
+        self.assertEqual(got_i[0].tolist(), [3, 7, 500])
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_topk_inf_rows(self):
+        from sglang.srt.layers.logsumexp import row_logsumexp_topk
+
+        logits = torch.full((3, 1000), float("-inf"), device="cuda")
+        logits[1, 3] = 2.5
+        _, _, got_v, got_i = row_logsumexp_topk(logits.bfloat16(), 2)
+        self.assertEqual(got_i[0].tolist(), [0, 1])
+        self.assertEqual(got_i[1, 0].item(), 3)
+        self.assertFalse(got_v.isnan().any().item())
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fast_path_end_to_end_on_cuda(self):
+        # Exercises the fused-kernel integration inside _forward_by_chunk
+        # (the CPU sweeps only cover the torch fallbacks), including the
+        # k > FUSED_TOPK_MAX_K fallback.
+        torch.manual_seed(0)
+        proc = InputLogprobProcessor()
+        for k_override in (None, 20):
+            # k=20 exceeds FUSED_TOPK_MAX_K, exercising the torch fallback;
+            # it needs a vocab that can supply 20 entries.
+            batch = _build_batch([(4, 1), (6, 2), (3, 0)], torch.float32, vocab=64)
+            pruned_states, sample_indices, lp_indices, t2s, metadata = batch
+            if k_override is not None:
+                metadata.top_logprobs_nums = [k_override] * 3
+            batch = (
+                pruned_states.cuda(),
+                sample_indices.cuda(),
+                lp_indices.cuda(),
+                t2s,
+                metadata,
+            )
+            metadata.extend_input_logprob_token_ids_gpu = (
+                metadata.extend_input_logprob_token_ids_gpu.cuda()
+            )
+            for chunk_size in (None, 2, 5):
+                ref, _ = _run(proc, batch, False, chunk_size)
+                got, _ = _run(proc, batch, True, chunk_size)
+                label = f"k_override={k_override} chunk={chunk_size}"
+                self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
+                _assert_nested_close(
+                    self, ref.top_logprobs_val, got.top_logprobs_val, label, 1e-5, 1e-5
+                )
+                torch.testing.assert_close(
+                    ref.token_logprobs, got.token_logprobs, rtol=1e-5, atol=1e-5
+                )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_row_logsumexp_kernel_matches_reference(self):
+        from sglang.srt.layers.logsumexp import row_logsumexp
+
+        torch.manual_seed(0)
+        for rows, cols in ((0, 128), (3, 0), (1, 1), (7, 1000), (64, 151936)):
+            for dtype in (torch.bfloat16, torch.float32):
+                logits = torch.randn(rows, cols, dtype=dtype, device="cuda") * 8
+                got_max, got_log_sum = row_logsumexp(logits)
+                self.assertEqual(got_max.dtype, torch.float32)
+                self.assertEqual(got_log_sum.dtype, torch.float32)
+                if not cols:
+                    self.assertTrue((got_max == float("-inf")).all())
+                    self.assertTrue((got_log_sum == 0).all())
+                    continue
+                self.assertTrue(torch.equal(got_max, logits.float().amax(-1)))
+                ref_log_sum = torch.logsumexp(
+                    logits.double() - got_max.double()[:, None], dim=-1
+                ).float()
+                torch.testing.assert_close(
+                    ref_log_sum,
+                    got_log_sum,
+                    rtol=1e-4,
+                    atol=1e-4,
+                    msg=f"{rows}x{cols} {dtype}",
+                )
+        # Rows dominated by -inf (masked-vocab shapes) must stay nan-free.
+        logits = torch.full((4, 1000), float("-inf"), device="cuda")
+        logits[1, 3] = 2.5
+        logits[2, :] = torch.randn(1000, device="cuda")
+        got_max, got_log_sum = row_logsumexp(logits.bfloat16())
+        self.assertEqual(got_max[0].item(), float("-inf"))
+        self.assertAlmostEqual((got_max[1] + got_log_sum[1]).item(), 2.5, delta=1e-2)
+        self.assertFalse(got_max.isnan().any() or got_log_sum.isnan().any())
+        # Non-contiguous input (sliced rows) exercises the stride args.
+        base = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16)
+        view = base[::2]
+        got_max, got_log_sum = row_logsumexp(view)
+        torch.testing.assert_close(
+            torch.logsumexp(view.double(), dim=-1).float(),
+            got_max + got_log_sum,
+            rtol=1e-4,
+            atol=1e-4,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()