chore: import upstream snapshot with attribution

python/sglang/jit_kernel/kv_canary/consts.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from enum import IntEnum, IntFlag
+from typing import Final
+
+CANARY_CHAIN_ANCHOR: Final[int] = 0xC0FFEE1234567890
+
+# Mirrors SGLang's ReqToTokenPool contract: req_pool_idx 0 is the CUDA-graph padding row, while real
+# request rows start at 1.
+REQ_POOL_IDX_PADDING: Final[int] = 0
+
+# Mirrors SGLang's TokenToKVPoolAllocator contract: token-to-KV slot 0 is reserved for padded-token dummy
+# writes. Since req_to_token stores token-to-KV slot ids and is zero-initialized, canary slot 0 is skipped
+# instead of treating unfilled entries as real KV slots.
+TOKEN_TO_KV_SLOT_PADDING: Final[int] = 0
+
+CANARY_FIELDS_PER_SLOT: Final[int] = 4
+CANARY_FIELD_TOKEN: Final[int] = 0
+CANARY_FIELD_POSITION: Final[int] = 1
+CANARY_FIELD_PREV_HASH: Final[int] = 2
+CANARY_FIELD_REAL_KV_HASH: Final[int] = 3
+
+VIOLATION_FIELDS: Final[int] = 8
+VIOLATION_FIELD_KERNEL_KIND: Final[int] = 0
+VIOLATION_FIELD_SLOT_IDX: Final[int] = 1
+VIOLATION_FIELD_POSITION: Final[int] = 2
+VIOLATION_FIELD_STORED_TOKEN: Final[int] = 3
+VIOLATION_FIELD_EXPECTED_TOKEN: Final[int] = 4
+VIOLATION_FIELD_STORED_CHAIN_HASH: Final[int] = 5
+VIOLATION_FIELD_EXPECTED_AUX: Final[int] = 6
+VIOLATION_FIELD_FAIL_REASON_BITS: Final[int] = 7
+
+
+class FailReason(IntFlag):
+    VERIFY_CHAIN_HASH_MISMATCH = 1 << 0
+    VERIFY_POSITION_MISMATCH = 1 << 1
+    VERIFY_REAL_KV_HASH_MISMATCH = 1 << 2
+    WRITE_TOKEN_MISMATCH = 1 << 3
+    WRITE_POSITION_MISMATCH = 1 << 4
+    VERIFY_TOKEN_MISMATCH = 1 << 5
+
+
+MAX_REAL_KV_SOURCES: Final[int] = 4
+
+REAL_KV_SOURCE_FIELDS_PER_ENTRY: Final[int] = 3
+REAL_KV_SOURCE_FIELD_PAGE_SIZE: Final[int] = 0
+REAL_KV_SOURCE_FIELD_NUM_BYTES_PER_TOKEN: Final[int] = 1
+REAL_KV_SOURCE_FIELD_READ_BYTES: Final[int] = 2
+
+
+class RealKvHashMode(IntEnum):
+    NONE = 0
+    PARTIAL = 1
+    ALL = 2
+
+
+_U64_MASK: int = (1 << 64) - 1
+
+
+def splitmix64(value: int) -> int:
+    x = value & _U64_MASK
+    x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & _U64_MASK
+    x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & _U64_MASK
+    return (x ^ (x >> 31)) & _U64_MASK
+
+
+def splitmix64_mix3(a: int, b: int, c: int) -> int:
+    h = splitmix64(a & _U64_MASK)
+    h = splitmix64(h ^ (b & _U64_MASK))
+    h = splitmix64(h ^ (c & _U64_MASK))
+    return h

python/sglang/jit_kernel/kv_canary/plan/__init__.py

@@ -0,0 +1 @@
+from sglang.jit_kernel.kv_canary.plan.api import launch_canary_plan_kernels

python/sglang/jit_kernel/kv_canary/plan/api.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.plan.entries_kernel import (
+    launch_plan_entries_kernel,
+)
+from sglang.jit_kernel.kv_canary.plan.offsets_kernel import (
+    _PLAN_BS_BLOCK_SIZE,
+    launch_plan_offsets_kernel,
+)
+from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+
+
+def launch_canary_plan_kernels(
+    *,
+    verify_plan_out: VerifyPlan,
+    write_plan_out: WritePlan,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    swa_window_size: int,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    verify_capacity: int,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor],
+    kv_token_id_vs_position_offset: int,
+) -> None:
+    """Fill verify_plan_out + write_plan_out from normalized canary plan inputs.
+
+    For each req r with req_pool_indices[r] != 0 (0 = padding sentinel):
+
+    - **Verify entries**: one per pos in [window_start, prefix_lens[r]), where window_start = max(0,
+      prefix_lens[r] - swa_window_size) if SWA else 0. slot_idx = req_to_token[req_pool_indices[r], pos]
+      (SWA-translated via full_to_swa_index_mapping if non-None); prev_slot_idx =
+      req_to_token[req_pool_indices[r], pos-1] for pos > 0, else -1. (SWA windows do NOT reset the chain —
+      the writer chains across the entire prefix; sweep verify within an SWA window dereferences the real
+      predecessor for chain-link reconstruction.) Expected-token gather: when
+      ``req_to_verify_expected_tokens`` is supplied, ``expected_input_id =
+      req_to_verify_expected_tokens[rp, pos + kv_token_id_vs_position_offset]`` when ``0 <= pos +
+      kv_token_id_vs_position_offset < req_to_verify_expected_tokens_valid_lens[r]``, else the ``-1``
+      sentinel (which the verify kernel treats as "skip token-id check").
+    - **Write metadata** (when extend_seq_lens[r] > 0): contribute extend_seq_lens[r] to the per-req
+      write count (for write_offsets cumsum). Per-req chain seed = req_to_token[req_pool_indices[r],
+      prefix_lens[r]-1] (SWA-translated), or -1 if prefix_lens[r] == 0. Per-token write data
+      (input_ids / positions / out_cache_loc) is NOT materialized here — launch_canary_write_kernel
+      reads it directly from ForwardBatch via write_offsets.
+
+    Args:
+        verify_plan_out: Pre-allocated VerifyPlan; filled in-place.
+        write_plan_out: Pre-allocated WritePlan; filled in-place.
+        req_pool_indices: Per-row ReqToTokenPool row index, shape [bs], int64. 0 is the padding sentinel.
+        prefix_lens: Per-req prefix length already written before this step, shape [bs], int64.
+        extend_seq_lens: Per-req tokens being written this step, shape [bs], int64.
+        req_to_token: ReqToTokenPool.req_to_token; full-pool slot index table, shape [max_reqs, max_seq_len],
+            int32.
+        swa_window_size: 0 for the FULL canary group; positive window length for the SWA group.
+        full_to_swa_index_mapping: SWA LUT, shape [full_pool_size + 1], int64, or None. Required (non-None) iff
+            swa_window_size > 0. Used to translate verify slot indices and chain-seed slot indices at plan time.
+            Loaded element-typed via Triton ``tl.load``; intermediate translated slot values are int64 inside the
+            kernel and stored in the int64 plan schema.
+        verify_capacity: Length of verify_plan_out.verify_*; on overflow the offsets kernel clears
+            verify_enable and plan_entries skips the scatter.
+        req_to_verify_expected_tokens: Optional source-of-truth token pool, shape [max_reqs, max_context_len],
+            int32. When supplied, the plan kernel gathers expected_input_id for each verify entry from
+            ``[rp, pos + kv_token_id_vs_position_offset]``; when None, every entry gets the ``-1`` sentinel.
+        req_to_verify_expected_tokens_valid_lens: Per-req snapshot length on ``req_to_verify_expected_tokens``,
+            shape [bs], int64. Required iff ``req_to_verify_expected_tokens`` is set. Reads past
+            ``valid_lens[r]`` skip the gather (emit ``-1``) — this is what makes the plan kernel correct in the
+            presence of EAGLE draft / verify positions written past the committed history, and across pool
+            rows recycled from a longer previous owner whose stale tail still lives at high indices.
+        kv_token_id_vs_position_offset: Per-buffer-group logical-position offset applied to ``pos`` before
+            indexing ``req_to_verify_expected_tokens``. 0 for target pools; +1 for EAGLE draft.
+
+    Implementation:
+        - Two sub-kernels launched in sequence:
+          1. Triton ``_plan_offsets_kernel`` (1-D grid ``(1,)``, single program over all ``bs`` reqs):
+             reads req_pool_indices[r], prefix_lens[r], extend_seq_lens[r] for each r; computes
+             verify_count = (prefix_lens - window_start) and write_count = extend_seq_lens (both 0 if rp == 0
+             padding); gathers seed_slot_full = req_to_token[rp, prefix_lens - 1] (or -1 if prefix_lens == 0),
+             SWA-translates seed_slot via full_to_swa_index_mapping[seed_slot_full] if non-None; runs
+             block-level cumsum (``tl.cumsum``) to produce verify_offsets[_PLAN_BS_BLOCK_SIZE + 1] and
+             write_plan_out.write_offsets[write_req_capacity + 1] in-place; scatters write_seed slots; writes
+             scalar totals ``verify_plan_out.verify_num_valid`` and ``write_plan_out.write_num_valid_reqs``.
+          2. CUDA ``plan_entries_persistent_kernel`` (1-D persistent grid sized to ``num_sms *
+             kBlocksPerSm`` blocks of ``kBlockSize`` threads), wrapped by Python
+             ``launch_plan_entries_kernel``: each thread grid-strides over ``tid ∈ [0, total_verify)``,
+             locates its owning req via ``find_req_id`` (binary search on verify_offsets), computes
+             out_position = window_start[req_id] + (tid - verify_offsets[req_id]), gathers slot =
+             req_to_token[rp, out_position] (SWA-translated when ``HAS_SWA_LUT``), prev_slot =
+             req_to_token[rp, out_position - 1] when out_position > 0 (also translated) else -1, and
+             scatters (slot, position, prev_slot) into verify_plan_out at flat index tid.
+        - All output tensors are addressed at addresses baked into the cuda-graph capture.
+
+    Calling contract:
+        - Pure side-effect; no host work, no D2H.
+        - Safe in cuda-graph capture; caller refills all input tensors in-place before replay.
+        - The wrapper launches the plan sub-kernels needed to fill both plans end-to-end.
+        - Padding rows contribute zero entries.
+
+    Pinned by Python reference
+    :func:`sglang.jit_kernel.kv_canary.plan_ref.launch_canary_plan_kernels_torch_reference`; both the Triton
+    offsets kernel and the CUDA JIT entries kernel must match byte-for-byte.
+    """
+    bs = int(req_pool_indices.shape[0])
+    if bs > _PLAN_BS_BLOCK_SIZE:
+        raise ValueError(
+            f"kv-canary: launch_canary_plan_kernels supports at most bs={_PLAN_BS_BLOCK_SIZE} reqs per launch, "
+            f"got bs={bs}. Bump _PLAN_BS_BLOCK_SIZE if real workloads need this."
+        )
+    if swa_window_size > 0 and full_to_swa_index_mapping is None:
+        raise ValueError(
+            "kv-canary: launch_canary_plan_kernels requires full_to_swa_index_mapping when swa_window_size > 0"
+        )
+
+    device = verify_plan_out.verify_slot_indices.device
+    verify_offsets_scratch = torch.empty(
+        _PLAN_BS_BLOCK_SIZE + 1, dtype=torch.int64, device=device
+    )
+
+    plan_verify_capacity = int(verify_plan_out.verify_slot_indices.shape[0])
+    if verify_capacity != plan_verify_capacity:
+        raise ValueError(
+            f"kv-canary: launch_canary_plan_kernels verify_capacity={verify_capacity} does not match "
+            f"verify_plan_out.verify_slot_indices.shape[0]={plan_verify_capacity}"
+        )
+
+    write_plan_out.write_offsets.zero_()
+
+    launch_plan_offsets_kernel(
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        extend_seq_lens=extend_seq_lens,
+        req_to_token=req_to_token,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        out_verify_offsets_scratch=verify_offsets_scratch,
+        out_write_offsets=write_plan_out.write_offsets,
+        out_write_seed_slot_indices=write_plan_out.write_seed_slot_indices,
+        out_verify_num_valid=verify_plan_out.verify_num_valid,
+        out_verify_enable=verify_plan_out.enable,
+        out_write_num_valid_reqs=write_plan_out.write_num_valid_reqs,
+        swa_window_size=int(swa_window_size),
+        verify_capacity=verify_capacity,
+    )
+
+    launch_plan_entries_kernel(
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        req_to_token=req_to_token,
+        full_to_swa_index_mapping=full_to_swa_index_mapping,
+        verify_offsets_scratch=verify_offsets_scratch,
+        verify_enable=verify_plan_out.enable,
+        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
+        out_verify_slot_indices=verify_plan_out.verify_slot_indices,
+        out_verify_expected_tokens=verify_plan_out.verify_expected_tokens,
+        out_verify_expected_positions=verify_plan_out.verify_expected_positions,
+        out_verify_prev_slot_indices=verify_plan_out.verify_prev_slot_indices,
+        kv_token_id_vs_position_offset=int(kv_token_id_vs_position_offset),
+        swa_window_size=int(swa_window_size),
+    )

