405 lines
17 KiB
Python
405 lines
17 KiB
Python
"""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()
|