chore: import upstream snapshot with attribution

python/sglang/srt/kv_canary/api.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.perturb.config import PerturbConfig
+from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
+from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
+from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
+from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+def install_canary(
+    *,
+    server_args: ServerArgs,
+    model_runner: ModelRunner,
+    token_oracle_manager: Optional[TokenOracleManager] = None,
+) -> Optional[CanaryManager]:
+    config = CanaryConfig.from_env(server_args)
+    if config.mode is CanaryMode.NONE:
+        return None
+
+    assert not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE), (
+        "kv-canary: piecewise cuda graph is not supported by the current "
+        "SingleForwardManager design; set --cuda-graph-backend-prefill=disabled "
+        "(or =breakable) when canary is enabled"
+    )
+
+    perturb_config = PerturbConfig.from_env()
+    device = torch.device(model_runner.device)
+    # EAGLE draft worker pools rotate input_ids so slot ``p`` stores K/V for the token at position ``p+1``;
+    # target pools have no such shift. Threaded into the plan-side expected-token gather kernel.
+    kv_token_id_vs_position_offset = 1 if model_runner.is_draft_worker else 0
+    buffer_groups = attach_canary_buffers(
+        pool=model_runner.token_to_kv_pool,
+        config=config,
+        device=device,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    allocator = model_runner.token_to_kv_pool_allocator
+    swa_allocator = (
+        allocator if isinstance(allocator, SWATokenToKVPoolAllocator) else None
+    )
+    launch_capacities = CanaryLaunchCapacities.from_args(
+        server_args=model_runner.server_args,
+        req_to_token_pool_size=model_runner.req_to_token_pool.size,
+        max_seq_len_per_req=model_runner.req_to_token_pool.req_to_token.shape[1],
+        pool_slot_count=model_runner.max_total_num_tokens,
+    )
+    swa_window_size = model_runner.sliding_window_size or 0
+    speculative_num_steps = int(server_args.speculative_num_steps or 1)
+    manager = CanaryManager(
+        config=config,
+        perturb_config=perturb_config,
+        buffer_groups=buffer_groups,
+        device=device,
+        req_to_token_pool=model_runner.req_to_token_pool,
+        launch_capacities=launch_capacities,
+        swa_window_size=swa_window_size,
+        token_oracle_manager=token_oracle_manager,
+        swa_allocator=swa_allocator,
+        speculative_num_steps=speculative_num_steps,
+        is_eagle_draft_decode=model_runner.is_draft_worker,
+    )
+
+    _patch_model_forward(model_runner=model_runner, manager=manager)
+
+    # Single-line summary of every knob that controls canary behavior at boot time.
+    # Disaggregation mode is included so PD logs are unambiguous about which side this is.
+    logger.info(
+        "install_canary: disaggregation_mode=%s config=%s perturb_config=%s "
+        "launch_capacities=%s n_buffer_groups=%d buffer_group_kinds=%s "
+        "swa_window_size=%d speculative_num_steps=%d",
+        server_args.disaggregation_mode,
+        config,
+        perturb_config,
+        launch_capacities,
+        len(buffer_groups),
+        [g.kind.name for g in buffer_groups],
+        swa_window_size,
+        speculative_num_steps,
+    )
+    return manager
+
+
+def _patch_model_forward(*, model_runner: ModelRunner, manager: CanaryManager) -> None:
+    def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any:
+        with manager.model_forward_bracket_scope() as should_bracket:
+            if not should_bracket:
+                # Nested model.forward calls share the active SingleForwardManager.
+                # Only the outermost call may run kv-canary pre/post ops; otherwise
+                # the phase checker sees a second pre-op before the first post-op.
+                return original(*args, **kwargs)
+
+            forward_batch = _extract_forward_batch(args, kwargs)
+            assert (
+                forward_batch is not None
+            ), "kv-canary: patched model.forward called without a ForwardBatch"
+
+            canary_pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch)
+            output = original(*args, **kwargs)
+            manager.post_ops_maybe_inside_graph(forward_batch, canary_pre_ops_output)
+            return output
+
+    wrap_method(model_runner.model, "forward", wrapper=_with_canary_bracketing)
+
+
+def _extract_forward_batch(args, kwargs) -> Optional[ForwardBatch]:
+    if "forward_batch" in kwargs and isinstance(kwargs["forward_batch"], ForwardBatch):
+        return kwargs["forward_batch"]
+    for arg in args:
+        if isinstance(arg, ForwardBatch):
+            return arg
+    return None

python/sglang/srt/kv_canary/buffer_group.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import RealKvSource
+
+
+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 + real-KV sources for one (PoolKind × K-half | V-half) on a pool.
+
+    Each (head | tail) launch sees a single 2-D uint8 buf for the canary, plus a list of RealKvSource for the
+    real-KV mixin. 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; real_kv_sources_v is empty). 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.
+        real_kv_sources_k: Real KV pieces folded into the K-half canary's real_kv_hash. Tuple length is
+            pool-specific (1 for simple MHA, more for multi-layer / weird-layout pools). Empty tuple =
+            real-KV mixin disabled for this half.
+        real_kv_sources_v: Same for V-half. Empty tuple iff v_head is None or the mixin is disabled.
+        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]
+    real_kv_sources_k: tuple[RealKvSource, ...]
+    real_kv_sources_v: tuple[RealKvSource, ...]
+    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,117 @@
+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_config = server_args.cuda_graph_config
+        cuda_graph_max_bs = (
+            cuda_graph_config.decode.max_bs if cuda_graph_config is not None else 0
+        ) 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,80 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from sglang.jit_kernel.kv_canary.consts import (
+    RealKvHashMode,
+)
+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, token_oracle_manager) 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.
+        sweep_interval: 0 disables sweep entirely; positive N means every N-th forward step the runner
+            additionally walks all radix-tree-held slots (overlap with per-forward HEAD/TAIL is harmless
+            redundancy) and verifies them.
+        real_kv_hash_mode: RealKvHashMode (NONE / PARTIAL / ALL). Uniform across head/tail/sweep launches;
+            PARTIAL (first 16B, hard cap) is cheap enough for production defaults.
+        enable_write_input_assert: bool. True = launch_canary_write_kernel additionally compares
+            forward_batch.input_ids[i] / positions[i] against caller-supplied expected_input_tokens[i] /
+            expected_input_positions[i]; mismatch records a violation. Only useful when something else
+            (e.g. token_oracle.oracle_manager.fill_expected_inputs) is feeding the expected_* placeholders
+            per forward — canary itself knows no oracle.
+        enable_verify_token_assert: bool. True = real-model token-id validator: build
+            expected_tokens from each req's ``origin_input_ids + output_ids`` (snapshotted at
+            ForwardBatch.init_new) and compare against the canary's stored tokens at verify time.
+            Independent of ``enable_write_input_assert``.
+        stats_print_every_n_steps: 0 disables periodic stats logging; positive N prints
+            "canary protected N tokens, ran M sweep passes, K violations so far" every N forward steps.
+    """
+
+    mode: CanaryMode
+    ring_capacity: int
+    sweep_interval: int
+    real_kv_hash_mode: RealKvHashMode
+    enable_write_input_assert: bool
+    enable_verify_token_assert: bool
+    stats_print_every_n_steps: 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}"
+            )
+
+        real_kv_raw = server_args.kv_canary_real_data.strip().upper()
+
+        return cls(
+            mode=CanaryMode(mode_raw),
+            ring_capacity=envs.SGLANG_KV_CANARY_RING_CAPACITY.get(),
+            sweep_interval=server_args.kv_canary_sweep_interval,
+            real_kv_hash_mode=RealKvHashMode[real_kv_raw],
+            enable_write_input_assert=envs.SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT.get(),
+            enable_verify_token_assert=envs.SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT.get(),
+            stats_print_every_n_steps=envs.SGLANG_KV_CANARY_STATS_PRINT_EVERY_N_STEPS.get(),
+        )

python/sglang/srt/kv_canary/endpoint.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.consts import (
+    RealKvHashMode,
+)
+from sglang.jit_kernel.kv_canary.verify import (
+    CanaryLaunchTag,
+    RealKvSource,
+    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]
+    real_kv_sources: tuple[RealKvSource, ...]
+    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,
+        real_kv_hash_mode: RealKvHashMode,
+    ) -> None:
+        if _is_sweep_tag(self.kernel_kind):
+            raise NotImplementedError(
+                f"kv-canary: launch_per_forward not supported on sweep endpoint {self.kernel_kind.name}"
+            )
+
+        context = self._make_verify_or_write_context(
+            violation_log=violation_log,
+            real_kv_hash_mode=real_kv_hash_mode,
+        )
+        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 launch_sweep(
+        self,
+        *,
+        verify_plan: VerifyPlan,
+        violation_log: ViolationLog,
+        real_kv_hash_mode: RealKvHashMode,
+    ) -> None:
+        if not _is_sweep_tag(self.kernel_kind):
+            raise NotImplementedError(
+                f"kv-canary: launch_sweep not supported on non-sweep endpoint {self.kernel_kind.name}"
+            )
+
+        launch_canary_verify_kernel(
+            context=self._make_verify_or_write_context(
+                violation_log=violation_log,
+                real_kv_hash_mode=real_kv_hash_mode,
+            ),
+            plan=verify_plan,
+            check_verify_expected_token=False,
+        )
+
+    def _make_verify_or_write_context(
+        self,
+        *,
+        violation_log: ViolationLog,
+        real_kv_hash_mode: RealKvHashMode,
+    ) -> 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,
+            real_kv_sources=self.real_kv_sources,
+            real_kv_hash_mode=real_kv_hash_mode,
+            enable_chain_position_assert=self.enable_chain_position_assert,
+        )
+
+
+def _is_sweep_tag(tag: CanaryLaunchTag) -> bool:
+    return tag in (
+        CanaryLaunchTag.SWEEP_K_FULL,
+        CanaryLaunchTag.SWEEP_V_FULL,
+        CanaryLaunchTag.SWEEP_K_SWA,
+        CanaryLaunchTag.SWEEP_V_SWA,
+    )
+
+
+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
+
+
+def _resolve_real_kv_sources(
+    *,
+    half: str,
+    group: CanaryBufferGroup,
+) -> tuple[RealKvSource, ...]:
+    if half == "K":
+        return group.real_kv_sources_k
+    return group.real_kv_sources_v
+
+
+_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"),
+    (CanaryLaunchTag.SWEEP_K_FULL, "SWEEP", "K"),
+    (CanaryLaunchTag.SWEEP_V_FULL, "SWEEP", "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"),
+    (CanaryLaunchTag.SWEEP_K_SWA, "SWEEP", "K"),
+    (CanaryLaunchTag.SWEEP_V_SWA, "SWEEP", "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
+
+        buf_slot = "TAIL" if slot == "SWEEP" else slot
+        canary_buf = _resolve_canary_buf(slot=buf_slot, half=half, group=group)
+        real_kv_sources = (
+            () if slot == "HEAD" else _resolve_real_kv_sources(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,
+                real_kv_sources=real_kv_sources,
+                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/perturb/config.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+
+from sglang.srt.environ import envs
+from sglang.srt.kv_canary.buffer_group import PoolKind
+
+
+class TargetGroupKind(IntEnum):
+    FULL = PoolKind.FULL.value
+    SWA = PoolKind.SWA.value
+
+    def __str__(self) -> str:
+        return self.name.lower()
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class PerturbConfig:
+    req_to_token_prob: float
+    real_kv_used_prob: float
+    real_kv_unused_cache_prob: float
+    real_kv_post_forward_prob: float
+    target_group_kind: TargetGroupKind | None
+    warmup_steps: int
+
+    @classmethod
+    def from_env(cls) -> PerturbConfig:
+        real_kv_used_prob = envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB.get()
+        real_kv_unused_cache_prob = (
+            envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB.get()
+        )
+        real_kv_post_forward_prob = (
+            envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_POST_FORWARD_PROB.get()
+        )
+        return cls(
+            req_to_token_prob=envs.SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB.get(),
+            real_kv_used_prob=real_kv_used_prob,
+            real_kv_unused_cache_prob=real_kv_unused_cache_prob,
+            real_kv_post_forward_prob=real_kv_post_forward_prob,
+            target_group_kind=_parse_target_group_kind_from_env(
+                raw=envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.get(),
+                real_kv_used_prob=real_kv_used_prob,
+                real_kv_unused_cache_prob=real_kv_unused_cache_prob,
+                real_kv_post_forward_prob=real_kv_post_forward_prob,
+            ),
+            warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(),
+        )
+
+
+def _parse_target_group_kind_from_env(
+    *,
+    raw: str | None,
+    real_kv_used_prob: float,
+    real_kv_unused_cache_prob: float,
+    real_kv_post_forward_prob: float,
+) -> TargetGroupKind | None:
+    if raw is not None and raw.strip():
+        return _parse_target_group_kind(raw)
+    if (
+        real_kv_used_prob > 0.0
+        or real_kv_unused_cache_prob > 0.0
+        or real_kv_post_forward_prob > 0.0
+    ):
+        return _parse_target_group_kind(raw)
+    return None
+
+
+def require_target_group_kind(
+    *, target_group_kind: TargetGroupKind | None, perturb_name: str
+) -> TargetGroupKind:
+    if target_group_kind is None:
+        raise ValueError(
+            "SGLANG_KV_CANARY_PERTURB_TARGET_GROUP must be explicitly set to "
+            f"'full' or 'swa' when {perturb_name} perturbation is enabled"
+        )
+    return target_group_kind
+
+
+def _parse_target_group_kind(raw: str | None) -> TargetGroupKind:
+    if raw is None or not raw.strip():
+        raise ValueError(
+            "SGLANG_KV_CANARY_PERTURB_TARGET_GROUP must be explicitly set to "
+            "'full' or 'swa'"
+        )
+
+    value = raw.strip().lower()
+    try:
+        return TargetGroupKind[value.upper()]
+    except KeyError:
+        raise ValueError(
+            "SGLANG_KV_CANARY_PERTURB_TARGET_GROUP must be one of 'full' / "
+            f"'swa', got {raw!r}"
+        ) from None

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

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.perturb import (
+    real_kv_post_forward,
+    real_kv_unused_cache,
+    real_kv_used,
+    req_to_token,
+)
+from sglang.srt.kv_canary.perturb.config import PerturbConfig
+from sglang.srt.kv_canary.perturb.utils import WarmupGate
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class PerturbManager:
+    def __init__(
+        self,
+        *,
+        config: PerturbConfig,
+        req_to_token_pool: ReqToTokenPool,
+        buffer_groups: tuple[CanaryBufferGroup, ...],
+        outer_step_counter_getter: Callable[[], int],
+        swa_window_size: int = 0,
+        sweep_interval: int = 0,
+    ) -> None:
+        self._config = config
+        self._req_to_token_pool = req_to_token_pool
+        self._buffer_groups = buffer_groups
+        self._outer_step_counter_getter = outer_step_counter_getter
+        self._swa_window_size = swa_window_size
+        self._sweep_interval = sweep_interval
+        self._radix_cache: Optional[BasePrefixCache] = None
+        self._warmup_gate = WarmupGate(
+            config=config, outer_step_counter_getter=outer_step_counter_getter
+        )
+
+    def attach_radix_cache(self, radix_cache: BasePrefixCache) -> None:
+        self._radix_cache = radix_cache
+
+    def perturb(
+        self,
+        *,
+        maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    ) -> None:
+        self.perturb_req_to_token(maybe_inaccurate_forward_batch)
+        self.perturb_real_kv_used(maybe_inaccurate_forward_batch)
+        self.perturb_real_kv_unused_cache(maybe_inaccurate_forward_batch)
+
+    def perturb_post_forward(
+        self,
+        *,
+        maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    ) -> None:
+        self.perturb_real_kv_post_forward(maybe_inaccurate_forward_batch)
+
+    def perturb_req_to_token(
+        self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
+    ) -> None:
+        req_to_token.run(
+            maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            config=self._config,
+            req_to_token_pool=self._req_to_token_pool,
+            warmup_gate=self._warmup_gate,
+        )
+
+    def perturb_real_kv_used(
+        self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
+    ) -> None:
+        real_kv_used.run(
+            maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            config=self._config,
+            req_to_token_pool=self._req_to_token_pool,
+            buffer_groups=self._buffer_groups,
+            swa_window_size=self._swa_window_size,
+            warmup_gate=self._warmup_gate,
+        )
+
+    def perturb_real_kv_unused_cache(
+        self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
+    ) -> None:
+        real_kv_unused_cache.run(
+            maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            config=self._config,
+            buffer_groups=self._buffer_groups,
+            radix_cache=self._radix_cache,
+            swa_window_size=self._swa_window_size,
+            sweep_interval=self._sweep_interval,
+            outer_step_counter=self._outer_step_counter_getter(),
+            warmup_gate=self._warmup_gate,
+        )
+
+    def perturb_real_kv_post_forward(
+        self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
+    ) -> None:
+        real_kv_post_forward.run(
+            maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            config=self._config,
+            buffer_groups=self._buffer_groups,
+            warmup_gate=self._warmup_gate,
+        )

python/sglang/srt/kv_canary/perturb/next_token_swap.py

@@ -0,0 +1,83 @@
+"""Swap two requests' sampled next tokens at the sampler exit.
+
+KV path is untouched, so kv_canary KV-side fail_reasons stay silent. The
+token-oracle input check downstream MUST report fail_reason=write_token — this
+validates that the input-check link is genuinely active.
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class NextTokenSwapConfig:
+    prob: float
+    warmup_steps: int
+
+    @classmethod
+    def from_env(cls) -> NextTokenSwapConfig:
+        return cls(
+            prob=envs.SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB.get(),
+            warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(),
+        )
+
+
+_config: Optional[NextTokenSwapConfig] = None
+_step_counter: int = 0
+
+
+def _get_config() -> NextTokenSwapConfig:
+    global _config
+    if _config is None:
+        _config = NextTokenSwapConfig.from_env()
+    return _config
+
+
+def maybe_perturb_swap_next_tokens(
+    batch_next_token_ids: torch.Tensor,
+) -> torch.Tensor:
+    global _step_counter
+
+    config = _get_config()
+    step = _step_counter
+    _step_counter += 1
+
+    if config.prob <= 0.0:
+        return batch_next_token_ids
+    if step < config.warmup_steps:
+        return batch_next_token_ids
+    if batch_next_token_ids.shape[0] < 2:
+        return batch_next_token_ids
+
+    if random.random() >= config.prob:
+        return batch_next_token_ids
+
+    batch_size = batch_next_token_ids.shape[0]
+    i = random.randrange(batch_size)
+    j = random.randrange(batch_size)
+    while j == i:
+        j = random.randrange(batch_size)
+
+    swapped = batch_next_token_ids.clone()
+    swapped[i], swapped[j] = (
+        batch_next_token_ids[j].clone(),
+        batch_next_token_ids[i].clone(),
+    )
+
+    logger.info(
+        "kv_canary perturb next_token_swap: swapped i=%d j=%d step=%d",
+        i,
+        j,
+        step,
+    )
+    return swapped

