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

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,

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

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

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(

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