chore: import upstream snapshot with attribution

python/sglang/jit_kernel/tests/deepseek_v4/common.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Literal, Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.dsv4 import CompressorDecodePlan, CompressorPrefillPlan
+
+
+@dataclass
+class LegacyContext:
+    """Per-request ring buffer (no req_to_token / full_to_swa).
+
+    `req_pool_indices[i]` directly maps to the request's ring base slot.
+    """
+
+    bs: int
+    head_dim: int
+    compress_ratio: int
+    req_pool_indices: torch.Tensor  # int64 [bs] on cuda
+    pages_per_req: int
+
+    @property
+    def num_pages(self) -> int:
+        # Reserve enough pages to hold all batched requests' rings.
+        return int(self.req_pool_indices.max().item() + 1) * self.pages_per_req
+
+    def state_loc(self, b: int, position: int) -> int:
+        rid = int(self.req_pool_indices[b].item())
+        if self.compress_ratio == 4:
+            page = rid * 2 + (position // 4) % 2
+        else:
+            page = rid
+        return page * self.compress_ratio + position % self.compress_ratio
+
+    def make_prefill_plan(
+        self,
+        seq_lens_cpu: torch.Tensor,
+        extend_lens_cpu: torch.Tensor,
+        num_q_tokens: int,
+    ) -> CompressorPrefillPlan:
+        return CompressorPrefillPlan.generate_legacy(
+            compress_ratio=self.compress_ratio,  # type: ignore
+            req_pool_indices=self.req_pool_indices,
+            seq_lens=seq_lens_cpu,
+            extend_lens=extend_lens_cpu,
+            num_q_tokens=num_q_tokens,
+            device=torch.device("cuda"),
+        )
+
+    def make_decode_plan(self, seq_lens_gpu: torch.Tensor) -> CompressorDecodePlan:
+        return CompressorDecodePlan.generate_legacy(
+            compress_ratio=self.compress_ratio,  # type: ignore
+            req_pool_indices=self.req_pool_indices,
+            seq_lens=seq_lens_gpu,
+        )
+
+
+@dataclass
+class PagedContext:
+    """SWA paged layout with identity req_to_token + identity full_to_swa.
+
+    Each request occupies `num_swa_pages_per_req` contiguous swa_pages, so
+    `req_to_token[r, p] = r * (num_swa_pages_per_req * swa_page_size) + p`.
+    """
+
+    bs: int
+    head_dim: int
+    compress_ratio: int
+    swa_page_size: int
+    ring_size: int
+    num_swa_pages_per_req: int
+    req_pool_indices: torch.Tensor  # int64 [bs] on cuda
+    req_to_token: torch.Tensor  # int64 [num_reqs_capacity, max_tokens_per_req] on cuda
+    full_to_swa: torch.Tensor  # int64 [num_swa_slots] on cuda
+
+    @property
+    def num_pages(self) -> int:
+        # Upper bound: every (request, position) state slot fits.
+        max_state_loc = (
+            self.bs * self.num_swa_pages_per_req * self.ring_size
+            + self.swa_page_size  # slack for the largest tail
+        )
+        return max_state_loc // self.compress_ratio + 1
+
+    def state_loc(self, b: int, position: int) -> int:
+        rid = int(self.req_pool_indices[b].item())
+        loc = int(self.req_to_token[rid, position].item())
+        swa_loc = int(self.full_to_swa[loc].item())
+        swa_page = swa_loc // self.swa_page_size
+        return swa_page * self.ring_size + swa_loc % self.ring_size
+
+    def make_prefill_plan(
+        self,
+        seq_lens_cpu: torch.Tensor,
+        extend_lens_cpu: torch.Tensor,
+        num_q_tokens: int,
+    ) -> CompressorPrefillPlan:
+        return CompressorPrefillPlan.generate(
+            compress_ratio=self.compress_ratio,  # type: ignore
+            req_pool_indices=self.req_pool_indices,
+            seq_lens=seq_lens_cpu,
+            extend_lens=extend_lens_cpu,
+            req_to_token=self.req_to_token,
+            full_to_state=self.full_to_swa,
+            swa_page_size=self.swa_page_size,
+            ring_size=self.ring_size,
+            num_q_tokens=num_q_tokens,
+        )
+
+    def make_decode_plan(self, seq_lens_gpu: torch.Tensor) -> CompressorDecodePlan:
+        return CompressorDecodePlan.generate(
+            compress_ratio=self.compress_ratio,  # type: ignore
+            req_pool_indices=self.req_pool_indices,
+            req_to_token=self.req_to_token,
+            full_to_state=self.full_to_swa,
+            seq_lens=seq_lens_gpu,
+            swa_page_size=self.swa_page_size,
+            ring_size=self.ring_size,
+        )
+
+
+def make_legacy_context(
+    bs: int,
+    compress_ratio: Literal[4, 128],
+    head_dim: int = 512,
+) -> LegacyContext:
+    pages_per_req = 2 if compress_ratio == 4 else 1
+    req_pool_indices = torch.arange(bs, dtype=torch.int64, device="cuda")
+    return LegacyContext(
+        bs=bs,
+        head_dim=head_dim,
+        compress_ratio=compress_ratio,
+        req_pool_indices=req_pool_indices,
+        pages_per_req=pages_per_req,
+    )
+
+
+def make_paged_context(
+    bs: int,
+    compress_ratio: Literal[4, 128],
+    head_dim: int = 512,
+    swa_page_size: int = 256,
+    ring_size: Optional[int] = None,
+    num_swa_pages_per_req: int = 8,
+    max_tokens_per_req: int = 8192,
+    num_reqs_capacity: int = 16,
+) -> PagedContext:
+    if ring_size is None:
+        ring_size = 8 if compress_ratio == 4 else 128
+    assert swa_page_size % ring_size == 0
+    assert ring_size % compress_ratio == 0
+    assert num_swa_pages_per_req * swa_page_size <= max_tokens_per_req
+
+    stride = num_swa_pages_per_req * swa_page_size
+    req_to_token = torch.zeros(
+        (num_reqs_capacity, max_tokens_per_req), dtype=torch.int32
+    )
+    for r in range(bs):
+        req_to_token[r, :stride] = torch.arange(
+            r * stride, (r + 1) * stride, dtype=torch.int32
+        )
+    total_swa_slots = num_reqs_capacity * stride
+    full_to_swa = torch.arange(total_swa_slots, dtype=torch.int64)
+    req_pool_indices = torch.arange(bs, dtype=torch.int64)
+    return PagedContext(
+        bs=bs,
+        head_dim=head_dim,
+        compress_ratio=compress_ratio,
+        swa_page_size=swa_page_size,
+        ring_size=ring_size,
+        num_swa_pages_per_req=num_swa_pages_per_req,
+        req_pool_indices=req_pool_indices.cuda(),
+        req_to_token=req_to_token.cuda(),
+        full_to_swa=full_to_swa.cuda(),
+    )
+
+
+def make_state_pool(num_pages: int, compress_ratio: int, head_dim: int) -> torch.Tensor:
+    last_dim = head_dim * (4 if compress_ratio == 4 else 2)
+    return torch.zeros(
+        (num_pages, compress_ratio, last_dim),
+        dtype=torch.float32,
+        device="cuda",
+    )
+
+
+def to_seq_extend(
+    seq_extend_pairs: List[Tuple[int, int]],
+) -> Tuple[torch.Tensor, torch.Tensor, int]:
+    seq_lens = torch.tensor([s for s, _ in seq_extend_pairs], dtype=torch.int64)
+    extend_lens = torch.tensor([e for _, e in seq_extend_pairs], dtype=torch.int64)
+    num_q = int(extend_lens.sum().item())
+    return seq_lens, extend_lens, num_q

python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py

@@ -0,0 +1,344 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
+from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.jit_kernel.tests.kv_canary._constants import (
+    _I64_SIGN_BIT,
+    _U64_MASK,
+    DEFAULT_NUM_SLOTS,
+    DEFAULT_RING_CAPACITY,
+    DEFAULT_SLOT_STRIDE_BYTES,
+)
+from sglang.jit_kernel.tests.kv_canary._fixtures import (
+    make_real_kv_source,
+    make_real_kv_sources,
+)
+
+__all__ = [
+    "FakeViolationLog",
+    "assert_canary_buf_equal",
+    "assert_canary_state_equal",
+    "assert_only_bits_set",
+    "chain_anchor_signed",
+    "make_canary_buf",
+    "make_canary_buf_pair",
+    "make_log_pair",
+    "make_real_kv_source",
+    "make_real_kv_sources",
+    "make_verify_plan",
+    "make_verify_plan_pair",
+    "make_write_plan",
+    "make_write_plan_pair",
+    "read_slot_fields",
+    "stamp_clean_chain",
+    "stamp_pair",
+    "to_signed_int64",
+    "write_slot_fields",
+]
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class FakeViolationLog:
+    ring: torch.Tensor
+    write_index: torch.Tensor
+    slot_run_counter: torch.Tensor
+    kernel_run_counter: torch.Tensor
+    enable_chain_position_assert: torch.Tensor
+
+    @classmethod
+    def allocate(
+        cls, *, capacity: int = DEFAULT_RING_CAPACITY, device: torch.device
+    ) -> FakeViolationLog:
+        return cls(
+            ring=torch.zeros(
+                capacity, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device
+            ),
+            write_index=torch.zeros(1, dtype=torch.int32, device=device),
+            slot_run_counter=torch.zeros(1, dtype=torch.int64, device=device),
+            kernel_run_counter=torch.zeros(1, dtype=torch.int64, device=device),
+            enable_chain_position_assert=torch.ones(
+                1, dtype=torch.int32, device=device
+            ),
+        )
+
+
+def make_canary_buf(
+    *,
+    num_slots: int = DEFAULT_NUM_SLOTS,
+    slot_stride_bytes: int = DEFAULT_SLOT_STRIDE_BYTES,
+    device: torch.device,
+) -> torch.Tensor:
+    return torch.zeros(num_slots, slot_stride_bytes, dtype=torch.uint8, device=device)
+
+
+def make_canary_buf_pair(
+    *,
+    num_slots: int = DEFAULT_NUM_SLOTS,
+    slot_stride_bytes: int = DEFAULT_SLOT_STRIDE_BYTES,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cuda_buf = make_canary_buf(
+        num_slots=num_slots, slot_stride_bytes=slot_stride_bytes, device=device
+    )
+    return cuda_buf, cuda_buf.clone()
+
+
+def make_log_pair(
+    *,
+    capacity: int = DEFAULT_RING_CAPACITY,
+    device: torch.device,
+) -> tuple[FakeViolationLog, FakeViolationLog]:
+    return (
+        FakeViolationLog.allocate(capacity=capacity, device=device),
+        FakeViolationLog.allocate(capacity=capacity, device=device),
+    )
+
+
+def make_verify_plan(
+    *,
+    slot_indices: list[int],
+    positions: list[int],
+    prev_slot_indices: list[int],
+    expected_input_ids: Optional[list[int]] = None,
+    capacity: Optional[int] = None,
+    device: torch.device,
+) -> VerifyPlan:
+    """Build a VerifyPlan whose active prefix matches the three input lists.
+
+    Active prefix mirrors the input lists. Tail entries are left at the
+    allocate-time defaults; ``verify_num_valid = len(slot_indices)``.
+
+    ``expected_input_ids`` defaults to ``[-1] * n_active`` (the verify-kernel
+    "skip token check" sentinel) so existing tests that only exercise the
+    chain / position / real-kv-hash paths keep working unchanged.
+    """
+    n_active = len(slot_indices)
+    if not (len(positions) == n_active and len(prev_slot_indices) == n_active):
+        raise ValueError(
+            "make_verify_plan: slot_indices, positions, and prev_slot_indices must all have the same length"
+        )
+    if expected_input_ids is None:
+        expected_input_ids = [-1] * n_active
+    if len(expected_input_ids) != n_active:
+        raise ValueError(
+            "make_verify_plan: expected_input_ids must match len(slot_indices)"
+        )
+    cap = capacity if capacity is not None else max(n_active, 1)
+    plan = VerifyPlan.allocate(verify_capacity=cap, device=device)
+    if n_active > 0:
+        plan.verify_slot_indices[:n_active] = torch.tensor(
+            slot_indices, dtype=torch.int64, device=device
+        )
+        plan.verify_expected_tokens[:n_active] = torch.tensor(
+            expected_input_ids, dtype=torch.int64, device=device
+        )
+        plan.verify_expected_positions[:n_active] = torch.tensor(
+            positions, dtype=torch.int64, device=device
+        )
+        plan.verify_prev_slot_indices[:n_active] = torch.tensor(
+            prev_slot_indices, dtype=torch.int64, device=device
+        )
+    plan.verify_num_valid[0] = n_active
+    return plan
+
+
+def make_verify_plan_pair(
+    *,
+    slot_indices: list[int],
+    positions: list[int],
+    prev_slot_indices: list[int],
+    expected_input_ids: Optional[list[int]] = None,
+    capacity: Optional[int] = None,
+    device: torch.device,
+) -> tuple[VerifyPlan, VerifyPlan]:
+    return (
+        make_verify_plan(
+            slot_indices=slot_indices,
+            positions=positions,
+            prev_slot_indices=prev_slot_indices,
+            expected_input_ids=expected_input_ids,
+            capacity=capacity,
+            device=device,
+        ),
+        make_verify_plan(
+            slot_indices=slot_indices,
+            positions=positions,
+            prev_slot_indices=prev_slot_indices,
+            expected_input_ids=expected_input_ids,
+            capacity=capacity,
+            device=device,
+        ),
+    )
+
+
+def make_write_plan(
+    *,
+    write_offsets: list[int],
+    seed_slot_indices: list[int],
+    num_valid_reqs: int,
+    req_capacity: Optional[int] = None,
+    device: torch.device,
+) -> WritePlan:
+    """Build a WritePlan from raw offsets and seed slot lists.
+
+    ``write_offsets`` must have length ``len(seed_slot_indices) + 1`` (the trailing total entry count).
+    """
+    n_active = len(seed_slot_indices)
+    if len(write_offsets) != n_active + 1:
+        raise ValueError(
+            "make_write_plan: write_offsets must have length len(seed_slot_indices) + 1"
+        )
+    cap = req_capacity if req_capacity is not None else max(n_active, 1)
+    plan = WritePlan.allocate(write_req_capacity=cap, device=device)
+    if n_active > 0:
+        plan.write_seed_slot_indices[:n_active] = torch.tensor(
+            seed_slot_indices, dtype=torch.int64, device=device
+        )
+    plan.write_offsets[: n_active + 1] = torch.tensor(
+        write_offsets, dtype=torch.int64, device=device
+    )
+    plan.write_num_valid_reqs[0] = num_valid_reqs
+    return plan
+
+
+def make_write_plan_pair(
+    *,
+    write_offsets: list[int],
+    seed_slot_indices: list[int],
+    num_valid_reqs: int,
+    req_capacity: Optional[int] = None,
+    device: torch.device,
+) -> tuple[WritePlan, WritePlan]:
+    return (
+        make_write_plan(
+            write_offsets=write_offsets,
+            seed_slot_indices=seed_slot_indices,
+            num_valid_reqs=num_valid_reqs,
+            req_capacity=req_capacity,
+            device=device,
+        ),
+        make_write_plan(
+            write_offsets=write_offsets,
+            seed_slot_indices=seed_slot_indices,
+            num_valid_reqs=num_valid_reqs,
+            req_capacity=req_capacity,
+            device=device,
+        ),
+    )
+
+
+def to_signed_int64(value: int) -> int:
+    value &= _U64_MASK
+    if value >= _I64_SIGN_BIT:
+        value -= 1 << 64
+    return value
+
+
+def chain_anchor_signed() -> int:
+    return to_signed_int64(splitmix64(consts.CANARY_CHAIN_ANCHOR))
+
+
+def write_slot_fields(
+    *,
+    canary_buf: torch.Tensor,
+    slot_idx: int,
+    token: int,
+    position: int,
+    prev_hash: int,
+    real_kv_hash: int,
+) -> None:
+    view = canary_buf.view(torch.int64)
+    view[slot_idx, 0] = token
+    view[slot_idx, 1] = position
+    view[slot_idx, 2] = prev_hash
+    view[slot_idx, 3] = real_kv_hash
+
+
+def stamp_pair(
+    buf_pair: tuple[torch.Tensor, torch.Tensor],
+    *,
+    slot_idx: int,
+    token: int,
+    position: int,
+    prev_hash: int,
+    real_kv_hash: int = 0,
+) -> None:
+    """Stamp the same slot fields into both (cuda, ref) canary buffers."""
+    for buf in buf_pair:
+        write_slot_fields(
+            canary_buf=buf,
+            slot_idx=slot_idx,
+            token=token,
+            position=position,
+            prev_hash=prev_hash,
+            real_kv_hash=real_kv_hash,
+        )
+
+
+def read_slot_fields(
+    *, canary_buf: torch.Tensor, slot_idx: int
+) -> tuple[int, int, int, int]:
+    row = canary_buf.view(torch.int64)[slot_idx, :4].detach().cpu().tolist()
+    return int(row[0]), int(row[1]), int(row[2]), int(row[3])
+
+
+def stamp_clean_chain(
+    *,
+    cuda_buf: torch.Tensor,
+    ref_buf: torch.Tensor,
+    slot_indices: list[int],
+    tokens: list[int],
+    positions: list[int],
+    real_kv_hashes: Optional[list[int]] = None,
+) -> list[int]:
+    n = len(tokens)
+    real_kv_hashes = real_kv_hashes if real_kv_hashes is not None else [0] * n
+    running_prev_hash = splitmix64(consts.CANARY_CHAIN_ANCHOR)
+    stored_prev_hashes: list[int] = []
+    for slot_idx, token, position, real_kv_hash in zip(
+        slot_indices, tokens, positions, real_kv_hashes
+    ):
+        signed_prev = to_signed_int64(running_prev_hash)
+        for buf in (cuda_buf, ref_buf):
+            write_slot_fields(
+                canary_buf=buf,
+                slot_idx=slot_idx,
+                token=token,
+                position=position,
+                prev_hash=signed_prev,
+                real_kv_hash=to_signed_int64(real_kv_hash),
+            )
+        stored_prev_hashes.append(signed_prev)
+        running_prev_hash = splitmix64_mix3(running_prev_hash, token, position)
+    return stored_prev_hashes
+
+
+def assert_canary_state_equal(
+    *, log_a: FakeViolationLog, log_b: FakeViolationLog
+) -> None:
+    for name in ("ring", "write_index", "slot_run_counter", "kernel_run_counter"):
+        assert torch.equal(
+            getattr(log_a, name), getattr(log_b, name)
+        ), f"{name} diverged (CUDA vs ref)"
+
+
+def assert_canary_buf_equal(*, buf_a: torch.Tensor, buf_b: torch.Tensor) -> None:
+    assert torch.equal(buf_a, buf_b), "canary_buf diverged (CUDA vs ref)"
+
+
+def assert_only_bits_set(fail_bits: int, expected_bits: int) -> None:
+    assert (
+        fail_bits & expected_bits
+    ) == expected_bits, (
+        f"missing expected bits: expected {expected_bits:#b} got {fail_bits:#b}"
+    )
+    assert (
+        fail_bits & ~expected_bits
+    ) == 0, f"unexpected extra bits: got {fail_bits:#b} extras {fail_bits & ~expected_bits:#b}"

python/sglang/jit_kernel/tests/kv_canary/_constants.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
+
+# Default fixture sizes — small enough for fast tests, large enough that ring overflow / multi-req cases
+# stay realistic without bloating the assertion surface.
+DEFAULT_RING_CAPACITY: int = 64
+DEFAULT_NUM_SLOTS: int = 32
+DEFAULT_SLOT_STRIDE_BYTES: int = CANARY_SLOT_BYTES
+
+_U64_MASK: int = (1 << 64) - 1
+_I64_SIGN_BIT: int = 1 << 63

python/sglang/jit_kernel/tests/kv_canary/_differential.py

@@ -0,0 +1,546 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any, Callable, Iterator, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
+from sglang.jit_kernel.kv_canary.plan_ref import (
+    launch_canary_plan_kernels_torch_reference,
+)
+from sglang.jit_kernel.kv_canary.verify import (
+    CanaryLaunchTag,
+    RealKvSource,
+    VerifyOrWriteContext,
+    VerifyPlan,
+    launch_canary_verify_kernel,
+)
+from sglang.jit_kernel.kv_canary.verify_ref import (
+    launch_canary_verify_kernel_torch_reference,
+)
+from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
+from sglang.jit_kernel.kv_canary.write_ref import (
+    launch_canary_write_kernel_torch_reference,
+)
+from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
+    FakeViolationLog,
+    assert_canary_buf_equal,
+    assert_canary_state_equal,
+    make_log_pair,
+)
+
+_DEVICE = torch.device("cuda")
+
+
+def _run_both_plan(
+    *,
+    triton_verify: VerifyPlan,
+    triton_write: WritePlan,
+    ref_verify: VerifyPlan,
+    ref_write: WritePlan,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+    swa_window_size: int,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    assert_equal: bool = True,
+    active_verify_entries: Optional[int] = None,
+    active_write_reqs: Optional[int] = None,
+    req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
+    kv_token_id_vs_position_offset: int = 0,
+) -> None:
+    _ = extras
+    verify_capacity = int(triton_verify.verify_slot_indices.shape[0])
+    # Default lens to "no tighter bound than pool width" so existing kernel tests that
+    # only care about gather wiring keep their old semantics without each call site
+    # explicitly building a per-req lens tensor.
+    if (
+        req_to_verify_expected_tokens is not None
+        and req_to_verify_expected_tokens_valid_lens is None
+    ):
+        req_to_verify_expected_tokens_valid_lens = torch.full(
+            (int(req_pool_indices.shape[0]),),
+            int(req_to_verify_expected_tokens.shape[1]),
+            dtype=torch.int64,
+            device=req_pool_indices.device,
+        )
+    launch_canary_plan_kernels(
+        verify_plan_out=triton_verify,
+        write_plan_out=triton_write,
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        extend_seq_lens=extend_seq_lens,
+        req_to_token=req_to_token,
+        swa_window_size=swa_window_size,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        verify_capacity=verify_capacity,
+        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    launch_canary_plan_kernels_torch_reference(
+        verify_plan_out=ref_verify,
+        write_plan_out=ref_write,
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        extend_seq_lens=extend_seq_lens,
+        req_to_token=req_to_token,
+        swa_window_size=swa_window_size,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        verify_capacity=int(ref_verify.verify_slot_indices.shape[0]),
+        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    torch.cuda.synchronize()
+
+    if assert_equal:
+        _assert_plans_byte_equal(
+            triton_verify=triton_verify,
+            triton_write=triton_write,
+            ref_verify=ref_verify,
+            ref_write=ref_write,
+            active_verify_entries=active_verify_entries,
+            active_write_reqs=active_write_reqs,
+        )
+
+
+def _assert_plans_byte_equal(
+    *,
+    triton_verify: VerifyPlan,
+    triton_write: WritePlan,
+    ref_verify: VerifyPlan,
+    ref_write: WritePlan,
+    active_verify_entries: Optional[int] = None,
+    active_write_reqs: Optional[int] = None,
+) -> None:
+    """Byte-equal check on (Triton vs ref) plan outputs.
+
+    Optional ``active_verify_entries`` / ``active_write_reqs`` truncate the comparison to the meaningful
+    prefix; tail entries past the active count are kernel-undefined and need not match byte-equal.
+    """
+    n_verify = (
+        active_verify_entries
+        if active_verify_entries is not None
+        else int(triton_verify.verify_num_valid[0].item())
+    )
+    n_verify_ref = int(ref_verify.verify_num_valid[0].item())
+    assert (
+        n_verify == n_verify_ref
+    ), f"verify_num_valid diverged: triton={n_verify} ref={n_verify_ref}"
+    # When total_verify > VERIFY_CAPACITY the offsets kernel clears verify_enable and
+    # plan_entries skips its scatter — leaving verify_slot_indices/positions/prev_slot_indices
+    # as whatever the (torch.empty) allocation contained. Skip the byte-equal probe in that
+    # case; verify_num_valid being clamped + verify_enable=0 is the contract here.
+    triton_enable = int(triton_verify.enable[0].item())
+    ref_enable = int(ref_verify.enable[0].item())
+    assert (
+        triton_enable == ref_enable
+    ), f"verify_enable diverged: triton={triton_enable} ref={ref_enable}"
+    if n_verify > 0 and triton_enable != 0:
+        assert torch.equal(
+            triton_verify.verify_slot_indices[:n_verify],
+            ref_verify.verify_slot_indices[:n_verify],
+        )
+        assert torch.equal(
+            triton_verify.verify_expected_tokens[:n_verify],
+            ref_verify.verify_expected_tokens[:n_verify],
+        )
+        assert torch.equal(
+            triton_verify.verify_expected_positions[:n_verify],
+            ref_verify.verify_expected_positions[:n_verify],
+        )
+        assert torch.equal(
+            triton_verify.verify_prev_slot_indices[:n_verify],
+            ref_verify.verify_prev_slot_indices[:n_verify],
+        )
+
+    n_write = (
+        active_write_reqs
+        if active_write_reqs is not None
+        else int(triton_write.write_num_valid_reqs[0].item())
+    )
+    n_write_ref = int(ref_write.write_num_valid_reqs[0].item())
+    assert (
+        n_write == n_write_ref
+    ), f"write_num_valid_reqs diverged: triton={n_write} ref={n_write_ref}"
+    assert torch.equal(
+        triton_write.write_offsets[: n_write + 1],
+        ref_write.write_offsets[: n_write + 1],
+    )
+    if n_write > 0:
+        assert torch.equal(
+            triton_write.write_seed_slot_indices[:n_write],
+            ref_write.write_seed_slot_indices[:n_write],
+        )
+
+
+def _run_both_verify(
+    *,
+    cuda_canary_buf: torch.Tensor,
+    ref_canary_buf: torch.Tensor,
+    plan_cuda,
+    plan_ref,
+    cuda_log: FakeViolationLog,
+    ref_log: FakeViolationLog,
+    real_kv_sources_cuda: tuple[RealKvSource, ...],
+    real_kv_sources_ref: tuple[RealKvSource, ...],
+    real_kv_hash_mode: consts.RealKvHashMode,
+    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
+    assert_equal: bool = True,
+    check_verify_expected_token: bool = True,
+) -> None:
+    launch_canary_verify_kernel(
+        context=VerifyOrWriteContext(
+            canary_buf=cuda_canary_buf,
+            kernel_kind=kernel_kind,
+            violation_ring=cuda_log.ring,
+            violation_write_index=cuda_log.write_index,
+            slot_run_counter=cuda_log.slot_run_counter,
+            kernel_run_counter=cuda_log.kernel_run_counter,
+            enable_chain_position_assert=cuda_log.enable_chain_position_assert,
+            real_kv_sources=real_kv_sources_cuda,
+            real_kv_hash_mode=real_kv_hash_mode,
+        ),
+        plan=plan_cuda,
+        check_verify_expected_token=check_verify_expected_token,
+    )
+    launch_canary_verify_kernel_torch_reference(
+        context=VerifyOrWriteContext(
+            canary_buf=ref_canary_buf,
+            kernel_kind=kernel_kind,
+            violation_ring=ref_log.ring,
+            violation_write_index=ref_log.write_index,
+            slot_run_counter=ref_log.slot_run_counter,
+            kernel_run_counter=ref_log.kernel_run_counter,
+            enable_chain_position_assert=ref_log.enable_chain_position_assert,
+            real_kv_sources=real_kv_sources_ref,
+            real_kv_hash_mode=real_kv_hash_mode,
+        ),
+        plan=plan_ref,
+        check_verify_expected_token=check_verify_expected_token,
+    )
+    torch.cuda.synchronize()
+
+    if assert_equal:
+        assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
+
+
+def _run_both_write(
+    *,
+    cuda_canary_buf: torch.Tensor,
+    ref_canary_buf: torch.Tensor,
+    plan_cuda,
+    plan_ref,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    enable_write_verify_inputs: bool,
+    expected_input_tokens: torch.Tensor,
+    expected_input_positions: torch.Tensor,
+    cuda_log: FakeViolationLog,
+    ref_log: FakeViolationLog,
+    real_kv_sources_cuda: tuple[RealKvSource, ...],
+    real_kv_sources_ref: tuple[RealKvSource, ...],
+    real_kv_hash_mode: consts.RealKvHashMode,
+    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
+    assert_equal: bool = True,
+) -> None:
+    expected_tokens_for_launch = (
+        expected_input_tokens if enable_write_verify_inputs else None
+    )
+    expected_positions_for_launch = (
+        expected_input_positions if enable_write_verify_inputs else None
+    )
+    launch_canary_write_kernel(
+        context=VerifyOrWriteContext(
+            canary_buf=cuda_canary_buf,
+            kernel_kind=kernel_kind,
+            violation_ring=cuda_log.ring,
+            violation_write_index=cuda_log.write_index,
+            slot_run_counter=cuda_log.slot_run_counter,
+            kernel_run_counter=cuda_log.kernel_run_counter,
+            enable_chain_position_assert=cuda_log.enable_chain_position_assert,
+            real_kv_sources=real_kv_sources_cuda,
+            real_kv_hash_mode=real_kv_hash_mode,
+        ),
+        plan=plan_cuda,
+        input_ids=input_ids,
+        positions=positions,
+        out_cache_loc=out_cache_loc,
+        enable_write_input_assert=enable_write_verify_inputs,
+        expected_input_tokens=expected_tokens_for_launch,
+        expected_input_positions=expected_positions_for_launch,
+    )
+    launch_canary_write_kernel_torch_reference(
+        context=VerifyOrWriteContext(
+            canary_buf=ref_canary_buf,
+            kernel_kind=kernel_kind,
+            violation_ring=ref_log.ring,
+            violation_write_index=ref_log.write_index,
+            slot_run_counter=ref_log.slot_run_counter,
+            kernel_run_counter=ref_log.kernel_run_counter,
+            enable_chain_position_assert=ref_log.enable_chain_position_assert,
+            real_kv_sources=real_kv_sources_ref,
+            real_kv_hash_mode=real_kv_hash_mode,
+        ),
+        plan=plan_ref,
+        input_ids=input_ids,
+        positions=positions,
+        out_cache_loc=out_cache_loc,
+        enable_write_input_assert=enable_write_verify_inputs,
+        expected_input_tokens=expected_tokens_for_launch,
+        expected_input_positions=expected_positions_for_launch,
+    )
+    torch.cuda.synchronize()
+
+    if assert_equal:
+        assert_canary_buf_equal(buf_a=cuda_canary_buf, buf_b=ref_canary_buf)
+        assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ShrinkResult:
+    inputs: Any
+    mutations_applied: list[str]
+
+
+def shrink_inputs(
+    inputs: Any,
+    *,
+    check_fn: Callable[[Any], bool],
+    max_iterations: int = 50,
+) -> ShrinkResult:
+    """Greedy 1-step minify for a fuzz inputs dataclass.
+
+    ``check_fn(candidate)`` returns True when ``candidate`` still reproduces the failure. Each round
+    yields candidate-simpler-than-current mutations through ``_yield_simpler``; the first accepted
+    candidate becomes the new current. Iteration stops when no mutation is accepted or ``max_iterations``
+    is reached.
+    """
+    current = inputs
+    applied: list[str] = []
+    for _ in range(max_iterations):
+        improved = False
+        for label, candidate in _yield_simpler(current):
+            try:
+                still_fails = check_fn(candidate)
+            except Exception:
+                still_fails = False
+            if still_fails:
+                current = candidate
+                applied.append(label)
+                improved = True
+                break
+        if not improved:
+            break
+    return ShrinkResult(inputs=current, mutations_applied=applied)
+
+
+def _yield_simpler(inputs: Any) -> Iterator[tuple[str, Any]]:
+    """Yield (label, simpler_candidate) tuples for generic fuzz-input minifiers.
+
+    The candidates touch only well-known field names; an inputs dataclass that lacks a field will simply
+    have that mutation skipped. No kernel-specific knowledge is encoded here so the same shrinker drives
+    Plan / Verify / Write fuzz failures uniformly.
+    """
+    fields = {
+        f: getattr(inputs, f) for f in inputs.__dataclass_fields__  # type: ignore[attr-defined]
+    }
+
+    def emit(label: str, **overrides: Any) -> Iterator[tuple[str, Any]]:
+        candidate = replace(inputs, **overrides)
+        yield label, candidate
+
+    bs_field = (
+        "req_pool_indices"
+        if "req_pool_indices" in fields
+        else ("input_ids" if "input_ids" in fields else None)
+    )
+    if bs_field is not None and isinstance(fields[bs_field], torch.Tensor):
+        tensor = fields[bs_field]
+        if tensor.numel() > 1:
+            new_len = tensor.numel() - 1
+            related_tensors_overrides: dict[str, Any] = {}
+            for name in (
+                "req_pool_indices",
+                "prefix_lens",
+                "extend_seq_lens",
+                "input_ids",
+                "positions",
+                "out_cache_loc",
+                "expected_input_tokens",
+                "expected_input_positions",
+            ):
+                t = fields.get(name)
+                if (
+                    isinstance(t, torch.Tensor)
+                    and t.numel() >= new_len
+                    and t.dim() == 1
+                ):
+                    related_tensors_overrides[name] = t[:new_len].contiguous()
+            if related_tensors_overrides:
+                yield from emit("drop_last_row", **related_tensors_overrides)
+
+    if "swa_window_size" in fields and isinstance(fields["swa_window_size"], int):
+        if fields["swa_window_size"] != 0:
+            yield from emit(
+                "swa_off", swa_window_size=0, full_to_swa_index_mapping=None
+            )
+
+    if "extras_count" in fields and isinstance(fields["extras_count"], int):
+        if fields["extras_count"] > 0:
+            yield from emit("extras_zero", extras_count=0)
+
+    if "real_kv_hash_mode" in fields:
+        cur = fields["real_kv_hash_mode"]
+        if hasattr(cur, "value"):
+            cls = cur.__class__
+            if int(cur) == 2:
+                yield from emit("hash_mode_bit", real_kv_hash_mode=cls(1))
+            elif int(cur) == 1:
+                yield from emit("hash_mode_off", real_kv_hash_mode=cls(0))
+
+    if "real_kv_sources" in fields:
+        srcs = fields["real_kv_sources"]
+        if isinstance(srcs, tuple) and len(srcs) > 1:
+            yield from emit("sources_to_one", real_kv_sources=srcs[:1])
+
+    if "enable_write_verify_inputs" in fields:
+        cur = fields["enable_write_verify_inputs"]
+        if hasattr(cur, "value") and int(cur) != 0:
+            cls = cur.__class__
+            yield from emit("pseudo_off", enable_write_verify_inputs=cls(0))
+
+    for name in ("verify_capacity", "write_req_capacity"):
+        if name in fields and isinstance(fields[name], int):
+            current_value = fields[name]
+            if current_value > 8:
+                yield from emit(f"shrink_{name}", **{name: max(8, current_value // 2)})
+
+
+def run_verify_diff(
+    *,
+    buf_pair: tuple[torch.Tensor, torch.Tensor],
+    plan_pair: tuple[VerifyPlan, VerifyPlan],
+    real_kv_sources_pair: tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] = (
+        (),
+        (),
+    ),
+    real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
+    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
+    device: torch.device = _DEVICE,
+    assert_equal: bool = True,
+    check_verify_expected_token: bool = True,
+) -> tuple[FakeViolationLog, FakeViolationLog]:
+    """Thin wrapper around ``_run_both_verify`` that creates a fresh log pair and packs (cuda, ref)
+    buf/plan/source arguments into 2-tuples to drop ~8 lines of boilerplate per call site.
+    """
+    cuda_log, ref_log = make_log_pair(device=device)
+    _run_both_verify(
+        cuda_canary_buf=buf_pair[0],
+        ref_canary_buf=buf_pair[1],
+        plan_cuda=plan_pair[0],
+        plan_ref=plan_pair[1],
+        cuda_log=cuda_log,
+        ref_log=ref_log,
+        real_kv_sources_cuda=real_kv_sources_pair[0],
+        real_kv_sources_ref=real_kv_sources_pair[1],
+        real_kv_hash_mode=real_kv_hash_mode,
+        kernel_kind=kernel_kind,
+        assert_equal=assert_equal,
+        check_verify_expected_token=check_verify_expected_token,
+    )
+    return cuda_log, ref_log
+
+
+def run_write_diff(
+    *,
+    buf_pair: tuple[torch.Tensor, torch.Tensor],
+    plan_pair: tuple[WritePlan, WritePlan],
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    expected_input_tokens: torch.Tensor,
+    expected_input_positions: torch.Tensor,
+    enable_write_verify_inputs: bool = False,
+    real_kv_sources_pair: tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] = (
+        (),
+        (),
+    ),
+    real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
+    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
+    device: torch.device = _DEVICE,
+    assert_equal: bool = True,
+) -> tuple[FakeViolationLog, FakeViolationLog]:
+    """Thin wrapper around ``_run_both_write`` that creates a fresh log pair and packs (cuda, ref)
+    buf/plan/source arguments into 2-tuples to drop ~10 lines of boilerplate per call site.
+    """
+    cuda_log, ref_log = make_log_pair(device=device)
+    _run_both_write(
+        cuda_canary_buf=buf_pair[0],
+        ref_canary_buf=buf_pair[1],
+        plan_cuda=plan_pair[0],
+        plan_ref=plan_pair[1],
+        input_ids=input_ids,
+        positions=positions,
+        out_cache_loc=out_cache_loc,
+        enable_write_verify_inputs=enable_write_verify_inputs,
+        expected_input_tokens=expected_input_tokens,
+        expected_input_positions=expected_input_positions,
+        cuda_log=cuda_log,
+        ref_log=ref_log,
+        real_kv_sources_cuda=real_kv_sources_pair[0],
+        real_kv_sources_ref=real_kv_sources_pair[1],
+        real_kv_hash_mode=real_kv_hash_mode,
+        kernel_kind=kernel_kind,
+        assert_equal=assert_equal,
+    )
+    return cuda_log, ref_log
+
+
+def run_plan_diff(
+    *,
+    plan_pair: tuple[tuple[VerifyPlan, WritePlan], tuple[VerifyPlan, WritePlan]],
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+    swa_window_size: int = 0,
+    full_to_swa_index_mapping: Optional[torch.Tensor] = None,
+    assert_equal: bool = True,
+    active_verify_entries: Optional[int] = None,
+    active_write_reqs: Optional[int] = None,
+    req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
+    kv_token_id_vs_position_offset: int = 0,
+) -> None:
+    """Thin wrapper around ``_run_both_plan`` that unpacks ``((triton_v, triton_w), (ref_v, ref_w))``
+    plan pairs to drop the per-call-site ``triton_verify=.../triton_write=.../ref_verify=...`` block.
+    """
+    (triton_verify, triton_write), (ref_verify, ref_write) = plan_pair
+    _run_both_plan(
+        triton_verify=triton_verify,
+        triton_write=triton_write,
+        ref_verify=ref_verify,
+        ref_write=ref_write,
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        extend_seq_lens=extend_seq_lens,
+        req_to_token=req_to_token,
+        extras=extras,
+        swa_window_size=swa_window_size,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        assert_equal=assert_equal,
+        active_verify_entries=active_verify_entries,
+        active_write_reqs=active_write_reqs,
+        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )

python/sglang/jit_kernel/tests/kv_canary/_fixtures.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import random
+from typing import Literal, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import (
+    RealKvSource,
+    VerifyPlan,
+)
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS
+
+_DEVICE = torch.device("cuda")
+
+
+LutKind = Literal["identity", "shift", "permutation", "with_oob"]
+
+
+def make_lut(
+    *,
+    kind: LutKind,
+    pool_size: int,
+    device: torch.device,
+    rng: Optional[random.Random] = None,
+) -> torch.Tensor:
+    base = torch.arange(pool_size + 1, dtype=torch.int64, device=device)
+    if kind == "identity":
+        return base.contiguous()
+    if kind == "shift":
+        return (base + 100).contiguous()
+    if kind in ("permutation", "with_oob"):
+        if rng is None:
+            rng = random.Random(0)
+        perm = list(range(pool_size + 1))
+        rng.shuffle(perm)
+        out = torch.tensor(perm, dtype=torch.int64, device=device)
+        if kind == "with_oob":
+            out[-1] = pool_size + 999
+        return out.contiguous()
+    raise ValueError(f"unknown LutKind: {kind}")
+
+
+ReqToTokenKind = Literal["linear", "sparse_permuted"]
+
+
+def make_req_to_token(
+    *,
+    kind: ReqToTokenKind,
+    max_reqs: int,
+    max_seq_len: int,
+    device: torch.device,
+    rng: Optional[random.Random] = None,
+) -> torch.Tensor:
+    if kind == "linear":
+        rp_axis = torch.arange(max_reqs, device=device, dtype=torch.int32).unsqueeze(1)
+        pos_axis = torch.arange(
+            max_seq_len, device=device, dtype=torch.int32
+        ).unsqueeze(0)
+        return (rp_axis * max_seq_len + pos_axis).contiguous()
+    if rng is None:
+        rng = random.Random(0)
+    pool_size = max_reqs * max_seq_len
+    # Slots index into a full_to_swa LUT sized [pool_size + 1], so values must stay
+    # in [0, pool_size]. The universe spans [1, pool_size] (skipping 0 as reserved),
+    # giving exactly max_reqs * max_seq_len unique slots — one per (rp, pos) cell.
+    slot_universe = list(range(1, pool_size + 1))
+    rng.shuffle(slot_universe)
+    rtt = torch.zeros((max_reqs, max_seq_len), dtype=torch.int32, device=device)
+    cursor = 0
+    for rp in range(max_reqs):
+        per_req = slot_universe[cursor : cursor + max_seq_len]
+        cursor += max_seq_len
+        rtt[rp, :] = torch.tensor(per_req, dtype=torch.int32, device=device)
+    return rtt.contiguous()
+
+
+def make_real_kv_source(
+    *,
+    num_slots: int = DEFAULT_NUM_SLOTS,
+    num_bytes_per_token: int = 16,
+    page_size: int = 1,
+    read_bytes: Optional[int] = None,
+    pad_dim1: int = 0,
+    device: torch.device,
+    fill: int = 0,
+) -> RealKvSource:
+    """Allocate one RealKvSource with the canonical [num_rows, dim1_bytes] uint8 shape.
+
+    ``pad_dim1`` adds trailing per-row bytes the canary should skip — used by the "holey dim 1" case to
+    confirm the kernel never reads past ``page_size * num_bytes_per_token``.
+    """
+    num_rows = (num_slots + page_size - 1) // page_size
+    cols = page_size * num_bytes_per_token + pad_dim1
+    tensor = torch.full(
+        (num_rows, cols), fill_value=fill, dtype=torch.uint8, device=device
+    )
+    effective_read = read_bytes if read_bytes is not None else num_bytes_per_token
+    return RealKvSource(
+        tensor=tensor,
+        page_size=page_size,
+        num_bytes_per_token=num_bytes_per_token,
+        read_bytes=effective_read,
+    )
+
+
+FillStrategy = Literal["constant_per_source", "random_bytes"]
+
+
+def make_real_kv_sources(
+    *,
+    count: int,
+    num_bytes_per_token: int = 16,
+    page_size: int = 1,
+    num_slots: int = DEFAULT_NUM_SLOTS,
+    device: torch.device,
+    rng: Optional[random.Random] = None,
+    fill_strategy: FillStrategy = "constant_per_source",
+) -> tuple[RealKvSource, ...]:
+    sources: list[RealKvSource] = []
+    for i in range(count):
+        read_bytes_eff = num_bytes_per_token
+        src = make_real_kv_source(
+            num_slots=num_slots,
+            num_bytes_per_token=num_bytes_per_token,
+            page_size=page_size,
+            read_bytes=read_bytes_eff,
+            device=device,
+            fill=(i + 1) * 17,
+        )
+        if fill_strategy == "random_bytes":
+            if rng is None:
+                rng = random.Random(0)
+            seed = rng.randint(0, 0xFFFFFFFF)
+            gen = torch.Generator(device=device).manual_seed(seed)
+            src.tensor.random_(generator=gen)
+        sources.append(src)
+    return tuple(sources)
+
+
+def clone_real_kv_sources(
+    sources: tuple[RealKvSource, ...],
+) -> tuple[RealKvSource, ...]:
+    return tuple(
+        RealKvSource(
+            tensor=src.tensor.clone(),
+            page_size=src.page_size,
+            num_bytes_per_token=src.num_bytes_per_token,
+            read_bytes=src.read_bytes,
+        )
+        for src in sources
+    )
+
+
+PaddingKind = Literal["none", "trailing", "interleaved"]
+
+
+def make_padding_mask(
+    *,
+    bs: int,
+    kind: PaddingKind,
+    rng: Optional[random.Random] = None,
+    padding_fraction: float = 0.25,
+) -> list[bool]:
+    if bs == 0:
+        return []
+    if kind == "none":
+        return [False] * bs
+    n_pad = max(1, int(bs * padding_fraction)) if bs > 0 else 0
+    n_pad = min(n_pad, bs)
+    if kind == "trailing":
+        return [False] * (bs - n_pad) + [True] * n_pad
+    if kind == "interleaved":
+        if rng is None:
+            rng = random.Random(0)
+        mask = [False] * bs
+        chosen = rng.sample(range(bs), k=n_pad)
+        for idx in chosen:
+            mask[idx] = True
+        return mask
+    raise ValueError(f"unknown PaddingKind: {kind}")
+
+
+CapacityKind = Literal["loose", "tight_match", "under_by_one"]
+
+
+def derive_plan_capacity(
+    *,
+    kind: CapacityKind,
+    total_verify: int,
+    extras_count: int,
+    bs: int,
+) -> tuple[int, int]:
+    needed = total_verify + extras_count
+    if kind == "loose":
+        return max(needed + 64, 128), max(bs + 4, 8)
+    if kind == "tight_match":
+        return max(needed, 1), max(bs + 4, 8)
+    if kind == "under_by_one":
+        return max(needed - 1, 1), max(bs + 4, 8)
+    raise ValueError(f"unknown CapacityKind: {kind}")
+
+
+def allocate_plan_pair(
+    *,
+    verify_capacity: int,
+    write_req_capacity: int,
+) -> tuple[VerifyPlan, WritePlan, VerifyPlan, WritePlan]:
+    return (
+        VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE),
+        WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE),
+        VerifyPlan.allocate(verify_capacity=verify_capacity, device=_DEVICE),
+        WritePlan.allocate(write_req_capacity=write_req_capacity, device=_DEVICE),
+    )
+
+
+def empty_extras() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    return (
+        torch.zeros(1, dtype=torch.int64, device=_DEVICE),
+        torch.zeros(1, dtype=torch.int64, device=_DEVICE),
+        torch.zeros(1, dtype=torch.int64, device=_DEVICE),
+        torch.zeros(1, dtype=torch.int32, device=_DEVICE),
+    )
+
+
+def dummy_pseudo_tensors(num_tokens: int) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.zeros(num_tokens, dtype=torch.int64, device=_DEVICE),
+        torch.zeros(num_tokens, dtype=torch.int64, device=_DEVICE),
+    )

python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import random
+from typing import Any, Callable
+
+from sglang.jit_kernel.tests.kv_canary._differential import (
+    ShrinkResult,
+    shrink_inputs,
+)
+
+FUZZ_SEEDS_PR: tuple[int, ...] = (0,)
+
+
+def check_repro(inputs: Any, *, run_one_fn: Callable[[Any], Any]) -> bool:
+    try:
+        run_one_fn(inputs)
+    except (AssertionError, RuntimeError, ValueError):
+        return True
+    return False
+
+
+def run_fuzz_combo(
+    seed: int,
+    *,
+    draw_fn: Callable[[random.Random], Any],
+    run_one_fn: Callable[[Any], Any],
+    summarize_fn: Callable[[Any], str],
+    n_iter: int,
+) -> None:
+    rng = random.Random(seed)
+    for iteration in range(n_iter):
+        inputs = draw_fn(rng)
+        try:
+            run_one_fn(inputs)
+        except AssertionError as exc:
+            shrunk: ShrinkResult = shrink_inputs(
+                inputs,
+                check_fn=lambda i: check_repro(i, run_one_fn=run_one_fn),
+            )
+            raise AssertionError(
+                f"seed={seed} iter={iteration} failure: {exc}\n"
+                f"original: {summarize_fn(inputs)}\n"
+                f"shrunk:   {summarize_fn(shrunk.inputs)}\n"
+                f"mutations applied: {shrunk.mutations_applied}"
+            ) from exc

python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py

@@ -0,0 +1,31 @@
+"""Hand-computed Python re-implementation of the real-kv-source fold, kept independent from
+``verify_ref._splitmix64_fold_bytes_scalar`` so a ref / kernel co-regression cannot silently fix the
+diff comparison."""
+
+from __future__ import annotations
+
+from sglang.jit_kernel.kv_canary.consts import splitmix64
+
+
+def _fold_words(padded: bytes) -> int:
+    """Pack padded bytes little-endian into 8-byte words, fold each via splitmix64 from acc=0."""
+    num_words = len(padded) // 8
+    acc = 0
+    for w in range(num_words):
+        chunk = padded[w * 8 : (w + 1) * 8]
+        word = sum(b << (8 * k) for k, b in enumerate(chunk))
+        acc = splitmix64(acc ^ word)
+    return splitmix64(0 ^ acc)
+
+
+def _hand_fold_partial(raw_bytes: bytes) -> int:
+    """PARTIAL-mode fold: first min(16, len) bytes, little-endian word-pack + splitmix64, same as ALL."""
+    truncated = raw_bytes[: min(16, len(raw_bytes))]
+    pad = (8 - len(truncated) % 8) % 8
+    return _fold_words(bytes(truncated) + bytes(pad))
+
+
+def _hand_fold_all(raw_bytes: bytes) -> int:
+    """ALL-mode fold: pack bytes little-endian into 8-byte words, fold each via splitmix64, then mix into acc=0."""
+    pad = (8 - len(raw_bytes) % 8) % 8
+    return _fold_words(raw_bytes + bytes(pad))

python/sglang/jit_kernel/tests/kv_canary/_invariants.py

@@ -0,0 +1,520 @@
+"""Ref/real-independent invariant assertions for kv_canary kernel tests.
+
+Each invariant only looks at the kernel's inputs and outputs (shape relationships, monotonicity, tail
+positions, etc.) — it must never re-implement the reference algorithm. Hand and fuzz tests both call
+into this module so a single contract violation surfaces consistently.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.jit_kernel.tests.kv_canary._canary_helpers import FakeViolationLog
+
+
+class PlanInvariants:
+    @staticmethod
+    def assert_all(
+        *,
+        verify_plan: VerifyPlan,
+        write_plan: WritePlan,
+        req_pool_indices: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        swa_window_size: int,
+        extras_slot_indices: torch.Tensor,
+        extras_positions: torch.Tensor,
+        extras_prev_slot_indices: torch.Tensor,
+        extras_count: int,
+    ) -> None:
+        PlanInvariants._assert_write_offsets_monotone(write_plan)
+        PlanInvariants._assert_write_offsets_total_matches_active_extend_sum(
+            write_plan=write_plan,
+            extend_seq_lens=extend_seq_lens,
+            req_pool_indices=req_pool_indices,
+        )
+        derived = PlanInvariants._assert_verify_num_valid_equals_derived_plus_extras(
+            verify_plan=verify_plan,
+            prefix_lens=prefix_lens,
+            req_pool_indices=req_pool_indices,
+            swa_window_size=swa_window_size,
+            extras_count=extras_count,
+        )
+        PlanInvariants._assert_padding_row_seed_is_minus_one(
+            write_plan=write_plan,
+            req_pool_indices=req_pool_indices,
+        )
+        # In overflow (derived + extras > verify_capacity) the plan kernel disables
+        # verify (enable=0) and the verify entries buffer is partially populated; the
+        # downstream entry-shape invariants only hold when the kernel actually emitted
+        # the full set, so guard them on enable=1.
+        verify_enabled = int(verify_plan.enable[0].item()) == 1
+        if verify_enabled:
+            PlanInvariants._assert_extras_land_at_tail(
+                verify_plan=verify_plan,
+                derived_verify_count=derived,
+                extras_slot_indices=extras_slot_indices,
+                extras_positions=extras_positions,
+                extras_prev_slot_indices=extras_prev_slot_indices,
+                extras_count=extras_count,
+            )
+            PlanInvariants._assert_prev_slot_minus_one_iff_chain_head(
+                verify_plan=verify_plan,
+                swa_window_size=swa_window_size,
+                derived_verify_count=derived,
+            )
+
+    @staticmethod
+    def _assert_write_offsets_monotone(write_plan: WritePlan) -> None:
+        n_active = int(write_plan.write_num_valid_reqs[0].item())
+        if n_active < 0:
+            raise AssertionError(f"write_num_valid_reqs negative: {n_active}")
+        offsets = write_plan.write_offsets[: n_active + 1].detach().cpu().tolist()
+        for i in range(len(offsets) - 1):
+            assert (
+                offsets[i] <= offsets[i + 1]
+            ), f"write_offsets non-monotone at {i}: {offsets[i]} > {offsets[i + 1]}"
+
+    @staticmethod
+    def _assert_write_offsets_total_matches_active_extend_sum(
+        *,
+        write_plan: WritePlan,
+        extend_seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+    ) -> None:
+        n_active = int(write_plan.write_num_valid_reqs[0].item())
+        total = int(write_plan.write_offsets[n_active].item())
+        rpi_cpu = req_pool_indices.detach().cpu().tolist()
+        ext_cpu = extend_seq_lens.detach().cpu().tolist()
+        expected_total = sum(ext for rpi, ext in zip(rpi_cpu, ext_cpu) if rpi != 0)
+        assert (
+            total == expected_total
+        ), f"write_offsets total {total} != active extend sum {expected_total}"
+
+    @staticmethod
+    def _assert_extras_land_at_tail(
+        *,
+        verify_plan: VerifyPlan,
+        derived_verify_count: int,
+        extras_slot_indices: torch.Tensor,
+        extras_positions: torch.Tensor,
+        extras_prev_slot_indices: torch.Tensor,
+        extras_count: int,
+    ) -> None:
+        if extras_count == 0:
+            return
+        tail_start = derived_verify_count
+        tail_end = derived_verify_count + extras_count
+        n_valid = int(verify_plan.verify_num_valid[0].item())
+        assert (
+            tail_end <= n_valid
+        ), f"extras tail {tail_end} exceeds verify_num_valid {n_valid}"
+        plan_slots = verify_plan.verify_slot_indices[tail_start:tail_end]
+        plan_positions = verify_plan.verify_expected_positions[tail_start:tail_end]
+        plan_prevs = verify_plan.verify_prev_slot_indices[tail_start:tail_end]
+        assert torch.equal(plan_slots, extras_slot_indices[:extras_count])
+        assert torch.equal(plan_positions, extras_positions[:extras_count])
+        assert torch.equal(plan_prevs, extras_prev_slot_indices[:extras_count])
+
+    @staticmethod
+    def _assert_padding_row_seed_is_minus_one(
+        *,
+        write_plan: WritePlan,
+        req_pool_indices: torch.Tensor,
+    ) -> None:
+        n_active = int(write_plan.write_num_valid_reqs[0].item())
+        if n_active == 0:
+            return
+        rpi_cpu = req_pool_indices.detach().cpu().tolist()
+        seeds_cpu = (
+            write_plan.write_seed_slot_indices[:n_active].detach().cpu().tolist()
+        )
+        for r in range(min(n_active, len(rpi_cpu))):
+            if rpi_cpu[r] == 0:
+                assert (
+                    seeds_cpu[r] == -1
+                ), f"padding row {r} has seed {seeds_cpu[r]} != -1"
+
+    @staticmethod
+    def _assert_prev_slot_minus_one_iff_chain_head(
+        *,
+        verify_plan: VerifyPlan,
+        swa_window_size: int,
+        derived_verify_count: int,
+    ) -> None:
+        if derived_verify_count == 0:
+            return
+        positions_cpu = (
+            verify_plan.verify_expected_positions[:derived_verify_count]
+            .detach()
+            .cpu()
+            .tolist()
+        )
+        prevs_cpu = (
+            verify_plan.verify_prev_slot_indices[:derived_verify_count]
+            .detach()
+            .cpu()
+            .tolist()
+        )
+        for i, (pos, prev) in enumerate(zip(positions_cpu, prevs_cpu)):
+            if pos == 0:
+                assert (
+                    prev == -1
+                ), f"entry {i} at position 0 must have prev=-1, got {prev}"
+            else:
+                if swa_window_size == 0:
+                    assert (
+                        prev != -1
+                    ), f"FULL entry {i} at position {pos} must have prev != -1, got {prev}"
+
+    @staticmethod
+    def _assert_verify_num_valid_equals_derived_plus_extras(
+        *,
+        verify_plan: VerifyPlan,
+        prefix_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        swa_window_size: int,
+        extras_count: int,
+    ) -> int:
+        rpi_cpu = req_pool_indices.detach().cpu().tolist()
+        pfx_cpu = prefix_lens.detach().cpu().tolist()
+        derived = 0
+        for rpi, pfx in zip(rpi_cpu, pfx_cpu):
+            if rpi == 0:
+                continue
+            if swa_window_size > 0:
+                window_start = max(0, pfx - swa_window_size)
+                derived += max(0, pfx - window_start)
+            else:
+                derived += max(0, pfx)
+        # The plan kernel clamps verify_num_valid to verify_capacity and turns enable
+        # off when (derived + extras) overflows the slot indices buffer. The invariant
+        # must match that: on overflow the kernel records the capacity, on no overflow
+        # it records the exact derived total (so the verify kernel scans every row).
+        verify_capacity = int(verify_plan.verify_slot_indices.shape[0])
+        expected_unclamped = derived + extras_count
+        expected = min(expected_unclamped, verify_capacity)
+        overflow = expected_unclamped > verify_capacity
+        actual = int(verify_plan.verify_num_valid[0].item())
+        assert actual == expected, (
+            f"verify_num_valid {actual} != min(derived {derived} + extras {extras_count}, "
+            f"verify_capacity {verify_capacity}) = {expected}"
+        )
+        enable = int(verify_plan.enable[0].item())
+        expected_enable = 0 if overflow else 1
+        assert enable == expected_enable, (
+            f"verify_plan.enable {enable} != expected {expected_enable} "
+            f"(overflow={overflow}; derived+extras={expected_unclamped}, "
+            f"verify_capacity={verify_capacity})"
+        )
+        return derived
+
+
+class VerifyInvariants:
+    @staticmethod
+    def assert_all(
+        *,
+        canary_buf_before: torch.Tensor,
+        canary_buf_after: torch.Tensor,
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+        plan: VerifyPlan,
+        kernel_kind: CanaryLaunchTag,
+    ) -> None:
+        VerifyInvariants._assert_canary_buf_unchanged(
+            canary_buf_before=canary_buf_before, canary_buf_after=canary_buf_after
+        )
+        VerifyInvariants._assert_violation_count_le_active_entries(
+            log_after=log_after, log_before=log_before, plan=plan
+        )
+        VerifyInvariants._assert_violation_rows_have_valid_slot_and_kernel_kind(
+            log_after=log_after,
+            log_before=log_before,
+            plan=plan,
+            kernel_kind=kernel_kind,
+        )
+        VerifyInvariants._assert_slot_run_counter_incremented_by_active_entries(
+            log_before=log_before, log_after=log_after, plan=plan
+        )
+        VerifyInvariants._assert_kernel_run_counter_incremented_by_one(
+            log_before=log_before, log_after=log_after
+        )
+
+    @staticmethod
+    def _assert_canary_buf_unchanged(
+        *,
+        canary_buf_before: torch.Tensor,
+        canary_buf_after: torch.Tensor,
+    ) -> None:
+        assert torch.equal(
+            canary_buf_before, canary_buf_after
+        ), "verify kernel mutated canary_buf (must be read-only)"
+
+    @staticmethod
+    def _assert_violation_count_le_active_entries(
+        *,
+        log_after: FakeViolationLog,
+        log_before: FakeViolationLog,
+        plan: VerifyPlan,
+    ) -> None:
+        delta = int(log_after.write_index[0].item()) - int(
+            log_before.write_index[0].item()
+        )
+        n_active = int(plan.verify_num_valid[0].item())
+        assert (
+            0 <= delta <= n_active
+        ), f"violation_write_index delta {delta} out of [0, {n_active}]"
+
+    @staticmethod
+    def _assert_violation_rows_have_valid_slot_and_kernel_kind(
+        *,
+        log_after: FakeViolationLog,
+        log_before: FakeViolationLog,
+        plan: VerifyPlan,
+        kernel_kind: CanaryLaunchTag,
+    ) -> None:
+        write_idx_after = int(log_after.write_index[0].item())
+        write_idx_before = int(log_before.write_index[0].item())
+        if write_idx_after == write_idx_before:
+            return
+        ring_capacity = log_after.ring.shape[0]
+        visible_start = write_idx_before
+        visible_end = min(write_idx_after, ring_capacity)
+        if visible_end <= visible_start:
+            return
+        n_active = int(plan.verify_num_valid[0].item())
+        plan_slots = set(plan.verify_slot_indices[:n_active].detach().cpu().tolist())
+        rows = log_after.ring[visible_start:visible_end].detach().cpu()
+        for i in range(rows.shape[0]):
+            kind = int(rows[i, consts.VIOLATION_FIELD_KERNEL_KIND].item())
+            assert kind == int(
+                kernel_kind
+            ), f"row {visible_start + i} kernel_kind {kind} != expected {int(kernel_kind)}"
+            slot = int(rows[i, consts.VIOLATION_FIELD_SLOT_IDX].item())
+            assert (
+                slot in plan_slots
+            ), f"row {visible_start + i} slot {slot} not in plan_slots"
+
+    @staticmethod
+    def _assert_slot_run_counter_incremented_by_active_entries(
+        *,
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+        plan: VerifyPlan,
+    ) -> None:
+        n_active = int(plan.verify_num_valid[0].item())
+        delta = int(log_after.slot_run_counter[0].item()) - int(
+            log_before.slot_run_counter[0].item()
+        )
+        assert (
+            delta == n_active
+        ), f"slot_run_counter delta {delta} != active entries {n_active}"
+
+    @staticmethod
+    def _assert_kernel_run_counter_incremented_by_one(
+        *,
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+    ) -> None:
+        delta = int(log_after.kernel_run_counter[0].item()) - int(
+            log_before.kernel_run_counter[0].item()
+        )
+        assert delta == 1, f"kernel_run_counter delta {delta} != 1"
+
+
+class WriteInvariants:
+    @staticmethod
+    def assert_all(
+        *,
+        canary_buf_before: torch.Tensor,
+        canary_buf_after: torch.Tensor,
+        plan: WritePlan,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+        enable_write_verify_inputs: bool,
+        expected_input_tokens: Optional[torch.Tensor],
+        expected_input_positions: Optional[torch.Tensor],
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+    ) -> None:
+        WriteInvariants._assert_written_slots_token_position_match_input(
+            canary_buf_after=canary_buf_after,
+            plan=plan,
+            input_ids=input_ids,
+            positions=positions,
+            out_cache_loc=out_cache_loc,
+        )
+        WriteInvariants._assert_slot_minus_one_skipped(
+            canary_buf_before=canary_buf_before,
+            canary_buf_after=canary_buf_after,
+            plan=plan,
+            out_cache_loc=out_cache_loc,
+        )
+        WriteInvariants._assert_pseudo_violation_only_on_mismatch(
+            enable_write_verify_inputs=enable_write_verify_inputs,
+            log_before=log_before,
+            log_after=log_after,
+            expected_input_tokens=expected_input_tokens,
+            expected_input_positions=expected_input_positions,
+            input_ids=input_ids,
+            positions=positions,
+            out_cache_loc=out_cache_loc,
+            plan=plan,
+        )
+        WriteInvariants._assert_write_slot_run_counter_incremented(
+            log_before=log_before,
+            log_after=log_after,
+            plan=plan,
+            out_cache_loc=out_cache_loc,
+        )
+        WriteInvariants._assert_write_kernel_run_counter_incremented_by_one(
+            log_before=log_before, log_after=log_after
+        )
+
+    @staticmethod
+    def _assert_written_slots_token_position_match_input(
+        *,
+        canary_buf_after: torch.Tensor,
+        plan: WritePlan,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+    ) -> None:
+        n_active = int(plan.write_num_valid_reqs[0].item())
+        if n_active == 0:
+            return
+        offsets = plan.write_offsets[: n_active + 1].detach().cpu().tolist()
+        total = offsets[n_active]
+        slots_cpu = out_cache_loc[:total].detach().cpu().tolist()
+        tokens_cpu = input_ids[:total].detach().cpu().tolist()
+        pos_cpu = positions[:total].detach().cpu().tolist()
+        view = canary_buf_after.view(torch.int64)
+        for i in range(total):
+            slot = slots_cpu[i]
+            if slot < 0:
+                continue
+            stored_token = int(view[slot, 0].item())
+            stored_position = int(view[slot, 1].item())
+            assert (
+                stored_token == tokens_cpu[i]
+            ), f"slot {slot}: stored token {stored_token} != input {tokens_cpu[i]}"
+            assert (
+                stored_position == pos_cpu[i]
+            ), f"slot {slot}: stored position {stored_position} != input {pos_cpu[i]}"
+
+    @staticmethod
+    def _assert_slot_minus_one_skipped(
+        *,
+        canary_buf_before: torch.Tensor,
+        canary_buf_after: torch.Tensor,
+        plan: WritePlan,
+        out_cache_loc: torch.Tensor,
+    ) -> None:
+        n_active = int(plan.write_num_valid_reqs[0].item())
+        if n_active == 0:
+            return
+        total = int(plan.write_offsets[n_active].item())
+        slots_cpu = out_cache_loc[:total].detach().cpu().tolist()
+        written_slots = {s for s in slots_cpu if s >= 0}
+        view_before = canary_buf_before.view(torch.int64)
+        view_after = canary_buf_after.view(torch.int64)
+        num_slots = canary_buf_after.shape[0]
+        for slot in range(num_slots):
+            if slot in written_slots:
+                continue
+            assert torch.equal(
+                view_before[slot], view_after[slot]
+            ), f"slot {slot} not in out_cache_loc but canary_buf changed"
+
+    @staticmethod
+    def _assert_pseudo_violation_only_on_mismatch(
+        *,
+        enable_write_verify_inputs: bool,
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+        expected_input_tokens: Optional[torch.Tensor],
+        expected_input_positions: Optional[torch.Tensor],
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+        plan: WritePlan,
+    ) -> None:
+        delta = int(log_after.write_index[0].item()) - int(
+            log_before.write_index[0].item()
+        )
+        if not enable_write_verify_inputs:
+            assert (
+                delta == 0
+            ), f"enable_write_verify_inputs=OFF must produce no violations, got {delta}"
+            return
+        if expected_input_tokens is None or expected_input_positions is None:
+            return
+        n_active = int(plan.write_num_valid_reqs[0].item())
+        if n_active == 0:
+            assert delta == 0, f"empty plan produced {delta} violations"
+            return
+        total = int(plan.write_offsets[n_active].item())
+        tok = input_ids[:total].detach().cpu().tolist()
+        pos = positions[:total].detach().cpu().tolist()
+        exp_tok = expected_input_tokens[:total].detach().cpu().tolist()
+        exp_pos = expected_input_positions[:total].detach().cpu().tolist()
+        slots_cpu = out_cache_loc[:total].detach().cpu().tolist()
+        mismatch_entries = sum(
+            1
+            for i in range(total)
+            if slots_cpu[i] >= 0 and (tok[i] != exp_tok[i] or pos[i] != exp_pos[i])
+        )
+        no_mismatch = mismatch_entries == 0
+        if no_mismatch:
+            assert (
+                delta == 0
+            ), f"enable_write_verify_inputs=ON with no mismatch produced {delta} violations"
+        else:
+            assert (
+                delta == mismatch_entries
+            ), f"write input mismatch count {mismatch_entries} produced {delta} violations"
+
+    @staticmethod
+    def _assert_write_slot_run_counter_incremented(
+        *,
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+        plan: WritePlan,
+        out_cache_loc: torch.Tensor,
+    ) -> None:
+        n_active = int(plan.write_num_valid_reqs[0].item())
+        if n_active == 0:
+            delta = int(log_after.slot_run_counter[0].item()) - int(
+                log_before.slot_run_counter[0].item()
+            )
+            assert delta == 0, f"empty plan incremented slot_run_counter by {delta}"
+            return
+        total = int(plan.write_offsets[n_active].item())
+        # The write kernel skips entries where out_cache_loc < 0 (the documented "mark
+        # skip" path used by SWA-translated callers), so the slot_run_counter delta
+        # tracks the count of writeable entries, not the planned total.
+        writeable = int((out_cache_loc[:total] >= 0).sum().item())
+        delta = int(log_after.slot_run_counter[0].item()) - int(
+            log_before.slot_run_counter[0].item()
+        )
+        assert delta == writeable, (
+            f"slot_run_counter delta {delta} != writeable entries {writeable} "
+            f"(total={total}, skipped={total - writeable})"
+        )
+
+    @staticmethod
+    def _assert_write_kernel_run_counter_incremented_by_one(
+        *,
+        log_before: FakeViolationLog,
+        log_after: FakeViolationLog,
+    ) -> None:
+        delta = int(log_after.kernel_run_counter[0].item()) - int(
+            log_before.kernel_run_counter[0].item()
+        )
+        assert delta == 1, f"kernel_run_counter delta {delta} != 1"

python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reference-vs-fused unit tests for the MiniMax-M3 ROCm native MXFP8 ops.
+
+Each fused kernel has a slow PyTorch / dequant-to-bf16 reference; these assert
+the two agree within tolerance:
+
+  * Fused MXFP8 activation quant (Triton)        -> torch reference
+  * Native MXFP8 linear (tl.dot_scaled)          -> dequant-to-bf16 @ matmul
+  * Native MXFP8 MoE (dot_scaled grouped GEMM)   -> dequant-to-bf16 MoE math
+
+ROCm-only. The pure quant test runs on any ROCm arch; the native MXFP8
+``dot_scaled`` linear/MoE tests are gated to CDNA4 gfx95x (the hardware
+microscaling matrix cores) -- gfx942 has no native ``dot_scaled`` MX path.
+
+Run:  pytest python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py -v
+"""
+
+import pytest
+import torch
+
+from sglang.srt.utils import is_hip
+
+if not is_hip():
+    pytest.skip(
+        "MiniMax-M3 native MXFP8 ops are the ROCm path.", allow_module_level=True
+    )
+if not torch.cuda.is_available():
+    pytest.skip("Requires a GPU.", allow_module_level=True)
+
+from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import (  # noqa: E402
+    _mxfp8_dot_scaled_linear,
+    _mxfp8_e4m3_quantize_torch,
+    _mxfp8_e4m3_quantize_triton,
+    dequant_mxfp8_to_bf16,
+)
+
+DEVICE = "cuda"
+
+
+def _gcn_arch() -> str:
+    try:
+        return torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:  # pragma: no cover - no device / non-AMD
+        return ""
+
+
+requires_gfx950 = pytest.mark.skipif(
+    "gfx95" not in _gcn_arch(),
+    reason="native MXFP8 dot_scaled is a CDNA4 (gfx95x) feature; "
+    "gfx942 has no native dot_scaled MX path.",
+)
+
+
+def _relerr(a: torch.Tensor, b: torch.Tensor) -> float:
+    a = a.float()
+    b = b.float()
+    return ((a - b).norm() / (b.norm() + 1e-8)).item()
+
+
+# --------------------------------------------------------------------------- #
+# Fused MXFP8 activation quant (Triton vs torch reference)
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize("shape", [(64, 4096), (1, 6144), (333, 2048)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_mxfp8_quant_triton_matches_torch(shape, dtype):
+    torch.manual_seed(0)
+    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    xq_t, s_t = _mxfp8_e4m3_quantize_torch(x)
+    xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
+    assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
+    # E8M0 block exponents share the round-up ceil(log2(amax/e4m3_max))+127
+    # algorithm; allow at most a 1-step difference at exact powers of two.
+    assert (s_k.int() - s_t.int()).abs().max().item() <= 1
+    # Dequantized values agree to fp8 granularity.
+    deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
+    deq_k = dequant_mxfp8_to_bf16(xq_k, s_k)
+    assert _relerr(deq_k, deq_t) < 1e-2
+
+
+@pytest.mark.parametrize("m,inter", [(8, 512), (65, 2048)])
+@torch.inference_mode()
+def test_minimax_swiglu_mxfp8_quant_matches_unfused_fp32(m, inter):
+    # The fused swiglu+quant kernel keeps the activation in fp32 through the
+    # E8M0 scale selection (no bf16 round-trip; matches the vLLM/ame kernel), so
+    # the reference is the unfused fp32 swiglu followed by MXFP8 quant. Not
+    # bit-identical because the reference quant runs in torch vs the fused triton
+    # path, but numerically equivalent (tight relerr, scales agree within 1 ulp).
+    from sglang.jit_kernel.minimax_m3 import (
+        swiglu_oai_mxfp8_quant,
+        swiglu_oai_split,
+    )
+    from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize
+
+    torch.manual_seed(0)
+    alpha, beta, limit = 1.702, 1.0, 7.0
+    gate_up = torch.randn(m, 2 * inter, device=DEVICE, dtype=torch.bfloat16) * 0.5
+
+    act = swiglu_oai_split(
+        gate_up, alpha=alpha, beta=beta, limit=limit, out_dtype=torch.float32
+    )
+    q_ref, s_ref = mxfp8_e4m3_quantize(act)
+    q, s = swiglu_oai_mxfp8_quant(gate_up, alpha=alpha, beta=beta, limit=limit)
+
+    assert q.shape == q_ref.shape
+    assert s.shape == s_ref.shape
+    # E8M0 block scales agree within one exponent step (last-bit amax differences).
+    assert (s.int() - s_ref.int()).abs().max().item() <= 1
+    assert (
+        _relerr(dequant_mxfp8_to_bf16(q, s), dequant_mxfp8_to_bf16(q_ref, s_ref)) < 1e-2
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Native MXFP8 linear (dot_scaled) vs dequant-to-bf16 matmul
+# --------------------------------------------------------------------------- #
+@requires_gfx950
+@pytest.mark.parametrize("m,n,k", [(64, 256, 128), (37, 512, 256), (1, 6144, 4096)])
+@torch.inference_mode()
+def test_mxfp8_native_linear(m, n, k):
+    torch.manual_seed(0)
+    w_bf16 = torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16)
+    x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) * 0.5
+
+    got = _mxfp8_dot_scaled_linear(x, w_fp8, w_scale)
+    # Reference consumes the SAME quantized weights (isolates activation-quant
+    # noise) -> dequant to bf16, plain matmul.
+    w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
+    ref = torch.nn.functional.linear(x, w_deq).to(x.dtype)
+    assert got.shape == (m, n)
+    assert _relerr(got, ref) < 5e-2
+
+
+# --------------------------------------------------------------------------- #
+# Native MXFP8 MoE (dot_scaled grouped GEMM) vs dequant-to-bf16 MoE math
+# --------------------------------------------------------------------------- #
+def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
+    T, H = x.shape
+    inter = w2.shape[-1]
+    top_k = topk_ids.shape[1]
+    out = torch.zeros(T, H, device=x.device, dtype=torch.float32)
+    for t in range(T):
+        for j in range(top_k):
+            e = int(topk_ids[t, j].item())
+            if e < 0 or e >= w13.shape[0]:
+                continue
+            g1 = x[t].float() @ w13[e].float().T  # [2I]
+            gate = g1[:inter]
+            up = g1[inter:]
+            if limit is not None:
+                gate = gate.clamp(max=limit)
+                up = up.clamp(min=-limit, max=limit)
+            act = gate * torch.sigmoid(alpha * gate) * (up + beta)
+            g2 = act @ w2[e].float().T  # [H]
+            out[t] += topk_weights[t, j].float() * g2
+    return out.to(x.dtype)
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    "T,H,inter,E,top_k", [(8, 256, 512, 8, 2), (1, 512, 256, 16, 4)]
+)
+@torch.inference_mode()
+def test_mxfp8_native_moe(T, H, inter, E, top_k):
+    from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
+        fused_moe_mxfp8_native,
+    )
+
+    torch.manual_seed(0)
+    alpha, beta, limit = 1.702, 1.0, 7.0
+    w13_bf16 = torch.randn(E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w2_bf16 = torch.randn(E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(w13_bf16)
+    w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16)
+
+    x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
+    logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
+    topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
+    topk_weights = topk_weights.to(torch.float32)
+    topk_ids = topk_ids.to(torch.int32)
+
+    got = fused_moe_mxfp8_native(
+        x,
+        w13_fp8,
+        w13_scale,
+        w2_fp8,
+        w2_scale,
+        topk_weights,
+        topk_ids,
+        alpha=alpha,
+        beta=beta,
+        limit=limit,
+    )
+    # Reference consumes the dequantized weights (same bits the kernel reads).
+    w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
+    w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
+    ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids, alpha, beta, limit)
+    assert got.shape == (T, H)
+    assert _relerr(got, ref) < 5e-2
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_mxfp8_native_moe_ep_expert_map_filters_non_local_routes():
+    from sglang.srt.layers.moe.moe_runner.triton_utils.mxfp8_moe_amd_gfx95 import (
+        fused_moe_mxfp8_native,
+    )
+
+    torch.manual_seed(0)
+    T, H, inter = 4, 256, 512
+    local_E = 3
+    alpha, beta, limit = 1.702, 1.0, 7.0
+
+    w13_bf16 = (
+        torch.randn(local_E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    )
+    w2_bf16 = torch.randn(local_E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(w13_bf16)
+    w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16)
+
+    x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
+    topk_ids_global = torch.tensor(
+        [[0, 1, 4], [2, 3, 5], [4, 0, 3], [5, 1, 2]],
+        device=DEVICE,
+        dtype=torch.int32,
+    )
+    topk_weights = torch.tensor(
+        [
+            [0.50, 0.25, 0.25],
+            [0.40, 0.30, 0.30],
+            [0.70, 0.20, 0.10],
+            [0.60, 0.30, 0.10],
+        ],
+        device=DEVICE,
+        dtype=torch.float32,
+    )
+    expert_map = torch.tensor([0, -1, 1, -1, 2, -1], device=DEVICE, dtype=torch.int32)
+
+    got = fused_moe_mxfp8_native(
+        x,
+        w13_fp8,
+        w13_scale,
+        w2_fp8,
+        w2_scale,
+        topk_weights,
+        topk_ids_global,
+        alpha=alpha,
+        beta=beta,
+        limit=limit,
+        expert_map=expert_map,
+    )
+
+    topk_ids_local = expert_map[topk_ids_global.long()]
+    w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
+    w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
+    ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids_local, alpha, beta, limit)
+    assert got.shape == (T, H)
+    assert _relerr(got, ref) < 5e-2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reference tests for MiniMax-M3 ROCm Gemma RMSNorm Triton kernels."""
+
+import pytest
+import torch
+
+from sglang.srt.utils import is_hip
+
+if not is_hip():
+    pytest.skip(
+        "MiniMax-M3 Gemma RMSNorm Triton kernels are ROCm-only.",
+        allow_module_level=True,
+    )
+if not torch.cuda.is_available():
+    pytest.skip("Requires a GPU.", allow_module_level=True)
+
+from sglang.jit_kernel.minimax_m3.rmsnorm import (  # noqa: E402
+    gemma_fused_add_rmsnorm,
+    gemma_rmsnorm,
+)
+
+DEVICE = "cuda"
+EPS = 1e-6
+
+
+def _gemma_rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    orig_dtype = x.dtype
+    x_f = x.float()
+    variance = x_f.pow(2).mean(dim=-1, keepdim=True)
+    out = x_f * torch.rsqrt(variance + EPS)
+    out = out * (1.0 + weight.float())
+    return out.to(orig_dtype)
+
+
+@pytest.mark.parametrize("shape", [(1, 512), (64, 6144), (257, 6144)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_gemma_rmsnorm_matches_reference(shape, dtype):
+    torch.manual_seed(0)
+    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    weight = torch.randn(shape[-1], device=DEVICE, dtype=torch.float32)
+
+    got = gemma_rmsnorm(x, weight, EPS)
+    ref = _gemma_rmsnorm_ref(x, weight)
+
+    torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_gemma_rmsnorm_accepts_strided_2d_input(dtype):
+    torch.manual_seed(0)
+    base = torch.randn(128, 1024, device=DEVICE, dtype=dtype)
+    x = base[:, ::2]
+    weight = torch.randn(x.shape[-1], device=DEVICE, dtype=torch.float32)
+
+    assert not x.is_contiguous()
+    got = gemma_rmsnorm(x, weight, EPS)
+    ref = _gemma_rmsnorm_ref(x, weight)
+
+    torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("shape", [(1, 512), (64, 6144), (257, 6144)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_gemma_fused_add_rmsnorm_matches_reference(shape, dtype):
+    torch.manual_seed(0)
+    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    residual = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    weight = torch.randn(shape[-1], device=DEVICE, dtype=torch.float32)
+
+    got, residual_out = gemma_fused_add_rmsnorm(x, residual, weight, EPS)
+    ref_residual = x + residual
+    ref = _gemma_rmsnorm_ref(ref_residual, weight)
+
+    torch.testing.assert_close(residual_out, ref_residual, atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)

python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py

@@ -0,0 +1,493 @@
+"""
+Correctness tests for the moe_topk_sigmoid JIT kernel.
+
+Validates against a pure-PyTorch reference and, when sgl_kernel is available,
+cross-checks against the AOT implementation.
+"""
+
+import itertools
+import os
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.jit_kernel.moe_topk_sigmoid import topk_sigmoid
+
+try:
+    from sgl_kernel import topk_sigmoid as topk_sigmoid_aot
+
+    AOT_AVAILABLE = True
+except ImportError:
+    AOT_AVAILABLE = False
+
+# ---------------------------------------------------------------------------
+# CI / full-range helpers
+# ---------------------------------------------------------------------------
+
+_is_ci = (
+    os.getenv("CI", "false").lower() == "true"
+    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+)
+
+# Power-of-2 configs covered by static dispatch (num_experts 1–256)
+# Plus 48 (non-power-of-2) to exercise the fallback path
+NUM_TOKENS_FULL = [1, 16, 128, 512, 1024, 2048]
+NUM_TOKENS_CI = [1, 128, 1024]
+
+NUM_EXPERTS_FULL = [16, 32, 64, 128, 256, 48]  # 48 = fallback path
+NUM_EXPERTS_CI = [16, 64, 48]
+
+TOPK_FULL = [1, 2, 4, 8]
+TOPK_CI = [1, 4]
+
+DTYPES_FULL = [torch.float32]
+DTYPES_CI = [torch.float32, torch.bfloat16]
+
+NUM_TOKENS = NUM_TOKENS_CI if _is_ci else NUM_TOKENS_FULL
+NUM_EXPERTS = NUM_EXPERTS_CI if _is_ci else NUM_EXPERTS_FULL
+TOPK_LIST = TOPK_CI if _is_ci else TOPK_FULL
+DTYPES = DTYPES_CI if _is_ci else DTYPES_FULL
+
+
+# ---------------------------------------------------------------------------
+# Pure-PyTorch reference
+# ---------------------------------------------------------------------------
+
+
+def grouped_topk_gpu(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: Optional[int] = None,
+    topk_group: Optional[int] = None,
+    num_fused_shared_experts: int = 0,
+    routed_scaling_factor: Optional[float] = None,
+    apply_routed_scaling_factor_on_output: Optional[bool] = False,
+    scoring_func: str = "softmax",
+):
+
+    # Scoring function: softmax or sigmoid
+    if scoring_func == "softmax":
+        scores = torch.softmax(gating_output, dim=-1)
+    elif scoring_func == "sigmoid":
+        scores = gating_output.sigmoid()
+    else:
+        raise ValueError(f"Unsupported scoring function: {scoring_func}")
+
+    num_token = scores.shape[0]
+    num_experts = scores.shape[1]
+    group_scores = (
+        scores.view(num_token, num_expert_group, -1).max(dim=-1).values
+    )  # [n, n_group]
+
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
+        1
+    ]  # [n, top_k_group]
+
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
+        .reshape(num_token, -1)
+    )  # [n, e]
+    tmp_scores = scores.masked_fill(
+        ~score_mask.bool(), float("-inf")
+    )  # [n, e] - use -inf like VLLM
+
+    topk_weights, topk_ids = torch.topk(
+        tmp_scores,
+        k=topk,
+        dim=-1,
+        sorted=(True if num_fused_shared_experts > 0 else True),
+    )
+
+    if num_fused_shared_experts:
+        topk_ids[:, -1] = torch.randint(
+            low=num_experts,
+            high=num_experts + num_fused_shared_experts,
+            size=(topk_ids.size(0),),
+            dtype=topk_ids.dtype,
+            device=topk_ids.device,
+        )
+        if routed_scaling_factor is not None:
+            topk_weights[:, -1] = (
+                topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+            )
+
+    if renormalize:
+        topk_weights_sum = (
+            topk_weights.sum(dim=-1, keepdim=True)
+            if num_fused_shared_experts == 0
+            else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
+        )
+        topk_weights = topk_weights / topk_weights_sum
+        if apply_routed_scaling_factor_on_output:
+            topk_weights *= routed_scaling_factor
+
+    topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+    return topk_weights, topk_ids
+
+
+def biased_grouped_topk_impl(
+    gating_output: torch.Tensor,
+    correction_bias: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: Optional[int] = None,
+    topk_group: Optional[int] = None,
+    num_fused_shared_experts: int = 0,
+    routed_scaling_factor: Optional[float] = None,
+    apply_routed_scaling_factor_on_output: Optional[bool] = False,
+):
+    scores = gating_output.sigmoid()
+
+    num_token = scores.shape[0]
+    num_experts = scores.shape[1]
+    scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
+    group_scores = (
+        scores_for_choice.view(num_token, num_expert_group, -1)
+        .topk(2, dim=-1)[0]
+        .sum(dim=-1)
+    )  # [n, n_group]
+
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
+        1
+    ]  # [n, top_k_group]
+
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
+        .reshape(num_token, -1)
+    )  # [n, e]
+    tmp_scores = scores_for_choice.masked_fill(
+        ~score_mask.bool(), float("-inf")
+    )  # [n, e]
+
+    _, topk_ids = torch.topk(
+        tmp_scores,
+        k=topk,
+        dim=-1,
+        sorted=(True if num_fused_shared_experts > 0 else True),
+    )
+    topk_weights = scores.gather(1, topk_ids)
+
+    if num_fused_shared_experts:
+        topk_ids[:, -1] = torch.randint(
+            low=num_experts,
+            high=num_experts + num_fused_shared_experts,
+            size=(topk_ids.size(0),),
+            dtype=topk_ids.dtype,
+            device=topk_ids.device,
+        )
+        if routed_scaling_factor is not None:
+            topk_weights[:, -1] = (
+                topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+            )
+    if renormalize:
+        topk_weights_sum = (
+            topk_weights.sum(dim=-1, keepdim=True)
+            if num_fused_shared_experts == 0
+            else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
+        )
+        topk_weights = topk_weights / topk_weights_sum
+        if apply_routed_scaling_factor_on_output:
+            topk_weights *= routed_scaling_factor
+
+    topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+    return topk_weights, topk_ids
+
+
+def topk_sigmoid_torch_ref(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    correction_bias: torch.Tensor | None,
+    num_fused_shared_experts: int = 0,
+    routed_scaling_factor: float = 1.0,
+    apply_routed_scaling_factor_on_output: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Reference: sigmoid → (add bias) → topk → (renormalize).
+    Indices are selected on biased scores; weights are the unbiased sigmoid values.
+    """
+    num_experts = gating_output.shape[1]
+    scores = gating_output.float().sigmoid()
+    biased = scores if correction_bias is None else scores + correction_bias.float()
+    _, ref_ids = torch.topk(biased, k=topk, dim=-1)
+    ref_weights = scores.gather(1, ref_ids)
+    if num_fused_shared_experts > 0:
+        ref_ids[:, -1] = torch.randint(
+            low=num_experts,
+            high=num_experts + num_fused_shared_experts,
+            size=(ref_ids.size(0),),
+            dtype=ref_ids.dtype,
+            device=ref_ids.device,
+        )
+        ref_weights[:, -1] = ref_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+    if renormalize:
+        topk_weights_sum = (
+            ref_weights.sum(dim=-1, keepdim=True)
+            if num_fused_shared_experts == 0
+            else ref_weights[:, :-1].sum(dim=-1, keepdim=True)
+        )
+        ref_weights = ref_weights / topk_weights_sum
+        if apply_routed_scaling_factor_on_output:
+            ref_weights *= routed_scaling_factor
+    return ref_weights.float(), ref_ids.int()
+
+
+def topk_sigmoid_grouped_ref(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    correction_bias: torch.Tensor | None,
+    num_fused_shared_experts: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if correction_bias is not None:
+        return biased_grouped_topk_impl(
+            gating_output,
+            correction_bias,
+            topk,
+            renormalize,
+            num_expert_group=1,
+            topk_group=1,
+            num_fused_shared_experts=num_fused_shared_experts,
+            routed_scaling_factor=1.0,
+            apply_routed_scaling_factor_on_output=True,
+        )
+    else:
+        return grouped_topk_gpu(
+            gating_output,
+            topk,
+            renormalize,
+            num_expert_group=1,
+            topk_group=1,
+            num_fused_shared_experts=num_fused_shared_experts,
+            routed_scaling_factor=1.0,
+            apply_routed_scaling_factor_on_output=True,
+            scoring_func="sigmoid",
+        )
+
+
+def topk_sigmoid_ref(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    correction_bias: torch.Tensor | None,
+    num_fused_shared_experts: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return topk_sigmoid_torch_ref(
+        gating_output, topk, renormalize, correction_bias, num_fused_shared_experts
+    )
+
+
+# ---------------------------------------------------------------------------
+# Correctness: JIT vs PyTorch reference
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(itertools.product(NUM_TOKENS, NUM_EXPERTS, TOPK_LIST)),
+)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("renormalize", [False, True])
+def test_topk_sigmoid_vs_ref(num_tokens, num_experts, topk, dtype, renormalize):
+    if topk > num_experts:
+        pytest.skip("topk > num_experts")
+
+    torch.manual_seed(num_tokens * num_experts)
+    gating = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+
+    topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid(topk_w, topk_i, gating, renormalize=renormalize)
+
+    ref_w, ref_i = topk_sigmoid_ref(gating, topk, renormalize, correction_bias=None)
+
+    # Compare sorted weights (indices may differ for ties when dtype != float32)
+    assert torch.allclose(
+        topk_w.sort(dim=-1)[0],
+        ref_w.sort(dim=-1)[0],
+        atol=1e-3,
+        rtol=1e-3,
+    ), f"Weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk}, renorm={renormalize})"
+    # Exact index match is only reliable for float32 (fp16/bf16 tie-breaking may differ)
+    if dtype == torch.float32:
+        assert torch.equal(
+            topk_i, ref_i
+        ), f"Index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
+
+
+# ---------------------------------------------------------------------------
+# Correctness: with correction_bias
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(itertools.product(NUM_TOKENS, NUM_EXPERTS, TOPK_LIST)),
+)
+@pytest.mark.parametrize("renormalize", [False, True])
+def test_topk_sigmoid_with_correction_bias(num_tokens, num_experts, topk, renormalize):
+    if topk > num_experts:
+        pytest.skip("topk > num_experts")
+
+    torch.manual_seed(num_tokens + num_experts + topk)
+    gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
+    bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
+
+    topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid(topk_w, topk_i, gating, renormalize=renormalize, correction_bias=bias)
+
+    ref_w, ref_i = topk_sigmoid_ref(gating, topk, renormalize, correction_bias=bias)
+
+    assert torch.allclose(
+        topk_w, ref_w, atol=1e-3, rtol=1e-3
+    ), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
+    assert torch.equal(
+        topk_i, ref_i
+    ), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
+
+
+# ---------------------------------------------------------------------------
+# Correctness: with fused shared experts
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(itertools.product(NUM_TOKENS, NUM_EXPERTS, TOPK_LIST)),
+)
+@pytest.mark.parametrize("renormalize", [False, True])
+def test_topk_sigmoid_with_fused_shared_experts(
+    num_tokens, num_experts, topk, renormalize
+):
+    if topk + 1 > num_experts:
+        pytest.skip("topk > num_experts")
+
+    torch.manual_seed(num_tokens + num_experts)
+    gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
+    bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
+
+    topk_w = torch.empty((num_tokens, topk + 1), dtype=torch.float32, device="cuda")
+    topk_i = torch.empty((num_tokens, topk + 1), dtype=torch.int32, device="cuda")
+    topk_sigmoid(
+        topk_w,
+        topk_i,
+        gating,
+        renormalize=renormalize,
+        correction_bias=bias,
+        num_fused_shared_experts=1,
+    )
+
+    ref_w, ref_i = topk_sigmoid_ref(
+        gating, topk + 1, renormalize, correction_bias=bias, num_fused_shared_experts=1
+    )
+
+    assert torch.allclose(
+        topk_w, ref_w, atol=1e-3, rtol=1e-3
+    ), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
+    assert torch.equal(
+        topk_i, ref_i
+    ), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
+
+
+# ---------------------------------------------------------------------------
+# Renormalization: weights should sum to 1 per row
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("num_tokens, num_experts, topk", [(128, 64, 4), (1, 8, 2)])
+def test_renormalize_sums_to_one(num_tokens, num_experts, topk):
+    gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
+    topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid(topk_w, topk_i, gating, renormalize=True)
+    row_sums = topk_w.sum(dim=-1)
+    torch.testing.assert_close(
+        row_sums, torch.ones(num_tokens, device="cuda"), rtol=1e-4, atol=1e-4
+    )
+
+
+# ---------------------------------------------------------------------------
+# Output shape and dtype
+# ---------------------------------------------------------------------------
+
+
+def test_output_shapes_and_dtypes():
+    num_tokens, num_experts, topk = 64, 128, 4
+    gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
+    topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid(topk_w, topk_i, gating)
+
+    assert topk_w.shape == (num_tokens, topk)
+    assert topk_i.shape == (num_tokens, topk)
+    assert topk_w.dtype == torch.float32
+    assert topk_i.dtype == torch.int32
+
+
+# ---------------------------------------------------------------------------
+# Fallback path (non-power-of-2 experts)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("num_experts", [48, 96])
+def test_fallback_non_power_of_two(num_experts):
+    num_tokens, topk = 64, 2
+    gating = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda")
+    topk_w = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid(topk_w, topk_i, gating, renormalize=True)
+
+    # Weights should be positive and sum to 1
+    assert torch.all(topk_w > 0)
+    torch.testing.assert_close(
+        topk_w.sum(dim=-1), torch.ones(num_tokens, device="cuda"), rtol=1e-4, atol=1e-4
+    )
+
+
+# ---------------------------------------------------------------------------
+# Cross-validation against AOT sgl_kernel
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available")
+@pytest.mark.parametrize(
+    "num_tokens, num_experts, topk",
+    list(itertools.product([1, 128, 1024], [8, 64, 128], [1, 4])),
+)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("renormalize", [False, True])
+def test_topk_sigmoid_vs_aot(num_tokens, num_experts, topk, dtype, renormalize):
+    if topk > num_experts:
+        pytest.skip("topk > num_experts")
+
+    torch.manual_seed(42)
+    gating = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+
+    topk_w_jit = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i_jit = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid(topk_w_jit, topk_i_jit, gating, renormalize=renormalize)
+
+    topk_w_aot = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_i_aot = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    topk_sigmoid_aot(topk_w_aot, topk_i_aot, gating, renormalize=renormalize)
+
+    assert torch.allclose(
+        topk_w_jit, topk_w_aot, atol=1e-3, rtol=1e-3
+    ), f"JIT vs AOT weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
+    assert torch.equal(
+        topk_i_jit, topk_i_aot
+    ), f"JIT vs AOT index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/jit_kernel/tests/utils.py