python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py

@@ -0,0 +1,78 @@
+"""Pick a random real-KV source byte derived from out_cache_loc and flip it
+in-place AFTER the TAIL kernel has captured its canary hash.
+
+The slot id is taken from maybe_inaccurate_forward_batch.out_cache_loc and used
+only as a lookup into the target group's real-KV source buffer; the actual flip
+happens inside that buffer. The flip is a PyTorch indexed write on the current
+CUDA stream; because TAIL is launched on the same stream, stream ordering
+guarantees it happens-after TAIL's canary write.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.perturb.config import (
+    PerturbConfig,
+    require_target_group_kind,
+)
+from sglang.srt.kv_canary.perturb.slot_picker import pick_out_cache_loc_slot
+from sglang.srt.kv_canary.perturb.utils import (
+    WarmupGate,
+    flip_random_source_byte_and_log,
+    pick_target_group,
+    should_run_perturbation,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+def run(
+    *,
+    maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    config: PerturbConfig,
+    buffer_groups: tuple[CanaryBufferGroup, ...],
+    warmup_gate: WarmupGate,
+) -> None:
+    if not should_run_perturbation(
+        perturb_name="real_kv_post_forward",
+        probability=config.real_kv_post_forward_prob,
+        warmup_gate=warmup_gate,
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+    ):
+        return
+
+    slot = pick_out_cache_loc_slot(
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch
+    )
+    if slot is None:
+        logger.info(
+            "kv_canary perturb real_kv_post_forward: skipped because maybe_inaccurate_forward_batch.out_cache_loc "
+            "had no valid slot"
+        )
+        return
+    group = pick_target_group(
+        buffer_groups=buffer_groups,
+        target_kind=require_target_group_kind(
+            target_group_kind=config.target_group_kind,
+            perturb_name="real_kv_post_forward",
+        ),
+    )
+    if group is None:
+        logger.info(
+            "kv_canary perturb real_kv_post_forward: skipped because no target group matched "
+            "target_group_kind=%s slot=%d",
+            config.target_group_kind,
+            slot,
+        )
+        return
+    flip_random_source_byte_and_log(
+        perturb_name="real_kv_post_forward",
+        group=group,
+        slot_idx=slot,
+    )

python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py

@@ -0,0 +1,162 @@
+"""Flip the first byte of a radix-cached but currently-unused (orphan) slot.
+
+Detection should come from sweep (per-forward verify won't even look at this
+slot). Designed to surface bugs where cached KV is silently corrupted and
+sleeps until much later when a prefix happens to match.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.perturb.config import (
+    PerturbConfig,
+    require_target_group_kind,
+)
+from sglang.srt.kv_canary.perturb.utils import (
+    WarmupGate,
+    flip_first_byte_in_source,
+    pick_target_group,
+    should_run_perturbation,
+)
+from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+def run(
+    *,
+    maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    config: PerturbConfig,
+    buffer_groups: tuple[CanaryBufferGroup, ...],
+    radix_cache: Optional[BasePrefixCache],
+    swa_window_size: int,
+    sweep_interval: int,
+    outer_step_counter: int,
+    warmup_gate: WarmupGate,
+) -> None:
+    if sweep_interval <= 0 or outer_step_counter % sweep_interval != 0:
+        return
+
+    if not should_run_perturbation(
+        perturb_name="real_kv_unused_cache",
+        probability=config.real_kv_unused_cache_prob,
+        warmup_gate=warmup_gate,
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+        require_forward_batch=False,
+    ):
+        return
+
+    group = pick_target_group(
+        buffer_groups=buffer_groups,
+        target_kind=require_target_group_kind(
+            target_group_kind=config.target_group_kind,
+            perturb_name="real_kv_unused_cache",
+        ),
+    )
+    if group is None:
+        logger.info(
+            "kv_canary perturb real_kv_unused_cache: skipped because no target group with "
+            "real_kv_sources_k matched target_group_kind=%s",
+            config.target_group_kind,
+        )
+        return
+    slot = _pick_sweep_slot_for_group(
+        radix_cache=radix_cache,
+        group=group,
+        swa_window_size=swa_window_size,
+    )
+    if slot is None:
+        logger.info(
+            "kv_canary perturb real_kv_unused_cache: skipped because no orphan sweep slot "
+            "was found for group=%s",
+            group.kind.name,
+        )
+        return
+    source_pick = int(torch.randint(0, len(group.real_kv_sources_k), (1,)).item())
+    source = group.real_kv_sources_k[source_pick]
+    flip_result = flip_first_byte_in_source(
+        group=group,
+        source=source,
+        slot_idx=slot,
+        slot_is_physical=True,
+    )
+    if flip_result is None:
+        logger.info(
+            "kv_canary perturb real_kv_unused_cache: skipped because slot=%d could not be mapped "
+            "into group=%s source_idx=%d",
+            slot,
+            group.kind.name,
+            source_pick,
+        )
+        return
+    row, col, original_byte = flip_result
+    logger.info(
+        "kv_canary perturb real_kv_unused_cache: group=%s source_idx=%d slot=%d row=%d col=%d "
+        "original_byte=0x%02X new_byte=0x%02X",
+        group.kind.name,
+        source_pick,
+        slot,
+        row,
+        col,
+        original_byte,
+        original_byte ^ 0xFF,
+    )
+
+
+def _pick_sweep_slot_for_group(
+    *,
+    radix_cache: Optional[BasePrefixCache],
+    group: CanaryBufferGroup,
+    swa_window_size: int,
+) -> Optional[int]:
+    if radix_cache is None:
+        return None
+
+    walk_result = walk_radix_cache_for_canary(
+        radix_cache=radix_cache,
+        unlocked_only=True,
+        swa_resident_only=group.kind is PoolKind.SWA,
+    )
+    slots = [
+        int(raw_slot)
+        for raw_slot in walk_result.slot_indices.detach().to("cpu").tolist()
+        if int(raw_slot) >= 0
+    ]
+    if group.kind is PoolKind.SWA:
+        slots = _translate_full_slots_to_swa_slots(
+            slots=slots,
+            full_to_swa_index_mapping=group.swa_index_lut,
+        )
+    if not slots:
+        return None
+
+    pick = int(torch.randint(0, len(slots), (1,)).item())
+    return slots[pick]
+
+
+def _translate_full_slots_to_swa_slots(
+    *,
+    slots: list[int],
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+) -> list[int]:
+    if full_to_swa_index_mapping is None:
+        return []
+
+    lut = full_to_swa_index_mapping.detach().to("cpu").to(torch.int64)
+    translated: list[int] = []
+    for slot in slots:
+        if slot >= int(lut.shape[0]):
+            continue
+        physical_slot = int(lut[slot].item())
+        if physical_slot >= 0:
+            translated.append(physical_slot)
+    return translated

