minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 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