Add the KV-canary core: data layer, MHA KV-pool patcher, and per-forward runner (#26808)

python/sglang/srt/environ.py

@@ -749,6 +749,12 @@ class Envs:
     SGLANG_PLATFORM = EnvStr("")
     SGLANG_PLUGINS = EnvStr("")
 
+    # ===================================================================
+    # KV-Canary (testing-only)
+    # ===================================================================
+    SGLANG_KV_CANARY_RING_CAPACITY = EnvInt(1024)
+    SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False)
+
 
 envs = Envs()
 EnvField._allow_set_name = False

python/sglang/srt/kv_canary/buffer_group.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Optional
+
+import torch
+
+
+class PoolKind(IntEnum):
+    """Which attention regime a canary group belongs to.
+
+    - ``FULL`` covers ``[0, K_req)``. Attached to plain MHA/MLA pools and as one of the two canaries on
+      every SWA system.
+    - ``SWA`` covers ``[max(0, K_req - window), K_req)``. Attached as the second canary on every
+      ``BaseSWAKVPool``.
+    """
+
+    FULL = 0
+    SWA = 1
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class CanaryBufferGroup:
+    """Canary buffers for one (PoolKind × K-half | V-half) on a pool.
+
+    Each (head | tail) launch sees a single 2-D uint8 buf for the canary. Head and tail use separate canary
+    buffers so they can be staged at different points in the forward pass without overwriting each other.
+
+    MLA-style pools have no V half (v_head / v_tail = None). SWA pools have two
+    CanaryBufferGroup instances (FULL sized to the full sub-pool, SWA sized to the swa sub-pool).
+
+    Fields:
+        kind: PoolKind.FULL or PoolKind.SWA.
+        k_head: Head canary buffer for K-half launches, shape [num_slots, CANARY_SLOT_BYTES], uint8.
+        k_tail: Tail canary buffer for K-half launches, same shape, uint8.
+        v_head: Same for V-half, or None for MLA-style pools.
+        v_tail: Same for V-half, or None.
+        swa_index_lut: SWA full-to-swa index mapping LUT, shape [full_pool_size + 1], int64, or None for FULL
+            groups. Used by launch_canary_plan_kernels to translate verify/seed slot indices at plan time, and by
+            launch_canary_write_kernel to translate write slots inline. None iff kind == PoolKind.FULL.
+        kv_token_id_vs_position_offset: Logical-position offset between a canary slot and the source-of-truth token it
+            fingerprints. 0 for target-style pools (slot ``p`` stores K/V for token at position ``p``); 1 for
+            EAGLE draft pools where the input_ids rotation makes slot ``p`` store K/V for token at position
+            ``p + 1``.
+    """
+
+    kind: PoolKind
+    k_head: torch.Tensor
+    k_tail: torch.Tensor
+    v_head: Optional[torch.Tensor]
+    v_tail: Optional[torch.Tensor]
+    swa_index_lut: Optional[torch.Tensor]
+    kv_token_id_vs_position_offset: int
+
+    @property
+    def has_v_half(self) -> bool:
+        return self.v_head is not None

python/sglang/srt/kv_canary/capacities.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class CanaryLaunchCapacities:
+    """Pre-allocation sizes for the per-forward tensors a SingleForwardManager owns. Computed
+    once at install_canary from ServerArgs + ModelRunner metadata; all fields are upper
+    bounds - actual per-step usage may be smaller but never larger.
+
+    Fields:
+        per_forward_verify_capacity: VerifyPlan row capacity for the per-forward HEAD/TAIL
+            launches. Sized to pool_slot_count * 3 (3x headroom; radix prefix sharing across
+            running reqs can cause sum_r prefix_lens[r] to exceed the pool slot count). When
+            the per-step actual count exceeds this, the plan kernel sets VerifyPlan.enable=0
+            and the verify kernel skips the step; host logs a warn (no install-time raise).
+        per_forward_write_req_capacity: WritePlan row capacity for per-forward writes, also used
+            to size the static PlanInput buffers (= max batch size under cuda graph).
+        per_forward_write_entry_capacity: Capacity for the expected_input_* placeholder tensors,
+            one entry per token written in a single forward.
+    """
+
+    per_forward_verify_capacity: int
+    per_forward_write_req_capacity: int
+    per_forward_write_entry_capacity: int
+
+    def __post_init__(self) -> None:
+        for name, value in (
+            ("per_forward_verify_capacity", self.per_forward_verify_capacity),
+            ("per_forward_write_req_capacity", self.per_forward_write_req_capacity),
+            ("per_forward_write_entry_capacity", self.per_forward_write_entry_capacity),
+        ):
+            if value <= 0:
+                raise ValueError(f"kv-canary: {name} must be positive, got {value}")
+
+    @classmethod
+    def from_args(
+        cls,
+        *,
+        server_args: "ServerArgs",
+        req_to_token_pool_size: int,
+        max_seq_len_per_req: int,
+        pool_slot_count: int,
+    ) -> "CanaryLaunchCapacities":
+        if req_to_token_pool_size <= 0:
+            raise ValueError(
+                "kv-canary: req_to_token_pool_size must be positive, "
+                f"got {req_to_token_pool_size}"
+            )
+        if max_seq_len_per_req <= 0:
+            raise ValueError(
+                "kv-canary: max_seq_len_per_req must be positive, "
+                f"got {max_seq_len_per_req}"
+            )
+        if pool_slot_count <= 0:
+            raise ValueError(
+                f"kv-canary: pool_slot_count must be positive, got {pool_slot_count}"
+            )
+
+        cuda_graph_max_bs = server_args.cuda_graph_max_bs or 0
+        if cuda_graph_max_bs < 0:
+            raise ValueError(
+                f"kv-canary: cuda_graph_max_bs must be non-negative, got {cuda_graph_max_bs}"
+            )
+
+        spec_num_draft_tokens = server_args.speculative_num_draft_tokens
+        if spec_num_draft_tokens is None:
+            spec_num_draft_tokens = 0
+        if spec_num_draft_tokens < 0:
+            raise ValueError(
+                "kv-canary: speculative_num_draft_tokens must be non-negative, "
+                f"got {spec_num_draft_tokens}"
+            )
+
+        max_prefill_tokens = server_args.max_prefill_tokens
+        if max_prefill_tokens <= 0:
+            raise ValueError(
+                f"kv-canary: max_prefill_tokens must be positive, got {max_prefill_tokens}"
+            )
+
+        num_tokens_per_bs = 1
+        if spec_num_draft_tokens:
+            num_tokens_per_bs = max(num_tokens_per_bs, spec_num_draft_tokens)
+
+        max_bs = max(cuda_graph_max_bs, req_to_token_pool_size)
+
+        chunked_prefill_size = server_args.chunked_prefill_size
+        chunked_limit = (
+            chunked_prefill_size
+            if chunked_prefill_size is not None and chunked_prefill_size >= 0
+            else math.inf
+        )
+        max_extend_tokens_per_forward = min(max_prefill_tokens, chunked_limit)
+
+        write_entry_capacity = max(
+            max_bs * num_tokens_per_bs, max_extend_tokens_per_forward
+        )
+
+        # Radix prefix sharing lets sum_r prefix_lens[r] exceed pool_slot_count; observed up to ~2x
+        # on 20 parallel token-oracle prompts. 3x headroom keeps the partial-fallback path
+        # (plan kernel enable=0 + host warn) exceptional. Overflow does not raise at install time.
+        per_forward_verify_capacity = int(pool_slot_count * 3)
+
+        return cls(
+            per_forward_verify_capacity=per_forward_verify_capacity,
+            per_forward_write_req_capacity=max_bs,
+            per_forward_write_entry_capacity=write_entry_capacity,
+        )

python/sglang/srt/kv_canary/config.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+
+class CanaryMode(str, Enum):
+    NONE = "none"
+    LOG = "log"
+    RAISE = "raise"
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class CanaryConfig:
+    """Top-level canary configuration. All knobs live here; nothing reads env vars deeper in the stack.
+
+    Constructed once inside install_canary(server_args, model_runner) via
+    CanaryConfig.from_env(server_args), then frozen and threaded through the canary stack.
+    Subsequent runtime never mutates it.
+
+    Fields:
+        mode: CanaryMode value. none = no canary installed; log = canary runs, violations are logged
+            but do NOT raise (used for production observability + canary self-test perturb); raise =
+            violations propagate to host as RuntimeError after the next D2H pump.
+        ring_capacity: Violation ring capacity (rows in ViolationLog.violation_ring). Sized generously;
+            overflow only drops detail beyond row N, the monotonic counter still grows.
+    """
+
+    mode: CanaryMode
+    ring_capacity: int
+
+    @classmethod
+    def from_env(cls, server_args: "ServerArgs") -> "CanaryConfig":
+        mode_raw = server_args.kv_canary.strip().lower()
+        if mode_raw not in ("none", "log", "raise"):
+            raise ValueError(
+                f"kv-canary: kv_canary must be one of none/log/raise, got {mode_raw!r}"
+            )
+
+        return cls(
+            mode=CanaryMode(mode_raw),
+            ring_capacity=envs.SGLANG_KV_CANARY_RING_CAPACITY.get(),
+        )

python/sglang/srt/kv_canary/endpoint.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import (
+    CanaryLaunchTag,
+    VerifyOrWriteContext,
+    VerifyPlan,
+    launch_canary_verify_kernel,
+)
+from sglang.jit_kernel.kv_canary.write import (
+    WritePlan,
+    launch_canary_write_kernel,
+)
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.state import (
+    CanaryDeviceState,
+    ViolationLog,
+)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class CanaryEndpoint:
+    kernel_kind: CanaryLaunchTag
+    canary_buf: torch.Tensor
+    full_to_swa_index_mapping: Optional[torch.Tensor]
+    slot_run_counter_view: torch.Tensor
+    kernel_run_counter_view: torch.Tensor
+    enable_chain_position_assert: torch.Tensor
+
+    def launch_per_forward(
+        self,
+        *,
+        verify_plan: VerifyPlan,
+        write_plan: WritePlan,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+        enable_write_input_assert: bool,
+        enable_verify_token_assert: bool,
+        expected_inputs: ExpectedInputs,
+        violation_log: ViolationLog,
+    ) -> None:
+        context = self._make_verify_or_write_context(
+            violation_log=violation_log,
+        )
+        launch_canary_verify_kernel(
+            context=context,
+            plan=verify_plan,
+            check_verify_expected_token=enable_verify_token_assert,
+        )
+
+        # SWA endpoints translate the per-token slot indices via a device tensor index op before invoking the write kernel.
+        if self.full_to_swa_index_mapping is not None:
+            out_cache_loc_for_canary = self.full_to_swa_index_mapping[out_cache_loc]
+        else:
+            out_cache_loc_for_canary = out_cache_loc
+        if enable_write_input_assert:
+            expected_input_tokens = expected_inputs.tokens
+            expected_input_positions = expected_inputs.positions
+        else:
+            expected_input_tokens = None
+            expected_input_positions = None
+
+        launch_canary_write_kernel(
+            context=context,
+            plan=write_plan,
+            input_ids=input_ids,
+            positions=positions,
+            out_cache_loc=out_cache_loc_for_canary,
+            enable_write_input_assert=enable_write_input_assert,
+            expected_input_tokens=expected_input_tokens,
+            expected_input_positions=expected_input_positions,
+        )
+
+    def _make_verify_or_write_context(
+        self,
+        *,
+        violation_log: ViolationLog,
+    ) -> VerifyOrWriteContext:
+        return VerifyOrWriteContext(
+            canary_buf=self.canary_buf,
+            kernel_kind=self.kernel_kind,
+            violation_ring=violation_log.violation_ring,
+            violation_write_index=violation_log.violation_write_index,
+            slot_run_counter=self.slot_run_counter_view,
+            kernel_run_counter=self.kernel_run_counter_view,
+            enable_chain_position_assert=self.enable_chain_position_assert,
+        )
+
+
+def _resolve_canary_buf(
+    *,
+    slot: str,
+    half: str,
+    group: CanaryBufferGroup,
+) -> torch.Tensor:
+    if half == "K":
+        if slot == "HEAD":
+            return group.k_head
+        return group.k_tail
+    if slot == "HEAD":
+        return group.v_head
+    return group.v_tail
+
+
+_FULL_LAYOUT: tuple[tuple[CanaryLaunchTag, str, str], ...] = (
+    (CanaryLaunchTag.HEAD_K_FULL, "HEAD", "K"),
+    (CanaryLaunchTag.HEAD_V_FULL, "HEAD", "V"),
+    (CanaryLaunchTag.TAIL_K_FULL, "TAIL", "K"),
+    (CanaryLaunchTag.TAIL_V_FULL, "TAIL", "V"),
+)
+
+
+_SWA_LAYOUT: tuple[tuple[CanaryLaunchTag, str, str], ...] = (
+    (CanaryLaunchTag.HEAD_K_SWA, "HEAD", "K"),
+    (CanaryLaunchTag.HEAD_V_SWA, "HEAD", "V"),
+    (CanaryLaunchTag.TAIL_K_SWA, "TAIL", "K"),
+    (CanaryLaunchTag.TAIL_V_SWA, "TAIL", "V"),
+)
+
+
+def build_endpoints_from_group(
+    *,
+    group: CanaryBufferGroup,
+    device_state: CanaryDeviceState,
+) -> tuple[CanaryEndpoint, ...]:
+    """Enumerate (slot × half) endpoints for one CanaryBufferGroup."""
+    pool_kind = group.kind
+    layout = _FULL_LAYOUT if pool_kind is PoolKind.FULL else _SWA_LAYOUT
+
+    endpoints: list[CanaryEndpoint] = []
+    for tag, slot, half in layout:
+        if half == "V" and not group.has_v_half:
+            continue
+
+        canary_buf = _resolve_canary_buf(slot=slot, half=half, group=group)
+        lut = group.swa_index_lut if pool_kind is PoolKind.SWA else None
+        slot_view = device_state.slot_run_counters[tag.value : tag.value + 1]
+        kernel_view = device_state.kernel_run_counters[tag.value : tag.value + 1]
+        endpoints.append(
+            CanaryEndpoint(
+                kernel_kind=tag,
+                canary_buf=canary_buf,
+                full_to_swa_index_mapping=lut,
+                slot_run_counter_view=slot_view,
+                kernel_run_counter_view=kernel_view,
+                enable_chain_position_assert=device_state.enable_chain_position_assert,
+            )
+        )
+
+    return tuple(endpoints)

python/sglang/srt/kv_canary/expected_inputs.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ExpectedInputs:
+    tokens: torch.Tensor
+    positions: torch.Tensor
+
+    @classmethod
+    def allocate(cls, *, capacity: int, device: torch.device) -> "ExpectedInputs":
+        return cls(
+            tokens=torch.empty(capacity, dtype=torch.int64, device=device),
+            positions=torch.empty(capacity, dtype=torch.int64, device=device),
+        )
+
+    def slice(self, num_tokens: int) -> "ExpectedInputs":
+        return ExpectedInputs(
+            tokens=self.tokens[:num_tokens],
+            positions=self.positions[:num_tokens],
+        )