python/sglang/srt/kv_canary/perturb/real_kv_used.py

@@ -0,0 +1,134 @@
+"""Flip the first byte of a slot currently being used by an active req.
+
+Detection should come from per-forward verify (HEAD/TAIL kernel), NOT from
+sweep. Designed to surface CUDA-graph-idle-class bugs where production reads
+a slot whose KV byte was silently overwritten.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.perturb.config import (
+    PerturbConfig,
+    require_target_group_kind,
+)
+from sglang.srt.kv_canary.perturb.slot_picker import (
+    ReqToTokenEntry,
+    collect_active_slots,
+)
+from sglang.srt.kv_canary.perturb.utils import (
+    WarmupGate,
+    flip_first_byte_in_source,
+    pick_target_group,
+    should_run_perturbation,
+)
+
+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__)
+
+
+def run(
+    *,
+    maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    config: PerturbConfig,
+    req_to_token_pool: ReqToTokenPool,
+    buffer_groups: tuple[CanaryBufferGroup, ...],
+    swa_window_size: int,
+    warmup_gate: WarmupGate,
+) -> None:
+    if not should_run_perturbation(
+        perturb_name="real_kv_used",
+        probability=config.real_kv_used_prob,
+        warmup_gate=warmup_gate,
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+    ):
+        return
+
+    group = pick_target_group(
+        buffer_groups=buffer_groups,
+        target_kind=require_target_group_kind(
+            target_group_kind=config.target_group_kind,
+            perturb_name="real_kv_used",
+        ),
+    )
+    if group is None:
+        logger.info(
+            "kv_canary perturb real_kv_used: skipped because no target group with real_kv_sources_k "
+            "matched target_group_kind=%s",
+            config.target_group_kind,
+        )
+        return
+
+    target = _pick_active_slot_for_group(
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+        req_to_token_pool=req_to_token_pool,
+        group=group,
+        swa_window_size=swa_window_size,
+    )
+    if target is None:
+        logger.info(
+            "kv_canary perturb real_kv_used: skipped because no active slot was found "
+            "for group=%s",
+            group.kind.name,
+        )
+        return
+
+    source_pick = int(torch.randint(0, len(group.real_kv_sources_k), (1,)).item())
+    source = group.real_kv_sources_k[source_pick]
+    flip_result = flip_first_byte_in_source(
+        group=group, source=source, slot_idx=target.value
+    )
+    if flip_result is None:
+        logger.info(
+            "kv_canary perturb real_kv_used: skipped because slot=%d could not be mapped into "
+            "group=%s source_idx=%d",
+            target.value,
+            group.kind.name,
+            source_pick,
+        )
+        return
+    row, col, original_byte = flip_result
+    logger.info(
+        "kv_canary perturb real_kv_used: group=%s source_idx=%d slot=%d row=%d col=%d "
+        "original_byte=0x%02X new_byte=0x%02X",
+        group.kind.name,
+        source_pick,
+        target.value,
+        row,
+        col,
+        original_byte,
+        original_byte ^ 0xFF,
+    )
+
+
+def _pick_active_slot_for_group(
+    *,
+    maybe_inaccurate_forward_batch: ForwardBatch,
+    req_to_token_pool: ReqToTokenPool,
+    group: CanaryBufferGroup,
+    swa_window_size: int,
+) -> Optional[ReqToTokenEntry]:
+    candidates = collect_active_slots(
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+        req_to_token_pool=req_to_token_pool,
+        exclude_out_cache_loc=True,
+    )
+    if group.kind is PoolKind.SWA:
+        candidates = [
+            entry
+            for entry in candidates
+            if entry.position >= max(0, entry.seq_len - swa_window_size)
+        ]
+    if not candidates:
+        return None
+
+    pick = int(torch.randint(0, len(candidates), (1,)).item())
+    return candidates[pick]

python/sglang/srt/kv_canary/perturb/req_to_token.py

@@ -0,0 +1,75 @@
+"""Flip the req_to_token pointer of a currently-active req.
+
+The hook picks a random (req_pool_idx, position, value) from active reqs,
+filtering out entries whose value is 0 (so slot 0 is excluded), and overwrites
+req_to_token[req_pool_idx, position] with another active req's slot id.
+KV bytes are not touched.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.kv_canary.perturb.config import PerturbConfig
+from sglang.srt.kv_canary.perturb.slot_picker import collect_active_slots
+from sglang.srt.kv_canary.perturb.utils import WarmupGate, should_run_perturbation
+
+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__)
+
+
+def run(
+    *,
+    maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    config: PerturbConfig,
+    req_to_token_pool: ReqToTokenPool,
+    warmup_gate: WarmupGate,
+) -> None:
+    if not should_run_perturbation(
+        perturb_name="req_to_token",
+        probability=config.req_to_token_prob,
+        warmup_gate=warmup_gate,
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+    ):
+        return
+
+    entries = collect_active_slots(
+        maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+        req_to_token_pool=req_to_token_pool,
+        exclude_out_cache_loc=True,
+    )
+    entries = [entry for entry in entries if entry.value >= 1]
+    if not entries:
+        logger.info(
+            "kv_canary perturb req_to_token: skipped because no active nonzero slots were found"
+        )
+        return
+
+    pick = int(torch.randint(0, len(entries), (1,)).item())
+    target = entries[pick]
+    replacement_values = [item.value for item in entries if item.value != target.value]
+    if not replacement_values:
+        logger.info(
+            "kv_canary perturb req_to_token: skipped because no replacement slot differs from "
+            "original_slot=%d",
+            target.value,
+        )
+        return
+    replacement_pick = int(torch.randint(0, len(replacement_values), (1,)).item())
+    new_value = replacement_values[replacement_pick]
+
+    req_to_token = req_to_token_pool.req_to_token
+    logger.info(
+        "kv_canary perturb req_to_token: req_pool_idx=%d position=%d original_slot=%d new_slot=%d",
+        target.req_pool_idx,
+        target.position,
+        target.value,
+        new_value,
+    )
+    req_to_token[target.req_pool_idx, target.position] = new_value

