chore: import upstream snapshot with attribution

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"