python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_plan_entries_module(
+    has_swa_lut: bool, has_verify_expected_token_pool: bool
+) -> Module:
+    args = make_cpp_args(has_swa_lut, has_verify_expected_token_pool)
+    return load_jit(
+        "kv_canary_plan_entries",
+        *args,
+        cuda_files=["kv_canary/canary_plan_entries.cuh"],
+        cuda_wrappers=[
+            ("plan_entries", f"PlanEntriesKernel<{args}>::run"),
+        ],
+    )
+
+
+def launch_plan_entries_kernel(
+    *,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    verify_offsets_scratch: torch.Tensor,
+    verify_enable: torch.Tensor,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor],
+    out_verify_slot_indices: torch.Tensor,
+    out_verify_expected_tokens: torch.Tensor,
+    out_verify_expected_positions: torch.Tensor,
+    out_verify_prev_slot_indices: torch.Tensor,
+    kv_token_id_vs_position_offset: int,
+    swa_window_size: int,
+) -> None:
+    has_swa_lut = full_to_swa_index_mapping is not None
+    has_verify_expected_token_pool = req_to_verify_expected_tokens is not None
+    if (
+        has_verify_expected_token_pool
+        and req_to_verify_expected_tokens_valid_lens is None
+    ):
+        raise ValueError(
+            "kv-canary: launch_plan_entries_kernel requires "
+            "req_to_verify_expected_tokens_valid_lens when req_to_verify_expected_tokens is set"
+        )
+    module = _jit_plan_entries_module(has_swa_lut, has_verify_expected_token_pool)
+    module.plan_entries(
+        req_pool_indices,
+        prefix_lens,
+        req_to_token,
+        full_to_swa_index_mapping,
+        verify_offsets_scratch,
+        verify_enable,
+        req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens,
+        out_verify_slot_indices,
+        out_verify_expected_tokens,
+        out_verify_expected_positions,
+        out_verify_prev_slot_indices,
+        int(kv_token_id_vs_position_offset),
+        int(swa_window_size),
+    )

python/sglang/jit_kernel/kv_canary/plan/offsets_kernel.py

@@ -0,0 +1,441 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.kv_canary.consts import (
+    REQ_POOL_IDX_PADDING,
+    TOKEN_TO_KV_SLOT_PADDING,
+)
+from sglang.jit_kernel.kv_canary.plan.utils import (
+    _compute_window_start,
+    _require_1d,
+    _require_2d,
+    _require_dtype,
+    _require_len,
+    _require_min_len,
+    _require_same_device,
+    _resolve_swa_lut,
+    _swa_translate_tile,
+)
+from sglang.jit_kernel.kv_canary.verify import _assert_contiguous
+
+# Upper bound on bs for _plan_offsets_kernel's block-level cumsum. Reqs larger than this exceed Triton's
+# single-program tl.cumsum reach. Increase if real workloads ever push past it; the cap is intentionally
+# generous so the wrapper never silently truncates.
+_PLAN_BS_BLOCK_SIZE: int = 4096
+
+
+def launch_plan_offsets_kernel(
+    *,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    out_verify_offsets_scratch: torch.Tensor,
+    out_write_offsets: torch.Tensor,
+    out_write_seed_slot_indices: torch.Tensor,
+    out_verify_num_valid: torch.Tensor,
+    out_verify_enable: torch.Tensor,
+    out_write_num_valid_reqs: torch.Tensor,
+    swa_window_size: int,
+    verify_capacity: int,
+) -> None:
+    bs = int(req_pool_indices.shape[0])
+    lut_tensor, lut_len, has_swa_lut = _resolve_swa_lut(
+        full_to_swa_index_mapping, out_verify_offsets_scratch.device
+    )
+    req_to_token_stride0 = int(req_to_token.stride(0))
+    write_offsets_len = int(out_write_offsets.shape[0])
+    write_req_capacity = int(out_write_seed_slot_indices.shape[0])
+
+    _validate_offsets_kernel_inputs(
+        req_pool_indices=req_pool_indices,
+        prefix_lens=prefix_lens,
+        extend_seq_lens=extend_seq_lens,
+        req_to_token=req_to_token,
+        lut_tensor=lut_tensor,
+        out_verify_offsets_scratch=out_verify_offsets_scratch,
+        out_write_offsets=out_write_offsets,
+        out_write_seed_slot_indices=out_write_seed_slot_indices,
+        out_verify_num_valid=out_verify_num_valid,
+        out_verify_enable=out_verify_enable,
+        out_write_num_valid_reqs=out_write_num_valid_reqs,
+        bs=bs,
+        req_to_token_stride0=req_to_token_stride0,
+        lut_len=lut_len,
+        has_swa_lut=has_swa_lut,
+        write_offsets_len=write_offsets_len,
+        write_req_capacity=write_req_capacity,
+        verify_capacity=verify_capacity,
+    )
+
+    _plan_offsets_kernel[(1,)](
+        req_pool_indices,
+        prefix_lens,
+        extend_seq_lens,
+        req_to_token,
+        lut_tensor,
+        out_verify_offsets_scratch,
+        out_write_offsets,
+        out_write_seed_slot_indices,
+        out_verify_num_valid,
+        out_verify_enable,
+        out_write_num_valid_reqs,
+        bs,
+        req_to_token_stride0,
+        lut_len,
+        BS_BLOCK=_PLAN_BS_BLOCK_SIZE,
+        SWA_WINDOW=int(swa_window_size),
+        HAS_SWA_LUT=has_swa_lut,
+        WRITE_OFFSETS_LEN=write_offsets_len,
+        WRITE_REQ_CAPACITY=write_req_capacity,
+        VERIFY_CAPACITY=verify_capacity,
+        REQ_POOL_IDX_PADDING=REQ_POOL_IDX_PADDING,
+        TOKEN_TO_KV_SLOT_PADDING=TOKEN_TO_KV_SLOT_PADDING,
+    )
+
+
+def _validate_offsets_kernel_inputs(
+    *,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    lut_tensor: torch.Tensor,
+    out_verify_offsets_scratch: torch.Tensor,
+    out_write_offsets: torch.Tensor,
+    out_write_seed_slot_indices: torch.Tensor,
+    out_verify_num_valid: torch.Tensor,
+    out_verify_enable: torch.Tensor,
+    out_write_num_valid_reqs: torch.Tensor,
+    bs: int,
+    req_to_token_stride0: int,
+    lut_len: int,
+    has_swa_lut: bool,
+    write_offsets_len: int,
+    write_req_capacity: int,
+    verify_capacity: int,
+) -> None:
+    _assert_contiguous(req_pool_indices, "req_pool_indices")
+    _assert_contiguous(prefix_lens, "prefix_lens")
+    _assert_contiguous(extend_seq_lens, "extend_seq_lens")
+    _assert_contiguous(req_to_token, "req_to_token")
+    _assert_contiguous(lut_tensor, "lut_tensor")
+    _assert_contiguous(out_verify_offsets_scratch, "out_verify_offsets_scratch")
+    _assert_contiguous(out_write_offsets, "out_write_offsets")
+    _assert_contiguous(out_write_seed_slot_indices, "out_write_seed_slot_indices")
+    _assert_contiguous(out_verify_num_valid, "out_verify_num_valid")
+    _assert_contiguous(out_verify_enable, "out_verify_enable")
+    _assert_contiguous(out_write_num_valid_reqs, "out_write_num_valid_reqs")
+
+    _require_dtype(req_pool_indices, "req_pool_indices", torch.int64)
+    _require_dtype(prefix_lens, "prefix_lens", torch.int64)
+    _require_dtype(extend_seq_lens, "extend_seq_lens", torch.int64)
+    _require_dtype(req_to_token, "req_to_token", torch.int32)
+    _require_dtype(lut_tensor, "lut_tensor", torch.int64)
+    _require_dtype(
+        out_verify_offsets_scratch, "out_verify_offsets_scratch", torch.int64
+    )
+    _require_dtype(out_write_offsets, "out_write_offsets", torch.int64)
+    _require_dtype(
+        out_write_seed_slot_indices, "out_write_seed_slot_indices", torch.int64
+    )
+    _require_dtype(out_verify_num_valid, "out_verify_num_valid", torch.int32)
+    _require_dtype(out_verify_enable, "out_verify_enable", torch.int32)
+    _require_dtype(out_write_num_valid_reqs, "out_write_num_valid_reqs", torch.int32)
+
+    if bs < 0 or bs > _PLAN_BS_BLOCK_SIZE:
+        raise ValueError(
+            f"kv-canary: offsets kernel bs must be in [0, {_PLAN_BS_BLOCK_SIZE}], got {bs}"
+        )
+    if write_offsets_len <= 0:
+        raise ValueError(
+            f"kv-canary: write_offsets_len must be positive, got {write_offsets_len}"
+        )
+    if write_req_capacity < 0:
+        raise ValueError(
+            f"kv-canary: write_req_capacity must be non-negative, got {write_req_capacity}"
+        )
+    if verify_capacity < 0:
+        raise ValueError(
+            f"kv-canary: verify_capacity must be non-negative, got {verify_capacity}"
+        )
+    if req_to_token_stride0 <= 0:
+        raise ValueError(
+            f"kv-canary: req_to_token_stride0 must be positive, got {req_to_token_stride0}"
+        )
+    if lut_len < 0:
+        raise ValueError(f"kv-canary: lut_len must be non-negative, got {lut_len}")
+    if not isinstance(has_swa_lut, bool):
+        raise ValueError(
+            f"kv-canary: has_swa_lut must be bool, got {type(has_swa_lut).__name__}"
+        )
+    if has_swa_lut and lut_len <= 0:
+        raise ValueError("kv-canary: lut_len must be positive when has_swa_lut is True")
+    if not has_swa_lut and lut_len != 0:
+        raise ValueError("kv-canary: lut_len must be 0 when has_swa_lut is False")
+
+    _require_len(req_pool_indices, "req_pool_indices", bs)
+    _require_len(prefix_lens, "prefix_lens", bs)
+    _require_len(extend_seq_lens, "extend_seq_lens", bs)
+    _require_2d(req_to_token, "req_to_token")
+    _require_min_len(lut_tensor, "lut_tensor", max(lut_len, 1))
+    _require_min_len(
+        out_verify_offsets_scratch,
+        "out_verify_offsets_scratch",
+        _PLAN_BS_BLOCK_SIZE + 1,
+    )
+    _require_len(out_write_offsets, "out_write_offsets", write_offsets_len)
+    _require_len(
+        out_write_seed_slot_indices,
+        "out_write_seed_slot_indices",
+        write_req_capacity,
+    )
+    _require_len(out_verify_num_valid, "out_verify_num_valid", 1)
+    _require_len(out_verify_enable, "out_verify_enable", 1)
+    _require_len(out_write_num_valid_reqs, "out_write_num_valid_reqs", 1)
+    _require_1d(lut_tensor, "lut_tensor")
+
+    if write_offsets_len != write_req_capacity + 1:
+        raise ValueError(
+            f"kv-canary: write_offsets_len must equal write_req_capacity + 1, got "
+            f"{write_offsets_len} and {write_req_capacity}"
+        )
+    if bs > write_req_capacity:
+        raise ValueError(
+            f"kv-canary: bs={bs} exceeds write_req_capacity={write_req_capacity}"
+        )
+    if req_to_token_stride0 != int(req_to_token.stride(0)):
+        raise ValueError(
+            f"kv-canary: req_to_token_stride0={req_to_token_stride0} does not match "
+            f"req_to_token.stride(0)={int(req_to_token.stride(0))}"
+        )
+
+    _require_same_device(
+        out_verify_offsets_scratch,
+        "out_verify_offsets_scratch",
+        (
+            (req_pool_indices, "req_pool_indices"),
+            (prefix_lens, "prefix_lens"),
+            (extend_seq_lens, "extend_seq_lens"),
+            (req_to_token, "req_to_token"),
+            (lut_tensor, "lut_tensor"),
+            (out_write_offsets, "out_write_offsets"),
+            (out_write_seed_slot_indices, "out_write_seed_slot_indices"),
+            (out_verify_num_valid, "out_verify_num_valid"),
+            (out_verify_enable, "out_verify_enable"),
+            (out_write_num_valid_reqs, "out_write_num_valid_reqs"),
+        ),
+    )
+
+
+@triton.jit
+def _plan_offsets_kernel(
+    # Input pointers.
+    req_pool_indices_ptr,
+    prefix_lens_ptr,
+    extend_seq_lens_ptr,
+    req_to_token_ptr,
+    full_to_swa_lut_ptr,
+    # Output pointers.
+    out_verify_offsets_ptr,
+    out_write_offsets_ptr,
+    out_write_seed_slot_indices_ptr,
+    out_verify_num_valid_ptr,
+    out_verify_enable_ptr,
+    out_write_num_valid_reqs_ptr,
+    # Runtime sizes.
+    bs,
+    req_to_token_stride0,
+    swa_lut_len,
+    # Compile-time constants.
+    BS_BLOCK: tl.constexpr,
+    SWA_WINDOW: tl.constexpr,
+    HAS_SWA_LUT: tl.constexpr,
+    WRITE_OFFSETS_LEN: tl.constexpr,
+    WRITE_REQ_CAPACITY: tl.constexpr,
+    VERIFY_CAPACITY: tl.constexpr,
+    REQ_POOL_IDX_PADDING: tl.constexpr,
+    TOKEN_TO_KV_SLOT_PADDING: tl.constexpr,
+):
+    bs_offs = tl.arange(0, BS_BLOCK)  # [BS_BLOCK]
+    bs_mask = bs_offs < bs  # [BS_BLOCK] bool
+
+    # Per-req inputs (int64 for canary-owned metadata; req_to_token keeps its pool dtype).
+    rpi = tl.load(
+        req_pool_indices_ptr + bs_offs, mask=bs_mask, other=REQ_POOL_IDX_PADDING
+    )  # [BS_BLOCK]
+    prefix_lens = tl.load(
+        prefix_lens_ptr + bs_offs, mask=bs_mask, other=0
+    )  # [BS_BLOCK]
+    extend_lens = tl.load(
+        extend_seq_lens_ptr + bs_offs, mask=bs_mask, other=0
+    )  # [BS_BLOCK]
+
+    is_active = (rpi != REQ_POOL_IDX_PADDING) & bs_mask  # [BS_BLOCK] bool
+    has_prefix = is_active & (prefix_lens > 0)  # [BS_BLOCK] bool
+
+    window_starts = _compute_window_start(prefix_lens, SWA_WINDOW)  # [BS_BLOCK]
+
+    verify_lens = prefix_lens - window_starts  # [BS_BLOCK]
+    verify_lens = tl.where(verify_lens > 0, verify_lens, 0)
+    verify_lens = tl.where(is_active, verify_lens, 0)
+    verify_exclusive, total_verify = _exclusive_offsets_and_total(verify_lens)
+
+    write_lens = tl.where(extend_lens > 0, extend_lens, 0)  # [BS_BLOCK]
+    write_lens = tl.where(is_active, write_lens, 0)
+    write_exclusive, total_write = _exclusive_offsets_and_total(write_lens)
+
+    _plan_verify_offsets(
+        verify_exclusive,
+        total_verify,
+        bs_offs,
+        bs_mask,
+        out_verify_offsets_ptr,
+        out_verify_num_valid_ptr,
+        out_verify_enable_ptr,
+        bs,
+        VERIFY_CAPACITY,
+    )
+    _plan_write_offsets(
+        rpi,
+        prefix_lens,
+        write_lens,
+        write_exclusive,
+        total_write,
+        has_prefix,
+        bs_offs,
+        bs_mask,
+        req_to_token_ptr,
+        full_to_swa_lut_ptr,
+        out_write_offsets_ptr,
+        out_write_seed_slot_indices_ptr,
+        out_write_num_valid_reqs_ptr,
+        bs,
+        req_to_token_stride0,
+        swa_lut_len,
+        BS_BLOCK,
+        HAS_SWA_LUT,
+        WRITE_OFFSETS_LEN,
+        WRITE_REQ_CAPACITY,
+        TOKEN_TO_KV_SLOT_PADDING,
+    )
+
+
+@triton.jit
+def _exclusive_offsets_and_total(lens):
+    inclusive = tl.cumsum(lens, axis=0)
+    return inclusive - lens, tl.sum(lens, axis=0)
+
+
+@triton.jit
+def _plan_verify_offsets(
+    verify_exclusive,
+    total_verify,
+    bs_offs,
+    bs_mask,
+    out_verify_offsets_ptr,
+    out_verify_num_valid_ptr,
+    out_verify_enable_ptr,
+    bs,
+    VERIFY_CAPACITY: tl.constexpr,
+):
+    tl.store(
+        out_verify_offsets_ptr + bs_offs,
+        verify_exclusive.to(tl.int64),
+        mask=bs_mask,
+    )
+    tl.store(out_verify_offsets_ptr + bs, total_verify.to(tl.int64))
+
+    # Scalar writes: out_verify_num_valid is clamped to the verify_capacity tensor extent so the verify kernel
+    # never indexes past the buffer; enable carries the overflow bit (0 when total_verify > capacity) so the
+    # verify kernel skips the whole launch and the host can warn-log this step.
+    overflow = total_verify > VERIFY_CAPACITY  # scalar bool
+    enable = tl.where(overflow, 0, 1)  # scalar
+    clamped = tl.where(overflow, VERIFY_CAPACITY, total_verify)  # scalar
+    tl.store(out_verify_num_valid_ptr, clamped.to(tl.int32))
+    tl.store(out_verify_enable_ptr, tl.full((), enable, tl.int32))
+
+
+@triton.jit
+def _plan_write_offsets(
+    rpi,
+    prefix_lens,
+    write_lens,
+    write_exclusive,
+    total_write,
+    has_prefix,
+    bs_offs,
+    bs_mask,
+    req_to_token_ptr,
+    full_to_swa_lut_ptr,
+    out_write_offsets_ptr,
+    out_write_seed_slot_indices_ptr,
+    out_write_num_valid_reqs_ptr,
+    bs,
+    req_to_token_stride0,
+    swa_lut_len,
+    BS_BLOCK: tl.constexpr,
+    HAS_SWA_LUT: tl.constexpr,
+    WRITE_OFFSETS_LEN: tl.constexpr,
+    WRITE_REQ_CAPACITY: tl.constexpr,
+    TOKEN_TO_KV_SLOT_PADDING: tl.constexpr,
+):
+    has_write_contribution = has_prefix & (write_lens > 0)  # [BS_BLOCK] bool
+
+    # Seed slot per req. prefix_lens == 0 means no prefix → -1 sentinel. Padding row → no write contribution
+    # → -1 sentinel either way; we also mask write_lens onto seed below to match the ref's "no write → -1".
+    safe_prefix_pos = tl.where(prefix_lens > 0, prefix_lens - 1, 0)  # [BS_BLOCK]
+    stride_i64 = req_to_token_stride0  # scalar
+    seed_full = tl.load(  # [BS_BLOCK]
+        req_to_token_ptr + rpi.to(tl.int64) * stride_i64 + safe_prefix_pos.to(tl.int64),
+        mask=has_prefix,
+        other=TOKEN_TO_KV_SLOT_PADDING,
+    )
+
+    if HAS_SWA_LUT:
+        seed_translated = _swa_translate_tile(  # [BS_BLOCK]
+            seed_full,
+            has_prefix,
+            full_to_swa_lut_ptr,
+            swa_lut_len,
+        )
+    else:
+        seed_translated = seed_full
+
+    # Reqs with no write contribution should expose seed = -1 (ref's _seed_slot is masked by write_lens > 0).
+    minus_one = tl.full((BS_BLOCK,), -1, dtype=seed_translated.dtype)  # [BS_BLOCK]
+    seed_slot = tl.where(
+        has_write_contribution, seed_translated, minus_one
+    )  # [BS_BLOCK]
+
+    write_offsets_mask = bs_offs < WRITE_OFFSETS_LEN  # [BS_BLOCK] bool
+    tl.store(
+        out_write_offsets_ptr + bs_offs,
+        write_exclusive.to(tl.int64),
+        mask=write_offsets_mask & bs_mask,
+    )
+
+    # Store the [bs] slot of out_write_offsets (one element past the last per-req entry).
+    # out_write_offsets has length WRITE_OFFSETS_LEN = write_req_capacity + 1; only store if in range.
+    write_tail_in_range = bs < WRITE_OFFSETS_LEN  # scalar bool
+    tl.store(
+        out_write_offsets_ptr + bs,
+        total_write.to(tl.int64),
+        mask=write_tail_in_range,
+    )
+
+    # Scatter seed slots (capped to write_req_capacity).
+    seed_mask = bs_mask & (bs_offs < WRITE_REQ_CAPACITY)  # [BS_BLOCK] bool
+    tl.store(
+        out_write_seed_slot_indices_ptr + bs_offs,
+        seed_slot.to(tl.int64),
+        mask=seed_mask,
+    )
+
+    tl.store(out_write_num_valid_reqs_ptr, tl.full((), bs, tl.int32))