python/sglang/srt/kv_canary/perturb/slot_picker.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ReqToTokenEntry:
+    req_pool_idx: int
+    position: int
+    value: int
+    seq_len: int = 0
+
+
+def collect_active_slots(
+    *,
+    maybe_inaccurate_forward_batch: ForwardBatch,
+    req_to_token_pool: ReqToTokenPool,
+    exclude_out_cache_loc: bool = True,
+) -> list[ReqToTokenEntry]:
+    """Collect every (req_pool_idx, position, value) triple for currently-active reqs.
+
+    Excludes slots in ``maybe_inaccurate_forward_batch.out_cache_loc`` when ``exclude_out_cache_loc=True``
+    so a slot the current forward is about to write isn't picked (write race).
+    """
+    req_pool_indices = maybe_inaccurate_forward_batch.req_pool_indices
+    seq_lens = maybe_inaccurate_forward_batch.seq_lens
+    if req_pool_indices is None or seq_lens is None:
+        return []
+
+    req_to_token = req_to_token_pool.req_to_token
+    if not isinstance(req_to_token, torch.Tensor) or req_to_token.numel() == 0:
+        return []
+
+    excluded: set[int] = set()
+    if exclude_out_cache_loc:
+        out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc
+        if out_cache_loc is not None:
+            valid_num_tokens = maybe_inaccurate_forward_batch.num_token_non_padded_cpu
+            if valid_num_tokens is None:
+                valid_num_tokens = int(out_cache_loc.shape[0])
+            excluded = set(
+                int(x)
+                for x in out_cache_loc[:valid_num_tokens].detach().to("cpu").tolist()
+            )
+
+    req_pool_indices_list = req_pool_indices.detach().to("cpu").tolist()
+    seq_lens_list = seq_lens.detach().to("cpu").tolist()
+    rows, cols = int(req_to_token.shape[0]), int(req_to_token.shape[1])
+
+    candidates: list[ReqToTokenEntry] = []
+    for req_pool_idx, seq_len in zip(req_pool_indices_list, seq_lens_list):
+        req_pool_idx_int = int(req_pool_idx)
+        seq_len_int = int(seq_len)
+        if req_pool_idx_int < 0 or req_pool_idx_int >= rows:
+            continue
+        upper = min(seq_len_int, cols)
+        if upper <= 0:
+            continue
+        row_values = req_to_token[req_pool_idx_int, :upper].detach().to("cpu").tolist()
+        candidates.extend(
+            ReqToTokenEntry(
+                req_pool_idx=req_pool_idx_int,
+                position=pos,
+                value=value,
+                seq_len=seq_len_int,
+            )
+            for pos, raw_value in enumerate(row_values)
+            if (value := int(raw_value)) >= 0 and value not in excluded
+        )
+    return candidates
+
+
+def pick_out_cache_loc_slot(
+    *, maybe_inaccurate_forward_batch: ForwardBatch
+) -> Optional[int]:
+    out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc
+    if out_cache_loc is None:
+        return None
+    total = int(out_cache_loc.shape[0])
+    if total <= 0:
+        return None
+    valid_num_tokens = maybe_inaccurate_forward_batch.num_token_non_padded_cpu
+    if valid_num_tokens is None:
+        valid_num_tokens = total
+    valid_num_tokens = int(valid_num_tokens)
+    if valid_num_tokens <= 0:
+        return None
+    pick = random.randrange(valid_num_tokens)
+    slot = int(out_cache_loc[pick].item())
+    if slot < 0:
+        return None
+    return slot

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

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import logging
+import random
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import RealKvSource
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.perturb.config import PerturbConfig, TargetGroupKind
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class WarmupGate:
+    """Per-hook warmup window check + once-per-lifetime disable/enable log emission.
+
+    Shared across the four perturb-point hooks so warmup state is decided in one place
+    rather than duplicated per hook.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: PerturbConfig,
+        outer_step_counter_getter: Callable[[], int],
+    ) -> None:
+        self._config = config
+        self._outer_step_counter_getter = outer_step_counter_getter
+        self._warmup_disable_logged: bool = False
+        self._warmup_enable_logged: bool = False
+
+    def is_in_warmup(self) -> bool:
+        step = self._outer_step_counter_getter()
+        warmup_steps = self._config.warmup_steps
+
+        if step < warmup_steps:
+            self._log_warmup_disabled_once(warmup_steps)
+            return True
+
+        self._log_warmup_enabled_once(step)
+        return False
+
+    def _log_warmup_disabled_once(self, warmup_steps: int) -> None:
+        if self._warmup_disable_logged:
+            return
+
+        logger.info(
+            "kv_canary perturb: disabled during warmup window "
+            "(first %d forward steps)",
+            warmup_steps,
+        )
+        self._warmup_disable_logged = True
+
+    def _log_warmup_enabled_once(self, step: int) -> None:
+        if self._warmup_enable_logged:
+            return
+
+        logger.info("kv_canary perturb: enabled after warmup window at step=%d", step)
+        self._warmup_enable_logged = True
+
+
+def should_run_perturbation(
+    *,
+    perturb_name: str,
+    probability: float,
+    warmup_gate: WarmupGate,
+    maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    require_forward_batch: bool = True,
+) -> bool:
+    if probability <= 0.0:
+        return False
+    if warmup_gate.is_in_warmup():
+        return False
+    if require_forward_batch and maybe_inaccurate_forward_batch is None:
+        logger.info(
+            "kv_canary perturb %s: skipped because maybe_inaccurate_forward_batch is unavailable",
+            perturb_name,
+        )
+        return False
+    return torch.rand((), device="cpu").item() < probability
+
+
+def pick_target_group(
+    *,
+    buffer_groups: tuple[CanaryBufferGroup, ...],
+    target_kind: TargetGroupKind,
+) -> Optional[CanaryBufferGroup]:
+    """Filter buffer_groups by target_kind restricted to groups with non-empty real_kv_sources_k.
+
+    Returns None if no group matches.
+    """
+    eligible = [group for group in buffer_groups if group.real_kv_sources_k]
+    if not eligible:
+        return None
+    if target_kind == TargetGroupKind.FULL:
+        want = PoolKind.FULL
+    elif target_kind == TargetGroupKind.SWA:
+        want = PoolKind.SWA
+    else:
+        raise ValueError(f"Unsupported target_group_kind: {target_kind!r}")
+    filtered = [group for group in eligible if group.kind == want]
+    if not filtered:
+        return None
+    pick = random.randrange(len(filtered))
+    return filtered[pick]
+
+
+def flip_random_source_byte_and_log(
+    *,
+    perturb_name: str,
+    group: CanaryBufferGroup,
+    slot_idx: int,
+) -> None:
+    """Pick a random K-half real_kv source on group, flip byte 0 of slot_idx's tile
+    in it, and log the result. Logs and returns silently when the group has no
+    real_kv_sources_k or the slot cannot be mapped into the chosen source."""
+    if not group.real_kv_sources_k:
+        logger.info(
+            "kv_canary perturb %s: skipped because group=%s has no real_kv_sources_k",
+            perturb_name,
+            group.kind.name,
+        )
+        return
+    source_pick = random.randrange(len(group.real_kv_sources_k))
+    source = group.real_kv_sources_k[source_pick]
+    flip_result = flip_first_byte_in_source(
+        group=group, source=source, slot_idx=slot_idx
+    )
+    if flip_result is None:
+        logger.info(
+            "kv_canary perturb %s: skipped because slot=%d could not be mapped "
+            "into group=%s source_idx=%d",
+            perturb_name,
+            slot_idx,
+            group.kind.name,
+            source_pick,
+        )
+        return
+    row, col, original_byte = flip_result
+    logger.info(
+        "kv_canary perturb %s: group=%s source_idx=%d slot=%d row=%d col=%d "
+        "original_byte=0x%02X new_byte=0x%02X",
+        perturb_name,
+        group.kind.name,
+        source_pick,
+        slot_idx,
+        row,
+        col,
+        original_byte,
+        original_byte ^ 0xFF,
+    )
+
+
+def flip_first_byte_in_source(
+    *,
+    group: CanaryBufferGroup,
+    source: RealKvSource,
+    slot_idx: int,
+    slot_is_physical: bool = False,
+) -> Optional[tuple[int, int, int]]:
+    """XOR 0xFF on byte 0 of slot_idx's tile in source.tensor (column
+    `(physical_slot % page_size) * num_bytes_per_token`, not row offset 0).
+
+    For SWA groups, slot_idx is translated through group.swa_index_lut before computing
+    (row, col). Returns (row, col, original_byte) for logging, or None if the slot is
+    out-of-range / source is degenerate.
+    """
+    if source.num_bytes_per_token <= 0 or source.read_bytes <= 0:
+        return None
+
+    physical_slot = slot_idx
+    if (
+        not slot_is_physical
+        and group.kind == PoolKind.SWA
+        and group.swa_index_lut is not None
+    ):
+        lut = group.swa_index_lut
+        if slot_idx < 0 or slot_idx >= int(lut.shape[0]):
+            return None
+        physical_slot = int(lut[slot_idx].detach().to("cpu").item())
+        if physical_slot < 0:
+            return None
+
+    page_size = max(1, source.page_size)
+    row = physical_slot // page_size
+    col = (physical_slot % page_size) * source.num_bytes_per_token
+    if row < 0 or row >= int(source.tensor.shape[0]):
+        return None
+    if col < 0 or col >= int(source.tensor.shape[1]):
+        return None
+
+    flat = source.tensor
+    original_byte = int(flat[row, col].item())
+    flat[row, col] = original_byte ^ 0xFF
+    return row, col, original_byte

python/sglang/srt/kv_canary/plan_input.py

@@ -0,0 +1,140 @@
+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.
+            Set only when ``CanaryConfig.enable_verify_token_assert`` is on.
+
+    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,
+        )
+
+        req_all_ids_lens = forward_batch.req_all_ids_lens
+        if req_all_ids_lens is not None:
+            self.req_to_verify_expected_tokens_valid_lens[:bs].copy_(
+                req_all_ids_lens.to(torch.int64), non_blocking=True
+            )
+
+
+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_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: EagleDraftWorkerBase.prepare_for_draft_extend 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/dsv4.py

@@ -0,0 +1,53 @@
+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_dsv4(
+    *,
+    pool: object,
+    device: torch.device,
+    read_bytes: int,
+    kv_token_id_vs_position_offset: int,
+) -> tuple[CanaryBufferGroup, ...]:
+    """Attach canary buffers to a DeepSeekV4TokenToKVPool.
+
+    TODO: only the swa_kv_pool sub-pool is wired; c4_kv_pool / c128_kv_pool /
+    c4_indexer_kv_pool / compress state pools are left uncovered.
+    TODO: even on swa_kv_pool, real-KV fingerprint is disabled (read_bytes is
+    ignored). DSV4 stores 584 B/token which is not 16-aligned (584 % 16 == 8),
+    so num_bytes_per_token cannot satisfy the 128-bit load alignment precondition.
+    """
+    del read_bytes
+
+    sub_pool = pool.swa_kv_pool
+    num_slots = int(sub_pool.size)
+
+    k_head = alloc_canary_buf(num_slots=num_slots, device=device)
+    k_tail = alloc_canary_buf(num_slots=num_slots, device=device)
+
+    group = CanaryBufferGroup(
+        kind=PoolKind.SWA,
+        k_head=k_head,
+        k_tail=k_tail,
+        v_head=None,
+        v_tail=None,
+        real_kv_sources_k=(),
+        real_kv_sources_v=(),
+        swa_index_lut=pool.full_to_swa_index_mapping,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+
+    patch_buf_info_method(
+        pool,
+        method_name="get_state_buf_infos",
+        group=group,
+        has_v_half=False,
+        page_size=sub_pool.page_size,
+    )
+
+    return (group,)

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

@@ -0,0 +1,48 @@
+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,
+    make_row_source,
+)
+
+
+def attach_mha(
+    *,
+    pool: object,
+    device: torch.device,
+    read_bytes: int,
+    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,
+        real_kv_sources_k=make_row_source(
+            layer_buffer=pool.k_buffer[0], read_bytes=read_bytes
+        ),
+        real_kv_sources_v=make_row_source(
+            layer_buffer=pool.v_buffer[0], read_bytes=read_bytes
+        ),
+        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/adapters/swa.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import Optional
+
+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,
+    make_row_source,
+)
+
+
+def attach_swa(
+    *,
+    pool: object,
+    device: torch.device,
+    read_bytes: int,
+    kv_token_id_vs_position_offset: int,
+) -> tuple[CanaryBufferGroup, ...]:
+    full_group = _build_subpool_group(
+        sub_pool=pool.full_kv_pool,
+        kind=PoolKind.FULL,
+        device=device,
+        read_bytes=read_bytes,
+        swa_lut=None,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    swa_group = _build_subpool_group(
+        sub_pool=pool.swa_kv_pool,
+        kind=PoolKind.SWA,
+        device=device,
+        read_bytes=read_bytes,
+        swa_lut=pool.full_to_swa_index_mapping,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+
+    patch_buf_info_method(
+        pool,
+        method_name="get_contiguous_buf_infos",
+        group=full_group,
+        has_v_half=True,
+        page_size=pool.page_size,
+    )
+    patch_buf_info_method(
+        pool,
+        method_name="get_state_buf_infos",
+        group=swa_group,
+        has_v_half=True,
+        page_size=pool.page_size,
+    )
+    return (full_group, swa_group)
+
+
+def _build_subpool_group(
+    *,
+    sub_pool: object,
+    kind: PoolKind,
+    device: torch.device,
+    read_bytes: int,
+    swa_lut: Optional[torch.Tensor],
+    kv_token_id_vs_position_offset: int,
+) -> CanaryBufferGroup:
+    num_slots = int(sub_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)
+    return CanaryBufferGroup(
+        kind=kind,
+        k_head=k_head,
+        k_tail=k_tail,
+        v_head=v_head,
+        v_tail=v_tail,
+        real_kv_sources_k=make_row_source(
+            layer_buffer=sub_pool.k_buffer[0], read_bytes=read_bytes
+        ),
+        real_kv_sources_v=make_row_source(
+            layer_buffer=sub_pool.v_buffer[0], read_bytes=read_bytes
+        ),
+        swa_index_lut=swa_lut,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )

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

@@ -0,0 +1,74 @@
+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.dsv4 import attach_dsv4
+from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha
+from sglang.srt.kv_canary.pool_patcher.adapters.swa import attach_swa
+from sglang.srt.kv_canary.pool_patcher.buffer_alloc import resolve_real_kv_read_bytes
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.mem_cache.memory_pool import (
+    KVCache,
+    MHATokenToKVPool,
+    MHATokenToKVPoolFP4,
+)
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+
+logger = logging.getLogger(__name__)
+
+PoolAttacher = Callable[..., tuple[CanaryBufferGroup, ...]]
+
+_POOL_ATTACHERS: Dict[Type, PoolAttacher] = {
+    MHATokenToKVPool: attach_mha,
+    MHATokenToKVPoolFP4: attach_mha,
+    SWAKVPool: attach_swa,
+    DeepSeekV4TokenToKVPool: attach_dsv4,
+}
+
+
+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)}"
+        )
+
+    read_bytes = resolve_real_kv_read_bytes(config)
+    groups = attacher(
+        pool=pool,
+        device=device,
+        read_bytes=read_bytes,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+    logger.info(
+        "attach_canary_buffers: pool=%s attacher=%s read_bytes=%d n_groups=%d kinds=%s "
+        "kv_token_id_vs_position_offset=%d",
+        type(pool).__name__,
+        attacher.__name__,
+        read_bytes,
+        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,115 @@
+from __future__ import annotations
+
+import sys
+from typing import Tuple
+
+import torch
+
+from sglang.jit_kernel.kv_canary.consts import RealKvHashMode
+from sglang.jit_kernel.kv_canary.verify import (
+    CANARY_SLOT_BYTES,
+    RealKvSource,
+)
+from sglang.srt.kv_canary.config import CanaryConfig
+
+_PARTIAL_REAL_KV_READ_BYTES = 16
+_REAL_KV_READ_ALIGN = 16
+
+
+def resolve_real_kv_read_bytes(config: CanaryConfig) -> int:
+    if config.real_kv_hash_mode is RealKvHashMode.NONE:
+        return 0
+    if config.real_kv_hash_mode is RealKvHashMode.ALL:
+        return sys.maxsize
+    return _PARTIAL_REAL_KV_READ_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)
+
+
+def _clip_read_bytes_aligned(*, requested: int, num_bytes_per_token: int) -> int:
+    """Validate and clip read_bytes for the CUDA fold kernel's 128-bit aligned loads.
+
+    Normalizes sentinels (``sys.maxsize`` -> ``num_bytes_per_token``, ``0`` -> ``0``) and
+    rejects negative / unaligned / oversized requests.
+    """
+    if num_bytes_per_token <= 0 or num_bytes_per_token % _REAL_KV_READ_ALIGN != 0:
+        raise ValueError(
+            "kv-canary: num_bytes_per_token must be a positive multiple of "
+            f"{_REAL_KV_READ_ALIGN}, got {num_bytes_per_token}"
+        )
+    if requested == 0:
+        return 0
+    if requested == sys.maxsize:
+        return num_bytes_per_token
+    if requested < 0:
+        raise ValueError(f"kv-canary: read_bytes must be non-negative, got {requested}")
+    if requested > num_bytes_per_token:
+        raise ValueError(
+            "kv-canary: read_bytes must be <= num_bytes_per_token "
+            f"({num_bytes_per_token}), got {requested}"
+        )
+    if requested % _REAL_KV_READ_ALIGN != 0:
+        raise ValueError(
+            "kv-canary: read_bytes must be a multiple of "
+            f"{_REAL_KV_READ_ALIGN}, got {requested}"
+        )
+    return requested
+
+
+def make_row_source(
+    *,
+    layer_buffer: torch.Tensor,
+    read_bytes: int,
+) -> Tuple[RealKvSource, ...]:
+    contiguous = layer_buffer.contiguous()
+    num_slots = int(contiguous.shape[0])
+    if num_slots == 0 or read_bytes == 0:
+        return ()
+    flat = contiguous.view(torch.uint8).reshape(num_slots, -1)
+    num_bytes_per_token = int(flat.shape[1])
+    clipped = _clip_read_bytes_aligned(
+        requested=read_bytes, num_bytes_per_token=num_bytes_per_token
+    )
+    if clipped == 0:
+        return ()
+    return (
+        RealKvSource(
+            tensor=flat,
+            page_size=1,
+            num_bytes_per_token=num_bytes_per_token,
+            read_bytes=clipped,
+        ),
+    )
+
+
+def make_packed_source(
+    *,
+    page_buffer: torch.Tensor,
+    page_size: int,
+    bytes_per_token: int,
+    read_bytes: int,
+) -> Tuple[RealKvSource, ...]:
+    if read_bytes == 0 or page_buffer.numel() == 0:
+        return ()
+    flat = page_buffer.contiguous().view(torch.uint8)
+    if flat.ndim == 1:
+        flat = flat.reshape(1, -1)
+    clipped = _clip_read_bytes_aligned(
+        requested=read_bytes, num_bytes_per_token=bytes_per_token
+    )
+    if clipped == 0:
+        return ()
+    return (
+        RealKvSource(
+            tensor=flat,
+            page_size=page_size,
+            num_bytes_per_token=bytes_per_token,
+            read_bytes=clipped,
+        ),
+    )

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/radix_cache_walker.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from sglang.srt.mem_cache.radix_cache import RadixCache
+from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
+from sglang.srt.mem_cache.unified_cache_components import (
+    BASE_COMPONENT_TYPE,
+    ComponentType,
+)
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.mem_cache.radix_cache import TreeNode
+    from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class RadixCacheWalkResult:
+    slot_indices: torch.Tensor
+    positions: torch.Tensor
+    prev_slot_indices: torch.Tensor
+
+
+def walk_radix_cache_for_canary(
+    *,
+    radix_cache: BasePrefixCache,
+    unlocked_only: bool = False,
+    swa_resident_only: bool = False,
+) -> RadixCacheWalkResult:
+    """Walk the radix tree and emit flat (slot_indices, positions, prev_slot_indices) tensors.
+
+    With both flags False (default), emits every slot held by the radix cache (including slots
+    also referenced by a currently-running req — that overlap is harmless redundancy with the
+    per-forward HEAD/TAIL path). ``unlocked_only=True`` skips nodes still locked by a running
+    req. ``swa_resident_only=True`` skips SWA-tombstoned nodes (slots evicted from the SWA
+    window)."""
+    cache_type = type(radix_cache)
+    if (
+        cache_type is not RadixCache
+        and cache_type is not SWARadixCache
+        and cache_type is not UnifiedRadixCache
+    ):
+        raise NotImplementedError(
+            f"walk_radix_cache_for_canary does not support {cache_type.__name__}"
+        )
+
+    slot_buf: list[int] = []
+    position_buf: list[int] = []
+    prev_slot_buf: list[int] = []
+
+    _walk_radix_subtree(
+        node=radix_cache.root_node,
+        radix_cache=radix_cache,
+        depth=0,
+        parent_last_slot=-1,
+        slot_buf=slot_buf,
+        position_buf=position_buf,
+        prev_slot_buf=prev_slot_buf,
+        is_root=True,
+        unlocked_only=unlocked_only,
+        swa_resident_only=swa_resident_only,
+    )
+
+    slot_tensor = torch.tensor(slot_buf, dtype=torch.int64)
+    position_tensor = torch.tensor(position_buf, dtype=torch.int64)
+    prev_slot_tensor = torch.tensor(prev_slot_buf, dtype=torch.int64)
+    return RadixCacheWalkResult(
+        slot_indices=slot_tensor,
+        positions=position_tensor,
+        prev_slot_indices=prev_slot_tensor,
+    )
+
+
+def _walk_radix_subtree(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+    depth: int,
+    parent_last_slot: int,
+    slot_buf: list[int],
+    position_buf: list[int],
+    prev_slot_buf: list[int],
+    is_root: bool,
+    unlocked_only: bool,
+    swa_resident_only: bool,
+) -> None:
+    node_slots = _node_slots_for_canary(node=node, radix_cache=radix_cache)
+
+    if unlocked_only:
+        emit_slots = not is_root and _node_is_unlocked_for_canary(
+            node=node, radix_cache=radix_cache
+        )
+    else:
+        emit_slots = not is_root
+    if swa_resident_only:
+        emit_slots = emit_slots and _node_is_swa_resident_for_canary(
+            node=node,
+            radix_cache=radix_cache,
+        )
+
+    chain_last_slot = parent_last_slot
+    for j, slot in enumerate(node_slots):
+        prev = parent_last_slot if j == 0 else node_slots[j - 1]
+        if emit_slots:
+            slot_buf.append(slot)
+            position_buf.append(depth + j)
+            prev_slot_buf.append(prev)
+        chain_last_slot = slot
+
+    child_depth = depth + _node_len_for_canary(
+        node=node,
+        radix_cache=radix_cache,
+        node_slots=node_slots,
+        is_root=is_root,
+    )
+    for child in node.children.values():
+        _walk_radix_subtree(
+            node=child,
+            radix_cache=radix_cache,
+            depth=child_depth,
+            parent_last_slot=chain_last_slot,
+            slot_buf=slot_buf,
+            position_buf=position_buf,
+            prev_slot_buf=prev_slot_buf,
+            is_root=False,
+            unlocked_only=unlocked_only,
+            swa_resident_only=swa_resident_only,
+        )
+
+
+def _node_slots_for_canary(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+) -> list[int]:
+    value: Any
+    if type(radix_cache) is UnifiedRadixCache:
+        value = node.component_data[BASE_COMPONENT_TYPE].value
+    else:
+        value = node.value
+
+    if isinstance(value, torch.Tensor):
+        return [int(s) for s in value.tolist()]
+    return []
+
+
+def _node_len_for_canary(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+    node_slots: list[int],
+    is_root: bool,
+) -> int:
+    if type(radix_cache) is not UnifiedRadixCache:
+        return len(node_slots)
+
+    if is_root or node.key is None:
+        return len(node_slots)
+    return len(node.key)
+
+
+def _node_is_unlocked_for_canary(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+) -> bool:
+    if type(radix_cache) is RadixCache:
+        return node.lock_ref == 0
+
+    if type(radix_cache) is SWARadixCache:
+        return node.full_lock_ref == 0
+
+    if type(radix_cache) is UnifiedRadixCache:
+        return node.component_data[BASE_COMPONENT_TYPE].lock_ref == 0
+
+    raise NotImplementedError(
+        f"walk_radix_cache_for_canary does not support {type(radix_cache).__name__}"
+    )
+
+
+def _node_is_swa_resident_for_canary(
+    *,
+    node: TreeNode | UnifiedTreeNode,
+    radix_cache: BasePrefixCache,
+) -> bool:
+    if type(radix_cache) is SWARadixCache:
+        return not node.swa_tombstone
+
+    if type(radix_cache) is UnifiedRadixCache:
+        if not radix_cache.supports_swa():
+            return True
+        return node.component_data[ComponentType.SWA].value is not None
+
+    return True

python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
+    launch_scatter_req_token_ids_kernel,
+)
+from sglang.srt.utils.common import flatten_arrays_to_int64_tensor
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def compute_req_all_ids_info(
+    reqs: list[Req],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Snapshot per-req (origin_input_ids + output_ids) as pinned CPU int64 tensors.
+
+    Returns:
+        ``(req_all_ids_flat, req_all_ids_lens)`` — both pinned CPU int64. ``flat`` is the
+        flattened ``cat(r.origin_input_ids, r.output_ids) for r in reqs``; ``lens`` is
+        per-req ``len(origin_input_ids) + len(output_ids)``.
+    """
+    parts = [arr for req in reqs for arr in (req.origin_input_ids, req.output_ids)]
+    req_all_ids_flat = flatten_arrays_to_int64_tensor(
+        parts, device=torch.device("cpu"), pin=True
+    )
+    req_all_ids_lens = torch.tensor(
+        [len(req.origin_input_ids) + len(req.output_ids) for req in reqs],
+        dtype=torch.int64,
+        pin_memory=True,
+    )
+    return req_all_ids_flat, req_all_ids_lens
+
+
+def populate_req_to_expected_token_ids(
+    *,
+    forward_batch: ForwardBatch,
+    req_to_verify_expected_tokens: Optional[torch.Tensor],
+) -> None:
+    """Scatter the forward batch's per-req token-id snapshot into the device-side pool."""
+    req_all_ids_flat_cpu = forward_batch.req_all_ids_flat
+    req_all_ids_lens_cpu = forward_batch.req_all_ids_lens
+    if req_all_ids_flat_cpu is None or req_all_ids_lens_cpu is None:
+        return
+    if req_to_verify_expected_tokens is None:
+        return
+
+    bs = int(forward_batch.req_pool_indices.shape[0])
+    if bs == 0:
+        return
+    if int(req_all_ids_lens_cpu.shape[0]) != bs:
+        raise RuntimeError(
+            f"kv-canary: req_all_ids_lens length {int(req_all_ids_lens_cpu.shape[0])} != "
+            f"batch_size {bs}; ForwardBatch snapshot diverged"
+        )
+
+    offsets_cpu = torch.zeros(bs + 1, dtype=torch.int64, pin_memory=True)
+    offsets_cpu[1:] = torch.cumsum(req_all_ids_lens_cpu, dim=0)
+    total_tokens = int(offsets_cpu[bs].item())
+    if total_tokens != int(req_all_ids_flat_cpu.shape[0]):
+        raise RuntimeError(
+            f"kv-canary: cumsum(req_all_ids_lens)={total_tokens} != "
+            f"req_all_ids_flat.numel()={int(req_all_ids_flat_cpu.shape[0])}; snapshot inconsistent"
+        )
+    if total_tokens == 0:
+        return
+
+    device = req_to_verify_expected_tokens.device
+    req_all_ids_flat_dev = req_all_ids_flat_cpu.to(device, non_blocking=True)
+    offsets_dev = offsets_cpu.to(device, non_blocking=True)
+    req_pool_indices_dev = forward_batch.req_pool_indices.to(
+        device=device, dtype=torch.int64
+    )
+
+    launch_scatter_req_token_ids_kernel(
+        flat_in=req_all_ids_flat_dev,
+        offsets=offsets_dev,
+        req_pool_indices=req_pool_indices_dev,
+        pool_out=req_to_verify_expected_tokens,
+    )

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

