sgl-project--sglang/python/sglang/jit_kernel/kv_canary/verify.py

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