python/sglang/jit_kernel/kv_canary/plan/utils.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+def _resolve_swa_lut(
+    lut: Optional[torch.Tensor], device: torch.device
+) -> tuple[torch.Tensor, int, bool]:
+    """Return the (tensor, length, has_lut) triple to launch the plan kernel with.
+
+    Triton requires a valid tensor pointer at every kernel-arg slot even when ``HAS_SWA_LUT`` is False, so
+    when the caller passes ``None`` we substitute a one-element sentinel tensor and set ``lut_len=0``;
+    the kernel's constexpr branch guarantees no dereference happens. Dtype matches the production LUT
+    (int64) so Triton ``tl.load`` element typing stays consistent.
+    """
+    if lut is not None:
+        return lut, int(lut.shape[0]), True
+    return torch.zeros(1, dtype=torch.int64, device=device), 0, False
+
+
+def _require_dtype(tensor: torch.Tensor, name: str, dtype: torch.dtype) -> None:
+    if tensor.dtype != dtype:
+        raise ValueError(
+            f"kv-canary: {name} must have dtype {dtype}, got {tensor.dtype}"
+        )
+
+
+def _require_1d(tensor: torch.Tensor, name: str) -> None:
+    if tensor.ndim != 1:
+        raise ValueError(
+            f"kv-canary: {name} must be 1-D, got shape {tuple(tensor.shape)}"
+        )
+
+
+def _require_2d(tensor: torch.Tensor, name: str) -> None:
+    if tensor.ndim != 2:
+        raise ValueError(
+            f"kv-canary: {name} must be 2-D, got shape {tuple(tensor.shape)}"
+        )
+
+
+def _require_len(tensor: torch.Tensor, name: str, expected: int) -> None:
+    _require_1d(tensor=tensor, name=name)
+    actual = int(tensor.shape[0])
+    if actual != expected:
+        raise ValueError(f"kv-canary: {name} length must be {expected}, got {actual}")
+
+
+def _require_min_len(tensor: torch.Tensor, name: str, minimum: int) -> None:
+    _require_1d(tensor=tensor, name=name)
+    actual = int(tensor.shape[0])
+    if actual < minimum:
+        raise ValueError(f"kv-canary: {name} length must be >= {minimum}, got {actual}")
+
+
+def _require_same_device(
+    reference: torch.Tensor,
+    reference_name: str,
+    tensors: tuple[tuple[torch.Tensor, str], ...],
+) -> None:
+    for tensor, name in tensors:
+        if tensor.device != reference.device:
+            raise ValueError(
+                f"kv-canary: {name} must be on {reference_name}'s device "
+                f"{reference.device}, got {tensor.device}"
+            )
+
+
+@triton.jit
+def _compute_window_start(prefix_lens, SWA_WINDOW: tl.constexpr):
+    """Per-req window start: max(prefix_lens - SWA_WINDOW, 0) when SWA, else 0.
+    Works for tile and scalar inputs (broadcasts via prefix_lens shape).
+    """
+    if SWA_WINDOW > 0:
+        clipped = prefix_lens - SWA_WINDOW
+        return tl.where(clipped > 0, clipped, 0)
+    else:
+        return prefix_lens - prefix_lens
+
+
+@triton.jit
+def _swa_translate_tile(raw, mask, lut_ptr, lut_len):
+    """SWA-translate a tile of slot indices. Sentinels (raw < 0) are passed through unchanged.
+
+    ``lut_len`` is the LUT's length (Python int from the host wrapper); when 0 the LUT is unused (the caller
+    will only enter this branch when HAS_SWA_LUT is True, so lut_len is always > 0 in practice).
+    """
+    sentinel = raw < 0
+    safe = tl.where(sentinel, 0, raw)
+    if lut_len > 0:
+        safe = tl.where(safe >= lut_len, lut_len - 1, safe)
+    xlat = tl.load(lut_ptr + safe, mask=mask & (~sentinel), other=0)
+    return tl.where(sentinel, raw, xlat)