@@ -0,0 +1,295 @@
+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.perturb.config import PerturbConfig
+from sglang.srt.kv_canary.perturb.manager import PerturbManager
+from sglang.srt.kv_canary.runner.health_checker import KernelRunCounterHealthChecker
+from sglang.srt.kv_canary.runner.stats_logger import PeriodicCanaryStatsLogger
+from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceReporter
+from sglang.srt.kv_canary.runner.sweep import SweepOrchestrator
+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
+from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    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,
+        perturb_config: PerturbConfig,
+        buffer_groups: tuple[CanaryBufferGroup, ...],
+        device: torch.device,
+        req_to_token_pool: ReqToTokenPool,
+        launch_capacities: CanaryLaunchCapacities,
+        swa_window_size: int = 0,
+        token_oracle_manager: Optional[TokenOracleManager] = None,
+        swa_allocator: Optional[SWATokenToKVPoolAllocator] = None,
+        speculative_num_steps: int = 1,
+        is_eagle_draft_decode: bool = False,
+    ) -> None:
+        self.config = config
+        self._req_to_token_pool = req_to_token_pool
+        self._swa_window_size = swa_window_size
+        self._swa_allocator: Optional[SWATokenToKVPoolAllocator] = swa_allocator
+        self._outer_step_counter: int = 0
+        self._active_single_forward_manager_index: Optional[int] = None
+        self._model_forward_bracket_depth: int = 0
+
+        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)
+
+        swa_divergence_interval = (
+            envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.get()
+        )
+        if swa_divergence_interval > 0:
+            self._swa_divergence_report: Optional[SwaDivergenceReporter] = (
+                SwaDivergenceReporter(
+                    device=device,
+                    d2h_stream=self._d2h_stream,
+                    interval=swa_divergence_interval,
+                    swa_allocator=self._swa_allocator,
+                    req_to_token_pool=self._req_to_token_pool,
+                )
+            )
+        else:
+            self._swa_divergence_report = None
+
+        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._sweep_orchestrator = SweepOrchestrator(
+            config=config,
+            device_state=self._device_state,
+            buffer_groups=self._buffer_groups,
+            endpoints=self._endpoints,
+            swa_window_size=self._swa_window_size,
+            outer_step_counter_getter=self._get_outer_step_counter,
+        )
+        self._perturb_manager = PerturbManager(
+            config=perturb_config,
+            req_to_token_pool=req_to_token_pool,
+            buffer_groups=self._buffer_groups,
+            outer_step_counter_getter=self._get_outer_step_counter,
+            swa_window_size=self._swa_window_size,
+            sweep_interval=config.sweep_interval,
+        )
+        self._health_checker = KernelRunCounterHealthChecker(
+            config=config,
+            device_state=self._device_state,
+            active_tags=self._active_tags,
+            outer_step_counter_getter=self._get_outer_step_counter,
+            d2h_stream=self._d2h_stream,
+        )
+        self._stats_logger = PeriodicCanaryStatsLogger(
+            config=config,
+            device_state=self._device_state,
+            active_tags=self._active_tags,
+            outer_step_counter_getter=self._get_outer_step_counter,
+            sweep_orchestrator=self._sweep_orchestrator,
+            d2h_stream=self._d2h_stream,
+        )
+
+        num_sfms = max(1, speculative_num_steps - 1)
+        self._single_forward_managers: tuple[SingleForwardManager, ...] = tuple(
+            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,
+                token_oracle_manager=token_oracle_manager,
+                swa_divergence_report=self._swa_divergence_report,
+                is_eagle_draft_decode=is_eagle_draft_decode,
+            )
+            for _ in range(num_sfms)
+        )
+
+    @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
+
+    @contextlib.contextmanager
+    def model_forward_bracket_scope(self) -> Iterator[bool]:
+        """Return whether this is the outermost patched ``model.forward`` call.
+
+        Some model implementations enter another patched forward from inside the
+        top-level forward (for example, a vision-language model calling its inner
+        language model). Kv-canary owns one pre/post bracket per active
+        SingleForwardManager; nested brackets would run a second pre-op while the
+        phase checker is already in the first bracket.
+        """
+        self._model_forward_bracket_depth += 1
+        try:
+            yield self._model_forward_bracket_depth == 1
+        finally:
+            self._model_forward_bracket_depth -= 1
+
+    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
+            )
+        self._perturb_manager.perturb(
+            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._perturb_manager.perturb_post_forward(
+            maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch
+        )
+        self._sweep_orchestrator.maybe_run_sweep()
+        self._outer_step_counter += 1
+        self._violation_manager.step()
+        self._health_checker.step()
+        self._stats_logger.step()
+        if self._swa_divergence_report is not None:
+            self._swa_divergence_report.step(
+                outer_step_counter=self._outer_step_counter,
+                maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            )
+
+    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 attach_radix_cache(self, radix_cache: BasePrefixCache) -> None:
+        self._sweep_orchestrator.attach_radix_cache(radix_cache)
+        self._perturb_manager.attach_radix_cache(radix_cache)
+
+    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/health_checker.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
+from sglang.srt.kv_canary.runner.kernel_launcher import passes_v_half_gate
+from sglang.srt.kv_canary.state import CanaryDeviceState
+
+logger = logging.getLogger(__name__)
+
+_HEALTH_CHECK_EVERY_N_STEPS: int = 100
+_HEALTH_CHECK_WARMUP_STEPS: int = 100
+_SWEEP_TAGS: frozenset[CanaryLaunchTag] = frozenset(
+    (
+        CanaryLaunchTag.SWEEP_K_FULL,
+        CanaryLaunchTag.SWEEP_V_FULL,
+        CanaryLaunchTag.SWEEP_K_SWA,
+        CanaryLaunchTag.SWEEP_V_SWA,
+    )
+)
+
+
+class KernelRunCounterHealthChecker:
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        device_state: CanaryDeviceState,
+        active_tags: tuple[CanaryLaunchTag, ...],
+        outer_step_counter_getter: Callable[[], int],
+        d2h_stream: torch.cuda.Stream,
+    ) -> None:
+        self._config = config
+        self._device_state = device_state
+        self._active_tags = active_tags
+        self._outer_step_counter_getter = outer_step_counter_getter
+        self._handler = DelayedDeviceHostHandler(d2h_stream=d2h_stream)
+        self._prev_counters_host: torch.Tensor = torch.zeros_like(
+            device_state.kernel_run_counters, device="cpu"
+        )
+
+    def step(self) -> None:
+        self._handler.step(
+            compute_on_device=self._compute_on_device,
+            postprocess_on_host=self._postprocess_on_host,
+        )
+
+    def _compute_on_device(self) -> Optional[torch.Tensor]:
+        outer_step_counter = self._outer_step_counter_getter()
+        if outer_step_counter < _HEALTH_CHECK_WARMUP_STEPS:
+            return None
+        if outer_step_counter % _HEALTH_CHECK_EVERY_N_STEPS != 0:
+            return None
+        if not self._active_tags:
+            return None
+        return self._device_state.kernel_run_counters
+
+    def _postprocess_on_host(self, new_counter_host: torch.Tensor) -> None:
+        delta = new_counter_host - self._prev_counters_host
+        self._prev_counters_host = new_counter_host
+        expected_tags = self._expected_active_tags_for_health_check()
+        stalled = [tag for tag in expected_tags if int(delta[tag.value]) == 0]
+        if stalled:
+            names = ", ".join(tag.name for tag in stalled)
+            raise RuntimeError(
+                f"kv-canary: kernel_run_counter did not increase since previous check "
+                f"for tags=[{names}] at step={self._outer_step_counter_getter()}; "
+                f"canary path is not executing"
+            )
+
+    def _expected_active_tags_for_health_check(self) -> tuple[CanaryLaunchTag, ...]:
+        tags = self._active_tags
+        if self._config.sweep_interval <= 0:
+            tags = tuple(tag for tag in tags if tag not in _SWEEP_TAGS)
+        return tuple(tag for tag in tags if passes_v_half_gate(tag))

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

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.consts import (
+    RealKvHashMode,
+)
+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,
+    real_kv_hash_mode: RealKvHashMode,
+    enable_write_input_assert: bool,
+    enable_verify_token_assert: bool,
+) -> 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 not _is_sweep_tag(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,
+            real_kv_hash_mode=real_kv_hash_mode,
+        )
+
+
+def launch_endpoints_sweep(
+    *,
+    endpoints: tuple[CanaryEndpoint, ...],
+    group: CanaryBufferGroup,
+    verify_plan: VerifyPlan,
+    violation_log: ViolationLog,
+    real_kv_hash_mode: RealKvHashMode,
+) -> None:
+    active_endpoints = [
+        endpoint
+        for endpoint in endpoints
+        if _endpoint_belongs_to_group(endpoint, group)
+        and _is_sweep_tag(endpoint.kernel_kind)
+        and passes_v_half_gate(endpoint.kernel_kind)
+    ]
+    assert len(active_endpoints) > 0
+
+    for endpoint in active_endpoints:
+        endpoint.launch_sweep(
+            verify_plan=verify_plan,
+            violation_log=violation_log,
+            real_kv_hash_mode=real_kv_hash_mode,
+        )
+
+
+def _is_sweep_tag(tag: CanaryLaunchTag) -> bool:
+    return tag in (
+        CanaryLaunchTag.SWEEP_K_FULL,
+        CanaryLaunchTag.SWEEP_V_FULL,
+        CanaryLaunchTag.SWEEP_K_SWA,
+        CanaryLaunchTag.SWEEP_V_SWA,
+    )
+
+
+def _is_v_half_tag(tag: CanaryLaunchTag) -> bool:
+    return tag in (
+        CanaryLaunchTag.HEAD_V_FULL,
+        CanaryLaunchTag.TAIL_V_FULL,
+        CanaryLaunchTag.SWEEP_V_FULL,
+        CanaryLaunchTag.HEAD_V_SWA,
+        CanaryLaunchTag.TAIL_V_SWA,
+        CanaryLaunchTag.SWEEP_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/stats_logger.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
+from sglang.srt.kv_canary.runner.sweep import SweepOrchestrator
+from sglang.srt.kv_canary.state import CanaryDeviceState
+
+logger = logging.getLogger(__name__)
+
+
+class PeriodicCanaryStatsLogger:
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        device_state: CanaryDeviceState,
+        active_tags: tuple[CanaryLaunchTag, ...],
+        outer_step_counter_getter: Callable[[], int],
+        sweep_orchestrator: SweepOrchestrator,
+        d2h_stream: torch.cuda.Stream,
+    ) -> None:
+        self._config = config
+        self._device_state = device_state
+        self._active_tags = active_tags
+        self._outer_step_counter_getter = outer_step_counter_getter
+        self._sweep_orchestrator = sweep_orchestrator
+        self._handler = DelayedDeviceHostHandler(d2h_stream=d2h_stream)
+
+    def step(self) -> None:
+        self._handler.step(
+            compute_on_device=self._compute_on_device,
+            postprocess_on_host=self._postprocess_on_host,
+        )
+
+    def _compute_on_device(self) -> Optional[dict[str, Any]]:
+        period = self._config.stats_print_every_n_steps
+        if period <= 0:
+            return None
+        outer_step_counter = self._outer_step_counter_getter()
+        if outer_step_counter == 0 or outer_step_counter % period != 0:
+            return None
+        device_state = self._device_state
+        return {
+            "step": outer_step_counter,
+            "slot_sum": device_state.slot_run_counters.sum().view(1),
+            "write_index": device_state.violation_log.violation_write_index,
+        }
+
+    def _postprocess_on_host(self, host_data: dict[str, Any]) -> None:
+        logger.info(
+            "[canary] step=%d protected_tokens=%d sweep_passes=%d violations=%d "
+            "launch_tags_active=%d/%d",
+            int(host_data["step"]),
+            int(host_data["slot_sum"].item()),
+            self._sweep_orchestrator.sweep_passes,
+            int(host_data["write_index"].item()),
+            len(self._active_tags),
+            len(CanaryLaunchTag),
+        )

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

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import json
+import logging
+import re
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+_SWA_DIVERGENCE_LOG_PREFIX: str = "kv_canary_swa_divergence="
+_SWA_DIVERGENCE_LINE_RE = re.compile(re.escape(_SWA_DIVERGENCE_LOG_PREFIX) + r"(\S+)")
+_FULL_IDX = 0
+_SWA_IDX = 1
+
+
+class SwaDivergenceReporter:
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        d2h_stream: torch.cuda.Stream,
+        interval: int,
+        swa_allocator: Optional[SWATokenToKVPoolAllocator] = None,
+        req_to_token_pool: Optional[ReqToTokenPool] = None,
+    ) -> None:
+        self._interval = interval
+        self._swa_allocator = swa_allocator
+        self._req_to_token_pool = req_to_token_pool
+        self._forward_ct: int = 0
+        # Per-group running total of verify entries (shape ``[2]``, int32).
+        self.verify_total_count_device: torch.Tensor = torch.zeros(
+            2, dtype=torch.int32, device=device
+        )
+        self._handler = DelayedDeviceHostHandler(d2h_stream=d2h_stream)
+
+    def observe_after_invoke_plan(
+        self, *, group: CanaryBufferGroup, verify_plan: VerifyPlan
+    ) -> None:
+        idx = _FULL_IDX if group.kind is PoolKind.FULL else _SWA_IDX
+        # verify_num_valid is shape [1]; slice to a length-1 view so the in-place add
+        # has matching ranks (else torch refuses the broadcast into shape []).
+        self.verify_total_count_device[idx : idx + 1].add_(verify_plan.verify_num_valid)
+
+    def step(
+        self,
+        *,
+        outer_step_counter: int,
+        maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    ) -> None:
+        self._forward_ct += 1
+        self._handler.step(
+            compute_on_device=lambda: self._compute_on_device(
+                outer_step_counter=outer_step_counter,
+                maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            ),
+            postprocess_on_host=self._postprocess_on_host,
+        )
+
+    def _compute_on_device(
+        self,
+        *,
+        outer_step_counter: int,
+        maybe_inaccurate_forward_batch: Optional[ForwardBatch],
+    ) -> Optional[dict[str, Any]]:
+        if outer_step_counter == 0 or outer_step_counter % self._interval != 0:
+            return None
+
+        result: dict[str, Any] = {
+            "forward_ct": self._forward_ct,
+            "verify_total_count": self.verify_total_count_device,
+        }
+        if (
+            self._swa_allocator is not None
+            and maybe_inaccurate_forward_batch is not None
+        ):
+            result["swa_full_idx_divergence"] = compute_swa_full_idx_divergence(
+                swa_allocator=self._swa_allocator,
+                req_to_token_pool=self._req_to_token_pool,
+                maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            )
+            result["swa_out_of_window_tokens"] = compute_swa_out_of_window_tokens(
+                swa_allocator=self._swa_allocator,
+                req_to_token_pool=self._req_to_token_pool,
+                maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
+            )
+        return result
+
+    def _postprocess_on_host(self, host_data: dict[str, Any]) -> None:
+        verify_totals = host_data["verify_total_count"].tolist()
+        swa_full_idx_divergence = (
+            int(x.item())
+            if (x := host_data.get("swa_full_idx_divergence")) is not None
+            else 0
+        )
+        swa_out_of_window_tokens = (
+            int(x.item())
+            if (x := host_data.get("swa_out_of_window_tokens")) is not None
+            else 0
+        )
+        logger.info(
+            SwaDivergenceLog(
+                forward_ct=host_data["forward_ct"],
+                verify_full=int(verify_totals[_FULL_IDX]),
+                verify_swa=int(verify_totals[_SWA_IDX]),
+                swa_full_idx_divergence=swa_full_idx_divergence,
+                swa_out_of_window_tokens=swa_out_of_window_tokens,
+            ).format()
+        )
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class SwaDivergenceLog:
+    forward_ct: int
+    verify_full: int
+    verify_swa: int
+    swa_full_idx_divergence: int
+    swa_out_of_window_tokens: int = 0
+
+    def format(self) -> str:
+        return _SWA_DIVERGENCE_LOG_PREFIX + json.dumps(
+            asdict(self), separators=(",", ":")
+        )
+
+    @classmethod
+    def parse(cls, line: str) -> Optional[SwaDivergenceLog]:
+        match = _SWA_DIVERGENCE_LINE_RE.search(line)
+        if match is None:
+            return None
+        return cls(**json.loads(match.group(1)))
+
+    @classmethod
+    def find_last(cls, text: str) -> Optional[tuple[SwaDivergenceLog, str]]:
+        last_match: Optional[re.Match] = None
+        for match in _SWA_DIVERGENCE_LINE_RE.finditer(text):
+            last_match = match
+        if last_match is None:
+            return None
+        return cls(**json.loads(last_match.group(1))), last_match.group(0)
+
+    @classmethod
+    def find_all(cls, text: str) -> list[tuple[SwaDivergenceLog, str]]:
+        return [
+            (cls(**json.loads(match.group(1))), match.group(0))
+            for match in _SWA_DIVERGENCE_LINE_RE.finditer(text)
+        ]
+
+
+def compute_swa_out_of_window_tokens(
+    *,
+    swa_allocator: SWATokenToKVPoolAllocator,
+    req_to_token_pool: ReqToTokenPool,
+    maybe_inaccurate_forward_batch: ForwardBatch,
+) -> torch.Tensor:
+    """Count tokens in the live req_to_token range whose SWA mapping is 0 (out-of-window)."""
+    full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping
+    device = full_to_swa_index_mapping.device
+    req_pool_indices = maybe_inaccurate_forward_batch.req_pool_indices
+    seq_lens = maybe_inaccurate_forward_batch.seq_lens
+    if req_pool_indices.numel() == 0:
+        return torch.zeros(1, dtype=torch.int32, device=device)
+    req_to_token = req_to_token_pool.req_to_token
+    rows = req_to_token[req_pool_indices]
+    positions = torch.arange(rows.shape[1], device=rows.device)
+    mask = positions[None, :] < seq_lens[:, None]
+    swa_indices = full_to_swa_index_mapping[rows]
+    return ((swa_indices == 0) & mask).sum().to(torch.int32).view(1)
+
+
+def compute_swa_full_idx_divergence(
+    *,
+    swa_allocator: SWATokenToKVPoolAllocator,
+    req_to_token_pool: ReqToTokenPool,
+    maybe_inaccurate_forward_batch: ForwardBatch,
+) -> torch.Tensor:
+    """Count non-identity (full, swa) index pairs in the live req_to_token range."""
+    full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping
+    device = full_to_swa_index_mapping.device
+    req_pool_indices = maybe_inaccurate_forward_batch.req_pool_indices
+    seq_lens = maybe_inaccurate_forward_batch.seq_lens
+
+    if req_pool_indices.numel() == 0:
+        return torch.zeros(1, dtype=torch.int32, device=device)
+
+    req_to_token = req_to_token_pool.req_to_token
+    rows = req_to_token[req_pool_indices]
+    positions = torch.arange(rows.shape[1], device=rows.device)
+    mask = positions[None, :] < seq_lens[:, None]
+    swa_indices = full_to_swa_index_mapping[rows]
+    # FULL pool slots beyond the sliding window have their SWA mapping written
+    # to 0 (see SWATokenToKVPoolAllocator.alloc_extend); skip those so they
+    # don't get counted as divergence.
+    return (
+        ((swa_indices != rows) & mask & (swa_indices != 0))
+        .sum()
+        .to(torch.int32)
+        .view(1)
+    )

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

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Optional
+
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
+from sglang.srt.kv_canary.config import CanaryConfig
+from sglang.srt.kv_canary.endpoint import CanaryEndpoint
+from sglang.srt.kv_canary.runner.kernel_launcher import launch_endpoints_sweep
+from sglang.srt.kv_canary.state import CanaryDeviceState
+from sglang.srt.kv_canary.sweep_plan_builder import build_verify_plan_radix_sweep
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+
+logger = logging.getLogger(__name__)
+
+
+class SweepOrchestrator:
+    def __init__(
+        self,
+        *,
+        config: CanaryConfig,
+        device_state: CanaryDeviceState,
+        buffer_groups: tuple[CanaryBufferGroup, ...],
+        endpoints: tuple[CanaryEndpoint, ...],
+        swa_window_size: int,
+        outer_step_counter_getter: Callable[[], int],
+    ) -> None:
+        self._config = config
+        self._device_state = device_state
+        self._buffer_groups = buffer_groups
+        self._endpoints = endpoints
+        self._swa_window_size = swa_window_size
+        self._outer_step_counter_getter = outer_step_counter_getter
+        self._radix_cache: Optional[BasePrefixCache] = None
+
+        self._last_sweep_step: int = -1
+        self._sweep_passes: int = 0
+
+    @property
+    def sweep_passes(self) -> int:
+        return self._sweep_passes
+
+    def attach_radix_cache(self, radix_cache: BasePrefixCache) -> None:
+        self._radix_cache = radix_cache
+
+    def maybe_run_sweep(self) -> None:
+        if self._config.sweep_interval == 0:
+            return
+        outer_step_counter = self._outer_step_counter_getter()
+        if (
+            self._last_sweep_step >= 0
+            and outer_step_counter - self._last_sweep_step < self._config.sweep_interval
+        ):
+            return
+        self._last_sweep_step = outer_step_counter
+
+        if self._radix_cache is None:
+            return
+
+        violation_log = self._device_state.violation_log
+        for group in self._buffer_groups:
+            window = self._swa_window_size if group.kind is PoolKind.SWA else 0
+            verify_plan = build_verify_plan_radix_sweep(
+                radix_cache=self._radix_cache,
+                swa_window_size=window,
+                full_to_swa_index_mapping=group.swa_index_lut,
+            )
+            launch_endpoints_sweep(
+                endpoints=self._endpoints,
+                group=group,
+                verify_plan=verify_plan,
+                violation_log=violation_log,
+                real_kv_hash_mode=self._config.real_kv_hash_mode,
+            )
+
+        self._sweep_passes += 1
+        logger.info(
+            "[canary] sweep succeeded %d times (last_step=%d)",
+            self._sweep_passes,
+            outer_step_counter,
+        )

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,163 @@
+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()
+    if tag in (
+        CanaryLaunchTag.SWEEP_K_FULL,
+        CanaryLaunchTag.SWEEP_V_FULL,
+        CanaryLaunchTag.SWEEP_K_SWA,
+        CanaryLaunchTag.SWEEP_V_SWA,
+    ):
+        return 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,60 @@
+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
+    swa_verify_total_count: torch.Tensor | None
+
+    @classmethod
+    def allocate(
+        cls,
+        *,
+        num_kernel_tags: int,
+        num_slot_tags: int,
+        swa_verify_total_count_shape: tuple[int, ...] | None,
+        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),
+            swa_verify_total_count=(
+                None
+                if swa_verify_total_count_shape is None
+                else torch.zeros(
+                    swa_verify_total_count_shape, 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,
+        swa_verify_total_count: torch.Tensor | None,
+    ) -> 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)
+        assert (self.swa_verify_total_count is not None) == (
+            swa_verify_total_count is not None
+        )
+        if self.swa_verify_total_count is not None:
+            self.swa_verify_total_count.copy_(swa_verify_total_count)

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

