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

2026-05-31 09:58:51 +08:00
parent 0ca610a6df
commit 6be4b32d8d
18 changed files with 983 additions and 13 deletions
@@ -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()