[Test] Consolidate kernel tests under plural kernels tree (#39966)

test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_bounds.py

@@ -0,0 +1,72 @@
+"""Boundary tests for the packed indices in the DSV4 prefill write plan."""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kernels.deepseek_v4.common import (
+    make_legacy_context,
+    make_paged_context,
+    to_seq_extend,
+)
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestCompressWritePlanBounds(CustomTestCase):
+    def test_64k_prefill_preserves_last_token(self):
+        """65536 tokens fit uint16 indices; the last token must not wrap or vanish."""
+        for cr in (4, 128):
+            paged = make_paged_context(
+                bs=16, compress_ratio=cr, num_swa_pages_per_req=16
+            )
+            legacy = make_legacy_context(bs=16, compress_ratio=cr)
+            seq_lens, extend_lens, num_q = to_seq_extend([(4096, 4096)] * 16)
+            for ctx, on_gpu in ((paged, False), (paged, True), (legacy, False)):
+                with self.subTest(cr=cr, paged=ctx is paged, on_gpu=on_gpu):
+                    device = "cuda" if on_gpu else "cpu"
+                    plan = ctx.make_prefill_plan(
+                        seq_lens.to(device), extend_lens.to(device), num_q
+                    )
+                    c = plan.plan_c.cpu().view(torch.int32).view(-1, 4)
+                    valid_c = c[:, 0] != -1
+                    ids = c[valid_c, 1].bitwise_and(0xFFFF).sort().values
+                    torch.testing.assert_close(
+                        ids, torch.arange(cr - 1, num_q, cr, dtype=torch.int32)
+                    )
+                    w = plan.plan_w.cpu().view(torch.int32).view(-1, 2)
+                    last = w[w[:, 0] == 65535]
+                    if cr == 4:
+                        self.assertEqual(len(last), 1)
+                        self.assertEqual(int(last[0, 1]), ctx.state_loc(15, 4095))
+                    else:
+                        # Non-overlapping C128 consumed the complete final block;
+                        # no raw tail remains to persist into the state ring.
+                        self.assertEqual(len(w[w[:, 0] != -1]), 0)
+
+    def test_prefill_rejects_uint16_index_overflow(self):
+        for ctx in (
+            make_paged_context(bs=16, compress_ratio=4, num_swa_pages_per_req=17),
+            make_legacy_context(bs=16, compress_ratio=4),
+        ):
+            seq_lens, extend_lens, num_q = to_seq_extend(
+                [(4096, 4096)] * 15 + [(4097, 4097)]
+            )
+            with self.assertRaisesRegex(RuntimeError, "plan_compress_prefill"):
+                ctx.make_prefill_plan(seq_lens, extend_lens, num_q)
+
+    def test_prefill_rejects_packed_invalid_sentinel(self):
+        # A 65536-request, one-token-per-request batch makes the last packed
+        # (batch_id, ragged_id) equal (65535, 65535), the invalid write sentinel.
+        ctx = make_legacy_context(bs=65536, compress_ratio=4)
+        seq_lens, extend_lens, num_q = to_seq_extend([(1, 1)] * 65536)
+        with self.assertRaisesRegex(RuntimeError, "plan_compress_prefill"):
+            ctx.make_prefill_plan(seq_lens, extend_lens, num_q)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py

@@ -11,10 +11,10 @@
 #
 # Test command:
 #   python3 -m pytest -q \
-#     test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py
+#     test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py
 #
 # Runtime-shape benchmark command:
-#   python3 test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py \
+#   python3 test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py \
 #     --benchmark \
 #     --shape-source runtime \
 #     --warmup 20 \
@@ -22,7 +22,7 @@
 #     --csv /data00/eval_results/operator_bench/runtime_shape_bench.csv
 #
 # Synthetic Flash/Pro shape benchmark command:
-#   python3 test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py \
+#   python3 test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py \
 #     --benchmark \
 #     --shape-source preset \
 #     --shape-presets all \

test/registered/kernels/ops/attention/test_dsa_kpool_metadata_fusion.py

@@ -0,0 +1,98 @@
+"""KPool fused metadata must retain live tails and refresh captured buffers."""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.dsa_kpool_metadata.verify import (
+    fused_dsa_target_verify_metadata,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestKPoolMetadataFusion(CustomTestCase):
+    def test_verify_replay_boundaries_and_request_remapping(self):
+        device = "cuda"
+        bs, next_n, width, topk, pool_size = 4, 6, 16384, 2048, 4
+        seq = torch.tensor([1, 61, 2047, 8191], device=device, dtype=torch.int64)
+        req = torch.tensor([3, 1, 6, 0], device=device, dtype=torch.int64)
+        table = torch.arange(8 * width, device=device, dtype=torch.int32).view(8, width)
+
+        def empty(*shape):
+            return torch.full(shape, -1, device=device, dtype=torch.int32)
+
+        buffers = dict(
+            cache_seqlens=empty(bs),
+            cu_seqlens_k=empty(bs + 1),
+            page_table_1=empty(bs * next_n, width),
+            seqlens_expanded=empty(bs * next_n),
+            dsa_cache_seqlens=empty(bs * next_n),
+            dsa_cu_seqlens_k=empty(bs * next_n + 1),
+            real_page_table=empty(bs * next_n, width // 64),
+            paged_mqa_ctx_lens_2d=empty(bs, next_n),
+        )
+        addresses = {key: value.data_ptr() for key, value in buffers.items()}
+
+        def refresh():
+            fused_dsa_target_verify_metadata(
+                seq_lens=seq,
+                req_pool_indices=req,
+                req_to_token=table,
+                bs=bs,
+                max_seqlen_k=width,
+                dsa_index_topk=topk,
+                real_page_size=64,
+                next_n=next_n,
+                index_kpool=pool_size,
+                **buffers,
+            )
+
+        refresh()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            refresh()
+        for lengths, requests in [
+            ([2, 63, 2048, 8193], [0, 6, 1, 3]),
+            ([64, 128, 2051, 8190], [5, 2, 7, 4]),
+            ([0, 3, 2045, 9000], [7, 0, 3, 2]),
+        ]:
+            seq.copy_(torch.tensor(lengths, device=device))
+            req.copy_(torch.tensor(requests, device=device))
+            graph.replay()
+            expanded = (
+                seq[:, None] + torch.arange(1, next_n + 1, device=device)
+            ).flatten()
+            expected = torch.minimum(expanded, topk + expanded % pool_size).int()
+            torch.testing.assert_close(buffers["seqlens_expanded"], expanded.int())
+            torch.testing.assert_close(buffers["dsa_cache_seqlens"], expected)
+            torch.testing.assert_close(
+                buffers["dsa_cu_seqlens_k"][1:], expected.cumsum(0).int()
+            )
+            torch.testing.assert_close(buffers["cache_seqlens"], (seq + next_n).int())
+            torch.testing.assert_close(
+                buffers["paged_mqa_ctx_lens_2d"],
+                (seq + next_n).int()[:, None].expand(bs, next_n),
+            )
+            expected_pages = table[req].repeat_interleave(next_n, dim=0)
+            row_lens = (seq + next_n).repeat_interleave(next_n)
+            live = torch.arange(width, device=device)[None, :] < row_lens[:, None]
+            torch.testing.assert_close(
+                buffers["page_table_1"][live], expected_pages[live]
+            )
+            real_live = (
+                torch.arange(0, width, 64, device=device)[None, :] < row_lens[:, None]
+            )
+            torch.testing.assert_close(
+                buffers["real_page_table"][real_live],
+                (expected_pages[:, ::64] // 64)[real_live],
+            )
+            self.assertEqual(
+                addresses, {key: value.data_ptr() for key, value in buffers.items()}
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_dsa_metadata_replay.py

@@ -0,0 +1,117 @@
+"""Fused KPool replay and MTP sibling copies preserve captured buffer identity."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.dsa_metadata_kit import (
+    BS,
+    NEXT_N,
+    ROUNDS,
+    addresses,
+    apply_metadata,
+    assert_metadata_equal,
+    inputs,
+    make_backend,
+)
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestDSAMetadataReplay(CustomTestCase):
+    def test_fusion_matches_ordinary_metadata(self):
+        for mode in (
+            ForwardMode.DECODE,
+            ForwardMode.TARGET_VERIFY,
+            ForwardMode.DRAFT_EXTEND_V2,
+        ):
+            with self.subTest(mode=mode):
+                seq, req = inputs(*ROUNDS[0])
+                fused = make_backend(mode, seq, req)
+                ordinary = make_backend(mode, seq, req, fusion=False)
+                pointers = addresses(fused.forward_metadata)
+                for lengths, requests in ROUNDS:
+                    seq.copy_(torch.tensor(lengths, device="cuda"))
+                    req.copy_(torch.tensor(requests, device="cuda"))
+                    spec = None
+                    if mode.is_draft_extend_v2():
+                        spec = SimpleNamespace(
+                            num_accept_tokens=torch.tensor(
+                                [1, 2, 5, NEXT_N], device="cuda", dtype=torch.int32
+                            )
+                        )
+                    apply_metadata(fused, mode, seq, req, spec)
+                    apply_metadata(ordinary, mode, seq, req, spec)
+                    assert_metadata_equal(
+                        self, fused.forward_metadata, ordinary.forward_metadata
+                    )
+                    self.assertEqual(pointers, addresses(fused.forward_metadata))
+
+    def test_precomputed_verify_retains_live_tail(self):
+        mode = ForwardMode.TARGET_VERIFY
+        seq, req = inputs(*ROUNDS[0])
+        fused = make_backend(mode, seq, req)
+        ordinary = make_backend(mode, seq, req, fusion=False)
+        pointers = addresses(fused.forward_metadata)
+        for lengths, requests in ROUNDS[1:]:
+            seq.copy_(torch.tensor(lengths, device="cuda"))
+            req.copy_(torch.tensor(requests, device="cuda"))
+            precomputed = fused._precompute_replay_metadata(
+                BS, req, seq, seq.cpu(), mode
+            )
+            fused.init_forward_metadata_replay_cuda_graph_from_precomputed(
+                BS, precomputed, mode
+            )
+            apply_metadata(ordinary, mode, seq, req)
+            assert_metadata_equal(
+                self, fused.forward_metadata, ordinary.forward_metadata
+            )
+            self.assertEqual(pointers, addresses(fused.forward_metadata))
+
+    def test_precomputed_and_sibling_copy_refresh_derived_metadata(self):
+        mode = ForwardMode.DECODE
+        seq, req = inputs(*ROUNDS[0])
+        source = make_backend(mode, seq, req)
+        sibling = make_backend(mode, seq, req)
+        ordinary = make_backend(mode, seq, req, fusion=False)
+        pointers = addresses(sibling.forward_metadata)
+        for lengths, requests in ROUNDS[1:]:
+            seq.copy_(torch.tensor(lengths, device="cuda"))
+            req.copy_(torch.tensor(requests, device="cuda"))
+            precomputed = source._precompute_replay_metadata(
+                BS, req, seq, seq.cpu(), mode
+            )
+            source.init_forward_metadata_replay_cuda_graph_from_precomputed(
+                BS, precomputed, mode
+            )
+            # An eligible sibling must reuse the derived results, not silently
+            # fall through to the full recomputation path.
+            with patch.object(
+                sibling,
+                "init_forward_metadata_replay_cuda_graph_from_precomputed",
+                side_effect=AssertionError("unexpected sibling fallback"),
+            ):
+                sibling._copy_replay_metadata_from_sibling(
+                    source, BS, precomputed, mode
+                )
+            apply_metadata(ordinary, mode, seq, req)
+            assert_metadata_equal(
+                self, source.forward_metadata, ordinary.forward_metadata
+            )
+            assert_metadata_equal(
+                self, sibling.forward_metadata, ordinary.forward_metadata
+            )
+            self.assertEqual(pointers, addresses(sibling.forward_metadata))
+            self.assertIsNot(
+                sibling.forward_metadata.kpool_write_plan,
+                source.forward_metadata.kpool_write_plan,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_dsv4_q_rope_store.py

@@ -0,0 +1,94 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.elementwise import fused_rope_inplace
+from sglang.kernels.ops.attention.dsv4.q_rope_store import q_rope_store
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestQRopeStore(CustomTestCase):
+    def test_exact_output_and_padding(self):
+        torch.manual_seed(911)
+        freqs = torch.polar(
+            torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda")
+        )
+        for rows in (1, 2, 5, 6, 8):
+            for heads in (8, 16, 32):
+                for dtype in (torch.int32, torch.int64):
+                    q = torch.randn(
+                        rows, heads + 1, 512, device="cuda", dtype=torch.bfloat16
+                    )[:, :heads]
+                    padding = torch.full(
+                        (rows, 64, 512), 7.0, device="cuda", dtype=q.dtype
+                    )
+                    output = padding[:, :heads]
+                    positions = torch.randint(
+                        0, 8192, (rows,), device="cuda", dtype=dtype
+                    )
+                    original = q.clone()
+                    expected = q.clone()
+                    fused_rope_inplace(expected[..., 448:], None, freqs, positions)
+                    q_rope_store(q, output, freqs, positions)
+                    torch.testing.assert_close(output, expected, rtol=0, atol=0)
+                    torch.testing.assert_close(q, original, rtol=0, atol=0)
+                    self.assertTrue((padding[:, heads:] == 7).all().item())
+
+    def test_large_prefill_exact_output_and_padding(self):
+        torch.manual_seed(911)
+        freqs = torch.polar(
+            torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda")
+        )
+        for rows in (4096, 4097, 65536):
+            for dtype in (torch.int32, torch.int64):
+                with self.subTest(rows=rows, dtype=dtype):
+                    q = torch.randn(rows, 17, 512, device="cuda", dtype=torch.bfloat16)[
+                        :, :16
+                    ]
+                    original = q.clone()
+                    expected = q.clone()
+                    padding = torch.full(
+                        (rows, 64, 512), 7.0, device="cuda", dtype=q.dtype
+                    )
+                    positions = torch.randint(
+                        0, 8192, (rows,), device="cuda", dtype=dtype
+                    )
+                    fused_rope_inplace(expected[..., 448:], None, freqs, positions)
+                    q_rope_store(q, padding[:, :16], freqs, positions)
+                    torch.testing.assert_close(
+                        padding[:, :16], expected, rtol=0, atol=0
+                    )
+                    torch.testing.assert_close(q, original, rtol=0, atol=0)
+                    self.assertTrue((padding[:, 16:] == 7).all().item())
+
+    def _check_graph_replay(self, rows):
+        q = torch.randn(rows, 16, 512, device="cuda", dtype=torch.bfloat16)
+        output = torch.zeros(rows, 64, 512, device="cuda", dtype=q.dtype)[:, :16]
+        freqs = torch.polar(
+            torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda")
+        )
+        positions = torch.arange(rows, device="cuda") % 8192
+        q_rope_store(q, output, freqs, positions)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            q_rope_store(q, output, freqs, positions)
+        for _ in range(3):
+            q.normal_()
+            positions.random_(0, 8192)
+            graph.replay()
+            expected = q.clone()
+            fused_rope_inplace(expected[..., 448:], None, freqs, positions)
+            torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+    def test_graph_replay(self):
+        self._check_graph_replay(6)
+
+    def test_large_prefill_graph_replay(self):
+        self._check_graph_replay(4097)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_flashmla_sched_meta.py

@@ -0,0 +1,167 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.runtime_context import get_platform
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.version.cuda is None
+    or not get_platform().is_blackwell,
+    reason="the FlashMLA split-KV schedule is Blackwell-only here",
+)
+
+H_Q, D_QK, D_V = 64, 512, 512
+# The only cache format this build's sparse decode takes for both caches, and a
+# page size that keeps page * bytes_per_token a multiple of 576.
+BYTES_PER_TOKEN, PAGE = 584, 288
+BLOCK_SIZE_N, FIXED_OVERHEAD = 64, 5
+
+
+def _cache():
+    n = 8192 // PAGE + 16
+    return torch.randint(
+        0, 200, (n, PAGE, 1, BYTES_PER_TOKEN), device="cuda", dtype=torch.uint8
+    )
+
+
+def _call(kv, b, s_q, topk, topk_length, *, extra=None, meta=None):
+    import sgl_kernel.flash_mla as flash_mla
+
+    g = torch.Generator(device="cuda").manual_seed(b * 7 + s_q * 13 + topk)
+    q = torch.randn(
+        (b, s_q, H_Q, D_QK), device="cuda", dtype=torch.bfloat16, generator=g
+    )
+    indices = torch.randint(
+        0, 5120, (b, s_q, topk), device="cuda", dtype=torch.int32, generator=g
+    )
+    sink = torch.randn((H_Q,), device="cuda", dtype=torch.float32, generator=g)
+    kwargs = {}
+    if extra is not None:
+        extra_kv, extra_indices, extra_topk_length = extra
+        kwargs = dict(
+            extra_k_cache=extra_kv,
+            extra_indices_in_kvcache=extra_indices,
+            extra_topk_length=extra_topk_length,
+        )
+    sched = flash_mla.FlashMLASchedMeta()
+    if meta is not None:
+        sched.tile_scheduler_metadata, sched.num_splits = meta
+    out, lse = flash_mla.flash_mla_with_kvcache(
+        q,
+        kv,
+        None,
+        None,
+        D_V,
+        sched,
+        indices=indices,
+        is_fp8_kvcache=True,
+        softmax_scale=0.1,
+        causal=False,
+        topk_length=topk_length,
+        attn_sink=sink,
+        **kwargs,
+    )
+    torch.cuda.synchronize()
+    return out, lse, sched
+
+
+def _ours(
+    like_meta, like_splits, topk_length, topk, *, extra_topk_length=None, extra_topk=0
+):
+    from sglang.kernels.ops.attention.dsv4.flashmla_sched_meta import (
+        flashmla_sched_meta,
+    )
+
+    meta = torch.empty_like(like_meta)
+    splits = torch.empty_like(like_splits)
+    flashmla_sched_meta(
+        meta,
+        splits,
+        topk_length=topk_length,
+        extra_topk_length=extra_topk_length,
+        block_size_n=BLOCK_SIZE_N,
+        fixed_overhead_num_blocks=FIXED_OVERHEAD,
+        topk=topk,
+        extra_topk=extra_topk,
+    )
+    return meta, splits
+
+
+def _lengths(b, topk, mode, seed):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    if mode == "full":
+        return torch.full((b,), topk, device="cuda", dtype=torch.int32)
+    if mode == "ones":
+        return torch.ones((b,), device="cuda", dtype=torch.int32)
+    if mode == "zeros":
+        return torch.zeros((b,), device="cuda", dtype=torch.int32)
+    lengths = torch.randint(
+        0, topk + 1, (b,), device="cuda", dtype=torch.int32, generator=g
+    )
+    if mode == "mixed":
+        lengths[0] = 0
+        lengths[-1] = topk
+    return lengths
+
+
+# FlashMLA's DecodingSchedMeta ends in a `_pad` word it never writes, so the
+# reference carries whatever torch::empty left there.
+DEFINED = slice(0, 7)
+
+
+@pytest.mark.parametrize("b", [1, 6, 64])
+@pytest.mark.parametrize("s_q", [1, 6])
+@pytest.mark.parametrize("topk", [2048])
+@pytest.mark.parametrize("mode", ["full", "random", "zeros", "ones", "mixed"])
+def test_matches_flashmla_schedule(b: int, s_q: int, topk: int, mode: str):
+    kv = _cache()
+    topk_length = _lengths(b, topk, mode, b * 1000 + s_q * 37 + topk + len(mode))
+    _, _, sched = _call(kv, b, s_q, topk, topk_length)
+    if sched.tile_scheduler_metadata is None:
+        pytest.skip("FlashMLA did not split the KV for this shape")
+    meta, splits = _ours(
+        sched.tile_scheduler_metadata, sched.num_splits, topk_length, topk
+    )
+    assert torch.equal(meta[:, DEFINED], sched.tile_scheduler_metadata[:, DEFINED])
+    assert torch.equal(splits, sched.num_splits)
+
+
+@pytest.mark.parametrize("b", [1, 8])
+@pytest.mark.parametrize("topk,extra_topk", [(512, 512), (2048, 512), (512, 2048)])
+def test_extra_cache_schedule(b: int, topk: int, extra_topk: int):
+    kv, extra_kv = _cache(), _cache()
+    s_q = 1
+    g = torch.Generator(device="cuda").manual_seed(b + topk + extra_topk)
+    topk_length = _lengths(b, topk, "random", b + topk)
+    extra_topk_length = torch.randint(
+        1, extra_topk + 1, (b,), device="cuda", dtype=torch.int32, generator=g
+    )
+    extra_indices = torch.randint(
+        0, 4096, (b, s_q, extra_topk), device="cuda", dtype=torch.int32, generator=g
+    )
+    extra = (extra_kv, extra_indices, extra_topk_length)
+    ref_out, ref_lse, sched = _call(kv, b, s_q, topk, topk_length, extra=extra)
+    if sched.tile_scheduler_metadata is None:
+        pytest.skip("FlashMLA did not split the KV for this shape")
+    meta, splits = _ours(
+        sched.tile_scheduler_metadata,
+        sched.num_splits,
+        topk_length,
+        topk,
+        extra_topk_length=extra_topk_length,
+        extra_topk=extra_topk,
+    )
+    assert torch.equal(meta[:, DEFINED], sched.tile_scheduler_metadata[:, DEFINED])
+    assert torch.equal(splits, sched.num_splits)
+    out, lse, _ = _call(kv, b, s_q, topk, topk_length, extra=extra, meta=(meta, splits))
+    assert torch.equal(out.view(torch.int16), ref_out.view(torch.int16))
+    assert torch.equal(lse.view(torch.int32), ref_lse.view(torch.int32))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/attention/test_kda_fused_verify_backend.py

@@ -0,0 +1,322 @@
+"""KDA backend dispatch and ReplaySSM verify -> commit -> verify parity."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import torch
+
+from sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify import (
+    fused_kda_conv_gating_verify,
+)
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    HybridLinearAttnBackend,
+)
+from sglang.srt.layers.attention.linear.kda_backend import (
+    KDAAttnBackend,
+    KDAKernelDispatcher,
+)
+from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend
+from sglang.srt.mem_cache.memory_pool import MambaPool
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.runtime_context import override_platform
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+# One bf16 ulp: the fused and unfused kernels reduce K in different orders.
+_OUTPUT_TOL = dict(rtol=2**-7, atol=1e-7)
+# Fold and snapshot recurrence differ by ~1 fp32 ulp; a wrong-step commit
+# moves elements by >= 1e-3, so 1e-5 still catches it.
+_SNAPSHOT_ORACLE_TOL = dict(rtol=0, atol=1e-5)
+
+
+class TestKDAFusedVerifyBackend(CustomTestCase):
+    def _make_case(self, batch_size=1, heads=2, v_heads=4, lower_bound=-5.0):
+        torch.manual_seed(36821)
+        steps, head_dim, num_layers = 4, 128, 2
+        num_slots = batch_size + 3
+        dim = (2 * heads + v_heads) * head_dim
+
+        def randn(*shape, dtype=torch.bfloat16):
+            return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2
+
+        layers = [
+            SimpleNamespace(
+                layer_id=i,
+                num_q_heads=heads,
+                num_k_heads=heads,
+                num_v_heads=v_heads,
+                head_q_dim=head_dim,
+                head_k_dim=head_dim,
+                head_v_dim=head_dim,
+                q_dim=heads * head_dim,
+                k_dim=heads * head_dim,
+                v_dim=v_heads * head_dim,
+                conv_weights=randn(dim, 4),
+                bias=randn(dim),
+                A_log=randn(v_heads, dtype=torch.float32),
+                dt_bias=randn(v_heads * head_dim, dtype=torch.float32),
+                lower_bound=lower_bound,
+            )
+            for i in range(num_layers)
+        ]
+        state = MambaPool.SpeculativeState(
+            conv=[randn(num_layers, num_slots, 3, dim)],
+            temporal=randn(
+                num_layers,
+                num_slots,
+                v_heads,
+                head_dim,
+                head_dim,
+                dtype=torch.float32,
+            ),
+            intermediate_ssm=None,
+            intermediate_conv_window=[randn(num_layers, batch_size, steps, 3, dim)],
+            replayssm_rawv=randn(num_layers, num_slots, v_heads, 16, head_dim),
+            replayssm_rawk=randn(num_layers, num_slots, heads, 16, head_dim),
+            replayssm_g=randn(
+                num_layers, num_slots, v_heads, 16, head_dim, dtype=torch.float32
+            ),
+            replayssm_beta=randn(
+                num_layers, num_slots, v_heads, 16, dtype=torch.float32
+            ),
+        )
+        # Physical slots differ from scratch rows; reverse them to catch callers
+        # accidentally committing by request index instead of mamba slot.
+        slots = torch.arange(batch_size + 1, 1, -1, device="cuda", dtype=torch.int32)
+        batch = SimpleNamespace(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            spec_info=SimpleNamespace(draft_token_num=steps, ragged_verify_layout=None),
+        )
+        rounds = [
+            [
+                (
+                    randn(batch_size * steps, dim),
+                    randn(1, batch_size * steps, v_heads * head_dim),
+                    randn(1, batch_size * steps, v_heads),
+                )
+                for _ in layers
+            ]
+            for _ in range(2)
+        ]
+        return layers, state, slots, batch, rounds
+
+    def _make_backend(self, template, slots, steps, *, fused, ring=True):
+        state = MambaPool.SpeculativeState(
+            conv=[template.conv[0].clone()],
+            temporal=template.temporal.clone(),
+            intermediate_conv_window=[template.intermediate_conv_window[0].clone()],
+            intermediate_ssm=(
+                None
+                if ring
+                else template.temporal.new_zeros(
+                    template.temporal.shape[0],
+                    slots.numel(),
+                    steps,
+                    *template.temporal.shape[2:],
+                )
+            ),
+            **{
+                name: getattr(template, name).clone() if ring else None
+                for name in (
+                    "replayssm_rawv",
+                    "replayssm_rawk",
+                    "replayssm_g",
+                    "replayssm_beta",
+                )
+            },
+        )
+        # Only the model/pool setup is a fixture. Verify, dispatch, ring fold and
+        # conv rollback below all use the production backend and GPU kernels.
+        backend = KDAAttnBackend.__new__(KDAAttnBackend)
+        backend.req_to_token_pool = SimpleNamespace(
+            mamba2_layer_cache=state.at_layer_idx,
+            get_speculative_mamba2_params_all_layers=lambda: state,
+            mamba_pool=SimpleNamespace(replayssm_is_kda=ring),
+        )
+        backend.forward_metadata = SimpleNamespace(
+            query_start_loc=torch.arange(
+                slots.numel() + 1, device="cuda", dtype=torch.int32
+            )
+            * steps,
+            mamba_cache_indices=slots,
+            retrieve_next_token=None,
+            retrieve_next_sibling=None,
+            retrieve_parent_token=None,
+        )
+        backend.verify_intermediate_state_indices = torch.arange(
+            slots.numel(), device="cuda", dtype=torch.int32
+        )
+        backend.accept_lens_pool = None
+        backend.kernel_dispatcher = KDAKernelDispatcher(
+            LinearAttnKernelBackend.TRITON,
+            LinearAttnKernelBackend.TRITON,
+            LinearAttnKernelBackend.TRITON,
+        )
+        backend._fused_chain_verify_fn = (
+            Mock(wraps=fused_kda_conv_gating_verify) if fused else None
+        )
+        hybrid = HybridLinearAttnBackend.__new__(HybridLinearAttnBackend)
+        hybrid.linear_attn_backend = backend
+        return backend, hybrid, state
+
+    @staticmethod
+    def _verify(backend, layers, batch, inputs):
+        return [
+            backend.forward_extend(layer, batch, mixed, a, b)
+            for layer, (mixed, a, b) in zip(layers, inputs)
+        ]
+
+    def test_verify_commit_verify(self):
+        # B=1 exercises the enabled path. Platform override makes the dispatch
+        # testable on any CUDA CI runner; it does not replace a GPU kernel.
+        for platform, (heads, v_heads, lower_bound), num_accept_tokens in (
+            ({"is_sm90": True}, (2, 2, None), 1),
+            ({"is_sm90": True}, (2, 2, None), 2),
+            ({"is_sm90": True}, (2, 2, None), 4),
+            ({"is_sm90": True}, (2, 4, -5.0), 1),
+            ({"is_sm90": True}, (2, 4, -5.0), 2),
+            ({"is_sm90": True}, (2, 4, -5.0), 4),
+            ({"is_sm90": False, "is_sm100": True}, (2, 4, -5.0), 2),
+        ):
+            with (
+                self.subTest(
+                    platform=platform,
+                    heads=heads,
+                    v_heads=v_heads,
+                    lower_bound=lower_bound,
+                    num_accept_tokens=num_accept_tokens,
+                ),
+                override_platform(**platform),
+            ):
+                layers, initial, slots, batch, rounds = self._make_case(
+                    heads=heads, v_heads=v_heads, lower_bound=lower_bound
+                )
+                fused, fused_hybrid, fused_state = self._make_backend(
+                    initial, slots, 4, fused=True
+                )
+                reference, ref_hybrid, ref_state = self._make_backend(
+                    initial, slots, 4, fused=False
+                )
+                snapshots, _, snapshot_state = self._make_backend(
+                    initial, slots, 4, fused=False, ring=False
+                )
+                out_fused = self._verify(fused, layers, batch, rounds[0])
+                out_ref = self._verify(reference, layers, batch, rounds[0])
+                self._verify(snapshots, layers, batch, rounds[0])
+                for actual, expected in zip(out_fused, out_ref):
+                    torch.testing.assert_close(actual, expected, **_OUTPUT_TOL)
+                for state in (fused_state, ref_state):
+                    torch.testing.assert_close(
+                        state.temporal, initial.temporal, rtol=0, atol=0
+                    )
+
+                last_steps = torch.full_like(slots, num_accept_tokens - 1)
+                for hybrid in (fused_hybrid, ref_hybrid):
+                    hybrid.update_mamba_state_after_mtp_verify(
+                        last_correct_step_indices=last_steps,
+                        mamba_track_indices=None,
+                        mamba_steps_to_track=None,
+                        model=None,
+                    )
+                torch.testing.assert_close(
+                    fused_state.temporal, ref_state.temporal, rtol=0, atol=0
+                )
+                torch.testing.assert_close(
+                    fused_state.conv[0], ref_state.conv[0], rtol=0, atol=0
+                )
+                # Independent snapshot oracle: equality between two ring
+                # arms alone would miss a shared no-op / wrong-step commit.
+                expected_ssm = initial.temporal.clone()
+                expected_ssm[:, slots.long()] = snapshot_state.intermediate_ssm[
+                    :, :, num_accept_tokens - 1
+                ]
+                torch.testing.assert_close(
+                    fused_state.temporal, expected_ssm, **_SNAPSHOT_ORACLE_TOL
+                )
+                expected_conv = initial.conv[0].clone()
+                for i, (mixed, _, _) in enumerate(rounds[0]):
+                    history = torch.cat(
+                        (
+                            initial.conv[0][i, slots.long()],
+                            mixed.view(1, 4, -1)[:, :num_accept_tokens],
+                        ),
+                        dim=1,
+                    )
+                    expected_conv[i, slots.long()] = history[:, -3:]
+                torch.testing.assert_close(
+                    fused_state.conv[0], expected_conv, rtol=0, atol=0
+                )
+
+                out_fused = self._verify(fused, layers, batch, rounds[1])
+                out_ref = self._verify(reference, layers, batch, rounds[1])
+                for actual, expected in zip(out_fused, out_ref):
+                    torch.testing.assert_close(actual, expected, **_OUTPUT_TOL)
+                self.assertEqual(fused._fused_chain_verify_fn.call_count, 4)
+
+    def test_ring_dispatch_falls_back(self):
+        # B=2 and the measured regression sizes on the enabled architectures,
+        # plus B=1 on an architecture without ring measurements.
+        sm90 = {"is_sm90": True, "is_sm100": False}
+        sm100 = {"is_sm90": False, "is_sm100": True}
+        other = {"is_sm90": False, "is_sm100": False}
+        for platform, batch_size in (
+            (sm90, 2),
+            (sm90, 4),
+            (sm90, 16),
+            (sm90, 64),
+            (sm100, 2),
+            (sm100, 4),
+            (sm100, 16),
+            (other, 1),
+            (other, 4),
+        ):
+            with (
+                self.subTest(platform=platform, batch_size=batch_size),
+                override_platform(**platform),
+            ):
+                layers, initial, slots, batch, rounds = self._make_case(batch_size)
+                backend, _, state = self._make_backend(initial, slots, 4, fused=True)
+                reference, _, ref_state = self._make_backend(
+                    initial, slots, 4, fused=False
+                )
+                out = self._verify(backend, layers, batch, rounds[0])
+                ref = self._verify(reference, layers, batch, rounds[0])
+                backend._fused_chain_verify_fn.assert_not_called()
+                for actual, expected in zip(out, ref):
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+                for name in (
+                    "replayssm_rawv",
+                    "replayssm_rawk",
+                    "replayssm_g",
+                    "replayssm_beta",
+                ):
+                    torch.testing.assert_close(
+                        getattr(state, name), getattr(ref_state, name), rtol=0, atol=0
+                    )
+
+    def test_snapshot_dispatch_is_unchanged(self):
+        for platform in (
+            {"is_sm90": True, "is_sm100": False},
+            {"is_sm90": False, "is_sm100": True},
+            {"is_sm90": False, "is_sm100": False},
+        ):
+            with self.subTest(platform=platform), override_platform(**platform):
+                layers, initial, slots, batch, rounds = self._make_case(batch_size=4)
+                backend, _, _ = self._make_backend(
+                    initial, slots, 4, fused=True, ring=False
+                )
+                reference, _, _ = self._make_backend(
+                    initial, slots, 4, fused=False, ring=False
+                )
+                out = self._verify(backend, layers, batch, rounds[0])
+                ref = self._verify(reference, layers, batch, rounds[0])
+                self.assertEqual(backend._fused_chain_verify_fn.call_count, 2)
+                for actual, expected in zip(out, ref):
+                    torch.testing.assert_close(actual, expected, **_OUTPUT_TOL)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_kda_track_state.py

@@ -0,0 +1,183 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.fla.kda import chunk_kda
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=180, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+CHUNK_SIZE = 64
+
+_BACKENDS = {"triton": chunk_kda}
+HELION_AVAILABLE = True
+try:
+    import helion  # noqa: F401
+except ModuleNotFoundError as error:
+    # A broken install (transitive import failure) must stay loud; only the
+    # absent package downgrades the run to triton-only.
+    if error.name != "helion":
+        raise
+    HELION_AVAILABLE = False
+if HELION_AVAILABLE:
+    from sglang.kernels.ops.attention.helion.kda_prefill import (
+        chunk_kda as helion_chunk_kda,
+    )
+
+    _BACKENDS["helion"] = helion_chunk_kda
+
+
+def _make_varlen_inputs(seed, lens, num_heads=2, head_dim=128):
+    """Packed varlen KDA inputs: [1, sum(lens), H, D] plus a zero fp32 state pool."""
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    total = sum(lens)
+
+    def randn(*shape, dtype=torch.bfloat16):
+        return torch.randn(*shape, generator=generator, device="cuda", dtype=dtype)
+
+    q = randn(1, total, num_heads, head_dim)
+    k = randn(1, total, num_heads, head_dim)
+    v = (0.1 * randn(1, total, num_heads, head_dim, dtype=torch.float32)).to(
+        torch.bfloat16
+    )
+    gate = randn(1, total, num_heads, head_dim)
+    beta = torch.sigmoid(randn(1, total, num_heads, dtype=torch.float32)).to(
+        torch.bfloat16
+    )
+    a_log = randn(num_heads, dtype=torch.float32)
+    dt_bias = randn(num_heads * head_dim, dtype=torch.float32)
+    state = torch.zeros(
+        len(lens), num_heads, head_dim, head_dim, device="cuda", dtype=torch.float32
+    )
+    cu_seqlens = torch.tensor(
+        [0, *torch.tensor(lens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
+    )
+    return q, k, v, gate, beta, a_log, dt_bias, state, cu_seqlens
+
+
+def _run_chunk_kda(
+    chunk_kda_fn, q, k, v, gate, beta, a_log, dt_bias, state, cu_seqlens, **kwargs
+):
+    return chunk_kda_fn(
+        # chunk_kda writes in place (the attention output lands in v, the gate
+        # cumsum in g); hand every run fresh copies so runs stay independent.
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=gate.clone(),
+        beta=beta.clone(),
+        scale=q.shape[-1] ** -0.5,
+        initial_state=state,
+        initial_state_indices=torch.arange(
+            state.shape[0], device="cuda", dtype=torch.int32
+        ),
+        use_qk_l2norm_in_kernel=True,
+        cu_seqlens=cu_seqlens,
+        A_log=a_log,
+        dt_bias=dt_bias,
+        lower_bound=-5.0,
+        **kwargs,
+    )
+
+
+class TestKdaTrackState(CustomTestCase):
+    def test_helion_backend_ran(self):
+        """Visibility hook: without helion installed the snapshot check above
+        runs triton-only and the Helion track configs go untested — surface
+        that as an explicit skip instead of a silent pass."""
+        if not HELION_AVAILABLE:
+            self.skipTest("helion is not installed; triton backend only")
+
+    @torch.inference_mode()
+    def test_track_state_snapshots_fp32_accumulator(self):
+        """Bug regression: the mamba radix track path snapshots the SSM state at
+        the last chunk boundary of unaligned sequences into the fp32 state pool.
+        It used to read the per-chunk states `h` (activation dtype, bf16), so a
+        prefix-cache hit restored a bf16-rounded state while a cache miss kept
+        fp32. `track_state` must carry the in-kernel fp32 accumulator: identical
+        to the fp32 final state of a run truncated at the boundary, and strictly
+        more precise than the bf16 `h` row for the same boundary.
+        """
+        if not torch.cuda.is_available():
+            self.skipTest("requires CUDA")
+        for backend, chunk_kda_fn in _BACKENDS.items():
+            # num_heads=2 exercises the Helion small-head track config; 16
+            # crosses _PREFILL_SMALL_HEAD_THRESHOLD (12) to exercise the
+            # large-head varlen track config that real models take.
+            for num_heads in (2, 16):
+                with self.subTest(backend=backend, num_heads=num_heads):
+                    self._check_track_state(chunk_kda_fn, num_heads)
+
+    def _check_track_state(self, chunk_kda_fn, num_heads):
+        # seq0: 100 tokens, unaligned -> snapshot at the 64-token boundary
+        # (start of chunk 1). seq1: 64 tokens, aligned -> not tracked.
+        lens = [100, 64]
+        q, k, v, gate, beta, a_log, dt_bias, state, cu_seqlens = _make_varlen_inputs(
+            0, lens, num_heads=num_heads
+        )
+        num_heads, head_dim = q.shape[2], q.shape[3]
+
+        track_state = torch.full(
+            (len(lens), num_heads, head_dim, head_dim),
+            float("nan"),
+            device="cuda",
+            dtype=torch.float32,
+        )
+        track_chunk_idx = torch.tensor([1, -1], dtype=torch.int32, device="cuda")
+        _, h = _run_chunk_kda(
+            chunk_kda_fn,
+            q,
+            k,
+            v,
+            gate,
+            beta,
+            a_log,
+            dt_bias,
+            state,
+            cu_seqlens,
+            output_intermediate_states=True,
+            track_state=track_state,
+            track_chunk_idx=track_chunk_idx,
+        )
+
+        # The untracked row must stay untouched; the tracked row must be finite.
+        self.assertTrue(torch.all(torch.isnan(track_state[1])))
+        self.assertFalse(torch.any(torch.isnan(track_state[0])))
+
+        # Reference: truncate seq0 at the boundary; the pool's fp32 row then
+        # receives the in-place final state for the same prefix — the
+        # established fp32 path the snapshot must agree with.
+        ref_state = torch.zeros(
+            1, num_heads, head_dim, head_dim, device="cuda", dtype=torch.float32
+        )
+        ref_cu_seqlens = torch.tensor([0, CHUNK_SIZE], dtype=torch.int32, device="cuda")
+        _run_chunk_kda(
+            chunk_kda_fn,
+            q[:, :CHUNK_SIZE],
+            k[:, :CHUNK_SIZE],
+            v[:, :CHUNK_SIZE],
+            gate[:, :CHUNK_SIZE],
+            beta[:, :CHUNK_SIZE],
+            a_log,
+            dt_bias,
+            ref_state,
+            ref_cu_seqlens,
+        )
+        torch.testing.assert_close(track_state[0], ref_state[0], rtol=1e-5, atol=1e-5)
+
+        # The guard: h packs one row per (seq, chunk); row 1 is seq0's state at
+        # the boundary, rounded to bf16. If the snapshot were re-routed through
+        # h, it could not match the fp32 reference above.
+        self.assertTrue(
+            torch.equal(h[0, 1].float(), track_state[0].to(torch.bfloat16).float()),
+            "h row should be exactly the bf16 rounding of the fp32 snapshot",
+        )
+        self.assertFalse(
+            torch.equal(track_state[0], track_state[0].to(torch.bfloat16).float()),
+            "test inputs must make bf16 rounding lossy",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_subblock_sage_fp8_sm90.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Correctness tests for native SM90 SubBlock Sage FP8 attention."""
+
+import math
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _native_sm90_sage_available() -> bool:
+    try:
+        import spas_sage_attn._qattn as qattn
+    except (ImportError, OSError):
+        return False
+    return hasattr(
+        qattn,
+        "qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90",
+    )
+
+
+requires_native_sm90_sage = unittest.skipUnless(
+    torch.cuda.is_available()
+    and torch.cuda.get_device_capability() == (9, 0)
+    and _native_sm90_sage_available(),
+    "requires SM90 and a compiled SpargeAttention installation",
+)
+
+
+def _masked_reference(q, k, v, index, counts, scale, key_block_size=128):
+    logits = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * scale
+    mask = torch.zeros_like(logits, dtype=torch.bool)
+    for b in range(q.shape[0]):
+        for h in range(q.shape[2]):
+            for qb in range(index.shape[2]):
+                q_slice = slice(qb * 64, min((qb + 1) * 64, q.shape[1]))
+                for slot in range(int(counts[b, h, qb])):
+                    kb = int(index[b, h, qb, slot])
+                    k_slice = slice(
+                        kb * key_block_size,
+                        min((kb + 1) * key_block_size, k.shape[1]),
+                    )
+                    mask[b, h, q_slice, k_slice] = True
+    logits.masked_fill_(~mask, -float("inf"))
+    p = torch.softmax(logits, dim=-1)
+    return torch.einsum("bhqk,bkhd->bqhd", p, v.float()).to(torch.bfloat16)
+
+
+def _cosine(a, b):
+    return float(
+        torch.nn.functional.cosine_similarity(
+            a.float().flatten(), b.float().flatten(), dim=0
+        )
+    )
+
+
+@requires_native_sm90_sage
+class TestSubBlockSageFp8NativeSm90(CustomTestCase):
+    def test_full_budget_ragged_tail_reproduces_dense(self):
+        from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
+            subblock_sage_fp8_sm90_attention,
+        )
+
+        torch.manual_seed(19)
+        seq_len = 1024 + 37
+        heads = 4
+        shape = (1, seq_len, heads, 128)
+        q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn_like(q)
+        v = torch.randn_like(q)
+        scale = 1.0 / math.sqrt(128)
+
+        query_blocks = math.ceil(seq_len / 64)
+        key_blocks = math.ceil(seq_len / 128)
+        index = (
+            torch.arange(key_blocks, device="cuda", dtype=torch.int32)
+            .view(1, 1, 1, key_blocks)
+            .expand(1, heads, query_blocks, key_blocks)
+            .contiguous()
+        )
+        output = subblock_sage_fp8_sm90_attention(q, k, v, index, key_blocks, scale)
+        reference = torch.nn.functional.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            scale=scale,
+        ).transpose(1, 2)
+
+        self.assertTrue(torch.isfinite(output.float()).all())
+        self.assertGreater(_cosine(output, reference), 0.998)
+
+    def test_sparse_k128_plan_and_variable_counts(self):
+        from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
+            subblock_sage_fp8_sm90_attention,
+        )
+
+        torch.manual_seed(23)
+        seq_len = 1024 + 37
+        heads = 4
+        shape = (1, seq_len, heads, 128)
+        q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn_like(q)
+        v = torch.randn_like(q)
+        scale = 1.0 / math.sqrt(128)
+
+        query_blocks = math.ceil(seq_len / 64)
+        key_blocks = math.ceil(seq_len / 128)
+        width = 4
+        index = torch.stack(
+            [
+                torch.roll(torch.arange(key_blocks), shifts=query_block)[:width]
+                for query_block in range(query_blocks)
+            ]
+        )
+        index = (
+            index.to(device="cuda", dtype=torch.int32)
+            .view(1, 1, query_blocks, width)
+            .expand(1, heads, query_blocks, width)
+            .contiguous()
+        )
+        counts = (
+            ((torch.arange(query_blocks, device="cuda", dtype=torch.int32) % width) + 1)
+            .view(1, 1, query_blocks)
+            .expand(1, heads, query_blocks)
+            .contiguous()
+        )
+
+        output = subblock_sage_fp8_sm90_attention(q, k, v, index, width, scale, counts)
+        reference = _masked_reference(q, k, v, index, counts, scale)
+
+        self.assertTrue(torch.isfinite(output.float()).all())
+        self.assertGreater(_cosine(output, reference), 0.997)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)

test/registered/kernels/ops/communication/test_deterministic_reduce_scatter.py

@@ -0,0 +1,80 @@
+"""A token must retain its bits when batch size changes its receiving TP rank.
+
+Run with ``python test_deterministic_reduce_scatter.py --num-gpu 4``.
+This exercises the collective used before DeepEP MoE without a model or DeepEP.
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from sglang.srt.distributed import parallel_state as ps
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kernels.utils import multigpu_pytest_main
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+@pytest.fixture(scope="module")
+def group():
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    # Match --enable-deterministic-inference's CUDA collective settings.
+    os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"] = "1"
+    os.environ["NCCL_ALGO"] = "allreduce:tree"
+    nchannels = str(envs.SGLANG_DETERMINISTIC_NCCL_NCHANNELS.get())
+    os.environ["NCCL_MIN_NCHANNELS"] = nchannels
+    os.environ["NCCL_MAX_NCHANNELS"] = nchannels
+    torch.cuda.set_device(local_rank)
+    ps.set_custom_all_reduce(False)
+    ps.init_distributed_environment(
+        world_size=world_size,
+        rank=rank,
+        local_rank=local_rank,
+        distributed_init_method="env://",
+    )
+    ps.initialize_model_parallel(tensor_model_parallel_size=world_size)
+    yield ps.get_tp_group()
+    ps.destroy_model_parallel()
+    ps.destroy_distributed_environment()
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("use_graph", [False, True])
+def test_batch_and_destination_invariance(group, dtype, use_graph):
+    torch.manual_seed(42 + group.rank_in_group)
+    # Non-integer contributions expose changes in floating-point sum order.
+    token = torch.randn(2048, device="cuda", dtype=dtype)
+    reference = None
+    for local_tokens in (1, 2, 3, 8, 32, 128):
+        input_ = token.repeat(local_tokens * group.world_size, 1)
+        original = input_.clone()
+        output = torch.empty((local_tokens, 2048), device="cuda", dtype=dtype)
+        if use_graph:
+            with group.graph_capture() as capture:
+                for _ in range(3):
+                    group.reduce_scatter_tensor(output, input_)
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph, stream=capture.stream):
+                    group.reduce_scatter_tensor(output, input_)
+            graph.replay()
+        else:
+            group.reduce_scatter_tensor(output, input_)
+
+        # Compare all destinations to catch rank-dependent reduction order.
+        all_outputs = torch.empty_like(input_)
+        dist.all_gather_into_tensor(all_outputs, output, group=group.device_group)
+        if reference is None:
+            reference = all_outputs[0].clone()
+        torch.testing.assert_close(
+            all_outputs, reference.expand_as(all_outputs), atol=0, rtol=0
+        )
+        torch.testing.assert_close(input_, original, atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    multigpu_pytest_main(__name__, __file__, num_gpus=(4,))

test/registered/kernels/ops/communication/test_symm_mem_all_gather.py

@@ -7,12 +7,12 @@ NCCL all-gather for a sweep of token counts, hidden widths, and the
 Usage::
 
     # Run on the default world sizes (2, 4, 8 GPUs):
-    python test/registered/jit/test_symm_mem_all_gather.py
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py
     # Pick a specific world size (or comma-separated list):
-    python test/registered/jit/test_symm_mem_all_gather.py --num-gpu 4
-    python test/registered/jit/test_symm_mem_all_gather.py --num-gpu 2,4,8
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py --num-gpu 4
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py --num-gpu 2,4,8
     # Extra pytest args (forwarded to each torchrun worker):
-    python test/registered/jit/test_symm_mem_all_gather.py -k 16384
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py -k 16384
 """
 
 from __future__ import annotations

test/registered/kernels/ops/cuda_graph/test_cuda_graph_dedup.py

@@ -0,0 +1,76 @@
+"""CUDA graph executable reuse across capture sizes and rejected updates."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.model_executor.runner_backend import cuda_graph_dedup_mixin as dedup
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestCudaGraphDedup(CustomTestCase):
+    @staticmethod
+    def capture(size, *, extra_node=False):
+        inputs = torch.zeros(size, device="cuda")
+        outputs = torch.empty_like(inputs)
+        graph = torch.cuda.CUDAGraph(keep_graph=True)
+        with torch.cuda.graph(graph):
+            torch.add(inputs, 1, out=outputs)
+            if extra_node:
+                outputs.mul_(2)
+        return graph, inputs, outputs
+
+    def test_grid_sizes_share_one_executable(self):
+        registry = dedup.DedupedCudaGraphRegistry()
+        self.addCleanup(registry.close)
+        captures = [self.capture(size) for size in (4096, 8192)]
+        self.assertEqual(
+            dedup.graph_signature(captures[0][0].raw_cuda_graph()),
+            dedup.graph_signature(captures[1][0].raw_cuda_graph()),
+        )
+        with patch.object(
+            registry, "instantiate", wraps=registry.instantiate
+        ) as instantiate:
+            graphs = [registry.register(capture[0]) for capture in captures]
+        self.assertEqual(instantiate.call_count, 1)
+        self.assertEqual(registry.stats(), (2, 1))
+        self.assertEqual(graphs[0].group.current_raw_graph, graphs[1].raw_graph)
+        registry.seal()
+        registry.seal()
+        for value in (3, 7):
+            for graph, (_, inputs, outputs) in zip(graphs, captures):
+                inputs.fill_(value)
+                graph.replay()
+                torch.cuda.synchronize()
+                self.assertTrue(
+                    torch.equal(outputs, torch.full_like(outputs, value + 1))
+                )
+
+    def test_rejected_registration_preserves_live_executable(self):
+        registry = dedup.DedupedCudaGraphRegistry()
+        self.addCleanup(registry.close)
+        original, inputs, outputs = self.capture(4096)
+        incompatible, _, _ = self.capture(4096, extra_node=True)
+        self.addCleanup(incompatible.reset)
+        graph = registry.register(original)
+        signature = dedup.graph_signature(graph.raw_graph)
+        with (
+            patch.object(dedup, "graph_signature", return_value=signature),
+            self.assertRaisesRegex(AssertionError, "register update failed"),
+        ):
+            registry.register(incompatible)
+        self.assertEqual(registry.stats(), (1, 1))
+        self.assertEqual(graph.group.current_raw_graph, graph.raw_graph)
+        inputs.fill_(11)
+        graph.replay()
+        torch.cuda.synchronize()
+        self.assertTrue(torch.equal(outputs, torch.full_like(outputs, 12)))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/diffusion/test_flux2_gated_resnorm.py

@@ -0,0 +1,68 @@
+"""Blackwell coverage for the FLUX.2 gated residual normalization fast path."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+import torch.nn as nn
+
+import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
+from sglang.kernels.ops.diffusion import residual_gate_add
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+def _is_blackwell() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10
+
+
+@unittest.skipUnless(_is_blackwell(), "FLUX.2 gated residual fusion requires SM100+")
+class TestFlux2GatedResnorm(CustomTestCase):
+    def test_gated_residual_norm_modulate_is_bit_exact(self):
+        torch.manual_seed(0)
+        hidden = 6144
+        shape = (1, 64, hidden)
+        residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+        update = torch.randn_like(residual)
+        params = torch.randn(1, 1, 3 * hidden, device="cuda").bfloat16()
+        gate, scale, shift = params.chunk(3, dim=-1)
+        norm = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6, device="cuda")
+
+        expected_residual = residual_gate_add(residual, update, gate)
+        expected = norm(expected_residual) * (1 + scale) + shift
+        actual, actual_residual = flux2._flux2_gated_resnorm(
+            norm, residual, update, gate, scale, shift
+        )
+
+        self.assertIsInstance(
+            flux2._defer_gated_residual(residual, update, gate), tuple
+        )
+        self.assertTrue(torch.equal(actual_residual, expected_residual))
+        self.assertTrue(torch.equal(actual, expected))
+
+    def test_gated_residual_defer_rejects_unsupported_inputs(self):
+        torch.manual_seed(0)
+        for batch, dtype in ((1, torch.float16), (2, torch.bfloat16)):
+            residual = torch.randn(batch, 17, 6144, device="cuda", dtype=dtype)
+            update = torch.randn_like(residual)
+            gate = torch.randn(batch, 1, 6144, device="cuda", dtype=dtype)
+            expected = residual_gate_add(residual, update, gate)
+            actual = flux2._defer_gated_residual(residual, update, gate)
+
+            self.assertIsInstance(actual, torch.Tensor)
+            self.assertTrue(torch.equal(actual, expected))
+
+        residual = torch.randn(1, 17, 6144, device="cuda", dtype=torch.bfloat16)
+        update = torch.randn_like(residual)
+        gate = torch.randn(1, 1, 6144, device="cuda", dtype=torch.bfloat16)
+        with patch("torch.compiler.is_compiling", return_value=True):
+            actual = flux2._defer_gated_residual(residual, update, gate)
+
+        self.assertIsInstance(actual, torch.Tensor)
+        self.assertTrue(torch.equal(actual, residual_gate_add(residual, update, gate)))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/diffusion/test_longcat_image_norm_modulate.py

@@ -0,0 +1,131 @@
+"""LongCat normalization parity and graph-safe fusion dispatch."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+from diffusers.models.normalization import AdaLayerNormZero, AdaLayerNormZeroSingle
+
+import sglang.multimodal_gen.runtime.models.dits.longcat_image as longcat
+from sglang.kernels.ops.diffusion import BitExactFusionGate, modulate_scale_shift
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestLongCatNormModulation(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        self.original_gate = longcat._LONGCAT_LN_MOD
+        longcat._LONGCAT_LN_MOD = BitExactFusionGate("test", per_signature=True)
+        torch.manual_seed(42)
+
+    def tearDown(self):
+        longcat._LONGCAT_LN_MOD = self.original_gate
+        super().tearDown()
+
+    def require_cuda(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+
+    @torch.inference_mode()
+    def test_adaln_checkpoint_and_output_parity(self):
+        for device, dtype, dim, seq in [
+            ("cpu", torch.float32, 64, 17),
+            ("cuda", torch.bfloat16, 3072, 512),
+            ("cuda", torch.bfloat16, 3072, 4608),
+        ]:
+            if device == "cuda" and not torch.cuda.is_available():
+                continue
+            for reference_cls, candidate_cls in [
+                (AdaLayerNormZero, longcat._LongCatAdaLayerNormZero),
+                (AdaLayerNormZeroSingle, longcat._LongCatAdaLayerNormZeroSingle),
+            ]:
+                with self.subTest(device=device, seq=seq, cls=reference_cls.__name__):
+                    reference = reference_cls(dim).to(device=device, dtype=dtype)
+                    candidate = candidate_cls(dim).to(device=device, dtype=dtype)
+                    candidate.load_state_dict(reference.state_dict(), strict=True)
+                    x = torch.randn(1, seq, dim, device=device, dtype=dtype)
+                    emb = torch.randn(1, dim, device=device, dtype=dtype)
+                    expected, actual = reference(x, emb=emb), candidate(x, emb=emb)
+                    for a, b in zip(expected, actual, strict=True):
+                        self.assertTrue(torch.equal(a, b))
+                    if device == "cuda":
+                        self.assertTrue(longcat._LONGCAT_LN_MOD.verified)
+                        self.assertFalse(longcat._LONGCAT_LN_MOD.disabled)
+
+    def inputs(self, seq=4096):
+        self.require_cuda()
+        x = torch.randn(1, seq, 3072, device="cuda", dtype=torch.bfloat16)
+        modulation = torch.randn(1, 6 * 3072, device="cuda", dtype=torch.bfloat16)
+        shift, scale, *_ = modulation.chunk(6, dim=-1)
+        norm = torch.nn.LayerNorm(3072, elementwise_affine=False, eps=1e-6).cuda()
+        return norm, x, scale, shift
+
+    def test_grad_enabled_uses_differentiable_reference(self):
+        norm, x, scale, shift = self.inputs(seq=17)
+        with torch.inference_mode():
+            longcat._longcat_norm_modulate(norm, x, scale, shift)
+        self.assertTrue(longcat._LONGCAT_LN_MOD.verified)
+        leaves = [t.detach().clone().requires_grad_() for t in (x, scale, shift)]
+        refs = [t.detach().clone().requires_grad_() for t in leaves]
+        with patch.object(longcat.diffusion_ops, "fused_layernorm_modulate") as fused:
+            actual = longcat._longcat_norm_modulate(norm, *leaves)
+            actual.float().sum().backward()
+            fused.assert_not_called()
+        expected = norm(refs[0]) * (1 + refs[1][:, None]) + refs[2][:, None]
+        expected.float().sum().backward()
+        self.assertTrue(torch.equal(actual, expected))
+        for a, b in zip(leaves, refs, strict=True):
+            self.assertIsNotNone(a.grad)
+            self.assertTrue(torch.equal(a.grad, b.grad))
+
+    @torch.inference_mode()
+    def test_changed_inputs_are_used_by_graph_replay(self):
+        norm, x, scale, shift = self.inputs()
+        longcat._longcat_norm_modulate(norm, x, scale, shift)
+        self.assertTrue(longcat._LONGCAT_LN_MOD.verified)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
+        x.add_(0.25)
+        scale.neg_()
+        shift.mul_(0.5)
+        graph.replay()
+        expected = norm(x) * (1 + scale[:, None]) + shift[:, None]
+        self.assertTrue(torch.equal(actual, expected))
+
+    @torch.inference_mode()
+    def test_unverified_capture_uses_eager_reference(self):
+        norm, x, scale, shift = self.inputs(seq=17)
+        expected = modulate_scale_shift(norm(x), scale, shift)
+        with patch.object(longcat.diffusion_ops, "fused_layernorm_modulate") as fused:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
+            graph.replay()
+            fused.assert_not_called()
+        self.assertFalse(longcat._LONGCAT_LN_MOD.verified)
+        self.assertTrue(torch.equal(actual, expected))
+
+    @torch.inference_mode()
+    def test_mismatch_disables_fusion_and_returns_reference(self):
+        norm, x, scale, shift = self.inputs(seq=17)
+        expected = norm(x) * (1 + scale[:, None]) + shift[:, None]
+        with patch.object(
+            longcat.diffusion_ops,
+            "fused_layernorm_modulate",
+            return_value=torch.zeros_like(x),
+        ):
+            actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
+        self.assertTrue(longcat._LONGCAT_LN_MOD.disabled)
+        self.assertTrue(torch.equal(actual, expected))
+        with patch.object(longcat.diffusion_ops, "fused_layernorm_modulate") as fused:
+            actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
+            fused.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/diffusion/test_mxfp8_swizzled.py

@@ -0,0 +1,78 @@
+"""MXFP8 producers against ``flashinfer.mxfp8_quantize`` of the bf16 tensor the
+unfused kernel stores: payload and the swizzled E8M0 scale buffer, padding
+included, byte for byte."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_mxfp8_swizzled,
+    can_use_silu_mul_mxfp8,
+    indexed_scale_shift_bf16_,
+    indexed_scale_shift_mxfp8_,
+    mxfp8_quantize_swizzled,
+    silu_mul_mxfp8,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+ROWS, HIDDEN = 333, 5376  # rows and scale columns both need padding
+
+
+def _assert_matches_flashinfer(
+    got: tuple[torch.Tensor, torch.Tensor], bf16: torch.Tensor
+) -> None:
+    import flashinfer
+
+    q, s = flashinfer.mxfp8_quantize(bf16.contiguous(), True)
+    assert torch.equal(got[0].view(torch.uint8), q.view(torch.uint8))
+    assert torch.equal(got[1].view(torch.uint8), s.view(torch.uint8))
+
+
+def test_quantize_swizzled_is_byte_exact() -> None:
+    g = torch.Generator(device="cuda").manual_seed(1)
+    x = (torch.randn((ROWS, HIDDEN), device="cuda", generator=g) * 3).to(torch.bfloat16)
+    x[0, :32] = 0  # an all-zero block takes the minimum exponent
+    assert can_use_mxfp8_swizzled(x)
+    _assert_matches_flashinfer(mxfp8_quantize_swizzled(x), x)
+
+
+def test_silu_mul_mxfp8_is_byte_exact() -> None:
+    g = torch.Generator(device="cuda").manual_seed(2)
+    x = (torch.randn((ROWS, 2 * HIDDEN), device="cuda", generator=g) * 2).to(
+        torch.bfloat16
+    )
+    assert can_use_silu_mul_mxfp8(x)
+    ref = torch.nn.functional.silu(x[:, :HIDDEN]) * x[:, HIDDEN:]
+    _assert_matches_flashinfer(silu_mul_mxfp8(x), ref)
+
+
+@pytest.mark.parametrize("keep_bf16", [True, False])
+def test_indexed_scale_shift_mxfp8_is_byte_exact(keep_bf16: bool) -> None:
+    g = torch.Generator(device="cuda").manual_seed(3)
+    x = torch.randn((ROWS, HIDDEN), device="cuda", generator=g).to(torch.bfloat16)
+    shift = torch.randn((3, HIDDEN), device="cuda", generator=g).to(torch.bfloat16)
+    scale = torch.randn((3, HIDDEN), device="cuda", generator=g).to(torch.bfloat16)
+    indices = torch.randint(0, 3, (ROWS,), device="cuda", generator=g)
+    ref = indexed_scale_shift_bf16_(x.clone(), shift, scale, indices)
+    kept, q, s = indexed_scale_shift_mxfp8_(
+        x, shift, scale, indices, keep_bf16=keep_bf16
+    )
+    _assert_matches_flashinfer((q, s), ref)
+    assert (kept is x and torch.equal(x, ref)) if keep_bf16 else kept is None
+
+
+def test_predicates_reject_unsupported_input() -> None:
+    assert not can_use_mxfp8_swizzled(
+        torch.randn(4, 64, device="cuda", dtype=torch.float16)
+    )
+    assert not can_use_silu_mul_mxfp8(
+        torch.randn(4, 96, device="cuda", dtype=torch.bfloat16)
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_qknorm_rope_out_of_place.py

@@ -0,0 +1,48 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_fused_inplace_qknorm_rope,
+    fused_inplace_qknorm_rope,
+    fused_qknorm_rope_out_of_place,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+def test_out_of_place_qknorm_rope_matches_inplace_and_keeps_inputs() -> None:
+    """The out-of-place variant (strided fused-qkv views in, contiguous copies
+    out) is bit-equal to the in-place kernel and leaves its inputs untouched;
+    VDN-H3's linear branch reads the raw q/k after it."""
+    T, H, D, R = 512, 4, 128, 96
+    if not can_use_fused_inplace_qknorm_rope(
+        D, R, True, torch.bfloat16, torch.bfloat16, True
+    ):
+        pytest.skip("fused qknorm+rope JIT kernel unavailable")
+    g = torch.Generator(device="cpu").manual_seed(0)
+    qkv = torch.randn(T, 3 * H * D, generator=g).to("cuda", torch.bfloat16)
+    q = qkv[:, : H * D].view(T, H, D)
+    k = qkv[:, H * D : 2 * H * D].view(T, H, D)
+    qw = (torch.rand(D, generator=g) + 0.5).to("cuda", torch.bfloat16)
+    kw = (torch.rand(D, generator=g) + 0.5).to("cuda", torch.bfloat16)
+    freqs = torch.randn(T, R // 2, generator=g).to("cuda")
+    cache = torch.cat((freqs.cos(), freqs.sin()), -1).to(torch.bfloat16).contiguous()
+    pos = torch.arange(T, device="cuda")
+    kwargs = dict(
+        is_neox=True, eps=1e-5, head_dim=D, rope_dim=R, round_norm_before_rope=True
+    )
+    q_ref, k_ref = q.clone(), k.clone()
+    fused_inplace_qknorm_rope(q_ref, k_ref, qw, kw, cache, pos, **kwargs)
+    q_out = torch.empty(T, H, D, device="cuda", dtype=torch.bfloat16)
+    k_out = torch.empty_like(q_out)
+    before = qkv.clone()
+    fused_qknorm_rope_out_of_place(q, k, q_out, k_out, qw, kw, cache, pos, **kwargs)
+    assert torch.equal(qkv, before)
+    assert torch.equal(q_out, q_ref) and torch.equal(k_out, k_ref)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_sana_wm_conv_post.py

@@ -0,0 +1,141 @@
+"""SANA-WM conv post-processing parity, rounding, and graph replay."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+import sglang.multimodal_gen.runtime.models.dits.sana_wm_components as wm
+from sglang.kernels.ops.diffusion import BitExactFusionGate
+from sglang.kernels.ops.diffusion.activation.sana_conv_post_triton import (
+    can_use_fused_bias_glu,
+    can_use_fused_bias_silu,
+    fused_bias_glu,
+    fused_bias_silu,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestSanaWMConvPost(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+        self.original_gate = wm._SANA_WM_CONV_POST
+        wm._SANA_WM_CONV_POST = BitExactFusionGate("test", per_signature=True)
+        self.addCleanup(setattr, wm, "_SANA_WM_CONV_POST", self.original_gate)
+        torch.manual_seed(42)
+
+    @torch.inference_mode()
+    def test_post_kernels_both_layouts_and_bias_modes(self):
+        for shape in [(3, 34, 17, 23), (2, 13440, 22, 40), (2, 34, 1, 1)]:
+            for layout in [torch.contiguous_format, torch.channels_last]:
+                x = torch.randn(shape, device="cuda", dtype=torch.bfloat16).to(
+                    memory_format=layout
+                )
+                bias = torch.randn(shape[1], device="cuda", dtype=x.dtype)
+                biased = x + bias[None, :, None, None]
+                self.assertTrue(torch.equal(fused_bias_silu(x, bias), F.silu(biased)))
+                for b, z in [(bias, biased), (None, x)]:
+                    a, g = z.chunk(2, dim=1)
+                    actual = fused_bias_glu(x, b)
+                    self.assertTrue(torch.equal(actual, a * F.silu(g)))
+                    self.assertTrue(actual.is_contiguous(memory_format=layout))
+        # Use a real spatial slice; width-one tensors remain contiguous.
+        sliced = torch.empty(2, 34, 7, 10, device="cuda", dtype=x.dtype)[:, :, :, ::2]
+        self.assertFalse(can_use_fused_bias_silu(sliced, bias))
+        self.assertFalse(can_use_fused_bias_glu(sliced, None))
+        self.assertFalse(can_use_fused_bias_glu(x.float(), None))
+
+    @staticmethod
+    def reference(module, x):
+        z = module.depth_conv(F.silu(module.inverted_conv.conv(x)))
+        a, g = z.chunk(2, dim=1)
+        return a * F.silu(g)
+
+    @torch.inference_mode()
+    def test_native_convolution_and_depthwise_bias_rounding(self):
+        for channels, hidden, shape in [
+            (32, 96, (3, 32, 17, 23)),
+            (2240, 6720, (14, 2240, 22, 40)),
+        ]:
+            module = wm.GLUMBConvTemp(channels, hidden).cuda().bfloat16()
+            for seed in (0, 1):
+                torch.manual_seed(seed)
+                x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+                self.assertTrue(
+                    torch.equal(module._spatial_glu(x), self.reference(module, x))
+                )
+                self.assertTrue(wm._SANA_WM_CONV_POST.verified)
+                self.assertFalse(wm._SANA_WM_CONV_POST.disabled)
+            # Guard the numerical distinction which forbids bias extraction.
+            z = module.inverted_conv(x)
+            conv = module.depth_conv.conv
+            split = F.conv2d(z, conv.weight, None, padding=1, groups=conv.groups)
+            split = split + conv.bias[None, :, None, None]
+            self.assertFalse(torch.equal(conv(z), split))
+
+    @torch.inference_mode()
+    def test_graph_replay_updates_inputs_and_streaming_tail(self):
+        module = wm.GLUMBConvTemp(32, 96).cuda().bfloat16()
+        # Exercise a nonzero temporal convolution instead of its zero init.
+        module.t_conv.weight.normal_(std=0.01)
+        x = torch.randn(2, 3 * 7 * 11, 32, device="cuda", dtype=torch.bfloat16)
+        tail = torch.randn(2, 32, 1, 77, device="cuda", dtype=x.dtype)
+        module(x, (3, 7, 11), ffn_tail=tail, save_ffn_tail=True)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual, actual_tail = module(
+                x, (3, 7, 11), ffn_tail=tail, save_ffn_tail=True
+            )
+        x.add_(0.25)
+        tail.mul_(0.5)
+        graph.replay()
+        with patch.object(wm._SANA_WM_CONV_POST, "disabled", True):
+            expected, expected_tail = module(
+                x, (3, 7, 11), ffn_tail=tail, save_ffn_tail=True
+            )
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertTrue(torch.equal(actual_tail, expected_tail))
+
+    @torch.inference_mode()
+    def test_unverified_capture_and_mismatch_use_reference(self):
+        module = wm.GLUMBConvTemp(32, 96).cuda().bfloat16()
+        x = torch.randn(3, 32, 17, 23, device="cuda", dtype=torch.bfloat16)
+        expected = self.reference(module, x)
+        with patch.object(wm.diffusion_ops, "fused_bias_silu") as fused:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = module._spatial_glu(x)
+            graph.replay()
+            fused.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+        with patch.object(
+            wm.diffusion_ops, "fused_bias_glu", return_value=torch.zeros_like(expected)
+        ):
+            actual = module._spatial_glu(x)
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertTrue(wm._SANA_WM_CONV_POST.disabled)
+        with patch.object(wm.diffusion_ops, "fused_bias_silu") as fused:
+            actual = module._spatial_glu(x)
+            fused.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+
+    def test_grad_enabled_uses_differentiable_reference(self):
+        module = wm.GLUMBConvTemp(32, 96).cuda().bfloat16()
+        x = torch.randn(
+            2, 32, 7, 11, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        with patch.object(wm.diffusion_ops, "fused_bias_silu") as fused:
+            actual = module._spatial_glu(x)
+            actual.float().sum().backward()
+            fused.assert_not_called()
+        self.assertIsNotNone(x.grad)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/diffusion/test_sana_wm_reverse_scan.py

@@ -0,0 +1,160 @@
+"""Exact streaming GDN output/state parity without materialized reversed video."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+import sglang.multimodal_gen.runtime.models.dits.sana_wm_components as wm
+from sglang.kernels.ops.diffusion import BitExactFusionGate
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestSanaWMReverseScan(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+        original = wm._SANA_WM_GDN_REVERSE
+        wm._SANA_WM_GDN_REVERSE = BitExactFusionGate("test", per_signature=True)
+        self.addCleanup(setattr, wm, "_SANA_WM_GDN_REVERSE", original)
+        torch.manual_seed(42)
+
+    def inputs(self, frames, transposed=False, frame_beta=False):
+        b, h, d, s = 1, 20, 112, 880
+        shape = (b, h, frames * s, d) if transposed else (b, h, d, frames * s)
+        xs = [torch.randn(shape, device="cuda") * 0.03 for _ in range(5)]
+        if transposed:
+            xs = [x.transpose(-1, -2) for x in xs]
+        beta_shape = (b, h, frames) if frame_beta else (b, h, frames, s)
+        return (
+            *xs,
+            torch.rand(beta_shape, device="cuda") * 0.02,
+            torch.rand(b, h, frames, device="cuda") * 0.5 + 0.4,
+        )
+
+    def assert_nested_equal(self, a, b):
+        if isinstance(a, tuple):
+            for x, y in zip(a, b, strict=True):
+                self.assert_nested_equal(x, y)
+        else:
+            self.assertTrue(torch.equal(a, b))
+
+    @torch.inference_mode()
+    def test_multiple_chunks_outputs_and_carried_states(self):
+        for transposed in (False, True):
+            for frame_beta in (False, True):
+                reference_state = candidate_state = (None, None)
+                reference_cam = candidate_cam = None
+                for frames in (1, 4, 3, 2):
+                    q, k, v, qr, kr, beta, decay = self.inputs(
+                        frames, transposed, frame_beta
+                    )
+                    with patch.object(wm._SANA_WM_GDN_REVERSE, "disabled", True):
+                        a, reference_state = wm._gdn_scan_cached(
+                            q,
+                            k,
+                            v,
+                            qr,
+                            kr,
+                            beta,
+                            decay,
+                            init_state_kv=reference_state[0],
+                            init_state_z=reference_state[1],
+                        )
+                        c, reference_cam = wm._single_path_delta_scan_cached(
+                            qr, kr, v, beta, decay, init_state_kv=reference_cam
+                        )
+                    b, candidate_state = wm._gdn_scan_cached(
+                        q,
+                        k,
+                        v,
+                        qr,
+                        kr,
+                        beta,
+                        decay,
+                        init_state_kv=candidate_state[0],
+                        init_state_z=candidate_state[1],
+                    )
+                    d, candidate_cam = wm._single_path_delta_scan_cached(
+                        qr, kr, v, beta, decay, init_state_kv=candidate_cam
+                    )
+                    self.assert_nested_equal(
+                        (a, reference_state, c, reference_cam),
+                        (b, candidate_state, d, candidate_cam),
+                    )
+                    self.assertFalse(wm._SANA_WM_GDN_REVERSE.disabled)
+        self.assertTrue(wm._SANA_WM_GDN_REVERSE.verified)
+
+    @torch.inference_mode()
+    def test_changed_graph_inputs_and_initial_states(self):
+        q, k, v, qr, kr, beta, decay = self.inputs(3, True)
+        state = torch.randn(1, 20, 112, 112, device="cuda") * 0.001
+        z = torch.randn(1, 20, 112, 1, device="cuda") * 0.001
+
+        def run():
+            return (
+                wm._gdn_scan_cached(
+                    q, k, v, qr, kr, beta, decay, init_state_kv=state, init_state_z=z
+                ),
+                wm._single_path_delta_scan_cached(
+                    qr, kr, v, beta, decay, init_state_kv=state
+                ),
+            )
+
+        run()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = run()
+        q.add_(0.01)
+        qr.mul_(0.7)
+        beta.mul_(0.8)
+        state.add_(0.001)
+        z.neg_()
+        graph.replay()
+        with patch.object(wm._SANA_WM_GDN_REVERSE, "disabled", True):
+            expected = run()
+        self.assert_nested_equal(actual, expected)
+
+    @torch.inference_mode()
+    def test_unverified_capture_and_mismatch_fallback(self):
+        _, _, v, qr, kr, beta, decay = self.inputs(2)
+
+        def reference():
+            return wm._single_path_delta_scan_backward_reference(qr, kr, v, beta, decay)
+
+        expected = reference()
+        with patch.object(wm, "_sana_wm_reverse_scan_impl") as fast:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = wm._sana_wm_reverse_scan(
+                    qr, kr, v, beta, decay, reference=reference
+                )
+            graph.replay()
+            fast.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+        with patch.object(
+            wm, "_sana_wm_reverse_scan_impl", return_value=torch.ones_like(expected)
+        ):
+            actual = wm._sana_wm_reverse_scan(
+                qr, kr, v, beta, decay, reference=reference
+            )
+        self.assertTrue(wm._SANA_WM_GDN_REVERSE.disabled)
+        self.assertTrue(torch.equal(actual, expected))
+
+    @torch.inference_mode()
+    def test_synthetic_zero_update_keeps_nonfinite_query_behavior(self):
+        q, k, v, qr, kr, beta, decay = self.inputs(1)
+        qr[..., 0] = float("nan")
+        q[..., 1] = float("inf")
+        expected = wm._gdn_scan_backward_reference(q, k, v, qr, kr, beta, decay, 1e-6)
+        actual = wm._sana_wm_reverse_scan_impl(qr, kr, v, beta, decay, q, k)
+        for a, b in zip(actual, expected, strict=True):
+            torch.testing.assert_close(a, b, rtol=0, atol=0, equal_nan=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/diffusion/test_vdn_delta_factors.py

@@ -0,0 +1,134 @@
+"""Fused VDN-H3 delta-rule factors against the eager Cholesky chain: both fp32 paths are
+held to the same cond(I + A)-dominated error band vs fp64, on model-shaped inputs."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import can_use_vdn_delta_factors, vdn_delta_factors
+from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as vdn
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+HEAD_DIM = 128
+
+
+def _inputs(
+    frames: int, heads: int, tokens: int, beta_scale: float = 1.0, seed: int = 0
+):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    k = torch.nn.functional.normalize(
+        torch.nn.functional.silu(
+            torch.randn(frames, heads, tokens, HEAD_DIM, device="cuda", generator=g)
+        ),
+        dim=-1,
+    )
+    v = torch.nn.functional.silu(
+        torch.randn(frames, heads, tokens, HEAD_DIM, device="cuda", generator=g)
+    )
+    beta = (
+        torch.sigmoid(torch.randn(frames, heads, tokens, device="cuda", generator=g))
+        * beta_scale
+    )
+    alpha = torch.rand(frames, heads, HEAD_DIM, device="cuda", generator=g) * 0.9 + 0.1
+    alpha[0] = 1.0
+    A = (k * beta.unsqueeze(-1)).transpose(-1, -2) @ k
+    A = 0.5 * (A + A.transpose(-1, -2))
+    B = (v * beta.unsqueeze(-1)).transpose(-1, -2) @ k
+    return A.float().contiguous(), B.float().contiguous(), alpha.float().contiguous()
+
+
+def _rel(x: torch.Tensor, ref: torch.Tensor) -> float:
+    return ((x.double() - ref).norm() / ref.norm()).item()
+
+
+def _fp64(A, B, alpha):
+    inv = torch.linalg.inv(
+        torch.eye(HEAD_DIM, device=A.device, dtype=torch.float64) + A.double()
+    )
+    return alpha.double().unsqueeze(-1) * inv, B.double() @ inv
+
+
+@pytest.mark.parametrize("frames,heads,tokens", [(1, 1, 16), (5, 3, 64), (11, 7, 1008)])
+@pytest.mark.parametrize("beta_scale", [1.0, 50.0])
+def test_matches_eager_and_fp64(frames, heads, tokens, beta_scale):
+    A, B, alpha = _inputs(frames, heads, tokens, beta_scale)
+    assert can_use_vdn_delta_factors(A, B, alpha)
+    t_ref, j_ref = _fp64(A, B, alpha)
+    t_eager, j_eager = vdn.delta_factor_apply(
+        "vdn_solve", alpha, A, B, tokens_per_frame=tokens
+    )
+    t_fused, j_fused = vdn_delta_factors(A, B, alpha)
+    assert t_fused.shape == A.shape and j_fused.shape == B.shape
+    assert t_fused.dtype is torch.float32 and j_fused.dtype is torch.float32
+    assert torch.isfinite(t_fused).all() and torch.isfinite(j_fused).all()
+    # same accuracy class as the Cholesky chain (both fp32, cond-dominated)
+    assert _rel(t_fused, t_ref) <= 1.5 * _rel(t_eager, t_ref) + 1e-7
+    assert _rel(j_fused, j_ref) <= 1.5 * _rel(j_eager, j_ref) + 1e-7
+    assert _rel(t_fused, t_ref) < 1e-5 and _rel(j_fused, j_ref) < 3e-5
+    # elementwise the two fp32 paths differ by a few 1e-5 on ill-conditioned inputs (cancellation)
+    torch.testing.assert_close(t_fused, t_eager, rtol=1e-4, atol=1e-5)
+    torch.testing.assert_close(j_fused, j_eager, rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("rule", ["vdn_solve", "vdn_scaled"])
+def test_delta_factor_apply_fused_path(rule):
+    A, B, alpha = _inputs(4, 2, 48)
+    eager = vdn.delta_factor_apply(rule, alpha, A, B, tokens_per_frame=48, fused=False)
+    fused = vdn.delta_factor_apply(rule, alpha, A, B, tokens_per_frame=48, fused=True)
+    for x, y in zip(fused, eager):
+        torch.testing.assert_close(x, y, rtol=2e-5, atol=2e-5)
+
+
+def test_sana_scaled_ignores_fused():
+    A, B, alpha = _inputs(2, 2, 32)
+    eager = vdn.delta_factor_apply(
+        "sana_scaled", alpha, A, B, tokens_per_frame=32, fused=False
+    )
+    fused = vdn.delta_factor_apply(
+        "sana_scaled", alpha, A, B, tokens_per_frame=32, fused=True
+    )
+    for x, y in zip(fused, eager):
+        assert torch.equal(x, y)
+
+
+def _storage_offset_copy(t: torch.Tensor) -> torch.Tensor:
+    # contiguous, but one element past a 16-byte boundary
+    flat = torch.empty(t.numel() + 1, dtype=t.dtype, device=t.device)
+    out = flat[1:].view(t.shape)
+    out.copy_(t)
+    assert out.is_contiguous() and out.data_ptr() % 16 != 0
+    return out
+
+
+@pytest.mark.parametrize("which", ["A", "B", "alpha"])
+def test_storage_offset_input_matches_aligned(which):
+    """A contiguous input with a storage offset must not fault in the float4 loads."""
+    A, B, alpha = _inputs(3, 2, 64)
+    ref = vdn_delta_factors(A, B, alpha)
+    inputs = {"A": A, "B": B, "alpha": alpha}
+    inputs[which] = _storage_offset_copy(inputs[which])
+    assert can_use_vdn_delta_factors(inputs["A"], inputs["B"], inputs["alpha"])
+    out = vdn_delta_factors(inputs["A"], inputs["B"], inputs["alpha"])
+    for got, want in zip(out, ref):
+        assert torch.equal(got, want)
+
+
+def test_can_use_rejects_unsupported():
+    A, B, alpha = _inputs(2, 2, 32)
+    assert can_use_vdn_delta_factors(A, B, alpha)
+    assert not can_use_vdn_delta_factors(
+        A[..., :64, :64].contiguous(),
+        B[..., :64, :64].contiguous(),
+        alpha[..., :64].contiguous(),
+    )
+    assert not can_use_vdn_delta_factors(A.bfloat16(), B, alpha)
+    assert not can_use_vdn_delta_factors(A.transpose(-1, -2), B, alpha)
+    assert not can_use_vdn_delta_factors(A, B, alpha[..., :1].expand_as(alpha))
+    assert not can_use_vdn_delta_factors(A.cpu(), B.cpu(), alpha.cpu())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_vdn_linear_branch.py

@@ -0,0 +1,139 @@
+"""VDN-H3 linear-branch kernels against the eager chains they replace: the data
+movers bit-exact, the three activation kernels within one bf16 ulp."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_vdn_frame_stats_prep,
+    can_use_vdn_gather_linear_state,
+    can_use_vdn_linear_epilogue,
+    can_use_vdn_silu_l2norm,
+    can_use_vdn_temporal_conv_act,
+    vdn_frame_stats_prep,
+    vdn_linear_epilogue,
+    vdn_silu_l2norm,
+    vdn_temporal_conv_act,
+)
+from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import (
+    VDNHybridAttentionArchConfig,
+)
+from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as vdn
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+FRAMES, TOKENS, HEADS, HEAD_DIM = 6, 24, 3, 32
+
+
+def _ulp_close(got: torch.Tensor, ref: torch.Tensor) -> bool:
+    scale = max(1.0, ref.float().abs().max().item())
+    return (got.float() - ref.float()).abs().max().item() <= 2e-2 * scale
+
+
+def test_temporal_conv_act_matches_eager_chain() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    x = torch.randn(FRAMES, TOKENS, HEADS * HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    w = (torch.randn(HEADS * HEAD_DIM, 5, generator=g) * 0.4).to("cuda", torch.bfloat16)
+    assert can_use_vdn_temporal_conv_act(x, HEADS, HEAD_DIM)
+    ref = vdn._activate(vdn._temporal_shift(x, w).reshape(-1, HEADS, HEAD_DIM), True)
+    assert _ulp_close(vdn_temporal_conv_act(x, w, HEADS, HEAD_DIM, True), ref)
+    frame_major = vdn_temporal_conv_act(x, w, HEADS, HEAD_DIM, True, frame_major=True)
+    assert (
+        frame_major.shape == (FRAMES, HEADS, TOKENS, HEAD_DIM)
+        and frame_major.is_contiguous()
+    )
+    assert torch.equal(
+        frame_major, ref.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    ) or _ulp_close(
+        frame_major, ref.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    )
+
+
+def test_silu_l2norm_reads_strided_qkv_views() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    tokens = torch.randn(FRAMES * TOKENS, 3 * HEADS * HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    strided = tokens[:, : HEADS * HEAD_DIM].view(FRAMES * TOKENS, HEADS, HEAD_DIM)
+    assert can_use_vdn_silu_l2norm(strided)
+    got = vdn_silu_l2norm(strided, True)
+    assert got.is_contiguous() and _ulp_close(got, vdn._activate(strided, True))
+    got_v = vdn_silu_l2norm(strided, False)
+    assert _ulp_close(got_v, torch.nn.functional.silu(strided))
+    frame_major = vdn_silu_l2norm(strided, True, per_frame=TOKENS)
+    assert frame_major.shape == (FRAMES, HEADS, TOKENS, HEAD_DIM)
+    assert torch.equal(
+        got.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3), frame_major
+    )
+    with pytest.raises(ValueError):
+        vdn_silu_l2norm(strided, True, per_frame=TOKENS + 1)
+
+
+def test_frame_stats_prep_is_bit_exact() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    key = torch.randn(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    value = torch.randn(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    beta = torch.rand(FRAMES * TOKENS, HEADS, generator=g).to("cuda", torch.bfloat16)
+    assert can_use_vdn_frame_stats_prep(key, value)
+    k16, k32, kb32, vb = vdn_frame_stats_prep(key, value, beta, FRAMES, TOKENS)
+    kf = key.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    vf = value.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    bf = beta.view(FRAMES, TOKENS, HEADS).permute(0, 2, 1)
+    assert torch.equal(k16, kf.contiguous())
+    assert torch.equal(k32, kf.float().contiguous())
+    assert torch.equal(kb32, (kf.float() * bf.unsqueeze(-1).float()).contiguous())
+    assert torch.equal(vb, (vf * bf.unsqueeze(-1).to(vf.dtype)).contiguous())
+
+
+def test_linear_epilogue_matches_eager_chain() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    readout = torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    weight = (1 + 0.1 * torch.randn(HEAD_DIM, generator=g)).to("cuda", torch.bfloat16)
+    gate = torch.rand(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    assert can_use_vdn_linear_epilogue(readout)
+    got = vdn_linear_epilogue(readout, weight, gate, 1e-6)
+    assert _ulp_close(got, vdn.linear_epilogue(readout, weight, gate, 1e-6))
+
+
+@pytest.mark.parametrize("bridge", ["alpha", "none"])
+@pytest.mark.parametrize("with_text_state", [False, True])
+def test_gather_linear_state_matches_eager(bridge: str, with_text_state: bool) -> None:
+    g = torch.Generator(device="cpu").manual_seed(5)
+    frames, heads, dim = 9, 2, 32
+    hybrid = VDNHybridAttentionArchConfig(chunk=3, radius=1, anchor_frames="none")
+    bounds = hybrid.window_bounds(frames)
+    prefix = torch.randn(frames, heads, dim, dim, generator=g).cuda()
+    suffix = torch.randn(frames, heads, dim, dim, generator=g).cuda()
+    alpha = (torch.rand(frames, heads, dim, generator=g) * 0.5 + 0.5).cuda()
+    text = torch.randn(heads, dim, dim, generator=g).cuda() if with_text_state else None
+    assert can_use_vdn_gather_linear_state(prefix)
+    kwargs = dict(bridge=bridge, text_state=text, out_dtype=torch.float32)
+    ref = vdn.gather_linear_state(prefix, suffix, alpha, bounds, fused=False, **kwargs)
+    got = vdn.gather_linear_state(prefix, suffix, alpha, bounds, **kwargs)
+    torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5)
+
+
+def test_predicates_reject_unsupported_inputs() -> None:
+    fp16 = torch.randn(4, 2, 32, device="cuda", dtype=torch.float16)
+    assert not can_use_vdn_silu_l2norm(fp16)
+    odd = torch.randn(4, 2, 48, device="cuda", dtype=torch.bfloat16)
+    assert not can_use_vdn_silu_l2norm(odd)
+    with pytest.raises(ValueError):
+        vdn_silu_l2norm(odd, True)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/disaggregation/test_kv_checksum.py

@@ -0,0 +1,345 @@
+"""Tests for KV checksum integration in PD disaggregation."""
+
+import unittest
+import zlib
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.disaggregation.checksum import (
+    KvChecksumComputer,
+    is_health_check_req,
+)
+from sglang.srt.disaggregation.decode import SchedulerDisaggregationDecodeMixin
+from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
+from sglang.srt.disaggregation.utils import MetadataBuffers
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _ref_strided_adler32(tensors, indices, strides) -> int:
+    parts = []
+    for tensor, idx, stride in zip(tensors, indices, strides):
+        raw = tensor.cpu().contiguous().flatten().view(torch.uint8)
+        for i in idx.cpu().tolist():
+            parts.append(raw[i * stride : (i + 1) * stride].numpy().tobytes())
+    return zlib.adler32(b"".join(parts))
+
+
+def _make_buf(size=4, *, kv_checksum_enabled=True, output_dsa_topk_indices_dim=0):
+    return MetadataBuffers(
+        size=size,
+        hidden_size=16,
+        hidden_states_dtype=torch.float32,
+        max_sampling_mask_tokens=128,
+        output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
+        kv_checksum_enabled=kv_checksum_enabled,
+    )
+
+
+class TestMetadataBuffers(unittest.TestCase):
+    def test_extends_aux_with_kv_checksum(self):
+        for sampling_mask in (False, True):
+            for seed_dim in (0, 32):
+                with (
+                    self.subTest(sampling_mask=sampling_mask, seed_dim=seed_dim),
+                    envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(sampling_mask),
+                ):
+                    buf = _make_buf(size=8, output_dsa_topk_indices_dim=seed_dim)
+                    disabled = _make_buf(
+                        size=8,
+                        kv_checksum_enabled=False,
+                        output_dsa_topk_indices_dim=seed_dim,
+                    )
+                    ptrs, data_lens, item_lens = buf.get_buf_infos()
+                    base_ptrs, base_data_lens, base_item_lens = disabled.get_buf_infos()
+                    self.assertEqual(len(ptrs), len(base_ptrs) + 1)
+                    self.assertEqual(data_lens[:-1], base_data_lens)
+                    self.assertEqual(item_lens[:-1], base_item_lens)
+                    self.assertEqual(ptrs[-1], buf.kv_checksum.data_ptr())
+                    self.assertEqual(data_lens[-1], buf.kv_checksum.nbytes)
+                    self.assertEqual(item_lens[-1], buf.kv_checksum[0].nbytes)
+                    self.assertEqual(item_lens[-1], item_lens[-2])
+                    self.assertEqual(len(buf.get_buf(2)), len(disabled.get_buf(2)))
+
+    def test_set_get_kv_checksum_roundtrip(self):
+        buf = _make_buf()
+        buf.set_kv_checksum(SimpleNamespace(metadata_buffer_index=1), 0xDEADBEEF)
+        self.assertEqual(buf.get_kv_checksum(1), 0xDEADBEEF)
+        self.assertEqual(buf.get_kv_checksum(0), 0)
+
+
+class TestKvChecksumComputerConfig(unittest.TestCase):
+    def test_flattens_nested_state_descriptor_components(self):
+        computer = KvChecksumComputer(
+            torch.device("cpu"),
+            kv_data_ptrs=[11, 22],
+            kv_item_lens=[33, 44],
+            state_data_ptrs=[[55, 66], [77]],
+            state_item_lens=[[88, 99], [111]],
+        )
+
+        self.assertEqual(computer._state_data_ptrs, [55, 66, 77])
+        self.assertEqual(computer._state_item_lens, [88, 99, 111])
+
+
+class TestKvChecksumHealthCheck(unittest.TestCase):
+    def test_detects_health_check_request(self):
+        self.assertTrue(is_health_check_req(SimpleNamespace(rid="HEALTH_CHECK_1")))
+        self.assertFalse(is_health_check_req(SimpleNamespace(rid="user_req")))
+        self.assertFalse(is_health_check_req(SimpleNamespace(rid=None)))
+
+
+def _make_kv(num_layers, num_pages, page_elems, dtype=torch.float16):
+    return [
+        torch.randn(num_pages, page_elems, dtype=dtype, device="cuda:0")
+        for _ in range(2 * num_layers)
+    ]
+
+
+def _make_computer(kv, item_len, state=None, state_item_lens=None):
+    return KvChecksumComputer(
+        torch.device("cuda:0"),
+        kv_data_ptrs=[t.data_ptr() for t in kv],
+        kv_item_lens=[item_len] * len(kv),
+        state_data_ptrs=[t.data_ptr() for t in (state or [])],
+        state_item_lens=state_item_lens or [],
+    )
+
+
+class TestKvChecksumComputer(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        self.device = torch.device("cuda:0")
+
+    def test_kv_only_matches_reference(self):
+        kv = _make_kv(num_layers=4, num_pages=32, page_elems=128, dtype=torch.bfloat16)
+        idx = torch.tensor([3, 7, 8, 15, 31], dtype=torch.int64, device=self.device)
+        item_len = 128 * 2
+        value = _make_computer(kv, item_len).compute(idx)
+        expected = _ref_strided_adler32(kv, [idx] * len(kv), [item_len] * len(kv))
+        self.assertEqual(value, expected)
+
+    def test_kv_corruption_detected(self):
+        kv = _make_kv(num_layers=2, num_pages=16, page_elems=64)
+        idx = torch.tensor([5], dtype=torch.int64, device=self.device)
+        computer = _make_computer(kv, 64 * 2)
+        v1 = computer.compute(idx)
+        kv[0][5, 0] += 1
+        self.assertNotEqual(v1, computer.compute(idx))
+
+    def test_kv_plus_state_matches_reference(self):
+        kv = _make_kv(num_layers=2, num_pages=8, page_elems=32)
+        state = [
+            torch.randn(4, 16, dtype=torch.float16, device=self.device)
+            for _ in range(2)
+        ]
+        kv_idx = torch.tensor([0, 1, 2], dtype=torch.int64, device=self.device)
+        state_idx = torch.tensor([1], dtype=torch.int64, device=self.device)
+        kv_len, state_lens = 32 * 2, [16 * 2, 16 * 2]
+        computer = _make_computer(kv, kv_len, state, state_lens)
+        value = computer.compute(kv_idx, state_idx)
+        expected = _ref_strided_adler32(
+            kv + state,
+            [kv_idx] * len(kv) + [state_idx] * len(state),
+            [kv_len] * len(kv) + state_lens,
+        )
+        self.assertEqual(value, expected)
+        state[0][1, 0] += 1
+        self.assertNotEqual(value, computer.compute(kv_idx, state_idx))
+
+    def test_nested_state_components_match_reference(self):
+        kv = _make_kv(num_layers=1, num_pages=8, page_elems=32)
+        state_components = [
+            [torch.randn(4, 16, dtype=torch.float16, device=self.device)],
+            [
+                torch.randn(4, 8, dtype=torch.float16, device=self.device),
+                torch.randn(4, 12, dtype=torch.float16, device=self.device),
+            ],
+        ]
+        kv_idx = torch.tensor([0, 3, 7], dtype=torch.int64, device=self.device)
+        state_idx = torch.tensor([1, 2], dtype=torch.int64, device=self.device)
+        state_tensors = [tensor for comp in state_components for tensor in comp]
+        kv_len = 32 * 2
+        state_lens = [[16 * 2], [8 * 2, 12 * 2]]
+        computer = KvChecksumComputer(
+            self.device,
+            kv_data_ptrs=[t.data_ptr() for t in kv],
+            kv_item_lens=[kv_len] * len(kv),
+            state_data_ptrs=[
+                [tensor.data_ptr() for tensor in comp] for comp in state_components
+            ],
+            state_item_lens=state_lens,
+        )
+
+        value = computer.compute(kv_idx, state_idx)
+        expected = _ref_strided_adler32(
+            kv + state_tensors,
+            [kv_idx] * len(kv) + [state_idx] * len(state_tensors),
+            [kv_len] * len(kv) + [item for comp in state_lens for item in comp],
+        )
+        self.assertEqual(value, expected)
+
+
+class _FakeScheduler(SchedulerDisaggregationDecodeMixin):
+    def __init__(self, computer, req_to_token):
+        self.kv_checksum_computer = computer
+        self.waiting_queue = []
+        self.token_to_kv_pool_allocator = SimpleNamespace(
+            page_size=1, get_kvcache=lambda: SimpleNamespace()
+        )
+        self.req_to_token_pool = SimpleNamespace(req_to_token=req_to_token)
+        self.tree_cache = None
+        self.output_streamer = SimpleNamespace(stream_output=self.stream_output)
+        self.metrics_reporter = SimpleNamespace(enable_metrics=True)
+        self.metrics_collector = Mock()
+        self.streamed_aborts = []
+
+    def stream_output(self, reqs, return_logprob):
+        self.streamed_aborts.extend(reqs)
+
+
+class _FakePrefillScheduler(SchedulerDisaggregationPrefillMixin):
+    def __init__(self):
+        self.kv_checksum_computer = object()
+        self.disagg_metadata_buffers = SimpleNamespace(set_kv_checksum=Mock())
+
+
+def _make_req(expected_chksum, num_input_tokens, rid="r0"):
+    return SimpleNamespace(
+        rid=rid,
+        bootstrap_room=12345,
+        kv=SimpleNamespace(req_pool_idx=0),
+        origin_input_ids=list(range(num_input_tokens)),
+        fill_ids=list(range(num_input_tokens)),
+        expected_kv_checksum=expected_chksum,
+        return_logprob=False,
+    )
+
+
+class TestPrefillHealthCheckChecksum(unittest.TestCase):
+    def test_health_check_clears_metadata_checksum(self):
+        sched = _FakePrefillScheduler()
+        req = _make_req(0xDEADBEEF, 1, rid="HEALTH_CHECK_1")
+        with patch.object(
+            SchedulerDisaggregationPrefillMixin,
+            "_send_kv_chunk",
+            lambda *args, **kwargs: None,
+        ):
+            sched.send_kv_chunk(req, last_chunk=True)
+        sched.disagg_metadata_buffers.set_kv_checksum.assert_called_once_with(req, 0)
+
+
+class TestGetNewPrebuiltBatchChecksum(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        self.device = torch.device("cuda:0")
+        self.num_pages = 8
+        self.kv = _make_kv(num_layers=2, num_pages=self.num_pages, page_elems=32)
+        self.item_len = 32 * 2
+        self.req_to_token = torch.arange(
+            self.num_pages, dtype=torch.int64, device=self.device
+        ).view(1, self.num_pages)
+        self.true_chksum = _ref_strided_adler32(
+            self.kv,
+            [torch.arange(self.num_pages, dtype=torch.int64, device=self.device)]
+            * len(self.kv),
+            [self.item_len] * len(self.kv),
+        )
+
+    def _make_sched(self, computer=None):
+        if computer is _SENTINEL:
+            computer = _make_computer(self.kv, self.item_len)
+        return _FakeScheduler(computer, self.req_to_token)
+
+    def _run_once(self, sched, batch_ret=None):
+        running_batch = SimpleNamespace()
+        with patch.object(
+            SchedulerDisaggregationDecodeMixin,
+            "_get_new_prebuilt_batch",
+            lambda s, rb: batch_ret,
+        ):
+            return sched.get_new_prebuilt_batch(running_batch)
+
+    def test_match_keeps_req(self):
+        sched = self._make_sched(_SENTINEL)
+        sched.waiting_queue = [_make_req(self.true_chksum, self.num_pages)]
+        self._run_once(sched)
+        self.assertEqual(len(sched.waiting_queue), 1)
+        self.assertEqual(sched.streamed_aborts, [])
+
+    def test_mismatch_aborts(self):
+        sched = self._make_sched(_SENTINEL)
+        req = _make_req(0xDEADBEEF, self.num_pages)
+        sched.waiting_queue = [req]
+        with (
+            envs.SGLANG_IS_IN_CI.override(False),
+            patch("sglang.srt.disaggregation.decode.prepare_abort") as mock_abort,
+            patch("sglang.srt.disaggregation.decode.release_kv_cache") as mock_release,
+        ):
+            self._run_once(sched)
+            self._run_once(sched)
+        self.assertEqual(sched.waiting_queue, [])
+        self.assertEqual(sched.streamed_aborts, [req])
+        mock_abort.assert_called_once()
+        mock_release.assert_called_once()
+        sched.metrics_collector.increment_transfer_failed_reqs.assert_called_once_with()
+
+    def test_mismatch_raises_in_ci(self):
+        sched = self._make_sched(_SENTINEL)
+        sched.waiting_queue = [_make_req(0xDEADBEEF, self.num_pages)]
+        with (
+            envs.SGLANG_IS_IN_CI.override(True),
+            patch("sglang.srt.disaggregation.decode.prepare_abort") as mock_abort,
+            self.assertRaisesRegex(RuntimeError, "KV checksum mismatch"),
+        ):
+            self._run_once(sched)
+        mock_abort.assert_not_called()
+        sched.metrics_collector.increment_transfer_failed_reqs.assert_not_called()
+
+    def test_health_check_skips_checksum(self):
+        sched = self._make_sched(_SENTINEL)
+        req = _make_req(0xDEADBEEF, self.num_pages, rid="HEALTH_CHECK_1")
+        sched.waiting_queue = [req]
+        with (
+            patch("sglang.srt.disaggregation.decode.prepare_abort") as mock_abort,
+            patch("sglang.srt.disaggregation.decode.release_kv_cache") as mock_release,
+        ):
+            self._run_once(sched)
+        self.assertEqual(sched.waiting_queue, [req])
+        self.assertEqual(sched.streamed_aborts, [])
+        mock_abort.assert_not_called()
+        mock_release.assert_not_called()
+
+    def test_retract_re_verifies(self):
+        sched = self._make_sched(_SENTINEL)
+        req = _make_req(self.true_chksum, self.num_pages)
+        sched.waiting_queue = [req]
+        self._run_once(sched)
+        self._run_once(sched)
+        self.assertEqual(sched.waiting_queue, [req])
+        self.assertEqual(sched.streamed_aborts, [])
+
+    def test_disabled_delegates_to_batch_builder(self):
+        sched = self._make_sched(computer=None)
+        sched.waiting_queue = [_make_req(0xABCD, 4)]
+        sentinel = object()
+        self.assertIs(self._run_once(sched, batch_ret=sentinel), sentinel)
+
+    def test_zero_expected_skips_checksum(self):
+        sched = self._make_sched(_SENTINEL)
+        sched.waiting_queue = [_make_req(0, self.num_pages)]
+        self._run_once(sched)
+        self.assertEqual(len(sched.waiting_queue), 1)
+
+
+_SENTINEL = object()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/elementwise/test_fast_topk.py

@@ -0,0 +1,141 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.elementwise.fast_topk import fast_topk
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _check_topk_values(score, lengths, indices, topk, row_starts):
+    """fast_topk leaves order and tie-breaking unspecified,
+    so compare the sorted top-k values rather than index sets."""
+    for b in range(score.shape[0]):
+        start = int(row_starts[b]) if row_starts is not None else 0
+        length = int(lengths[b])
+        section = score[b, start : start + length]
+        row = indices[b]
+        if length <= topk:
+            # naive path: identity indices, then -1 fill
+            assert torch.equal(
+                row[:length].cpu(), torch.arange(length, dtype=torch.int32)
+            )
+            assert (row[length:] == -1).all()
+            continue
+        assert (row >= 0).all(), "long rows must fill every slot"
+        picked = section[row.long()]
+        expected = torch.topk(section, topk).values
+        assert torch.equal(
+            picked.sort(descending=True).values, expected.sort(descending=True).values
+        ), f"row {b}: top-{topk} value multiset mismatch"
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+@pytest.mark.parametrize(
+    "batch,length",
+    [
+        (1, 4096),
+        (7, 3000),
+        (33, 32768),
+        (128, 2050),
+    ],
+)
+def test_fast_topk_long_rows(topk, batch, length):
+    torch.manual_seed(0)
+    score = torch.randn(batch, length, dtype=torch.float32, device="cuda")
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_short_and_mixed_rows(topk):
+    torch.manual_seed(0)
+    max_len = topk + 128
+    batch = 8
+    score = torch.randn(batch, max_len, dtype=torch.float32, device="cuda")
+    # rows shorter than k (naive path), exactly k, and longer than k
+    lens = [1, topk // 3, topk - 1, topk, topk + 1, topk + 7, 17, max_len]
+    lengths = torch.tensor(lens[:batch], dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_ragged_with_row_starts(topk):
+    torch.manual_seed(0)
+    batch, width = 16, 8192
+    score = torch.randn(batch, width, dtype=torch.float32, device="cuda")
+    row_starts = torch.randint(0, 2048, (batch,), dtype=torch.int32, device="cuda")
+    lengths = torch.randint(1, 2048, (batch,), dtype=torch.int32, device="cuda")
+    lengths = torch.minimum(lengths, width - row_starts).to(torch.int32)
+    # ensure some rows are longer than k
+    lengths[0] = min(width - int(row_starts[0]), topk + 100)
+
+    indices = fast_topk(score, lengths, topk, row_starts=row_starts)
+    _check_topk_values(score, lengths, indices, topk, row_starts)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_row_stride(topk):
+    torch.manual_seed(0)
+    batch, length = 8, 4096
+    base = torch.randn(batch, 2 * length, dtype=torch.float32, device="cuda")
+    score = base[:, :length]  # stride(0) == 2*length, stride(1) == 1
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+@pytest.mark.parametrize(
+    "fill",
+    [
+        "binary",  # only 0s and 1s: extreme duplication at the threshold bin
+        "few_levels",  # a handful of distinct levels incl. negatives
+        "constant",  # whole rows of one value
+    ],
+)
+def test_fast_topk_duplicate_heavy(topk, fill):
+    torch.manual_seed(0)
+    batch, length = 16, 8192
+    if fill == "binary":
+        score = torch.randint(0, 2, (batch, length), dtype=torch.float32, device="cuda")
+    elif fill == "few_levels":
+        levels = torch.tensor([-5.0, -1.0, 0.0, 0.5, 2.0], device="cuda")
+        score = levels[torch.randint(0, 5, (batch, length), device="cuda")]
+    else:
+        score = torch.full((batch, length), 3.25, dtype=torch.float32, device="cuda")
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+@pytest.mark.parametrize("topk", [512, 2048])
+def test_fast_topk_negative_and_zero(topk):
+    torch.manual_seed(0)
+    batch, length = 8, 16384
+    score = torch.randn(batch, length, dtype=torch.float32, device="cuda") * 100
+    score[:, : length // 3] = 0.0  # long zero prefix
+    score[:, length // 3 : length // 2] = -1e30  # very negative block
+    lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
+
+    indices = fast_topk(score, lengths, topk)
+    _check_topk_values(score, lengths, indices, topk, None)
+
+
+def test_fast_topk_unsupported_k():
+    score = torch.randn(2, 4096, dtype=torch.float32, device="cuda")
+    lengths = torch.full((2,), 4096, dtype=torch.int32, device="cuda")
+    with pytest.raises(RuntimeError, match="topk"):
+        fast_topk(score, lengths, 1024)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/elementwise/test_hc_combine.py

@@ -0,0 +1,161 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.kernels.ops.elementwise.hc_combine import hc_combine
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+HC_COUNT = 4
+HIDDEN_SIZE = 2560
+
+
+def _reference_hc_combine(
+    block_output: torch.Tensor,
+    residual: torch.Tensor,
+    normed_residual: torch.Tensor,
+    inject_weight: torch.Tensor,
+    hc: int,
+    hs: int,
+    compute_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Eager reference mirroring ``GatedResidual._combine_compute``;
+    ``compute_dtype=torch.float64`` is the near-exact reference."""
+    R = residual.to(compute_dtype).unflatten(-1, (hc, hs))
+    gates = 2 * torch.sigmoid(
+        F.linear(normed_residual.to(compute_dtype), inject_weight.to(compute_dtype))
+        / hc
+    )
+    injection = block_output.to(compute_dtype).unsqueeze(-2) * gates.unsqueeze(-1)
+    return (R + injection).flatten(-2)
+
+
+def _make_inputs(
+    num_tokens: int, dtype: torch.dtype, hc: int = HC_COUNT, hs: int = HIDDEN_SIZE
+):
+    torch.manual_seed(0)
+    block_output = torch.randn(num_tokens, hs, dtype=dtype, device="cuda")
+    residual = torch.randn(num_tokens, hc * hs, dtype=dtype, device="cuda")
+    normed_residual = torch.randn(num_tokens, hc * hs, dtype=dtype, device="cuda")
+    inject_weight = torch.randn(hc, hc * hs, dtype=dtype, device="cuda") * 0.02
+    return block_output, residual, normed_residual, inject_weight
+
+
+# Worst case over M in {1, 7, 128, 8192} against the fp64 reference on B300 (sm103):
+# bf16 max rel err 7.8e-3 (1 ulp at a binade edge), fp16 below 1e-3 (1 ulp = 9.8e-4);
+# the residual is fp32 reordering flipping the final rounding.
+_TOLERANCES = {
+    torch.bfloat16: dict(rtol=1e-2, atol=5e-3),
+    torch.float16: dict(rtol=1e-3, atol=1e-3),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 8192])
+def test_hc_combine_correctness(dtype, num_tokens):
+    block_output, residual, normed_residual, inject_weight = _make_inputs(
+        num_tokens, dtype
+    )
+
+    out = hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+    )
+    expected = _reference_hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        compute_dtype=torch.float64,
+    ).to(dtype)
+
+    torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_hc_combine_out_param(dtype):
+    block_output, residual, normed_residual, inject_weight = _make_inputs(64, dtype)
+    out = torch.empty_like(residual)
+
+    result = hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        out=out,
+    )
+    expected = _reference_hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        compute_dtype=torch.float64,
+    ).to(dtype)
+
+    assert result.data_ptr() == out.data_ptr()
+    torch.testing.assert_close(result, expected, **_TOLERANCES[dtype])
+
+
+def test_hc_combine_3d_input():
+    dtype = torch.bfloat16
+    block_output, residual, normed_residual, inject_weight = _make_inputs(32, dtype)
+    block_output = block_output.reshape(4, 8, HIDDEN_SIZE)
+    residual = residual.reshape(4, 8, HC_COUNT * HIDDEN_SIZE)
+    normed_residual = normed_residual.reshape(4, 8, HC_COUNT * HIDDEN_SIZE)
+
+    out = hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+    )
+    expected = _reference_hc_combine(
+        block_output,
+        residual,
+        normed_residual,
+        inject_weight,
+        HC_COUNT,
+        HIDDEN_SIZE,
+        compute_dtype=torch.float64,
+    ).to(dtype)
+
+    assert out.shape == residual.shape
+    torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
+
+
+def test_hc_combine_bad_hidden_size():
+    dtype = torch.bfloat16
+    block_output, residual, normed_residual, inject_weight = _make_inputs(
+        4,
+        dtype,
+        hc=4,
+        hs=1000,  # 4 * 1000 = 4000, not a multiple of 2048
+    )
+    with pytest.raises(RuntimeError, match="2048"):
+        hc_combine(
+            block_output,
+            residual,
+            normed_residual,
+            inject_weight,
+            4,
+            1000,
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/embeddings/test_qwen4_ple_offload.py

@@ -0,0 +1,269 @@
+import os
+import tempfile
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
+from sglang.srt.layers.vocab_parallel_embedding import (
+    VocabParallelEmbeddingShardIndices,
+)
+from sglang.srt.models import qwen4_exp as qwen4_exp_module
+from sglang.srt.models.qwen4_exp import (
+    Qwen4ExpPinnedHostEmbedding,
+    Qwen4ExpPLELayer,
+)
+from sglang.srt.utils import set_weight_attrs
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required for this test."
+)
+
+
+def _make_source_embedding(
+    *,
+    dtype=torch.bfloat16,
+    embedding_dim=7,
+    vocab_start=0,
+    vocab_end=8,
+    org_vocab_size=8,
+    tp_size=1,
+    num_added_embeddings=0,
+):
+    local_rows = vocab_end - vocab_start
+    weight = nn.Parameter(
+        torch.empty((local_rows, embedding_dim), dtype=dtype, device="cuda"),
+        requires_grad=False,
+    )
+    set_weight_attrs(
+        weight,
+        {
+            "input_dim": 1,
+            "output_dim": 0,
+            "weight_loader": lambda *_args, **_kwargs: None,
+        },
+    )
+    shard_indices = VocabParallelEmbeddingShardIndices(
+        padded_org_vocab_start_index=vocab_start,
+        padded_org_vocab_end_index=vocab_end,
+        padded_added_vocab_start_index=org_vocab_size,
+        padded_added_vocab_end_index=org_vocab_size,
+        org_vocab_start_index=vocab_start,
+        org_vocab_end_index=vocab_end,
+        added_vocab_start_index=org_vocab_size,
+        added_vocab_end_index=org_vocab_size,
+    )
+    return SimpleNamespace(
+        weight=weight,
+        quant_config=None,
+        enable_tp=True,
+        use_attn_tp_group=False,
+        tp_size=tp_size,
+        num_embeddings=org_vocab_size + num_added_embeddings,
+        org_vocab_size=org_vocab_size,
+        padding_size=1,
+        num_added_embeddings=num_added_embeddings,
+        use_presharded_weights=False,
+        org_vocab_size_padded=org_vocab_size,
+        num_embeddings_padded=org_vocab_size + num_added_embeddings,
+        shard_indices=shard_indices,
+        embedding_dim=embedding_dim,
+        weight_scale=None,
+        quant_method=UnquantizedEmbeddingMethod(),
+        num_embeddings_per_partition=local_rows,
+        num_org_embeddings_per_partition=local_rows,
+        num_added_embeddings_per_partition=0,
+    )
+
+
+def _load_rows(offloaded, rows, *, pinned=True):
+    pointer = offloaded.weight.data_ptr()
+    offloaded.weight_loader(offloaded.weight, rows)
+    assert offloaded.weight.data_ptr() == pointer
+    assert offloaded.weight.is_pinned() == pinned
+    assert offloaded.weight.weight_loader.__self__ is offloaded
+    assert offloaded.quant_method is None
+
+
+@pytest.mark.parametrize("input_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("embedding_dim", [7, 64, 257])
+def test_qwen4_ple_pinned_gather_tp1(input_dtype, embedding_dim):
+    source = _make_source_embedding(embedding_dim=embedding_dim)
+    offloaded = Qwen4ExpPinnedHostEmbedding(source)
+    rows = torch.arange(8 * embedding_dim, dtype=torch.bfloat16, device="cuda").reshape(
+        8, embedding_dim
+    )
+    _load_rows(offloaded, rows)
+
+    ids = torch.tensor([[0, 7, 3], [4, 1, 6]], dtype=input_dtype, device="cuda")
+    expected = rows.index_select(0, ids.long().flatten()).reshape(
+        *ids.shape, embedding_dim
+    )
+    actual = offloaded(ids)
+
+    assert actual.shape == expected.shape
+    assert actual.is_contiguous()
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_qwen4_ple_pinned_gather_shard_boundaries_and_out_buffer():
+    embedding_dim = 13
+    source = _make_source_embedding(
+        embedding_dim=embedding_dim,
+        vocab_start=4,
+        vocab_end=8,
+        org_vocab_size=8,
+        tp_size=2,
+    )
+    offloaded = Qwen4ExpPinnedHostEmbedding(source)
+    rows = torch.arange(8 * embedding_dim, dtype=torch.bfloat16, device="cuda").reshape(
+        8, embedding_dim
+    )
+    _load_rows(offloaded, rows)
+
+    ids = torch.tensor([[-1, 3, 4], [7, 8, 100]], device="cuda")
+    output = torch.full(
+        (*ids.shape, embedding_dim),
+        torch.nan,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    actual = offloaded.gather(ids, out=output)
+    expected = torch.zeros_like(output)
+    expected[0, 2] = rows[4]
+    expected[1, 0] = rows[7]
+
+    assert actual.data_ptr() == output.data_ptr()
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_qwen4_ple_pinned_gather_empty_input():
+    offloaded = Qwen4ExpPinnedHostEmbedding(_make_source_embedding())
+    _load_rows(offloaded, torch.zeros((8, 7), dtype=torch.bfloat16, device="cuda"))
+    ids = torch.empty((0, 3), dtype=torch.int64, device="cuda")
+    actual = offloaded.gather(ids)
+    assert actual.shape == (0, 3, 7)
+    assert actual.numel() == 0
+
+
+def test_qwen4_ple_pinned_embedding_rejects_unsupported_weights():
+    with pytest.raises(TypeError, match="requires bfloat16"):
+        Qwen4ExpPinnedHostEmbedding(_make_source_embedding(dtype=torch.float16))
+    with pytest.raises(NotImplementedError, match="added vocabulary"):
+        Qwen4ExpPinnedHostEmbedding(_make_source_embedding(num_added_embeddings=1))
+
+
+def test_qwen4_ple_prefetch_buffer_lifecycle(monkeypatch):
+    layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
+    nn.Module.__init__(layer)
+    layer.ple_embed_dim = 7
+    layer.ple_embedding = SimpleNamespace(
+        ngram_embedding=Qwen4ExpPinnedHostEmbedding(
+            _make_source_embedding(embedding_dim=layer.ple_embed_dim)
+        )
+    )
+    layer._graph_prefetch_buffers = {}
+    layer._eager_prefetch_buffer = None
+    lookup_ids = torch.empty((0,), dtype=torch.int64, device="cuda")
+
+    monkeypatch.setattr(qwen4_exp_module, "get_is_capture_mode", lambda: False)
+    eager_large = layer._get_prefetch_buffer(8, lookup_ids)
+    eager_small = layer._get_prefetch_buffer(3, lookup_ids)
+    assert eager_small.data_ptr() == eager_large.data_ptr()
+    assert layer._eager_prefetch_buffer.shape == (8, layer.ple_embed_dim)
+
+    eager_grown = layer._get_prefetch_buffer(12, lookup_ids)
+    eager_grown_small = layer._get_prefetch_buffer(4, lookup_ids)
+    assert eager_grown_small.data_ptr() == eager_grown.data_ptr()
+    assert layer._eager_prefetch_buffer.shape == (12, layer.ple_embed_dim)
+
+    monkeypatch.setattr(qwen4_exp_module, "get_is_capture_mode", lambda: True)
+    graph_three = layer._get_prefetch_buffer(3, lookup_ids)
+    graph_five = layer._get_prefetch_buffer(5, lookup_ids)
+    graph_three_reused = layer._get_prefetch_buffer(3, lookup_ids)
+    assert graph_three_reused.data_ptr() == graph_three.data_ptr()
+    assert graph_five.data_ptr() != graph_three.data_ptr()
+    assert set(layer._graph_prefetch_buffers) == {3, 5}
+
+
+def _file_backend_supported() -> bool:
+    from sglang.srt.models.qwen4_exp_ple_table import device_uses_host_page_tables
+
+    return (
+        torch.cuda.is_available()
+        and device_uses_host_page_tables(torch.cuda.current_device()) is True
+    )
+
+
+@pytest.mark.skipif(
+    not _file_backend_supported(),
+    reason="the file backend needs pageable host memory reachable through host page tables",
+)
+@pytest.mark.parametrize("embedding_dim", [7, 160])
+def test_qwen4_ple_file_backend_matches_pinned(embedding_dim):
+    with tempfile.TemporaryDirectory() as table_dir:
+        pinned = Qwen4ExpPinnedHostEmbedding(
+            _make_source_embedding(embedding_dim=embedding_dim)
+        )
+        filed = Qwen4ExpPinnedHostEmbedding(
+            _make_source_embedding(embedding_dim=embedding_dim),
+            backend="file",
+            table_dir=table_dir,
+        )
+        assert pinned._file_prefetcher is None and filed._file_prefetcher is not None
+        (name,) = os.listdir(table_dir)
+        assert "rows0-8" in name  # this rank's vocabulary shard
+        rows = torch.arange(
+            8 * embedding_dim, dtype=torch.bfloat16, device="cuda"
+        ).reshape(8, embedding_dim)
+        _load_rows(pinned, rows)
+        _load_rows(filed, rows, pinned=False)
+        ids = torch.tensor([[0, 7, 3], [4, 1, 6]], dtype=torch.int64, device="cuda")
+        torch.testing.assert_close(filed(ids), pinned(ids), rtol=0, atol=0)
+        # A prefill-sized gather goes through the page-cache hint path.
+        big = torch.randint(0, 8, (4096,), device="cuda")
+        torch.testing.assert_close(
+            filed(big), rows.index_select(0, big), rtol=0, atol=0
+        )
+
+
+@pytest.mark.skipif(
+    not _file_backend_supported(),
+    reason="the file backend needs pageable host memory reachable through host page tables",
+)
+def test_qwen4_ple_file_backend_fp8_table():
+    embedding_dim = 160
+    with tempfile.TemporaryDirectory() as table_dir:
+        filed = Qwen4ExpPinnedHostEmbedding(
+            _make_source_embedding(
+                embedding_dim=embedding_dim, dtype=torch.float8_e4m3fn
+            ),
+            backend="file",
+            table_dir=table_dir,
+        )
+        assert filed.weight.dtype == torch.float8_e4m3fn
+        rows = (
+            torch.arange(8 * embedding_dim, dtype=torch.float32, device="cuda").reshape(
+                8, embedding_dim
+            )
+            / 64
+        ).to(torch.float8_e4m3fn)
+        _load_rows(filed, rows, pinned=False)
+        ids = torch.tensor([[0, 7, 3]], dtype=torch.int64, device="cuda")
+        expected = (
+            rows.index_select(0, ids.flatten())
+            .to(torch.bfloat16)
+            .reshape(1, 3, embedding_dim)
+        )
+        torch.testing.assert_close(filed(ids), expected, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/hyperconnection/test_hc_mix_triton.py

@@ -0,0 +1,84 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.layers.hc_mix_triton import (
+    _FUSED_MIX_MAX_ROWS,
+    fused_hc_mix,
+    fused_hc_mix_supported,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+HC_COUNT = 4
+HIDDEN_SIZE = 2560
+LOWRANK = 320
+
+
+def _reference_mix(
+    hyper_input_normed: torch.Tensor,
+    w_down: torch.Tensor,
+    w_up: torch.Tensor,
+    hc: int,
+    hs: int,
+    compute_dtype: torch.dtype = torch.float64,
+) -> torch.Tensor:
+    """Mirrors GatedResidual._mix_compute in hyperconnection.py."""
+    x = hyper_input_normed.to(compute_dtype)
+    t = F.silu(F.linear(x, w_down.to(compute_dtype)) / hc)
+    u = torch.sigmoid(F.linear(t, w_up.to(compute_dtype)))
+    return (u.unflatten(-1, (hc, hs)) * x.unflatten(-1, (hc, hs))).mean(dim=-2)
+
+
+def _make_inputs(num_tokens: int, dtype: torch.dtype):
+    torch.manual_seed(0)
+    x = torch.randn(num_tokens, HC_COUNT * HIDDEN_SIZE, dtype=dtype, device="cuda")
+    w_down = (
+        torch.randn(LOWRANK, HC_COUNT * HIDDEN_SIZE, dtype=dtype, device="cuda") * 0.02
+    )
+    w_up = (
+        torch.randn(HC_COUNT * HIDDEN_SIZE, LOWRANK, dtype=dtype, device="cuda") * 0.02
+    )
+    return x, w_down, w_up
+
+
+_TOLERANCES = {
+    torch.bfloat16: dict(rtol=1e-2, atol=5e-3),
+    torch.float16: dict(rtol=2e-3, atol=1e-3),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_tokens", [1, 4, 7, _FUSED_MIX_MAX_ROWS])
+def test_fused_hc_mix_matches_reference(dtype, num_tokens):
+    x, w_down, w_up = _make_inputs(num_tokens, dtype)
+    assert fused_hc_mix_supported(x, w_down, w_up)
+    out = fused_hc_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    ref = _reference_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    torch.testing.assert_close(out.to(torch.float64), ref, **_TOLERANCES[dtype])
+
+
+def test_fused_hc_mix_no_less_accurate_than_eager():
+    """The fused kernel (fp32 accumulation throughout) must not be farther
+    from the fp64 reference than the eager bf16 chain it replaces."""
+    x, w_down, w_up = _make_inputs(8, torch.bfloat16)
+    ref = _reference_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    fused = fused_hc_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
+    eager = _reference_mix(
+        x, w_down, w_up, HC_COUNT, HIDDEN_SIZE, compute_dtype=torch.bfloat16
+    )
+    fused_err = (fused.to(torch.float64) - ref).abs().max()
+    eager_err = (eager.to(torch.float64) - ref).abs().max()
+    assert fused_err <= eager_err * 1.5 + 1e-6
+
+
+def test_fused_hc_mix_gate_rejects_prefill_rows():
+    x, w_down, w_up = _make_inputs(_FUSED_MIX_MAX_ROWS + 1, torch.bfloat16)
+    assert not fused_hc_mix_supported(x, w_down, w_up)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/kimi_k3/test_ar_fusion.py

@@ -10,7 +10,7 @@ semaphore window cycling.
 
 Usage::
 
-    python test/registered/jit/kimi_k3/test_ar_fusion.py   # relaunches under torchrun (8 GPUs)
+    python test/registered/kernels/ops/kimi_k3/test_ar_fusion.py   # relaunches under torchrun (8 GPUs)
 """
 
 from __future__ import annotations

test/registered/kernels/ops/kvcache/test_hisparse_spec.py

@@ -0,0 +1,550 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import pytest
+import torch
+
+from sglang.kernels.ops.kvcache.hisparse import (
+    HiSparseSpecState,
+    copy_cache_planned_mla,
+    load_cache_to_device_buffer_spec_mla,
+)
+from sglang.srt.utils import is_npu, is_xpu
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or is_npu() or is_xpu(),
+    reason="HiSparse speculative swap tests require a CUDA GPU.",
+)
+
+DEVICE = "cuda"
+TOKEN_SCALE = 1_000_003
+
+
+class _SwapState(NamedTuple):
+    device_buffer_tokens: torch.Tensor
+    device_buffer_locs: torch.Tensor
+    host_cache_locs: torch.Tensor
+    host_cache: torch.Tensor
+    device_buffer: torch.Tensor
+    swap_state: HiSparseSpecState
+
+
+def _make_cache_index(num_reqs: int, hot_buffer_size: int) -> torch.Tensor:
+    hash_size = 1 << (2 * hot_buffer_size - 1).bit_length()
+    cache_index = torch.full(
+        (num_reqs, 2, hash_size), -1, dtype=torch.int64, device=DEVICE
+    )
+    tokens = torch.arange(hot_buffer_size, dtype=torch.int64, device=DEVICE)
+    hash_slots = ((tokens * 2654435761) & (hash_size - 1)).to(torch.long)
+    packed_entries = (tokens << 32) | tokens
+    cache_index[:, 0, hash_slots] = packed_entries
+    return cache_index
+
+
+def _make_state(
+    *,
+    num_reqs: int,
+    hot_buffer_size: int,
+    page_size: int,
+    scratch_size: int,
+    seq_len: int,
+    item_words: int,
+    metadata_occurrences: int,
+) -> _SwapState:
+    buffer_size = hot_buffer_size + page_size
+    device_buffer_tokens = torch.full(
+        (num_reqs, buffer_size), -1, dtype=torch.int32, device=DEVICE
+    )
+    device_buffer_tokens[:, :hot_buffer_size] = torch.arange(
+        hot_buffer_size, dtype=torch.int32, device=DEVICE
+    )
+
+    physical_tokens_per_req = buffer_size + scratch_size
+    request_bases = (
+        torch.arange(num_reqs, dtype=torch.int32, device=DEVICE).view(-1, 1)
+        * physical_tokens_per_req
+    )
+    device_buffer_locs = (
+        request_bases
+        + torch.arange(buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1)
+    ).contiguous()
+    scratch_locs = (
+        request_bases
+        + buffer_size
+        + torch.arange(scratch_size, dtype=torch.int32, device=DEVICE).view(1, -1)
+    ).contiguous()
+
+    host_cache_locs = torch.arange(seq_len, dtype=torch.int64, device=DEVICE)
+    host_cache_locs = host_cache_locs.view(1, -1).repeat(num_reqs, 1).contiguous()
+    host_cache = torch.empty((seq_len, item_words), dtype=torch.int64, pin_memory=True)
+    host_cache.copy_(
+        torch.arange(seq_len, dtype=torch.int64).view(-1, 1) * TOKEN_SCALE
+        + torch.arange(item_words, dtype=torch.int64).view(1, -1)
+    )
+
+    device_buffer = torch.full(
+        (num_reqs * physical_tokens_per_req, item_words),
+        -1,
+        dtype=torch.int64,
+        device=DEVICE,
+    )
+    hot_locs = device_buffer_locs[:, :hot_buffer_size].to(torch.long)
+    device_buffer[hot_locs] = host_cache[:hot_buffer_size].to(DEVICE)
+
+    scratch_state = torch.full(
+        (num_reqs + 1, max(4 * num_reqs, 5 * metadata_occurrences)),
+        -1,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    scratch_state[0].zero_()
+    swap_state = HiSparseSpecState(
+        cache_index=_make_cache_index(num_reqs, hot_buffer_size),
+        cache_policy=torch.zeros(
+            (num_reqs + 1, hot_buffer_size),
+            dtype=torch.int32,
+            device=DEVICE,
+        ),
+        scratch_locs=scratch_locs,
+        scratch_state=scratch_state,
+    )
+    return _SwapState(
+        device_buffer_tokens=device_buffer_tokens,
+        device_buffer_locs=device_buffer_locs,
+        host_cache_locs=host_cache_locs,
+        host_cache=host_cache,
+        device_buffer=device_buffer,
+        swap_state=swap_state,
+    )
+
+
+def _run_swap(
+    *,
+    top_k_tokens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    state: _SwapState,
+    out: torch.Tensor | None = None,
+    req_pool_indices: torch.Tensor | None = None,
+    num_real_reqs: torch.Tensor | None = None,
+    miss_src: torch.Tensor | None = None,
+    miss_dst: torch.Tensor | None = None,
+    miss_count: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.full_like(top_k_tokens, -1)
+    else:
+        out.fill_(-1)
+    num_reqs = top_k_tokens.size(0)
+    if req_pool_indices is None:
+        req_pool_indices = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
+    if num_real_reqs is None:
+        num_real_reqs = torch.tensor([num_reqs], dtype=torch.int32, device=DEVICE)
+    load_cache_to_device_buffer_spec_mla(
+        top_k_tokens=top_k_tokens,
+        device_buffer_tokens=state.device_buffer_tokens,
+        host_cache_locs=state.host_cache_locs,
+        device_buffer_locs=state.device_buffer_locs,
+        host_cache=state.host_cache,
+        device_buffer=state.device_buffer,
+        top_k_device_locs=out,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        state=state.swap_state,
+        num_real_reqs=num_real_reqs,
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+    )
+    return out
+
+
+def _assert_output_matches_tokens(
+    state: _SwapState, out: torch.Tensor, tokens: torch.Tensor
+) -> None:
+    actual = state.device_buffer[out.to(torch.long)]
+    expected = tokens.to(torch.int64).unsqueeze(-1) * TOKEN_SCALE + torch.arange(
+        state.device_buffer.size(-1), dtype=torch.int64, device=DEVICE
+    )
+    torch.testing.assert_close(actual, expected)
+
+
+class TestHiSparseSpec(CustomTestCase):
+    def test_deduplicates_repeated_misses_and_copies_full_items(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k, item_words = 4, 2048, 72
+        total_occurrences = num_steps * top_k
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=16384,
+            item_words=item_words,
+            metadata_occurrences=total_occurrences,
+        )
+
+        miss_count = 196
+        hits = torch.arange(top_k - miss_count, dtype=torch.int32, device=DEVICE)
+        shared_misses = hot_size + torch.arange(
+            miss_count, dtype=torch.int32, device=DEVICE
+        )
+        step = torch.cat((hits, shared_misses))
+        top_k_tokens = step.view(1, 1, -1).repeat(1, num_steps, 1)
+        seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
+
+        out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
+        torch.cuda.synchronize()
+
+        _assert_output_matches_tokens(state, out, top_k_tokens)
+        self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), miss_count)
+        repeated_miss_locs = out[0, :, -miss_count:]
+        self.assertTrue(torch.all(repeated_miss_locs == repeated_miss_locs[0]).item())
+
+    def test_copies_782_cross_step_unique_misses(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k, item_words = 4, 2048, 72
+        total_occurrences = num_steps * top_k
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=16384,
+            item_words=item_words,
+            metadata_occurrences=total_occurrences,
+        )
+
+        steps = []
+        next_miss = hot_size
+        for step_idx, miss_count in enumerate((196, 196, 195, 195)):
+            hits = torch.roll(
+                torch.arange(hot_size, dtype=torch.int32, device=DEVICE),
+                step_idx * 137,
+            )[: top_k - miss_count]
+            misses = torch.arange(
+                next_miss,
+                next_miss + miss_count,
+                dtype=torch.int32,
+                device=DEVICE,
+            )
+            next_miss += miss_count
+            steps.append(torch.cat((hits, misses)))
+        top_k_tokens = torch.stack(steps).unsqueeze(0).contiguous()
+        seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
+        out = _run_swap(
+            top_k_tokens=top_k_tokens,
+            seq_lens=seq_lens,
+            state=state,
+        )
+        torch.cuda.synchronize()
+
+        _assert_output_matches_tokens(state, out, top_k_tokens)
+        self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 782)
+
+    def test_records_union_plan_for_shared_layer_io(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k, item_words = 4, 2048, 72
+        total_occurrences = num_steps * top_k
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=16384,
+            item_words=item_words,
+            metadata_occurrences=total_occurrences,
+        )
+
+        steps = []
+        next_miss = hot_size
+        for step_idx, step_miss_count in enumerate((196, 196, 195, 195)):
+            hits = torch.roll(
+                torch.arange(hot_size, dtype=torch.int32, device=DEVICE),
+                step_idx * 137,
+            )[: top_k - step_miss_count]
+            misses = torch.arange(
+                next_miss,
+                next_miss + step_miss_count,
+                dtype=torch.int32,
+                device=DEVICE,
+            )
+            next_miss += step_miss_count
+            steps.append(torch.cat((hits, misses)))
+        top_k_tokens = torch.stack(steps).unsqueeze(0).contiguous()
+        seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
+        miss_src = torch.full(
+            (1, total_occurrences), -1, dtype=torch.int64, device=DEVICE
+        )
+        miss_dst = torch.full(
+            (1, total_occurrences), -1, dtype=torch.int32, device=DEVICE
+        )
+        miss_count = torch.full((1,), -1, dtype=torch.int32, device=DEVICE)
+
+        _run_swap(
+            top_k_tokens=top_k_tokens,
+            seq_lens=seq_lens,
+            state=state,
+            miss_src=miss_src,
+            miss_dst=miss_dst,
+            miss_count=miss_count,
+        )
+        shared_layer_buffer = torch.full_like(state.device_buffer, -1)
+        copy_cache_planned_mla(
+            miss_src=miss_src,
+            miss_dst=miss_dst,
+            miss_count=miss_count,
+            num_real_reqs=torch.ones(1, dtype=torch.int32, device=DEVICE),
+            host_cache=state.host_cache,
+            device_buffer=shared_layer_buffer,
+            item_size_bytes=state.host_cache.stride(0)
+            * state.host_cache.element_size(),
+        )
+        torch.cuda.synchronize()
+
+        self.assertEqual(int(miss_count.item()), 782)
+        count = int(miss_count.item())
+        src = miss_src[0, :count].to(torch.long)
+        dst = miss_dst[0, :count].to(torch.long)
+        torch.testing.assert_close(
+            shared_layer_buffer[dst], state.host_cache[src.cpu()].to(DEVICE)
+        )
+        torch.testing.assert_close(shared_layer_buffer[dst], state.device_buffer[dst])
+
+    def test_padded_request_clears_stale_plan_count(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k = 4, 2048
+        total_occurrences = num_steps * top_k
+        state = _make_state(
+            num_reqs=2,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=8192,
+            item_words=1,
+            metadata_occurrences=total_occurrences,
+        )
+        top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).view(
+            1, 1, -1
+        )
+        top_k_tokens = top_k_tokens.repeat(2, num_steps, 1).contiguous()
+        seq_lens = torch.full((2 * num_steps,), 8192, dtype=torch.int32, device=DEVICE)
+        miss_src = torch.full(
+            (2, total_occurrences), -1, dtype=torch.int64, device=DEVICE
+        )
+        miss_dst = torch.full(
+            (2, total_occurrences), -1, dtype=torch.int32, device=DEVICE
+        )
+        miss_count = torch.full((2,), 123, dtype=torch.int32, device=DEVICE)
+
+        _run_swap(
+            top_k_tokens=top_k_tokens,
+            seq_lens=seq_lens,
+            state=state,
+            num_real_reqs=torch.ones(1, dtype=torch.int32, device=DEVICE),
+            miss_src=miss_src,
+            miss_dst=miss_dst,
+            miss_count=miss_count,
+        )
+        torch.cuda.synchronize()
+
+        self.assertEqual(int(miss_count[1].item()), 0)
+
+    def test_resolves_all_speculative_extra_page_slots_without_host_io(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k = 4, 2048
+        total_occurrences = num_steps * top_k
+        seq_len = 8192
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=seq_len,
+            item_words=72,
+            metadata_occurrences=total_occurrences,
+        )
+
+        draft_tokens = torch.arange(
+            seq_len - num_steps, seq_len, dtype=torch.int32, device=DEVICE
+        )
+        extra_offsets = torch.tensor([0, 7, 31, 63], device=DEVICE)
+        extra_locs = state.device_buffer_locs[0, hot_size + extra_offsets].to(
+            torch.long
+        )
+        state.device_buffer_tokens[0, hot_size + extra_offsets] = draft_tokens
+        state.device_buffer[extra_locs] = state.host_cache[
+            draft_tokens.to(device="cpu", dtype=torch.long)
+        ].to(DEVICE)
+        state.host_cache_locs[0, draft_tokens.to(torch.long)] = -1
+
+        hits = torch.arange(top_k - 1, dtype=torch.int32, device=DEVICE)
+        top_k_tokens = torch.stack(
+            [torch.cat((hits, draft_tokens[step : step + 1])) for step in range(4)]
+        ).unsqueeze(0)
+        seq_lens = draft_tokens + 1
+
+        out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
+        torch.cuda.synchronize()
+
+        _assert_output_matches_tokens(state, out, top_k_tokens)
+        torch.testing.assert_close(out[0, :, -1].to(torch.long), extra_locs)
+        self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 0)
+
+    def test_full_union_overflow_preserves_all_8192_outputs(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k = 4, 2048
+        total_occurrences = num_steps * top_k
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=total_occurrences - hot_size,
+            seq_len=16384,
+            item_words=72,
+            metadata_occurrences=total_occurrences,
+        )
+        top_k_tokens = (
+            hot_size + torch.arange(total_occurrences, dtype=torch.int32, device=DEVICE)
+        ).view(1, num_steps, top_k)
+        seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
+
+        miss_src = torch.full(
+            (1, total_occurrences), -1, dtype=torch.int64, device=DEVICE
+        )
+        miss_dst = torch.full(
+            (1, total_occurrences), -1, dtype=torch.int32, device=DEVICE
+        )
+        miss_count = torch.full((1,), -1, dtype=torch.int32, device=DEVICE)
+        out = _run_swap(
+            top_k_tokens=top_k_tokens,
+            seq_lens=seq_lens,
+            state=state,
+            miss_src=miss_src,
+            miss_dst=miss_dst,
+            miss_count=miss_count,
+        )
+        torch.cuda.synchronize()
+
+        _assert_output_matches_tokens(state, out, top_k_tokens)
+        self.assertEqual(torch.unique(out).numel(), total_occurrences)
+        self.assertEqual(
+            int(state.swap_state.scratch_state[0, 0].item()), total_occurrences
+        )
+        self.assertEqual(int(miss_count.item()), total_occurrences)
+        self.assertTrue(miss_src.ge(0).all().item())
+        self.assertTrue(miss_dst.ge(0).all().item())
+
+    def test_packed_ring_supports_glm52_native_context_length(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k = 4, 2048
+        total_occurrences = num_steps * top_k
+        seq_len = 1_048_648
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=seq_len,
+            item_words=1,
+            metadata_occurrences=total_occurrences,
+        )
+        top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).view(
+            1, 1, -1
+        )
+        top_k_tokens = top_k_tokens.repeat(1, num_steps, 1)
+        high_token = seq_len - 1
+        top_k_tokens[:, :, -1] = high_token
+        seq_lens = torch.full((num_steps,), seq_len, dtype=torch.int32, device=DEVICE)
+
+        out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
+        torch.cuda.synchronize()
+
+        _assert_output_matches_tokens(state, out, top_k_tokens)
+        self.assertTrue(out.ge(0).all().item())
+
+        # The first call admits the high token into the packed hash. The
+        # second call must resolve it as a hot hit rather than truncating the
+        # packed int64 entry and repeating Host-to-GPU IO.
+        out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
+        torch.cuda.synchronize()
+        _assert_output_matches_tokens(state, out, top_k_tokens)
+        self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 0)
+
+    def test_cuda_graph_replay_preserves_valid_locations(self) -> None:
+        hot_size, page_size = 4096, 64
+        num_steps, top_k = 4, 2048
+        total_occurrences = num_steps * top_k
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=hot_size,
+            page_size=page_size,
+            scratch_size=hot_size,
+            seq_len=65536,
+            item_words=72,
+            metadata_occurrences=total_occurrences,
+        )
+        top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).repeat(
+            num_steps, 1
+        )
+        for step, miss_count in enumerate((164, 102, 61, 20)):
+            top_k_tokens[step, -miss_count:] = torch.arange(
+                8192 + step * top_k,
+                8192 + step * top_k + miss_count,
+                dtype=torch.int32,
+                device=DEVICE,
+            )
+        top_k_tokens = top_k_tokens.unsqueeze(0).contiguous()
+        seq_lens = torch.tensor(
+            [65533, 65534, 65535, 65536], dtype=torch.int32, device=DEVICE
+        )
+
+        _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
+        torch.cuda.synchronize()
+        graph_out = torch.full_like(top_k_tokens, -1)
+        req_pool_indices = torch.arange(1, dtype=torch.int64, device=DEVICE)
+        num_real_reqs = torch.tensor([1], dtype=torch.int32, device=DEVICE)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            _run_swap(
+                top_k_tokens=top_k_tokens,
+                seq_lens=seq_lens,
+                state=state,
+                out=graph_out,
+                req_pool_indices=req_pool_indices,
+                num_real_reqs=num_real_reqs,
+            )
+
+        for _ in range(4):
+            graph.replay()
+        torch.cuda.synchronize()
+
+        _assert_output_matches_tokens(state, graph_out, top_k_tokens)
+        self.assertTrue(graph_out.ge(0).all().item())
+
+    def test_rejects_invalid_step_shape_before_compilation(self) -> None:
+        state = _make_state(
+            num_reqs=1,
+            hot_buffer_size=4096,
+            page_size=64,
+            scratch_size=4096,
+            seq_len=8192,
+            item_words=1,
+            metadata_occurrences=8192,
+        )
+        with self.assertRaisesRegex(ValueError, "2-4 steps"):
+            _run_swap(
+                top_k_tokens=torch.zeros(
+                    (1, 1, 2048), dtype=torch.int32, device=DEVICE
+                ),
+                seq_lens=torch.tensor([8192], dtype=torch.int32, device=DEVICE),
+                state=state,
+            )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/layernorm/test_grouped_gemma_rmsnorm.py

@@ -0,0 +1,87 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.layernorm.grouped_gemma_rmsnorm import grouped_gemma_rmsnorm
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _reference_grouped_gemma_rmsnorm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    group_size: int,
+    eps: float,
+    compute_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Mirrors GroupedGemmaRMSNorm.forward (hyperconnection.py); keep them in sync."""
+    x_float = x.to(compute_dtype)
+    hidden = x_float.shape[-1]
+    x_grouped = x_float.reshape(*x_float.shape[:-1], hidden // group_size, group_size)
+    variance = x_grouped.pow(2).mean(dim=-1, keepdim=True)
+    x_norm = (x_grouped * torch.rsqrt(variance + eps)).flatten(-2)
+    return x_norm * (1.0 + weight.to(compute_dtype))
+
+
+# Tolerances are at the output-dtype quantization floor, measured against the
+# fp64 reference on 4xB300 (sm103): bf16 max rel err 3.9e-3 (1 ulp), fp16
+# 4.9e-4 (0.5 ulp). The kernel computes in fp32 like the eager reference.
+_TOLERANCES = {
+    torch.bfloat16: dict(rtol=5e-3, atol=5e-3),
+    torch.float16: dict(rtol=1e-3, atol=1e-3),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize(
+    "num_tokens,hidden_size,group_size",
+    [
+        (1, 10240, 2560),  # production shape (HC 4 x 2560)
+        (7, 10240, 2560),
+        (128, 10240, 2560),
+        (33, 1024, 512),
+        (5, 512, 512),  # single group == plain gemma rmsnorm
+        (1024, 2048, 1024),
+    ],
+)
+@pytest.mark.parametrize("eps", [1e-6, 1e-5])
+def test_grouped_gemma_rmsnorm_correctness(
+    dtype, num_tokens, hidden_size, group_size, eps
+):
+    torch.manual_seed(0)
+    x = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda")
+    weight = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2
+
+    out = grouped_gemma_rmsnorm(x, weight, group_size, eps)
+    expected = _reference_grouped_gemma_rmsnorm(
+        x, weight, group_size, eps, compute_dtype=torch.float64
+    ).to(dtype)
+
+    torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
+
+
+def test_grouped_gemma_rmsnorm_out_param():
+    x = torch.randn(64, 10240, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(10240, dtype=torch.bfloat16, device="cuda") * 0.2
+    out = torch.empty_like(x)
+
+    result = grouped_gemma_rmsnorm(x, weight, 2560, 1e-6, out=out)
+    expected = _reference_grouped_gemma_rmsnorm(
+        x, weight, 2560, 1e-6, compute_dtype=torch.float64
+    ).to(x.dtype)
+
+    assert result.data_ptr() == out.data_ptr()
+    torch.testing.assert_close(result, expected, **_TOLERANCES[x.dtype])
+
+
+def test_grouped_gemma_rmsnorm_bad_group_size():
+    x = torch.randn(4, 10240, dtype=torch.bfloat16, device="cuda")
+    weight = torch.zeros(10240, dtype=torch.bfloat16, device="cuda")
+    with pytest.raises(RuntimeError, match="group_size"):
+        grouped_gemma_rmsnorm(x, weight, 1000, 1e-6)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/layernorm/test_mxfp8_epilogue.py

@@ -0,0 +1,153 @@
+"""MXFP8 epilogues stay bitwise identical to FlashInfer's standalone quantizer."""
+
+import sys
+
+import flashinfer
+import pytest
+import torch
+from flashinfer import mxfp8_quantize
+
+from sglang.kernels.ops.attention.dsv4.wo_a_bf16 import (
+    _quantize_partial,
+    _wo_a_reduce,
+    wo_a_bf16_small_batch,
+    wo_a_bf16_small_batch_mxfp8,
+)
+from sglang.kernels.ops.layernorm.mxfp8_epilogue import rmsnorm_mxfp8
+from sglang.srt.runtime_context import get_platform
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.version.cuda is None
+    or not get_platform().is_blackwell,
+    reason="the MXFP8 reference quantizer is Blackwell-only",
+)
+
+DEVICE = "cuda"
+
+HIDDEN = 5120
+STREAMS = 4
+
+
+@pytest.mark.parametrize("m", [1, 2, 5, 6, 8])
+@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3])
+@pytest.mark.parametrize("backend", ["cuda", "cute-dsl"])
+def test_bitwise_identical_to_norm_then_quantize(m: int, scale: float, backend: str):
+    from sglang.kernels.ops.layernorm.hc_combine_norm import (
+        hc_combine_norm,
+        hc_combine_norm_mxfp8,
+    )
+    from sglang.srt.layers.quantization.fp8_utils import flashinfer_mxfp8_quantize
+
+    g = torch.Generator(device="cuda").manual_seed(m * 31 + int(scale * 1000))
+    x = (
+        torch.randn(
+            (m, STREAMS * HIDDEN), device="cuda", dtype=torch.bfloat16, generator=g
+        )
+        * scale
+    )
+    pre = torch.randn(
+        (m, STREAMS), device="cuda", dtype=torch.bfloat16, generator=g
+    ).contiguous()
+    w = torch.randn((HIDDEN,), device="cuda", dtype=torch.bfloat16, generator=g)
+    eps = 1e-6
+
+    y_ref = hc_combine_norm(x, pre, w, eps)
+    q_ref, sf_ref = flashinfer_mxfp8_quantize(y_ref, True, 32, backend)
+    y, q, sf = hc_combine_norm_mxfp8(x, pre, w, eps)
+
+    assert torch.equal(y, y_ref)
+    assert torch.equal(
+        q.reshape(-1).view(torch.uint8), q_ref.reshape(-1).view(torch.uint8)
+    )
+    assert sf.shape == sf_ref.reshape(-1).shape
+    assert torch.equal(sf, sf_ref.reshape(-1))
+
+
+def check(x, w, got):
+    y, q, sf = got
+    expected = flashinfer.norm.rmsnorm(x, w, 1e-6)
+    eq, esf = flashinfer.mxfp8_quantize(expected, is_sf_swizzled_layout=True)
+    assert torch.equal(y.view(torch.int16), expected.view(torch.int16))
+    assert torch.equal(q.view(torch.uint8), eq.view(torch.uint8))
+    m = x.shape[0]
+    g = torch.arange(40, device=x.device)
+    row = torch.arange(m, device=x.device)[:, None]
+    offsets = (g // 4) * 512 + ((row % 32) * 4 + row // 32) * 4 + g % 4
+    assert torch.equal(sf[offsets], esf.flatten()[offsets])
+    pad = torch.ones_like(sf, dtype=torch.bool)
+    pad[offsets] = False
+    assert torch.count_nonzero(sf[pad]) == 0
+
+
+@pytest.mark.parametrize("m", [1, 5, 6, 8])
+@pytest.mark.parametrize("stride", [1280, 1792])
+def test_dynamic_graph(m, stride):
+    torch.manual_seed(941)
+    x = torch.randn((m, stride), device="cuda", dtype=torch.bfloat16)[:, :1280]
+    w = torch.randn(1280, device="cuda", dtype=torch.bfloat16)
+    for _ in range(3):
+        check(x, w, rmsnorm_mxfp8(x, w, 1e-6))
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        got = rmsnorm_mxfp8(x, w, 1e-6)
+    for scale in [0, 1e-4, 1, 100]:
+        x.copy_(torch.randn_like(x) * scale)
+        w.copy_(torch.randn_like(w))
+        graph.replay()
+        check(x, w, got)
+
+
+def test_wo_a_partial_quant_matches_quantize():
+    torch.manual_seed(0)
+    for rows in range(2, 9):
+        for magnitude in (0.0, 1e-37, 1e-7, 1.0, 448.0, 1e10):
+            partial = torch.randn(8, rows, 2, 1024, device=DEVICE) * magnitude
+            bf16 = torch.empty(rows, 2048, dtype=torch.bfloat16, device=DEVICE)
+            _wo_a_reduce[(rows * 8,)](partial, bf16, rows * 2048, num_warps=4)
+            expected_q, expected_s = mxfp8_quantize(bf16, True, alignment=32)
+            actual_q, actual_s = _quantize_partial(partial)
+            torch.testing.assert_close(
+                actual_q.view(torch.uint8),
+                expected_q.view(torch.uint8),
+                rtol=0,
+                atol=0,
+            )
+            torch.testing.assert_close(actual_s, expected_s, rtol=0, atol=0)
+
+        x = torch.randn(rows, 64, 512, device=DEVICE, dtype=torch.bfloat16)[
+            :, :16
+        ].view(rows, 2, 4096)
+        wo_a = torch.randn(2, 1024, 4096, device=DEVICE, dtype=torch.bfloat16) * 0.02
+        bf16 = wo_a_bf16_small_batch(x, wo_a).flatten(1)
+        q, s = wo_a_bf16_small_batch_mxfp8(x, wo_a)
+        expected_q, expected_s = mxfp8_quantize(bf16, True, alignment=32)
+        torch.testing.assert_close(
+            q.view(torch.uint8), expected_q.view(torch.uint8), rtol=0, atol=0
+        )
+        torch.testing.assert_close(s, expected_s, rtol=0, atol=0)
+
+    partial = torch.randn(8, 6, 2, 1024, device=DEVICE)
+    _quantize_partial(partial)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        q, s = _quantize_partial(partial)
+    for _ in range(3):
+        partial.normal_()
+        # The replay must regenerate scale padding as well as live rows.
+        s.fill_(255)
+        graph.replay()
+        bf16 = torch.empty(6, 2048, dtype=torch.bfloat16, device=DEVICE)
+        _wo_a_reduce[(48,)](partial, bf16, 6 * 2048, num_warps=4)
+        eq, es = mxfp8_quantize(bf16, True, alignment=32)
+        torch.testing.assert_close(
+            q.view(torch.uint8), eq.view(torch.uint8), rtol=0, atol=0
+        )
+        torch.testing.assert_close(s, es, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/memory/test_adler32_checksum.py

@@ -0,0 +1,95 @@
+"""Tests for Adler-32 GPU checksum against Python zlib.adler32."""
+
+import unittest
+import zlib
+
+import torch
+
+from sglang.kernels.ops.memory.adler32 import (
+    adler32_checksum,
+    adler32_regions_checksum,
+    adler32_strided_checksum,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _ref_adler32(tensor: torch.Tensor) -> int:
+    return zlib.adler32(tensor.cpu().contiguous().view(torch.uint8).numpy().tobytes())
+
+
+def _ref_strided_adler32(tensors, indices, strides) -> int:
+    parts = []
+    for tensor, idx, stride in zip(tensors, indices, strides):
+        raw = tensor.cpu().contiguous().flatten().view(torch.uint8)
+        for i in idx.cpu().tolist():
+            parts.append(raw[i * stride : (i + 1) * stride].numpy().tobytes())
+    return zlib.adler32(b"".join(parts))
+
+
+class TestAdler32(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        torch.manual_seed(42)
+
+    def _check_whole(self, tensor):
+        self.assertEqual(adler32_checksum(tensor), _ref_adler32(tensor))
+
+    def test_whole_small(self):
+        self._check_whole(torch.tensor([1.0, 2, 3, 4], device="cuda"))
+
+    def test_whole_single_element(self):
+        self._check_whole(torch.tensor([42.0], device="cuda"))
+
+    def test_whole_dtypes(self):
+        for dtype, shape in [
+            (torch.bfloat16, (1024, 128)),
+            (torch.float16, (512, 64)),
+            (torch.float32, (4096, 256)),
+        ]:
+            self._check_whole(torch.randn(*shape, dtype=dtype, device="cuda"))
+
+    def _check_strided(self, tensors, indices, strides):
+        actual = adler32_strided_checksum(
+            [t.data_ptr() for t in tensors], strides, indices
+        )
+        expected = _ref_strided_adler32(tensors, indices, strides)
+        self.assertEqual(actual, expected)
+
+    def test_strided_single_tensor(self):
+        t = torch.randn(100, 64, device="cuda")
+        idx = torch.tensor([0, 5, 10, 50, 99], dtype=torch.int64, device="cuda")
+        self._check_strided([t], [idx], [64 * 4])
+
+    def test_strided_multi_tensor_different_strides(self):
+        t1 = torch.randn(50, 32, dtype=torch.float16, device="cuda")
+        t2 = torch.randn(80, 64, dtype=torch.float16, device="cuda")
+        idx1 = torch.tensor([0, 10, 49], dtype=torch.int64, device="cuda")
+        idx2 = torch.tensor([30, 79], dtype=torch.int64, device="cuda")
+        self._check_strided([t1, t2], [idx1, idx2], [32 * 2, 64 * 2])
+
+    def test_strided_many_items(self):
+        t = torch.randn(1000, 128, dtype=torch.bfloat16, device="cuda")
+        idx = torch.arange(1000, dtype=torch.int64, device="cuda")
+        self._check_strided([t], [idx], [128 * 2])
+
+    def test_regions(self):
+        first = torch.randint(
+            0, 256, (5 * 1024 * 1024,), dtype=torch.uint8, device="cuda"
+        )
+        second = torch.randint(0, 256, (12345,), dtype=torch.uint8, device="cuda")
+        actual = adler32_regions_checksum(
+            [first.data_ptr(), second.data_ptr()],
+            [first.numel(), second.numel()],
+            first.device,
+        )
+        expected = zlib.adler32(
+            first.cpu().numpy().tobytes() + second.cpu().numpy().tobytes()
+        )
+        self.assertEqual(actual, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/moe/test_deepep_v2_contig_scatter.py

@@ -0,0 +1,115 @@
+"""Tests for the DeepEP v2 contiguous-layout scatter kernel."""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.moe.ep_moe_kernels import ep_scatter_from_psum
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+HIDDEN = 256
+SCALE_HIDDEN = HIDDEN // 128
+# Match the scatter kernel's BLOCK_E alignment.
+ALIGN = 128
+
+
+class TestDeepEPv2ContigScatter(CustomTestCase):
+    """BF16 and FP8 scatter must preserve routed rows and mask invalid experts."""
+
+    # Local expert ids per (token, slot); -1 marks a route to a remote expert,
+    # and 7 is out of this rank's expert range.
+    RECV_TOPK = [
+        [0, -1],
+        [0, -1],
+        [0, 1],
+        [1, 7],
+        [-1, -1],
+    ]
+    NUM_LOCAL_EXPERTS = 2
+
+    def _run(self, dtype, with_scale):
+        num_recv = len(self.RECV_TOPK)
+        recv_x = (
+            torch.arange(num_recv * HIDDEN, dtype=torch.float32, device=DEVICE).reshape(
+                num_recv, HIDDEN
+            )
+            % 100
+        ).to(dtype)
+        recv_topk = torch.tensor(self.RECV_TOPK, dtype=torch.int64, device=DEVICE)
+        psum = torch.tensor(
+            [ALIGN * (e + 1) for e in range(self.NUM_LOCAL_EXPERTS)],
+            dtype=torch.int32,
+            device=DEVICE,
+        )
+        all_tokens = int(psum[-1].item())
+
+        recv_x_scale = None
+        output_tensor_scale = None
+        if with_scale:
+            recv_x_scale = torch.arange(
+                num_recv * SCALE_HIDDEN, dtype=torch.float32, device=DEVICE
+            ).reshape(num_recv, SCALE_HIDDEN)
+            output_tensor_scale = torch.zeros(
+                (all_tokens, SCALE_HIDDEN), dtype=torch.float32, device=DEVICE
+            )
+
+        output_tensor = torch.zeros((all_tokens, HIDDEN), device=DEVICE, dtype=dtype)
+        m_indices = torch.empty(all_tokens, device=DEVICE, dtype=torch.int32)
+        output_index = torch.empty_like(recv_topk)
+        expert_start_loc = torch.empty_like(psum)
+
+        ep_scatter_from_psum(
+            recv_x,
+            recv_x_scale,
+            recv_topk,
+            psum,
+            expert_start_loc,
+            output_tensor,
+            output_tensor_scale,
+            m_indices,
+            output_index,
+        )
+        return recv_x, recv_x_scale, output_tensor, output_tensor_scale, output_index
+
+    def _check(self, dtype, with_scale):
+        recv_x, recv_x_scale, out, out_scale, output_index = self._run(
+            dtype, with_scale
+        )
+        index = output_index.tolist()
+
+        for token, slots in enumerate(self.RECV_TOPK):
+            for slot, expert in enumerate(slots):
+                dest = index[token][slot]
+                if not 0 <= expert < self.NUM_LOCAL_EXPERTS:
+                    # -1 suppresses this route in the post-permute gather.
+                    self.assertEqual(dest, -1, msg=f"{token=} {slot=} {expert=}")
+                    continue
+                self.assertTrue(
+                    ALIGN * expert <= dest < ALIGN * (expert + 1),
+                    msg=f"{token=} {slot=} {expert=} {dest=}",
+                )
+                torch.testing.assert_close(out[dest].float(), recv_x[token].float())
+                if with_scale:
+                    torch.testing.assert_close(out_scale[dest], recv_x_scale[token])
+
+        accepted = [
+            index[t][s]
+            for t, slots in enumerate(self.RECV_TOPK)
+            for s, e in enumerate(slots)
+            if 0 <= e < self.NUM_LOCAL_EXPERTS
+        ]
+        self.assertEqual(len(set(accepted)), len(accepted))
+
+    def test_scatter_bf16_without_scales(self):
+        self._check(torch.bfloat16, with_scale=False)
+
+    def test_scatter_fp8_with_scales(self):
+        self._check(torch.float8_e4m3fn, with_scale=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/moe/test_jit_grouped_topk.py

@@ -0,0 +1,186 @@
+"""The unified Triton router, admitted on ROCm and for single-group routing.
+
+Pins that the router selects what the torch reference selects, and that the
+shared expert appears exactly once -- two places can emit it (this router, or
+_post_process_topk_ids), and if both do, the id is written twice and evicts a
+real routed expert while the model keeps producing plausible logits.
+
+The reference is `biased_grouped_topk_impl`, not `select_experts` with the flag
+off: on ROCm that is the aiter path, which casts the correction bias down to the
+gating dtype, and GLM-5.2 keeps that bias where bf16 cannot separate neighbours.
+It reorders routing on its own.
+"""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=20, suite="jit-kernel-unit-test-amd")
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+
+E, TOPK_ROUTED, SHARED, SCALE = 256, 8, 1, 2.5
+HIDDEN = 512
+
+
+def _jit_routed(monkeypatch, logits, hidden, bias, groups):
+    """Routed ids from select_experts with the unified router on, sorted."""
+    monkeypatch.setenv("SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK", "1")
+    from sglang.srt.layers.moe.topk import TopKConfig, select_experts
+
+    cfg = TopKConfig(
+        top_k=TOPK_ROUTED + SHARED,
+        renormalize=True,
+        use_grouped_topk=True,
+        num_expert_group=groups,
+        num_fused_shared_experts=SHARED,
+        topk_group=1,
+        scoring_func="sigmoid",
+        correction_bias=bias,
+        routed_scaling_factor=SCALE,
+        apply_routed_scaling_factor_on_output=False,
+    )
+    ids = select_experts(hidden, logits, cfg).topk_ids.long()
+    return ids, ids[ids < E].view(ids.shape[0], TOPK_ROUTED).sort(-1).values
+
+
+def _reference_routed(logits, hidden, bias, groups):
+    """Routed ids from the torch reference, sorted.
+
+    It overwrites the last slot with the shared id, so the routed experts are
+    what survives below E.
+    """
+    from sglang.srt.layers.moe.topk import biased_grouped_topk_impl
+
+    _, ids = biased_grouped_topk_impl(
+        hidden_states=hidden,
+        gating_output=logits,
+        correction_bias=bias,
+        topk=TOPK_ROUTED + SHARED,
+        renormalize=True,
+        num_expert_group=groups,
+        topk_group=1,
+        num_fused_shared_experts=SHARED,
+        routed_scaling_factor=SCALE,
+    )
+    ids = ids.long()
+    return ids[ids < E].view(ids.shape[0], TOPK_ROUTED).sort(-1).values
+
+
+def _inputs(tokens, dev="cuda"):
+    torch.manual_seed(0)
+    hidden = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device=dev)
+    # A narrow band at a large offset, as GLM-5.2's bias is: near-equal values
+    # are what a router has to keep apart.
+    bias = (7.0 + 0.04 * torch.randn(E, device=dev)).float()
+    # fp32 logits: in bf16 the two round the sigmoid differently and near-equal
+    # rows would flip for that reason alone.
+    logits = torch.randn(tokens, E, dtype=torch.float32, device=dev)
+    return hidden, bias, logits
+
+
+def _groups(request_groups: int) -> int:
+    """Skip single-group cases where the gate does not admit them.
+
+    biased_grouped_topk_gpu admits one group on ROCm only; CUDA still requires
+    num_expert_group > 1, so a groups=1 case there never reaches the router and
+    the test would assert against a path it did not exercise.
+    """
+    from sglang.srt.layers.moe import topk as topk_mod
+
+    if request_groups == 1 and not topk_mod._is_hip:
+        pytest.skip("single-group routing is admitted on ROCm only")
+    return request_groups
+
+
+@pytest.mark.parametrize("groups", [1, 8])
+@pytest.mark.parametrize("tokens", [6, 48, 256])
+def test_jit_router_selects_what_the_reference_selects(monkeypatch, tokens, groups):
+    groups = _groups(groups)
+    hidden, bias, logits = _inputs(tokens)
+
+    want = _reference_routed(logits, hidden, bias, groups)
+    _, got = _jit_routed(monkeypatch, logits, hidden, bias, groups)
+
+    score = logits.sigmoid() + bias
+    for r in (want != got).any(-1).nonzero().flatten().tolist():
+        only_want = sorted(set(want[r].tolist()) - set(got[r].tolist()))
+        only_got = sorted(set(got[r].tolist()) - set(want[r].tolist()))
+        # Exact ties may break either way; nothing else may.
+        for x, y in zip(only_want, only_got):
+            assert score[r, x].item() == score[r, y].item(), (
+                f"row {r}: reference took {x} (score {score[r, x].item():.9f}) "
+                f"but the router took {y} (score {score[r, y].item():.9f})"
+            )
+
+
+@pytest.mark.parametrize("groups", [1, 8])
+def test_shared_expert_appears_exactly_once(monkeypatch, groups):
+    groups = _groups(groups)
+    hidden, bias, logits = _inputs(48)
+    ids, routed = _jit_routed(monkeypatch, logits, hidden, bias, groups)
+
+    assert ids.shape[-1] == TOPK_ROUTED + SHARED
+    shared = (ids >= E).sum(-1)
+    assert torch.equal(shared, torch.full_like(shared, SHARED)), (
+        f"shared expert appears {shared.tolist()} times a row, expected {SHARED}"
+    )
+    assert (routed[:, 1:] != routed[:, :-1]).all(), "a routed expert repeats"
+
+
+@pytest.mark.parametrize("use_aiter", [True, False])
+def test_router_is_asked_for_the_total_width(monkeypatch, use_aiter):
+    """The width handed to the kernel, for both callers.
+
+    select_experts passes `num_routed_topk if _use_aiter else top_k`; the kernel
+    wants the total either way. Whichever GPU runs the suite fixes `_use_aiter`
+    and can only exercise one half, so pin the arithmetic here.
+    """
+    from sglang.srt.layers.moe import topk as topk_mod
+
+    seen = {}
+
+    class _Captured(Exception):
+        pass
+
+    def _capture(scores, bias, topk, **kwargs):
+        seen["topk"] = topk
+        raise _Captured
+
+    monkeypatch.setenv("SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK", "1")
+    monkeypatch.setattr(topk_mod, "_use_aiter", use_aiter, raising=False)
+    monkeypatch.setattr(
+        "sglang.kernels.ops.moe.moe_fused_gate.moe_fused_gate", _capture
+    )
+
+    hidden, bias, logits = _inputs(4)
+    # The aiter caller hands over routed-only; everyone else the total.
+    topk_in = TOPK_ROUTED if use_aiter else TOPK_ROUTED + SHARED
+    with pytest.raises(_Captured):
+        topk_mod.biased_grouped_topk_gpu(
+            hidden_states=hidden,
+            gating_output=logits,
+            correction_bias=bias,
+            topk=topk_in,
+            renormalize=True,
+            # 8 groups, not 1: the arithmetic under test is the same either
+            # way, and only this value reaches the router on both platforms.
+            num_expert_group=8,
+            topk_group=1,
+            num_fused_shared_experts=SHARED,
+            routed_scaling_factor=SCALE,
+        )
+
+    assert seen["topk"] == TOPK_ROUTED + SHARED, (
+        f"_use_aiter={use_aiter}: caller passed topk={topk_in}, kernel was asked "
+        f"for {seen['topk']} slots, expected {TOPK_ROUTED + SHARED} "
+        f"(routed {TOPK_ROUTED} + shared {SHARED})"
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/qsa/test_qsa_indexer.py

@@ -0,0 +1,273 @@
+"""Fused QSA indexer-prep kernels must match the eager indexer path bit-for-bit,
+up to rare last-ulp RMSNorm flips (see assert_bit_comparable)."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+from sglang.srt.layers.attention.qsa.kernel import (
+    average_pool_qsa_keys,
+    expand_qsa_block_indices,
+    torch_expand_qsa_block_indices,
+)
+from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer
+from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
+
+# MRotaryEmbedding reads the exec config bag at init; publish a minimal
+# process context for the bare pytest process.
+from sglang.srt.runtime_context import publish
+from sglang.srt.server_args import ServerArgs
+
+publish(ServerArgs(model_path="dummy"), role="test")
+
+HEAD_DIM = 128
+NUM_Q_HEADS = 4
+RATIO = 4
+HIDDEN = 2560
+EPS = 1e-6
+
+
+def _make_config():
+    return SimpleNamespace(
+        indexer_n_heads=NUM_Q_HEADS,
+        indexer_kv_heads=1,
+        indexer_head_dim=HEAD_DIM,
+        indexer_budget=2048,
+        indexer_compress_ratio=RATIO,
+        hidden_size=HIDDEN,
+        rms_norm_eps=EPS,
+    )
+
+
+def _make_rotary(mrope_section, mrope_interleaved, device, dtype=torch.bfloat16):
+    return MRotaryEmbedding(
+        head_size=HEAD_DIM,
+        rotary_dim=HEAD_DIM,
+        max_position_embeddings=32768,
+        base=1000000,
+        is_neox_style=True,
+        dtype=dtype,
+        mrope_section=mrope_section,
+        mrope_interleaved=mrope_interleaved,
+    )
+
+
+def _make_indexer(rotary, device, dtype=torch.bfloat16):
+    # Build under the model dtype like ModelRunner does; device-only .to()
+    # afterwards so the fp32 cos_sin_cache buffer keeps its dtype.
+    prev_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        indexer = QSAIndexer(
+            _make_config(), layer_id=0, quant_config=None, rotary_emb=rotary
+        )
+        indexer.to(device=device)
+    finally:
+        torch.set_default_dtype(prev_dtype)
+    with torch.no_grad():
+        out_features = (NUM_Q_HEADS + 1) * HEAD_DIM
+        indexer.index_qk_proj.weight.data.copy_(
+            torch.randn(out_features, HIDDEN, device=device, dtype=dtype) * 0.02
+        )
+        for norm in (indexer.q_layernorm, indexer.k_layernorm):
+            w = torch.randn(HEAD_DIM, device=device, dtype=dtype) * 0.1
+            norm._weight_loader(norm.weight, w)
+    return indexer
+
+
+class FakePool:
+    """Minimal stand-in for the QSA KV pool buffers used by the indexer."""
+
+    index_state_dtype = torch.bfloat16
+
+    def __init__(self, num_slots, num_compressed, device, dtype=torch.bfloat16):
+        self.key_state = torch.zeros(num_slots, 1, HEAD_DIM, dtype=dtype, device=device)
+        self.qsa_rope_position_buffer = torch.zeros(
+            num_slots, 3, dtype=torch.int64, device=device
+        )
+        self.compressed = torch.zeros(
+            num_compressed, 1, HEAD_DIM, dtype=dtype, device=device
+        )
+
+    def get_qsa_key_state_buffer(self, layer_id):
+        return self.key_state
+
+    def set_qsa_key_state_buffer(self, layer_id, loc, token_k):
+        self.key_state[loc.long()] = token_k.to(self.key_state.dtype)
+
+    def set_qsa_rope_position_buffer(self, loc, positions):
+        positions = positions.long()
+        if positions.ndim == 1:
+            positions = positions.unsqueeze(0).expand(3, -1)
+        self.qsa_rope_position_buffer[loc.long()] = positions.transpose(0, 1)
+
+    def get_qsa_rope_position_buffer(self, loc):
+        return self.qsa_rope_position_buffer[loc.long()]
+
+    def get_qsa_compressed_k_buffer(self, layer_id):
+        return self.compressed
+
+    def set_qsa_compressed_k_buffer(self, layer_id, loc, compressed_k):
+        self.compressed[loc.long()] = compressed_k.to(self.compressed.dtype)
+
+
+def assert_bit_comparable(actual, expected, max_frac=1e-5, max_abs=0.02):
+    """Eager RMSNorm (flashinfer CuTe DSL) reduces in an unreproducible order,
+    so ~1 row in 30k flips by 1-2 bf16 ulp; max_frac and max_abs bound that."""
+    diff = (actual.float() - expected.float()).abs()
+    mismatches = int((diff > 0).sum())
+    allowed = max(16, int(max_frac * actual.numel()))
+    assert mismatches <= allowed, f"{mismatches} mismatched elements"
+    if mismatches:
+        peak = diff.max().item()
+        assert peak <= max_abs, f"largest deviation {peak} exceeds {max_abs}"
+
+
+def _eager_compress_reference(indexer, pool, group_locs, write_locs):
+    """The pre-fusion compression chain, via the indexer's own helpers."""
+    key_groups = pool.get_qsa_key_state_buffer(0)[group_locs.long()]
+    pooled = average_pool_qsa_keys(key_groups)
+    rope_positions = indexer._rope_from_matrix(
+        pool.get_qsa_rope_position_buffer(group_locs[:, 0])
+    )
+    normalized = indexer.normalize_compressed_keys(pooled, rope_positions)
+    pool.set_qsa_compressed_k_buffer(0, write_locs, normalized)
+
+
+@pytest.mark.parametrize("num_groups", [1, 5, 2000])
+@pytest.mark.parametrize(
+    "mrope_section, mrope_interleaved",
+    [([24, 20, 20], True), ([24, 20, 20], False), (None, False)],
+)
+def test_fused_compress_matches_eager(num_groups, mrope_section, mrope_interleaved):
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    torch.manual_seed(num_groups)
+    rotary = _make_rotary(mrope_section, mrope_interleaved, device, dtype)
+    indexer = _make_indexer(rotary, device, dtype)
+
+    pool_ref = FakePool(8192, 4096, device, dtype)
+    pool_new = FakePool(8192, 4096, device, dtype)
+    pool_new.key_state.copy_(
+        pool_ref.key_state.copy_(
+            torch.randn(8192, 1, HEAD_DIM, device=device, dtype=dtype)
+        )
+    )
+    positions = torch.randint(0, 30000, (8192, 3), device=device)
+    pool_new.qsa_rope_position_buffer.copy_(positions)
+    pool_ref.qsa_rope_position_buffer.copy_(positions)
+
+    # Random groups; slot 0 doubles as the CUDA-graph dummy write target, so
+    # allow repeats there too.
+    group_locs = torch.randint(0, 8192, (num_groups, RATIO), device=device).to(
+        torch.int32
+    )
+    write_locs = torch.randperm(4096, device=device)[:num_groups].to(torch.int32)
+
+    _eager_compress_reference(indexer, pool_ref, group_locs, write_locs)
+    indexer._fused_compress_store(pool_new, group_locs, write_locs)
+
+    assert_bit_comparable(pool_new.compressed, pool_ref.compressed)
+
+
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_expand_block_indices_int_inputs(dtype):
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    rows, block_topk, token_topk, ratio = 37, 512, 2048, 4
+    query_positions = torch.randint(0, 8000, (rows,), dtype=dtype, device=device)
+    sequence_lengths = (
+        query_positions + torch.randint(1, 9, (rows,), dtype=dtype, device=device)
+    ).to(dtype)
+    # Production contract: top-k only selects blocks inside [0, seq_len//4),
+    # so no selected block ever masks out against sequence_lengths.
+    counts = torch.randint(0, block_topk + 1, (rows,))
+    block_indices = torch.full((rows, block_topk), -1, dtype=torch.int32)
+    seq_lens_host = sequence_lengths.cpu()
+    for r in range(rows):
+        limit = max(int(seq_lens_host[r]) // ratio, 1)
+        count = min(int(counts[r]), limit)
+        if count:
+            block_indices[r, :count] = torch.randperm(limit)[:count].to(torch.int32)
+    block_indices = block_indices.to(device)
+    out = expand_qsa_block_indices(
+        block_indices, query_positions, sequence_lengths, ratio, token_topk
+    )
+    ref = torch_expand_qsa_block_indices(
+        block_indices.cpu(),
+        query_positions.cpu(),
+        sequence_lengths.cpu(),
+        ratio,
+        token_topk,
+    )
+    assert torch.equal(out.cpu(), ref)
+
+
+def test_decode_selection_equivalent():
+    """Last-ulp norm flips must not change the selected blocks:
+    scores are fp32 sums of 128-dim dots, so a 1-ulp flip only matters on exact ties."""
+    from sglang.srt.layers.attention.qsa.kernel import qsa_fast_topk
+    from sglang.srt.layers.attention.qsa.mqa import torch_qsa_mqa_decode
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    torch.manual_seed(7)
+    rotary = _make_rotary([24, 20, 20], True, device, dtype)
+    indexer = _make_indexer(rotary, device, dtype)
+
+    batch, max_pages, page_size = 4, 32, 64
+    max_model_len = max_pages * page_size
+    hidden = torch.randn(batch, HIDDEN, device=device, dtype=dtype)
+    positions = (
+        torch.arange(8000, 8000 + batch, device=device)
+        .unsqueeze(0)
+        .expand(3, -1)
+        .contiguous()
+    )
+    qk, _ = indexer.index_qk_proj(hidden)
+
+    # Eager index q.
+    q_ref = indexer.q_layernorm(qk[:, : NUM_Q_HEADS * HEAD_DIM].reshape(-1, HEAD_DIM))
+    q_ref = q_ref.reshape(batch, NUM_Q_HEADS, HEAD_DIM)
+    q_ref = indexer.apply_rope(positions, q_ref)
+
+    # Fused index q.
+    pool = FakePool(64, 4096, device, dtype)
+    cache_loc = torch.arange(1, batch + 1, device=device)
+    q_new, _, stored = indexer.project_qk(
+        hidden, positions, pool=pool, cache_loc=cache_loc
+    )
+    assert stored
+
+    compressed_cache = torch.randn(
+        64, page_size, 1, HEAD_DIM, device=device, dtype=dtype
+    )
+    page_table = torch.arange(max_pages, dtype=torch.int32, device=device).repeat(
+        batch, 1
+    )
+    context_lens = torch.full((batch,), 1500, dtype=torch.int32, device=device)
+
+    def select(q):
+        logits = torch_qsa_mqa_decode(
+            q, compressed_cache, page_table, context_lens, max_model_len
+        )
+        row_starts = torch.zeros_like(context_lens)
+        return qsa_fast_topk(logits, row_starts, context_lens, topk=512)
+
+    idx_ref = select(q_ref)
+    idx_new = select(q_new[:, :NUM_Q_HEADS].contiguous())
+    for row in range(batch):
+        ref_set = set(idx_ref[row][idx_ref[row] >= 0].tolist())
+        new_set = set(idx_new[row][idx_new[row] >= 0].tolist())
+        assert ref_set == new_set, f"row {row}: selection mismatch"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/qsa/test_qsa_strided_zero_fill.py

@@ -0,0 +1,192 @@
+"""Regression test for the QSA strided sparse-decode scratch zero-fill.
+
+Poison the packed scratch with NaN, gather with the strided layout used by
+`_forward_trtllm_sparse`, and require that (a) valid rows are copied exactly and
+(b) every slot in [valid_count, stride) is zero, so the paged decode kernel can never
+multiply masked probabilities into stale NaN/Inf bytes. Also checks the compact
+(FA2 fallback) layout is unchanged. Intended for test/registered/kernels/ops/qsa/.
+"""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+from sglang.srt.layers.attention.qsa.sparse_attn import (
+    qwen_sparse_fa2_cu_seqlens_triton,
+    qwen_sparse_kv_extraction_compact_triton,
+)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+def test_strided_gather_zero_fills_tail(dtype):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    batch, topk, page, heads, dim = 3, 2051, 64, 2, 256
+    pages_per_row = (topk + page - 1) // page
+    stride = pages_per_row * page
+    pool_rows = 8192
+    k_pool = torch.randn(pool_rows, heads, dim, device=device, dtype=torch.bfloat16).to(
+        dtype
+    )
+    v_pool = torch.randn(pool_rows, heads, dim, device=device, dtype=torch.bfloat16).to(
+        dtype
+    )
+    seq_lens = torch.tensor([733, 109, 2500], device=device, dtype=torch.int32)
+    req_to_token = (
+        torch.randperm(pool_rows, device=device)[: batch * 2600]
+        .reshape(batch, 2600)
+        .to(torch.int32)
+    )
+    req_indices = torch.arange(batch, device=device, dtype=torch.int32)
+    # top-k rows: the first min(seq_len, topk) logical positions, then -1 padding
+    indices = torch.full((batch, topk), -1, device=device, dtype=torch.int32)
+    for b in range(batch):
+        n = min(int(seq_lens[b]), topk)
+        indices[b, :n] = torch.arange(n, device=device, dtype=torch.int32)
+    cu_strided = torch.arange(batch + 1, device=device, dtype=torch.int32) * stride
+    # the scratch is always in the compute dtype (bf16); an FP8 pool is dequantized on the way in
+    packed_k = torch.full(
+        (batch * stride, heads, dim), float("nan"), device=device, dtype=torch.bfloat16
+    )
+    packed_v = packed_k.clone()
+
+    qwen_sparse_kv_extraction_compact_triton(
+        k_pool,
+        v_pool,
+        req_to_token,
+        req_indices,
+        indices,
+        seq_lens,
+        cu_strided,
+        packed_k,
+        packed_v,
+        batch,
+        topk,
+        zero_fill_cols=stride,
+    )
+    pk, pv = (
+        packed_k.float().view(batch, stride, heads, dim),
+        packed_v.float().view(batch, stride, heads, dim),
+    )
+    assert torch.isfinite(pk).all() and torch.isfinite(pv).all()
+    for b in range(batch):
+        n = min(int(seq_lens[b]), topk)
+        slots = req_to_token[b, :n].long()
+        torch.testing.assert_close(pk[b, :n], k_pool[slots].to(torch.bfloat16).float())
+        torch.testing.assert_close(pv[b, :n], v_pool[slots].to(torch.bfloat16).float())
+        assert (pk[b, n:] == 0).all() and (pv[b, n:] == 0).all()
+
+
+def test_compact_gather_unchanged():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    batch, topk, heads, dim = 2, 2051, 2, 256
+    k_pool = torch.randn(4096, heads, dim, device=device, dtype=torch.bfloat16)
+    v_pool = torch.randn(4096, heads, dim, device=device, dtype=torch.bfloat16)
+    seq_lens = torch.tensor([300, 50], device=device, dtype=torch.int32)
+    req_to_token = torch.arange(batch * 512, device=device, dtype=torch.int32).reshape(
+        batch, 512
+    )
+    req_indices = torch.arange(batch, device=device, dtype=torch.int32)
+    indices = torch.full((batch, topk), -1, device=device, dtype=torch.int32)
+    for b in range(batch):
+        indices[b, : int(seq_lens[b])] = torch.arange(
+            int(seq_lens[b]), device=device, dtype=torch.int32
+        )
+    counts = torch.empty(batch, device=device, dtype=torch.int32)
+    cu_k = torch.empty(batch + 1, device=device, dtype=torch.int32)
+    qwen_sparse_fa2_cu_seqlens_triton(seq_lens, indices, counts, cu_k, batch, topk)
+    assert cu_k.tolist() == [0, 300, 350]
+    packed_k = torch.full(
+        (batch * topk, heads, dim), float("nan"), device=device, dtype=torch.bfloat16
+    )
+    packed_v = packed_k.clone()
+    qwen_sparse_kv_extraction_compact_triton(
+        k_pool,
+        v_pool,
+        req_to_token,
+        req_indices,
+        indices,
+        seq_lens,
+        cu_k,
+        packed_k,
+        packed_v,
+        batch,
+        topk,
+    )
+    torch.testing.assert_close(packed_k[:300], k_pool[req_to_token[0, :300].long()])
+    torch.testing.assert_close(packed_k[300:350], k_pool[req_to_token[1, :50].long()])
+    # compact layout leaves the region past the packed rows untouched (still NaN)
+    assert torch.isnan(packed_k[350:]).all()
+
+
+def test_strided_gather_addresses_pool_beyond_int32_elements():
+    """Slots past 2^31 / (heads * dim) must be addressed with 64-bit offsets.
+
+    An FP8 KV pool on one GB300 holds ~7.6M tokens for Qwen3.8-Flash-Next (2 kv heads x 256),
+    so slot indices above 4,194,304 occur in production; int32 element offsets wrap there.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if torch.cuda.get_device_properties(0).total_memory < 6 * 1024**3:
+        pytest.skip("needs ~2.5 GB of device memory")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    heads, dim = 2, 256
+    threshold = (1 << 31) // (heads * dim)  # 4,194,304
+    pool_rows = threshold + 4096
+    k_pool = torch.zeros(
+        pool_rows, heads, dim, device=device, dtype=torch.float8_e4m3fn
+    )
+    v_pool = torch.zeros(
+        pool_rows, heads, dim, device=device, dtype=torch.float8_e4m3fn
+    )
+    hi = torch.arange(threshold + 64, threshold + 64 + 300, device=device)
+    k_pool[hi] = torch.randn(300, heads, dim, device=device, dtype=torch.bfloat16).to(
+        torch.float8_e4m3fn
+    )
+    v_pool[hi] = torch.randn(300, heads, dim, device=device, dtype=torch.bfloat16).to(
+        torch.float8_e4m3fn
+    )
+    batch, topk, page = 1, 2051, 64
+    stride = ((topk + page - 1) // page) * page
+    seq_lens = torch.tensor([300], device=device, dtype=torch.int32)
+    req_to_token = torch.zeros(batch, 512, device=device, dtype=torch.int32)
+    req_to_token[0, :300] = hi.to(torch.int32)
+    indices = torch.full((batch, topk), -1, device=device, dtype=torch.int32)
+    indices[0, :300] = torch.arange(300, device=device, dtype=torch.int32)
+    cu_strided = torch.arange(batch + 1, device=device, dtype=torch.int32) * stride
+    packed_k = torch.full(
+        (batch * stride, heads, dim), float("nan"), device=device, dtype=torch.bfloat16
+    )
+    packed_v = packed_k.clone()
+    qwen_sparse_kv_extraction_compact_triton(
+        k_pool,
+        v_pool,
+        req_to_token,
+        torch.zeros(1, device=device, dtype=torch.int32),
+        indices,
+        seq_lens,
+        cu_strided,
+        packed_k,
+        packed_v,
+        batch,
+        topk,
+        zero_fill_cols=stride,
+    )
+    torch.testing.assert_close(packed_k[:300], k_pool[hi].to(torch.bfloat16))
+    torch.testing.assert_close(packed_v[:300], v_pool[hi].to(torch.bfloat16))
+    assert (packed_k[300:] == 0).all() and (packed_v[300:] == 0).all()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/quantization/test_fp8_utils_aiter.py

@@ -0,0 +1,50 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.fp8_utils import unshuffle_aiter_fp8_weight
+from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=5, stage="jit-kernel-unit", runner_config="amd")
+
+
+@unittest.skipUnless(is_hip(), "requires ROCm AITER")
+class TestAiterFp8Utils(CustomTestCase):
+    def test_unshuffle_weight_round_trip(self):
+        from aiter.ops.shuffle import shuffle_weight
+
+        for shape in ((32, 64), (2, 32, 64)):
+            with self.subTest(shape=shape):
+                logical = (
+                    torch.arange(
+                        torch.Size(shape).numel(), device="cuda", dtype=torch.float32
+                    )
+                    .remainder(7)
+                    .to(torch.float8_e4m3fn)
+                    .reshape(shape)
+                )
+                shuffled = shuffle_weight(logical, layout=(16, 16))
+
+                torch.testing.assert_close(
+                    unshuffle_aiter_fp8_weight(shuffled), logical
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/speculative/test_dflash_domino.py

@@ -0,0 +1,125 @@
+import unittest
+
+import torch
+from torch import nn
+
+from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler
+from sglang.srt.speculative.domino_utils import _domino_gru_cell, domino_greedy_rollout
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
+class TestDFlashDominoRollout(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+        self.embedding = nn.Embedding(31, 8, device="cuda", dtype=torch.bfloat16)
+        self.prefix_gru = nn.GRU(8, 4, batch_first=True, bias=False).cuda().bfloat16()
+        self.embed_proj = (
+            nn.Sequential(
+                nn.Linear(12, 5, bias=False), nn.SiLU(), nn.Linear(5, 31, bias=False)
+            )
+            .cuda()
+            .bfloat16()
+        )
+        self.lm_head_weight = torch.randn(31, 8, device="cuda", dtype=torch.bfloat16)
+        self.hidden = torch.randn(3, 16, 8, device="cuda", dtype=torch.bfloat16)
+        self.bonus_tokens = torch.tensor([1, 4, 9], device="cuda")
+
+    def rollout(self, hidden, bonus_tokens, pool_size=5, shift_label=True):
+        return domino_greedy_rollout(
+            draft_hidden=hidden,
+            bonus_tokens=bonus_tokens,
+            target_embedding=self.embedding,
+            lm_head_weight=self.lm_head_weight,
+            prefix_gru=self.prefix_gru,
+            embed_proj=self.embed_proj,
+            vocab_size=31,
+            shift_label=shift_label,
+            candidate_pool_size=pool_size,
+        )
+
+    def test_gru_feedback_matches_sequence(self):
+        embeddings = self.embedding(torch.tensor([[1, 2, 3], [4, 5, 6]], device="cuda"))
+        _, expected = self.prefix_gru(embeddings)
+        state = torch.zeros(2, 4, device="cuda", dtype=torch.bfloat16)
+        for step in embeddings.unbind(dim=1):
+            state = _domino_gru_cell(self.prefix_gru, step, state)
+        torch.testing.assert_close(state, expected[0], rtol=0.02, atol=0.002)
+
+    def test_candidate_pool_boundaries(self):
+        for shift_label in (True, False):
+            with self.subTest(shift_label=shift_label):
+                full = self.rollout(self.hidden, self.bonus_tokens, 0, shift_label)
+                for pool_size in (31, 32):
+                    actual = self.rollout(
+                        self.hidden, self.bonus_tokens, pool_size, shift_label
+                    )
+                    torch.testing.assert_close(actual, full, rtol=0, atol=0)
+                first_hidden = self.hidden[:, 0 if shift_label else 1]
+                expected_first = (first_hidden @ self.lm_head_weight.T).argmax(dim=-1)
+                for block_size in (2, 16):
+                    limited = self.rollout(
+                        self.hidden[:, :block_size], self.bonus_tokens, 1, shift_label
+                    )
+                    self.assertEqual(limited.shape, (3, block_size - 1))
+                    torch.testing.assert_close(limited[:, 0], expected_first)
+                    if block_size > 2:
+                        torch.testing.assert_close(
+                            limited[:, 1:], limited[:, 1:2].expand_as(limited[:, 1:])
+                        )
+
+    def test_batch_matches_individual_requests(self):
+        for pool_size in (0, 5):
+            with self.subTest(pool_size=pool_size):
+                batched = self.rollout(self.hidden, self.bonus_tokens, pool_size)
+                individual = torch.cat(
+                    [
+                        self.rollout(hidden[None], bonus[None], pool_size)
+                        for hidden, bonus in zip(self.hidden, self.bonus_tokens)
+                    ]
+                )
+                torch.testing.assert_close(batched, individual, rtol=0, atol=0)
+
+    def test_sampler_replays_with_new_inputs(self):
+        sampler = _DominoDraftSampler(
+            target_embedding=self.embedding,
+            lm_head_weight=self.lm_head_weight,
+            prefix_gru=self.prefix_gru,
+            embed_proj=self.embed_proj,
+            vocab_size=31,
+            block_size=16,
+            shift_label=True,
+            max_bs=3,
+            candidate_pool_size=5,
+        )
+        block_ids = torch.zeros(3, 16, device="cuda", dtype=torch.long)
+        block_ids[:, 0].copy_(self.bonus_tokens)
+
+        def sample():
+            sampler(self.hidden.flatten(0, 1), block_ids.flatten())
+
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            sample()
+        torch.cuda.current_stream().wait_stream(warmup)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            sample()
+
+        for _ in range(2):
+            self.hidden.copy_(torch.randn_like(self.hidden))
+            block_ids[:, 0].copy_(torch.randint(31, (3,), device="cuda"))
+            expected = self.rollout(self.hidden, block_ids[:, 0])
+            graph.replay()
+            torch.cuda.synchronize()
+            torch.testing.assert_close(
+                sampler.out.view(3, 15), expected, rtol=0, atol=0
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/speculative/test_spec_kv_indices_grid.py

@@ -0,0 +1,195 @@
+"""2D (token-block) launches of the KV index-copy kernels must be
+bit-identical to the historical 1D launches, with no unwritten or
+overwritten bytes (sentinel-checked over the full buffers), and the draft
+kernel's output must match a Python reference."""
+
+import unittest
+
+import torch
+import triton
+
+from sglang.kernels.ops.attention.utils import (
+    create_flashinfer_kv_indices_triton,
+    kv_indices_num_token_blocks,
+)
+from sglang.kernels.ops.speculative.cache_locs import (
+    generate_draft_decode_kv_indices,
+)
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=25, stage="jit-kernel-unit", runner_config="amd")
+
+SENTINEL = 0x7EADBEEF
+POOL_LEN = 262_144
+LENSETS = [
+    [0, 1, 511, 512, 513, 8191, 8192, 8193],
+    [100_000, 33, 4096],
+    [1000] * 7 + [100_000],
+]
+
+
+def _npo2(x: int) -> int:
+    return max(1, 1 << (max(1, x) - 1).bit_length())
+
+
+def _draft_inputs(seqs, topk, steps, device, idx_dtype=torch.int64):
+    bs = len(seqs)
+    req_pool = torch.arange(bs, dtype=idx_dtype, device=device)
+    r2t = torch.randint(
+        0, 6_000_000, (bs + 1, POOL_LEN), dtype=torch.int32, device=device
+    )
+    lens = torch.tensor(seqs, dtype=idx_dtype, device=device)
+    width = topk * (max(seqs) + steps) + 64
+    return bs, req_pool, r2t, lens, width, bs * topk
+
+
+def _run_draft(kern_inputs, topk, steps, page_size, nb, kw):
+    bs, req_pool, r2t, lens, width, tot = kern_inputs
+    dev = lens.device
+    kv_i = torch.full((steps, bs * width), SENTINEL, dtype=torch.int32, device=dev)
+    kv_p = torch.full((steps, tot + 1), SENTINEL, dtype=torch.int32, device=dev)
+    # positions is per draft token in production (bs * topk entries); the
+    # kernel reads positions[:bs * topk] for the kv_indptr prefix sums.
+    positions = torch.repeat_interleave(lens, topk)
+    generate_draft_decode_kv_indices[(steps * nb, bs, topk)](
+        req_pool,
+        r2t,
+        lens,
+        kv_i,
+        kv_p,
+        positions,
+        POOL_LEN,
+        kv_i.shape[1],
+        kv_p.shape[1],
+        _npo2(bs),
+        _npo2(steps),
+        _npo2(tot),
+        page_size,
+        **kw,
+    )
+    torch.cuda.synchronize()
+    return kv_i, kv_p
+
+
+class TestSpecKvIndicesGrid(CustomTestCase):
+    def test_draft_grid_equivalence(self):
+        torch.manual_seed(0)
+        for seqs in LENSETS:
+            for topk, page_size in [(1, 1), (4, 1), (4, 16)]:
+                for steps, idx_dtype in [
+                    (2, torch.int64),
+                    (3, torch.int32),
+                    (4, torch.int64),
+                ]:
+                    inputs = _draft_inputs(seqs, topk, steps, "cuda", idx_dtype)
+                    ref = _run_draft(inputs, topk, steps, page_size, 1, {})
+                    for nb in [
+                        1,
+                        kv_indices_num_token_blocks(POOL_LEN, steps * len(seqs) * topk),
+                        triton.cdiv(POOL_LEN, 8192),
+                    ]:
+                        out = _run_draft(
+                            inputs, topk, steps, page_size, nb, {"NUM_STEPS": steps}
+                        )
+                        self.assertTrue(torch.equal(ref[0], out[0]), (seqs, topk, nb))
+                        self.assertTrue(torch.equal(ref[1], out[1]), (seqs, topk, nb))
+
+    def test_draft_reference(self):
+        torch.manual_seed(1)
+        seqs, steps = [100_000, 33, 4096, 16], 3
+        for topk, page_size in [(1, 1), (4, 1), (4, 16)]:
+            inputs = _draft_inputs(seqs, topk, steps, "cuda")
+            bs, _, r2t, _, _, tot = inputs
+            nb = kv_indices_num_token_blocks(POOL_LEN, steps * bs * topk)
+            kv_i, kv_p = _run_draft(
+                inputs, topk, steps, page_size, nb, {"NUM_STEPS": steps}
+            )
+            for it in range(steps):
+                iters = it + 1
+                for s in range(bs):
+                    ln = seqs[s]
+                    for k in range(topk):
+                        off = sum(seqs[:s]) * topk + s * iters * topk + k * (ln + iters)
+                        self.assertTrue(
+                            torch.equal(r2t[s, :ln], kv_i[it, off : off + ln]),
+                            (it, s, k, topk, page_size),
+                        )
+                        if page_size == 1 or topk == 1:
+                            src = ln + k * steps
+                        else:
+                            last = ln % page_size
+                            pages = -(-(last + steps) // page_size)
+                            src = (
+                                ln // page_size * page_size
+                                + k * pages * page_size
+                                + last
+                            )
+                        self.assertTrue(
+                            torch.equal(
+                                r2t[s, src : src + iters],
+                                kv_i[it, off + ln : off + ln + iters],
+                            ),
+                            (it, s, k, topk, page_size),
+                        )
+                positions = [ln for ln in seqs for _ in range(topk)]
+                for z in range(1, tot + 1):
+                    self.assertEqual(
+                        int(kv_p[it, z]),
+                        sum(positions[:z]) + z * iters,
+                        (it, z, topk, page_size),
+                    )
+
+    def test_flat_grid_equivalence(self):
+        torch.manual_seed(0)
+        dev = "cuda"
+        for seqs in LENSETS:
+            for use_start, entry_page_size in [(False, 1), (True, 1), (False, 16)]:
+                bs = len(seqs)
+                req_pool = torch.arange(bs, dtype=torch.int64, device=dev)
+                r2t = torch.randint(
+                    0, 6_000_000, (bs + 1, POOL_LEN), dtype=torch.int32, device=dev
+                )
+                lens = torch.tensor(
+                    seqs,
+                    dtype=torch.int32 if use_start else torch.int64,
+                    device=dev,
+                )
+                indptr = torch.zeros(bs + 1, dtype=torch.int32, device=dev)
+                indptr[1:] = torch.cumsum(lens, 0).to(torch.int32)
+                start = (
+                    torch.full((bs,), 7, dtype=torch.int32, device=dev)
+                    if use_start
+                    else None
+                )
+                n = int(indptr[-1]) + (7 * bs if use_start else 0) + 8
+                outs = []
+                for grid, parallel in [
+                    ((bs,), False),
+                    ((bs, 1), True),
+                    ((bs, kv_indices_num_token_blocks(POOL_LEN, bs)), True),
+                    ((bs, triton.cdiv(POOL_LEN, 8192)), True),
+                ]:
+                    kv_i = torch.full((n,), SENTINEL, dtype=torch.int32, device=dev)
+                    create_flashinfer_kv_indices_triton[grid](
+                        r2t,
+                        req_pool,
+                        lens,
+                        indptr,
+                        start,
+                        kv_i,
+                        r2t.shape[1],
+                        ENTRY_PAGE_SIZE=entry_page_size,
+                        TOKEN_BLOCK_PARALLEL=parallel,
+                    )
+                    torch.cuda.synchronize()
+                    outs.append(kv_i)
+                for kv_i in outs[1:]:
+                    self.assertTrue(
+                        torch.equal(outs[0], kv_i), (seqs, use_start, entry_page_size)
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/speculative/test_verify_commit_triton.py

@@ -0,0 +1,81 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+    fused_commit_track_indices,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def _reference(accept_index, accept_lens, seq_lens, draft_token_num, track_interval):
+    """Mirrors the eager branch of spec_utils._verify_commit_step_indices."""
+    bs = accept_lens.shape[0]
+    offset = torch.arange(
+        0,
+        bs * draft_token_num,
+        step=draft_token_num,
+        dtype=accept_lens.dtype,
+        device=accept_lens.device,
+    )
+    req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device)
+    last = accept_index[req_idx, (accept_lens - 1).to(torch.int64)] - offset
+    if track_interval <= 0:
+        return last, None
+    pre = seq_lens
+    post = seq_lens + accept_lens
+    mask = pre // track_interval != post // track_interval
+    point = post // track_interval * track_interval
+    ith = torch.clamp(point - pre - 1, min=0).to(torch.int64)
+    cand = accept_index[req_idx, ith] - offset
+    track = torch.where(mask, cand, torch.full_like(cand, -1))
+    return last, track
+
+
+@pytest.mark.parametrize("bs", [1, 3, 48, 257])
+@pytest.mark.parametrize("track_interval", [0, 64])
+@pytest.mark.parametrize("tree_depth", [4, 3])
+def test_verify_commit_steps_matches_eager(bs, track_interval, tree_depth):
+    """The fused kernel must match eager on both outputs near tracking boundaries
+    and when accept_index rows (max_tree_depth) are narrower than draft_token_num."""
+    if not torch.cuda.is_available():
+        pytest.skip("needs CUDA")
+    torch.manual_seed(bs + track_interval + tree_depth)
+    device = "cuda"
+    draft_token_num = 4
+    accept_lens = torch.randint(
+        1, tree_depth + 1, (bs,), device=device, dtype=torch.int32
+    )
+    tree_nodes = torch.argsort(torch.rand(bs, draft_token_num, device=device), dim=1)[
+        :, :tree_depth
+    ]
+    accept_index = (
+        torch.arange(bs, device=device, dtype=torch.int64).unsqueeze(1)
+        * draft_token_num
+        + tree_nodes
+    ).to(torch.int32)
+    # Cluster seq lens around tracking boundaries to exercise the crossing.
+    seq_lens = torch.randint(60, 70, (bs,), device=device, dtype=torch.int64)
+
+    exp_last, exp_track = _reference(
+        accept_index, accept_lens, seq_lens, draft_token_num, track_interval
+    )
+    got_last, got_track = fused_commit_track_indices(
+        accept_index,
+        accept_lens,
+        seq_lens if track_interval > 0 else None,
+        draft_token_num,
+        track_interval,
+    )
+    assert torch.equal(got_last, exp_last)
+    if track_interval > 0:
+        assert torch.equal(got_track, exp_track)
+    else:
+        assert got_track is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))