python/sglang/jit_kernel/kv_canary/plan_ref.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.consts import REQ_POOL_IDX_PADDING
+from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+
+
+def launch_canary_plan_kernels_torch_reference(
+    *,
+    verify_plan_out: VerifyPlan,
+    write_plan_out: WritePlan,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    swa_window_size: int,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    verify_capacity: int,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor],
+    kv_token_id_vs_position_offset: int,
+) -> None:
+    """Python reference for :func:`launch_canary_plan_kernels`. Same signature & byte-equal semantics."""
+    bs = int(req_pool_indices.shape[0])
+    work_device = torch.device("cpu")
+
+    plan_verify_capacity = int(verify_plan_out.verify_slot_indices.shape[0])
+    if verify_capacity != plan_verify_capacity:
+        raise ValueError(
+            f"kv-canary: launch_canary_plan_kernels_torch_reference verify_capacity={verify_capacity} does not "
+            f"match verify_plan_out.verify_slot_indices.shape[0]={plan_verify_capacity}"
+        )
+    write_req_capacity = int(write_plan_out.write_seed_slot_indices.shape[0])
+
+    req_pool_indices_host = req_pool_indices.detach().to(
+        device=work_device, dtype=torch.int64
+    )
+    prefix_lens_host = prefix_lens.detach().to(device=work_device, dtype=torch.int64)
+    extend_seq_lens_host = extend_seq_lens.detach().to(
+        device=work_device, dtype=torch.int64
+    )
+    req_to_token_host = req_to_token.detach().to(device=work_device, dtype=torch.int64)
+
+    lut: Optional[torch.Tensor] = None
+    if full_to_swa_index_mapping is not None:
+        lut = full_to_swa_index_mapping.detach().to(device=work_device)
+
+    expected_token_pool_host: Optional[torch.Tensor] = None
+    req_to_verify_expected_tokens_valid_lens_host: Optional[torch.Tensor] = None
+    if req_to_verify_expected_tokens is not None:
+        expected_token_pool_host = req_to_verify_expected_tokens.detach().to(
+            device=work_device, dtype=torch.int64
+        )
+        if req_to_verify_expected_tokens_valid_lens is None:
+            raise ValueError(
+                "kv-canary: launch_canary_plan_kernels_torch_reference requires "
+                "req_to_verify_expected_tokens_valid_lens when req_to_verify_expected_tokens is set"
+            )
+        req_to_verify_expected_tokens_valid_lens_host = (
+            req_to_verify_expected_tokens_valid_lens.detach().to(
+                device=work_device, dtype=torch.int64
+            )
+        )
+
+    total_verify = _materialize_verify_entries(
+        verify_plan_out=verify_plan_out,
+        req_pool_indices_host=req_pool_indices_host,
+        prefix_lens_host=prefix_lens_host,
+        req_to_token_host=req_to_token_host,
+        swa_window_size=swa_window_size,
+        lut=lut,
+        verify_capacity=verify_capacity,
+        work_device=work_device,
+        bs=bs,
+        expected_token_pool_host=expected_token_pool_host,
+        req_to_verify_expected_tokens_valid_lens_host=req_to_verify_expected_tokens_valid_lens_host,
+        kv_token_id_vs_position_offset=int(kv_token_id_vs_position_offset),
+    )
+
+    _materialize_write_metadata(
+        write_plan_out=write_plan_out,
+        req_pool_indices_host=req_pool_indices_host,
+        prefix_lens_host=prefix_lens_host,
+        extend_seq_lens_host=extend_seq_lens_host,
+        req_to_token_host=req_to_token_host,
+        lut=lut,
+        write_req_capacity=write_req_capacity,
+        work_device=work_device,
+        bs=bs,
+    )
+
+    _write_num_valid_and_enable(
+        verify_plan_out=verify_plan_out,
+        requested=total_verify,
+        verify_capacity=verify_capacity,
+    )
+
+
+def _write_num_valid_and_enable(
+    *,
+    verify_plan_out: VerifyPlan,
+    requested: int,
+    verify_capacity: int,
+) -> None:
+    overflow = requested > verify_capacity
+    clamped = verify_capacity if overflow else requested
+    enable = 0 if overflow else 1
+    verify_plan_out.verify_num_valid.fill_(int(clamped))
+    verify_plan_out.enable.fill_(int(enable))
+
+
+def _swa_translate_slot(*, slot: int, lut: torch.Tensor) -> int:
+    if slot < 0:
+        return slot
+    lut_len = int(lut.shape[0])
+    if slot >= lut_len:
+        raise ValueError(
+            f"kv-canary: SWA slot {slot} is outside full_to_swa_index_mapping length {lut_len}"
+        )
+    return int(lut[slot].item())
+
+
+def _materialize_verify_entries(
+    *,
+    verify_plan_out: VerifyPlan,
+    req_pool_indices_host: torch.Tensor,
+    prefix_lens_host: torch.Tensor,
+    req_to_token_host: torch.Tensor,
+    swa_window_size: int,
+    lut: Optional[torch.Tensor],
+    verify_capacity: int,
+    work_device: torch.device,
+    bs: int,
+    expected_token_pool_host: Optional[torch.Tensor],
+    req_to_verify_expected_tokens_valid_lens_host: Optional[torch.Tensor],
+    kv_token_id_vs_position_offset: int,
+) -> int:
+    out_slots: list[int] = []
+    out_positions: list[int] = []
+    out_expected_input_ids: list[int] = []
+    out_prev_slots: list[int] = []
+
+    for r in range(bs):
+        rpi = int(req_pool_indices_host[r].item())
+        prefix_len = int(prefix_lens_host[r].item())
+
+        if rpi == REQ_POOL_IDX_PADDING:
+            continue
+
+        if swa_window_size > 0:
+            window_start = max(0, prefix_len - swa_window_size)
+        else:
+            window_start = 0
+        verify_len = max(0, prefix_len - window_start)
+
+        valid_len_r = (
+            int(req_to_verify_expected_tokens_valid_lens_host[r].item())
+            if req_to_verify_expected_tokens_valid_lens_host is not None
+            else 0
+        )
+
+        for j in range(verify_len):
+            position = window_start + j
+            slot_full = int(req_to_token_host[rpi, position].item())
+
+            if lut is not None:
+                slot = _swa_translate_slot(slot=slot_full, lut=lut)
+            else:
+                slot = slot_full
+
+            prev_position = position - 1
+            if prev_position < 0:
+                prev_slot = -1
+            else:
+                prev_slot_full = int(req_to_token_host[rpi, prev_position].item())
+                if lut is not None:
+                    prev_slot = _swa_translate_slot(slot=prev_slot_full, lut=lut)
+                else:
+                    prev_slot = prev_slot_full
+
+            expected_input_id = -1
+            if expected_token_pool_host is not None:
+                sot_pos = position + kv_token_id_vs_position_offset
+                if 0 <= sot_pos < valid_len_r:
+                    expected_input_id = int(
+                        expected_token_pool_host[rpi, sot_pos].item()
+                    )
+
+            out_slots.append(slot)
+            out_positions.append(position)
+            out_expected_input_ids.append(expected_input_id)
+            out_prev_slots.append(prev_slot)
+
+    total_verify = len(out_slots)
+    if total_verify == 0:
+        return 0
+
+    # On overflow CUDA plan_entries skips scatter (verify_enable=0); mirror that.
+    if total_verify > verify_capacity:
+        return total_verify
+
+    slots_t = torch.tensor(out_slots, dtype=torch.int64, device=work_device)
+    positions_t = torch.tensor(out_positions, dtype=torch.int64, device=work_device)
+    expected_input_ids_t = torch.tensor(
+        out_expected_input_ids, dtype=torch.int64, device=work_device
+    )
+    prev_slots_t = torch.tensor(out_prev_slots, dtype=torch.int64, device=work_device)
+
+    verify_plan_out.verify_slot_indices[:total_verify].copy_(
+        slots_t.to(verify_plan_out.verify_slot_indices.dtype).to(
+            verify_plan_out.verify_slot_indices.device
+        )
+    )
+    verify_plan_out.verify_expected_tokens[:total_verify].copy_(
+        expected_input_ids_t.to(verify_plan_out.verify_expected_tokens.dtype).to(
+            verify_plan_out.verify_expected_tokens.device
+        )
+    )
+    verify_plan_out.verify_expected_positions[:total_verify].copy_(
+        positions_t.to(verify_plan_out.verify_expected_positions.dtype).to(
+            verify_plan_out.verify_expected_positions.device
+        )
+    )
+    verify_plan_out.verify_prev_slot_indices[:total_verify].copy_(
+        prev_slots_t.to(verify_plan_out.verify_prev_slot_indices.dtype).to(
+            verify_plan_out.verify_prev_slot_indices.device
+        )
+    )
+
+    return total_verify
+
+
+def _materialize_write_metadata(
+    *,
+    write_plan_out: WritePlan,
+    req_pool_indices_host: torch.Tensor,
+    prefix_lens_host: torch.Tensor,
+    extend_seq_lens_host: torch.Tensor,
+    req_to_token_host: torch.Tensor,
+    lut: Optional[torch.Tensor],
+    write_req_capacity: int,
+    work_device: torch.device,
+    bs: int,
+) -> None:
+    out_write_offsets_len = int(write_plan_out.write_offsets.shape[0])
+    max_seq_len = int(req_to_token_host.shape[1])
+
+    write_offsets_list: list[int] = []
+    seed_slots_list: list[int] = []
+
+    running_offset = 0
+    for r in range(bs):
+        write_offsets_list.append(running_offset)
+
+        rpi = int(req_pool_indices_host[r].item())
+        extend_len = int(extend_seq_lens_host[r].item())
+
+        if rpi == REQ_POOL_IDX_PADDING or extend_len <= 0:
+            write_len = 0
+        else:
+            write_len = max(0, extend_len)
+
+        running_offset += write_len
+
+    write_offsets_list.append(running_offset)
+
+    copy_len = min(bs + 1, out_write_offsets_len)
+    write_offsets_t = torch.tensor(
+        write_offsets_list[:copy_len], dtype=torch.int64, device=work_device
+    )
+    write_plan_out.write_offsets[:copy_len].copy_(
+        write_offsets_t.to(write_plan_out.write_offsets.dtype).to(
+            write_plan_out.write_offsets.device
+        )
+    )
+    if copy_len < out_write_offsets_len:
+        write_plan_out.write_offsets[copy_len:].zero_()
+
+    capped_reqs = min(bs, write_req_capacity)
+    for r in range(capped_reqs):
+        rpi = int(req_pool_indices_host[r].item())
+        prefix_len = int(prefix_lens_host[r].item())
+        extend_len = int(extend_seq_lens_host[r].item())
+
+        if rpi == REQ_POOL_IDX_PADDING or extend_len <= 0:
+            seed_slots_list.append(-1)
+            continue
+
+        if prefix_len <= 0:
+            seed_slots_list.append(-1)
+            continue
+
+        safe_seed_pos = min(prefix_len - 1, max(max_seq_len - 1, 0))
+        seed_slot_full = int(req_to_token_host[rpi, safe_seed_pos].item())
+
+        if lut is not None:
+            seed_slot = _swa_translate_slot(slot=seed_slot_full, lut=lut)
+        else:
+            seed_slot = seed_slot_full
+
+        seed_slots_list.append(seed_slot)
+
+    if len(seed_slots_list) > 0:
+        seed_slots_t = torch.tensor(
+            seed_slots_list, dtype=torch.int64, device=work_device
+        )
+        write_plan_out.write_seed_slot_indices[:capped_reqs].copy_(
+            seed_slots_t.to(write_plan_out.write_seed_slot_indices.dtype).to(
+                write_plan_out.write_seed_slot_indices.device
+            )
+        )
+
+    write_plan_out.write_num_valid_reqs.fill_(int(bs))

python/sglang/jit_kernel/kv_canary/scatter_req_token_ids.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_SCATTER_TOKEN_BLOCK: int = 256
+# Upper bound on bs+1 the kernel can scan per program. Owner-req lookup uses an
+# outer-product tile of shape ``[TOKEN_BLOCK, BATCH_BLOCK]``; keep this small so
+# the tile stays in registers (256 x 512 = 128 KiB i1, well below the SM trap-
+# inducing budget that bites at the 1M cell mark).
+_SCATTER_BATCH_BLOCK: int = 512
+
+
+def launch_scatter_req_token_ids_kernel(
+    *,
+    flat_in: torch.Tensor,
+    offsets: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    pool_out: torch.Tensor,
+) -> None:
+    """Scatter a flat per-req int64 object sequence into a 2-D int32 pool.
+
+    For each global object index ``t`` in ``[0, total_tokens)``:
+
+    - find ``r`` = largest req index s.t. ``offsets[r] <= t``
+    - ``pos = t - offsets[r]``
+    - ``rp = req_pool_indices[r]``
+    - if ``pos < pool_max_context_len``:
+      ``pool_out[rp, pos] = flat_in[t].to(int32)``
+
+    Args:
+        flat_in: ``[total_tokens]`` int64 device tensor of objects, flattened
+            per-req in req order.
+        offsets: ``[bs + 1]`` int64 device tensor (host-computed cumsum of per-req
+            lengths). ``offsets[bs] == total_tokens``.
+        req_pool_indices: ``[bs]`` int64 device tensor of pool row indices.
+        pool_out: ``[max_reqs, max_context_len]`` int32 device tensor of objects.
+            Mutated in-place; rows not addressed by ``req_pool_indices`` are untouched.
+
+    Implementation notes:
+        - Linear scan over ``offsets`` (``BATCH_BLOCK >= bs + 1``); fits easily in
+          registers for the workloads kv-canary handles (``bs <= a few thousand``).
+    """
+    if flat_in.dim() != 1:
+        raise ValueError(
+            f"kv-canary: scatter_req_token_ids flat_in must be 1-D, got shape "
+            f"{tuple(flat_in.shape)}"
+        )
+    if offsets.dim() != 1:
+        raise ValueError(
+            f"kv-canary: scatter_req_token_ids offsets must be 1-D, got shape "
+            f"{tuple(offsets.shape)}"
+        )
+    if req_pool_indices.dim() != 1:
+        raise ValueError(
+            f"kv-canary: scatter_req_token_ids req_pool_indices must be 1-D, got shape "
+            f"{tuple(req_pool_indices.shape)}"
+        )
+    if pool_out.dim() != 2:
+        raise ValueError(
+            f"kv-canary: scatter_req_token_ids pool_out must be 2-D, got shape "
+            f"{tuple(pool_out.shape)}"
+        )
+    if flat_in.dtype != torch.int64:
+        raise TypeError(
+            f"kv-canary: scatter_req_token_ids flat_in must be int64, got "
+            f"{flat_in.dtype}"
+        )
+    if offsets.dtype != torch.int64:
+        raise TypeError(
+            f"kv-canary: scatter_req_token_ids offsets must be int64, got "
+            f"{offsets.dtype}"
+        )
+    if req_pool_indices.dtype != torch.int64:
+        raise TypeError(
+            f"kv-canary: scatter_req_token_ids req_pool_indices must be int64, got "
+            f"{req_pool_indices.dtype}"
+        )
+    if pool_out.dtype != torch.int32:
+        raise TypeError(
+            f"kv-canary: scatter_req_token_ids pool_out must be int32, got "
+            f"{pool_out.dtype}"
+        )
+
+    bs = int(req_pool_indices.shape[0])
+    if int(offsets.shape[0]) != bs + 1:
+        raise ValueError(
+            f"kv-canary: scatter_req_token_ids offsets length {offsets.shape[0]} != "
+            f"bs+1 ({bs + 1})"
+        )
+    if bs + 1 > _SCATTER_BATCH_BLOCK:
+        raise ValueError(
+            f"kv-canary: scatter_req_token_ids bs+1={bs + 1} exceeds BATCH_BLOCK="
+            f"{_SCATTER_BATCH_BLOCK}; bump _SCATTER_BATCH_BLOCK if real workloads need this"
+        )
+
+    num_tokens = int(flat_in.shape[0])
+    if num_tokens == 0:
+        return
+
+    pool_stride0 = int(pool_out.stride(0))
+    pool_max_context_len = int(pool_out.shape[1])
+
+    grid = (triton.cdiv(num_tokens, _SCATTER_TOKEN_BLOCK),)
+    _scatter_req_token_ids_kernel[grid](
+        flat_in,
+        offsets,
+        req_pool_indices,
+        pool_out,
+        num_tokens=num_tokens,
+        num_batch=bs,
+        pool_stride0=pool_stride0,
+        pool_max_context_len=pool_max_context_len,
+        TOKEN_BLOCK=_SCATTER_TOKEN_BLOCK,
+        BATCH_BLOCK=_SCATTER_BATCH_BLOCK,
+    )
+
+
+@triton.jit
+def _scatter_req_token_ids_kernel(
+    flat_in_ptr,  # [num_tokens] int64
+    offsets_ptr,  # [num_batch + 1] int64
+    req_pool_indices_ptr,  # [num_batch] int64
+    pool_out_ptr,  # [num_rows, pool_max_context_len] int32, row stride = pool_stride0
+    num_tokens,  # scalar int32
+    num_batch,  # scalar int32
+    pool_stride0,  # scalar int32 (row stride of pool_out in elements)
+    pool_max_context_len,  # scalar int32 (dim-1 length of pool_out)
+    TOKEN_BLOCK: tl.constexpr,
+    BATCH_BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    tids = pid * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK)  # [TOKEN_BLOCK] int32
+    tid_mask = tids < num_tokens  # [TOKEN_BLOCK] bool
+
+    bs_offs = tl.arange(0, BATCH_BLOCK)  # [BATCH_BLOCK] int32
+    bs_mask = bs_offs < (num_batch + 1)  # [BATCH_BLOCK] bool
+    offs_vals = tl.load(  # [BATCH_BLOCK] int64
+        offsets_ptr + bs_offs,
+        mask=bs_mask,
+        other=(1 << 62),
+    )
+
+    # find owning req for each tid via reduce-sum: req_idx = (count of offsets <= tid) - 1
+    le = offs_vals[None, :] <= tids[:, None]  # [TOKEN_BLOCK, BATCH_BLOCK] bool
+    req_idx = tl.sum(le.to(tl.int32), axis=1) - 1  # [TOKEN_BLOCK] int32
+
+    safe_req_idx = tl.where(tid_mask, req_idx, 0)  # [TOKEN_BLOCK] int32
+    starts = tl.load(
+        offsets_ptr + safe_req_idx, mask=tid_mask, other=0
+    )  # [TOKEN_BLOCK] int64
+    pos = tids - starts  # [TOKEN_BLOCK] int64
+    rp = tl.load(
+        req_pool_indices_ptr + safe_req_idx, mask=tid_mask, other=0
+    )  # [TOKEN_BLOCK] int64
+
+    # Bound writes by the pool's max_context_len so a token sequence longer than the
+    # ReqToTokenPool row never spills into an adjacent row.
+    in_row = pos < pool_max_context_len  # [TOKEN_BLOCK] bool
+    write_mask = tid_mask & in_row  # [TOKEN_BLOCK] bool
+
+    val = tl.load(flat_in_ptr + tids, mask=tid_mask, other=0).to(
+        tl.int32
+    )  # [TOKEN_BLOCK] int32
+    tl.store(pool_out_ptr + rp * pool_stride0 + pos, val, mask=write_mask)
+
+
+def scatter_req_token_ids_torch_reference(
+    *,
+    flat_in: torch.Tensor,
+    offsets: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    pool_out: torch.Tensor,
+) -> None:
+    """Plain-PyTorch reference for :func:`launch_scatter_req_token_ids_kernel`."""
+    bs = int(req_pool_indices.shape[0])
+    offsets_host = offsets.detach().cpu().tolist()
+    req_pool_indices_host = req_pool_indices.detach().cpu().tolist()
+    flat_host = flat_in.detach().cpu()
+    pool_max_context_len = int(pool_out.shape[1])
+
+    for r in range(bs):
+        start = int(offsets_host[r])
+        end = int(offsets_host[r + 1])
+        if end <= start:
+            continue
+        rp = int(req_pool_indices_host[r])
+        seg = flat_host[start:end].to(torch.int32)
+        write_len = min(int(seg.shape[0]), pool_max_context_len)
+        if write_len <= 0:
+            continue
+        pool_out[rp, :write_len] = seg[:write_len].to(pool_out.device)