@@ -0,0 +1,325 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import TYPE_CHECKING, Optional
+
+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.req_to_expected_token_ids_manager import (
+    populate_req_to_expected_token_ids,
+)
+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.runner.swa_divergence import SwaDivergenceReporter
+from sglang.srt.kv_canary.single_forward_manager.data import (
+    PostOpsInsideGraphOutputBuffer,
+)
+from sglang.srt.kv_canary.state import CanaryDeviceState
+from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
+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,
+        token_oracle_manager: Optional[TokenOracleManager],
+        swa_divergence_report: Optional[SwaDivergenceReporter],
+        is_eagle_draft_decode: bool,
+    ) -> 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._token_oracle_manager: Optional[TokenOracleManager] = token_oracle_manager
+        self._swa_divergence_report: Optional[SwaDivergenceReporter] = (
+            swa_divergence_report
+        )
+        self._is_eagle_draft_decode: bool = is_eagle_draft_decode
+
+        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]),
+            swa_verify_total_count_shape=(
+                None
+                if swa_divergence_report is None
+                else tuple(swa_divergence_report.verify_total_count_device.shape)
+            ),
+            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"
+            )
+
+        if self._config.enable_verify_token_assert:
+            populate_req_to_expected_token_ids(
+                forward_batch=maybe_inaccurate_forward_batch,
+                req_to_verify_expected_tokens=self._device_state.req_to_verify_expected_tokens,
+            )
+
+    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
+        )
+
+        enable_write_input_assert = self._should_enable_write_input_assert_for_launch(
+            forward_batch
+        )
+        if enable_write_input_assert:
+            manager = self._token_oracle_manager
+            if manager is None:
+                raise RuntimeError(
+                    "kv-canary: enable_write_input_assert=True requires a TokenOracleManager; pass "
+                    "token_oracle_manager=install_oracle_sampler(oracle=...) into "
+                    "install_canary(...)"
+                )
+            manager.fill_expected_inputs(
+                forward_batch=forward_batch,
+                expected_inputs_out=expected_inputs,
+            )
+
+        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,
+            )
+            if self._swa_divergence_report is not None:
+                self._swa_divergence_report.observe_after_invoke_plan(
+                    group=group,
+                    verify_plan=verify_plan,
+                )
+            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,
+                real_kv_hash_mode=self._config.real_kv_hash_mode,
+                enable_write_input_assert=enable_write_input_assert,
+                enable_verify_token_assert=self._config.enable_verify_token_assert,
+            )
+
+        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)
+        enable_write_input_assert = self._should_enable_write_input_assert_for_launch(
+            forward_batch
+        )
+        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,
+                real_kv_hash_mode=self._config.real_kv_hash_mode,
+                enable_write_input_assert=enable_write_input_assert,
+                enable_verify_token_assert=self._config.enable_verify_token_assert,
+            )
+
+        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,
+            swa_verify_total_count=(
+                None
+                if self._swa_divergence_report is None
+                else self._swa_divergence_report.verify_total_count_device
+            ),
+        )
+
+    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 _should_enable_write_input_assert_for_launch(
+        self, forward_batch: ForwardBatch
+    ) -> bool:
+        if not self._config.enable_write_input_assert:
+            return False
+        forward_mode = forward_batch.forward_mode
+        if (
+            self._is_eagle_draft_decode
+            and forward_mode is not None
+            and forward_mode.is_decode()
+        ):
+            return False
+        return True
+
+
+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,135 @@
+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 / sweep, 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|sweep, 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 / sweep × 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``. Allocated only when
+            ``CanaryConfig.enable_verify_token_assert`` is True. 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.
+    """
+
+    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)
+        if config.enable_verify_token_assert:
+            if req_to_token_alloc_size is None or max_context_len is None:
+                raise ValueError(
+                    "kv-canary: CanaryDeviceState.allocate requires req_to_token_alloc_size "
+                    "and max_context_len when CanaryConfig.enable_verify_token_assert is on"
+                )
+            req_to_verify_expected_tokens = torch.empty(
+                (req_to_token_alloc_size, max_context_len),
+                dtype=torch.int32,
+                device=device,
+            )
+        else:
+            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/kv_canary/sweep_plan_builder.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import VerifyPlan
+from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+
+
+def build_verify_plan_radix_sweep(
+    *,
+    radix_cache: BasePrefixCache,
+    swa_window_size: int,
+    full_to_swa_index_mapping: Optional[torch.Tensor],
+    unlocked_only: bool = False,
+) -> VerifyPlan:
+    """Build a sweep VerifyPlan directly from the radix-cache walker.
+
+    The walker covers every slot held by the radix tree. There is no exclusion for slots also owned
+    by running requests because the overlap with per-forward HEAD/TAIL coverage is harmless
+    redundancy. The helper applies the SWA LUT before writing the plan.
+    """
+    device = radix_cache.req_to_token_pool.req_to_token.device
+
+    walk_result = walk_radix_cache_for_canary(
+        radix_cache=radix_cache,
+        unlocked_only=unlocked_only,
+    )
+    slot_indices = walk_result.slot_indices.to(device)
+    positions = walk_result.positions.to(device)
+    prev_slot_indices = walk_result.prev_slot_indices.to(device)
+
+    if swa_window_size > 0:
+        assert (
+            full_to_swa_index_mapping is not None
+        ), "full_to_swa_index_mapping is required when SWA is enabled"
+        slot_indices = _swa_translate(
+            indices=slot_indices, lut=full_to_swa_index_mapping
+        )
+        prev_slot_indices = _swa_translate(
+            indices=prev_slot_indices, lut=full_to_swa_index_mapping
+        )
+
+    num_valid = int(slot_indices.shape[0])
+    verify_plan = VerifyPlan.allocate(verify_capacity=max(1, num_valid), device=device)
+
+    verify_plan.verify_slot_indices[:num_valid].copy_(slot_indices)
+    verify_plan.verify_expected_positions[:num_valid].copy_(positions)
+    verify_plan.verify_prev_slot_indices[:num_valid].copy_(prev_slot_indices)
+    verify_plan.verify_num_valid.fill_(num_valid)
+    verify_plan.enable.fill_(1)
+
+    return verify_plan
+
+
+def _swa_translate(
+    *,
+    indices: torch.Tensor,
+    lut: torch.Tensor,
+) -> torch.Tensor:
+    # 0 is both SWAKVPool's evicted-sentinel and the kernel's kTokenToKvSlotPadding,
+    # so evicted indices propagate as the canonical "no real slot" value and the
+    # verify kernel handles them.
+    if indices.numel() == 0:
+        return indices
+    lut_dev = lut.to(indices.device).to(torch.int64)
+    anchor_mask = indices < 0
+    safe = torch.where(anchor_mask, torch.zeros_like(indices), indices).to(torch.int64)
+    looked_up = lut_dev[safe]
+    return torch.where(anchor_mask, indices.to(torch.int64), looked_up)

python/sglang/srt/kv_canary/token_oracle/install.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
+from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
+from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+
+def install_token_oracle_from_env(
+    *, server_args: ServerArgs, vocab_size: int
+) -> Optional[TokenOracleManager]:
+    # Must be called before create_sampler() so the factory is present when the
+    # Sampler is first constructed.
+    if server_args.sampling_backend != "token_oracle":
+        return None
+
+    oracle = HashOracle(vocab_size=vocab_size)
+    return install_oracle_sampler(oracle=oracle)

python/sglang/srt/kv_canary/token_oracle/oracle.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+
+
+class TokenOracle(Protocol):
+    """Deterministic (generalized_req_id, position) -> token_id mapping."""
+
+    def expected_tokens(
+        self, *, generalized_req_ids: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor: ...
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class HashOracle:
+    """token_id = splitmix64(generalized_req_id XOR position) % vocab_size."""
+
+    vocab_size: int
+
+    def expected_tokens(
+        self, *, generalized_req_ids: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        x = generalized_req_ids.to(torch.int64) ^ positions.to(torch.int64)
+        x = _splitmix64_tensor(x)
+        return _uint64_mod(x, self.vocab_size).to(torch.int32)
+
+
+_C1: int = -4658895280553007687  # 0xBF58476D1CE4E5B9 as signed int64
+_C2: int = -7723592293110705685  # 0x94D049BB133111EB as signed int64
+
+
+def _splitmix64_tensor(x: torch.Tensor) -> torch.Tensor:
+    x = (x ^ _logical_shr(x, 30)) * _C1
+    x = (x ^ _logical_shr(x, 27)) * _C2
+    x = x ^ _logical_shr(x, 31)
+    return x
+
+
+def _logical_shr(x: torch.Tensor, n: int) -> torch.Tensor:
+    return (x >> n) & ((1 << (64 - n)) - 1)
+
+
+def _uint64_mod(x: torch.Tensor, mod: int) -> torch.Tensor:
+    offset = (1 << 64) % mod
+    base = x % mod
+    correction = (x < 0).to(x.dtype) * offset
+    return (base + correction) % mod

python/sglang/srt/kv_canary/token_oracle/oracle_manager.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.token_oracle.oracle import TokenOracle
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class TokenOracleManager:
+    def __init__(self, *, oracle: TokenOracle) -> None:
+        self.oracle = oracle
+
+    def fill_expected_inputs(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        expected_inputs_out: ExpectedInputs,
+    ) -> None:
+        positions = forward_batch.positions
+        input_ids = forward_batch.input_ids
+        num_tokens = int(input_ids.shape[0])
+
+        if num_tokens == 0:
+            return
+
+        generalized_req_ids = _build_generalized_req_id_per_token(
+            forward_batch=forward_batch,
+            num_tokens=num_tokens,
+            generalized_req_ids_per_row=select_generalized_req_ids(
+                vanilla_req_ids=forward_batch.rids_int,
+                bootstrap_room_ids_int=forward_batch.bootstrap_room_ids_int,
+            ),
+        )
+        if forward_batch.forward_mode.is_extend():
+            expected_tokens = input_ids
+        else:
+            expected_tokens = self.oracle.expected_tokens(
+                generalized_req_ids=generalized_req_ids,
+                positions=positions.to(torch.int64),
+            )
+        expected_inputs_out.tokens[:num_tokens].copy_(expected_tokens.to(torch.int64))
+        expected_inputs_out.positions[:num_tokens].copy_(positions.to(torch.int64))
+
+    def sample_next_tokens(
+        self, *, generalized_req_ids: torch.Tensor, logits_positions: torch.Tensor
+    ) -> torch.Tensor:
+        return self.oracle.expected_tokens(
+            generalized_req_ids=generalized_req_ids,
+            positions=logits_positions.to(torch.int64) + 1,
+        )
+
+
+def _build_generalized_req_id_per_token(
+    *,
+    forward_batch: ForwardBatch,
+    num_tokens: int,
+    generalized_req_ids_per_row: torch.Tensor,
+) -> torch.Tensor:
+    forward_mode = forward_batch.forward_mode
+    if forward_mode.is_target_verify():
+        per_req = int(forward_batch.spec_info.draft_token_num)
+        result = _expand_uniform(generalized_req_ids_per_row, per_req)
+    elif forward_mode.is_draft_extend_v2():
+        per_req = int(forward_batch.spec_info.num_tokens_per_req)
+        result = _expand_uniform(generalized_req_ids_per_row, per_req)
+    elif forward_mode.is_extend():
+        extend_seq_lens = forward_batch.extend_seq_lens
+        if extend_seq_lens is None:
+            raise RuntimeError(
+                "_build_generalized_req_id_per_token: extend_seq_lens is None in extend mode"
+            )
+        lens = extend_seq_lens.to(torch.int64)
+        result = torch.repeat_interleave(generalized_req_ids_per_row, lens)
+    else:
+        result = generalized_req_ids_per_row
+
+    if int(result.shape[0]) != num_tokens:
+        raise RuntimeError(
+            f"fill_expected_inputs: sum(lens)={int(result.shape[0])} != num_tokens={num_tokens}"
+        )
+    return result
+
+
+def _expand_uniform(values: torch.Tensor, per_row: int) -> torch.Tensor:
+    bs = int(values.shape[0])
+    return values.unsqueeze(1).expand(bs, per_row).reshape(bs * per_row)
+
+
+def select_generalized_req_ids(
+    *,
+    vanilla_req_ids: torch.Tensor,
+    bootstrap_room_ids_int: torch.Tensor | None,
+) -> torch.Tensor:
+    if bootstrap_room_ids_int is None:
+        return vanilla_req_ids
+
+    bootstrap_room_ids_int = bootstrap_room_ids_int.to(
+        device=vanilla_req_ids.device,
+        dtype=torch.int64,
+    )
+    return torch.where(
+        bootstrap_room_ids_int >= 0,
+        bootstrap_room_ids_int,
+        vanilla_req_ids.to(torch.int64),
+    )

python/sglang/srt/kv_canary/token_oracle/sampler.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+import torch
+
+from sglang.srt.kv_canary.perturb.next_token_swap import maybe_perturb_swap_next_tokens
+from sglang.srt.kv_canary.token_oracle.oracle import TokenOracle
+from sglang.srt.kv_canary.token_oracle.oracle_manager import (
+    TokenOracleManager,
+    select_generalized_req_ids,
+)
+from sglang.srt.layers.sampler import Sampler, register_sampler_backend
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+    from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+
+
+def install_oracle_sampler(*, oracle: TokenOracle) -> TokenOracleManager:
+    manager = TokenOracleManager(oracle=oracle)
+    register_sampler_backend(
+        "token_oracle",
+        lambda: _OracleSampler(token_oracle_manager=manager),
+    )
+    return manager
+
+
+class _OracleSampler(Sampler):
+    def __init__(self, *, token_oracle_manager: TokenOracleManager) -> None:
+        super().__init__()
+        self._token_oracle_manager = token_oracle_manager
+
+    def forward(
+        self,
+        logits_output: LogitsProcessorOutput,
+        sampling_info: SamplingBatchInfo,
+        return_logprob: bool,
+        top_logprobs_nums: List[int],
+        token_ids_logprobs: List[List[int]],
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        vanilla_req_ids = sampling_info.rids_int
+        if vanilla_req_ids is None:
+            raise RuntimeError(
+                "_OracleSampler.forward: generalized_req_id source tensor is None; "
+                "token oracle requires a per-forward generalized_req_id source tensor "
+                "(set in ForwardBatch.init_new when SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE=1)"
+            )
+        batch_next_token_ids = self._token_oracle_manager.sample_next_tokens(
+            generalized_req_ids=select_generalized_req_ids(
+                vanilla_req_ids=vanilla_req_ids,
+                bootstrap_room_ids_int=sampling_info.bootstrap_room_ids_int,
+            ),
+            logits_positions=positions,
+        )
+
+        batch_next_token_ids = maybe_perturb_swap_next_tokens(batch_next_token_ids)
+        return batch_next_token_ids