python/sglang/srt/kv_canary/plan_input.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class PlanInput:
+    """Pre-staged input to launch_canary_plan_kernels for the per-forward path.
+
+    All tensors live on device.
+
+    Fields:
+        req_pool_indices: Per-row ReqToTokenPool row index, shape [bs_capacity], int64.
+            0 = padding sentinel.
+        prefix_lens: Per-req prefix length already written before this step, shape
+            [bs_capacity], int64. Extend → extend_prefix_lens; decode → seq_lens - 1.
+        extend_seq_lens: Per-req tokens being written this step, shape [bs_capacity], int64.
+            Extend length or all-ones for decode.
+        req_to_verify_expected_tokens_valid_lens: Per-req snapshot length on the verify-token pool,
+            shape [bs_capacity], int64. Equals
+            ``len(req.origin_input_ids) + len(req.output_ids)`` at the moment ``ForwardBatch``
+            was built. The plan kernel uses ``valid_lens[req_id]`` as the upper bound on
+            ``sot_pos`` when gathering the expected token; everything past the snapshot
+            (e.g. EAGLE draft / verify positions, or stale residue from a longer recycled
+            slot owner) returns the ``-1`` sentinel and the verify kernel skips the check.
+            The naive build never populates the verify-token-id cross-check, so this stays all
+            zeros and the plan kernel's gather degrades to the ``-1`` skip sentinel.
+
+    Allocated fresh per forward by :class:`SingleForwardManager`. The boundary
+    ForwardBatch token/position/slot tensors must already be int64
+    contiguous (upstream phase-1 hook is responsible).
+    """
+
+    req_pool_indices: torch.Tensor
+    prefix_lens: torch.Tensor
+    extend_seq_lens: torch.Tensor
+    req_to_verify_expected_tokens_valid_lens: torch.Tensor
+
+    def zero_(self) -> None:
+        self.req_pool_indices.zero_()
+        self.prefix_lens.zero_()
+        self.extend_seq_lens.zero_()
+        self.req_to_verify_expected_tokens_valid_lens.zero_()
+
+    @classmethod
+    def allocate(
+        cls,
+        *,
+        bs_capacity: int,
+        device: torch.device,
+    ) -> "PlanInput":
+        return cls(
+            req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
+            prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
+            extend_seq_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
+            req_to_verify_expected_tokens_valid_lens=torch.zeros(
+                bs_capacity, dtype=torch.int64, device=device
+            ),
+        )
+
+    def fill_from_forward_batch(self, *, forward_batch: "ForwardBatch") -> None:
+        req_pool_indices = forward_batch.req_pool_indices
+        bs = int(req_pool_indices.shape[0])
+        capacity = int(self.req_pool_indices.shape[0])
+        if bs > capacity:
+            raise RuntimeError(
+                f"kv-canary: per-forward batch size {bs} exceeds static capacity {capacity}; "
+                "raise the buffer size in CanaryLaunchCapacities"
+            )
+
+        self.zero_()
+        self.req_pool_indices[:bs].copy_(req_pool_indices)
+
+        _extract_prefix_lens_and_extend_seq_lens(
+            forward_batch=forward_batch,
+            out_prefix_lens=self.prefix_lens[:bs],
+            out_extend_seq_lens=self.extend_seq_lens[:bs],
+            bs=bs,
+        )
+
+
+def _extract_prefix_lens_and_extend_seq_lens(
+    *,
+    forward_batch: "ForwardBatch",
+    out_prefix_lens: torch.Tensor,
+    out_extend_seq_lens: torch.Tensor,
+    bs: int,
+) -> None:
+    # TODO: once ForwardMode is refactored upstream so every mode ships a canonical
+    # (prefix_lens, extend_seq_lens) pair on forward_batch, collapse this back to a single
+    # unconditional copy.
+    forward_mode = forward_batch.forward_mode
+    spec_info = forward_batch.spec_info
+    if forward_mode.is_decode_or_idle():
+        # Anchor on ``positions`` (canonical write position) — eagle draft leaves seq_lens
+        # pre-bump so deriving prefix_lens from seq_lens is off-by-one. Padding tail (positions
+        # shorter than bs under cuda-graph padding) keeps whatever stale data it had; the offsets
+        # kernel masks those rows via ``is_active`` before using prefix_lens.
+        positions = forward_batch.positions
+        out_prefix_lens[: positions.shape[0]].copy_(positions.to(torch.int64))
+        out_extend_seq_lens.fill_(1)
+    elif forward_mode.is_target_verify():
+        # Evidence: EagleVerifyInputV2Mixin.prepare_for_v2_verify assigns out_cache_loc in
+        # [seq_lens, seq_lens + draft_token_num) without bumping seq_lens. The target-verify
+        # branch in TRTLLMHAAttnBackend.init_forward_metadata uses seq_lens as the prefix and
+        # tokens_per_req as the query length, so mirror that as seq_lens plus draft_token_num.
+        out_prefix_lens.copy_(forward_batch.seq_lens[:bs].to(torch.int64))
+        out_extend_seq_lens.fill_(int(spec_info.draft_token_num))
+    elif forward_mode.is_draft_extend_v2():
+        # Evidence: EagleDraftInputV2Mixin.prepare_for_extend_to_fill_draft_kvcache bumps
+        # seq_lens by num_draft_tokens. FlashAttentionBackend.init_forward_metadata reads the
+        # draft-extend-v2 query length from spec_info.extend_seq_lens_tensor when available.
+        # CUDA-graph replay passes extend_seq_lens but omits extend_prefix_lens, so derive the
+        # prefix as seq_lens - extend_seq_lens.
+        extend_seq_lens = forward_batch.extend_seq_lens[:bs].to(torch.int64)
+        out_extend_seq_lens.copy_(extend_seq_lens)
+        out_prefix_lens.copy_(
+            forward_batch.seq_lens[:bs].to(torch.int64) - extend_seq_lens
+        )
+    elif forward_mode.is_extend():
+        # Evidence: ForwardBatch.init_new copies batch.prefix_lens and batch.extend_lens into
+        # extend_prefix_lens / extend_seq_lens for non-decode, non-idle modes, matching regular
+        # extend metadata builders that consume those tensors directly.
+        out_prefix_lens.copy_(forward_batch.extend_prefix_lens[:bs].to(torch.int64))
+        out_extend_seq_lens.copy_(forward_batch.extend_seq_lens[:bs].to(torch.int64))
+    else:
+        raise NotImplementedError(
+            f"Unsupported forward mode for kv-canary: {forward_mode}"
+        )

