diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/bench_scatter_req_token_ids.py b/python/sglang/jit_kernel/benchmark/kv_canary/bench_scatter_req_token_ids.py new file mode 100644 index 000000000..fcf88a112 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/kv_canary/bench_scatter_req_token_ids.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import triton +import triton.testing + +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + get_benchmark_range, + run_benchmark_no_cudagraph, +) +from sglang.jit_kernel.kv_canary.scatter_req_token_ids import ( + launch_scatter_req_token_ids_kernel, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True) + + +_BS_AXIS_FULL: list[int] = [1, 8, 64, 256] +_SEQ_LEN_AXIS_FULL: list[int] = [128, 512, 2048, 8192] +_BS_AXIS_CI: list[int] = [1, 64] +_SEQ_LEN_AXIS_CI: list[int] = [512, 2048] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _BenchCase: + bs: int + seq_len: int + + +def _build_cases() -> list[_BenchCase]: + bs_axis = get_benchmark_range(full_range=_BS_AXIS_FULL, ci_range=_BS_AXIS_CI) + seq_axis = get_benchmark_range( + full_range=_SEQ_LEN_AXIS_FULL, ci_range=_SEQ_LEN_AXIS_CI + ) + return [ + _BenchCase(bs=bs, seq_len=seq_len) for bs in bs_axis for seq_len in seq_axis + ] + + +_X_NAMES = ["bs", "seq_len"] +_X_VALS = [(c.bs, c.seq_len) for c in _build_cases()] + + +def _build_inputs(*, bs: int, seq_len: int, device: torch.device) -> dict: + max_reqs = max(bs + 1, 4) + max_context_len = max(seq_len + 1, 1) + total_tokens = bs * seq_len + + flat = torch.randint( + low=0, + high=1 << 30, + size=(total_tokens,), + dtype=torch.int64, + device=device, + ) + lens = torch.full((bs,), seq_len, dtype=torch.int64, device=device) + offsets = torch.zeros(bs + 1, dtype=torch.int64, device=device) + offsets[1:] = torch.cumsum(lens, dim=0) + + req_pool_indices = torch.arange(1, bs + 1, dtype=torch.int64, device=device) + pool = torch.zeros((max_reqs, max_context_len), dtype=torch.int32, device=device) + return dict( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=_X_NAMES, + x_vals=_X_VALS, + line_arg="provider", + line_vals=["triton"], + line_names=["Triton"], + styles=[("blue", "-")], + ylabel="time (us)", + plot_name="kv-canary-scatter-req-token-ids", + args={}, + ) +) +def benchmark(bs: int, seq_len: int, provider: str) -> tuple[float, float, float]: + inputs = _build_inputs(bs=bs, seq_len=seq_len, device=torch.device(DEFAULT_DEVICE)) + return run_benchmark_no_cudagraph( + lambda: launch_scatter_req_token_ids_kernel(**inputs) + ) + + +if __name__ == "__main__": + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/kv_canary/scatter_req_token_ids.py b/python/sglang/jit_kernel/kv_canary/scatter_req_token_ids.py new file mode 100644 index 000000000..34f829c87 --- /dev/null +++ b/python/sglang/jit_kernel/kv_canary/scatter_req_token_ids.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +_SCATTER_TOKEN_BLOCK: int = 256 +# Upper bound on bs+1 the kernel can scan per program. Owner-req lookup uses an +# outer-product tile of shape ``[TOKEN_BLOCK, BATCH_BLOCK]``; keep this small so +# the tile stays in registers (256 x 512 = 128 KiB i1, well below the SM trap- +# inducing budget that bites at the 1M cell mark). +_SCATTER_BATCH_BLOCK: int = 512 + + +def launch_scatter_req_token_ids_kernel( + *, + flat_in: torch.Tensor, + offsets: torch.Tensor, + req_pool_indices: torch.Tensor, + pool_out: torch.Tensor, +) -> None: + """Scatter a flat per-req int64 object sequence into a 2-D int32 pool. + + For each global object index ``t`` in ``[0, total_tokens)``: + + - find ``r`` = largest req index s.t. ``offsets[r] <= t`` + - ``pos = t - offsets[r]`` + - ``rp = req_pool_indices[r]`` + - if ``pos < pool_max_context_len``: + ``pool_out[rp, pos] = flat_in[t].to(int32)`` + + Args: + flat_in: ``[total_tokens]`` int64 device tensor of objects, flattened + per-req in req order. + offsets: ``[bs + 1]`` int64 device tensor (host-computed cumsum of per-req + lengths). ``offsets[bs] == total_tokens``. + req_pool_indices: ``[bs]`` int64 device tensor of pool row indices. + pool_out: ``[max_reqs, max_context_len]`` int32 device tensor of objects. + Mutated in-place; rows not addressed by ``req_pool_indices`` are untouched. + + Implementation notes: + - Linear scan over ``offsets`` (``BATCH_BLOCK >= bs + 1``); fits easily in + registers for the workloads kv-canary handles (``bs <= a few thousand``). + """ + if flat_in.dim() != 1: + raise ValueError( + f"kv-canary: scatter_req_token_ids flat_in must be 1-D, got shape " + f"{tuple(flat_in.shape)}" + ) + if offsets.dim() != 1: + raise ValueError( + f"kv-canary: scatter_req_token_ids offsets must be 1-D, got shape " + f"{tuple(offsets.shape)}" + ) + if req_pool_indices.dim() != 1: + raise ValueError( + f"kv-canary: scatter_req_token_ids req_pool_indices must be 1-D, got shape " + f"{tuple(req_pool_indices.shape)}" + ) + if pool_out.dim() != 2: + raise ValueError( + f"kv-canary: scatter_req_token_ids pool_out must be 2-D, got shape " + f"{tuple(pool_out.shape)}" + ) + if flat_in.dtype != torch.int64: + raise TypeError( + f"kv-canary: scatter_req_token_ids flat_in must be int64, got " + f"{flat_in.dtype}" + ) + if offsets.dtype != torch.int64: + raise TypeError( + f"kv-canary: scatter_req_token_ids offsets must be int64, got " + f"{offsets.dtype}" + ) + if req_pool_indices.dtype != torch.int64: + raise TypeError( + f"kv-canary: scatter_req_token_ids req_pool_indices must be int64, got " + f"{req_pool_indices.dtype}" + ) + if pool_out.dtype != torch.int32: + raise TypeError( + f"kv-canary: scatter_req_token_ids pool_out must be int32, got " + f"{pool_out.dtype}" + ) + + bs = int(req_pool_indices.shape[0]) + if int(offsets.shape[0]) != bs + 1: + raise ValueError( + f"kv-canary: scatter_req_token_ids offsets length {offsets.shape[0]} != " + f"bs+1 ({bs + 1})" + ) + if bs + 1 > _SCATTER_BATCH_BLOCK: + raise ValueError( + f"kv-canary: scatter_req_token_ids bs+1={bs + 1} exceeds BATCH_BLOCK=" + f"{_SCATTER_BATCH_BLOCK}; bump _SCATTER_BATCH_BLOCK if real workloads need this" + ) + + num_tokens = int(flat_in.shape[0]) + if num_tokens == 0: + return + + pool_stride0 = int(pool_out.stride(0)) + pool_max_context_len = int(pool_out.shape[1]) + + grid = (triton.cdiv(num_tokens, _SCATTER_TOKEN_BLOCK),) + _scatter_req_token_ids_kernel[grid]( + flat_in, + offsets, + req_pool_indices, + pool_out, + num_tokens=num_tokens, + num_batch=bs, + pool_stride0=pool_stride0, + pool_max_context_len=pool_max_context_len, + TOKEN_BLOCK=_SCATTER_TOKEN_BLOCK, + BATCH_BLOCK=_SCATTER_BATCH_BLOCK, + ) + + +@triton.jit +def _scatter_req_token_ids_kernel( + flat_in_ptr, # [num_tokens] int64 + offsets_ptr, # [num_batch + 1] int64 + req_pool_indices_ptr, # [num_batch] int64 + pool_out_ptr, # [num_rows, pool_max_context_len] int32, row stride = pool_stride0 + num_tokens, # scalar int32 + num_batch, # scalar int32 + pool_stride0, # scalar int32 (row stride of pool_out in elements) + pool_max_context_len, # scalar int32 (dim-1 length of pool_out) + TOKEN_BLOCK: tl.constexpr, + BATCH_BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + tids = pid * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK) # [TOKEN_BLOCK] int32 + tid_mask = tids < num_tokens # [TOKEN_BLOCK] bool + + bs_offs = tl.arange(0, BATCH_BLOCK) # [BATCH_BLOCK] int32 + bs_mask = bs_offs < (num_batch + 1) # [BATCH_BLOCK] bool + offs_vals = tl.load( # [BATCH_BLOCK] int64 + offsets_ptr + bs_offs, + mask=bs_mask, + other=(1 << 62), + ) + + # find owning req for each tid via reduce-sum: req_idx = (count of offsets <= tid) - 1 + le = offs_vals[None, :] <= tids[:, None] # [TOKEN_BLOCK, BATCH_BLOCK] bool + req_idx = tl.sum(le.to(tl.int32), axis=1) - 1 # [TOKEN_BLOCK] int32 + + safe_req_idx = tl.where(tid_mask, req_idx, 0) # [TOKEN_BLOCK] int32 + starts = tl.load( + offsets_ptr + safe_req_idx, mask=tid_mask, other=0 + ) # [TOKEN_BLOCK] int64 + pos = tids - starts # [TOKEN_BLOCK] int64 + rp = tl.load( + req_pool_indices_ptr + safe_req_idx, mask=tid_mask, other=0 + ) # [TOKEN_BLOCK] int64 + + # Bound writes by the pool's max_context_len so a token sequence longer than the + # ReqToTokenPool row never spills into an adjacent row. + in_row = pos < pool_max_context_len # [TOKEN_BLOCK] bool + write_mask = tid_mask & in_row # [TOKEN_BLOCK] bool + + val = tl.load(flat_in_ptr + tids, mask=tid_mask, other=0).to( + tl.int32 + ) # [TOKEN_BLOCK] int32 + tl.store(pool_out_ptr + rp * pool_stride0 + pos, val, mask=write_mask) + + +def scatter_req_token_ids_torch_reference( + *, + flat_in: torch.Tensor, + offsets: torch.Tensor, + req_pool_indices: torch.Tensor, + pool_out: torch.Tensor, +) -> None: + """Plain-PyTorch reference for :func:`launch_scatter_req_token_ids_kernel`.""" + bs = int(req_pool_indices.shape[0]) + offsets_host = offsets.detach().cpu().tolist() + req_pool_indices_host = req_pool_indices.detach().cpu().tolist() + flat_host = flat_in.detach().cpu() + pool_max_context_len = int(pool_out.shape[1]) + + for r in range(bs): + start = int(offsets_host[r]) + end = int(offsets_host[r + 1]) + if end <= start: + continue + rp = int(req_pool_indices_host[r]) + seg = flat_host[start:end].to(torch.int32) + write_len = min(int(seg.shape[0]), pool_max_context_len) + if write_len <= 0: + continue + pool_out[rp, :write_len] = seg[:write_len].to(pool_out.device) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_scatter_req_token_ids.py b/python/sglang/jit_kernel/tests/kv_canary/test_scatter_req_token_ids.py new file mode 100644 index 000000000..1bdb5bd5c --- /dev/null +++ b/python/sglang/jit_kernel/tests/kv_canary/test_scatter_req_token_ids.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +import random +import unittest + +import torch + +from sglang.jit_kernel.kv_canary.scatter_req_token_ids import ( + _SCATTER_BATCH_BLOCK, + launch_scatter_req_token_ids_kernel, + scatter_req_token_ids_torch_reference, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large") + + +_DEVICE = torch.device("cuda") + + +def _build_pool(*, max_reqs: int, max_context_len: int) -> torch.Tensor: + return torch.zeros((max_reqs, max_context_len), dtype=torch.int32, device=_DEVICE) + + +def _build_offsets(lens: list[int]) -> torch.Tensor: + cumsum = [0] + for n in lens: + cumsum.append(cumsum[-1] + n) + return torch.tensor(cumsum, dtype=torch.int64, device=_DEVICE) + + +def _build_flat(seqs: list[list[int]]) -> torch.Tensor: + flat: list[int] = [] + for s in seqs: + flat.extend(s) + return torch.tensor(flat, dtype=torch.int64, device=_DEVICE) + + +class TestScatterReqTokenIds(CustomTestCase): + def test_scatter_byte_equal_basic(self) -> None: + """Triton output matches the PyTorch reference for a small mixed batch.""" + seqs = [[10, 20, 30], [40, 50], [60, 70, 80, 90]] + lens = [len(s) for s in seqs] + rp = [3, 1, 5] + + flat = _build_flat(seqs) + offsets = _build_offsets(lens) + req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE) + triton_pool = _build_pool(max_reqs=8, max_context_len=16) + ref_pool = _build_pool(max_reqs=8, max_context_len=16) + + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=triton_pool, + ) + scatter_req_token_ids_torch_reference( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=ref_pool, + ) + torch.cuda.synchronize() + + self.assertTrue(torch.equal(triton_pool, ref_pool)) + # Spot-check: req in slot 3 holds [10,20,30,0,0,...] etc. + self.assertEqual(triton_pool[3, :3].tolist(), [10, 20, 30]) + self.assertEqual(triton_pool[1, :2].tolist(), [40, 50]) + self.assertEqual(triton_pool[5, :4].tolist(), [60, 70, 80, 90]) + + def test_scatter_empty_batch_no_op(self) -> None: + """Empty input (num_tokens == 0) returns without touching the pool.""" + flat = torch.empty(0, dtype=torch.int64, device=_DEVICE) + offsets = torch.zeros(1, dtype=torch.int64, device=_DEVICE) + req_pool_indices = torch.empty(0, dtype=torch.int64, device=_DEVICE) + pool = _build_pool(max_reqs=8, max_context_len=16) + pool_before = pool.clone() + + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool, pool_before)) + + def test_scatter_single_req(self) -> None: + """One req with a full row of tokens writes byte-equal to the reference.""" + seqs = [list(range(10))] + lens = [10] + rp = [2] + + flat = _build_flat(seqs) + offsets = _build_offsets(lens) + req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE) + triton_pool = _build_pool(max_reqs=4, max_context_len=16) + ref_pool = _build_pool(max_reqs=4, max_context_len=16) + + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=triton_pool, + ) + scatter_req_token_ids_torch_reference( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=ref_pool, + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(triton_pool, ref_pool)) + + def test_scatter_truncates_at_max_context_len(self) -> None: + """Tokens past the pool's max_context_len are silently dropped (no row spill).""" + # Two reqs; first req longer than max_context_len. Second req must remain + # uncorrupted. + seqs = [list(range(20)), [777, 888, 999]] + lens = [len(s) for s in seqs] + rp = [1, 2] + max_context_len = 8 + + flat = _build_flat(seqs) + offsets = _build_offsets(lens) + req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE) + pool = _build_pool(max_reqs=4, max_context_len=max_context_len) + + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + torch.cuda.synchronize() + + self.assertEqual( + pool[1, :max_context_len].tolist(), list(range(max_context_len)) + ) + self.assertEqual(pool[2, :3].tolist(), [777, 888, 999]) + + def test_scatter_mixed_empty_and_nonempty_reqs(self) -> None: + """Middle req has length 0 between two non-empty reqs: pool rows are byte-equal and untouched rows stay zero.""" + seqs = [[1, 2], [], [3, 4, 5]] + lens = [len(s) for s in seqs] + rp = [2, 4, 6] + + flat = _build_flat(seqs) + offsets = _build_offsets(lens) + req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE) + triton_pool = _build_pool(max_reqs=8, max_context_len=8) + ref_pool = _build_pool(max_reqs=8, max_context_len=8) + + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=triton_pool, + ) + scatter_req_token_ids_torch_reference( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=ref_pool, + ) + torch.cuda.synchronize() + + self.assertTrue(torch.equal(triton_pool, ref_pool)) + # Middle req contributes nothing; its pool row stays zero. + zero_row = torch.zeros(8, dtype=torch.int32, device=_DEVICE) + self.assertTrue(torch.equal(triton_pool[4], zero_row)) + # First and third reqs are written to their respective rows. + self.assertEqual(triton_pool[2, :2].tolist(), [1, 2]) + self.assertEqual(triton_pool[6, :3].tolist(), [3, 4, 5]) + + def test_scatter_random_byte_equal(self) -> None: + """Randomized fuzz across bs, seq lengths, and req pool indices.""" + rng = random.Random(0) + max_reqs = 64 + max_context_len = 32 + + for _ in range(8): + bs = rng.randint(1, 16) + lens = [rng.randint(0, max_context_len) for _ in range(bs)] + # All distinct req pool indices in [1, max_reqs) + rp = rng.sample(range(1, max_reqs), k=bs) + seqs = [[rng.randint(0, 1 << 30) for _ in range(n)] for n in lens] + + flat = _build_flat(seqs) + offsets = _build_offsets(lens) + req_pool_indices = torch.tensor(rp, dtype=torch.int64, device=_DEVICE) + triton_pool = _build_pool( + max_reqs=max_reqs, max_context_len=max_context_len + ) + ref_pool = _build_pool(max_reqs=max_reqs, max_context_len=max_context_len) + + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=triton_pool, + ) + scatter_req_token_ids_torch_reference( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=ref_pool, + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(triton_pool, ref_pool)) + + +class TestScatterInputValidation(CustomTestCase): + """Cover the strict input checks in launch_scatter_req_token_ids_kernel.""" + + def test_raises_on_2d_flat_in(self) -> None: + """A 2-D flat_in tensor triggers a ValueError before any kernel launch.""" + flat = torch.zeros((2, 2), dtype=torch.int64, device=_DEVICE) + offsets = torch.tensor([0, 1, 2], dtype=torch.int64, device=_DEVICE) + req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE) + pool = _build_pool(max_reqs=4, max_context_len=4) + with self.assertRaises(ValueError): + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + + def test_raises_on_wrong_dtype_pool(self) -> None: + """A pool_out with non-int32 dtype triggers a TypeError.""" + flat = torch.tensor([10, 20], dtype=torch.int64, device=_DEVICE) + offsets = torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE) + req_pool_indices = torch.tensor([1], dtype=torch.int64, device=_DEVICE) + pool = torch.zeros((4, 4), dtype=torch.int64, device=_DEVICE) + with self.assertRaises(TypeError): + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + + def test_raises_on_offsets_len_mismatch(self) -> None: + """offsets.shape[0] must equal bs + 1; mismatch triggers a ValueError.""" + flat = torch.tensor([10, 20], dtype=torch.int64, device=_DEVICE) + # bs = 2 but offsets has length 2 instead of 3. + offsets = torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE) + req_pool_indices = torch.tensor([1, 2], dtype=torch.int64, device=_DEVICE) + pool = _build_pool(max_reqs=4, max_context_len=4) + with self.assertRaises(ValueError): + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + + def test_raises_on_bs_plus_one_exceeds_batch_block(self) -> None: + """bs+1 must fit in _SCATTER_BATCH_BLOCK; exceeding it triggers a ValueError.""" + bs = _SCATTER_BATCH_BLOCK + flat = torch.empty(0, dtype=torch.int64, device=_DEVICE) + offsets = torch.zeros(bs + 1, dtype=torch.int64, device=_DEVICE) + req_pool_indices = torch.zeros(bs, dtype=torch.int64, device=_DEVICE) + pool = _build_pool(max_reqs=4, max_context_len=4) + with self.assertRaises(ValueError): + launch_scatter_req_token_ids_kernel( + flat_in=flat, + offsets=offsets, + req_pool_indices=req_pool_indices, + pool_out=pool, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 97b2ddd4c..6d98a974c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -758,6 +758,7 @@ class Envs: SGLANG_KV_CANARY_PERTURB_TARGET_GROUP = EnvStr(None) SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB = EnvFloat(0.0) SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE = EnvBool(False) + SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT = EnvBool(False) SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False) diff --git a/python/sglang/srt/kv_canary/config.py b/python/sglang/srt/kv_canary/config.py index 7da1c2a34..2ae617ee5 100644 --- a/python/sglang/srt/kv_canary/config.py +++ b/python/sglang/srt/kv_canary/config.py @@ -43,6 +43,10 @@ class CanaryConfig: expected_input_positions[i]; mismatch records a violation. Only useful when something else (e.g. token_oracle.oracle_manager.fill_expected_inputs) is feeding the expected_* placeholders per forward — canary itself knows no oracle. + enable_verify_token_assert: bool. True = real-model token-id validator: build + expected_tokens from each req's ``origin_input_ids + output_ids`` (snapshotted at + ForwardBatch.init_new) and compare against the canary's stored tokens at verify time. + Independent of ``enable_write_input_assert``. """ mode: CanaryMode @@ -50,6 +54,7 @@ class CanaryConfig: sweep_interval: int real_kv_hash_mode: RealKvHashMode enable_write_input_assert: bool + enable_verify_token_assert: bool @classmethod def from_env(cls, server_args: "ServerArgs") -> "CanaryConfig": @@ -67,4 +72,5 @@ class CanaryConfig: sweep_interval=server_args.kv_canary_sweep_interval, real_kv_hash_mode=RealKvHashMode[real_kv_raw], enable_write_input_assert=envs.SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT.get(), + enable_verify_token_assert=envs.SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT.get(), ) diff --git a/python/sglang/srt/kv_canary/plan_input.py b/python/sglang/srt/kv_canary/plan_input.py index 477a9026f..bcbaddeea 100644 --- a/python/sglang/srt/kv_canary/plan_input.py +++ b/python/sglang/srt/kv_canary/plan_input.py @@ -29,8 +29,7 @@ class PlanInput: ``sot_pos`` when gathering the expected token; everything past the snapshot (e.g. EAGLE draft / verify positions, or stale residue from a longer recycled slot owner) returns the ``-1`` sentinel and the verify kernel skips the check. - The naive build never populates the verify-token-id cross-check, so this stays all - zeros and the plan kernel's gather degrades to the ``-1`` skip sentinel. + Set only when ``CanaryConfig.enable_verify_token_assert`` is on. Allocated fresh per forward by :class:`SingleForwardManager`. The boundary ForwardBatch token/position/slot tensors must already be int64 @@ -84,6 +83,12 @@ class PlanInput: bs=bs, ) + req_all_ids_lens = forward_batch.req_all_ids_lens + if req_all_ids_lens is not None: + self.req_to_verify_expected_tokens_valid_lens[:bs].copy_( + req_all_ids_lens.to(torch.int64), non_blocking=True + ) + def _extract_prefix_lens_and_extend_seq_lens( *, diff --git a/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py b/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py new file mode 100644 index 000000000..d36ea782e --- /dev/null +++ b/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.jit_kernel.kv_canary.scatter_req_token_ids import ( + launch_scatter_req_token_ids_kernel, +) +from sglang.srt.utils.common import flatten_arrays_to_int64_tensor + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +def compute_req_all_ids_info( + reqs: "list[Req]", +) -> tuple[torch.Tensor, torch.Tensor]: + """Snapshot per-req (origin_input_ids + output_ids) as pinned CPU int64 tensors. + + Returns: + ``(req_all_ids_flat, req_all_ids_lens)`` — both pinned CPU int64. ``flat`` is the + flattened ``cat(r.origin_input_ids, r.output_ids) for r in reqs``; ``lens`` is + per-req ``len(origin_input_ids) + len(output_ids)``. + """ + parts = [arr for req in reqs for arr in (req.origin_input_ids, req.output_ids)] + req_all_ids_flat = flatten_arrays_to_int64_tensor( + parts, device=torch.device("cpu"), pin=True + ) + req_all_ids_lens = torch.tensor( + [len(req.origin_input_ids) + len(req.output_ids) for req in reqs], + dtype=torch.int64, + pin_memory=True, + ) + return req_all_ids_flat, req_all_ids_lens + + +def populate_req_to_expected_token_ids( + *, + forward_batch: "ForwardBatch", + req_to_verify_expected_tokens: Optional[torch.Tensor], +) -> None: + """Scatter the forward batch's per-req token-id snapshot into the device-side pool.""" + req_all_ids_flat_cpu = forward_batch.req_all_ids_flat + req_all_ids_lens_cpu = forward_batch.req_all_ids_lens + if req_all_ids_flat_cpu is None or req_all_ids_lens_cpu is None: + return + if req_to_verify_expected_tokens is None: + return + + bs = int(forward_batch.req_pool_indices.shape[0]) + if bs == 0: + return + if int(req_all_ids_lens_cpu.shape[0]) != bs: + raise RuntimeError( + f"kv-canary: req_all_ids_lens length {int(req_all_ids_lens_cpu.shape[0])} != " + f"batch_size {bs}; ForwardBatch snapshot diverged" + ) + + offsets_cpu = torch.zeros(bs + 1, dtype=torch.int64, pin_memory=True) + offsets_cpu[1:] = torch.cumsum(req_all_ids_lens_cpu, dim=0) + total_tokens = int(offsets_cpu[bs].item()) + if total_tokens != int(req_all_ids_flat_cpu.shape[0]): + raise RuntimeError( + f"kv-canary: cumsum(req_all_ids_lens)={total_tokens} != " + f"req_all_ids_flat.numel()={int(req_all_ids_flat_cpu.shape[0])}; snapshot inconsistent" + ) + if total_tokens == 0: + return + + device = req_to_verify_expected_tokens.device + req_all_ids_flat_dev = req_all_ids_flat_cpu.to(device, non_blocking=True) + offsets_dev = offsets_cpu.to(device, non_blocking=True) + req_pool_indices_dev = forward_batch.req_pool_indices.to( + device=device, dtype=torch.int64 + ) + + launch_scatter_req_token_ids_kernel( + flat_in=req_all_ids_flat_dev, + offsets=offsets_dev, + req_pool_indices=req_pool_indices_dev, + pool_out=req_to_verify_expected_tokens, + ) diff --git a/python/sglang/srt/kv_canary/single_forward_manager/manager.py b/python/sglang/srt/kv_canary/single_forward_manager/manager.py index 5f4293176..9328d1d18 100644 --- a/python/sglang/srt/kv_canary/single_forward_manager/manager.py +++ b/python/sglang/srt/kv_canary/single_forward_manager/manager.py @@ -13,6 +13,9 @@ from sglang.srt.kv_canary.config import CanaryConfig from sglang.srt.kv_canary.endpoint import CanaryEndpoint from sglang.srt.kv_canary.expected_inputs import ExpectedInputs from sglang.srt.kv_canary.plan_input import PlanInput +from sglang.srt.kv_canary.req_to_expected_token_ids_manager import ( + populate_req_to_expected_token_ids, +) from sglang.srt.kv_canary.runner.enable_warner import CanaryEnableWarner from sglang.srt.kv_canary.runner.kernel_launcher import ( invoke_plan, @@ -130,6 +133,12 @@ class SingleForwardManager: f"CanaryLaunchCapacities.from_args" ) + if self._config.enable_verify_token_assert: + populate_req_to_expected_token_ids( + forward_batch=maybe_inaccurate_forward_batch, + req_to_verify_expected_tokens=self._device_state.req_to_verify_expected_tokens, + ) + def pre_ops_maybe_inside_graph( self, forward_batch: "ForwardBatch" ) -> "_PreOpsMaybeInsideGraphOutput": @@ -203,7 +212,7 @@ class SingleForwardManager: violation_log=violation_log, real_kv_hash_mode=self._config.real_kv_hash_mode, enable_write_input_assert=enable_write_input_assert, - enable_verify_token_assert=False, + enable_verify_token_assert=self._config.enable_verify_token_assert, ) return _PreOpsMaybeInsideGraphOutput( @@ -241,7 +250,7 @@ class SingleForwardManager: violation_log=violation_log, real_kv_hash_mode=self._config.real_kv_hash_mode, enable_write_input_assert=enable_write_input_assert, - enable_verify_token_assert=False, + enable_verify_token_assert=self._config.enable_verify_token_assert, ) verify_plan_enable_combined = _torch_reduce_minimum( diff --git a/python/sglang/srt/kv_canary/state.py b/python/sglang/srt/kv_canary/state.py index 8b9a5a35c..177a27246 100644 --- a/python/sglang/srt/kv_canary/state.py +++ b/python/sglang/srt/kv_canary/state.py @@ -80,11 +80,11 @@ class CanaryDeviceState: req_to_verify_expected_tokens: Optional int32 device tensor shape ``[req_to_token_alloc_size, max_context_len]``. Mirrors ReqToTokenPool layout; ``pool[req_idx, p]`` = source-of-truth token at logical position ``p`` for the - req in slot ``req_idx``. The plan-side entries kernel gathers from this pool (via - ``kv_token_id_vs_position_offset`` per buffer group) into - ``VerifyPlan.verify_expected_tokens``; the verify kernel then compares against each - canary slot's stored token. The naive build never populates the verify-token-id - cross-check, so this is always ``None`` and the gather degrades to the ``-1`` sentinel. + req in slot ``req_idx``. Allocated only when + ``CanaryConfig.enable_verify_token_assert`` is True. The plan-side entries + kernel gathers from this pool (via ``kv_token_id_vs_position_offset`` per buffer + group) into ``VerifyPlan.verify_expected_tokens``; the verify kernel then + compares against each canary slot's stored token. """ violation_log: ViolationLog @@ -113,10 +113,19 @@ class CanaryDeviceState: kernel_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device) slot_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device) enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device) - # The naive build does not run the verify-token-id cross-check, so the source-of-truth - # token pool is never allocated. The field is kept on the dataclass and downstream code - # (plan kernel gather) treats ``None`` as "emit the -1 skip sentinel". - req_to_verify_expected_tokens = None + if config.enable_verify_token_assert: + if req_to_token_alloc_size is None or max_context_len is None: + raise ValueError( + "kv-canary: CanaryDeviceState.allocate requires req_to_token_alloc_size " + "and max_context_len when CanaryConfig.enable_verify_token_assert is on" + ) + req_to_verify_expected_tokens = torch.empty( + (req_to_token_alloc_size, max_context_len), + dtype=torch.int32, + device=device, + ) + else: + req_to_verify_expected_tokens = None return cls( violation_log=violation_log, kernel_run_counters=kernel_run_counters, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 6e0644c0c..11acb8639 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -42,6 +42,9 @@ from sglang.srt.distributed.parallel_state import ( get_tensor_model_parallel_world_size, ) from sglang.srt.environ import envs +from sglang.srt.kv_canary.req_to_expected_token_ids_manager import ( + compute_req_all_ids_info, +) from sglang.srt.layers.dp_attention import ( DpPaddingMode, get_attention_cp_size, @@ -587,6 +590,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ret.rids_int = hashed ret.bootstrap_room_ids_int = bootstrap_room_ids + if envs.SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT.get(): + ret.req_all_ids_flat, ret.req_all_ids_lens = compute_req_all_ids_info( + batch.reqs + ) + if batch.extend_input_logprob_token_ids is not None: ret.extend_input_logprob_token_ids_gpu = ( batch.extend_input_logprob_token_ids.to(device, non_blocking=True) diff --git a/python/sglang/test/kv_canary/e2e_base.py b/python/sglang/test/kv_canary/e2e_base.py index 0c6cf059b..61c9ecca3 100644 --- a/python/sglang/test/kv_canary/e2e_base.py +++ b/python/sglang/test/kv_canary/e2e_base.py @@ -78,6 +78,7 @@ class CanaryE2EBase(CapturedServerE2EBase): def setUpClass(cls) -> None: cls._cfg = _MODE_CONFIGS[cls.model_mode] server_env = os.environ.copy() + server_env.setdefault("SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT", "1") server_env.update(cls.extra_env) if cls.model_mode == "swa": # SWA mode uses google/gemma-4-E2B-it, whose forward does a diff --git a/python/sglang/test/kv_canary/fixtures.py b/python/sglang/test/kv_canary/fixtures.py index bf2b361c8..6ac7f253a 100644 --- a/python/sglang/test/kv_canary/fixtures.py +++ b/python/sglang/test/kv_canary/fixtures.py @@ -139,6 +139,7 @@ def make_base_config() -> CanaryConfig: sweep_interval=0, real_kv_hash_mode=consts.RealKvHashMode.NONE, enable_write_input_assert=False, + enable_verify_token_assert=True, ) diff --git a/python/sglang/test/kv_canary/runner_test_base.py b/python/sglang/test/kv_canary/runner_test_base.py index e84a94904..077c09a47 100644 --- a/python/sglang/test/kv_canary/runner_test_base.py +++ b/python/sglang/test/kv_canary/runner_test_base.py @@ -29,6 +29,7 @@ def make_config( sweep_interval: int = 0, real_kv_hash_mode: RealKvHashMode = RealKvHashMode.NONE, enable_write_input_assert: bool = False, + enable_verify_token_assert: bool = True, ) -> CanaryConfig: return CanaryConfig( mode=mode, @@ -36,6 +37,7 @@ def make_config( sweep_interval=sweep_interval, real_kv_hash_mode=real_kv_hash_mode, enable_write_input_assert=enable_write_input_assert, + enable_verify_token_assert=enable_verify_token_assert, ) diff --git a/test/registered/kv_canary/test_self_e2e_pr_25015.py b/test/registered/kv_canary/test_self_e2e_pr_25015.py index 91be27c95..f1657eb88 100644 --- a/test/registered/kv_canary/test_self_e2e_pr_25015.py +++ b/test/registered/kv_canary/test_self_e2e_pr_25015.py @@ -14,6 +14,7 @@ register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") _SPEC_EAGLE_TOKEN_ORACLE_ENV = { "SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT": "0", "SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "1", + "SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT": "0", } _SPEC_EAGLE_REVERT_PR_ENV = { **_SPEC_EAGLE_TOKEN_ORACLE_ENV, diff --git a/test/registered/kv_canary/test_self_e2e_pr_26329.py b/test/registered/kv_canary/test_self_e2e_pr_26329.py new file mode 100644 index 000000000..832190029 --- /dev/null +++ b/test/registered/kv_canary/test_self_e2e_pr_26329.py @@ -0,0 +1,84 @@ +"""Regression for PR #26329 EAGLE chunked-prefill rotation.""" + +from __future__ import annotations + +import random +import string +import unittest +from typing import ClassVar + +from sglang.srt.kv_canary.config import CanaryMode +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.e2e_base import CanaryE2EBase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + + +_CHUNKED_PREFILL_SIZE = 2048 +_EAGLE_CHUNKED_SERVER_ARGS = ( + "--speculative-algorithm", + "EAGLE", + "--chunked-prefill-size", + str(_CHUNKED_PREFILL_SIZE), + "--cuda-graph-max-bs", + "1", + "--max-running-requests", + "4", +) + + +class _EagleChunkedRotationBase(CanaryE2EBase): + model_mode = "mha" + kv_canary_mode = CanaryMode.LOG + extra_server_args = _EAGLE_CHUNKED_SERVER_ARGS + + revert_pr: ClassVar[bool] + + @classmethod + def setUpClass(cls) -> None: + if cls is _EagleChunkedRotationBase: + raise unittest.SkipTest("abstract base; concrete subclasses set revert_pr") + cls.extra_env = {"SGLANG_DEBUG_REVERT_PR": "26329"} if cls.revert_pr else {} + super().setUpClass() + + def make_prompts(self, n: int) -> list[str]: + # Seeded random ASCII so the model can't predict the next prompt token + # — otherwise target's bonus token can accidentally match prompt[K1] + # at the chunk boundary and the validator never fires. + rng = random.Random(0) + # ~3K tokens after BPE — spans 2+ chunks at chunked_prefill_size=2048. + body = "".join(rng.choices(string.ascii_letters + string.digits + " ", k=8000)) + return [body] * n + + def test_chunked_rotation_token_id_mismatch(self) -> None: + self.send_parallel_requests( + n=1, + assert_all_success=not self.revert_pr, + max_new_tokens=8, + timeout=60.0, + ) + + if self.revert_pr: + self.assert_violation_logged_any( + launch_tag_patterns=("*",), + fail_reason="verify_token", + flush_wait_seconds=3.0, + ) + else: + self.assert_no_violation(wait_seconds=2.0) + + +class TestEagleChunkedRotationRegression(_EagleChunkedRotationBase): + """Revert PR #26329 fix; expect canary to fire a verify_token violation.""" + + revert_pr = True + + +class TestEagleChunkedRotationClean(_EagleChunkedRotationBase): + """With the PR #26329 fix in place, the same request runs clean.""" + + revert_pr = False + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_buffer_alloc.py b/test/registered/kv_canary/test_self_unit_buffer_alloc.py index 3e0ced43c..3872b6c7a 100644 --- a/test/registered/kv_canary/test_self_unit_buffer_alloc.py +++ b/test/registered/kv_canary/test_self_unit_buffer_alloc.py @@ -25,6 +25,7 @@ def _config(mode: RealKvHashMode) -> CanaryConfig: sweep_interval=0, real_kv_hash_mode=mode, enable_write_input_assert=False, + enable_verify_token_assert=False, ) diff --git a/test/registered/kv_canary/test_self_unit_plan_input.py b/test/registered/kv_canary/test_self_unit_plan_input.py index 9eb9a59f8..937b7221b 100644 --- a/test/registered/kv_canary/test_self_unit_plan_input.py +++ b/test/registered/kv_canary/test_self_unit_plan_input.py @@ -108,6 +108,28 @@ class TestSelfUnitPlanInput(CustomTestCase): self.assertEqual(plan.prefix_lens[:3].tolist(), [3, 6, 0]) self.assertEqual(plan.extend_seq_lens[:3].tolist(), [1, 1, 1]) + def test_plan_input_mirrors_req_all_ids_lens(self): + """req_to_verify_expected_tokens_valid_lens copies forward_batch.req_all_ids_lens for active rows.""" + fb = make_forward_batch( + self.device, + req_pool_indices=torch.tensor( + [1, 2], dtype=torch.int64, device=self.device + ), + seq_lens=torch.tensor([10, 12], dtype=torch.int32, device=self.device), + is_extend=False, + ) + fb.req_all_ids_lens = torch.tensor([7, 9], dtype=torch.int64, pin_memory=True) + plan = _make_static_plan_input(bs_capacity=4, device=self.device) + plan.fill_from_forward_batch(forward_batch=fb) + torch.cuda.synchronize() + self.assertEqual( + plan.req_to_verify_expected_tokens_valid_lens[:2].tolist(), [7, 9] + ) + # Padding tail stays at zero so the plan kernel reads "no in-range positions" for it. + self.assertEqual( + plan.req_to_verify_expected_tokens_valid_lens[2:].tolist(), [0, 0] + ) + def test_plan_input_padding_dummy_sentinel(self): """Verify padding sentinel rows remain valid plan input entries.""" fb = make_forward_batch( diff --git a/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py b/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py new file mode 100644 index 000000000..518c4fbb9 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import unittest +from array import array +from types import SimpleNamespace + +import torch + +from sglang.srt.kv_canary.req_to_expected_token_ids_manager import ( + compute_req_all_ids_info, + populate_req_to_expected_token_ids, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_forward_batch +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=15, stage="extra-a", runner_config="1-gpu-small") + + +def _make_req(*, origin: list[int], output: list[int]) -> SimpleNamespace: + return SimpleNamespace( + origin_input_ids=array("q", origin), + output_ids=array("q", output), + ) + + +class TestComputeReqAllIdsInfo(CustomTestCase): + def test_single_req_concatenates_origin_then_output(self) -> None: + """One req's flat is origin_input_ids followed by output_ids in that order.""" + req = _make_req(origin=[10, 20, 30], output=[40, 50]) + flat, lens = compute_req_all_ids_info([req]) + self.assertEqual(flat.tolist(), [10, 20, 30, 40, 50]) + self.assertEqual(lens.tolist(), [5]) + + def test_multi_req_flat_is_concat_across_reqs(self) -> None: + """Multi-req flat is per-req (origin+output) concatenated in req order.""" + reqs = [ + _make_req(origin=[1, 2], output=[3]), + _make_req(origin=[100], output=[]), + _make_req(origin=[7, 8, 9], output=[10, 11]), + ] + flat, lens = compute_req_all_ids_info(reqs) + self.assertEqual(flat.tolist(), [1, 2, 3, 100, 7, 8, 9, 10, 11]) + self.assertEqual(lens.tolist(), [3, 1, 5]) + + def test_returned_cpu_tensors_are_pinned(self) -> None: + """Snapshot tensors live on pinned CPU memory so the manager's async H2D actually overlaps.""" + req = _make_req(origin=[1, 2, 3], output=[4]) + flat, lens = compute_req_all_ids_info([req]) + self.assertTrue(flat.is_pinned()) + self.assertTrue(lens.is_pinned()) + self.assertEqual(flat.device, torch.device("cpu")) + self.assertEqual(lens.device, torch.device("cpu")) + + +class TestPopulateReqToExpectedTokenIds(CustomTestCase): + def setUp(self) -> None: + self.device = DEFAULT_DEVICE + + def _make_pool(self, *, max_reqs: int, max_context_len: int) -> torch.Tensor: + return torch.full( + (max_reqs, max_context_len), + -999, + dtype=torch.int32, + device=self.device, + ) + + def _fb_with_snapshot( + self, + *, + req_pool_indices: list[int], + lens: list[int], + flat: list[int], + ) -> SimpleNamespace: + fb = make_forward_batch( + self.device, + bs=len(req_pool_indices), + req_pool_indices=torch.tensor( + req_pool_indices, dtype=torch.int64, device=self.device + ), + ) + fb.req_all_ids_flat = torch.tensor(flat, dtype=torch.int64, pin_memory=True) + fb.req_all_ids_lens = torch.tensor(lens, dtype=torch.int64, pin_memory=True) + return fb + + def test_no_op_when_snapshot_is_none(self) -> None: + """Cuda-graph capture's dry-run leaves snapshot fields as None; manager must early-return.""" + fb = make_forward_batch(self.device, bs=2) + pool = self._make_pool(max_reqs=4, max_context_len=8) + original = pool.clone() + populate_req_to_expected_token_ids( + forward_batch=fb, req_to_verify_expected_tokens=pool + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool, original)) + + def test_no_op_when_pool_is_none(self) -> None: + """When the validator is off the device pool is None; manager must early-return without touching anything.""" + fb = self._fb_with_snapshot(req_pool_indices=[1], lens=[3], flat=[10, 20, 30]) + populate_req_to_expected_token_ids( + forward_batch=fb, req_to_verify_expected_tokens=None + ) + + def test_no_op_when_bs_zero(self) -> None: + """Empty batch (bs == 0) must early-return; no kernel launch.""" + fb = make_forward_batch( + self.device, + bs=0, + req_pool_indices=torch.zeros(0, dtype=torch.int64, device=self.device), + seq_lens=torch.zeros(0, dtype=torch.int32, device=self.device), + ) + fb.req_all_ids_flat = torch.zeros(0, dtype=torch.int64, pin_memory=True) + fb.req_all_ids_lens = torch.zeros(0, dtype=torch.int64, pin_memory=True) + pool = self._make_pool(max_reqs=4, max_context_len=8) + original = pool.clone() + populate_req_to_expected_token_ids( + forward_batch=fb, req_to_verify_expected_tokens=pool + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(pool, original)) + + def test_raises_when_lens_length_mismatches_batch_size(self) -> None: + """req_all_ids_lens length must equal forward_batch batch size; mismatch indicates a corrupted snapshot.""" + fb = self._fb_with_snapshot( + req_pool_indices=[1, 2], lens=[3], flat=[10, 20, 30] + ) + pool = self._make_pool(max_reqs=4, max_context_len=8) + with self.assertRaisesRegex(RuntimeError, "req_all_ids_lens length"): + populate_req_to_expected_token_ids( + forward_batch=fb, req_to_verify_expected_tokens=pool + ) + + def test_raises_when_cumsum_does_not_match_flat_numel(self) -> None: + """cumsum(lens) must equal flat.numel(); inconsistent snapshot raises.""" + fb = self._fb_with_snapshot( + req_pool_indices=[1, 2], lens=[3, 4], flat=[10, 20, 30, 40] + ) + pool = self._make_pool(max_reqs=4, max_context_len=8) + with self.assertRaisesRegex(RuntimeError, "snapshot inconsistent"): + populate_req_to_expected_token_ids( + forward_batch=fb, req_to_verify_expected_tokens=pool + ) + + def test_happy_path_scatters_each_req_into_its_pool_row(self) -> None: + """Scatter populates pool[rp, :len_r] = req_r's flattened tokens; other rows untouched.""" + fb = self._fb_with_snapshot( + req_pool_indices=[1, 3], + lens=[3, 2], + flat=[10, 20, 30, 40, 50], + ) + pool = self._make_pool(max_reqs=5, max_context_len=8) + original = pool.clone() + populate_req_to_expected_token_ids( + forward_batch=fb, req_to_verify_expected_tokens=pool + ) + torch.cuda.synchronize() + + pool_cpu = pool.cpu() + self.assertEqual(pool_cpu[1, :3].tolist(), [10, 20, 30]) + self.assertEqual(pool_cpu[3, :2].tolist(), [40, 50]) + # Other rows must be left at their pre-scatter sentinel. + for untouched in (0, 2, 4): + self.assertTrue( + torch.equal(pool_cpu[untouched], original[untouched].cpu()), + f"row {untouched} was unexpectedly modified", + ) + + +if __name__ == "__main__": + unittest.main()