python/sglang/jit_kernel/kv_canary/verify.py

@@ -0,0 +1,402 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import TYPE_CHECKING, Final
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+# Bytes per canary slot = CANARY_FIELDS_PER_SLOT * 8.
+CANARY_SLOT_BYTES: Final[int] = consts.CANARY_FIELDS_PER_SLOT * 8
+
+
+class CanaryLaunchTag(IntEnum):
+    """Unique tag per (head | tail | sweep) × (K | V) × (FULL | SWA) launch."""
+
+    HEAD_K_FULL = 0
+    HEAD_V_FULL = 1
+    TAIL_K_FULL = 2
+    TAIL_V_FULL = 3
+    SWEEP_K_FULL = 4
+    SWEEP_V_FULL = 5
+    HEAD_K_SWA = 6
+    HEAD_V_SWA = 7
+    TAIL_K_SWA = 8
+    TAIL_V_SWA = 9
+    SWEEP_K_SWA = 10
+    SWEEP_V_SWA = 11
+
+
+def _assert_contiguous(tensor: torch.Tensor, name: str) -> None:
+    if not tensor.is_contiguous():
+        raise ValueError(f"kv-canary: {name} must be contiguous")
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class RealKvSource:
+    """One piece of real KV the canary folds into its fingerprint.
+
+    Slot access invariant (must hold for every source, regardless of underlying layout) — for a given slot_idx,
+    the canary reads exactly these bytes:
+
+        tensor[
+            slot_idx // page_size,
+            (slot_idx % page_size) * num_bytes_per_token
+            : ((slot_idx % page_size) + 1) * num_bytes_per_token
+        ]
+
+    Note that ``tensor`` may have "holes" in dim 1 — ``tensor.shape[1]`` can exceed ``page_size *
+    num_bytes_per_token``. Trailing bytes of each row are ignored by the canary; this is exactly how the
+    abstraction accommodates pools whose per-row layout interleaves canary-relevant bytes with other metadata
+    (layer-split storage, K/V interleaving, ...). When ``page_size == 1`` the pattern
+    collapses to the simple ``tensor[slot_idx, :num_bytes_per_token]`` case.
+
+    A pool may expose multiple RealKvSource instances per (canary buffer × K/V half) — the launch wrappers
+    iterate the source list and fold each into the running real_kv_hash via splitmix64 (one int64 fingerprint
+    per slot, regardless of source count).
+
+    Pool patchers construct sources by:
+    - viewing / reshaping the underlying KV layer into the canonical [num_rows, dim1_bytes] form (no stage-copy
+      needed when the underlying storage is already row-major contiguous on dim 0),
+    - choosing ``page_size`` and ``num_bytes_per_token`` so that the access pattern above lands on the bytes
+      the canary should fingerprint,
+    - leaving any per-row padding / non-canary bytes in the trailing portion of each row (they will simply be
+      skipped).
+
+    16-byte alignment precondition: the CUDA fold kernel issues 128-bit aligned loads, so ``read_bytes``,
+    ``num_bytes_per_token``, and the row stride (``tensor.shape[1]`` in bytes) must all be positive
+    multiples of 16. There is no "skip this source" sentinel — callers omit the source from their
+    ``real_kv_sources`` tuple entirely (factory helpers return an empty tuple in that case).
+
+    Fields:
+        tensor: The source tensor, any shape such that the access pattern above yields ``num_bytes_per_token``
+            uint8 bytes per slot. Dtype is whatever the underlying pool uses; the canary views the relevant
+            bytes via ``.view(torch.uint8)``.
+        page_size: Number of slots packed into one row of dim 0. ``>= 1``.
+        num_bytes_per_token: Bytes per slot in the dim-1 strip the canary reads. Must be a positive
+            multiple of 16.
+        read_bytes: Leading bytes (out of ``num_bytes_per_token``) per slot folded into the fingerprint.
+            Must be a positive multiple of 16, ``<= num_bytes_per_token``.
+    """
+
+    tensor: torch.Tensor
+    page_size: int
+    num_bytes_per_token: int
+    read_bytes: int
+
+    def __post_init__(self) -> None:
+        if self.page_size < 1:
+            raise ValueError(
+                f"kv-canary: RealKvSource.page_size must be >= 1, got {self.page_size}"
+            )
+        if self.num_bytes_per_token <= 0 or self.num_bytes_per_token % 16 != 0:
+            raise ValueError(
+                f"kv-canary: RealKvSource.num_bytes_per_token must be a positive multiple of 16, "
+                f"got {self.num_bytes_per_token}"
+            )
+        if (
+            self.read_bytes <= 0
+            or self.read_bytes > self.num_bytes_per_token
+            or self.read_bytes % 16 != 0
+        ):
+            raise ValueError(
+                f"kv-canary: RealKvSource.read_bytes must be a positive multiple of 16 in "
+                f"(0, num_bytes_per_token={self.num_bytes_per_token}], got {self.read_bytes}"
+            )
+        if self.tensor.ndim < 2:
+            raise ValueError(
+                f"kv-canary: RealKvSource.tensor must be at least 2-D, got shape {tuple(self.tensor.shape)}"
+            )
+        row_stride_bytes = int(self.tensor.shape[1]) * self.tensor.element_size()
+        if row_stride_bytes % 16 != 0:
+            raise ValueError(
+                f"kv-canary: RealKvSource.tensor dim-1 byte width must be a multiple of 16, "
+                f"got {row_stride_bytes} bytes (shape={tuple(self.tensor.shape)}, "
+                f"dtype={self.tensor.dtype})"
+            )
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class VerifyOrWriteContext:
+    """Shared launch context for canary verify/write kernels.
+
+    Fields:
+        canary_buf: Canary buffer this launch verifies or writes, shape [num_slots, slot_stride_bytes], uint8.
+            slot_stride_bytes is read from canary_buf.shape[1].
+        kernel_kind: CanaryLaunchTag identifying which launch fired. Stamped (as int) into every violation row
+            so host can attribute a violation back to its source launch.
+        violation_ring: Global append-only sink, shape [ring_capacity, VIOLATION_FIELDS], int64. Shared across
+            all canary launches; fill-once.
+        violation_write_index: Global monotonic violation counter, shape [1], int32.
+        slot_run_counter: Health counter, shape [1], int64. Verify increments by active entries processed;
+            write increments by write entries processed.
+        kernel_run_counter: Health counter, shape [1], int64. Incremented by 1 per call.
+        real_kv_sources: Real KV pieces folded into each slot's real_kv_hash, as a tuple of RealKvSource. Empty
+            tuple disables the mixin. Multiple sources are folded sequentially via splitmix64 to produce one
+            int64 fingerprint per slot.
+        real_kv_hash_mode: RealKvHashMode (NONE / PARTIAL / ALL). Applies uniformly across all sources.
+        enable_chain_position_assert: int32 [1] device flag gating the write kernel's chain-step
+            write_position assert. 0 during warmup / cuda-graph capture; flipped to 1 in
+            CanaryManager.mark_init_finished().
+    """
+
+    canary_buf: torch.Tensor
+    kernel_kind: CanaryLaunchTag
+    violation_ring: torch.Tensor
+    violation_write_index: torch.Tensor
+    slot_run_counter: torch.Tensor
+    kernel_run_counter: torch.Tensor
+    real_kv_sources: tuple[RealKvSource, ...]
+    real_kv_hash_mode: consts.RealKvHashMode
+    enable_chain_position_assert: torch.Tensor
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class VerifyPlan:
+    """Flat verify entries consumed by launch_canary_verify_kernel.
+
+    Each row is a self-contained (slot_idx, position, prev_slot_idx) triple, so the verify kernel makes no
+    assumption about the entry's source — per-forward derivation, sweep over running reqs, and sweep over
+    radix-cache orphan slots all populate the same schema. prev_slot_idx == -1 flags a chain-seed entry (kernel
+    anchors on the hardcoded CANARY_CHAIN_ANCHOR constant instead of reading a predecessor).
+
+    Sized to a cuda-graph-captured capacity; active prefix is verify_num_valid[0]. Padding tail entries are
+    unspecified — kernel skips tid >= verify_num_valid[0].
+
+    Fields:
+        verify_slot_indices: Canary slot index per entry, shape [verify_capacity], int64. Already SWA-translated
+            for the SWA group.
+        verify_expected_tokens: Source-of-truth token id per entry, shape [verify_capacity], int64.
+            The plan-side entries kernel gathers from
+            ``CanaryDeviceState.req_to_verify_expected_tokens[rp, position + kv_token_id_vs_position_offset]``;
+            entries that fall outside the pool's row (e.g. EAGLE draft's last slot rotating in a bonus
+            token, or padding beyond the per-req length) get the ``-1`` sentinel. The verify kernel
+            compares against the stored canary token and skips when this is ``-1``.
+        verify_expected_positions: Expected sequence position per entry, shape [verify_capacity], int64.
+        verify_prev_slot_indices: Chain predecessor slot per entry, shape [verify_capacity], int64. -1 = chain
+            head (anchor on CANARY_CHAIN_ANCHOR). Explicit (not derived from verify_slot_indices[i-1])
+            because chain heads, SWA window starts, cross-req boundaries, and radix-orphan extras break the
+            "predecessor == previous array entry" assumption.
+        verify_num_valid: Active entry count, shape [1], int32. Clamped by the plan kernel to
+            min(total_requested, verify_capacity) so the verify kernel grid never reads past the buffer.
+        enable: Run-this-step flag, shape [1], int32. 1 = verify kernel runs as usual; 0 = the plan kernel
+            detected overflow (requested > verify_capacity) and the entire verify launch is skipped this step.
+            Allocated as 1 by default; the plan kernel rewrites it every step.
+    """
+
+    verify_slot_indices: torch.Tensor
+    verify_expected_tokens: torch.Tensor
+    verify_expected_positions: torch.Tensor
+    verify_prev_slot_indices: torch.Tensor
+    verify_num_valid: torch.Tensor
+    enable: torch.Tensor
+
+    @classmethod
+    def allocate(cls, *, verify_capacity: int, device: torch.device) -> VerifyPlan:
+        if verify_capacity <= 0:
+            raise ValueError(
+                f"kv-canary: VerifyPlan verify_capacity must be positive, got {verify_capacity}"
+            )
+        return cls(
+            verify_slot_indices=torch.empty(
+                verify_capacity, dtype=torch.int64, device=device
+            ),
+            verify_expected_tokens=torch.empty(
+                verify_capacity, dtype=torch.int64, device=device
+            ),
+            verify_expected_positions=torch.empty(
+                verify_capacity, dtype=torch.int64, device=device
+            ),
+            verify_prev_slot_indices=torch.empty(
+                verify_capacity, dtype=torch.int64, device=device
+            ),
+            verify_num_valid=torch.empty(1, dtype=torch.int32, device=device),
+            # enable defaults to 1 ("run verify") so test helpers that build a VerifyPlan
+            # directly (no plan kernel) don't have to remember to set it. Plan kernel always
+            # overwrites this so the default is observable only when no plan kernel runs.
+            enable=torch.ones(1, dtype=torch.int32, device=device),
+        )
+
+    def zero_for_testing_(self) -> VerifyPlan:
+        """WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or
+        write kernel to avoid hiding bugs."""
+        self.verify_slot_indices.zero_()
+        # Test helpers expect the "skip token check" sentinel after zero-out, matching
+        # the verify-kernel contract.
+        self.verify_expected_tokens.fill_(-1)
+        self.verify_expected_positions.zero_()
+        self.verify_prev_slot_indices.zero_()
+        self.verify_num_valid.zero_()
+        self.enable.zero_()
+        return self
+
+
+def launch_canary_verify_kernel(
+    *,
+    context: VerifyOrWriteContext,
+    plan: VerifyPlan,
+    check_verify_expected_token: bool,
+) -> None:
+    """Verify one canary buffer against a VerifyPlan.
+
+    A fixed persistent grid of `kPersistentBlocks * kVerifyBlockSize` CUDA threads grid-strides over active
+    verify entries. Each thread reads the slot's 4 stored int64 fields (token_id, position, prev_hash,
+    real_kv_hash), recomputes the expected prev_hash from the predecessor slot (or from
+    splitmix64(CANARY_CHAIN_ANCHOR) for chain heads, signaled by prev_slot_idx == -1), and atomically appends
+    any mismatch (chain hash / position / real_kv_hash) to violation_ring. Read-only on canary_buf.
+
+    Canary slot layout: each slot is canary_buf.shape[1] bytes holding 4 int64 fields (token_id, position,
+    prev_hash, real_kv_hash). Chain link: next.prev_hash == splitmix64_mix3(this.prev_hash, this.token_id,
+    this.position), where splitmix64_mix3 folds each input into a running accumulator
+    via ``acc = splitmix64(acc ^ next)`` starting from ``splitmix64(prev_hash)``. ``real_kv_hash`` is NOT
+    folded into the chain (see ``compute_slot_hash`` rationale: keeps chain content-only and immune to
+    legitimate radix prefix folding). Chain head anchors on
+    splitmix64(CANARY_CHAIN_ANCHOR), where CANARY_CHAIN_ANCHOR is a hardcoded module-level constant (no
+    runtime seed parameter — the canary is for bug detection, not adversarial security, so a fixed anchor
+    is sufficient).
+
+    Args:
+        context: Shared verify/write launch context, including canary buffer, launch tag, violation sink,
+            health counters, and real KV fingerprint sources.
+        plan: Pre-allocated VerifyPlan; addresses baked into cuda-graph capture.
+
+    Token-to-KV slot 0 is unconditionally skipped by the verify kernel: SGLang's TokenToKVPoolAllocator
+    reserves it for padded-token dummy writes, and zero-initialized req_to_token entries therefore point to
+    a non-real KV slot. Canary-attached pools mirror that contract by reserving canary slot 0.
+
+    Implementation:
+        - CUDA __global__ `canary_verify_kernel`: fixed 1-D grid `(kPersistentBlocks=64, 1, 1)` blocks ×
+          `(kVerifyBlockSize=512, 1, 1)` threads (= 32768 threads total). Each thread grid-strides over
+          verify entries `entry_idx ∈ [tid, tid + grid_threads, ...)` until
+          `min(plan.verify_num_valid[0], verify_capacity)`.
+        - Per thread, gather:
+          (a) self_slot fields: 4 separate ``canary_load_field`` int64 loads from
+              canary_buf[plan.verify_slot_indices[tid]] for (token, position, prev_hash, real_kv_hash).
+          (b) expected_prev_hash = compute_slot_hash(canary_buf, slot_stride_bytes, prev_slot_idx), which
+              folds only (token, position, prev_hash) from canary_buf[plan.verify_prev_slot_indices[tid]];
+              prev_slot_idx == -1 anchors at splitmix64(CANARY_CHAIN_ANCHOR).
+          (c) For each src in real_kv_sources: read src.read_bytes leading bytes from src.tensor[...] (per the
+              RealKvSource access invariant) and splitmix64-fold into running_real_kv_hash.
+        - Compare expected vs stored (chain hash, position, real_kv_hash) and accumulate fail_reason bits; if
+          non-zero → record_violation().
+        - record_violation(): idx = atomicAdd(violation_write_index, 1); if idx < ring_capacity, atomic-write
+          the 8 int64 fields to violation_ring[idx] (kernel_kind, slot_idx, position, stored vs expected
+          fields, fail_reason).
+        - Counters: each thread maintains a local count of active entries it processed, warp-reduces via
+          ``__shfl_down_sync`` (offsets 16..1), then the warp leader (lane 0) does a single atomicAdd of the
+          warp's summed count into slot_run_counter. kernel_run_counter += 1: single thread (tid == 0) does an
+          atomicAdd once per launch.
+
+    Calling contract:
+        - Pure side-effect; never raises. Host polls violation_write_index[0] > 0 for is_errored and
+          violation_ring[0] for the first violation.
+        - kernel_run_counter is bumped every call (canary-ran health signal).
+        - Safe in cuda-graph capture; caller refills plan in-place before replay.
+
+    Pinned by torch reference
+    :func:`sglang.jit_kernel.kv_canary.verify_ref.launch_canary_verify_kernel_torch_reference`; CUDA must match
+    byte-for-byte.
+    """
+    canary_buf = context.canary_buf
+    real_kv_sources = context.real_kv_sources
+    if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
+        raise ValueError(
+            f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
+            f"got {len(real_kv_sources)}"
+        )
+
+    _assert_contiguous(canary_buf, "canary_buf")
+    _assert_contiguous(plan.verify_slot_indices, "plan.verify_slot_indices")
+    _assert_contiguous(plan.verify_expected_tokens, "plan.verify_expected_tokens")
+    _assert_contiguous(plan.verify_expected_positions, "plan.verify_expected_positions")
+    _assert_contiguous(plan.verify_prev_slot_indices, "plan.verify_prev_slot_indices")
+    _assert_contiguous(plan.verify_num_valid, "plan.verify_num_valid")
+    _assert_contiguous(plan.enable, "plan.enable")
+    _assert_contiguous(context.violation_ring, "violation_ring")
+    _assert_contiguous(context.violation_write_index, "violation_write_index")
+    _assert_contiguous(context.slot_run_counter, "slot_run_counter")
+    _assert_contiguous(context.kernel_run_counter, "kernel_run_counter")
+
+    padded_bufs, source_params = _build_real_kv_source_abi(
+        real_kv_sources=real_kv_sources, device=canary_buf.device
+    )
+
+    module = _jit_canary_verify_module(check_verify_expected_token)
+    module.canary_verify_step_cuda(
+        canary_buf,
+        plan.verify_slot_indices,
+        plan.verify_expected_tokens,
+        plan.verify_expected_positions,
+        plan.verify_prev_slot_indices,
+        plan.verify_num_valid,
+        plan.enable,
+        int(context.kernel_kind),
+        context.violation_ring,
+        context.violation_write_index,
+        context.slot_run_counter,
+        context.kernel_run_counter,
+        padded_bufs[0],
+        padded_bufs[1],
+        padded_bufs[2],
+        padded_bufs[3],
+        source_params,
+        len(real_kv_sources),
+        int(context.real_kv_hash_mode),
+    )
+
+
+@cache_once
+def _jit_canary_verify_module(check_verify_expected_token: bool) -> Module:
+    args = make_cpp_args(check_verify_expected_token)
+    return load_jit(
+        "kv_canary_verify",
+        *args,
+        cuda_files=["kv_canary/canary_verify.cuh"],
+        cuda_wrappers=[
+            (
+                "canary_verify_step_cuda",
+                f"canary::CanaryVerifyKernel<{args}>::run",
+            ),
+        ],
+    )
+
+
+def _build_real_kv_source_abi(
+    *,
+    real_kv_sources: tuple[RealKvSource, ...],
+    device: torch.device,
+) -> tuple[list[torch.Tensor], torch.Tensor]:
+    padded_bufs: list[torch.Tensor] = []
+    params = torch.zeros(
+        (consts.MAX_REAL_KV_SOURCES, consts.REAL_KV_SOURCE_FIELDS_PER_ENTRY),
+        dtype=torch.int32,
+        device="cpu",
+    )
+
+    for i, source in enumerate(real_kv_sources):
+        _assert_contiguous(source.tensor, f"real_kv_sources[{i}].tensor")
+        source_u8 = source.tensor.view(torch.uint8)
+        if source_u8.dim() != 2:
+            raise ValueError(
+                f"kv-canary: real_kv_sources[{i}].tensor (viewed as uint8) must be 2-D, "
+                f"got {source_u8.dim()}-D"
+            )
+        padded_bufs.append(source_u8)
+        params[i, consts.REAL_KV_SOURCE_FIELD_PAGE_SIZE] = source.page_size
+        params[i, consts.REAL_KV_SOURCE_FIELD_NUM_BYTES_PER_TOKEN] = (
+            source.num_bytes_per_token
+        )
+        params[i, consts.REAL_KV_SOURCE_FIELD_READ_BYTES] = source.read_bytes
+
+    # Pad bufs (never read by the kernel — num_sources bounds the iteration); params already zero.
+    dummy = torch.empty((1, 1), dtype=torch.uint8, device=device)
+    for _ in range(len(real_kv_sources), consts.MAX_REAL_KV_SOURCES):
+        padded_bufs.append(dummy)
+
+    return padded_bufs, params