python/sglang/srt/kv_canary/pool_patcher/adapters/mha.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.pool_patcher.buf_info_splice import patch_buf_info_method
+from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf
+
+
+def attach_mha(
+    *,
+    pool: object,
+    device: torch.device,
+    kv_token_id_vs_position_offset: int,
+) -> tuple[CanaryBufferGroup, ...]:
+    num_slots = int(pool.k_buffer[0].shape[0])
+    k_head = alloc_canary_buf(num_slots=num_slots, device=device)
+    k_tail = alloc_canary_buf(num_slots=num_slots, device=device)
+    v_head = alloc_canary_buf(num_slots=num_slots, device=device)
+    v_tail = alloc_canary_buf(num_slots=num_slots, device=device)
+
+    group = CanaryBufferGroup(
+        kind=PoolKind.FULL,
+        k_head=k_head,
+        k_tail=k_tail,
+        v_head=v_head,
+        v_tail=v_tail,
+        swa_index_lut=None,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    patch_buf_info_method(
+        pool,
+        method_name="get_contiguous_buf_infos",
+        group=group,
+        has_v_half=True,
+        page_size=pool.page_size,
+    )
+    return (group,)

python/sglang/srt/kv_canary/pool_patcher/api.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import logging
+from typing import Callable, Dict, Type
+
+import torch
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha
+from sglang.srt.mem_cache.memory_pool import (
+    KVCache,
+    MHATokenToKVPool,
+    MHATokenToKVPoolFP4,
+)
+
+logger = logging.getLogger(__name__)
+
+PoolAttacher = Callable[..., tuple[CanaryBufferGroup, ...]]
+
+_POOL_ATTACHERS: Dict[Type, PoolAttacher] = {
+    MHATokenToKVPool: attach_mha,
+    MHATokenToKVPoolFP4: attach_mha,
+}
+
+
+def register_pool_attacher(pool_class: Type, attacher: PoolAttacher) -> None:
+    _POOL_ATTACHERS[pool_class] = attacher
+
+
+def attach_canary_buffers(
+    *,
+    pool: KVCache,
+    config: CanaryConfig,
+    device: torch.device,
+    kv_token_id_vs_position_offset: int,
+) -> tuple[CanaryBufferGroup, ...]:
+    """Install canary buffers on a KV pool and return the resulting CanaryBufferGroup tuple.
+
+    ``kv_token_id_vs_position_offset`` is propagated into every produced :class:`CanaryBufferGroup` (0 for target
+    pools; 1 for draft pools where the input-ids rotation shifts the slot-to-token mapping by one).
+    """
+    attacher = _POOL_ATTACHERS.get(type(pool))
+    if attacher is None:
+        raise NotImplementedError(
+            f"kv-canary: no attacher registered for pool class {type(pool).__name__}; "
+            f"supported: {sorted(cls.__name__ for cls in _POOL_ATTACHERS)}"
+        )
+
+    groups = attacher(
+        pool=pool,
+        device=device,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    logger.info(
+        "attach_canary_buffers: pool=%s attacher=%s n_groups=%d kinds=%s "
+        "kv_token_id_vs_position_offset=%d",
+        type(pool).__name__,
+        attacher.__name__,
+        len(groups),
+        [g.kind.name for g in groups],
+        kv_token_id_vs_position_offset,
+    )
+    return groups

python/sglang/srt/kv_canary/pool_patcher/buf_info_splice.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import Any, Callable, List, Tuple
+
+import torch
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
+
+BufInfoTriple = Tuple[List[int], List[int], List[int]]
+
+
+def patch_buf_info_method(
+    pool: object,
+    *,
+    method_name: str,
+    group: CanaryBufferGroup,
+    has_v_half: bool,
+    page_size: int,
+) -> None:
+    """Wrap ``pool.<method_name>()`` so its (ptrs, lens, item_lens) triple is spliced with K/V
+    head and tail entries from ``group``."""
+
+    def _with_splice(original: Callable, *args: Any, **kwargs: Any) -> BufInfoTriple:
+        ptrs, lens, item_lens = original(*args, **kwargs)
+        return splice_kv_buf_info(
+            ptrs=ptrs,
+            lens=lens,
+            item_lens=item_lens,
+            group=group,
+            has_v_half=has_v_half,
+            page_size=page_size,
+        )
+
+    wrap_method(pool, method_name, wrapper=_with_splice)
+
+
+def splice_kv_buf_info(
+    *,
+    ptrs: List[int],
+    lens: List[int],
+    item_lens: List[int],
+    group: CanaryBufferGroup,
+    has_v_half: bool,
+    page_size: int,
+) -> BufInfoTriple:
+    entries = list(zip(ptrs, lens, item_lens))
+    k_head = _entry_triple(group.k_head, page_size=page_size)
+    k_tail = _entry_triple(group.k_tail, page_size=page_size)
+
+    if not has_v_half:
+        out = [k_head, *entries, k_tail]
+    else:
+        assert group.v_head is not None and group.v_tail is not None
+        v_head = _entry_triple(group.v_head, page_size=page_size)
+        v_tail = _entry_triple(group.v_tail, page_size=page_size)
+        if len(entries) % 2 != 0:
+            raise RuntimeError(
+                f"kv-canary: K/V split adapter expects even-length buf_info list, got {len(entries)}"
+            )
+        mid = len(entries) // 2
+        out = [k_head, *entries[:mid], k_tail, v_head, *entries[mid:], v_tail]
+
+    return _untranspose_entries(out)
+
+
+def _entry_triple(buf: torch.Tensor, *, page_size: int) -> Tuple[int, int, int]:
+    return (
+        buf.data_ptr(),
+        buf.nbytes,
+        buf[0].nbytes * page_size,
+    )
+
+
+def _untranspose_entries(entries: List[Tuple[int, int, int]]) -> BufInfoTriple:
+    out_ptrs, out_lens, out_item_lens = (list(col) for col in zip(*entries))
+    return out_ptrs, out_lens, out_item_lens

python/sglang/srt/kv_canary/pool_patcher/buffer_alloc.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
+
+
+def alloc_canary_buf(
+    *,
+    num_slots: int,
+    device: torch.device,
+) -> torch.Tensor:
+    return torch.zeros(num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device)

python/sglang/srt/kv_canary/pool_patcher/utils.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import functools
+from typing import Any, Callable
+
+_WRAPPED_MARKER_ATTR = "_kv_canary_wrapped_by"
+
+
+def wrap_method(
+    obj: object,
+    method_name: str,
+    *,
+    wrapper: Callable[..., Any],
+) -> None:
+    """Replace ``obj.method_name`` with a closure that delegates to ``wrapper``.
+
+    ``wrapper(original, *args, **kwargs)`` receives the original bound method as its first arg and the
+    call-site args/kwargs as the rest. It decides when (and whether) to call ``original`` and what to
+    return. The patched callable is installed as a plain function; :func:`functools.wraps` preserves
+    ``__name__`` / ``__doc__`` but the bound-method nature of the original is not retained.
+
+    Raises:
+        AttributeError: ``obj`` has no attribute ``method_name``.
+        RuntimeError: ``obj.method_name`` has already been wrapped by ``wrap_method`` (idempotency
+            guard — re-wrapping silently would stack two transforms and corrupt return values).
+    """
+    if not hasattr(obj, method_name):
+        raise AttributeError(
+            f"kv-canary: {type(obj).__name__} missing required method {method_name!r}"
+        )
+    original = getattr(obj, method_name)
+    if getattr(original, _WRAPPED_MARKER_ATTR, None) is not None:
+        raise RuntimeError(
+            f"kv-canary: {type(obj).__name__}.{method_name} already wrapped by kv-canary"
+        )
+
+    @functools.wraps(original)
+    def patched(*args: Any, **kwargs: Any) -> Any:
+        return wrapper(original, *args, **kwargs)
+
+    setattr(patched, _WRAPPED_MARKER_ATTR, method_name)
+    setattr(obj, method_name, patched)

python/sglang/srt/kv_canary/runner/canary_manager.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import contextlib
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Iterator, Optional, Sequence
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.environ import envs
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.endpoint import (
+    CanaryEndpoint,
+    build_endpoints_from_group,
+)
+from sglang.srt.kv_canary.runner.violation_manager import ViolationManager
+from sglang.srt.kv_canary.single_forward_manager.manager import (
+    SingleForwardManager,
+    _PreOpsMaybeInsideGraphOutput,
+)
+from sglang.srt.kv_canary.state import CanaryDeviceState
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+class CanaryManager:
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        buffer_groups: tuple[CanaryBufferGroup, ...],
+        device: torch.device,
+        req_to_token_pool: "ReqToTokenPool",
+        launch_capacities: CanaryLaunchCapacities,
+        swa_window_size: int = 0,
+    ) -> None:
+        self.config = config
+        self._req_to_token_pool = req_to_token_pool
+        self._swa_window_size = swa_window_size
+        self._outer_step_counter: int = 0
+        self._active_single_forward_manager_index: Optional[int] = None
+
+        self._buffer_groups: tuple[CanaryBufferGroup, ...] = tuple(buffer_groups)
+
+        self._device_state = CanaryDeviceState.allocate(
+            config=config,
+            device=device,
+            num_tags=len(CanaryLaunchTag),
+            req_to_token_alloc_size=req_to_token_pool.req_to_token.shape[0],
+            max_context_len=req_to_token_pool.max_context_len,
+        )
+        # Disable the chain-step position assert until warmup / cuda-graph capture finishes
+        # (synthetic positions trip the +1 invariant). mark_init_finished() sets it to 1.
+        self._device_state.enable_chain_position_assert.fill_(0)
+
+        self._endpoints: tuple[CanaryEndpoint, ...] = tuple(
+            endpoint
+            for group in self._buffer_groups
+            for endpoint in build_endpoints_from_group(
+                group=group, device_state=self._device_state
+            )
+        )
+        self._active_tags: tuple[CanaryLaunchTag, ...] = tuple(
+            sorted(
+                {endpoint.kernel_kind for endpoint in self._endpoints},
+                key=lambda tag: tag.value,
+            )
+        )
+
+        self._d2h_stream: torch.cuda.Stream = torch.cuda.Stream(device=device)
+
+        self._violation_manager = ViolationManager(
+            config=config,
+            device_state=self._device_state,
+            d2h_stream=self._d2h_stream,
+            outer_step_counter_getter=self._get_outer_step_counter,
+        )
+        self._single_forward_managers: tuple[SingleForwardManager, ...] = (
+            SingleForwardManager(
+                config=config,
+                device=device,
+                device_state=self._device_state,
+                buffer_groups=self._buffer_groups,
+                endpoints=self._endpoints,
+                req_to_token_pool=req_to_token_pool,
+                swa_window_size=self._swa_window_size,
+                per_forward_verify_capacity=launch_capacities.per_forward_verify_capacity,
+                per_forward_write_req_capacity=launch_capacities.per_forward_write_req_capacity,
+                per_forward_write_entry_capacity=launch_capacities.per_forward_write_entry_capacity,
+                d2h_stream=self._d2h_stream,
+            ),
+        )
+
+    @contextlib.contextmanager
+    def with_active_single_forward_manager(self, index: int) -> Iterator[None]:
+        assert (
+            self._active_single_forward_manager_index is None
+        ), "kv-canary: nested with_active_single_forward_manager is forbidden"
+        self._active_single_forward_manager_index = index
+        try:
+            yield
+        finally:
+            assert self._active_single_forward_manager_index == index, (
+                f"kv-canary: with_active_single_forward_manager({index}) exited with "
+                f"_active_single_forward_manager_index="
+                f"{self._active_single_forward_manager_index}; nested or mismatched bracket"
+            )
+            self._active_single_forward_manager_index = None
+
+    def pre_ops_maybe_inside_graph(
+        self, forward_batch: "ForwardBatch"
+    ) -> _PreOpsMaybeInsideGraphOutput:
+        assert self._active_single_forward_manager_index is not None, (
+            "kv-canary: pre_ops_maybe_inside_graph called without active SingleForwardManager; "
+            "caller must wrap in CanaryManager.with_active_single_forward_manager(i)"
+        )
+        sfm = self._single_forward_managers[self._active_single_forward_manager_index]
+        return sfm.pre_ops_maybe_inside_graph(forward_batch)
+
+    def post_ops_maybe_inside_graph(
+        self,
+        forward_batch: "ForwardBatch",
+        pre_ops_output: _PreOpsMaybeInsideGraphOutput,
+    ) -> None:
+        assert self._active_single_forward_manager_index is not None, (
+            "kv-canary: post_ops_maybe_inside_graph called without active SingleForwardManager; "
+            "caller must wrap in CanaryManager.with_active_single_forward_manager(i)"
+        )
+        sfm = self._single_forward_managers[self._active_single_forward_manager_index]
+        sfm.post_ops_maybe_inside_graph(forward_batch, pre_ops_output)
+
+    @contextlib.contextmanager
+    def with_ops_outside_graph(
+        self,
+        *,
+        single_forward_indices: Sequence[int],
+        maybe_inaccurate_forward_batch: "ForwardBatch",
+    ) -> Iterator[None]:
+        self._pre_ops_outside_graph(
+            single_forward_indices=single_forward_indices,
+            maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+        )
+        try:
+            yield
+        finally:
+            self._post_ops_outside_graph(
+                single_forward_indices=single_forward_indices,
+                maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            )
+
+    def _pre_ops_outside_graph(
+        self,
+        *,
+        single_forward_indices: Sequence[int],
+        maybe_inaccurate_forward_batch: "ForwardBatch",
+    ) -> None:
+        for idx in single_forward_indices:
+            self._single_forward_managers[idx].pre_ops_outside_graph(
+                maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch
+            )
+
+    def _post_ops_outside_graph(
+        self,
+        *,
+        single_forward_indices: Sequence[int],
+        maybe_inaccurate_forward_batch: "ForwardBatch",
+    ) -> None:
+        for idx in single_forward_indices:
+            self._single_forward_managers[idx].post_ops_outside_graph()
+        self._outer_step_counter += 1
+        self._violation_manager.step()
+
+    def mark_init_finished(self) -> None:
+        for single_forward_manager in self._single_forward_managers:
+            single_forward_manager.phase_checker.enable_assert()
+        self._device_state.enable_chain_position_assert.fill_(1)
+
+    def _get_outer_step_counter(self) -> int:
+        return self._outer_step_counter
+
+
+@contextmanager
+def context_tuple(ctx_a, ctx_b):
+    with ctx_a, ctx_b:
+        yield

python/sglang/srt/kv_canary/runner/enable_warner.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
+
+logger = logging.getLogger(__name__)
+
+
+class CanaryEnableWarner:
+    def __init__(
+        self, *, verify_capacity: int, d2h_stream: Optional[torch.cuda.Stream]
+    ) -> None:
+        self._verify_capacity = verify_capacity
+        self._overflow_count_total: int = 0
+        self._handler = DelayedDeviceHostHandler(d2h_stream=d2h_stream)
+
+    def tick(self, enable_device: torch.Tensor) -> None:
+        self._handler.step(
+            compute_on_device=lambda: enable_device,
+            postprocess_on_host=self._postprocess_on_host,
+        )
+
+    def _postprocess_on_host(self, host_tensor: torch.Tensor) -> None:
+        if int(host_tensor.item()) == 0:
+            self._overflow_count_total += 1
+            logger.warning(
+                "kv-canary: per-forward verify skipped this step due to overflow "
+                "(total=%d, capacity=%d); check ServerArgs / pool sizing",
+                self._overflow_count_total,
+                self._verify_capacity,
+            )

python/sglang/srt/kv_canary/runner/future_tensor.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, Optional, Union
+
+import torch
+
+_PayloadDict = dict[str, Any]
+_TensorOrDict = Union[torch.Tensor, _PayloadDict]
+
+_DUMMY_DICT_KEY = "__dummy_key__"
+
+
+@dataclass(slots=True, kw_only=True)
+class FutureTensors:
+    _data: Optional[_PayloadDict]
+    _event: Optional[torch.cuda.Event]
+    # Device-source clones must outlive the async d2h copy.
+    _retained_device_clones: Optional[dict[str, torch.Tensor]] = None
+
+    @classmethod
+    def device_to_host(
+        cls, xs_device: _TensorOrDict, *, d2h_stream: torch.cuda.Stream
+    ) -> "FutureTensors":
+        assert not torch.cuda.is_current_stream_capturing(), (
+            "FutureTensors.device_to_host must not be called during cuda-graph "
+            "capture: the d2h side-stream copy + pinned-host alloc cannot be "
+            "captured. Upper-layer callers are responsible for placing the d2h "
+            "staging OUTSIDE the cuda graph (not inside it)."
+        )
+        if not isinstance(xs_device, dict):
+            xs_device = {_DUMMY_DICT_KEY: xs_device}
+
+        first_tensor = next(
+            (x for x in xs_device.values() if isinstance(x, torch.Tensor)), None
+        )
+        if first_tensor is None:
+            raise ValueError(
+                f"FutureTensors.device_to_host requires at least one tensor entry; "
+                f"got dict with keys={list(xs_device)} containing no Tensor"
+            )
+        device = first_tensor.device
+        del first_tensor
+
+        tensors_device = {
+            k: v for k, v in xs_device.items() if isinstance(v, torch.Tensor)
+        }
+        non_tensors_device = {
+            k: v for k, v in xs_device.items() if not isinstance(v, torch.Tensor)
+        }
+        del xs_device
+
+        # Must happen in current stream, not d2h stream
+        tensors_device_cloned = {
+            key: x.detach().clone() for key, x in tensors_device.items()
+        }
+
+        tensors_host = {
+            key: torch.empty(x.shape, dtype=x.dtype, pin_memory=True)
+            for key, x in tensors_device.items()
+        }
+
+        d2h_stream.wait_stream(torch.cuda.current_stream(device))
+        with torch.cuda.stream(d2h_stream):
+            for key in tensors_device_cloned:
+                tensors_host[key].copy_(tensors_device_cloned[key], non_blocking=True)
+            event = torch.cuda.Event()
+            event.record()
+
+        return cls(
+            _data=tensors_host | non_tensors_device,
+            _event=event,
+            _retained_device_clones=tensors_device_cloned,
+        )
+
+    def wait(self) -> _TensorOrDict:
+        data = self._data
+        event = self._event
+        retained_device_clones = self._retained_device_clones
+        self._data = None
+        self._event = None
+        self._retained_device_clones = None
+
+        if data is None or event is None:
+            raise RuntimeError("FutureTensors.wait() was called more than once")
+
+        # Releasing clones AFTER event.synchronize() so the d2h copy
+        # finishes reading from them before they become free-able.
+        event.synchronize()
+        del retained_device_clones
+
+        if _DUMMY_DICT_KEY in data:
+            data = data[_DUMMY_DICT_KEY]
+
+        return data
+
+
+@dataclass(slots=True, kw_only=True)
+class DelayedDeviceHostHandler:
+    """Stage device-side compute at step T, drain + postprocess host copy at step T+1."""
+
+    d2h_stream: torch.cuda.Stream
+    _future: Optional[FutureTensors] = field(default=None)
+
+    def step(
+        self,
+        *,
+        compute_on_device: Callable[[], Optional[_TensorOrDict]],
+        postprocess_on_host: Callable[[_TensorOrDict], None],
+    ) -> None:
+        if (pending := self._future) is not None:
+            postprocess_on_host(pending.wait())
+            self._future = None
+
+        # Must run on current stream, not d2h stream
+        device_data = compute_on_device()
+
+        if device_data is None:
+            self._future = None
+        else:
+            self._future = FutureTensors.device_to_host(
+                device_data, d2h_stream=self.d2h_stream
+            )

python/sglang/srt/kv_canary/runner/kernel_launcher.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
+from sglang.jit_kernel.kv_canary.verify import (
+    CanaryLaunchTag,
+    VerifyPlan,
+)
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.srt.environ import envs
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.endpoint import CanaryEndpoint
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.plan_input import PlanInput
+from sglang.srt.kv_canary.state import ViolationLog
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+_BOUNDARY_INT_DTYPES = (torch.int32, torch.int64)
+_INPUT_IDS = "forward_batch.input_ids"
+_OUT_LOC = "forward_batch.out_cache_loc"
+_POSITIONS = "forward_batch.positions"
+
+
+def invoke_plan(
+    *,
+    plan_input: PlanInput,
+    verify_plan: VerifyPlan,
+    write_plan: WritePlan,
+    group: CanaryBufferGroup,
+    req_to_token: torch.Tensor,
+    swa_window_size: int,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+) -> None:
+    window = swa_window_size if group.kind is PoolKind.SWA else 0
+    launch_canary_plan_kernels(
+        verify_plan_out=verify_plan,
+        write_plan_out=write_plan,
+        req_pool_indices=plan_input.req_pool_indices,
+        prefix_lens=plan_input.prefix_lens,
+        extend_seq_lens=plan_input.extend_seq_lens,
+        req_to_token=req_to_token,
+        swa_window_size=window,
+        full_to_swa_index_mapping=group.swa_index_lut,
+        verify_capacity=int(verify_plan.verify_slot_indices.shape[0]),
+        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        req_to_verify_expected_tokens_valid_lens=plan_input.req_to_verify_expected_tokens_valid_lens,
+        kv_token_id_vs_position_offset=group.kv_token_id_vs_position_offset,
+    )
+
+
+def launch_endpoints_per_forward(
+    *,
+    endpoints: tuple[CanaryEndpoint, ...],
+    group: CanaryBufferGroup,
+    tag_filter: Callable[[CanaryLaunchTag], bool],
+    verify_plan: VerifyPlan,
+    write_plan: WritePlan,
+    forward_batch: "ForwardBatch",
+    expected_inputs: ExpectedInputs,
+    violation_log: ViolationLog,
+    enable_write_input_assert: bool = False,
+    enable_verify_token_assert: bool = False,
+) -> None:
+    positions = _canonicalize_boundary_int64(forward_batch.positions, _POSITIONS)
+    out_cache_loc = _canonicalize_boundary_int64(forward_batch.out_cache_loc, _OUT_LOC)
+    input_ids = _canonicalize_boundary_int64(forward_batch.input_ids, _INPUT_IDS)
+
+    num_tokens = int(positions.shape[0])
+    if expected_inputs.tokens.shape[0] != num_tokens:
+        raise RuntimeError(
+            f"kv-canary: expected_inputs.tokens shape {expected_inputs.tokens.shape[0]} "
+            f"!= num_tokens {num_tokens}; caller must slice before invoking"
+        )
+    if expected_inputs.positions.shape[0] != num_tokens:
+        raise RuntimeError(
+            f"kv-canary: expected_inputs.positions shape {expected_inputs.positions.shape[0]} "
+            f"!= num_tokens {num_tokens}; caller must slice before invoking"
+        )
+
+    active_endpoints = [
+        endpoint
+        for endpoint in endpoints
+        if _endpoint_belongs_to_group(endpoint, group)
+        and tag_filter(endpoint.kernel_kind)
+        and passes_v_half_gate(endpoint.kernel_kind)
+    ]
+    assert len(active_endpoints) > 0
+
+    for endpoint in active_endpoints:
+        endpoint.launch_per_forward(
+            verify_plan=verify_plan,
+            write_plan=write_plan,
+            input_ids=input_ids,
+            positions=positions,
+            out_cache_loc=out_cache_loc,
+            enable_write_input_assert=enable_write_input_assert,
+            enable_verify_token_assert=enable_verify_token_assert,
+            expected_inputs=expected_inputs,
+            violation_log=violation_log,
+        )
+
+
+def _is_v_half_tag(tag: CanaryLaunchTag) -> bool:
+    return tag in (
+        CanaryLaunchTag.HEAD_V_FULL,
+        CanaryLaunchTag.TAIL_V_FULL,
+        CanaryLaunchTag.HEAD_V_SWA,
+        CanaryLaunchTag.TAIL_V_SWA,
+    )
+
+
+def passes_v_half_gate(tag: CanaryLaunchTag) -> bool:
+    if not _is_v_half_tag(tag):
+        return True
+    return envs.SGLANG_KV_CANARY_ENABLE_MHA_V.get()
+
+
+def _endpoint_belongs_to_group(
+    endpoint: CanaryEndpoint, group: CanaryBufferGroup
+) -> bool:
+    suffix = endpoint.kernel_kind.name.rsplit("_", 1)[1]
+    return suffix == group.kind.name
+
+
+def _canonicalize_boundary_int64(
+    tensor: torch.Tensor | None, name: str
+) -> torch.Tensor | None:
+    if tensor is None:
+        return None
+    if tensor.dtype not in _BOUNDARY_INT_DTYPES:
+        raise TypeError(
+            f"kv-canary: {name} must have dtype torch.int32 or torch.int64, got {tensor.dtype}"
+        )
+    return tensor.to(torch.int64).contiguous()

python/sglang/srt/kv_canary/runner/violation_manager.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
+from sglang.srt.kv_canary.runner.violation_reporter import ViolationReporter
+from sglang.srt.kv_canary.state import CanaryDeviceState
+
+
+class ViolationManager:
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        device_state: CanaryDeviceState,
+        d2h_stream: torch.cuda.Stream,
+        outer_step_counter_getter: Callable[[], int],
+    ) -> None:
+        self._device_state = device_state
+        self._outer_step_counter_getter = outer_step_counter_getter
+        self._violation_reporter = ViolationReporter(
+            config=config, device_state=device_state
+        )
+        self._handler = DelayedDeviceHostHandler(d2h_stream=d2h_stream)
+
+    def step(self) -> None:
+        drain_result: dict[str, bool] = {"errored": False}
+        self._handler.step(
+            compute_on_device=self._compute_on_device,
+            postprocess_on_host=lambda host: drain_result.update(
+                errored=bool(int(host.item()))
+            ),
+        )
+        if drain_result["errored"] and not self._violation_reporter.is_raised:
+            self._violation_reporter.log_or_raise_violation(
+                outer_step_counter=self._outer_step_counter_getter()
+            )
+
+    def _compute_on_device(self) -> torch.Tensor:
+        violation_log = self._device_state.violation_log
+        return (violation_log.violation_write_index > 0).to(torch.uint8).view(-1)[:1]

python/sglang/srt/kv_canary/runner/violation_reporter.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import logging
+
+from sglang.jit_kernel.kv_canary.consts import FailReason
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.state import CanaryDeviceState
+
+logger = logging.getLogger(__name__)
+
+_WRITE_BITS = FailReason.WRITE_TOKEN_MISMATCH | FailReason.WRITE_POSITION_MISMATCH
+_TOKEN_BITS = FailReason.WRITE_TOKEN_MISMATCH | FailReason.VERIFY_TOKEN_MISMATCH
+
+
+def _reason_label(bit: FailReason) -> str:
+    return bit.name.lower().removesuffix("_mismatch")
+
+
+class ViolationReporter:
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        device_state: CanaryDeviceState,
+    ) -> None:
+        self._config = config
+        self._device_state = device_state
+        self._raised: bool = False
+        self._last_logged_write_index: int = 0
+
+    @property
+    def is_raised(self) -> bool:
+        return self._raised
+
+    def log_or_raise_violation(self, *, outer_step_counter: int) -> None:
+        violation_log = self._device_state.violation_log
+        write_index = int(violation_log.violation_write_index.cpu().item())
+        if write_index == 0:
+            return
+        ring = violation_log.violation_ring.cpu()
+        ring_capacity = int(ring.shape[0])
+        valid_count = min(write_index, ring_capacity)
+        ring_overflow = write_index > ring_capacity
+
+        start = min(self._last_logged_write_index, valid_count)
+        if start >= valid_count:
+            return
+
+        messages: list[str] = [
+            _format_violation(
+                row=ring[i].tolist(),
+                total=write_index,
+                ring_overflow=ring_overflow,
+                step_when_pumped=outer_step_counter,
+            )
+            for i in range(start, valid_count)
+        ]
+        self._last_logged_write_index = valid_count
+
+        # log mode: always surface every violation as WARNING.
+        if self._config.mode is CanaryMode.LOG:
+            for message in messages:
+                logger.warning(message)
+            return
+        self._raised = True
+        raise RuntimeError("\n".join(messages))
+
+
+def _canary_kind_label(tag: CanaryLaunchTag) -> str:
+    name_lower = tag.name.lower()
+    return f"per_forward_{name_lower}"
+
+
+def _format_violation(
+    *,
+    row: list[int],
+    total: int,
+    ring_overflow: bool,
+    step_when_pumped: int,
+) -> str:
+    (
+        kernel_kind,
+        slot_idx,
+        position,
+        stored_token,
+        expected_token,
+        stored_chain_hash,
+        expected_aux,
+        fail_reason_bits,
+    ) = row
+    try:
+        tag_label = CanaryLaunchTag(int(kernel_kind)).name
+        canary_kind = _canary_kind_label(CanaryLaunchTag(int(kernel_kind)))
+    except ValueError:
+        tag_label = f"unknown({int(kernel_kind)})"
+        canary_kind = tag_label
+    bits_int = int(fail_reason_bits)
+    reasons = [_reason_label(bit) for bit in FailReason if bits_int & int(bit)]
+    is_write = bool(bits_int & int(_WRITE_BITS))
+    u64_mask = (1 << 64) - 1
+
+    # Stable single-line key=value summary, parsed by the regex in
+    # python/sglang/test/kv_canary/violation_log_utils.py and asserted by
+    # assert_violation_logged_any in python/sglang/test/kv_canary/violation_assert_mixin.py.
+    # Format frozen: do not reorder / rename / change separators without updating those helpers.
+    structured_line = (
+        f"kv_canary violation: "
+        f"launch_tag={tag_label} "
+        f"fail_reason={'+'.join(reasons) if reasons else 'none'} "
+        f"slot_idx={int(slot_idx)} "
+        f"position={int(position)} "
+        f"stored_token={int(stored_token)} "
+        f"expected_token={int(expected_token)} "
+        f"stored_chain_hash={int(stored_chain_hash) & u64_mask:#018x} "
+        f"expected_aux={int(expected_aux) & u64_mask:#018x}"
+    )
+
+    header = (
+        f"KV cache canary violation detected (kernel_kind={tag_label}, "
+        f"slot_idx={int(slot_idx)}, position={int(position)})"
+    )
+    kind_line = f"canary_kind:       {canary_kind}"
+    reasons_line = f"  fail_reasons: {' '.join(reasons) if reasons else 'none'}"
+    footer = (
+        f"  total_violations={total} ring_overflow={ring_overflow} "
+        f"step_when_pumped={step_when_pumped}"
+    )
+
+    has_token_check = bool(bits_int & int(_TOKEN_BITS))
+    if is_write:
+        running_prev_hash = int(stored_chain_hash) & u64_mask
+        body = [
+            (
+                f"  actual:   token_id={int(stored_token)}   position={int(position)} "
+                f"prev_hash={running_prev_hash:#018x}"
+            ),
+            (
+                f"  expected: token_id={int(expected_token)}   position={int(expected_aux)}"
+            ),
+        ]
+    else:
+        stored_prev_hash = int(stored_chain_hash) & u64_mask
+        expected_prev_hash = int(expected_aux) & u64_mask
+        stored_body = (
+            f"  stored:   token_id={int(stored_token)}   position={int(position)} "
+            f"prev_hash={stored_prev_hash:#018x}"
+        )
+        expected_body = (
+            f"  expected: token_id={int(expected_token)} prev_hash={expected_prev_hash:#018x}"
+            if has_token_check
+            else f"  expected: prev_hash={expected_prev_hash:#018x}"
+        )
+        body = [stored_body, expected_body]
+
+    return "\n".join([structured_line, header, kind_line, reasons_line, *body, footer])

