chore: import upstream snapshot with attribution

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