python/sglang/jit_kernel/kv_canary/verify_ref.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+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 (
+    RealKvSource,
+    VerifyOrWriteContext,
+    VerifyPlan,
+)
+
+_U64_MASK: int = (1 << 64) - 1
+_I64_SIGN_BIT: int = 1 << 63
+
+
+def launch_canary_verify_kernel_torch_reference(
+    *,
+    context: VerifyOrWriteContext,
+    plan: VerifyPlan,
+    check_verify_expected_token: bool,
+) -> None:
+    canary_buf = context.canary_buf
+    kernel_kind = context.kernel_kind
+    violation_ring = context.violation_ring
+    violation_write_index = context.violation_write_index
+    slot_run_counter = context.slot_run_counter
+    kernel_run_counter = context.kernel_run_counter
+    real_kv_sources = context.real_kv_sources
+    real_kv_hash_mode = context.real_kv_hash_mode
+
+    work_device = torch.device("cpu")
+
+    kernel_run_counter.add_(1)
+
+    enable = int(plan.enable.detach().to("cpu").item())
+    if enable == 0:
+        return
+
+    num_valid = int(
+        plan.verify_slot_indices.new_empty(()).copy_(plan.verify_num_valid[0]).item()
+    )
+    capacity = int(plan.verify_slot_indices.shape[0])
+    active = max(0, min(num_valid, capacity))
+    if active <= 0:
+        return
+
+    slot_indices_host = plan.verify_slot_indices[:active].to(
+        device=work_device, dtype=torch.int64
+    )
+    if check_verify_expected_token:
+        expected_input_ids_host = plan.verify_expected_tokens[:active].to(
+            device=work_device, dtype=torch.int64
+        )
+    else:
+        expected_input_ids_host = torch.full(
+            (active,), -1, dtype=torch.int64, device=work_device
+        )
+    expected_positions_host = plan.verify_expected_positions[:active].to(
+        device=work_device, dtype=torch.int64
+    )
+    prev_slot_indices_host = plan.verify_prev_slot_indices[:active].to(
+        device=work_device, dtype=torch.int64
+    )
+
+    slot_run_counter.add_(active)
+
+    kept_slots: list[int] = []
+    kept_expected_positions: list[int] = []
+    kept_expected_input_ids: list[int] = []
+    kept_prev_slots: list[int] = []
+    for k in range(active):
+        s = int(slot_indices_host[k].item())
+        # Skip SGLang's padded-token dummy KV slot so unfilled req_to_token entries (zero-initialized) do not
+        # produce spurious chain_hash / position violations.
+        if s != consts.TOKEN_TO_KV_SLOT_PADDING:
+            kept_slots.append(s)
+            kept_expected_positions.append(int(expected_positions_host[k].item()))
+            kept_expected_input_ids.append(int(expected_input_ids_host[k].item()))
+            kept_prev_slots.append(int(prev_slot_indices_host[k].item()))
+    active = len(kept_slots)
+    if active <= 0:
+        return
+    slot_indices_list: list[int] = kept_slots
+    expected_positions_list: list[int] = kept_expected_positions
+    expected_input_ids_list: list[int] = kept_expected_input_ids
+    prev_slot_indices_list: list[int] = kept_prev_slots
+
+    buf_i64 = canary_buf.detach().to(device=work_device).contiguous().view(torch.int64)
+    slot_stride_i64 = int(buf_i64.shape[1])
+    if slot_stride_i64 < 4:
+        raise ValueError(
+            f"kv-canary: canary_buf slot stride must hold at least 4 int64 fields, got {slot_stride_i64}"
+        )
+
+    violation_rows: list[list[int]] = []
+
+    for k in range(active):
+        slot_idx = slot_indices_list[k]
+        expected_position = expected_positions_list[k]
+        expected_input_id = expected_input_ids_list[k]
+        prev_slot = prev_slot_indices_list[k]
+
+        stored_token = int(buf_i64[slot_idx, consts.CANARY_FIELD_TOKEN].item())
+        stored_position = int(buf_i64[slot_idx, consts.CANARY_FIELD_POSITION].item())
+        stored_chain_hash = int(buf_i64[slot_idx, consts.CANARY_FIELD_PREV_HASH].item())
+        stored_real_kv_hash = int(
+            buf_i64[slot_idx, consts.CANARY_FIELD_REAL_KV_HASH].item()
+        )
+
+        prev_reachable = prev_slot != consts.TOKEN_TO_KV_SLOT_PADDING
+        if prev_reachable:
+            expected_chain_hash = _to_signed_int64(
+                compute_slot_hash(buf_i64, prev_slot)
+            )
+        else:
+            expected_chain_hash = stored_chain_hash
+
+        expected_real_kv_hash_u64 = _compute_real_kv_hash_scalar(
+            slot_idx=slot_idx,
+            real_kv_sources=real_kv_sources,
+            real_kv_hash_mode=real_kv_hash_mode,
+            work_device=work_device,
+        )
+        expected_real_kv_hash = _to_signed_int64(expected_real_kv_hash_u64)
+
+        fail_reason = consts.FailReason(0)
+        if prev_reachable and stored_chain_hash != expected_chain_hash:
+            fail_reason |= consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH
+        if check_verify_expected_token:
+            if expected_input_id != -1 and stored_token != expected_input_id:
+                fail_reason |= consts.FailReason.VERIFY_TOKEN_MISMATCH
+        if stored_position != expected_position:
+            fail_reason |= consts.FailReason.VERIFY_POSITION_MISMATCH
+        if stored_real_kv_hash != expected_real_kv_hash:
+            fail_reason |= consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH
+
+        if fail_reason != consts.FailReason(0):
+            row = [0] * consts.VIOLATION_FIELDS
+            row[consts.VIOLATION_FIELD_KERNEL_KIND] = int(kernel_kind)
+            row[consts.VIOLATION_FIELD_SLOT_IDX] = slot_idx
+            row[consts.VIOLATION_FIELD_POSITION] = stored_position
+            row[consts.VIOLATION_FIELD_STORED_TOKEN] = stored_token
+            row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = expected_input_id
+            row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = stored_chain_hash
+            row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_chain_hash
+            row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = int(fail_reason)
+            violation_rows.append(row)
+
+    if len(violation_rows) == 0:
+        return
+
+    num_new_violations = len(violation_rows)
+    base_idx = int(
+        violation_write_index.new_empty(()).copy_(violation_write_index[0]).item()
+    )
+    ring_capacity = int(violation_ring.shape[0])
+
+    new_rows = torch.zeros(
+        (num_new_violations, consts.VIOLATION_FIELDS), dtype=torch.int64
+    )
+    for v, row in enumerate(violation_rows):
+        for f in range(consts.VIOLATION_FIELDS):
+            new_rows[v, f] = row[f]
+
+    write_count_in_ring = max(0, min(num_new_violations, ring_capacity - base_idx))
+    if write_count_in_ring > 0:
+        ring_host = violation_ring.detach().to(device=work_device)
+        ring_host[base_idx : base_idx + write_count_in_ring, :] = new_rows[
+            :write_count_in_ring, :
+        ]
+        violation_ring.copy_(ring_host.to(violation_ring.device))
+
+    violation_write_index[0] = violation_write_index[0] + num_new_violations
+
+
+def _to_signed_int64(value: int) -> int:
+    value &= _U64_MASK
+    if value >= _I64_SIGN_BIT:
+        value -= 1 << 64
+    return value
+
+
+def compute_slot_hash(buf_i64: torch.Tensor, source_slot_idx: int) -> int:
+    if source_slot_idx < 0:
+        return splitmix64(consts.CANARY_CHAIN_ANCHOR)
+    token = int(buf_i64[source_slot_idx, consts.CANARY_FIELD_TOKEN].item())
+    position = int(buf_i64[source_slot_idx, consts.CANARY_FIELD_POSITION].item())
+    prev_hash = int(buf_i64[source_slot_idx, consts.CANARY_FIELD_PREV_HASH].item())
+    return splitmix64_mix3(prev_hash, token, position)
+
+
+def _compute_real_kv_hash_scalar(
+    *,
+    slot_idx: int,
+    real_kv_sources: tuple[RealKvSource, ...],
+    real_kv_hash_mode: consts.RealKvHashMode,
+    work_device: torch.device,
+) -> int:
+    mode = int(real_kv_hash_mode)
+    if mode == int(consts.RealKvHashMode.NONE) or len(real_kv_sources) == 0:
+        return 0
+
+    acc: int = 0
+
+    for source in real_kv_sources:
+        page_size = source.page_size
+        num_bytes_per_token = source.num_bytes_per_token
+        read_bytes = source.read_bytes
+        tensor_u8 = (
+            source.tensor.detach().to(device=work_device).contiguous().view(torch.uint8)
+        )
+
+        row = slot_idx // page_size
+        col_within_page = slot_idx % page_size
+        col_start = col_within_page * num_bytes_per_token
+
+        effective_read_bytes = (
+            16 if mode == int(consts.RealKvHashMode.PARTIAL) else read_bytes
+        )
+        raw_bytes: list[int] = []
+        for b in range(effective_read_bytes):
+            raw_bytes.append(int(tensor_u8[row, col_start + b].item()))
+
+        source_hash = _splitmix64_fold_bytes_scalar(raw_bytes=raw_bytes)
+
+        combined = acc ^ source_hash
+        acc = splitmix64(combined)
+
+    return acc
+
+
+def _splitmix64_fold_bytes_scalar(*, raw_bytes: list[int]) -> int:
+    read_bytes = len(raw_bytes)
+    pad = (8 - read_bytes % 8) % 8
+    padded = raw_bytes + [0] * pad
+    num_words = len(padded) // 8
+
+    acc: int = 0
+    for w in range(num_words):
+        word: int = 0
+        for k in range(8):
+            word |= padded[w * 8 + k] << (8 * k)
+        word &= _U64_MASK
+        acc = splitmix64(acc ^ word)
+
+    return acc