@@ -0,0 +1,64 @@
+import sys
+from typing import Callable, List, Optional, Sequence
+
+import pytest
+
+from sglang.jit_kernel.mp import multigpu_launch
+
+
+def multigpu_pytest_main(
+    name: str,
+    file: str,
+    num_gpus: Sequence[int],
+    *,
+    pre_launch_fn: Optional[Callable[[List[int]], None]] = None,
+    timeout: Optional[int] = 600,
+) -> None:
+    """cudalib-style multi-GPU pytest entry point.
+
+    Drop this at the bottom of a test file::
+
+        multigpu_pytest_main(__name__, __file__, num_gpus=range(2, 9))
+
+    When the file is run with ``python <file>``, it relaunches itself under
+    ``torchrun --nproc_per_node=N <file>`` for each N in ``num_gpus``. Inside
+    each worker, ``pytest.main([file, ...forwarded_args])`` runs the collected
+    tests. Pass ``--num-gpu 2,4`` on the command line to override ``num_gpus``.
+
+    ``pre_launch_fn`` (kw-only) runs once in the outer process before any
+    torchrun child starts, receiving the runnable world sizes. Use it for
+    parallel JIT precompilation so torchrun children hit a warm disk cache
+    instead of compiling kernels on first call.
+
+    ``timeout`` (kw-only, seconds) bounds each per-world-size torchrun
+    invocation. The default budget covers the cold-cache first invocation
+    (the worker pays the full triton + cutlass JIT compile cost, 60-180s
+    observed on H200) plus the nightly full sweep, which runs every size x
+    dtype x algo x graph-mode parametrisation rather than the reduced in-CI
+    range. A worker that exceeds the budget is killed and the run fails. Pass
+    ``None`` to wait indefinitely.
+    """
+
+    def inner() -> int:
+        # CI's run_unittest_files invokes `python3 <file> -f` (legacy
+        # unittest failfast). Translate to pytest's `-x` so it survives.
+        pytest_args = ["-x" if a == "-f" else a for a in sys.argv[1:]]
+        # Dump all thread stacks (every rank; stderr is not redirected) if a
+        # single test exceeds half the harness budget, so a hung collective
+        # is attributable from the CI log before the outer timeout kills the
+        # process group. Non-fatal: the test keeps running after the dump.
+        dump_after = (timeout // 2) if timeout else 300
+        return pytest.main(
+            [file, "-o", f"faulthandler_timeout={dump_after}"] + pytest_args
+        )
+
+    return multigpu_launch(
+        name,
+        file,
+        num_gpus,
+        env_key="_IS_TEST_MULTIGPU_SGLANG_JIT_KERNEL",
+        inner=inner,
+        kind="test",
+        pre_launch_fn=pre_launch_fn,
+        timeout=timeout,
+    )