minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Sam Shleifer
2026-08-07 14:01:08 -07:00
committed by GitHub
parent 8a22b8305d
commit 8e7d361def
4 changed files with 824 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+4
View File
@@ -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
+118 -10
View File
@@ -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)
values = ret.values.tolist()
indices = ret.indices.tolist()
if log_normalizer is not None:
row_max, row_log_sum = log_normalizer
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
+298
View File
@@ -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
+404
View File
@@ -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()