python/sglang/srt/kv_canary/single_forward_manager/data.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class PostOpsInsideGraphOutputBuffer:
+    verify_plan_enable: torch.Tensor
+    kernel_run_counters: torch.Tensor
+    slot_run_counters: torch.Tensor
+    violation_write_index: torch.Tensor
+
+    @classmethod
+    def allocate(
+        cls,
+        *,
+        num_kernel_tags: int,
+        num_slot_tags: int,
+        device: torch.device,
+    ) -> "PostOpsInsideGraphOutputBuffer":
+        return cls(
+            verify_plan_enable=torch.zeros(1, dtype=torch.int32, device=device),
+            kernel_run_counters=torch.zeros(
+                num_kernel_tags, dtype=torch.int64, device=device
+            ),
+            slot_run_counters=torch.zeros(
+                num_slot_tags, dtype=torch.int64, device=device
+            ),
+            violation_write_index=torch.zeros(1, dtype=torch.int32, device=device),
+        )
+
+    def copy_from(
+        self,
+        *,
+        verify_plan_enable: torch.Tensor,
+        kernel_run_counters: torch.Tensor,
+        slot_run_counters: torch.Tensor,
+        violation_write_index: torch.Tensor,
+    ) -> None:
+        self.verify_plan_enable.copy_(verify_plan_enable)
+        self.kernel_run_counters.copy_(kernel_run_counters)
+        self.slot_run_counters.copy_(slot_run_counters)
+        self.violation_write_index.copy_(violation_write_index)

python/sglang/srt/kv_canary/single_forward_manager/manager.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.endpoint import CanaryEndpoint
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.plan_input import PlanInput
+from sglang.srt.kv_canary.runner.enable_warner import CanaryEnableWarner
+from sglang.srt.kv_canary.runner.kernel_launcher import (
+    invoke_plan,
+    launch_endpoints_per_forward,
+)
+from sglang.srt.kv_canary.single_forward_manager.data import (
+    PostOpsInsideGraphOutputBuffer,
+)
+from sglang.srt.kv_canary.state import CanaryDeviceState
+from sglang.srt.utils.phase_checker import SimplePhaseChecker
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class _SingleForwardPhase(IntEnum):
+    IDLE = 0
+    AFTER_PRE_OUT = 1
+    AFTER_PRE_MAYBE_IN = 2
+    AFTER_POST_MAYBE_IN = 3
+
+
+def _torch_reduce_minimum(tensors: list[torch.Tensor]) -> torch.Tensor:
+    out = tensors[0]
+    for t in tensors[1:]:
+        out = torch.minimum(out, t)
+    return out
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class _PreOpsMaybeInsideGraphOutput:
+    verify_plans: tuple[VerifyPlan, ...]
+    write_plans: tuple[WritePlan, ...]
+    expected_inputs: ExpectedInputs
+
+
+class SingleForwardManager:
+    """Owns the state of one inner ``model.forward`` invocation."""
+
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        device: torch.device,
+        device_state: CanaryDeviceState,
+        buffer_groups: tuple[CanaryBufferGroup, ...],
+        endpoints: tuple[CanaryEndpoint, ...],
+        req_to_token_pool: "ReqToTokenPool",
+        swa_window_size: int,
+        per_forward_verify_capacity: int,
+        per_forward_write_req_capacity: int,
+        per_forward_write_entry_capacity: int,
+        d2h_stream: torch.cuda.Stream,
+    ) -> None:
+        self._config = config
+        self._device = device
+        self._device_state = device_state
+        self._buffer_groups = buffer_groups
+        self._endpoints = endpoints
+        self._req_to_token_pool = req_to_token_pool
+        self._swa_window_size = swa_window_size
+        self._d2h_stream = d2h_stream
+
+        self._write_req_capacity = per_forward_write_req_capacity
+        self._write_entry_capacity = per_forward_write_entry_capacity
+        self._verify_capacity = per_forward_verify_capacity
+
+        self._enable_warner = CanaryEnableWarner(
+            verify_capacity=self._verify_capacity,
+            d2h_stream=d2h_stream,
+        )
+
+        self._phase_checker = SimplePhaseChecker(
+            initial_phase=_SingleForwardPhase.IDLE, device=device
+        )
+
+        self._output_buffer = PostOpsInsideGraphOutputBuffer.allocate(
+            num_kernel_tags=int(device_state.kernel_run_counters.shape[0]),
+            num_slot_tags=int(device_state.slot_run_counters.shape[0]),
+            device=device,
+        )
+
+    @property
+    def phase_checker(self) -> SimplePhaseChecker:
+        return self._phase_checker
+
+    def pre_ops_outside_graph(
+        self, *, maybe_inaccurate_forward_batch: "ForwardBatch"
+    ) -> None:
+        self._phase_checker.update(
+            expect_phase=_SingleForwardPhase.IDLE,
+            next_phase=_SingleForwardPhase.AFTER_PRE_OUT,
+            caller_name="SingleForwardManager.pre_ops_outside_graph",
+        )
+
+        bs = int(maybe_inaccurate_forward_batch.batch_size)
+        num_tokens = int(maybe_inaccurate_forward_batch.positions.shape[0])
+        if bs > self._write_req_capacity:
+            raise RuntimeError(
+                f"kv-canary: forward_batch.batch_size={bs} exceeds pre-allocated "
+                f"write_req_capacity={self._write_req_capacity}; raise --cuda-graph-max-bs "
+                f"or check CanaryLaunchCapacities.from_args"
+            )
+        if num_tokens > self._write_entry_capacity:
+            raise RuntimeError(
+                f"kv-canary: forward_batch token count={num_tokens} exceeds pre-allocated "
+                f"write_entry_capacity={self._write_entry_capacity}; raise "
+                f"--chunked-prefill-size / --max-prefill-tokens or check "
+                f"CanaryLaunchCapacities.from_args"
+            )
+
+    def pre_ops_maybe_inside_graph(
+        self, forward_batch: "ForwardBatch"
+    ) -> "_PreOpsMaybeInsideGraphOutput":
+        self._phase_checker.update(
+            expect_phase=_SingleForwardPhase.AFTER_PRE_OUT,
+            next_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN,
+            caller_name="SingleForwardManager.pre_ops_maybe_inside_graph",
+        )
+
+        verify_plans = tuple(
+            VerifyPlan.allocate(
+                verify_capacity=self._verify_capacity, device=self._device
+            )
+            for _ in self._buffer_groups
+        )
+        write_plans = tuple(
+            WritePlan.allocate(
+                write_req_capacity=self._write_req_capacity, device=self._device
+            )
+            for _ in self._buffer_groups
+        )
+        expected_inputs = ExpectedInputs.allocate(
+            capacity=self._write_entry_capacity, device=self._device
+        )
+        plan_input = PlanInput.allocate(
+            bs_capacity=self._write_req_capacity, device=self._device
+        )
+
+        plan_input.fill_from_forward_batch(forward_batch=forward_batch)
+
+        violation_log = self._device_state.violation_log
+        num_tokens = int(forward_batch.positions.shape[0])
+        expected_inputs_slice = expected_inputs.slice(num_tokens)
+
+        for group_idx, group in enumerate(self._buffer_groups):
+            verify_plan = verify_plans[group_idx]
+            write_plan = write_plans[group_idx]
+            invoke_plan(
+                plan_input=plan_input,
+                verify_plan=verify_plan,
+                write_plan=write_plan,
+                group=group,
+                req_to_token=self._req_to_token_pool.req_to_token,
+                swa_window_size=self._swa_window_size,
+                req_to_verify_expected_tokens=self._device_state.req_to_verify_expected_tokens,
+            )
+            launch_endpoints_per_forward(
+                endpoints=self._endpoints,
+                group=group,
+                tag_filter=_is_head_tag,
+                verify_plan=verify_plan,
+                write_plan=write_plan,
+                forward_batch=forward_batch,
+                expected_inputs=expected_inputs_slice,
+                violation_log=violation_log,
+            )
+
+        return _PreOpsMaybeInsideGraphOutput(
+            verify_plans=verify_plans,
+            write_plans=write_plans,
+            expected_inputs=expected_inputs,
+        )
+
+    def post_ops_maybe_inside_graph(
+        self,
+        forward_batch: "ForwardBatch",
+        pre_ops_output: "_PreOpsMaybeInsideGraphOutput",
+    ) -> None:
+        self._phase_checker.update(
+            expect_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN,
+            next_phase=_SingleForwardPhase.AFTER_POST_MAYBE_IN,
+            caller_name="SingleForwardManager.post_ops_maybe_inside_graph",
+        )
+
+        violation_log = self._device_state.violation_log
+        num_tokens = int(forward_batch.positions.shape[0])
+        expected_inputs_slice = pre_ops_output.expected_inputs.slice(num_tokens)
+        for group_idx, group in enumerate(self._buffer_groups):
+            launch_endpoints_per_forward(
+                endpoints=self._endpoints,
+                group=group,
+                tag_filter=_is_tail_tag,
+                verify_plan=pre_ops_output.verify_plans[group_idx],
+                write_plan=pre_ops_output.write_plans[group_idx],
+                forward_batch=forward_batch,
+                expected_inputs=expected_inputs_slice,
+                violation_log=violation_log,
+            )
+
+        verify_plan_enable_combined = _torch_reduce_minimum(
+            [x.enable for x in pre_ops_output.verify_plans]
+        )
+        self._output_buffer.copy_from(
+            verify_plan_enable=verify_plan_enable_combined,
+            kernel_run_counters=self._device_state.kernel_run_counters,
+            slot_run_counters=self._device_state.slot_run_counters,
+            violation_write_index=self._device_state.violation_log.violation_write_index,
+        )
+
+    def post_ops_outside_graph(self) -> None:
+        self._phase_checker.update(
+            expect_phase=_SingleForwardPhase.AFTER_POST_MAYBE_IN,
+            next_phase=_SingleForwardPhase.IDLE,
+            caller_name="SingleForwardManager.post_ops_outside_graph",
+        )
+
+        self._enable_warner.tick(self._output_buffer.verify_plan_enable)
+
+
+def _is_head_tag(tag: CanaryLaunchTag) -> bool:
+    return tag in (
+        CanaryLaunchTag.HEAD_K_FULL,
+        CanaryLaunchTag.HEAD_V_FULL,
+        CanaryLaunchTag.HEAD_K_SWA,
+        CanaryLaunchTag.HEAD_V_SWA,
+    )
+
+
+def _is_tail_tag(tag: CanaryLaunchTag) -> bool:
+    return tag in (
+        CanaryLaunchTag.TAIL_K_FULL,
+        CanaryLaunchTag.TAIL_V_FULL,
+        CanaryLaunchTag.TAIL_K_SWA,
+        CanaryLaunchTag.TAIL_V_SWA,
+    )

python/sglang/srt/kv_canary/state.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.consts import VIOLATION_FIELDS
+from sglang.srt.kv_canary.config import CanaryConfig
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ViolationLog:
+    """Global violation sink shared across all canary launches.
+
+    One instance per canary runner — every launch (head / tail, K / V half, FULL / SWA group) writes
+    into the same ring. The kernel_kind field stamped into each violation row identifies which launch fired
+    (kernel_kind is a static IntEnum tag — :class:`CanaryLaunchTag` in
+    ``sglang.jit_kernel.kv_canary.verify`` — with a unique value per (head|tail, K|V, FULL|SWA) tuple).
+
+    Ring capacity is sized generously (≥ 1024) so overflow is a non-concern in practice — violations are
+    cold-path and the host raises at the first one anyway (or just logs it in mode="log"). atomicAdd
+    contention on a single counter is also negligible since violation events are rare.
+
+    Derived state (host computes on read; not stored):
+        is_errored = violation_write_index[0] > 0
+        first_violation = violation_ring[0]   (valid iff is_errored)
+        ring_valid_count = min(violation_write_index[0], ring_capacity)
+
+    The ring is fill-once: writes beyond ring_capacity are dropped but the counter still increments. Whoever
+    wins atomicAdd for idx == 0 permanently occupies row 0.
+
+    Fields:
+        violation_ring: Append-only violation sink, shape [ring_capacity, VIOLATION_FIELDS], int64. Row 0 is
+            the first violation; rows 1..min(write_index, capacity) follow in atomic order. Fill-once.
+        violation_write_index: Monotonic violation counter, shape [1], int32. Incremented on every violation
+            regardless of ring capacity.
+    """
+
+    violation_ring: torch.Tensor
+    violation_write_index: torch.Tensor
+
+    @classmethod
+    def allocate(cls, *, ring_capacity: int, device: torch.device) -> "ViolationLog":
+        if ring_capacity <= 0:
+            raise ValueError(
+                f"kv-canary: ViolationLog ring_capacity must be positive, got {ring_capacity}"
+            )
+        return cls(
+            violation_ring=torch.zeros(
+                ring_capacity, VIOLATION_FIELDS, dtype=torch.int64, device=device
+            ),
+            violation_write_index=torch.zeros(1, dtype=torch.int32, device=device),
+        )
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class CanaryDeviceState:
+    """Device-side state owned by one CanaryManager instance.
+
+    One instance per ModelRunner. Held on the same device as the KV pool. All tensors are allocated up
+    front (sizes fixed by CanaryConfig + cuda-graph capture capacity) and reused across forward steps —
+    no per-step allocation.
+
+    Fields:
+        violation_log: The single ViolationLog shared by every launch (head / tail × K / V ×
+            FULL / SWA). All kernels atomicAdd into violation_log.violation_write_index and stamp their
+            CanaryLaunchTag into each violation row.
+        kernel_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. The
+            kernel itself does NOT index this array; runner takes a 1-element view at tag's slot (via
+            CanaryEndpoint.kernel_run_counter_view) and hands a shape [1] tensor to the kernel,
+            which atomicAdds 1 regardless of whether the plan had any active entry. Health watchdog
+            reads this array to confirm "canary path actually ran".
+        slot_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. Same
+            view-handed-to-kernel pattern as kernel_run_counters; each launch adds its active entry
+            count to its slot. Used for periodic stats ("protected N tokens").
+        enable_chain_position_assert: int32 [1] device flag gating the write kernel's chain-step
+            write_position assert. allocate() defaults to 1; CanaryManager zeros it during
+            __init__ for the warmup window and mark_init_finished() flips it back to 1.
+        req_to_verify_expected_tokens: Optional int32 device tensor shape
+            ``[req_to_token_alloc_size, max_context_len]``. Mirrors ReqToTokenPool layout;
+            ``pool[req_idx, p]`` = source-of-truth token at logical position ``p`` for the
+            req in slot ``req_idx``. The plan-side entries kernel gathers from this pool (via
+            ``kv_token_id_vs_position_offset`` per buffer group) into
+            ``VerifyPlan.verify_expected_tokens``; the verify kernel then compares against each
+            canary slot's stored token. The naive build never populates the verify-token-id
+            cross-check, so this is always ``None`` and the gather degrades to the ``-1`` sentinel.
+    """
+
+    violation_log: ViolationLog
+    kernel_run_counters: torch.Tensor
+    slot_run_counters: torch.Tensor
+    enable_chain_position_assert: torch.Tensor
+    req_to_verify_expected_tokens: Optional[torch.Tensor]
+
+    @classmethod
+    def allocate(
+        cls,
+        *,
+        config: CanaryConfig,
+        device: torch.device,
+        num_tags: int,
+        req_to_token_alloc_size: Optional[int] = None,
+        max_context_len: Optional[int] = None,
+    ) -> "CanaryDeviceState":
+        if num_tags <= 0:
+            raise ValueError(
+                f"kv-canary: CanaryDeviceState num_tags must be positive, got {num_tags}"
+            )
+        violation_log = ViolationLog.allocate(
+            ring_capacity=config.ring_capacity, device=device
+        )
+        kernel_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device)
+        slot_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device)
+        enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device)
+        # The naive build does not run the verify-token-id cross-check, so the source-of-truth
+        # token pool is never allocated. The field is kept on the dataclass and downstream code
+        # (plan kernel gather) treats ``None`` as "emit the -1 skip sentinel".
+        req_to_verify_expected_tokens = None
+        return cls(
+            violation_log=violation_log,
+            kernel_run_counters=kernel_run_counters,
+            slot_run_counters=slot_run_counters,
+            enable_chain_position_assert=enable_chain_position_assert,
+            req_to_verify_expected_tokens=req_to_verify_expected_tokens,
+        )