python/sglang/jit_kernel/kv_canary/write.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.verify import (
+    VerifyOrWriteContext,
+    _assert_contiguous,
+    _build_real_kv_source_abi,
+)
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class WritePlan:
+    """Write plan consumed by launch_canary_write_kernel: per-token slot indices + per-req metadata.
+
+    Fully per-req — no per-token tile. launch_canary_write_kernel uses write_offsets to map each thread's
+    (req, j) into a flat index i, then reads token-level data from input_ids / positions /
+    out_cache_loc[i] directly.
+    SWA translation of per-token slots is done **host-side by the caller** (typically the endpoint) before
+    invoking launch_canary_write_kernel — the kernel is SWA-agnostic and only understands "slot ≥ 0 ⇒ write;
+    slot < 0 ⇒ skip this entry". Only the chain-seed slot (a per-req gather from req_to_token at plan time)
+    is SWA-translated by the plan kernel and lives in write_seed_slot_indices.
+
+    Req r's write entries occupy flat indices [write_offsets[r], write_offsets[r+1]). seed_slot_idx == -1 means
+    K_req_old == 0 (anchor on CANARY_CHAIN_ANCHOR).
+
+    Fields:
+        write_offsets: Exclusive prefix-sum offsets indexing into ForwardBatch's input_ids / positions /
+            out_cache_loc, shape [write_req_capacity + 1], int64. write_offsets[0] == 0;
+            write_offsets[write_num_valid_reqs[0]] == total_write_entries.
+        write_seed_slot_indices: Chain-seed slot per write req, shape [write_req_capacity], int64. Already
+            SWA-translated. -1 = no prefix (chain anchors on CANARY_CHAIN_ANCHOR).
+        write_num_valid_reqs: Active write-req count, shape [1], int32. launch_canary_write_kernel skips blocks
+            with block_id >= write_num_valid_reqs[0].
+    """
+
+    write_offsets: torch.Tensor
+    write_seed_slot_indices: torch.Tensor
+    write_num_valid_reqs: torch.Tensor
+
+    @classmethod
+    def allocate(
+        cls,
+        *,
+        write_req_capacity: int,
+        device: torch.device,
+    ) -> WritePlan:
+        if write_req_capacity <= 0:
+            raise ValueError(
+                f"kv-canary: WritePlan write_req_capacity must be positive, got {write_req_capacity}"
+            )
+        return cls(
+            write_offsets=torch.empty(
+                write_req_capacity + 1, dtype=torch.int64, device=device
+            ),
+            write_seed_slot_indices=torch.empty(
+                write_req_capacity, dtype=torch.int64, device=device
+            ),
+            write_num_valid_reqs=torch.empty(1, dtype=torch.int32, device=device),
+        )
+
+    def zero_for_testing_(self) -> WritePlan:
+        """WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or
+        write kernel to avoid hiding bugs."""
+        self.write_offsets.zero_()
+        self.write_seed_slot_indices.zero_()
+        self.write_num_valid_reqs.zero_()
+        return self
+
+
+def launch_canary_write_kernel(
+    *,
+    context: VerifyOrWriteContext,
+    plan: WritePlan,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    enable_write_input_assert: bool,
+    expected_input_tokens: torch.Tensor | None,
+    expected_input_positions: torch.Tensor | None,
+) -> None:
+    """Write canary fingerprints into one canary buffer per a WritePlan.
+
+    Grid: one CUDA block per active write req, single thread per block (chain is intrinsically serial).
+    Block r walks entries ``[plan.write_offsets[r], plan.write_offsets[r+1])``. Per chain step ``i``:
+
+    - ``slot`` = ``out_cache_loc[i]`` (caller-pre-translated for SWA groups; entries set to -1 are skipped).
+    - ``token / position`` = ``input_ids[i] / positions[i]``.
+    - ``real_kv_hash`` = ``real_kv_fold_sources(real_kv_sources, slot)`` if ``real_kv_hash_mode != NONE`` else 0.
+    - Store 4 int64s ``(token, position, running_prev_hash, real_kv_hash)`` into ``canary_buf[slot]``.
+    - Advance ``running_prev_hash = splitmix64_mix3(prev, token, position)``, where
+      splitmix64_mix3 folds each input via ``acc = splitmix64(acc ^ next)`` starting from ``splitmix64(prev)``.
+      ``real_kv_hash`` is intentionally not folded into the chain — see ``compute_slot_hash`` in
+      ``csrc/kv_canary/canary_common.cuh`` for the radix-folding rationale.
+
+    Initial ``running_prev_hash`` when ``seed_slot_idx >= 0``: load (token, position, prev_hash) from
+    ``canary_buf[plan.write_seed_slot_indices[r]]`` and set
+    ``running_prev_hash = splitmix64_mix3(seed.prev_hash, seed.token, seed.position)``
+    (i.e. apply the same advance step that produced ``seed``'s successor — this keeps slot[0]'s stored
+    ``prev_hash`` consistent with the chain link). Else
+    ``running_prev_hash = splitmix64(CANARY_CHAIN_ANCHOR)``. ``write_seed_slot_indices`` is already
+    SWA-translated by the plan kernel; ``CANARY_CHAIN_ANCHOR`` is hardcoded module-level (no runtime seed).
+
+    Write-time input verification (caller-driven, kernel is oracle-agnostic): when
+    ``enable_write_input_assert`` is True the kernel additionally compares ``input_ids[i]`` against
+    ``expected_input_tokens[i]`` and ``positions[i]`` against ``expected_input_positions[i]``; mismatch
+    on either field records a violation. The chain still advances on the actual values (not the expected
+    ones) so a downstream verify won't cascade. Whoever produced the expected tensors is responsible for
+    filling them; the kernel runs no oracle internally.
+
+    Write only writes canary_buf (reads only at seed slots). Block uses no shared memory.
+
+    The ForwardBatch-derived arguments are passed through unchanged from the source ForwardBatch — canary does not transform
+    them.
+
+    Args:
+        context: Shared verify/write launch context, including canary buffer, launch tag, violation sink,
+            health counters, and real KV fingerprint sources.
+        plan: Pre-allocated WritePlan.
+        input_ids: ForwardBatch.input_ids; token ids being written, shape [num_tokens_padded], int64.
+            Flattened across reqs in plan.write_offsets order; tail beyond
+            plan.write_offsets[plan.write_num_valid_reqs[0]] is cuda-graph padding.
+        positions: ForwardBatch.positions; sequence positions of input_ids, shape [num_tokens_padded], int64.
+        out_cache_loc: Per-token canary slot index, shape [num_tokens_padded], int64. The caller is
+            responsible for translating ForwardBatch.out_cache_loc into the canary's index space for SWA
+            groups (typically a host-side LUT gather in the endpoint); FULL groups pass it through
+            unchanged. A -1 entry signals skip-this-token (used for SWA out-of-window slots or padding).
+            The kernel does not consult any LUT.
+        enable_write_input_assert: bool toggle. False = expected_input_* tensors must be None. True = compare
+            each chain step's actual (token, position) against the caller-supplied expected tensors below.
+        expected_input_tokens: Expected token id per write entry, shape [num_tokens_padded], int64. Only read
+            when enable_write_input_assert is True; must be None when enable_write_input_assert is False.
+            Layout mirrors input_ids (flattened across reqs in plan.write_offsets order); padding tail
+            is ignored. Filled by the caller from whichever oracle produces expected inputs — the kernel
+            knows no oracle.
+        expected_input_positions: Expected position per write entry, shape [num_tokens_padded], int64, or None.
+            Same shape/layout/lifetime rules as expected_input_tokens.
+
+    Implementation:
+        - CUDA __global__ `canary_write_kernel`: 1-D grid `(write_req_capacity, 1, 1)` blocks × `(1, 1, 1)` thread
+          per block. block_id r = blockIdx.x = one write req; chains are intrinsically serial so a single thread
+          per block is optimal (warp-level parallelism would idle 31 lanes).
+        - Per block, early-exit on r >= plan.write_num_valid_reqs[0]. Else load entry_start = plan.write_offsets[r],
+          entry_count = plan.write_offsets[r+1] - entry_start, seed_slot_idx = plan.write_seed_slot_indices[r] into
+          registers.
+        - Initialize running_prev_hash: if seed_slot_idx >= 0, load (token, position, prev_hash) from
+          canary_buf[seed_slot_idx] and set running_prev_hash = splitmix64_mix3(prev_hash, token, position);
+          else running_prev_hash = splitmix64(kCanaryChainAnchor).
+        - Serial chain loop `for j in range(entry_count)`:
+              i = entry_start + j;
+              slot = out_cache_loc[i];  // caller-pre-translated; the kernel never consults a LUT
+              if (slot < 0) continue;       // -1 sentinel = skip (SWA out-of-window or padding)
+              token = input_ids[i]; position = positions[i];
+              real_kv_hash = (real_kv_hash_mode == NONE) ? 0 : real_kv_fold_sources(real_kv_sources, slot);
+                  // applies RealKvSource access invariant
+              if enable_write_input_assert:
+                  if token != expected_input_tokens[i] or position != expected_input_positions[i]:
+                      record_violation();  // chain still advances on the ACTUAL (token, position) below
+              store (token, position, running_prev_hash, real_kv_hash) to canary_buf[slot] as 4 int64 fields;
+              running_prev_hash = splitmix64_mix3(running_prev_hash, token, position);
+        - All chain state lives in the block's single thread's registers. No shared memory, no cross-block
+          coordination.
+        - record_violation() identical to verify (atomicAdd + atomic-write).
+        - Counters: thread of block 0 does atomicAdd(kernel_run_counter, 1); each block accumulates its
+          entry_count and atomicAdds to slot_run_counter once at exit.
+
+    Calling contract:
+        - Pure side-effect; never raises.
+        - Input-verification mismatch records violations but does NOT abort the chain.
+        - kernel_run_counter is bumped every call.
+        - Safe in cuda-graph capture; caller refills input_ids / positions / out_cache_loc / plan
+          in-place before replay.
+
+    Pinned by torch reference
+    :func:`sglang.jit_kernel.kv_canary.write_ref.launch_canary_write_kernel_torch_reference`; CUDA must match
+    byte-for-byte.
+    """
+    canary_buf = context.canary_buf
+    real_kv_sources = context.real_kv_sources
+    if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
+        raise ValueError(
+            f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
+            f"got {len(real_kv_sources)}"
+        )
+
+    _assert_contiguous(canary_buf, "canary_buf")
+    _assert_contiguous(plan.write_offsets, "plan.write_offsets")
+    _assert_contiguous(plan.write_seed_slot_indices, "plan.write_seed_slot_indices")
+    _assert_contiguous(plan.write_num_valid_reqs, "plan.write_num_valid_reqs")
+    _assert_contiguous(input_ids, "input_ids")
+    _assert_contiguous(positions, "positions")
+    _assert_contiguous(out_cache_loc, "out_cache_loc")
+    if enable_write_input_assert:
+        if expected_input_tokens is None or expected_input_positions is None:
+            raise ValueError(
+                "kv-canary: expected input tensors are required when enable_write_input_assert=True"
+            )
+        _assert_contiguous(expected_input_tokens, "expected_input_tokens")
+        _assert_contiguous(expected_input_positions, "expected_input_positions")
+    else:
+        if expected_input_tokens is not None or expected_input_positions is not None:
+            raise ValueError(
+                "kv-canary: expected input tensors must be None when enable_write_input_assert=False"
+            )
+    _assert_contiguous(context.violation_ring, "violation_ring")
+    _assert_contiguous(context.violation_write_index, "violation_write_index")
+    _assert_contiguous(context.slot_run_counter, "slot_run_counter")
+    _assert_contiguous(context.kernel_run_counter, "kernel_run_counter")
+    _assert_contiguous(
+        context.enable_chain_position_assert, "enable_chain_position_assert"
+    )
+
+    padded_bufs, source_params = _build_real_kv_source_abi(
+        real_kv_sources=real_kv_sources, device=canary_buf.device
+    )
+
+    module = _jit_canary_write_module()
+    module.canary_write_step_cuda(
+        canary_buf,
+        plan.write_offsets,
+        plan.write_seed_slot_indices,
+        plan.write_num_valid_reqs,
+        input_ids,
+        positions,
+        out_cache_loc,
+        int(context.kernel_kind),
+        int(enable_write_input_assert),
+        expected_input_tokens,
+        expected_input_positions,
+        context.violation_ring,
+        context.violation_write_index,
+        context.slot_run_counter,
+        context.kernel_run_counter,
+        context.enable_chain_position_assert,
+        padded_bufs[0],
+        padded_bufs[1],
+        padded_bufs[2],
+        padded_bufs[3],
+        source_params,
+        len(real_kv_sources),
+        int(context.real_kv_hash_mode),
+    )
+
+
+@cache_once
+def _jit_canary_write_module() -> Module:
+    return load_jit(
+        "kv_canary_write",
+        cuda_files=["kv_canary/canary_write.cuh"],
+        cuda_wrappers=[
+            ("canary_write_step_cuda", "canary::canary_write_step_cuda"),
+        ],
+    )

