diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f8c51025f..5edd4aca8 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index 5ea86a6a5..997695bd0 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/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) - 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 diff --git a/python/sglang/srt/layers/logsumexp.py b/python/sglang/srt/layers/logsumexp.py new file mode 100644 index 000000000..0d195dc71 --- /dev/null +++ b/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 diff --git a/test/registered/unit/layers/test_logprob_fast_input.py b/test/registered/unit/layers/test_logprob_fast_input.py new file mode 100644 index 000000000..72727f75e --- /dev/null +++ b/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()