python/sglang/srt/model_executor/forward_batch_info.py

@@ -40,6 +40,7 @@ from sglang.srt.distributed.parallel_state import (
     get_moe_expert_parallel_world_size,
     get_tensor_model_parallel_world_size,
 )
+from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
     get_attention_cp_size,
@@ -448,6 +449,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     # For ngram embedding
     ngram_embedding_info: Optional[NgramEmbeddingInfo] = None
 
+    # kv-canary token-id validator snapshot
+    req_all_ids_flat: Optional[torch.Tensor] = None
+    req_all_ids_lens: Optional[torch.Tensor] = None
+
     @classmethod
     def init_new(
         cls,
@@ -1314,3 +1319,5 @@ if is_cuda() or is_hip():
     clamp_position = clamp_position_cuda
 else:
     clamp_position = _clamp_position_native
+
+

python/sglang/srt/server_args.py

@@ -777,6 +777,7 @@ class ServerArgs:
     enable_attn_tp_input_scattered: bool = False
     disable_attn_tp_gather: bool = False
     gc_threshold: Optional[List[int]] = None
+    kv_canary: str = "none"
     # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2
     enable_dsa_prefill_context_parallel: bool = False
     dsa_prefill_cp_mode: str = "round-robin-split"
@@ -6305,6 +6306,18 @@ class ServerArgs:
             action="store_true",
             help="Disable RadixAttention for prefix caching.",
         )
+        parser.add_argument(
+            "--kv-canary",
+            type=str,
+            default=ServerArgs.kv_canary,
+            choices=["none", "log", "raise"],
+            help=(
+                "KV cache canary mode. "
+                "'none' disables the canary (default). "
+                "'log' prints them while the server keeps running (production-safe). "
+                "'raise' fails the server on the first detected mismatch (CI lane)."
+            ),
+        )
         parser.add_argument(
             "--cuda-graph-max-bs",
             type=int,

python/sglang/test/kv_canary/e2e_base.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import io
+import os
+import string
+from typing import ClassVar, Literal, Optional
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kv_canary.mode_config import _MODE_CONFIGS, _ModeConfig
+from sglang.test.kv_canary.utils import build_canary_server_args, post_parallel_generate
+from sglang.test.kv_canary.violation_assert_mixin import CanaryViolationAssertMixin
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+# Long prompt body shared by all canary e2e tests. The repetition count is chosen
+# so the tokenised prompt is comfortably long; token count is roughly 7k after BPE.
+_LONG_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 700).strip()
+_UNIQUE_PROMPT_FIRST_CHARS = string.ascii_letters + string.digits
+
+
+class CapturedServerE2EBase(CanaryViolationAssertMixin, CustomTestCase):
+    process: ClassVar[Optional[object]] = None
+    base_url: ClassVar[str] = DEFAULT_URL_FOR_TEST
+    _stdout_buf: ClassVar[Optional[io.StringIO]] = None
+    _stderr_buf: ClassVar[Optional[io.StringIO]] = None
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        if cls.process is not None:
+            kill_process_tree(cls.process.pid)
+        for buf in (cls._stdout_buf, cls._stderr_buf):
+            if buf is not None:
+                buf.close()
+        cls._stdout_buf = None
+        cls._stderr_buf = None
+
+    def _captured_log_text(
+        self, side: Optional[Literal["prefill", "decode"]] = None
+    ) -> str:
+        stdout_text = (
+            self._stdout_buf.getvalue() if self._stdout_buf is not None else ""
+        )
+        stderr_text = (
+            self._stderr_buf.getvalue() if self._stderr_buf is not None else ""
+        )
+        return stdout_text + stderr_text
+
+    def assert_log_contains(self, substring: str) -> None:
+        log_text = self._captured_log_text()
+        if substring not in log_text:
+            raise AssertionError(
+                f"Expected substring {substring!r} not found in captured log. "
+                f"Log tail:\n{log_text[-2000:]}"
+            )
+
+
+class CanaryE2EBase(CapturedServerE2EBase):
+    model_mode: ClassVar[Literal["mha"]]
+    kv_canary_mode: ClassVar[CanaryMode]
+    extra_env: ClassVar[dict[str, str]] = {}
+    extra_server_args: ClassVar[tuple[str, ...]] = ()
+    use_unique_prompts: ClassVar[bool] = False
+    # Number of sequential request batches each test method sends. Default 1 keeps tests fast.
+    workload_n_batches: ClassVar[int] = 1
+
+    _cfg: ClassVar[Optional[_ModeConfig]] = None
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls._cfg = _MODE_CONFIGS[cls.model_mode]
+        server_env = os.environ.copy()
+        server_env.update(cls.extra_env)
+
+        cls._stdout_buf = io.StringIO()
+        cls._stderr_buf = io.StringIO()
+
+        server_args = build_canary_server_args(
+            kv_canary_mode=cls.kv_canary_mode,
+            mode_cfg=cls._cfg,
+            extra_server_args=(
+                "--max-total-tokens",
+                "65536",
+                "--skip-server-warmup",
+                *cls.extra_server_args,
+            ),
+        )
+        cls.process = popen_launch_server(
+            cls._cfg.model_path,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=server_args,
+            env=server_env,
+            return_stdout_stderr=(cls._stdout_buf, cls._stderr_buf),
+        )
+
+    def make_prompts(self, n: int) -> list[str]:
+        if self.use_unique_prompts:
+            return _make_unique_prompts(n)
+        return [_LONG_PROMPT_BODY] * n
+
+    def send_parallel_requests(
+        self,
+        n: int = 8,
+        *,
+        assert_all_success: bool = True,
+        max_new_tokens: int = 2048,
+        timeout: float = 240.0,
+    ) -> list[dict]:
+        """Fan out n parallel /generate requests; return list of response dicts."""
+        results = post_parallel_generate(
+            url=self.base_url + "/generate",
+            prompts=self.make_prompts(n),
+            max_new_tokens=max_new_tokens,
+            timeout=timeout,
+        )
+        if assert_all_success:
+            for result in results:
+                self.assertEqual(result.get("status_code"), 200, result)
+        return results
+
+
+def _make_unique_prompts(n: int) -> list[str]:
+    if n > len(_UNIQUE_PROMPT_FIRST_CHARS):
+        raise ValueError(
+            f"unique prompt count {n} exceeds supported count "
+            f"{len(_UNIQUE_PROMPT_FIRST_CHARS)}"
+        )
+
+    return [
+        (
+            f"{_UNIQUE_PROMPT_FIRST_CHARS[i]}"
+            f"{hex(i * 0x9E3779B1 & 0xFFFFFFFF)[2:]} "
+            f"{_LONG_PROMPT_BODY}"
+        )
+        for i in range(n)
+    ]