python/sglang/jit_kernel/kv_canary/write_ref.py

@@ -0,0 +1,213 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.verify import (
+    VerifyOrWriteContext,
+)
+from sglang.jit_kernel.kv_canary.verify_ref import (
+    _compute_real_kv_hash_scalar,
+    _to_signed_int64,
+    compute_slot_hash,
+    splitmix64_mix3,
+)
+from sglang.jit_kernel.kv_canary.write import WritePlan
+
+
+def launch_canary_write_kernel_torch_reference(
+    *,
+    context: VerifyOrWriteContext,
+    plan: WritePlan,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    enable_write_input_assert: bool,
+    expected_input_tokens: torch.Tensor | None,
+    expected_input_positions: torch.Tensor | None,
+) -> None:
+    canary_buf = context.canary_buf
+    kernel_kind = context.kernel_kind
+    violation_ring = context.violation_ring
+    violation_write_index = context.violation_write_index
+    slot_run_counter = context.slot_run_counter
+    kernel_run_counter = context.kernel_run_counter
+    real_kv_sources = context.real_kv_sources
+    real_kv_hash_mode = context.real_kv_hash_mode
+    enable_chain_position_assert_value = int(
+        context.enable_chain_position_assert.detach().to("cpu").item()
+    )
+
+    work_device = torch.device("cpu")
+
+    kernel_run_counter.add_(1)
+
+    num_valid_reqs = int(plan.write_num_valid_reqs.detach().to("cpu").item())
+    req_capacity = int(plan.write_seed_slot_indices.shape[0])
+    active_reqs = max(0, min(num_valid_reqs, req_capacity))
+    if active_reqs <= 0:
+        return
+
+    write_offsets_host = plan.write_offsets.detach().to(
+        device=work_device, dtype=torch.int64
+    )
+    seed_slot_indices_host = plan.write_seed_slot_indices[:active_reqs].to(
+        device=work_device, dtype=torch.int64
+    )
+    input_ids_host = input_ids.detach().to(device=work_device, dtype=torch.int64)
+    positions_host = positions.detach().to(device=work_device, dtype=torch.int64)
+    out_cache_loc_host = out_cache_loc.detach().to(
+        device=work_device, dtype=torch.int64
+    )
+
+    total_entries = int(write_offsets_host[active_reqs].item())
+    if total_entries <= 0:
+        return
+
+    buf_i64 = (
+        canary_buf.detach()
+        .to(device=work_device)
+        .contiguous()
+        .view(torch.int64)
+        .clone()
+    )
+    slot_stride_i64 = int(buf_i64.shape[1])
+    if slot_stride_i64 < 4:
+        raise ValueError(
+            f"kv-canary: canary_buf slot stride must hold at least 4 int64 fields, got {slot_stride_i64}"
+        )
+
+    if enable_write_input_assert:
+        if expected_input_tokens is None or expected_input_positions is None:
+            raise ValueError(
+                "kv-canary: expected input tensors are required when enable_write_input_assert=True"
+            )
+        expected_input_tokens_host = expected_input_tokens.detach().to(
+            device=work_device, dtype=torch.int64
+        )
+        expected_input_positions_host = expected_input_positions.detach().to(
+            device=work_device, dtype=torch.int64
+        )
+    else:
+        if expected_input_tokens is not None or expected_input_positions is not None:
+            raise ValueError(
+                "kv-canary: expected input tensors must be None when enable_write_input_assert=False"
+            )
+        expected_input_tokens_host = None
+        expected_input_positions_host = None
+
+    violation_rows: list[list[int]] = []
+    total_slots_written = 0
+
+    for r in range(active_reqs):
+        entry_start = int(write_offsets_host[r].item())
+        entry_end = int(write_offsets_host[r + 1].item())
+        entry_count = entry_end - entry_start
+        if entry_count <= 0:
+            continue
+
+        seed_slot = int(seed_slot_indices_host[r].item())
+        running_prev_hash = compute_slot_hash(buf_i64, seed_slot)
+
+        do_chain_position_assert = (seed_slot >= 0) and (
+            enable_chain_position_assert_value != 0
+        )
+        if do_chain_position_assert:
+            running_prev_position = int(
+                buf_i64[seed_slot, consts.CANARY_FIELD_POSITION].item()
+            )
+        else:
+            running_prev_position = 0
+
+        for entry_offset in range(entry_count):
+            entry_idx = entry_start + entry_offset
+            slot = int(out_cache_loc_host[entry_idx].item())
+            if slot < 0:
+                continue
+            token = int(input_ids_host[entry_idx].item())
+            position = int(positions_host[entry_idx].item())
+
+            real_kv_hash_u64 = _compute_real_kv_hash_scalar(
+                slot_idx=slot,
+                real_kv_sources=real_kv_sources,
+                real_kv_hash_mode=real_kv_hash_mode,
+                work_device=work_device,
+            )
+
+            if enable_write_input_assert:
+                assert expected_input_tokens_host is not None
+                assert expected_input_positions_host is not None
+                mismatch_bits = consts.FailReason(0)
+                expected_token = int(expected_input_tokens_host[entry_idx].item())
+                expected_position = int(expected_input_positions_host[entry_idx].item())
+                if token != expected_token:
+                    mismatch_bits |= consts.FailReason.WRITE_TOKEN_MISMATCH
+                if position != expected_position:
+                    mismatch_bits |= consts.FailReason.WRITE_POSITION_MISMATCH
+                if mismatch_bits != consts.FailReason(0):
+                    row = [0] * consts.VIOLATION_FIELDS
+                    row[consts.VIOLATION_FIELD_KERNEL_KIND] = int(kernel_kind)
+                    row[consts.VIOLATION_FIELD_SLOT_IDX] = slot
+                    row[consts.VIOLATION_FIELD_POSITION] = position
+                    row[consts.VIOLATION_FIELD_STORED_TOKEN] = token
+                    row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = expected_token
+                    row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = _to_signed_int64(
+                        running_prev_hash
+                    )
+                    row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_position
+                    row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = int(mismatch_bits)
+                    violation_rows.append(row)
+
+            if do_chain_position_assert:
+                expected_position_chain = running_prev_position + 1
+                if position != expected_position_chain:
+                    row = [0] * consts.VIOLATION_FIELDS
+                    row[consts.VIOLATION_FIELD_KERNEL_KIND] = int(kernel_kind)
+                    row[consts.VIOLATION_FIELD_SLOT_IDX] = slot
+                    row[consts.VIOLATION_FIELD_POSITION] = position
+                    row[consts.VIOLATION_FIELD_STORED_TOKEN] = token
+                    row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = token
+                    row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = _to_signed_int64(
+                        running_prev_hash
+                    )
+                    row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_position_chain
+                    row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = int(
+                        consts.FailReason.WRITE_POSITION_MISMATCH
+                    )
+                    violation_rows.append(row)
+                running_prev_position = position
+
+            buf_i64[slot, consts.CANARY_FIELD_TOKEN] = token
+            buf_i64[slot, consts.CANARY_FIELD_POSITION] = position
+            buf_i64[slot, consts.CANARY_FIELD_PREV_HASH] = _to_signed_int64(
+                running_prev_hash
+            )
+            buf_i64[slot, consts.CANARY_FIELD_REAL_KV_HASH] = _to_signed_int64(
+                real_kv_hash_u64
+            )
+
+            running_prev_hash = splitmix64_mix3(running_prev_hash, token, position)
+
+            total_slots_written += 1
+
+    canary_buf.view(torch.int64).copy_(
+        buf_i64.to(canary_buf.device).view(canary_buf.shape[0], slot_stride_i64)
+    )
+
+    slot_run_counter.add_(total_slots_written)
+
+    if len(violation_rows) == 0:
+        return
+
+    base_idx = int(violation_write_index.detach().to("cpu").item())
+    ring_capacity = int(violation_ring.shape[0])
+    new_rows = torch.tensor(violation_rows, dtype=torch.int64, device=work_device)
+    write_count_in_ring = max(0, min(len(violation_rows), ring_capacity - base_idx))
+    if write_count_in_ring > 0:
+        ring_host = violation_ring.detach().to(device=work_device)
+        ring_host[base_idx : base_idx + write_count_in_ring, :] = new_rows[
+            :write_count_in_ring, :
+        ]
+        violation_ring.copy_(ring_host.to(violation_ring.device))
+
+    violation_write_index[0] = violation_write_index[0] + len(violation_rows)