Add token-id verification to the KV-canary (#26818)

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)

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)

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()

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)
 
 

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(),
         )

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(
     *,

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,
+    )

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(

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,

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)

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

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,
     )
 
 

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,
     )