python/sglang/test/kv_canary/fixtures.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import List, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha
+from sglang.srt.kv_canary.pool_patcher.api import register_pool_attacher
+from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+DEFAULT_DEVICE: torch.device = torch.device("cuda")
+
+
+@dataclass
+class FakeMHAPool:
+    layer_num: int
+    k_buffer: List[torch.Tensor]
+    v_buffer: List[torch.Tensor]
+    page_size: int = 1
+
+    def get_contiguous_buf_infos(self):
+        ptrs = [b.data_ptr() for b in self.k_buffer] + [
+            b.data_ptr() for b in self.v_buffer
+        ]
+        lens = [b.nbytes for b in self.k_buffer] + [b.nbytes for b in self.v_buffer]
+        item_lens = [b[0].nbytes * self.page_size for b in self.k_buffer] + [
+            b[0].nbytes * self.page_size for b in self.v_buffer
+        ]
+        return ptrs, lens, item_lens
+
+
+def make_mha_pool(
+    device: torch.device = DEFAULT_DEVICE,
+    *,
+    num_slots: int = 16,
+    dim: int = 8,
+    layer_num: int = 2,
+) -> FakeMHAPool:
+    k_layers = [
+        torch.zeros(num_slots, dim, dtype=torch.float16, device=device)
+        for _ in range(layer_num)
+    ]
+    v_layers = [
+        torch.zeros(num_slots, dim, dtype=torch.float16, device=device)
+        for _ in range(layer_num)
+    ]
+    return FakeMHAPool(layer_num=layer_num, k_buffer=k_layers, v_buffer=v_layers)
+
+
+def make_base_config() -> CanaryConfig:
+    return CanaryConfig(
+        mode=CanaryMode.RAISE,
+        ring_capacity=1024,
+    )
+
+
+def make_req_to_token_pool(
+    device: torch.device = DEFAULT_DEVICE,
+    *,
+    max_reqs: int = 8,
+    max_seq_len: int = 32,
+) -> SimpleNamespace:
+    req_to_token = torch.zeros(max_reqs, max_seq_len, dtype=torch.int32, device=device)
+    return SimpleNamespace(
+        req_to_token=req_to_token, size=max_reqs, max_context_len=max_seq_len
+    )
+
+
+def make_forward_batch(
+    device: torch.device = DEFAULT_DEVICE,
+    *,
+    bs: int = 2,
+    seq_lens_list: tuple[int, ...] = (3, 4),
+    req_pool_indices: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    seq_lens_sum: Optional[int] = None,
+    extend_prefix_lens: Optional[torch.Tensor] = None,
+    extend_prefix_lens_cpu: Optional[list] = None,
+    extend_seq_lens: Optional[torch.Tensor] = None,
+    extend_seq_lens_cpu: Optional[list] = None,
+    is_extend: bool = False,
+    is_target_verify: bool = False,
+    is_draft_extend_v2: bool = False,
+    spec_info: Optional[object] = None,
+    input_ids: Optional[torch.Tensor] = None,
+    positions: Optional[torch.Tensor] = None,
+    out_cache_loc: Optional[torch.Tensor] = None,
+    num_token_non_padded_cpu: Optional[int] = None,
+) -> SimpleNamespace:
+    seq_lens_default = list(seq_lens_list[:bs])
+    if req_pool_indices is None:
+        req_pool_indices = torch.tensor([1, 2][:bs], dtype=torch.int64, device=device)
+    if seq_lens is None:
+        seq_lens = torch.tensor(seq_lens_default, dtype=torch.int32, device=device)
+    if seq_lens_sum is None:
+        seq_lens_sum = int(sum(seq_lens_default))
+    if input_ids is None:
+        input_ids = torch.zeros(bs, dtype=torch.int32, device=device)
+    if positions is None:
+        # Default to decode-canonical: positions = seq_lens - 1 (one-token-per-req decode write
+        # at the post-bump tail). Plan input derives decode prefix_lens from positions
+        # directly, so the default must keep parity.
+        positions = (seq_lens.to(torch.int64) - 1).clamp(min=0).to(torch.int32)
+    if out_cache_loc is None:
+        out_cache_loc = torch.zeros(bs, dtype=torch.int32, device=device)
+
+    if is_extend:
+        forward_mode = ForwardMode.EXTEND
+    elif is_target_verify:
+        forward_mode = ForwardMode.TARGET_VERIFY
+    elif is_draft_extend_v2:
+        forward_mode = ForwardMode.DRAFT_EXTEND_V2
+    else:
+        forward_mode = ForwardMode.DECODE
+    return SimpleNamespace(
+        forward_mode=forward_mode,
+        batch_size=bs,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        seq_lens_sum=seq_lens_sum,
+        extend_prefix_lens=extend_prefix_lens,
+        extend_prefix_lens_cpu=extend_prefix_lens_cpu,
+        extend_seq_lens=extend_seq_lens,
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        spec_info=spec_info,
+        input_ids=input_ids,
+        positions=positions,
+        out_cache_loc=out_cache_loc,
+        num_token_non_padded_cpu=num_token_non_padded_cpu,
+        req_all_ids_flat=None,
+        req_all_ids_lens=None,
+    )
+
+
+def make_buffer_group(
+    *,
+    device: torch.device = DEFAULT_DEVICE,
+    kind: PoolKind = PoolKind.FULL,
+    has_v: bool = True,
+    swa_index_lut: Optional[torch.Tensor] = None,
+    num_slots: int = 4,
+    kv_token_id_vs_position_offset: int = 0,
+) -> CanaryBufferGroup:
+    def _zero() -> torch.Tensor:
+        return torch.zeros(
+            num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device
+        )
+
+    return CanaryBufferGroup(
+        kind=kind,
+        k_head=_zero(),
+        k_tail=_zero(),
+        v_head=_zero() if has_v else None,
+        v_tail=_zero() if has_v else None,
+        swa_index_lut=swa_index_lut,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+
+
+def make_radix_cache(
+    slot_lists: List[List[int]], device: torch.device = DEFAULT_DEVICE
+):
+    cache = RadixCache.__new__(RadixCache)
+    cache.device = device
+    cache.page_size = 1
+    cache.disable = False
+
+    root = TreeNode()
+    root.value = torch.tensor(
+        slot_lists[0] if slot_lists else [], dtype=torch.int32, device=device
+    )
+    cache.root_node = root
+
+    current = root
+    for child_slots in slot_lists[1:]:
+        child = TreeNode()
+        child.value = torch.tensor(child_slots, dtype=torch.int32, device=device)
+        child.parent = current
+        current.children[child.id] = child
+        current = child
+
+    return cache
+
+
+register_pool_attacher(FakeMHAPool, attach_mha)

python/sglang/test/kv_canary/mode_config.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class _ModeConfig:
+    model_path: str
+    json_model_override_args: Optional[str] = None
+
+
+_MODE_CONFIGS: dict[str, _ModeConfig] = {
+    "mha": _ModeConfig(
+        model_path="Qwen/Qwen3-0.6B",
+    ),
+}

python/sglang/test/kv_canary/runner_test_base.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.kv_canary import endpoint as endpoint_module
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.runner import kernel_launcher as kernel_launcher_module
+from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
+from sglang.test.kv_canary.fixtures import (
+    DEFAULT_DEVICE,
+    make_buffer_group,
+    make_req_to_token_pool,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+def make_config(
+    *,
+    mode: CanaryMode = CanaryMode.RAISE,
+    ring_capacity: int = 1024,
+) -> CanaryConfig:
+    return CanaryConfig(
+        mode=mode,
+        ring_capacity=ring_capacity,
+    )
+
+
+class RecordingEndpoint:
+    def __init__(self, *, kernel_kind: CanaryLaunchTag) -> None:
+        self.kernel_kind = kernel_kind
+        self.calls: list[dict[str, object]] = []
+
+    def launch_per_forward(self, **kwargs: object) -> None:
+        self.calls.append(kwargs)
+
+
+def make_manager(
+    *,
+    device: torch.device,
+    config: CanaryConfig | None = None,
+    group: CanaryBufferGroup | None = None,
+    req_pool: SimpleNamespace | None = None,
+    per_forward_verify_capacity: int = 16,
+) -> CanaryManager:
+    if config is None:
+        config = make_config()
+    if group is None:
+        group = make_buffer_group(device=device)
+    if req_pool is None:
+        req_pool = make_req_to_token_pool(device=device, max_reqs=4, max_seq_len=8)
+    return CanaryManager(
+        config=config,
+        buffer_groups=(group,),
+        device=device,
+        req_to_token_pool=req_pool,
+        launch_capacities=CanaryLaunchCapacities(
+            per_forward_verify_capacity=per_forward_verify_capacity,
+            per_forward_write_req_capacity=2,
+            per_forward_write_entry_capacity=8,
+        ),
+    )
+
+
+class CanaryManagerTestCase(CustomTestCase):
+    def setUp(self) -> None:
+        self.device = DEFAULT_DEVICE
+        # Stub plan/verify/write kernels so CPU runs don't need CUDA JIT.
+        self._patchers = [
+            patch.object(
+                kernel_launcher_module,
+                "launch_canary_plan_kernels",
+                lambda **kwargs: None,
+            ),
+            patch.object(
+                endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None
+            ),
+            patch.object(
+                endpoint_module, "launch_canary_write_kernel", lambda **kwargs: None
+            ),
+        ]
+        for patcher in self._patchers:
+            patcher.start()
+            self.addCleanup(patcher.stop)

python/sglang/test/kv_canary/utils.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+
+import requests
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.test.kv_canary.mode_config import _ModeConfig
+
+
+def build_canary_server_args(
+    *,
+    kv_canary_mode: CanaryMode,
+    mode_cfg: _ModeConfig,
+    extra_server_args: tuple[str, ...] = (),
+) -> list[str]:
+    args = [
+        "--kv-canary",
+        kv_canary_mode.value,
+        "--disable-piecewise-cuda-graph",
+        "--context-length",
+        "16384",
+        *extra_server_args,
+    ]
+    if mode_cfg.json_model_override_args is not None:
+        args.extend(["--json-model-override-args", mode_cfg.json_model_override_args])
+    return args
+
+
+def post_parallel_generate(
+    *,
+    url: str,
+    prompts: list[str],
+    max_new_tokens: int,
+    timeout: float,
+) -> list[dict]:
+    def _send(prompt: str) -> dict:
+        try:
+            resp = requests.post(
+                url,
+                json={
+                    "text": prompt,
+                    "sampling_params": {
+                        "max_new_tokens": max_new_tokens,
+                        "temperature": 0.0,
+                    },
+                },
+                timeout=timeout,
+            )
+            return {"status_code": resp.status_code, "body": resp.text}
+        except requests.RequestException as exc:
+            return {"status_code": -1, "error": repr(exc)}
+
+    with ThreadPoolExecutor(max_workers=max(1, len(prompts))) as pool:
+        return list(pool.map(_send, prompts))

python/sglang/test/kv_canary/violation_assert_mixin.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import time
+from typing import Literal, Optional
+
+from sglang.test.kv_canary.violation_log_utils import (
+    assert_no_violation_in_log,
+    find_violation_in_log,
+)
+
+_Side = Optional[Literal["prefill", "decode"]]
+
+
+class CanaryViolationAssertMixin:
+    def _captured_log_text(self, side: _Side = None) -> str:
+        raise NotImplementedError
+
+    def assert_per_forward_violation_reported(
+        self,
+        *,
+        fail_reason: str,
+        side: _Side = None,
+        flush_wait_seconds: float = 2.0,
+    ) -> None:
+        self.assert_violation_logged_any(
+            launch_tag_patterns=("HEAD_*", "TAIL_*"),
+            fail_reason=fail_reason,
+            side=side,
+            flush_wait_seconds=flush_wait_seconds,
+        )
+
+    def assert_any_launch_tag_violation_reported(
+        self,
+        *,
+        fail_reason: str,
+        side: _Side = None,
+        flush_wait_seconds: float = 3.0,
+        max_retries: int = 10,
+    ) -> None:
+        self.assert_violation_logged_any(
+            launch_tag_patterns=("*",),
+            fail_reason=fail_reason,
+            side=side,
+            flush_wait_seconds=flush_wait_seconds,
+            max_retries=max_retries,
+        )
+
+    def assert_any_launch_tag_violation_absent(
+        self, *, fail_reason: str, side: _Side = None
+    ) -> None:
+        self.assert_no_violation_matching(
+            launch_tag_patterns=("*",), fail_reason=fail_reason, side=side
+        )
+
+    def assert_violation_logged_any(
+        self,
+        *,
+        launch_tag_patterns: tuple[str, ...],
+        fail_reason: str,
+        side: _Side = None,
+        flush_wait_seconds: float = 2.0,
+        max_retries: int = 1,
+    ) -> None:
+        log_text = ""
+        for _ in range(max_retries):
+            time.sleep(flush_wait_seconds)
+            log_text = self._captured_log_text(side)
+            if find_violation_in_log(
+                log_text,
+                launch_tag_patterns=launch_tag_patterns,
+                fail_reason=fail_reason,
+            ):
+                return
+        side_label = "" if side is None else f" on side={side}"
+        other_side_diag = ""
+        if side in ("prefill", "decode"):
+            other = "decode" if side == "prefill" else "prefill"
+            try:
+                other_text = self._captured_log_text(other)
+                other_match = find_violation_in_log(
+                    other_text,
+                    launch_tag_patterns=launch_tag_patterns,
+                    fail_reason=fail_reason,
+                )
+                other_side_diag = (
+                    f"\n[diag] other side ({other}) buf len={len(other_text)} "
+                    f"contains_match={other_match}"
+                )
+            except (NotImplementedError, ValueError):
+                pass
+        raise AssertionError(
+            f"No canary violation matching launch_tag_patterns={launch_tag_patterns!r} "
+            f"fail_reason={fail_reason!r}{side_label} after max_retries={max_retries} "
+            f"(wait={flush_wait_seconds}s each). "
+            f"log_text len={len(log_text)}.{other_side_diag} Log tail:\n"
+            f"{log_text[-2000:]}"
+        )
+
+    def assert_no_violation_matching(
+        self,
+        *,
+        launch_tag_patterns: tuple[str, ...],
+        fail_reason: str,
+        side: _Side = None,
+    ) -> None:
+        log_text = self._captured_log_text(side)
+        if find_violation_in_log(
+            log_text,
+            launch_tag_patterns=launch_tag_patterns,
+            fail_reason=fail_reason,
+        ):
+            raise AssertionError(
+                f"Unexpected canary violation matching "
+                f"launch_tag_patterns={launch_tag_patterns!r} "
+                f"fail_reason={fail_reason!r}. Log tail:\n{log_text[-2000:]}"
+            )
+
+    def assert_no_violation(
+        self,
+        *,
+        side: _Side = None,
+        wait_seconds: float = 2.0,
+    ) -> None:
+        time.sleep(wait_seconds)
+        assert_no_violation_in_log(self._captured_log_text(side))

python/sglang/test/kv_canary/violation_log_utils.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import fnmatch
+import re
+
+_VIOLATION_LINE_RE = re.compile(
+    r"kv_canary violation: launch_tag=(\S+) fail_reason=(\S+)"
+)
+
+
+def find_violation_in_log(
+    log_text: str,
+    *,
+    launch_tag_patterns: tuple[str, ...],
+    fail_reason: str,
+) -> bool:
+    for match in _VIOLATION_LINE_RE.finditer(log_text):
+        tag = match.group(1)
+        reason_field = match.group(2)
+        if fail_reason not in reason_field.split("+"):
+            continue
+        if any(fnmatch.fnmatchcase(tag, pattern) for pattern in launch_tag_patterns):
+            return True
+    return False
+
+
+def assert_no_violation_in_log(log_text: str) -> None:
+    if "kv_canary violation:" in log_text:
+        raise AssertionError(
+            f"Unexpected canary violation found. Log tail:\n{log_text[-2000:]}"
+        )

test/registered/kv_canary/test_e2e_base.py

@@ -0,0 +1,31 @@
+import unittest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kv_canary.e2e_base import (
+    _LONG_PROMPT_BODY,
+    _UNIQUE_PROMPT_FIRST_CHARS,
+    _make_unique_prompts,
+)
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-b-test-cpu")
+
+
+class TestCanaryE2EBase(CustomTestCase):
+    def test_make_unique_prompts_have_distinct_first_characters(self) -> None:
+        """Verify generated prompts use distinct first characters and all end with the shared long body."""
+        prompts = _make_unique_prompts(8)
+
+        self.assertEqual(len({prompt[0] for prompt in prompts}), len(prompts))
+        self.assertTrue(all(prompt.endswith(_LONG_PROMPT_BODY) for prompt in prompts))
+
+    def test_make_unique_prompts_rejects_more_prompts_than_distinct_first_characters(
+        self,
+    ) -> None:
+        """Verify prompt generation rejects requests beyond the unique prefix budget."""
+        with self.assertRaisesRegex(ValueError, "unique prompt count"):
+            _make_unique_prompts(len(_UNIQUE_PROMPT_FIRST_CHARS) + 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_buffer_alloc.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
+from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
+
+
+class TestAllocCanaryBuf(CustomTestCase):
+    def test_alloc_canary_buf_shape_and_dtype(self) -> None:
+        """Verify alloc_canary_buf returns a zeroed uint8 buffer of [num_slots, CANARY_SLOT_BYTES]."""
+        buf = alloc_canary_buf(num_slots=8, device=DEFAULT_DEVICE)
+        self.assertEqual(buf.shape, (8, CANARY_SLOT_BYTES))
+        self.assertEqual(buf.dtype, torch.uint8)
+        self.assertEqual(buf.device.type, DEFAULT_DEVICE.type)
+        self.assertTrue(torch.equal(buf, torch.zeros_like(buf)))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_capacities.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
+
+
+class TestComputeLaunchCapacities(CustomTestCase):
+    @staticmethod
+    def _make_server_args(*, max_bs: int) -> SimpleNamespace:
+        return SimpleNamespace(
+            cuda_graph_max_bs=max_bs,
+            speculative_num_draft_tokens=0,
+            chunked_prefill_size=None,
+            max_prefill_tokens=128,
+        )
+
+    @staticmethod
+    def _from_args(
+        *,
+        max_bs: int,
+        max_seq_len: int,
+        max_total_num_tokens: int | None = None,
+    ) -> CanaryLaunchCapacities:
+        if max_total_num_tokens is None:
+            max_total_num_tokens = max_bs * max_seq_len
+        return CanaryLaunchCapacities.from_args(
+            server_args=TestComputeLaunchCapacities._make_server_args(max_bs=max_bs),
+            req_to_token_pool_size=max_bs,
+            max_seq_len_per_req=max_seq_len,
+            pool_slot_count=max_total_num_tokens,
+        )
+
+    def test_per_forward_verify_capacity_covers_multi_req_prefix_sum(self) -> None:
+        """Verify per-forward verify capacity equals max_total_num_tokens * 3."""
+        max_bs = 8
+        max_seq_len = 64
+        max_total_num_tokens = 1024
+        capacities = self._from_args(
+            max_bs=max_bs,
+            max_seq_len=max_seq_len,
+            max_total_num_tokens=max_total_num_tokens,
+        )
+        self.assertEqual(
+            capacities.per_forward_verify_capacity,
+            int(max_total_num_tokens * 3),
+        )
+
+    def test_from_args_treats_missing_speculative_draft_tokens_as_zero(self) -> None:
+        """per_forward_write_entry_capacity is floored by max_prefill_tokens when batch * tokens_per_bs is smaller."""
+        server_args = self._make_server_args(max_bs=2)
+        server_args.speculative_num_draft_tokens = None
+
+        capacities = CanaryLaunchCapacities.from_args(
+            server_args=server_args,
+            req_to_token_pool_size=2,
+            max_seq_len_per_req=32,
+            pool_slot_count=64,
+        )
+
+        self.assertEqual(capacities.per_forward_write_entry_capacity, 128)
+
+    def test_manual_capacities_reject_non_positive_fields(self) -> None:
+        """Verify manual launch capacities fail instead of being clamped."""
+        with self.assertRaisesRegex(ValueError, "per_forward_verify_capacity"):
+            CanaryLaunchCapacities(
+                per_forward_verify_capacity=0,
+                per_forward_write_req_capacity=1,
+                per_forward_write_entry_capacity=1,
+            )
+
+    def test_from_args_rejects_empty_pool_capacity(self) -> None:
+        """Verify derived launch capacities reject invalid pool sizing."""
+        with self.assertRaisesRegex(ValueError, "pool_slot_count"):
+            CanaryLaunchCapacities.from_args(
+                server_args=self._make_server_args(max_bs=1),
+                req_to_token_pool_size=1,
+                max_seq_len_per_req=1,
+                pool_slot_count=0,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_endpoint.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import (
+    CANARY_SLOT_BYTES,
+    CanaryLaunchTag,
+    VerifyPlan,
+)
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.srt.kv_canary import endpoint as endpoint_module
+from sglang.srt.kv_canary.endpoint import (
+    CanaryEndpoint,
+)
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.state import (
+    ViolationLog,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="extra-a", runner_config="1-gpu-small")
+
+
+def _make_endpoint(*, device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL, swa_lut=None):
+    canary_buf = torch.zeros(4, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device)
+    slot_view = torch.zeros(1, dtype=torch.int64, device=device)
+    kernel_view = torch.zeros(1, dtype=torch.int64, device=device)
+    enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device)
+    return CanaryEndpoint(
+        kernel_kind=kernel_kind,
+        canary_buf=canary_buf,
+        full_to_swa_index_mapping=swa_lut,
+        slot_run_counter_view=slot_view,
+        kernel_run_counter_view=kernel_view,
+        enable_chain_position_assert=enable_chain_position_assert,
+    )
+
+
+def _make_kernel_args(device):
+    verify_plan = VerifyPlan.allocate(verify_capacity=1, device=device)
+    write_plan = WritePlan.allocate(write_req_capacity=1, device=device)
+    log = ViolationLog.allocate(ring_capacity=2, device=device)
+    return SimpleNamespace(
+        verify_plan=verify_plan,
+        write_plan=write_plan,
+        violation_log=log,
+        input_ids=torch.zeros(1, dtype=torch.int64, device=device),
+        positions=torch.zeros(1, dtype=torch.int64, device=device),
+        out_cache_loc=torch.zeros(1, dtype=torch.int64, device=device),
+        enable_write_input_assert=False,
+        enable_verify_token_assert=False,
+        expected_inputs=ExpectedInputs.allocate(capacity=1, device=device),
+    )
+
+
+class TestSelfUnitEndpoint(CustomTestCase):
+    def setUp(self):
+        self.device = DEFAULT_DEVICE
+
+    def test_launch_per_forward_passes_kernel_kind(self):
+        """Verify per-forward launch passes the endpoint kernel kind."""
+        captured: list[tuple[str, CanaryLaunchTag]] = []
+        with patch.object(
+            endpoint_module,
+            "launch_canary_verify_kernel",
+            lambda **kwargs: captured.append(("verify", kwargs["context"].kernel_kind)),
+        ), patch.object(
+            endpoint_module,
+            "launch_canary_write_kernel",
+            lambda **kwargs: captured.append(("write", kwargs["context"].kernel_kind)),
+        ):
+            ep = _make_endpoint(
+                device=self.device, kernel_kind=CanaryLaunchTag.TAIL_V_SWA
+            )
+            args = _make_kernel_args(self.device)
+            ep.launch_per_forward(
+                verify_plan=args.verify_plan,
+                write_plan=args.write_plan,
+                input_ids=args.input_ids,
+                positions=args.positions,
+                out_cache_loc=args.out_cache_loc,
+                enable_write_input_assert=args.enable_write_input_assert,
+                enable_verify_token_assert=args.enable_verify_token_assert,
+                expected_inputs=args.expected_inputs,
+                violation_log=args.violation_log,
+            )
+        self.assertIn(("verify", CanaryLaunchTag.TAIL_V_SWA), captured)
+        self.assertIn(("write", CanaryLaunchTag.TAIL_V_SWA), captured)
+
+    def test_endpoint_shares_violation_log_across_launches(self):
+        """Verify endpoints can reuse the same violation log."""
+        captured_rings: list[int] = []
+        with patch.object(
+            endpoint_module,
+            "launch_canary_verify_kernel",
+            lambda **kwargs: captured_rings.append(
+                kwargs["context"].violation_ring.data_ptr()
+            ),
+        ), patch.object(
+            endpoint_module,
+            "launch_canary_write_kernel",
+            lambda **kwargs: None,
+        ):
+            shared_log = ViolationLog.allocate(ring_capacity=2, device=self.device)
+            ep_a = _make_endpoint(
+                device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL
+            )
+            ep_b = _make_endpoint(
+                device=self.device, kernel_kind=CanaryLaunchTag.HEAD_V_FULL
+            )
+
+            args = _make_kernel_args(self.device)
+            ep_a.launch_per_forward(
+                verify_plan=args.verify_plan,
+                write_plan=args.write_plan,
+                input_ids=args.input_ids,
+                positions=args.positions,
+                out_cache_loc=args.out_cache_loc,
+                enable_write_input_assert=args.enable_write_input_assert,
+                enable_verify_token_assert=args.enable_verify_token_assert,
+                expected_inputs=args.expected_inputs,
+                violation_log=shared_log,
+            )
+            ep_b.launch_per_forward(
+                verify_plan=args.verify_plan,
+                write_plan=args.write_plan,
+                input_ids=args.input_ids,
+                positions=args.positions,
+                out_cache_loc=args.out_cache_loc,
+                enable_write_input_assert=args.enable_write_input_assert,
+                enable_verify_token_assert=args.enable_verify_token_assert,
+                expected_inputs=args.expected_inputs,
+                violation_log=shared_log,
+            )
+        self.assertEqual(captured_rings[0], captured_rings[1])
+        self.assertEqual(captured_rings[0], shared_log.violation_ring.data_ptr())
+
+    def test_swa_endpoint_pre_translates_out_cache_loc(self):
+        """Verify SWA endpoints translate cache locations before write launch."""
+        captured: list[torch.Tensor] = []
+        with patch.object(
+            endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None
+        ), patch.object(
+            endpoint_module,
+            "launch_canary_write_kernel",
+            lambda **kwargs: captured.append(kwargs["out_cache_loc"]),
+        ):
+            # LUT maps full slot i → swa slot (i + 100) so we can verify the gather happened.
+            lut = torch.arange(8, dtype=torch.int64, device=self.device) + 100
+            swa_ep = _make_endpoint(
+                device=self.device,
+                kernel_kind=CanaryLaunchTag.HEAD_K_SWA,
+                swa_lut=lut,
+            )
+            full_ep = _make_endpoint(
+                device=self.device,
+                kernel_kind=CanaryLaunchTag.HEAD_K_FULL,
+                swa_lut=None,
+            )
+            args = _make_kernel_args(self.device)
+
+            swa_ep.launch_per_forward(
+                verify_plan=args.verify_plan,
+                write_plan=args.write_plan,
+                input_ids=args.input_ids,
+                positions=args.positions,
+                out_cache_loc=args.out_cache_loc,
+                enable_write_input_assert=args.enable_write_input_assert,
+                enable_verify_token_assert=args.enable_verify_token_assert,
+                expected_inputs=args.expected_inputs,
+                violation_log=args.violation_log,
+            )
+            full_ep.launch_per_forward(
+                verify_plan=args.verify_plan,
+                write_plan=args.write_plan,
+                input_ids=args.input_ids,
+                positions=args.positions,
+                out_cache_loc=args.out_cache_loc,
+                enable_write_input_assert=args.enable_write_input_assert,
+                enable_verify_token_assert=args.enable_verify_token_assert,
+                expected_inputs=args.expected_inputs,
+                violation_log=args.violation_log,
+            )
+        # SWA call: out_cache_loc was rewritten via lut gather (so identity-shifted by +100 here).
+        expected_swa = lut[args.out_cache_loc]
+        self.assertTrue(torch.equal(captured[0], expected_swa))
+        # FULL call: out_cache_loc keeps the same values and dtype.
+        self.assertIs(captured[1], args.out_cache_loc)
+
+    def test_swa_endpoint_trailing_sentinel_row_yields_skip(self):
+        """Verify SWA sentinel cache rows become write-skip markers."""
+        captured: list[torch.Tensor] = []
+        with patch.object(
+            endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None
+        ), patch.object(
+            endpoint_module,
+            "launch_canary_write_kernel",
+            lambda **kwargs: captured.append(kwargs["out_cache_loc"]),
+        ):
+            # 8 in-window rows + 1 trailing sentinel row at index 8.
+            lut = torch.arange(8, dtype=torch.int64, device=self.device)
+            lut = torch.cat(
+                [lut, torch.tensor([-1], dtype=torch.int64, device=self.device)]
+            )
+            swa_ep = _make_endpoint(
+                device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_SWA, swa_lut=lut
+            )
+            args = _make_kernel_args(self.device)
+            # Point out_cache_loc at the trailing-sentinel-row index — this is how sglang signals
+            # "this token is out-of-window for the SWA group" pre-cleanup, and the new host gather must
+            # produce -1 here.
+            args.out_cache_loc.fill_(8)
+
+            swa_ep.launch_per_forward(
+                verify_plan=args.verify_plan,
+                write_plan=args.write_plan,
+                input_ids=args.input_ids,
+                positions=args.positions,
+                out_cache_loc=args.out_cache_loc,
+                enable_write_input_assert=args.enable_write_input_assert,
+                enable_verify_token_assert=args.enable_verify_token_assert,
+                expected_inputs=args.expected_inputs,
+                violation_log=args.violation_log,
+            )
+
+        self.assertTrue(
+            torch.equal(
+                captured[0],
+                torch.tensor([-1], dtype=torch.int64, device=self.device),
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_future_tensor.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import unittest
+from typing import cast
+
+import torch
+
+from sglang.srt.kv_canary.runner.future_tensor import FutureTensors
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="extra-a", runner_config="1-gpu-small")
+
+
+class _FakeEvent:
+    def __init__(self) -> None:
+        self.synchronize_count = 0
+
+    def synchronize(self) -> None:
+        self.synchronize_count += 1
+
+
+class TestFutureTensors(CustomTestCase):
+    def test_cuda_stage_then_wait_returns_host_copy(self) -> None:
+        """Verify staged CUDA tensors are copied back on wait."""
+        device = torch.device("cuda")
+        alt_stream = torch.cuda.Stream(device=device)
+        default_stream = torch.cuda.current_stream(device)
+        self.assertNotEqual(alt_stream.stream_id, default_stream.stream_id)
+
+        src_first = torch.tensor([41], dtype=torch.int32, device=device)
+        future_first = FutureTensors.device_to_host(
+            xs_device=src_first, d2h_stream=alt_stream
+        )
+        result_first = future_first.wait()
+        self.assertEqual(int(result_first.item()), 41)
+
+        src_second = torch.tensor([97], dtype=torch.int32, device=device)
+        future_second = FutureTensors.device_to_host(
+            xs_device=src_second, d2h_stream=alt_stream
+        )
+        result_second = future_second.wait()
+        self.assertEqual(int(result_second.item()), 97)
+
+    def test_cuda_pinned_when_stream_is_provided(self) -> None:
+        """Verify CUDA staging uses pinned host memory with a stream."""
+        device = torch.device("cuda")
+        alt_stream = torch.cuda.Stream(device=device)
+        src = torch.tensor([5], dtype=torch.int32, device=device)
+        future = FutureTensors.device_to_host(xs_device=src, d2h_stream=alt_stream)
+        staged_tensors = [
+            v for v in future._data.values() if isinstance(v, torch.Tensor)
+        ]
+        self.assertTrue(staged_tensors)
+        self.assertTrue(all(t.is_pinned() for t in staged_tensors))
+        self.assertEqual(int(future.wait().item()), 5)
+
+    def test_cuda_each_call_allocates_fresh_host(self) -> None:
+        """Verify each CUDA staging call owns a fresh host buffer."""
+        device = torch.device("cuda")
+        alt_stream = torch.cuda.Stream(device=device)
+        src_a = torch.tensor([13], dtype=torch.int32, device=device)
+        src_b = torch.tensor([29], dtype=torch.int32, device=device)
+        future_a = FutureTensors.device_to_host(xs_device=src_a, d2h_stream=alt_stream)
+        future_b = FutureTensors.device_to_host(xs_device=src_b, d2h_stream=alt_stream)
+        ptrs_a = {
+            v.data_ptr() for v in future_a._data.values() if isinstance(v, torch.Tensor)
+        }
+        ptrs_b = {
+            v.data_ptr() for v in future_b._data.values() if isinstance(v, torch.Tensor)
+        }
+        self.assertTrue(ptrs_a and ptrs_b)
+        self.assertFalse(ptrs_a & ptrs_b)
+        self.assertEqual(int(future_a.wait().item()), 13)
+        self.assertEqual(int(future_b.wait().item()), 29)
+
+    def test_dict_of_all_tensors_roundtrip(self) -> None:
+        """Verify a dict of multiple tensors round-trips entry-by-entry."""
+        device = torch.device("cuda")
+        stream = torch.cuda.Stream(device=device)
+        src = {
+            "x": torch.tensor([11, 22], dtype=torch.int64, device=device),
+            "y": torch.tensor([99], dtype=torch.int32, device=device),
+        }
+        future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
+        out = future.wait()
+        self.assertIsInstance(out, dict)
+        self.assertEqual(out["x"].tolist(), [11, 22])
+        self.assertEqual(int(out["y"].item()), 99)
+        self.assertTrue(out["x"].is_pinned())
+        self.assertTrue(out["y"].is_pinned())
+
+    def test_dict_mixes_tensor_and_passthrough(self) -> None:
+        """Verify non-tensor dict entries ride through verbatim alongside staging."""
+        device = torch.device("cuda")
+        stream = torch.cuda.Stream(device=device)
+        sentinel_obj = {"nested": [1, 2, 3]}
+        src = {
+            "step": 42,
+            "label": "decode",
+            "extra": sentinel_obj,
+            "counter": torch.tensor([7], dtype=torch.int32, device=device),
+        }
+        future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
+        out = future.wait()
+        self.assertEqual(out["step"], 42)
+        self.assertEqual(out["label"], "decode")
+        # Identity (not deep-copy) — callers can rely on shared mutable references.
+        self.assertIs(out["extra"], sentinel_obj)
+        self.assertEqual(int(out["counter"].item()), 7)
+        self.assertTrue(out["counter"].is_pinned())
+
+    def test_dict_passthrough_preserves_tensor_value(self) -> None:
+        """Verify tensors share device memory but non-tensor types are not staged."""
+        device = torch.device("cuda")
+        stream = torch.cuda.Stream(device=device)
+        src_tensor = torch.tensor([3], dtype=torch.int32, device=device)
+        src = {"step": 100, "buf": src_tensor}
+        future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
+        out = future.wait()
+        # Tensor is staged to a fresh pinned-host buffer (different storage from src).
+        self.assertNotEqual(out["buf"].data_ptr(), src_tensor.data_ptr())
+        self.assertTrue(out["buf"].is_pinned())
+        # Non-tensor passes through with no copy.
+        self.assertEqual(out["step"], 100)
+        self.assertIsInstance(out["step"], int)
+
+    def test_dict_without_tensor_raises(self) -> None:
+        """Verify a tensor-less dict raises (no device to anchor the d2h sync)."""
+        device = torch.device("cuda")
+        stream = torch.cuda.Stream(device=device)
+        with self.assertRaises(ValueError):
+            FutureTensors.device_to_host(
+                xs_device={"step": 0, "label": "decode"}, d2h_stream=stream
+            )
+
+    def test_wait_called_twice_raises(self) -> None:
+        """Verify wait() after the first drain raises (state cleared)."""
+        device = torch.device("cuda")
+        stream = torch.cuda.Stream(device=device)
+        src = torch.tensor([3], dtype=torch.int32, device=device)
+        future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
+        self.assertEqual(int(future.wait().item()), 3)
+        with self.assertRaises(RuntimeError):
+            future.wait()
+
+    def test_wait_clears_fields_and_rejects_second_wait(self) -> None:
+        """Verify wait() syncs the event exactly once and clears internal state."""
+        tensor = torch.tensor([1, 2, 3])
+        event = _FakeEvent()
+        future = FutureTensors(
+            _data={"x": tensor}, _event=cast(torch.cuda.Event, event)
+        )
+
+        result = future.wait()
+        self.assertIs(result["x"], tensor)
+        self.assertEqual(event.synchronize_count, 1)
+        self.assertIsNone(future._data)
+        self.assertIsNone(future._event)
+
+        with self.assertRaisesRegex(RuntimeError, "called more than once"):
+            future.wait()
+        # Failed wait must not re-trigger event.synchronize.
+        self.assertEqual(event.synchronize_count, 1)
+
+    def test_dict_anchor_picked_from_first_tensor(self) -> None:
+        """Verify staging works when the first key is a non-tensor (anchor must scan)."""
+        device = torch.device("cuda")
+        stream = torch.cuda.Stream(device=device)
+        src = {
+            "step": 5,
+            "buf": torch.tensor([17], dtype=torch.int32, device=device),
+        }
+        out = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream).wait()
+        self.assertEqual(out["step"], 5)
+        self.assertEqual(int(out["buf"].item()), 17)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_plan_input.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.kv_canary.plan_input import PlanInput
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.fixtures import (
+    DEFAULT_DEVICE,
+    make_forward_batch,
+)
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small")
+
+
+def _make_static_plan_input(*, bs_capacity: int, device) -> PlanInput:
+    return PlanInput(
+        req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
+        prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
+        extend_seq_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
+        req_to_verify_expected_tokens_valid_lens=torch.zeros(
+            bs_capacity, dtype=torch.int64, device=device
+        ),
+    )
+
+
+class TestSelfUnitPlanInput(CustomTestCase):
+    def setUp(self):
+        self.device = DEFAULT_DEVICE
+
+    def test_plan_input_fill_from_forward_batch_extend(self):
+        """Verify extend batches populate per-forward plan inputs."""
+        fb = make_forward_batch(
+            self.device,
+            req_pool_indices=torch.tensor(
+                [1, 2], dtype=torch.int64, device=self.device
+            ),
+            seq_lens=torch.tensor([10, 12], dtype=torch.int32, device=self.device),
+            extend_prefix_lens=torch.tensor(
+                [3, 5], dtype=torch.int32, device=self.device
+            ),
+            extend_seq_lens=torch.tensor([7, 7], dtype=torch.int32, device=self.device),
+            is_extend=True,
+        )
+        plan = _make_static_plan_input(bs_capacity=4, device=self.device)
+        plan.fill_from_forward_batch(forward_batch=fb)
+        self.assertEqual(plan.req_pool_indices[:2].tolist(), [1, 2])
+        self.assertEqual(plan.req_pool_indices[2:].tolist(), [0, 0])
+        self.assertEqual(plan.prefix_lens[:2].tolist(), [3, 5])
+        self.assertEqual(plan.extend_seq_lens[:2].tolist(), [7, 7])
+        self.assertEqual(plan.prefix_lens.dtype, torch.int64)
+        self.assertEqual(plan.extend_seq_lens.dtype, torch.int64)
+
+    def test_plan_input_fill_from_forward_batch_target_verify(self):
+        """Verify target-verify batches derive draft verification spans."""
+        # TARGET_VERIFY writes spec_info.draft_token_num positions per req starting at the
+        # current (un-bumped) seq_lens; extend_prefix_lens / extend_seq_lens are deliberately
+        # NOT supplied because init_new does not populate them for this mode.
+        spec_info = SimpleNamespace(draft_token_num=4)
+        fb = make_forward_batch(
+            self.device,
+            req_pool_indices=torch.tensor(
+                [6, 7], dtype=torch.int64, device=self.device
+            ),
+            seq_lens=torch.tensor([10, 14], dtype=torch.int32, device=self.device),
+            is_target_verify=True,
+            spec_info=spec_info,
+        )
+        plan = _make_static_plan_input(bs_capacity=4, device=self.device)
+        plan.fill_from_forward_batch(forward_batch=fb)
+        self.assertEqual(plan.prefix_lens[:2].tolist(), [10, 14])
+        self.assertEqual(plan.extend_seq_lens[:2].tolist(), [4, 4])
+
+    def test_plan_input_fill_from_forward_batch_draft_extend_v2(self):
+        """Verify draft-extend-v2 batches derive prefix lengths from sequence lengths."""
+        # DRAFT_EXTEND_V2 has seq_lens already bumped by the per-req draft refill length;
+        # extend_prefix_lens is intentionally absent (cuda-graph replay does not set it). The
+        # builder must derive prefix as seq_lens - extend_seq_lens.
+        fb = make_forward_batch(
+            self.device,
+            req_pool_indices=torch.tensor(
+                [4, 5], dtype=torch.int64, device=self.device
+            ),
+            seq_lens=torch.tensor([14, 18], dtype=torch.int32, device=self.device),
+            extend_seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=self.device),
+            is_draft_extend_v2=True,
+        )
+        plan = _make_static_plan_input(bs_capacity=4, device=self.device)
+        plan.fill_from_forward_batch(forward_batch=fb)
+        self.assertEqual(plan.prefix_lens[:2].tolist(), [10, 14])
+        self.assertEqual(plan.extend_seq_lens[:2].tolist(), [4, 4])
+
+    def test_plan_input_fill_from_forward_batch_decode(self):
+        """Verify decode batches populate one-token verification spans."""
+        fb = make_forward_batch(
+            self.device,
+            req_pool_indices=torch.tensor(
+                [1, 2, 3], dtype=torch.int64, device=self.device
+            ),
+            seq_lens=torch.tensor([4, 7, 1], dtype=torch.int32, device=self.device),
+            is_extend=False,
+        )
+        plan = _make_static_plan_input(bs_capacity=4, device=self.device)
+        plan.fill_from_forward_batch(forward_batch=fb)
+        self.assertEqual(plan.prefix_lens[:3].tolist(), [3, 6, 0])
+        self.assertEqual(plan.extend_seq_lens[:3].tolist(), [1, 1, 1])
+
+    def test_plan_input_padding_dummy_sentinel(self):
+        """Verify padding sentinel rows remain valid plan input entries."""
+        fb = make_forward_batch(
+            self.device,
+            req_pool_indices=torch.tensor(
+                [0, 5, 0], dtype=torch.int64, device=self.device
+            ),
+            seq_lens=torch.tensor([0, 3, 0], dtype=torch.int32, device=self.device),
+            is_extend=False,
+        )
+        plan = _make_static_plan_input(bs_capacity=4, device=self.device)
+        plan.fill_from_forward_batch(forward_batch=fb)
+        self.assertEqual(plan.req_pool_indices[:3].tolist(), [0, 5, 0])
+        self.assertEqual(plan.req_pool_indices.dtype, torch.int64)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_pool_patcher.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
+from sglang.srt.kv_canary.buffer_group import PoolKind
+from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.fixtures import (
+    DEFAULT_DEVICE,
+    make_base_config,
+    make_mha_pool,
+)
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
+
+
+class PoolPatcherHelper:
+    def setUp(self):
+        self.device = DEFAULT_DEVICE
+        self.config = make_base_config()
+
+
+class TestAttachCanaryBuffers(PoolPatcherHelper, CustomTestCase):
+    def test_canary_buffer_group_allocate_full_only(self):
+        """Verify MHA pools allocate only full canary buffers."""
+        pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)
+        groups_tuple = attach_canary_buffers(
+            pool=pool,
+            config=self.config,
+            device=self.device,
+            kv_token_id_vs_position_offset=0,
+        )
+        groups = {g.kind: g for g in groups_tuple}
+        self.assertEqual(set(groups.keys()), {PoolKind.FULL})
+        group = groups[PoolKind.FULL]
+        self.assertEqual(group.k_head.shape, (16, CANARY_SLOT_BYTES))
+        self.assertEqual(group.k_tail.shape, (16, CANARY_SLOT_BYTES))
+        self.assertIsNotNone(group.v_head)
+        self.assertIsNotNone(group.v_tail)
+        self.assertEqual(group.v_head.shape, (16, CANARY_SLOT_BYTES))
+
+
+class TestPoolPatcherBufferInfos(PoolPatcherHelper, CustomTestCase):
+    def test_get_contiguous_buf_infos_inserts_canary_entries(self):
+        """Verify contiguous buffer metadata includes canary entries after patching."""
+        for patched in (False, True):
+            with self.subTest(patched=patched):
+                pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)
+                ptrs_before, _, _ = pool.get_contiguous_buf_infos()
+                n_before = len(ptrs_before)
+
+                if patched:
+                    attach_canary_buffers(
+                        pool=pool,
+                        config=self.config,
+                        device=self.device,
+                        kv_token_id_vs_position_offset=0,
+                    )
+                    ptrs_after, _, _ = pool.get_contiguous_buf_infos()
+                    self.assertEqual(len(ptrs_after), n_before + 4)
+                else:
+                    ptrs_after, _, _ = pool.get_contiguous_buf_infos()
+                    self.assertEqual(ptrs_after, ptrs_before)
+
+    def test_pd_layout_canary_inserted_correctly(self):
+        """Verify PD (prefill-decode disaggregation) canary buffers are inserted in layout order."""
+        pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)
+        k_ptrs_orig = [b.data_ptr() for b in pool.k_buffer]
+        v_ptrs_orig = [b.data_ptr() for b in pool.v_buffer]
+
+        groups_tuple = attach_canary_buffers(
+            pool=pool,
+            config=self.config,
+            device=self.device,
+            kv_token_id_vs_position_offset=0,
+        )
+        group = {g.kind: g for g in groups_tuple}[PoolKind.FULL]
+        ptrs_after, _, _ = pool.get_contiguous_buf_infos()
+
+        canary_k_ptrs = [group.k_head.data_ptr(), group.k_tail.data_ptr()]
+        canary_v_ptrs = [group.v_head.data_ptr(), group.v_tail.data_ptr()]
+        self.assertEqual(ptrs_after[0], canary_k_ptrs[0])
+        self.assertEqual(ptrs_after[1 : 1 + len(k_ptrs_orig)], k_ptrs_orig)
+        k_tail_idx = 1 + len(k_ptrs_orig)
+        self.assertEqual(ptrs_after[k_tail_idx], canary_k_ptrs[1])
+        self.assertEqual(ptrs_after[k_tail_idx + 1], canary_v_ptrs[0])
+        v_start = k_tail_idx + 2
+        self.assertEqual(ptrs_after[v_start : v_start + len(v_ptrs_orig)], v_ptrs_orig)
+        self.assertEqual(ptrs_after[-1], canary_v_ptrs[1])
+
+
+class TestCanaryBufferBudget(PoolPatcherHelper, CustomTestCase):
+    def test_canary_buf_per_token_bytes_within_budget(self):
+        """Verify canary per-token storage stays below the real KV budget."""
+        pool = make_mha_pool(self.device, num_slots=16, dim=64, layer_num=2)
+        groups_tuple = attach_canary_buffers(
+            pool=pool,
+            config=self.config,
+            device=self.device,
+            kv_token_id_vs_position_offset=0,
+        )
+        group = {g.kind: g for g in groups_tuple}[PoolKind.FULL]
+
+        slot_stride_bytes = group.k_head.stride(0) * group.k_head.element_size()
+        self.assertLessEqual(slot_stride_bytes, CANARY_SLOT_BYTES)
+
+        real_kv_per_token_bytes = (
+            pool.k_buffer[0].stride(0) * pool.k_buffer[0].element_size()
+        )
+        self.assertLess(slot_stride_bytes, real_kv_per_token_bytes)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_pool_patcher_utils.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
+
+
+class _FakeObj:
+    def greet(self, name: str) -> str:
+        return f"hello {name}"
+
+
+class TestPoolPatcherUtils(CustomTestCase):
+    def test_wrap_method_delegates_to_wrapper(self) -> None:
+        """Verify wrapped methods delegate through the wrapper."""
+        obj = _FakeObj()
+
+        def _with_validation(original, name: str) -> str:
+            return original(name) + "!"
+
+        wrap_method(obj, "greet", wrapper=_with_validation)
+        self.assertEqual(obj.greet("world"), "hello world!")
+
+    def test_wrap_method_missing_method_raises_attribute_error(self) -> None:
+        """Verify wrapping a missing method raises AttributeError."""
+        obj = _FakeObj()
+        with self.assertRaisesRegex(AttributeError, "missing required method"):
+            wrap_method(
+                obj, "nonexistent", wrapper=lambda orig, *a, **kw: orig(*a, **kw)
+            )
+
+    def test_wrap_method_double_wrap_raises_runtime_error(self) -> None:
+        """Verify wrapping the same method twice raises RuntimeError."""
+        obj = _FakeObj()
+        wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw))
+        with self.assertRaisesRegex(RuntimeError, "already wrapped by kv-canary"):
+            wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw))
+
+    def test_wrap_method_preserves_functools_wraps_metadata(self) -> None:
+        """Verify wrapping preserves method metadata."""
+        obj = _FakeObj()
+        original_name = obj.greet.__name__
+        wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw))
+        self.assertEqual(obj.greet.__name__, original_name)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_violation.py

@@ -0,0 +1,333 @@
+from __future__ import annotations
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.consts import FailReason
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.srt.kv_canary.runner import violation_reporter as violation_reporter_module
+from sglang.srt.kv_canary.runner.violation_reporter import (
+    ViolationReporter,
+    _format_violation,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=5, stage="extra-a", runner_config="1-gpu-small")
+
+
+def _make_row(
+    *,
+    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
+    slot_idx: int = 17,
+    position: int = 42,
+    stored_token: int = 111,
+    expected_token: int = 0,
+    stored_chain_hash: int = 0,
+    expected_aux: int = 0,
+    fail_reason_bits: int = 0,
+) -> list[int]:
+    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] = position
+    row[consts.VIOLATION_FIELD_STORED_TOKEN] = stored_token
+    row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = expected_token
+    row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = stored_chain_hash
+    row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_aux
+    row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = fail_reason_bits
+    return row
+
+
+class TestViolationReporter(CustomTestCase):
+    def test_format_violation_verify_path_labels_each_bit(self) -> None:
+        """Verify verify-path violations render each fail-reason bit."""
+        row = _make_row(
+            stored_chain_hash=0x1111111111111111,
+            expected_aux=0x2222222222222222,
+            fail_reason_bits=int(
+                FailReason.VERIFY_CHAIN_HASH_MISMATCH
+                | FailReason.VERIFY_POSITION_MISMATCH
+            ),
+        )
+        out = _format_violation(
+            row=row, total=1, ring_overflow=False, step_when_pumped=7
+        )
+        self.assertEqual(
+            out,
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash+verify_position "
+            "slot_idx=17 position=42 stored_token=111 expected_token=0 stored_chain_hash=0x1111111111111111 "
+            "expected_aux=0x2222222222222222\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: verify_chain_hash verify_position\n"
+            "  stored:   token_id=111   position=42 prev_hash=0x1111111111111111\n"
+            "  expected: prev_hash=0x2222222222222222\n"
+            "  total_violations=1 ring_overflow=False step_when_pumped=7",
+        )
+
+    def test_format_violation_write_token_mismatch_labels_and_position(self) -> None:
+        """Verify write-token violations render token and position details."""
+        row = _make_row(
+            position=42,
+            stored_token=999,
+            expected_token=888,
+            stored_chain_hash=0xDEADBEEFCAFEBABE,
+            expected_aux=43,
+            fail_reason_bits=int(FailReason.WRITE_TOKEN_MISMATCH),
+        )
+        out = _format_violation(
+            row=row, total=1, ring_overflow=False, step_when_pumped=0
+        )
+        self.assertEqual(
+            out,
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_token slot_idx=17 position=42 "
+            "stored_token=999 expected_token=888 stored_chain_hash=0xdeadbeefcafebabe "
+            "expected_aux=0x000000000000002b\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: write_token\n"
+            "  actual:   token_id=999   position=42 prev_hash=0xdeadbeefcafebabe\n"
+            "  expected: token_id=888   position=43\n"
+            "  total_violations=1 ring_overflow=False step_when_pumped=0",
+        )
+
+    def test_format_violation_write_position_mismatch_uses_expected_aux_as_position(
+        self,
+    ) -> None:
+        """Verify write-position violations render expected_aux as a position."""
+        row = _make_row(
+            position=42,
+            stored_token=111,
+            expected_token=111,
+            expected_aux=99,
+            fail_reason_bits=int(FailReason.WRITE_POSITION_MISMATCH),
+        )
+        out = _format_violation(
+            row=row, total=1, ring_overflow=False, step_when_pumped=0
+        )
+        self.assertEqual(
+            out,
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_position slot_idx=17 position=42 "
+            "stored_token=111 expected_token=111 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000063\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: write_position\n"
+            "  actual:   token_id=111   position=42 prev_hash=0x0000000000000000\n"
+            "  expected: token_id=111   position=99\n"
+            "  total_violations=1 ring_overflow=False step_when_pumped=0",
+        )
+
+    def test_format_violation_combined_write_bits_render_both_labels(self) -> None:
+        """Verify combined write violation bits render both labels."""
+        row = _make_row(
+            fail_reason_bits=int(
+                FailReason.WRITE_TOKEN_MISMATCH | FailReason.WRITE_POSITION_MISMATCH
+            ),
+        )
+        out = _format_violation(
+            row=row, total=1, ring_overflow=False, step_when_pumped=0
+        )
+        self.assertEqual(
+            out,
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_token+write_position "
+            "slot_idx=17 position=42 stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: write_token write_position\n"
+            "  actual:   token_id=111   position=42 prev_hash=0x0000000000000000\n"
+            "  expected: token_id=0   position=0\n"
+            "  total_violations=1 ring_overflow=False step_when_pumped=0",
+        )
+
+    def test_format_violation_unknown_kernel_kind_renders_unknown_label(self) -> None:
+        """Verify unknown kernel kinds render an unknown label."""
+        row = _make_row(fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH))
+        row[consts.VIOLATION_FIELD_KERNEL_KIND] = 9999
+        out = _format_violation(
+            row=row, total=1, ring_overflow=False, step_when_pumped=0
+        )
+        self.assertEqual(
+            out,
+            "kv_canary violation: launch_tag=unknown(9999) fail_reason=verify_chain_hash slot_idx=17 position=42 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=unknown(9999), slot_idx=17, position=42)\n"
+            "canary_kind:       unknown(9999)\n"
+            "  fail_reasons: verify_chain_hash\n"
+            "  stored:   token_id=111   position=42 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=1 ring_overflow=False step_when_pumped=0",
+        )
+
+
+def _make_reporter(
+    *,
+    rows: list[list[int]],
+    write_index: int,
+    ring_capacity: int,
+    mode: CanaryMode = CanaryMode.LOG,
+) -> ViolationReporter:
+    ring = torch.zeros(ring_capacity, consts.VIOLATION_FIELDS, dtype=torch.int64)
+    for i, row in enumerate(rows):
+        ring[i] = torch.tensor(row, dtype=torch.int64)
+    violation_log = SimpleNamespace(
+        violation_ring=ring,
+        violation_write_index=torch.tensor([write_index], dtype=torch.int32),
+    )
+    device_state = SimpleNamespace(violation_log=violation_log)
+    config = SimpleNamespace(mode=mode)
+    return ViolationReporter(config=config, device_state=device_state)
+
+
+class TestLogOrRaiseViolation(CustomTestCase):
+    def test_log_or_raise_violation_empty_ring_is_noop(self) -> None:
+        """Empty ring (write_index=0) emits no warning and leaves reporter non-raised."""
+        reporter = _make_reporter(
+            rows=[], write_index=0, ring_capacity=4, mode=CanaryMode.LOG
+        )
+        with patch.object(violation_reporter_module.logger, "warning") as mock_warning:
+            reporter.log_or_raise_violation(outer_step_counter=0)
+        mock_warning.assert_not_called()
+        self.assertFalse(reporter.is_raised)
+
+    def test_log_mode_emits_one_warning_per_violation(self) -> None:
+        """Log mode with 2 valid rows emits 2 warnings, each a full _format_violation snapshot for that row."""
+        rows = [
+            _make_row(
+                slot_idx=11,
+                position=101,
+                fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH),
+            ),
+            _make_row(
+                slot_idx=22,
+                position=202,
+                fail_reason_bits=int(FailReason.VERIFY_POSITION_MISMATCH),
+            ),
+        ]
+        reporter = _make_reporter(
+            rows=rows, write_index=2, ring_capacity=4, mode=CanaryMode.LOG
+        )
+        with patch.object(violation_reporter_module.logger, "warning") as mock_warning:
+            reporter.log_or_raise_violation(outer_step_counter=7)
+
+        self.assertEqual(mock_warning.call_count, 2)
+        messages: list[str] = [call.args[0] for call in mock_warning.call_args_list]
+        self.assertEqual(
+            messages[0],
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash slot_idx=11 position=101 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: verify_chain_hash\n"
+            "  stored:   token_id=111   position=101 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=2 ring_overflow=False step_when_pumped=7",
+        )
+        self.assertEqual(
+            messages[1],
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_position slot_idx=22 position=202 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: verify_position\n"
+            "  stored:   token_id=111   position=202 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=2 ring_overflow=False step_when_pumped=7",
+        )
+        self.assertFalse(reporter.is_raised)
+
+    def test_raise_mode_raises_one_error_containing_all_violations(self) -> None:
+        """Raise mode raises a single RuntimeError whose text is the 2 formatted rows joined with single newlines."""
+        rows = [
+            _make_row(
+                slot_idx=11,
+                position=101,
+                fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH),
+            ),
+            _make_row(
+                slot_idx=22,
+                position=202,
+                fail_reason_bits=int(FailReason.VERIFY_POSITION_MISMATCH),
+            ),
+        ]
+        reporter = _make_reporter(
+            rows=rows, write_index=2, ring_capacity=4, mode=CanaryMode.RAISE
+        )
+        with self.assertRaises(RuntimeError) as ctx:
+            reporter.log_or_raise_violation(outer_step_counter=5)
+
+        self.assertEqual(
+            str(ctx.exception),
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash slot_idx=11 position=101 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: verify_chain_hash\n"
+            "  stored:   token_id=111   position=101 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=2 ring_overflow=False step_when_pumped=5\n"
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_position slot_idx=22 position=202 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: verify_position\n"
+            "  stored:   token_id=111   position=202 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=2 ring_overflow=False step_when_pumped=5",
+        )
+        self.assertTrue(reporter.is_raised)
+
+    def test_log_mode_ring_overflow_marks_overflow_in_each_row(self) -> None:
+        """Log mode with write_index=5 but ring_capacity=2 emits 2 warnings, each a full snapshot with overflow footer."""
+        rows = [
+            _make_row(slot_idx=11, position=101),
+            _make_row(slot_idx=22, position=202),
+        ]
+        reporter = _make_reporter(
+            rows=rows, write_index=5, ring_capacity=2, mode=CanaryMode.LOG
+        )
+        with patch.object(violation_reporter_module.logger, "warning") as mock_warning:
+            reporter.log_or_raise_violation(outer_step_counter=0)
+
+        self.assertEqual(mock_warning.call_count, 2)
+        messages: list[str] = [call.args[0] for call in mock_warning.call_args_list]
+        self.assertEqual(
+            messages[0],
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=none slot_idx=11 position=101 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: none\n"
+            "  stored:   token_id=111   position=101 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=5 ring_overflow=True step_when_pumped=0",
+        )
+        self.assertEqual(
+            messages[1],
+            "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=none slot_idx=22 position=202 "
+            "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 "
+            "expected_aux=0x0000000000000000\n"
+            "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n"
+            "canary_kind:       per_forward_head_k_full\n"
+            "  fail_reasons: none\n"
+            "  stored:   token_id=111   position=202 prev_hash=0x0000000000000000\n"
+            "  expected: prev_hash=0x0000000000000000\n"
+            "  total_violations=5 ring_overflow=True step_when_pumped=0",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()