From 11391b2a1ca22401fffc6cf85a67615f1ed1531a Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 31 May 2026 09:54:24 +0800 Subject: [PATCH] Add the KV-canary core: data layer, MHA KV-pool patcher, and per-forward runner (#26808) --- python/sglang/srt/environ.py | 6 + python/sglang/srt/kv_canary/__init__.py | 0 python/sglang/srt/kv_canary/buffer_group.py | 58 +++ python/sglang/srt/kv_canary/capacities.py | 114 ++++++ python/sglang/srt/kv_canary/config.py | 49 +++ python/sglang/srt/kv_canary/endpoint.py | 156 ++++++++ .../sglang/srt/kv_canary/expected_inputs.py | 24 ++ python/sglang/srt/kv_canary/plan_input.py | 135 +++++++ .../srt/kv_canary/pool_patcher/__init__.py | 0 .../pool_patcher/adapters/__init__.py | 0 .../kv_canary/pool_patcher/adapters/mha.py | 38 ++ .../sglang/srt/kv_canary/pool_patcher/api.py | 64 ++++ .../kv_canary/pool_patcher/buf_info_splice.py | 77 ++++ .../kv_canary/pool_patcher/buffer_alloc.py | 13 + .../srt/kv_canary/pool_patcher/utils.py | 42 +++ .../sglang/srt/kv_canary/runner/__init__.py | 0 .../srt/kv_canary/runner/canary_manager.py | 192 ++++++++++ .../srt/kv_canary/runner/enable_warner.py | 35 ++ .../srt/kv_canary/runner/future_tensor.py | 124 +++++++ .../srt/kv_canary/runner/kernel_launcher.py | 140 ++++++++ .../srt/kv_canary/runner/violation_manager.py | 44 +++ .../kv_canary/runner/violation_reporter.py | 156 ++++++++ .../single_forward_manager/__init__.py | 0 .../kv_canary/single_forward_manager/data.py | 45 +++ .../single_forward_manager/manager.py | 252 +++++++++++++ python/sglang/srt/kv_canary/state.py | 126 +++++++ .../srt/model_executor/forward_batch_info.py | 7 + python/sglang/srt/server_args.py | 13 + python/sglang/test/kv_canary/__init__.py | 0 python/sglang/test/kv_canary/e2e_base.py | 141 ++++++++ python/sglang/test/kv_canary/fixtures.py | 191 ++++++++++ python/sglang/test/kv_canary/mode_config.py | 17 + .../sglang/test/kv_canary/runner_test_base.py | 89 +++++ python/sglang/test/kv_canary/utils.py | 55 +++ .../test/kv_canary/violation_assert_mixin.py | 125 +++++++ .../test/kv_canary/violation_log_utils.py | 31 ++ test/manual/kv_canary/__init__.py | 0 test/registered/kv_canary/test_e2e_base.py | 31 ++ .../kv_canary/test_self_unit_buffer_alloc.py | 27 ++ .../kv_canary/test_self_unit_capacities.py | 89 +++++ .../kv_canary/test_self_unit_endpoint.py | 241 +++++++++++++ .../kv_canary/test_self_unit_future_tensor.py | 180 ++++++++++ .../kv_canary/test_self_unit_plan_input.py | 128 +++++++ .../kv_canary/test_self_unit_pool_patcher.py | 116 ++++++ .../test_self_unit_pool_patcher_utils.py | 52 +++ .../kv_canary/test_self_unit_violation.py | 333 ++++++++++++++++++ 46 files changed, 3756 insertions(+) create mode 100644 python/sglang/srt/kv_canary/__init__.py create mode 100644 python/sglang/srt/kv_canary/buffer_group.py create mode 100644 python/sglang/srt/kv_canary/capacities.py create mode 100644 python/sglang/srt/kv_canary/config.py create mode 100644 python/sglang/srt/kv_canary/endpoint.py create mode 100644 python/sglang/srt/kv_canary/expected_inputs.py create mode 100644 python/sglang/srt/kv_canary/plan_input.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/__init__.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/adapters/__init__.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/adapters/mha.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/api.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/buf_info_splice.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/buffer_alloc.py create mode 100644 python/sglang/srt/kv_canary/pool_patcher/utils.py create mode 100644 python/sglang/srt/kv_canary/runner/__init__.py create mode 100644 python/sglang/srt/kv_canary/runner/canary_manager.py create mode 100644 python/sglang/srt/kv_canary/runner/enable_warner.py create mode 100644 python/sglang/srt/kv_canary/runner/future_tensor.py create mode 100644 python/sglang/srt/kv_canary/runner/kernel_launcher.py create mode 100644 python/sglang/srt/kv_canary/runner/violation_manager.py create mode 100644 python/sglang/srt/kv_canary/runner/violation_reporter.py create mode 100644 python/sglang/srt/kv_canary/single_forward_manager/__init__.py create mode 100644 python/sglang/srt/kv_canary/single_forward_manager/data.py create mode 100644 python/sglang/srt/kv_canary/single_forward_manager/manager.py create mode 100644 python/sglang/srt/kv_canary/state.py create mode 100644 python/sglang/test/kv_canary/__init__.py create mode 100644 python/sglang/test/kv_canary/e2e_base.py create mode 100644 python/sglang/test/kv_canary/fixtures.py create mode 100644 python/sglang/test/kv_canary/mode_config.py create mode 100644 python/sglang/test/kv_canary/runner_test_base.py create mode 100644 python/sglang/test/kv_canary/utils.py create mode 100644 python/sglang/test/kv_canary/violation_assert_mixin.py create mode 100644 python/sglang/test/kv_canary/violation_log_utils.py create mode 100644 test/manual/kv_canary/__init__.py create mode 100644 test/registered/kv_canary/test_e2e_base.py create mode 100644 test/registered/kv_canary/test_self_unit_buffer_alloc.py create mode 100644 test/registered/kv_canary/test_self_unit_capacities.py create mode 100644 test/registered/kv_canary/test_self_unit_endpoint.py create mode 100644 test/registered/kv_canary/test_self_unit_future_tensor.py create mode 100644 test/registered/kv_canary/test_self_unit_plan_input.py create mode 100644 test/registered/kv_canary/test_self_unit_pool_patcher.py create mode 100644 test/registered/kv_canary/test_self_unit_pool_patcher_utils.py create mode 100644 test/registered/kv_canary/test_self_unit_violation.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9ca64d7fc..4c2267621 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -749,6 +749,12 @@ class Envs: SGLANG_PLATFORM = EnvStr("") SGLANG_PLUGINS = EnvStr("") + # =================================================================== + # KV-Canary (testing-only) + # =================================================================== + SGLANG_KV_CANARY_RING_CAPACITY = EnvInt(1024) + SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False) + envs = Envs() EnvField._allow_set_name = False diff --git a/python/sglang/srt/kv_canary/__init__.py b/python/sglang/srt/kv_canary/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/buffer_group.py b/python/sglang/srt/kv_canary/buffer_group.py new file mode 100644 index 000000000..527b47bd3 --- /dev/null +++ b/python/sglang/srt/kv_canary/buffer_group.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum +from typing import Optional + +import torch + + +class PoolKind(IntEnum): + """Which attention regime a canary group belongs to. + + - ``FULL`` covers ``[0, K_req)``. Attached to plain MHA/MLA pools and as one of the two canaries on + every SWA system. + - ``SWA`` covers ``[max(0, K_req - window), K_req)``. Attached as the second canary on every + ``BaseSWAKVPool``. + """ + + FULL = 0 + SWA = 1 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CanaryBufferGroup: + """Canary buffers for one (PoolKind × K-half | V-half) on a pool. + + Each (head | tail) launch sees a single 2-D uint8 buf for the canary. Head and tail use separate canary + buffers so they can be staged at different points in the forward pass without overwriting each other. + + MLA-style pools have no V half (v_head / v_tail = None). SWA pools have two + CanaryBufferGroup instances (FULL sized to the full sub-pool, SWA sized to the swa sub-pool). + + Fields: + kind: PoolKind.FULL or PoolKind.SWA. + k_head: Head canary buffer for K-half launches, shape [num_slots, CANARY_SLOT_BYTES], uint8. + k_tail: Tail canary buffer for K-half launches, same shape, uint8. + v_head: Same for V-half, or None for MLA-style pools. + v_tail: Same for V-half, or None. + swa_index_lut: SWA full-to-swa index mapping LUT, shape [full_pool_size + 1], int64, or None for FULL + groups. Used by launch_canary_plan_kernels to translate verify/seed slot indices at plan time, and by + launch_canary_write_kernel to translate write slots inline. None iff kind == PoolKind.FULL. + kv_token_id_vs_position_offset: Logical-position offset between a canary slot and the source-of-truth token it + fingerprints. 0 for target-style pools (slot ``p`` stores K/V for token at position ``p``); 1 for + EAGLE draft pools where the input_ids rotation makes slot ``p`` store K/V for token at position + ``p + 1``. + """ + + kind: PoolKind + k_head: torch.Tensor + k_tail: torch.Tensor + v_head: Optional[torch.Tensor] + v_tail: Optional[torch.Tensor] + swa_index_lut: Optional[torch.Tensor] + kv_token_id_vs_position_offset: int + + @property + def has_v_half(self) -> bool: + return self.v_head is not None diff --git a/python/sglang/srt/kv_canary/capacities.py b/python/sglang/srt/kv_canary/capacities.py new file mode 100644 index 000000000..5b5f3696e --- /dev/null +++ b/python/sglang/srt/kv_canary/capacities.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CanaryLaunchCapacities: + """Pre-allocation sizes for the per-forward tensors a SingleForwardManager owns. Computed + once at install_canary from ServerArgs + ModelRunner metadata; all fields are upper + bounds - actual per-step usage may be smaller but never larger. + + Fields: + per_forward_verify_capacity: VerifyPlan row capacity for the per-forward HEAD/TAIL + launches. Sized to pool_slot_count * 3 (3x headroom; radix prefix sharing across + running reqs can cause sum_r prefix_lens[r] to exceed the pool slot count). When + the per-step actual count exceeds this, the plan kernel sets VerifyPlan.enable=0 + and the verify kernel skips the step; host logs a warn (no install-time raise). + per_forward_write_req_capacity: WritePlan row capacity for per-forward writes, also used + to size the static PlanInput buffers (= max batch size under cuda graph). + per_forward_write_entry_capacity: Capacity for the expected_input_* placeholder tensors, + one entry per token written in a single forward. + """ + + per_forward_verify_capacity: int + per_forward_write_req_capacity: int + per_forward_write_entry_capacity: int + + def __post_init__(self) -> None: + for name, value in ( + ("per_forward_verify_capacity", self.per_forward_verify_capacity), + ("per_forward_write_req_capacity", self.per_forward_write_req_capacity), + ("per_forward_write_entry_capacity", self.per_forward_write_entry_capacity), + ): + if value <= 0: + raise ValueError(f"kv-canary: {name} must be positive, got {value}") + + @classmethod + def from_args( + cls, + *, + server_args: "ServerArgs", + req_to_token_pool_size: int, + max_seq_len_per_req: int, + pool_slot_count: int, + ) -> "CanaryLaunchCapacities": + if req_to_token_pool_size <= 0: + raise ValueError( + "kv-canary: req_to_token_pool_size must be positive, " + f"got {req_to_token_pool_size}" + ) + if max_seq_len_per_req <= 0: + raise ValueError( + "kv-canary: max_seq_len_per_req must be positive, " + f"got {max_seq_len_per_req}" + ) + if pool_slot_count <= 0: + raise ValueError( + f"kv-canary: pool_slot_count must be positive, got {pool_slot_count}" + ) + + cuda_graph_max_bs = server_args.cuda_graph_max_bs or 0 + if cuda_graph_max_bs < 0: + raise ValueError( + f"kv-canary: cuda_graph_max_bs must be non-negative, got {cuda_graph_max_bs}" + ) + + spec_num_draft_tokens = server_args.speculative_num_draft_tokens + if spec_num_draft_tokens is None: + spec_num_draft_tokens = 0 + if spec_num_draft_tokens < 0: + raise ValueError( + "kv-canary: speculative_num_draft_tokens must be non-negative, " + f"got {spec_num_draft_tokens}" + ) + + max_prefill_tokens = server_args.max_prefill_tokens + if max_prefill_tokens <= 0: + raise ValueError( + f"kv-canary: max_prefill_tokens must be positive, got {max_prefill_tokens}" + ) + + num_tokens_per_bs = 1 + if spec_num_draft_tokens: + num_tokens_per_bs = max(num_tokens_per_bs, spec_num_draft_tokens) + + max_bs = max(cuda_graph_max_bs, req_to_token_pool_size) + + chunked_prefill_size = server_args.chunked_prefill_size + chunked_limit = ( + chunked_prefill_size + if chunked_prefill_size is not None and chunked_prefill_size >= 0 + else math.inf + ) + max_extend_tokens_per_forward = min(max_prefill_tokens, chunked_limit) + + write_entry_capacity = max( + max_bs * num_tokens_per_bs, max_extend_tokens_per_forward + ) + + # Radix prefix sharing lets sum_r prefix_lens[r] exceed pool_slot_count; observed up to ~2x + # on 20 parallel token-oracle prompts. 3x headroom keeps the partial-fallback path + # (plan kernel enable=0 + host warn) exceptional. Overflow does not raise at install time. + per_forward_verify_capacity = int(pool_slot_count * 3) + + return cls( + per_forward_verify_capacity=per_forward_verify_capacity, + per_forward_write_req_capacity=max_bs, + per_forward_write_entry_capacity=write_entry_capacity, + ) diff --git a/python/sglang/srt/kv_canary/config.py b/python/sglang/srt/kv_canary/config.py new file mode 100644 index 000000000..8a21139ad --- /dev/null +++ b/python/sglang/srt/kv_canary/config.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING + +from sglang.srt.environ import envs + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + + +class CanaryMode(str, Enum): + NONE = "none" + LOG = "log" + RAISE = "raise" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CanaryConfig: + """Top-level canary configuration. All knobs live here; nothing reads env vars deeper in the stack. + + Constructed once inside install_canary(server_args, model_runner) via + CanaryConfig.from_env(server_args), then frozen and threaded through the canary stack. + Subsequent runtime never mutates it. + + Fields: + mode: CanaryMode value. none = no canary installed; log = canary runs, violations are logged + but do NOT raise (used for production observability + canary self-test perturb); raise = + violations propagate to host as RuntimeError after the next D2H pump. + ring_capacity: Violation ring capacity (rows in ViolationLog.violation_ring). Sized generously; + overflow only drops detail beyond row N, the monotonic counter still grows. + """ + + mode: CanaryMode + ring_capacity: int + + @classmethod + def from_env(cls, server_args: "ServerArgs") -> "CanaryConfig": + mode_raw = server_args.kv_canary.strip().lower() + if mode_raw not in ("none", "log", "raise"): + raise ValueError( + f"kv-canary: kv_canary must be one of none/log/raise, got {mode_raw!r}" + ) + + return cls( + mode=CanaryMode(mode_raw), + ring_capacity=envs.SGLANG_KV_CANARY_RING_CAPACITY.get(), + ) diff --git a/python/sglang/srt/kv_canary/endpoint.py b/python/sglang/srt/kv_canary/endpoint.py new file mode 100644 index 000000000..1917879be --- /dev/null +++ b/python/sglang/srt/kv_canary/endpoint.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyOrWriteContext, + VerifyPlan, + launch_canary_verify_kernel, +) +from sglang.jit_kernel.kv_canary.write import ( + WritePlan, + launch_canary_write_kernel, +) +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.state import ( + CanaryDeviceState, + ViolationLog, +) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CanaryEndpoint: + kernel_kind: CanaryLaunchTag + canary_buf: torch.Tensor + full_to_swa_index_mapping: Optional[torch.Tensor] + slot_run_counter_view: torch.Tensor + kernel_run_counter_view: torch.Tensor + enable_chain_position_assert: torch.Tensor + + def launch_per_forward( + self, + *, + verify_plan: VerifyPlan, + write_plan: WritePlan, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + enable_write_input_assert: bool, + enable_verify_token_assert: bool, + expected_inputs: ExpectedInputs, + violation_log: ViolationLog, + ) -> None: + context = self._make_verify_or_write_context( + violation_log=violation_log, + ) + launch_canary_verify_kernel( + context=context, + plan=verify_plan, + check_verify_expected_token=enable_verify_token_assert, + ) + + # SWA endpoints translate the per-token slot indices via a device tensor index op before invoking the write kernel. + if self.full_to_swa_index_mapping is not None: + out_cache_loc_for_canary = self.full_to_swa_index_mapping[out_cache_loc] + else: + out_cache_loc_for_canary = out_cache_loc + if enable_write_input_assert: + expected_input_tokens = expected_inputs.tokens + expected_input_positions = expected_inputs.positions + else: + expected_input_tokens = None + expected_input_positions = None + + launch_canary_write_kernel( + context=context, + plan=write_plan, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc_for_canary, + enable_write_input_assert=enable_write_input_assert, + expected_input_tokens=expected_input_tokens, + expected_input_positions=expected_input_positions, + ) + + def _make_verify_or_write_context( + self, + *, + violation_log: ViolationLog, + ) -> VerifyOrWriteContext: + return VerifyOrWriteContext( + canary_buf=self.canary_buf, + kernel_kind=self.kernel_kind, + violation_ring=violation_log.violation_ring, + violation_write_index=violation_log.violation_write_index, + slot_run_counter=self.slot_run_counter_view, + kernel_run_counter=self.kernel_run_counter_view, + enable_chain_position_assert=self.enable_chain_position_assert, + ) + + +def _resolve_canary_buf( + *, + slot: str, + half: str, + group: CanaryBufferGroup, +) -> torch.Tensor: + if half == "K": + if slot == "HEAD": + return group.k_head + return group.k_tail + if slot == "HEAD": + return group.v_head + return group.v_tail + + +_FULL_LAYOUT: tuple[tuple[CanaryLaunchTag, str, str], ...] = ( + (CanaryLaunchTag.HEAD_K_FULL, "HEAD", "K"), + (CanaryLaunchTag.HEAD_V_FULL, "HEAD", "V"), + (CanaryLaunchTag.TAIL_K_FULL, "TAIL", "K"), + (CanaryLaunchTag.TAIL_V_FULL, "TAIL", "V"), +) + + +_SWA_LAYOUT: tuple[tuple[CanaryLaunchTag, str, str], ...] = ( + (CanaryLaunchTag.HEAD_K_SWA, "HEAD", "K"), + (CanaryLaunchTag.HEAD_V_SWA, "HEAD", "V"), + (CanaryLaunchTag.TAIL_K_SWA, "TAIL", "K"), + (CanaryLaunchTag.TAIL_V_SWA, "TAIL", "V"), +) + + +def build_endpoints_from_group( + *, + group: CanaryBufferGroup, + device_state: CanaryDeviceState, +) -> tuple[CanaryEndpoint, ...]: + """Enumerate (slot × half) endpoints for one CanaryBufferGroup.""" + pool_kind = group.kind + layout = _FULL_LAYOUT if pool_kind is PoolKind.FULL else _SWA_LAYOUT + + endpoints: list[CanaryEndpoint] = [] + for tag, slot, half in layout: + if half == "V" and not group.has_v_half: + continue + + canary_buf = _resolve_canary_buf(slot=slot, half=half, group=group) + lut = group.swa_index_lut if pool_kind is PoolKind.SWA else None + slot_view = device_state.slot_run_counters[tag.value : tag.value + 1] + kernel_view = device_state.kernel_run_counters[tag.value : tag.value + 1] + endpoints.append( + CanaryEndpoint( + kernel_kind=tag, + canary_buf=canary_buf, + full_to_swa_index_mapping=lut, + slot_run_counter_view=slot_view, + kernel_run_counter_view=kernel_view, + enable_chain_position_assert=device_state.enable_chain_position_assert, + ) + ) + + return tuple(endpoints) diff --git a/python/sglang/srt/kv_canary/expected_inputs.py b/python/sglang/srt/kv_canary/expected_inputs.py new file mode 100644 index 000000000..9cff403e3 --- /dev/null +++ b/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], + ) diff --git a/python/sglang/srt/kv_canary/plan_input.py b/python/sglang/srt/kv_canary/plan_input.py new file mode 100644 index 000000000..477a9026f --- /dev/null +++ b/python/sglang/srt/kv_canary/plan_input.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PlanInput: + """Pre-staged input to launch_canary_plan_kernels for the per-forward path. + + All tensors live on device. + + Fields: + req_pool_indices: Per-row ReqToTokenPool row index, shape [bs_capacity], int64. + 0 = padding sentinel. + prefix_lens: Per-req prefix length already written before this step, shape + [bs_capacity], int64. Extend → extend_prefix_lens; decode → seq_lens - 1. + extend_seq_lens: Per-req tokens being written this step, shape [bs_capacity], int64. + Extend length or all-ones for decode. + req_to_verify_expected_tokens_valid_lens: Per-req snapshot length on the verify-token pool, + shape [bs_capacity], int64. Equals + ``len(req.origin_input_ids) + len(req.output_ids)`` at the moment ``ForwardBatch`` + was built. The plan kernel uses ``valid_lens[req_id]`` as the upper bound on + ``sot_pos`` when gathering the expected token; everything past the snapshot + (e.g. EAGLE draft / verify positions, or stale residue from a longer recycled + slot owner) returns the ``-1`` sentinel and the verify kernel skips the check. + The naive build never populates the verify-token-id cross-check, so this stays all + zeros and the plan kernel's gather degrades to the ``-1`` skip sentinel. + + Allocated fresh per forward by :class:`SingleForwardManager`. The boundary + ForwardBatch token/position/slot tensors must already be int64 + contiguous (upstream phase-1 hook is responsible). + """ + + req_pool_indices: torch.Tensor + prefix_lens: torch.Tensor + extend_seq_lens: torch.Tensor + req_to_verify_expected_tokens_valid_lens: torch.Tensor + + def zero_(self) -> None: + self.req_pool_indices.zero_() + self.prefix_lens.zero_() + self.extend_seq_lens.zero_() + self.req_to_verify_expected_tokens_valid_lens.zero_() + + @classmethod + def allocate( + cls, + *, + bs_capacity: int, + device: torch.device, + ) -> "PlanInput": + return cls( + req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device), + prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device), + extend_seq_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device), + req_to_verify_expected_tokens_valid_lens=torch.zeros( + bs_capacity, dtype=torch.int64, device=device + ), + ) + + def fill_from_forward_batch(self, *, forward_batch: "ForwardBatch") -> None: + req_pool_indices = forward_batch.req_pool_indices + bs = int(req_pool_indices.shape[0]) + capacity = int(self.req_pool_indices.shape[0]) + if bs > capacity: + raise RuntimeError( + f"kv-canary: per-forward batch size {bs} exceeds static capacity {capacity}; " + "raise the buffer size in CanaryLaunchCapacities" + ) + + self.zero_() + self.req_pool_indices[:bs].copy_(req_pool_indices) + + _extract_prefix_lens_and_extend_seq_lens( + forward_batch=forward_batch, + out_prefix_lens=self.prefix_lens[:bs], + out_extend_seq_lens=self.extend_seq_lens[:bs], + bs=bs, + ) + + +def _extract_prefix_lens_and_extend_seq_lens( + *, + forward_batch: "ForwardBatch", + out_prefix_lens: torch.Tensor, + out_extend_seq_lens: torch.Tensor, + bs: int, +) -> None: + # TODO: once ForwardMode is refactored upstream so every mode ships a canonical + # (prefix_lens, extend_seq_lens) pair on forward_batch, collapse this back to a single + # unconditional copy. + forward_mode = forward_batch.forward_mode + spec_info = forward_batch.spec_info + if forward_mode.is_decode_or_idle(): + # Anchor on ``positions`` (canonical write position) — eagle draft leaves seq_lens + # pre-bump so deriving prefix_lens from seq_lens is off-by-one. Padding tail (positions + # shorter than bs under cuda-graph padding) keeps whatever stale data it had; the offsets + # kernel masks those rows via ``is_active`` before using prefix_lens. + positions = forward_batch.positions + out_prefix_lens[: positions.shape[0]].copy_(positions.to(torch.int64)) + out_extend_seq_lens.fill_(1) + elif forward_mode.is_target_verify(): + # Evidence: EagleVerifyInputV2Mixin.prepare_for_v2_verify assigns out_cache_loc in + # [seq_lens, seq_lens + draft_token_num) without bumping seq_lens. The target-verify + # branch in TRTLLMHAAttnBackend.init_forward_metadata uses seq_lens as the prefix and + # tokens_per_req as the query length, so mirror that as seq_lens plus draft_token_num. + out_prefix_lens.copy_(forward_batch.seq_lens[:bs].to(torch.int64)) + out_extend_seq_lens.fill_(int(spec_info.draft_token_num)) + elif forward_mode.is_draft_extend_v2(): + # Evidence: EagleDraftInputV2Mixin.prepare_for_extend_to_fill_draft_kvcache bumps + # seq_lens by num_draft_tokens. FlashAttentionBackend.init_forward_metadata reads the + # draft-extend-v2 query length from spec_info.extend_seq_lens_tensor when available. + # CUDA-graph replay passes extend_seq_lens but omits extend_prefix_lens, so derive the + # prefix as seq_lens - extend_seq_lens. + extend_seq_lens = forward_batch.extend_seq_lens[:bs].to(torch.int64) + out_extend_seq_lens.copy_(extend_seq_lens) + out_prefix_lens.copy_( + forward_batch.seq_lens[:bs].to(torch.int64) - extend_seq_lens + ) + elif forward_mode.is_extend(): + # Evidence: ForwardBatch.init_new copies batch.prefix_lens and batch.extend_lens into + # extend_prefix_lens / extend_seq_lens for non-decode, non-idle modes, matching regular + # extend metadata builders that consume those tensors directly. + out_prefix_lens.copy_(forward_batch.extend_prefix_lens[:bs].to(torch.int64)) + out_extend_seq_lens.copy_(forward_batch.extend_seq_lens[:bs].to(torch.int64)) + else: + raise NotImplementedError( + f"Unsupported forward mode for kv-canary: {forward_mode}" + ) diff --git a/python/sglang/srt/kv_canary/pool_patcher/__init__.py b/python/sglang/srt/kv_canary/pool_patcher/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/pool_patcher/adapters/__init__.py b/python/sglang/srt/kv_canary/pool_patcher/adapters/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/pool_patcher/adapters/mha.py b/python/sglang/srt/kv_canary/pool_patcher/adapters/mha.py new file mode 100644 index 000000000..fc8dd5d0c --- /dev/null +++ b/python/sglang/srt/kv_canary/pool_patcher/adapters/mha.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import torch + +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.pool_patcher.buf_info_splice import patch_buf_info_method +from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf + + +def attach_mha( + *, + pool: object, + device: torch.device, + kv_token_id_vs_position_offset: int, +) -> tuple[CanaryBufferGroup, ...]: + num_slots = int(pool.k_buffer[0].shape[0]) + k_head = alloc_canary_buf(num_slots=num_slots, device=device) + k_tail = alloc_canary_buf(num_slots=num_slots, device=device) + v_head = alloc_canary_buf(num_slots=num_slots, device=device) + v_tail = alloc_canary_buf(num_slots=num_slots, device=device) + + group = CanaryBufferGroup( + kind=PoolKind.FULL, + k_head=k_head, + k_tail=k_tail, + v_head=v_head, + v_tail=v_tail, + swa_index_lut=None, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + patch_buf_info_method( + pool, + method_name="get_contiguous_buf_infos", + group=group, + has_v_half=True, + page_size=pool.page_size, + ) + return (group,) diff --git a/python/sglang/srt/kv_canary/pool_patcher/api.py b/python/sglang/srt/kv_canary/pool_patcher/api.py new file mode 100644 index 000000000..90a9a217d --- /dev/null +++ b/python/sglang/srt/kv_canary/pool_patcher/api.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import logging +from typing import Callable, Dict, Type + +import torch + +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup +from sglang.srt.kv_canary.config import CanaryConfig +from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha +from sglang.srt.mem_cache.memory_pool import ( + KVCache, + MHATokenToKVPool, + MHATokenToKVPoolFP4, +) + +logger = logging.getLogger(__name__) + +PoolAttacher = Callable[..., tuple[CanaryBufferGroup, ...]] + +_POOL_ATTACHERS: Dict[Type, PoolAttacher] = { + MHATokenToKVPool: attach_mha, + MHATokenToKVPoolFP4: attach_mha, +} + + +def register_pool_attacher(pool_class: Type, attacher: PoolAttacher) -> None: + _POOL_ATTACHERS[pool_class] = attacher + + +def attach_canary_buffers( + *, + pool: KVCache, + config: CanaryConfig, + device: torch.device, + kv_token_id_vs_position_offset: int, +) -> tuple[CanaryBufferGroup, ...]: + """Install canary buffers on a KV pool and return the resulting CanaryBufferGroup tuple. + + ``kv_token_id_vs_position_offset`` is propagated into every produced :class:`CanaryBufferGroup` (0 for target + pools; 1 for draft pools where the input-ids rotation shifts the slot-to-token mapping by one). + """ + attacher = _POOL_ATTACHERS.get(type(pool)) + if attacher is None: + raise NotImplementedError( + f"kv-canary: no attacher registered for pool class {type(pool).__name__}; " + f"supported: {sorted(cls.__name__ for cls in _POOL_ATTACHERS)}" + ) + + groups = attacher( + pool=pool, + device=device, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + logger.info( + "attach_canary_buffers: pool=%s attacher=%s n_groups=%d kinds=%s " + "kv_token_id_vs_position_offset=%d", + type(pool).__name__, + attacher.__name__, + len(groups), + [g.kind.name for g in groups], + kv_token_id_vs_position_offset, + ) + return groups diff --git a/python/sglang/srt/kv_canary/pool_patcher/buf_info_splice.py b/python/sglang/srt/kv_canary/pool_patcher/buf_info_splice.py new file mode 100644 index 000000000..968572840 --- /dev/null +++ b/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.()`` 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 diff --git a/python/sglang/srt/kv_canary/pool_patcher/buffer_alloc.py b/python/sglang/srt/kv_canary/pool_patcher/buffer_alloc.py new file mode 100644 index 000000000..8e2f265ee --- /dev/null +++ b/python/sglang/srt/kv_canary/pool_patcher/buffer_alloc.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +import torch + +from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES + + +def alloc_canary_buf( + *, + num_slots: int, + device: torch.device, +) -> torch.Tensor: + return torch.zeros(num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device) diff --git a/python/sglang/srt/kv_canary/pool_patcher/utils.py b/python/sglang/srt/kv_canary/pool_patcher/utils.py new file mode 100644 index 000000000..946845979 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/kv_canary/runner/__init__.py b/python/sglang/srt/kv_canary/runner/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py new file mode 100644 index 000000000..98f25be53 --- /dev/null +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import contextlib +import logging +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator, Optional, Sequence + +import torch + +from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag +from sglang.srt.environ import envs +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup +from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities +from sglang.srt.kv_canary.config import CanaryConfig +from sglang.srt.kv_canary.endpoint import ( + CanaryEndpoint, + build_endpoints_from_group, +) +from sglang.srt.kv_canary.runner.violation_manager import ViolationManager +from sglang.srt.kv_canary.single_forward_manager.manager import ( + SingleForwardManager, + _PreOpsMaybeInsideGraphOutput, +) +from sglang.srt.kv_canary.state import CanaryDeviceState + +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +logger = logging.getLogger(__name__) + + +class CanaryManager: + def __init__( + self, + *, + config: CanaryConfig, + buffer_groups: tuple[CanaryBufferGroup, ...], + device: torch.device, + req_to_token_pool: "ReqToTokenPool", + launch_capacities: CanaryLaunchCapacities, + swa_window_size: int = 0, + ) -> None: + self.config = config + self._req_to_token_pool = req_to_token_pool + self._swa_window_size = swa_window_size + self._outer_step_counter: int = 0 + self._active_single_forward_manager_index: Optional[int] = None + + self._buffer_groups: tuple[CanaryBufferGroup, ...] = tuple(buffer_groups) + + self._device_state = CanaryDeviceState.allocate( + config=config, + device=device, + num_tags=len(CanaryLaunchTag), + req_to_token_alloc_size=req_to_token_pool.req_to_token.shape[0], + max_context_len=req_to_token_pool.max_context_len, + ) + # Disable the chain-step position assert until warmup / cuda-graph capture finishes + # (synthetic positions trip the +1 invariant). mark_init_finished() sets it to 1. + self._device_state.enable_chain_position_assert.fill_(0) + + self._endpoints: tuple[CanaryEndpoint, ...] = tuple( + endpoint + for group in self._buffer_groups + for endpoint in build_endpoints_from_group( + group=group, device_state=self._device_state + ) + ) + self._active_tags: tuple[CanaryLaunchTag, ...] = tuple( + sorted( + {endpoint.kernel_kind for endpoint in self._endpoints}, + key=lambda tag: tag.value, + ) + ) + + self._d2h_stream: torch.cuda.Stream = torch.cuda.Stream(device=device) + + self._violation_manager = ViolationManager( + config=config, + device_state=self._device_state, + d2h_stream=self._d2h_stream, + outer_step_counter_getter=self._get_outer_step_counter, + ) + self._single_forward_managers: tuple[SingleForwardManager, ...] = ( + SingleForwardManager( + config=config, + device=device, + device_state=self._device_state, + buffer_groups=self._buffer_groups, + endpoints=self._endpoints, + req_to_token_pool=req_to_token_pool, + swa_window_size=self._swa_window_size, + per_forward_verify_capacity=launch_capacities.per_forward_verify_capacity, + per_forward_write_req_capacity=launch_capacities.per_forward_write_req_capacity, + per_forward_write_entry_capacity=launch_capacities.per_forward_write_entry_capacity, + d2h_stream=self._d2h_stream, + ), + ) + + @contextlib.contextmanager + def with_active_single_forward_manager(self, index: int) -> Iterator[None]: + assert ( + self._active_single_forward_manager_index is None + ), "kv-canary: nested with_active_single_forward_manager is forbidden" + self._active_single_forward_manager_index = index + try: + yield + finally: + assert self._active_single_forward_manager_index == index, ( + f"kv-canary: with_active_single_forward_manager({index}) exited with " + f"_active_single_forward_manager_index=" + f"{self._active_single_forward_manager_index}; nested or mismatched bracket" + ) + self._active_single_forward_manager_index = None + + def pre_ops_maybe_inside_graph( + self, forward_batch: "ForwardBatch" + ) -> _PreOpsMaybeInsideGraphOutput: + assert self._active_single_forward_manager_index is not None, ( + "kv-canary: pre_ops_maybe_inside_graph called without active SingleForwardManager; " + "caller must wrap in CanaryManager.with_active_single_forward_manager(i)" + ) + sfm = self._single_forward_managers[self._active_single_forward_manager_index] + return sfm.pre_ops_maybe_inside_graph(forward_batch) + + def post_ops_maybe_inside_graph( + self, + forward_batch: "ForwardBatch", + pre_ops_output: _PreOpsMaybeInsideGraphOutput, + ) -> None: + assert self._active_single_forward_manager_index is not None, ( + "kv-canary: post_ops_maybe_inside_graph called without active SingleForwardManager; " + "caller must wrap in CanaryManager.with_active_single_forward_manager(i)" + ) + sfm = self._single_forward_managers[self._active_single_forward_manager_index] + sfm.post_ops_maybe_inside_graph(forward_batch, pre_ops_output) + + @contextlib.contextmanager + def with_ops_outside_graph( + self, + *, + single_forward_indices: Sequence[int], + maybe_inaccurate_forward_batch: "ForwardBatch", + ) -> Iterator[None]: + self._pre_ops_outside_graph( + single_forward_indices=single_forward_indices, + maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, + ) + try: + yield + finally: + self._post_ops_outside_graph( + single_forward_indices=single_forward_indices, + maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, + ) + + def _pre_ops_outside_graph( + self, + *, + single_forward_indices: Sequence[int], + maybe_inaccurate_forward_batch: "ForwardBatch", + ) -> None: + for idx in single_forward_indices: + self._single_forward_managers[idx].pre_ops_outside_graph( + maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch + ) + + def _post_ops_outside_graph( + self, + *, + single_forward_indices: Sequence[int], + maybe_inaccurate_forward_batch: "ForwardBatch", + ) -> None: + for idx in single_forward_indices: + self._single_forward_managers[idx].post_ops_outside_graph() + self._outer_step_counter += 1 + self._violation_manager.step() + + def mark_init_finished(self) -> None: + for single_forward_manager in self._single_forward_managers: + single_forward_manager.phase_checker.enable_assert() + self._device_state.enable_chain_position_assert.fill_(1) + + def _get_outer_step_counter(self) -> int: + return self._outer_step_counter + + +@contextmanager +def context_tuple(ctx_a, ctx_b): + with ctx_a, ctx_b: + yield diff --git a/python/sglang/srt/kv_canary/runner/enable_warner.py b/python/sglang/srt/kv_canary/runner/enable_warner.py new file mode 100644 index 000000000..ca922859b --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/kv_canary/runner/future_tensor.py b/python/sglang/srt/kv_canary/runner/future_tensor.py new file mode 100644 index 000000000..e13fffdbb --- /dev/null +++ b/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 + ) diff --git a/python/sglang/srt/kv_canary/runner/kernel_launcher.py b/python/sglang/srt/kv_canary/runner/kernel_launcher.py new file mode 100644 index 000000000..418727308 --- /dev/null +++ b/python/sglang/srt/kv_canary/runner/kernel_launcher.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Optional + +import torch + +from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels +from sglang.jit_kernel.kv_canary.verify import ( + CanaryLaunchTag, + VerifyPlan, +) +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.srt.environ import envs +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.endpoint import CanaryEndpoint +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.plan_input import PlanInput +from sglang.srt.kv_canary.state import ViolationLog + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +_BOUNDARY_INT_DTYPES = (torch.int32, torch.int64) +_INPUT_IDS = "forward_batch.input_ids" +_OUT_LOC = "forward_batch.out_cache_loc" +_POSITIONS = "forward_batch.positions" + + +def invoke_plan( + *, + plan_input: PlanInput, + verify_plan: VerifyPlan, + write_plan: WritePlan, + group: CanaryBufferGroup, + req_to_token: torch.Tensor, + swa_window_size: int, + req_to_verify_expected_tokens: Optional[torch.Tensor], +) -> None: + window = swa_window_size if group.kind is PoolKind.SWA else 0 + launch_canary_plan_kernels( + verify_plan_out=verify_plan, + write_plan_out=write_plan, + req_pool_indices=plan_input.req_pool_indices, + prefix_lens=plan_input.prefix_lens, + extend_seq_lens=plan_input.extend_seq_lens, + req_to_token=req_to_token, + swa_window_size=window, + full_to_swa_index_mapping=group.swa_index_lut, + verify_capacity=int(verify_plan.verify_slot_indices.shape[0]), + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + req_to_verify_expected_tokens_valid_lens=plan_input.req_to_verify_expected_tokens_valid_lens, + kv_token_id_vs_position_offset=group.kv_token_id_vs_position_offset, + ) + + +def launch_endpoints_per_forward( + *, + endpoints: tuple[CanaryEndpoint, ...], + group: CanaryBufferGroup, + tag_filter: Callable[[CanaryLaunchTag], bool], + verify_plan: VerifyPlan, + write_plan: WritePlan, + forward_batch: "ForwardBatch", + expected_inputs: ExpectedInputs, + violation_log: ViolationLog, + enable_write_input_assert: bool = False, + enable_verify_token_assert: bool = False, +) -> None: + positions = _canonicalize_boundary_int64(forward_batch.positions, _POSITIONS) + out_cache_loc = _canonicalize_boundary_int64(forward_batch.out_cache_loc, _OUT_LOC) + input_ids = _canonicalize_boundary_int64(forward_batch.input_ids, _INPUT_IDS) + + num_tokens = int(positions.shape[0]) + if expected_inputs.tokens.shape[0] != num_tokens: + raise RuntimeError( + f"kv-canary: expected_inputs.tokens shape {expected_inputs.tokens.shape[0]} " + f"!= num_tokens {num_tokens}; caller must slice before invoking" + ) + if expected_inputs.positions.shape[0] != num_tokens: + raise RuntimeError( + f"kv-canary: expected_inputs.positions shape {expected_inputs.positions.shape[0]} " + f"!= num_tokens {num_tokens}; caller must slice before invoking" + ) + + active_endpoints = [ + endpoint + for endpoint in endpoints + if _endpoint_belongs_to_group(endpoint, group) + and tag_filter(endpoint.kernel_kind) + and passes_v_half_gate(endpoint.kernel_kind) + ] + assert len(active_endpoints) > 0 + + for endpoint in active_endpoints: + endpoint.launch_per_forward( + verify_plan=verify_plan, + write_plan=write_plan, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + enable_write_input_assert=enable_write_input_assert, + enable_verify_token_assert=enable_verify_token_assert, + expected_inputs=expected_inputs, + violation_log=violation_log, + ) + + +def _is_v_half_tag(tag: CanaryLaunchTag) -> bool: + return tag in ( + CanaryLaunchTag.HEAD_V_FULL, + CanaryLaunchTag.TAIL_V_FULL, + CanaryLaunchTag.HEAD_V_SWA, + CanaryLaunchTag.TAIL_V_SWA, + ) + + +def passes_v_half_gate(tag: CanaryLaunchTag) -> bool: + if not _is_v_half_tag(tag): + return True + return envs.SGLANG_KV_CANARY_ENABLE_MHA_V.get() + + +def _endpoint_belongs_to_group( + endpoint: CanaryEndpoint, group: CanaryBufferGroup +) -> bool: + suffix = endpoint.kernel_kind.name.rsplit("_", 1)[1] + return suffix == group.kind.name + + +def _canonicalize_boundary_int64( + tensor: torch.Tensor | None, name: str +) -> torch.Tensor | None: + if tensor is None: + return None + if tensor.dtype not in _BOUNDARY_INT_DTYPES: + raise TypeError( + f"kv-canary: {name} must have dtype torch.int32 or torch.int64, got {tensor.dtype}" + ) + return tensor.to(torch.int64).contiguous() diff --git a/python/sglang/srt/kv_canary/runner/violation_manager.py b/python/sglang/srt/kv_canary/runner/violation_manager.py new file mode 100644 index 000000000..d533eac43 --- /dev/null +++ b/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] diff --git a/python/sglang/srt/kv_canary/runner/violation_reporter.py b/python/sglang/srt/kv_canary/runner/violation_reporter.py new file mode 100644 index 000000000..ed81b667a --- /dev/null +++ b/python/sglang/srt/kv_canary/runner/violation_reporter.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import logging + +from sglang.jit_kernel.kv_canary.consts import FailReason +from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag +from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode +from sglang.srt.kv_canary.state import CanaryDeviceState + +logger = logging.getLogger(__name__) + +_WRITE_BITS = FailReason.WRITE_TOKEN_MISMATCH | FailReason.WRITE_POSITION_MISMATCH +_TOKEN_BITS = FailReason.WRITE_TOKEN_MISMATCH | FailReason.VERIFY_TOKEN_MISMATCH + + +def _reason_label(bit: FailReason) -> str: + return bit.name.lower().removesuffix("_mismatch") + + +class ViolationReporter: + def __init__( + self, + *, + config: CanaryConfig, + device_state: CanaryDeviceState, + ) -> None: + self._config = config + self._device_state = device_state + self._raised: bool = False + self._last_logged_write_index: int = 0 + + @property + def is_raised(self) -> bool: + return self._raised + + def log_or_raise_violation(self, *, outer_step_counter: int) -> None: + violation_log = self._device_state.violation_log + write_index = int(violation_log.violation_write_index.cpu().item()) + if write_index == 0: + return + ring = violation_log.violation_ring.cpu() + ring_capacity = int(ring.shape[0]) + valid_count = min(write_index, ring_capacity) + ring_overflow = write_index > ring_capacity + + start = min(self._last_logged_write_index, valid_count) + if start >= valid_count: + return + + messages: list[str] = [ + _format_violation( + row=ring[i].tolist(), + total=write_index, + ring_overflow=ring_overflow, + step_when_pumped=outer_step_counter, + ) + for i in range(start, valid_count) + ] + self._last_logged_write_index = valid_count + + # log mode: always surface every violation as WARNING. + if self._config.mode is CanaryMode.LOG: + for message in messages: + logger.warning(message) + return + self._raised = True + raise RuntimeError("\n".join(messages)) + + +def _canary_kind_label(tag: CanaryLaunchTag) -> str: + name_lower = tag.name.lower() + return f"per_forward_{name_lower}" + + +def _format_violation( + *, + row: list[int], + total: int, + ring_overflow: bool, + step_when_pumped: int, +) -> str: + ( + kernel_kind, + slot_idx, + position, + stored_token, + expected_token, + stored_chain_hash, + expected_aux, + fail_reason_bits, + ) = row + try: + tag_label = CanaryLaunchTag(int(kernel_kind)).name + canary_kind = _canary_kind_label(CanaryLaunchTag(int(kernel_kind))) + except ValueError: + tag_label = f"unknown({int(kernel_kind)})" + canary_kind = tag_label + bits_int = int(fail_reason_bits) + reasons = [_reason_label(bit) for bit in FailReason if bits_int & int(bit)] + is_write = bool(bits_int & int(_WRITE_BITS)) + u64_mask = (1 << 64) - 1 + + # Stable single-line key=value summary, parsed by the regex in + # python/sglang/test/kv_canary/violation_log_utils.py and asserted by + # assert_violation_logged_any in python/sglang/test/kv_canary/violation_assert_mixin.py. + # Format frozen: do not reorder / rename / change separators without updating those helpers. + structured_line = ( + f"kv_canary violation: " + f"launch_tag={tag_label} " + f"fail_reason={'+'.join(reasons) if reasons else 'none'} " + f"slot_idx={int(slot_idx)} " + f"position={int(position)} " + f"stored_token={int(stored_token)} " + f"expected_token={int(expected_token)} " + f"stored_chain_hash={int(stored_chain_hash) & u64_mask:#018x} " + f"expected_aux={int(expected_aux) & u64_mask:#018x}" + ) + + header = ( + f"KV cache canary violation detected (kernel_kind={tag_label}, " + f"slot_idx={int(slot_idx)}, position={int(position)})" + ) + kind_line = f"canary_kind: {canary_kind}" + reasons_line = f" fail_reasons: {' '.join(reasons) if reasons else 'none'}" + footer = ( + f" total_violations={total} ring_overflow={ring_overflow} " + f"step_when_pumped={step_when_pumped}" + ) + + has_token_check = bool(bits_int & int(_TOKEN_BITS)) + if is_write: + running_prev_hash = int(stored_chain_hash) & u64_mask + body = [ + ( + f" actual: token_id={int(stored_token)} position={int(position)} " + f"prev_hash={running_prev_hash:#018x}" + ), + ( + f" expected: token_id={int(expected_token)} position={int(expected_aux)}" + ), + ] + else: + stored_prev_hash = int(stored_chain_hash) & u64_mask + expected_prev_hash = int(expected_aux) & u64_mask + stored_body = ( + f" stored: token_id={int(stored_token)} position={int(position)} " + f"prev_hash={stored_prev_hash:#018x}" + ) + expected_body = ( + f" expected: token_id={int(expected_token)} prev_hash={expected_prev_hash:#018x}" + if has_token_check + else f" expected: prev_hash={expected_prev_hash:#018x}" + ) + body = [stored_body, expected_body] + + return "\n".join([structured_line, header, kind_line, reasons_line, *body, footer]) diff --git a/python/sglang/srt/kv_canary/single_forward_manager/__init__.py b/python/sglang/srt/kv_canary/single_forward_manager/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/single_forward_manager/data.py b/python/sglang/srt/kv_canary/single_forward_manager/data.py new file mode 100644 index 000000000..29b5b60c3 --- /dev/null +++ b/python/sglang/srt/kv_canary/single_forward_manager/data.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PostOpsInsideGraphOutputBuffer: + verify_plan_enable: torch.Tensor + kernel_run_counters: torch.Tensor + slot_run_counters: torch.Tensor + violation_write_index: torch.Tensor + + @classmethod + def allocate( + cls, + *, + num_kernel_tags: int, + num_slot_tags: int, + device: torch.device, + ) -> "PostOpsInsideGraphOutputBuffer": + return cls( + verify_plan_enable=torch.zeros(1, dtype=torch.int32, device=device), + kernel_run_counters=torch.zeros( + num_kernel_tags, dtype=torch.int64, device=device + ), + slot_run_counters=torch.zeros( + num_slot_tags, dtype=torch.int64, device=device + ), + violation_write_index=torch.zeros(1, dtype=torch.int32, device=device), + ) + + def copy_from( + self, + *, + verify_plan_enable: torch.Tensor, + kernel_run_counters: torch.Tensor, + slot_run_counters: torch.Tensor, + violation_write_index: torch.Tensor, + ) -> None: + self.verify_plan_enable.copy_(verify_plan_enable) + self.kernel_run_counters.copy_(kernel_run_counters) + self.slot_run_counters.copy_(slot_run_counters) + self.violation_write_index.copy_(violation_write_index) diff --git a/python/sglang/srt/kv_canary/single_forward_manager/manager.py b/python/sglang/srt/kv_canary/single_forward_manager/manager.py new file mode 100644 index 000000000..c9e0ff23f --- /dev/null +++ b/python/sglang/srt/kv_canary/single_forward_manager/manager.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup +from sglang.srt.kv_canary.config import CanaryConfig +from sglang.srt.kv_canary.endpoint import CanaryEndpoint +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.plan_input import PlanInput +from sglang.srt.kv_canary.runner.enable_warner import CanaryEnableWarner +from sglang.srt.kv_canary.runner.kernel_launcher import ( + invoke_plan, + launch_endpoints_per_forward, +) +from sglang.srt.kv_canary.single_forward_manager.data import ( + PostOpsInsideGraphOutputBuffer, +) +from sglang.srt.kv_canary.state import CanaryDeviceState +from sglang.srt.utils.phase_checker import SimplePhaseChecker + +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class _SingleForwardPhase(IntEnum): + IDLE = 0 + AFTER_PRE_OUT = 1 + AFTER_PRE_MAYBE_IN = 2 + AFTER_POST_MAYBE_IN = 3 + + +def _torch_reduce_minimum(tensors: list[torch.Tensor]) -> torch.Tensor: + out = tensors[0] + for t in tensors[1:]: + out = torch.minimum(out, t) + return out + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _PreOpsMaybeInsideGraphOutput: + verify_plans: tuple[VerifyPlan, ...] + write_plans: tuple[WritePlan, ...] + expected_inputs: ExpectedInputs + + +class SingleForwardManager: + """Owns the state of one inner ``model.forward`` invocation.""" + + def __init__( + self, + *, + config: CanaryConfig, + device: torch.device, + device_state: CanaryDeviceState, + buffer_groups: tuple[CanaryBufferGroup, ...], + endpoints: tuple[CanaryEndpoint, ...], + req_to_token_pool: "ReqToTokenPool", + swa_window_size: int, + per_forward_verify_capacity: int, + per_forward_write_req_capacity: int, + per_forward_write_entry_capacity: int, + d2h_stream: torch.cuda.Stream, + ) -> None: + self._config = config + self._device = device + self._device_state = device_state + self._buffer_groups = buffer_groups + self._endpoints = endpoints + self._req_to_token_pool = req_to_token_pool + self._swa_window_size = swa_window_size + self._d2h_stream = d2h_stream + + self._write_req_capacity = per_forward_write_req_capacity + self._write_entry_capacity = per_forward_write_entry_capacity + self._verify_capacity = per_forward_verify_capacity + + self._enable_warner = CanaryEnableWarner( + verify_capacity=self._verify_capacity, + d2h_stream=d2h_stream, + ) + + self._phase_checker = SimplePhaseChecker( + initial_phase=_SingleForwardPhase.IDLE, device=device + ) + + self._output_buffer = PostOpsInsideGraphOutputBuffer.allocate( + num_kernel_tags=int(device_state.kernel_run_counters.shape[0]), + num_slot_tags=int(device_state.slot_run_counters.shape[0]), + device=device, + ) + + @property + def phase_checker(self) -> SimplePhaseChecker: + return self._phase_checker + + def pre_ops_outside_graph( + self, *, maybe_inaccurate_forward_batch: "ForwardBatch" + ) -> None: + self._phase_checker.update( + expect_phase=_SingleForwardPhase.IDLE, + next_phase=_SingleForwardPhase.AFTER_PRE_OUT, + caller_name="SingleForwardManager.pre_ops_outside_graph", + ) + + bs = int(maybe_inaccurate_forward_batch.batch_size) + num_tokens = int(maybe_inaccurate_forward_batch.positions.shape[0]) + if bs > self._write_req_capacity: + raise RuntimeError( + f"kv-canary: forward_batch.batch_size={bs} exceeds pre-allocated " + f"write_req_capacity={self._write_req_capacity}; raise --cuda-graph-max-bs " + f"or check CanaryLaunchCapacities.from_args" + ) + if num_tokens > self._write_entry_capacity: + raise RuntimeError( + f"kv-canary: forward_batch token count={num_tokens} exceeds pre-allocated " + f"write_entry_capacity={self._write_entry_capacity}; raise " + f"--chunked-prefill-size / --max-prefill-tokens or check " + f"CanaryLaunchCapacities.from_args" + ) + + def pre_ops_maybe_inside_graph( + self, forward_batch: "ForwardBatch" + ) -> "_PreOpsMaybeInsideGraphOutput": + self._phase_checker.update( + expect_phase=_SingleForwardPhase.AFTER_PRE_OUT, + next_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN, + caller_name="SingleForwardManager.pre_ops_maybe_inside_graph", + ) + + verify_plans = tuple( + VerifyPlan.allocate( + verify_capacity=self._verify_capacity, device=self._device + ) + for _ in self._buffer_groups + ) + write_plans = tuple( + WritePlan.allocate( + write_req_capacity=self._write_req_capacity, device=self._device + ) + for _ in self._buffer_groups + ) + expected_inputs = ExpectedInputs.allocate( + capacity=self._write_entry_capacity, device=self._device + ) + plan_input = PlanInput.allocate( + bs_capacity=self._write_req_capacity, device=self._device + ) + + plan_input.fill_from_forward_batch(forward_batch=forward_batch) + + violation_log = self._device_state.violation_log + num_tokens = int(forward_batch.positions.shape[0]) + expected_inputs_slice = expected_inputs.slice(num_tokens) + + for group_idx, group in enumerate(self._buffer_groups): + verify_plan = verify_plans[group_idx] + write_plan = write_plans[group_idx] + invoke_plan( + plan_input=plan_input, + verify_plan=verify_plan, + write_plan=write_plan, + group=group, + req_to_token=self._req_to_token_pool.req_to_token, + swa_window_size=self._swa_window_size, + req_to_verify_expected_tokens=self._device_state.req_to_verify_expected_tokens, + ) + launch_endpoints_per_forward( + endpoints=self._endpoints, + group=group, + tag_filter=_is_head_tag, + verify_plan=verify_plan, + write_plan=write_plan, + forward_batch=forward_batch, + expected_inputs=expected_inputs_slice, + violation_log=violation_log, + ) + + return _PreOpsMaybeInsideGraphOutput( + verify_plans=verify_plans, + write_plans=write_plans, + expected_inputs=expected_inputs, + ) + + def post_ops_maybe_inside_graph( + self, + forward_batch: "ForwardBatch", + pre_ops_output: "_PreOpsMaybeInsideGraphOutput", + ) -> None: + self._phase_checker.update( + expect_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN, + next_phase=_SingleForwardPhase.AFTER_POST_MAYBE_IN, + caller_name="SingleForwardManager.post_ops_maybe_inside_graph", + ) + + violation_log = self._device_state.violation_log + num_tokens = int(forward_batch.positions.shape[0]) + expected_inputs_slice = pre_ops_output.expected_inputs.slice(num_tokens) + for group_idx, group in enumerate(self._buffer_groups): + launch_endpoints_per_forward( + endpoints=self._endpoints, + group=group, + tag_filter=_is_tail_tag, + verify_plan=pre_ops_output.verify_plans[group_idx], + write_plan=pre_ops_output.write_plans[group_idx], + forward_batch=forward_batch, + expected_inputs=expected_inputs_slice, + violation_log=violation_log, + ) + + verify_plan_enable_combined = _torch_reduce_minimum( + [x.enable for x in pre_ops_output.verify_plans] + ) + self._output_buffer.copy_from( + verify_plan_enable=verify_plan_enable_combined, + kernel_run_counters=self._device_state.kernel_run_counters, + slot_run_counters=self._device_state.slot_run_counters, + violation_write_index=self._device_state.violation_log.violation_write_index, + ) + + def post_ops_outside_graph(self) -> None: + self._phase_checker.update( + expect_phase=_SingleForwardPhase.AFTER_POST_MAYBE_IN, + next_phase=_SingleForwardPhase.IDLE, + caller_name="SingleForwardManager.post_ops_outside_graph", + ) + + self._enable_warner.tick(self._output_buffer.verify_plan_enable) + + +def _is_head_tag(tag: CanaryLaunchTag) -> bool: + return tag in ( + CanaryLaunchTag.HEAD_K_FULL, + CanaryLaunchTag.HEAD_V_FULL, + CanaryLaunchTag.HEAD_K_SWA, + CanaryLaunchTag.HEAD_V_SWA, + ) + + +def _is_tail_tag(tag: CanaryLaunchTag) -> bool: + return tag in ( + CanaryLaunchTag.TAIL_K_FULL, + CanaryLaunchTag.TAIL_V_FULL, + CanaryLaunchTag.TAIL_K_SWA, + CanaryLaunchTag.TAIL_V_SWA, + ) diff --git a/python/sglang/srt/kv_canary/state.py b/python/sglang/srt/kv_canary/state.py new file mode 100644 index 000000000..7d19c7216 --- /dev/null +++ b/python/sglang/srt/kv_canary/state.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from sglang.jit_kernel.kv_canary.consts import VIOLATION_FIELDS +from sglang.srt.kv_canary.config import CanaryConfig + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ViolationLog: + """Global violation sink shared across all canary launches. + + One instance per canary runner — every launch (head / tail, K / V half, FULL / SWA group) writes + into the same ring. The kernel_kind field stamped into each violation row identifies which launch fired + (kernel_kind is a static IntEnum tag — :class:`CanaryLaunchTag` in + ``sglang.jit_kernel.kv_canary.verify`` — with a unique value per (head|tail, K|V, FULL|SWA) tuple). + + Ring capacity is sized generously (≥ 1024) so overflow is a non-concern in practice — violations are + cold-path and the host raises at the first one anyway (or just logs it in mode="log"). atomicAdd + contention on a single counter is also negligible since violation events are rare. + + Derived state (host computes on read; not stored): + is_errored = violation_write_index[0] > 0 + first_violation = violation_ring[0] (valid iff is_errored) + ring_valid_count = min(violation_write_index[0], ring_capacity) + + The ring is fill-once: writes beyond ring_capacity are dropped but the counter still increments. Whoever + wins atomicAdd for idx == 0 permanently occupies row 0. + + Fields: + violation_ring: Append-only violation sink, shape [ring_capacity, VIOLATION_FIELDS], int64. Row 0 is + the first violation; rows 1..min(write_index, capacity) follow in atomic order. Fill-once. + violation_write_index: Monotonic violation counter, shape [1], int32. Incremented on every violation + regardless of ring capacity. + """ + + violation_ring: torch.Tensor + violation_write_index: torch.Tensor + + @classmethod + def allocate(cls, *, ring_capacity: int, device: torch.device) -> "ViolationLog": + if ring_capacity <= 0: + raise ValueError( + f"kv-canary: ViolationLog ring_capacity must be positive, got {ring_capacity}" + ) + return cls( + violation_ring=torch.zeros( + ring_capacity, VIOLATION_FIELDS, dtype=torch.int64, device=device + ), + violation_write_index=torch.zeros(1, dtype=torch.int32, device=device), + ) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CanaryDeviceState: + """Device-side state owned by one CanaryManager instance. + + One instance per ModelRunner. Held on the same device as the KV pool. All tensors are allocated up + front (sizes fixed by CanaryConfig + cuda-graph capture capacity) and reused across forward steps — + no per-step allocation. + + Fields: + violation_log: The single ViolationLog shared by every launch (head / tail × K / V × + FULL / SWA). All kernels atomicAdd into violation_log.violation_write_index and stamp their + CanaryLaunchTag into each violation row. + kernel_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. The + kernel itself does NOT index this array; runner takes a 1-element view at tag's slot (via + CanaryEndpoint.kernel_run_counter_view) and hands a shape [1] tensor to the kernel, + which atomicAdds 1 regardless of whether the plan had any active entry. Health watchdog + reads this array to confirm "canary path actually ran". + slot_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. Same + view-handed-to-kernel pattern as kernel_run_counters; each launch adds its active entry + count to its slot. Used for periodic stats ("protected N tokens"). + enable_chain_position_assert: int32 [1] device flag gating the write kernel's chain-step + write_position assert. allocate() defaults to 1; CanaryManager zeros it during + __init__ for the warmup window and mark_init_finished() flips it back to 1. + req_to_verify_expected_tokens: Optional int32 device tensor shape + ``[req_to_token_alloc_size, max_context_len]``. Mirrors ReqToTokenPool layout; + ``pool[req_idx, p]`` = source-of-truth token at logical position ``p`` for the + req in slot ``req_idx``. The plan-side entries kernel gathers from this pool (via + ``kv_token_id_vs_position_offset`` per buffer group) into + ``VerifyPlan.verify_expected_tokens``; the verify kernel then compares against each + canary slot's stored token. The naive build never populates the verify-token-id + cross-check, so this is always ``None`` and the gather degrades to the ``-1`` sentinel. + """ + + violation_log: ViolationLog + kernel_run_counters: torch.Tensor + slot_run_counters: torch.Tensor + enable_chain_position_assert: torch.Tensor + req_to_verify_expected_tokens: Optional[torch.Tensor] + + @classmethod + def allocate( + cls, + *, + config: CanaryConfig, + device: torch.device, + num_tags: int, + req_to_token_alloc_size: Optional[int] = None, + max_context_len: Optional[int] = None, + ) -> "CanaryDeviceState": + if num_tags <= 0: + raise ValueError( + f"kv-canary: CanaryDeviceState num_tags must be positive, got {num_tags}" + ) + violation_log = ViolationLog.allocate( + ring_capacity=config.ring_capacity, device=device + ) + kernel_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device) + slot_run_counters = torch.zeros(num_tags, dtype=torch.int64, device=device) + enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device) + # The naive build does not run the verify-token-id cross-check, so the source-of-truth + # token pool is never allocated. The field is kept on the dataclass and downstream code + # (plan kernel gather) treats ``None`` as "emit the -1 skip sentinel". + req_to_verify_expected_tokens = None + return cls( + violation_log=violation_log, + kernel_run_counters=kernel_run_counters, + slot_run_counters=slot_run_counters, + enable_chain_position_assert=enable_chain_position_assert, + req_to_verify_expected_tokens=req_to_verify_expected_tokens, + ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a978df1fd..c7d27fe72 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -40,6 +40,7 @@ from sglang.srt.distributed.parallel_state import ( get_moe_expert_parallel_world_size, get_tensor_model_parallel_world_size, ) +from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( DpPaddingMode, get_attention_cp_size, @@ -448,6 +449,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For ngram embedding ngram_embedding_info: Optional[NgramEmbeddingInfo] = None + # kv-canary token-id validator snapshot + req_all_ids_flat: Optional[torch.Tensor] = None + req_all_ids_lens: Optional[torch.Tensor] = None + @classmethod def init_new( cls, @@ -1314,3 +1319,5 @@ if is_cuda() or is_hip(): clamp_position = clamp_position_cuda else: clamp_position = _clamp_position_native + + diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 76d9b440f..efd65e23f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -777,6 +777,7 @@ class ServerArgs: enable_attn_tp_input_scattered: bool = False disable_attn_tp_gather: bool = False gc_threshold: Optional[List[int]] = None + kv_canary: str = "none" # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 enable_dsa_prefill_context_parallel: bool = False dsa_prefill_cp_mode: str = "round-robin-split" @@ -6305,6 +6306,18 @@ class ServerArgs: action="store_true", help="Disable RadixAttention for prefix caching.", ) + parser.add_argument( + "--kv-canary", + type=str, + default=ServerArgs.kv_canary, + choices=["none", "log", "raise"], + help=( + "KV cache canary mode. " + "'none' disables the canary (default). " + "'log' prints them while the server keeps running (production-safe). " + "'raise' fails the server on the first detected mismatch (CI lane)." + ), + ) parser.add_argument( "--cuda-graph-max-bs", type=int, diff --git a/python/sglang/test/kv_canary/__init__.py b/python/sglang/test/kv_canary/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/test/kv_canary/e2e_base.py b/python/sglang/test/kv_canary/e2e_base.py new file mode 100644 index 000000000..21d244e87 --- /dev/null +++ b/python/sglang/test/kv_canary/e2e_base.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import io +import os +import string +from typing import ClassVar, Literal, Optional + +from sglang.srt.kv_canary.config import CanaryMode +from sglang.srt.utils import kill_process_tree +from sglang.test.kv_canary.mode_config import _MODE_CONFIGS, _ModeConfig +from sglang.test.kv_canary.utils import build_canary_server_args, post_parallel_generate +from sglang.test.kv_canary.violation_assert_mixin import CanaryViolationAssertMixin +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +# Long prompt body shared by all canary e2e tests. The repetition count is chosen +# so the tokenised prompt is comfortably long; token count is roughly 7k after BPE. +_LONG_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 700).strip() +_UNIQUE_PROMPT_FIRST_CHARS = string.ascii_letters + string.digits + + +class CapturedServerE2EBase(CanaryViolationAssertMixin, CustomTestCase): + process: ClassVar[Optional[object]] = None + base_url: ClassVar[str] = DEFAULT_URL_FOR_TEST + _stdout_buf: ClassVar[Optional[io.StringIO]] = None + _stderr_buf: ClassVar[Optional[io.StringIO]] = None + + @classmethod + def tearDownClass(cls) -> None: + if cls.process is not None: + kill_process_tree(cls.process.pid) + for buf in (cls._stdout_buf, cls._stderr_buf): + if buf is not None: + buf.close() + cls._stdout_buf = None + cls._stderr_buf = None + + def _captured_log_text( + self, side: Optional[Literal["prefill", "decode"]] = None + ) -> str: + stdout_text = ( + self._stdout_buf.getvalue() if self._stdout_buf is not None else "" + ) + stderr_text = ( + self._stderr_buf.getvalue() if self._stderr_buf is not None else "" + ) + return stdout_text + stderr_text + + def assert_log_contains(self, substring: str) -> None: + log_text = self._captured_log_text() + if substring not in log_text: + raise AssertionError( + f"Expected substring {substring!r} not found in captured log. " + f"Log tail:\n{log_text[-2000:]}" + ) + + +class CanaryE2EBase(CapturedServerE2EBase): + model_mode: ClassVar[Literal["mha"]] + kv_canary_mode: ClassVar[CanaryMode] + extra_env: ClassVar[dict[str, str]] = {} + extra_server_args: ClassVar[tuple[str, ...]] = () + use_unique_prompts: ClassVar[bool] = False + # Number of sequential request batches each test method sends. Default 1 keeps tests fast. + workload_n_batches: ClassVar[int] = 1 + + _cfg: ClassVar[Optional[_ModeConfig]] = None + + @classmethod + def setUpClass(cls) -> None: + cls._cfg = _MODE_CONFIGS[cls.model_mode] + server_env = os.environ.copy() + server_env.update(cls.extra_env) + + cls._stdout_buf = io.StringIO() + cls._stderr_buf = io.StringIO() + + server_args = build_canary_server_args( + kv_canary_mode=cls.kv_canary_mode, + mode_cfg=cls._cfg, + extra_server_args=( + "--max-total-tokens", + "65536", + "--skip-server-warmup", + *cls.extra_server_args, + ), + ) + cls.process = popen_launch_server( + cls._cfg.model_path, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=server_args, + env=server_env, + return_stdout_stderr=(cls._stdout_buf, cls._stderr_buf), + ) + + def make_prompts(self, n: int) -> list[str]: + if self.use_unique_prompts: + return _make_unique_prompts(n) + return [_LONG_PROMPT_BODY] * n + + def send_parallel_requests( + self, + n: int = 8, + *, + assert_all_success: bool = True, + max_new_tokens: int = 2048, + timeout: float = 240.0, + ) -> list[dict]: + """Fan out n parallel /generate requests; return list of response dicts.""" + results = post_parallel_generate( + url=self.base_url + "/generate", + prompts=self.make_prompts(n), + max_new_tokens=max_new_tokens, + timeout=timeout, + ) + if assert_all_success: + for result in results: + self.assertEqual(result.get("status_code"), 200, result) + return results + + +def _make_unique_prompts(n: int) -> list[str]: + if n > len(_UNIQUE_PROMPT_FIRST_CHARS): + raise ValueError( + f"unique prompt count {n} exceeds supported count " + f"{len(_UNIQUE_PROMPT_FIRST_CHARS)}" + ) + + return [ + ( + f"{_UNIQUE_PROMPT_FIRST_CHARS[i]}" + f"{hex(i * 0x9E3779B1 & 0xFFFFFFFF)[2:]} " + f"{_LONG_PROMPT_BODY}" + ) + for i in range(n) + ] diff --git a/python/sglang/test/kv_canary/fixtures.py b/python/sglang/test/kv_canary/fixtures.py new file mode 100644 index 000000000..5ade24525 --- /dev/null +++ b/python/sglang/test/kv_canary/fixtures.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import List, Optional + +import torch + +from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode +from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha +from sglang.srt.kv_canary.pool_patcher.api import register_pool_attacher +from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode +from sglang.srt.model_executor.forward_batch_info import ForwardMode + +DEFAULT_DEVICE: torch.device = torch.device("cuda") + + +@dataclass +class FakeMHAPool: + layer_num: int + k_buffer: List[torch.Tensor] + v_buffer: List[torch.Tensor] + page_size: int = 1 + + def get_contiguous_buf_infos(self): + ptrs = [b.data_ptr() for b in self.k_buffer] + [ + b.data_ptr() for b in self.v_buffer + ] + lens = [b.nbytes for b in self.k_buffer] + [b.nbytes for b in self.v_buffer] + item_lens = [b[0].nbytes * self.page_size for b in self.k_buffer] + [ + b[0].nbytes * self.page_size for b in self.v_buffer + ] + return ptrs, lens, item_lens + + +def make_mha_pool( + device: torch.device = DEFAULT_DEVICE, + *, + num_slots: int = 16, + dim: int = 8, + layer_num: int = 2, +) -> FakeMHAPool: + k_layers = [ + torch.zeros(num_slots, dim, dtype=torch.float16, device=device) + for _ in range(layer_num) + ] + v_layers = [ + torch.zeros(num_slots, dim, dtype=torch.float16, device=device) + for _ in range(layer_num) + ] + return FakeMHAPool(layer_num=layer_num, k_buffer=k_layers, v_buffer=v_layers) + + +def make_base_config() -> CanaryConfig: + return CanaryConfig( + mode=CanaryMode.RAISE, + ring_capacity=1024, + ) + + +def make_req_to_token_pool( + device: torch.device = DEFAULT_DEVICE, + *, + max_reqs: int = 8, + max_seq_len: int = 32, +) -> SimpleNamespace: + req_to_token = torch.zeros(max_reqs, max_seq_len, dtype=torch.int32, device=device) + return SimpleNamespace( + req_to_token=req_to_token, size=max_reqs, max_context_len=max_seq_len + ) + + +def make_forward_batch( + device: torch.device = DEFAULT_DEVICE, + *, + bs: int = 2, + seq_lens_list: tuple[int, ...] = (3, 4), + req_pool_indices: Optional[torch.Tensor] = None, + seq_lens: Optional[torch.Tensor] = None, + seq_lens_sum: Optional[int] = None, + extend_prefix_lens: Optional[torch.Tensor] = None, + extend_prefix_lens_cpu: Optional[list] = None, + extend_seq_lens: Optional[torch.Tensor] = None, + extend_seq_lens_cpu: Optional[list] = None, + is_extend: bool = False, + is_target_verify: bool = False, + is_draft_extend_v2: bool = False, + spec_info: Optional[object] = None, + input_ids: Optional[torch.Tensor] = None, + positions: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + num_token_non_padded_cpu: Optional[int] = None, +) -> SimpleNamespace: + seq_lens_default = list(seq_lens_list[:bs]) + if req_pool_indices is None: + req_pool_indices = torch.tensor([1, 2][:bs], dtype=torch.int64, device=device) + if seq_lens is None: + seq_lens = torch.tensor(seq_lens_default, dtype=torch.int32, device=device) + if seq_lens_sum is None: + seq_lens_sum = int(sum(seq_lens_default)) + if input_ids is None: + input_ids = torch.zeros(bs, dtype=torch.int32, device=device) + if positions is None: + # Default to decode-canonical: positions = seq_lens - 1 (one-token-per-req decode write + # at the post-bump tail). Plan input derives decode prefix_lens from positions + # directly, so the default must keep parity. + positions = (seq_lens.to(torch.int64) - 1).clamp(min=0).to(torch.int32) + if out_cache_loc is None: + out_cache_loc = torch.zeros(bs, dtype=torch.int32, device=device) + + if is_extend: + forward_mode = ForwardMode.EXTEND + elif is_target_verify: + forward_mode = ForwardMode.TARGET_VERIFY + elif is_draft_extend_v2: + forward_mode = ForwardMode.DRAFT_EXTEND_V2 + else: + forward_mode = ForwardMode.DECODE + return SimpleNamespace( + forward_mode=forward_mode, + batch_size=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_sum=seq_lens_sum, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + spec_info=spec_info, + input_ids=input_ids, + positions=positions, + out_cache_loc=out_cache_loc, + num_token_non_padded_cpu=num_token_non_padded_cpu, + req_all_ids_flat=None, + req_all_ids_lens=None, + ) + + +def make_buffer_group( + *, + device: torch.device = DEFAULT_DEVICE, + kind: PoolKind = PoolKind.FULL, + has_v: bool = True, + swa_index_lut: Optional[torch.Tensor] = None, + num_slots: int = 4, + kv_token_id_vs_position_offset: int = 0, +) -> CanaryBufferGroup: + def _zero() -> torch.Tensor: + return torch.zeros( + num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device + ) + + return CanaryBufferGroup( + kind=kind, + k_head=_zero(), + k_tail=_zero(), + v_head=_zero() if has_v else None, + v_tail=_zero() if has_v else None, + swa_index_lut=swa_index_lut, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + + +def make_radix_cache( + slot_lists: List[List[int]], device: torch.device = DEFAULT_DEVICE +): + cache = RadixCache.__new__(RadixCache) + cache.device = device + cache.page_size = 1 + cache.disable = False + + root = TreeNode() + root.value = torch.tensor( + slot_lists[0] if slot_lists else [], dtype=torch.int32, device=device + ) + cache.root_node = root + + current = root + for child_slots in slot_lists[1:]: + child = TreeNode() + child.value = torch.tensor(child_slots, dtype=torch.int32, device=device) + child.parent = current + current.children[child.id] = child + current = child + + return cache + + +register_pool_attacher(FakeMHAPool, attach_mha) diff --git a/python/sglang/test/kv_canary/mode_config.py b/python/sglang/test/kv_canary/mode_config.py new file mode 100644 index 000000000..1e60ad492 --- /dev/null +++ b/python/sglang/test/kv_canary/mode_config.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ModeConfig: + model_path: str + json_model_override_args: Optional[str] = None + + +_MODE_CONFIGS: dict[str, _ModeConfig] = { + "mha": _ModeConfig( + model_path="Qwen/Qwen3-0.6B", + ), +} diff --git a/python/sglang/test/kv_canary/runner_test_base.py b/python/sglang/test/kv_canary/runner_test_base.py new file mode 100644 index 000000000..a55c9805b --- /dev/null +++ b/python/sglang/test/kv_canary/runner_test_base.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag +from sglang.srt.kv_canary import endpoint as endpoint_module +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup +from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities +from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode +from sglang.srt.kv_canary.runner import kernel_launcher as kernel_launcher_module +from sglang.srt.kv_canary.runner.canary_manager import CanaryManager +from sglang.test.kv_canary.fixtures import ( + DEFAULT_DEVICE, + make_buffer_group, + make_req_to_token_pool, +) +from sglang.test.test_utils import CustomTestCase + + +def make_config( + *, + mode: CanaryMode = CanaryMode.RAISE, + ring_capacity: int = 1024, +) -> CanaryConfig: + return CanaryConfig( + mode=mode, + ring_capacity=ring_capacity, + ) + + +class RecordingEndpoint: + def __init__(self, *, kernel_kind: CanaryLaunchTag) -> None: + self.kernel_kind = kernel_kind + self.calls: list[dict[str, object]] = [] + + def launch_per_forward(self, **kwargs: object) -> None: + self.calls.append(kwargs) + + +def make_manager( + *, + device: torch.device, + config: CanaryConfig | None = None, + group: CanaryBufferGroup | None = None, + req_pool: SimpleNamespace | None = None, + per_forward_verify_capacity: int = 16, +) -> CanaryManager: + if config is None: + config = make_config() + if group is None: + group = make_buffer_group(device=device) + if req_pool is None: + req_pool = make_req_to_token_pool(device=device, max_reqs=4, max_seq_len=8) + return CanaryManager( + config=config, + buffer_groups=(group,), + device=device, + req_to_token_pool=req_pool, + launch_capacities=CanaryLaunchCapacities( + per_forward_verify_capacity=per_forward_verify_capacity, + per_forward_write_req_capacity=2, + per_forward_write_entry_capacity=8, + ), + ) + + +class CanaryManagerTestCase(CustomTestCase): + def setUp(self) -> None: + self.device = DEFAULT_DEVICE + # Stub plan/verify/write kernels so CPU runs don't need CUDA JIT. + self._patchers = [ + patch.object( + kernel_launcher_module, + "launch_canary_plan_kernels", + lambda **kwargs: None, + ), + patch.object( + endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None + ), + patch.object( + endpoint_module, "launch_canary_write_kernel", lambda **kwargs: None + ), + ] + for patcher in self._patchers: + patcher.start() + self.addCleanup(patcher.stop) diff --git a/python/sglang/test/kv_canary/utils.py b/python/sglang/test/kv_canary/utils.py new file mode 100644 index 000000000..0a4f76b36 --- /dev/null +++ b/python/sglang/test/kv_canary/utils.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor + +import requests + +from sglang.srt.kv_canary.config import CanaryMode +from sglang.test.kv_canary.mode_config import _ModeConfig + + +def build_canary_server_args( + *, + kv_canary_mode: CanaryMode, + mode_cfg: _ModeConfig, + extra_server_args: tuple[str, ...] = (), +) -> list[str]: + args = [ + "--kv-canary", + kv_canary_mode.value, + "--disable-piecewise-cuda-graph", + "--context-length", + "16384", + *extra_server_args, + ] + if mode_cfg.json_model_override_args is not None: + args.extend(["--json-model-override-args", mode_cfg.json_model_override_args]) + return args + + +def post_parallel_generate( + *, + url: str, + prompts: list[str], + max_new_tokens: int, + timeout: float, +) -> list[dict]: + def _send(prompt: str) -> dict: + try: + resp = requests.post( + url, + json={ + "text": prompt, + "sampling_params": { + "max_new_tokens": max_new_tokens, + "temperature": 0.0, + }, + }, + timeout=timeout, + ) + return {"status_code": resp.status_code, "body": resp.text} + except requests.RequestException as exc: + return {"status_code": -1, "error": repr(exc)} + + with ThreadPoolExecutor(max_workers=max(1, len(prompts))) as pool: + return list(pool.map(_send, prompts)) diff --git a/python/sglang/test/kv_canary/violation_assert_mixin.py b/python/sglang/test/kv_canary/violation_assert_mixin.py new file mode 100644 index 000000000..808203acc --- /dev/null +++ b/python/sglang/test/kv_canary/violation_assert_mixin.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import time +from typing import Literal, Optional + +from sglang.test.kv_canary.violation_log_utils import ( + assert_no_violation_in_log, + find_violation_in_log, +) + +_Side = Optional[Literal["prefill", "decode"]] + + +class CanaryViolationAssertMixin: + def _captured_log_text(self, side: _Side = None) -> str: + raise NotImplementedError + + def assert_per_forward_violation_reported( + self, + *, + fail_reason: str, + side: _Side = None, + flush_wait_seconds: float = 2.0, + ) -> None: + self.assert_violation_logged_any( + launch_tag_patterns=("HEAD_*", "TAIL_*"), + fail_reason=fail_reason, + side=side, + flush_wait_seconds=flush_wait_seconds, + ) + + def assert_any_launch_tag_violation_reported( + self, + *, + fail_reason: str, + side: _Side = None, + flush_wait_seconds: float = 3.0, + max_retries: int = 10, + ) -> None: + self.assert_violation_logged_any( + launch_tag_patterns=("*",), + fail_reason=fail_reason, + side=side, + flush_wait_seconds=flush_wait_seconds, + max_retries=max_retries, + ) + + def assert_any_launch_tag_violation_absent( + self, *, fail_reason: str, side: _Side = None + ) -> None: + self.assert_no_violation_matching( + launch_tag_patterns=("*",), fail_reason=fail_reason, side=side + ) + + def assert_violation_logged_any( + self, + *, + launch_tag_patterns: tuple[str, ...], + fail_reason: str, + side: _Side = None, + flush_wait_seconds: float = 2.0, + max_retries: int = 1, + ) -> None: + log_text = "" + for _ in range(max_retries): + time.sleep(flush_wait_seconds) + log_text = self._captured_log_text(side) + if find_violation_in_log( + log_text, + launch_tag_patterns=launch_tag_patterns, + fail_reason=fail_reason, + ): + return + side_label = "" if side is None else f" on side={side}" + other_side_diag = "" + if side in ("prefill", "decode"): + other = "decode" if side == "prefill" else "prefill" + try: + other_text = self._captured_log_text(other) + other_match = find_violation_in_log( + other_text, + launch_tag_patterns=launch_tag_patterns, + fail_reason=fail_reason, + ) + other_side_diag = ( + f"\n[diag] other side ({other}) buf len={len(other_text)} " + f"contains_match={other_match}" + ) + except (NotImplementedError, ValueError): + pass + raise AssertionError( + f"No canary violation matching launch_tag_patterns={launch_tag_patterns!r} " + f"fail_reason={fail_reason!r}{side_label} after max_retries={max_retries} " + f"(wait={flush_wait_seconds}s each). " + f"log_text len={len(log_text)}.{other_side_diag} Log tail:\n" + f"{log_text[-2000:]}" + ) + + def assert_no_violation_matching( + self, + *, + launch_tag_patterns: tuple[str, ...], + fail_reason: str, + side: _Side = None, + ) -> None: + log_text = self._captured_log_text(side) + if find_violation_in_log( + log_text, + launch_tag_patterns=launch_tag_patterns, + fail_reason=fail_reason, + ): + raise AssertionError( + f"Unexpected canary violation matching " + f"launch_tag_patterns={launch_tag_patterns!r} " + f"fail_reason={fail_reason!r}. Log tail:\n{log_text[-2000:]}" + ) + + def assert_no_violation( + self, + *, + side: _Side = None, + wait_seconds: float = 2.0, + ) -> None: + time.sleep(wait_seconds) + assert_no_violation_in_log(self._captured_log_text(side)) diff --git a/python/sglang/test/kv_canary/violation_log_utils.py b/python/sglang/test/kv_canary/violation_log_utils.py new file mode 100644 index 000000000..4d2e64d17 --- /dev/null +++ b/python/sglang/test/kv_canary/violation_log_utils.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import fnmatch +import re + +_VIOLATION_LINE_RE = re.compile( + r"kv_canary violation: launch_tag=(\S+) fail_reason=(\S+)" +) + + +def find_violation_in_log( + log_text: str, + *, + launch_tag_patterns: tuple[str, ...], + fail_reason: str, +) -> bool: + for match in _VIOLATION_LINE_RE.finditer(log_text): + tag = match.group(1) + reason_field = match.group(2) + if fail_reason not in reason_field.split("+"): + continue + if any(fnmatch.fnmatchcase(tag, pattern) for pattern in launch_tag_patterns): + return True + return False + + +def assert_no_violation_in_log(log_text: str) -> None: + if "kv_canary violation:" in log_text: + raise AssertionError( + f"Unexpected canary violation found. Log tail:\n{log_text[-2000:]}" + ) diff --git a/test/manual/kv_canary/__init__.py b/test/manual/kv_canary/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/kv_canary/test_e2e_base.py b/test/registered/kv_canary/test_e2e_base.py new file mode 100644 index 000000000..452dcfb99 --- /dev/null +++ b/test/registered/kv_canary/test_e2e_base.py @@ -0,0 +1,31 @@ +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.kv_canary.e2e_base import ( + _LONG_PROMPT_BODY, + _UNIQUE_PROMPT_FIRST_CHARS, + _make_unique_prompts, +) +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-b-test-cpu") + + +class TestCanaryE2EBase(CustomTestCase): + def test_make_unique_prompts_have_distinct_first_characters(self) -> None: + """Verify generated prompts use distinct first characters and all end with the shared long body.""" + prompts = _make_unique_prompts(8) + + self.assertEqual(len({prompt[0] for prompt in prompts}), len(prompts)) + self.assertTrue(all(prompt.endswith(_LONG_PROMPT_BODY) for prompt in prompts)) + + def test_make_unique_prompts_rejects_more_prompts_than_distinct_first_characters( + self, + ) -> None: + """Verify prompt generation rejects requests beyond the unique prefix budget.""" + with self.assertRaisesRegex(ValueError, "unique prompt count"): + _make_unique_prompts(len(_UNIQUE_PROMPT_FIRST_CHARS) + 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_buffer_alloc.py b/test/registered/kv_canary/test_self_unit_buffer_alloc.py new file mode 100644 index 000000000..db8930186 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_buffer_alloc.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import unittest + +import torch + +from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES +from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small") + + +class TestAllocCanaryBuf(CustomTestCase): + def test_alloc_canary_buf_shape_and_dtype(self) -> None: + """Verify alloc_canary_buf returns a zeroed uint8 buffer of [num_slots, CANARY_SLOT_BYTES].""" + buf = alloc_canary_buf(num_slots=8, device=DEFAULT_DEVICE) + self.assertEqual(buf.shape, (8, CANARY_SLOT_BYTES)) + self.assertEqual(buf.dtype, torch.uint8) + self.assertEqual(buf.device.type, DEFAULT_DEVICE.type) + self.assertTrue(torch.equal(buf, torch.zeros_like(buf))) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_capacities.py b/test/registered/kv_canary/test_self_unit_capacities.py new file mode 100644 index 000000000..510f6120b --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_capacities.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small") + + +class TestComputeLaunchCapacities(CustomTestCase): + @staticmethod + def _make_server_args(*, max_bs: int) -> SimpleNamespace: + return SimpleNamespace( + cuda_graph_max_bs=max_bs, + speculative_num_draft_tokens=0, + chunked_prefill_size=None, + max_prefill_tokens=128, + ) + + @staticmethod + def _from_args( + *, + max_bs: int, + max_seq_len: int, + max_total_num_tokens: int | None = None, + ) -> CanaryLaunchCapacities: + if max_total_num_tokens is None: + max_total_num_tokens = max_bs * max_seq_len + return CanaryLaunchCapacities.from_args( + server_args=TestComputeLaunchCapacities._make_server_args(max_bs=max_bs), + req_to_token_pool_size=max_bs, + max_seq_len_per_req=max_seq_len, + pool_slot_count=max_total_num_tokens, + ) + + def test_per_forward_verify_capacity_covers_multi_req_prefix_sum(self) -> None: + """Verify per-forward verify capacity equals max_total_num_tokens * 3.""" + max_bs = 8 + max_seq_len = 64 + max_total_num_tokens = 1024 + capacities = self._from_args( + max_bs=max_bs, + max_seq_len=max_seq_len, + max_total_num_tokens=max_total_num_tokens, + ) + self.assertEqual( + capacities.per_forward_verify_capacity, + int(max_total_num_tokens * 3), + ) + + def test_from_args_treats_missing_speculative_draft_tokens_as_zero(self) -> None: + """per_forward_write_entry_capacity is floored by max_prefill_tokens when batch * tokens_per_bs is smaller.""" + server_args = self._make_server_args(max_bs=2) + server_args.speculative_num_draft_tokens = None + + capacities = CanaryLaunchCapacities.from_args( + server_args=server_args, + req_to_token_pool_size=2, + max_seq_len_per_req=32, + pool_slot_count=64, + ) + + self.assertEqual(capacities.per_forward_write_entry_capacity, 128) + + def test_manual_capacities_reject_non_positive_fields(self) -> None: + """Verify manual launch capacities fail instead of being clamped.""" + with self.assertRaisesRegex(ValueError, "per_forward_verify_capacity"): + CanaryLaunchCapacities( + per_forward_verify_capacity=0, + per_forward_write_req_capacity=1, + per_forward_write_entry_capacity=1, + ) + + def test_from_args_rejects_empty_pool_capacity(self) -> None: + """Verify derived launch capacities reject invalid pool sizing.""" + with self.assertRaisesRegex(ValueError, "pool_slot_count"): + CanaryLaunchCapacities.from_args( + server_args=self._make_server_args(max_bs=1), + req_to_token_pool_size=1, + max_seq_len_per_req=1, + pool_slot_count=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_endpoint.py b/test/registered/kv_canary/test_self_unit_endpoint.py new file mode 100644 index 000000000..521d7a478 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_endpoint.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.jit_kernel.kv_canary.verify import ( + CANARY_SLOT_BYTES, + CanaryLaunchTag, + VerifyPlan, +) +from sglang.jit_kernel.kv_canary.write import WritePlan +from sglang.srt.kv_canary import endpoint as endpoint_module +from sglang.srt.kv_canary.endpoint import ( + CanaryEndpoint, +) +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.state import ( + ViolationLog, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=20, stage="extra-a", runner_config="1-gpu-small") + + +def _make_endpoint(*, device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL, swa_lut=None): + canary_buf = torch.zeros(4, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device) + slot_view = torch.zeros(1, dtype=torch.int64, device=device) + kernel_view = torch.zeros(1, dtype=torch.int64, device=device) + enable_chain_position_assert = torch.ones(1, dtype=torch.int32, device=device) + return CanaryEndpoint( + kernel_kind=kernel_kind, + canary_buf=canary_buf, + full_to_swa_index_mapping=swa_lut, + slot_run_counter_view=slot_view, + kernel_run_counter_view=kernel_view, + enable_chain_position_assert=enable_chain_position_assert, + ) + + +def _make_kernel_args(device): + verify_plan = VerifyPlan.allocate(verify_capacity=1, device=device) + write_plan = WritePlan.allocate(write_req_capacity=1, device=device) + log = ViolationLog.allocate(ring_capacity=2, device=device) + return SimpleNamespace( + verify_plan=verify_plan, + write_plan=write_plan, + violation_log=log, + input_ids=torch.zeros(1, dtype=torch.int64, device=device), + positions=torch.zeros(1, dtype=torch.int64, device=device), + out_cache_loc=torch.zeros(1, dtype=torch.int64, device=device), + enable_write_input_assert=False, + enable_verify_token_assert=False, + expected_inputs=ExpectedInputs.allocate(capacity=1, device=device), + ) + + +class TestSelfUnitEndpoint(CustomTestCase): + def setUp(self): + self.device = DEFAULT_DEVICE + + def test_launch_per_forward_passes_kernel_kind(self): + """Verify per-forward launch passes the endpoint kernel kind.""" + captured: list[tuple[str, CanaryLaunchTag]] = [] + with patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: captured.append(("verify", kwargs["context"].kernel_kind)), + ), patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: captured.append(("write", kwargs["context"].kernel_kind)), + ): + ep = _make_endpoint( + device=self.device, kernel_kind=CanaryLaunchTag.TAIL_V_SWA + ) + args = _make_kernel_args(self.device) + ep.launch_per_forward( + verify_plan=args.verify_plan, + write_plan=args.write_plan, + input_ids=args.input_ids, + positions=args.positions, + out_cache_loc=args.out_cache_loc, + enable_write_input_assert=args.enable_write_input_assert, + enable_verify_token_assert=args.enable_verify_token_assert, + expected_inputs=args.expected_inputs, + violation_log=args.violation_log, + ) + self.assertIn(("verify", CanaryLaunchTag.TAIL_V_SWA), captured) + self.assertIn(("write", CanaryLaunchTag.TAIL_V_SWA), captured) + + def test_endpoint_shares_violation_log_across_launches(self): + """Verify endpoints can reuse the same violation log.""" + captured_rings: list[int] = [] + with patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: captured_rings.append( + kwargs["context"].violation_ring.data_ptr() + ), + ), patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: None, + ): + shared_log = ViolationLog.allocate(ring_capacity=2, device=self.device) + ep_a = _make_endpoint( + device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL + ) + ep_b = _make_endpoint( + device=self.device, kernel_kind=CanaryLaunchTag.HEAD_V_FULL + ) + + args = _make_kernel_args(self.device) + ep_a.launch_per_forward( + verify_plan=args.verify_plan, + write_plan=args.write_plan, + input_ids=args.input_ids, + positions=args.positions, + out_cache_loc=args.out_cache_loc, + enable_write_input_assert=args.enable_write_input_assert, + enable_verify_token_assert=args.enable_verify_token_assert, + expected_inputs=args.expected_inputs, + violation_log=shared_log, + ) + ep_b.launch_per_forward( + verify_plan=args.verify_plan, + write_plan=args.write_plan, + input_ids=args.input_ids, + positions=args.positions, + out_cache_loc=args.out_cache_loc, + enable_write_input_assert=args.enable_write_input_assert, + enable_verify_token_assert=args.enable_verify_token_assert, + expected_inputs=args.expected_inputs, + violation_log=shared_log, + ) + self.assertEqual(captured_rings[0], captured_rings[1]) + self.assertEqual(captured_rings[0], shared_log.violation_ring.data_ptr()) + + def test_swa_endpoint_pre_translates_out_cache_loc(self): + """Verify SWA endpoints translate cache locations before write launch.""" + captured: list[torch.Tensor] = [] + with patch.object( + endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None + ), patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: captured.append(kwargs["out_cache_loc"]), + ): + # LUT maps full slot i → swa slot (i + 100) so we can verify the gather happened. + lut = torch.arange(8, dtype=torch.int64, device=self.device) + 100 + swa_ep = _make_endpoint( + device=self.device, + kernel_kind=CanaryLaunchTag.HEAD_K_SWA, + swa_lut=lut, + ) + full_ep = _make_endpoint( + device=self.device, + kernel_kind=CanaryLaunchTag.HEAD_K_FULL, + swa_lut=None, + ) + args = _make_kernel_args(self.device) + + swa_ep.launch_per_forward( + verify_plan=args.verify_plan, + write_plan=args.write_plan, + input_ids=args.input_ids, + positions=args.positions, + out_cache_loc=args.out_cache_loc, + enable_write_input_assert=args.enable_write_input_assert, + enable_verify_token_assert=args.enable_verify_token_assert, + expected_inputs=args.expected_inputs, + violation_log=args.violation_log, + ) + full_ep.launch_per_forward( + verify_plan=args.verify_plan, + write_plan=args.write_plan, + input_ids=args.input_ids, + positions=args.positions, + out_cache_loc=args.out_cache_loc, + enable_write_input_assert=args.enable_write_input_assert, + enable_verify_token_assert=args.enable_verify_token_assert, + expected_inputs=args.expected_inputs, + violation_log=args.violation_log, + ) + # SWA call: out_cache_loc was rewritten via lut gather (so identity-shifted by +100 here). + expected_swa = lut[args.out_cache_loc] + self.assertTrue(torch.equal(captured[0], expected_swa)) + # FULL call: out_cache_loc keeps the same values and dtype. + self.assertIs(captured[1], args.out_cache_loc) + + def test_swa_endpoint_trailing_sentinel_row_yields_skip(self): + """Verify SWA sentinel cache rows become write-skip markers.""" + captured: list[torch.Tensor] = [] + with patch.object( + endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None + ), patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: captured.append(kwargs["out_cache_loc"]), + ): + # 8 in-window rows + 1 trailing sentinel row at index 8. + lut = torch.arange(8, dtype=torch.int64, device=self.device) + lut = torch.cat( + [lut, torch.tensor([-1], dtype=torch.int64, device=self.device)] + ) + swa_ep = _make_endpoint( + device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_SWA, swa_lut=lut + ) + args = _make_kernel_args(self.device) + # Point out_cache_loc at the trailing-sentinel-row index — this is how sglang signals + # "this token is out-of-window for the SWA group" pre-cleanup, and the new host gather must + # produce -1 here. + args.out_cache_loc.fill_(8) + + swa_ep.launch_per_forward( + verify_plan=args.verify_plan, + write_plan=args.write_plan, + input_ids=args.input_ids, + positions=args.positions, + out_cache_loc=args.out_cache_loc, + enable_write_input_assert=args.enable_write_input_assert, + enable_verify_token_assert=args.enable_verify_token_assert, + expected_inputs=args.expected_inputs, + violation_log=args.violation_log, + ) + + self.assertTrue( + torch.equal( + captured[0], + torch.tensor([-1], dtype=torch.int64, device=self.device), + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_future_tensor.py b/test/registered/kv_canary/test_self_unit_future_tensor.py new file mode 100644 index 000000000..e593cd6d0 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_future_tensor.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import unittest +from typing import cast + +import torch + +from sglang.srt.kv_canary.runner.future_tensor import FutureTensors +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=20, stage="extra-a", runner_config="1-gpu-small") + + +class _FakeEvent: + def __init__(self) -> None: + self.synchronize_count = 0 + + def synchronize(self) -> None: + self.synchronize_count += 1 + + +class TestFutureTensors(CustomTestCase): + def test_cuda_stage_then_wait_returns_host_copy(self) -> None: + """Verify staged CUDA tensors are copied back on wait.""" + device = torch.device("cuda") + alt_stream = torch.cuda.Stream(device=device) + default_stream = torch.cuda.current_stream(device) + self.assertNotEqual(alt_stream.stream_id, default_stream.stream_id) + + src_first = torch.tensor([41], dtype=torch.int32, device=device) + future_first = FutureTensors.device_to_host( + xs_device=src_first, d2h_stream=alt_stream + ) + result_first = future_first.wait() + self.assertEqual(int(result_first.item()), 41) + + src_second = torch.tensor([97], dtype=torch.int32, device=device) + future_second = FutureTensors.device_to_host( + xs_device=src_second, d2h_stream=alt_stream + ) + result_second = future_second.wait() + self.assertEqual(int(result_second.item()), 97) + + def test_cuda_pinned_when_stream_is_provided(self) -> None: + """Verify CUDA staging uses pinned host memory with a stream.""" + device = torch.device("cuda") + alt_stream = torch.cuda.Stream(device=device) + src = torch.tensor([5], dtype=torch.int32, device=device) + future = FutureTensors.device_to_host(xs_device=src, d2h_stream=alt_stream) + staged_tensors = [ + v for v in future._data.values() if isinstance(v, torch.Tensor) + ] + self.assertTrue(staged_tensors) + self.assertTrue(all(t.is_pinned() for t in staged_tensors)) + self.assertEqual(int(future.wait().item()), 5) + + def test_cuda_each_call_allocates_fresh_host(self) -> None: + """Verify each CUDA staging call owns a fresh host buffer.""" + device = torch.device("cuda") + alt_stream = torch.cuda.Stream(device=device) + src_a = torch.tensor([13], dtype=torch.int32, device=device) + src_b = torch.tensor([29], dtype=torch.int32, device=device) + future_a = FutureTensors.device_to_host(xs_device=src_a, d2h_stream=alt_stream) + future_b = FutureTensors.device_to_host(xs_device=src_b, d2h_stream=alt_stream) + ptrs_a = { + v.data_ptr() for v in future_a._data.values() if isinstance(v, torch.Tensor) + } + ptrs_b = { + v.data_ptr() for v in future_b._data.values() if isinstance(v, torch.Tensor) + } + self.assertTrue(ptrs_a and ptrs_b) + self.assertFalse(ptrs_a & ptrs_b) + self.assertEqual(int(future_a.wait().item()), 13) + self.assertEqual(int(future_b.wait().item()), 29) + + def test_dict_of_all_tensors_roundtrip(self) -> None: + """Verify a dict of multiple tensors round-trips entry-by-entry.""" + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + src = { + "x": torch.tensor([11, 22], dtype=torch.int64, device=device), + "y": torch.tensor([99], dtype=torch.int32, device=device), + } + future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream) + out = future.wait() + self.assertIsInstance(out, dict) + self.assertEqual(out["x"].tolist(), [11, 22]) + self.assertEqual(int(out["y"].item()), 99) + self.assertTrue(out["x"].is_pinned()) + self.assertTrue(out["y"].is_pinned()) + + def test_dict_mixes_tensor_and_passthrough(self) -> None: + """Verify non-tensor dict entries ride through verbatim alongside staging.""" + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + sentinel_obj = {"nested": [1, 2, 3]} + src = { + "step": 42, + "label": "decode", + "extra": sentinel_obj, + "counter": torch.tensor([7], dtype=torch.int32, device=device), + } + future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream) + out = future.wait() + self.assertEqual(out["step"], 42) + self.assertEqual(out["label"], "decode") + # Identity (not deep-copy) — callers can rely on shared mutable references. + self.assertIs(out["extra"], sentinel_obj) + self.assertEqual(int(out["counter"].item()), 7) + self.assertTrue(out["counter"].is_pinned()) + + def test_dict_passthrough_preserves_tensor_value(self) -> None: + """Verify tensors share device memory but non-tensor types are not staged.""" + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + src_tensor = torch.tensor([3], dtype=torch.int32, device=device) + src = {"step": 100, "buf": src_tensor} + future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream) + out = future.wait() + # Tensor is staged to a fresh pinned-host buffer (different storage from src). + self.assertNotEqual(out["buf"].data_ptr(), src_tensor.data_ptr()) + self.assertTrue(out["buf"].is_pinned()) + # Non-tensor passes through with no copy. + self.assertEqual(out["step"], 100) + self.assertIsInstance(out["step"], int) + + def test_dict_without_tensor_raises(self) -> None: + """Verify a tensor-less dict raises (no device to anchor the d2h sync).""" + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + with self.assertRaises(ValueError): + FutureTensors.device_to_host( + xs_device={"step": 0, "label": "decode"}, d2h_stream=stream + ) + + def test_wait_called_twice_raises(self) -> None: + """Verify wait() after the first drain raises (state cleared).""" + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + src = torch.tensor([3], dtype=torch.int32, device=device) + future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream) + self.assertEqual(int(future.wait().item()), 3) + with self.assertRaises(RuntimeError): + future.wait() + + def test_wait_clears_fields_and_rejects_second_wait(self) -> None: + """Verify wait() syncs the event exactly once and clears internal state.""" + tensor = torch.tensor([1, 2, 3]) + event = _FakeEvent() + future = FutureTensors( + _data={"x": tensor}, _event=cast(torch.cuda.Event, event) + ) + + result = future.wait() + self.assertIs(result["x"], tensor) + self.assertEqual(event.synchronize_count, 1) + self.assertIsNone(future._data) + self.assertIsNone(future._event) + + with self.assertRaisesRegex(RuntimeError, "called more than once"): + future.wait() + # Failed wait must not re-trigger event.synchronize. + self.assertEqual(event.synchronize_count, 1) + + def test_dict_anchor_picked_from_first_tensor(self) -> None: + """Verify staging works when the first key is a non-tensor (anchor must scan).""" + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + src = { + "step": 5, + "buf": torch.tensor([17], dtype=torch.int32, device=device), + } + out = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream).wait() + self.assertEqual(out["step"], 5) + self.assertEqual(int(out["buf"].item()), 17) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_plan_input.py b/test/registered/kv_canary/test_self_unit_plan_input.py new file mode 100644 index 000000000..9eb9a59f8 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_plan_input.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.kv_canary.plan_input import PlanInput +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import ( + DEFAULT_DEVICE, + make_forward_batch, +) +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small") + + +def _make_static_plan_input(*, bs_capacity: int, device) -> PlanInput: + return PlanInput( + req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device), + prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device), + extend_seq_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device), + req_to_verify_expected_tokens_valid_lens=torch.zeros( + bs_capacity, dtype=torch.int64, device=device + ), + ) + + +class TestSelfUnitPlanInput(CustomTestCase): + def setUp(self): + self.device = DEFAULT_DEVICE + + def test_plan_input_fill_from_forward_batch_extend(self): + """Verify extend batches populate per-forward plan inputs.""" + fb = make_forward_batch( + self.device, + req_pool_indices=torch.tensor( + [1, 2], dtype=torch.int64, device=self.device + ), + seq_lens=torch.tensor([10, 12], dtype=torch.int32, device=self.device), + extend_prefix_lens=torch.tensor( + [3, 5], dtype=torch.int32, device=self.device + ), + extend_seq_lens=torch.tensor([7, 7], dtype=torch.int32, device=self.device), + is_extend=True, + ) + plan = _make_static_plan_input(bs_capacity=4, device=self.device) + plan.fill_from_forward_batch(forward_batch=fb) + self.assertEqual(plan.req_pool_indices[:2].tolist(), [1, 2]) + self.assertEqual(plan.req_pool_indices[2:].tolist(), [0, 0]) + self.assertEqual(plan.prefix_lens[:2].tolist(), [3, 5]) + self.assertEqual(plan.extend_seq_lens[:2].tolist(), [7, 7]) + self.assertEqual(plan.prefix_lens.dtype, torch.int64) + self.assertEqual(plan.extend_seq_lens.dtype, torch.int64) + + def test_plan_input_fill_from_forward_batch_target_verify(self): + """Verify target-verify batches derive draft verification spans.""" + # TARGET_VERIFY writes spec_info.draft_token_num positions per req starting at the + # current (un-bumped) seq_lens; extend_prefix_lens / extend_seq_lens are deliberately + # NOT supplied because init_new does not populate them for this mode. + spec_info = SimpleNamespace(draft_token_num=4) + fb = make_forward_batch( + self.device, + req_pool_indices=torch.tensor( + [6, 7], dtype=torch.int64, device=self.device + ), + seq_lens=torch.tensor([10, 14], dtype=torch.int32, device=self.device), + is_target_verify=True, + spec_info=spec_info, + ) + plan = _make_static_plan_input(bs_capacity=4, device=self.device) + plan.fill_from_forward_batch(forward_batch=fb) + self.assertEqual(plan.prefix_lens[:2].tolist(), [10, 14]) + self.assertEqual(plan.extend_seq_lens[:2].tolist(), [4, 4]) + + def test_plan_input_fill_from_forward_batch_draft_extend_v2(self): + """Verify draft-extend-v2 batches derive prefix lengths from sequence lengths.""" + # DRAFT_EXTEND_V2 has seq_lens already bumped by the per-req draft refill length; + # extend_prefix_lens is intentionally absent (cuda-graph replay does not set it). The + # builder must derive prefix as seq_lens - extend_seq_lens. + fb = make_forward_batch( + self.device, + req_pool_indices=torch.tensor( + [4, 5], dtype=torch.int64, device=self.device + ), + seq_lens=torch.tensor([14, 18], dtype=torch.int32, device=self.device), + extend_seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=self.device), + is_draft_extend_v2=True, + ) + plan = _make_static_plan_input(bs_capacity=4, device=self.device) + plan.fill_from_forward_batch(forward_batch=fb) + self.assertEqual(plan.prefix_lens[:2].tolist(), [10, 14]) + self.assertEqual(plan.extend_seq_lens[:2].tolist(), [4, 4]) + + def test_plan_input_fill_from_forward_batch_decode(self): + """Verify decode batches populate one-token verification spans.""" + fb = make_forward_batch( + self.device, + req_pool_indices=torch.tensor( + [1, 2, 3], dtype=torch.int64, device=self.device + ), + seq_lens=torch.tensor([4, 7, 1], dtype=torch.int32, device=self.device), + is_extend=False, + ) + plan = _make_static_plan_input(bs_capacity=4, device=self.device) + plan.fill_from_forward_batch(forward_batch=fb) + self.assertEqual(plan.prefix_lens[:3].tolist(), [3, 6, 0]) + self.assertEqual(plan.extend_seq_lens[:3].tolist(), [1, 1, 1]) + + def test_plan_input_padding_dummy_sentinel(self): + """Verify padding sentinel rows remain valid plan input entries.""" + fb = make_forward_batch( + self.device, + req_pool_indices=torch.tensor( + [0, 5, 0], dtype=torch.int64, device=self.device + ), + seq_lens=torch.tensor([0, 3, 0], dtype=torch.int32, device=self.device), + is_extend=False, + ) + plan = _make_static_plan_input(bs_capacity=4, device=self.device) + plan.fill_from_forward_batch(forward_batch=fb) + self.assertEqual(plan.req_pool_indices[:3].tolist(), [0, 5, 0]) + self.assertEqual(plan.req_pool_indices.dtype, torch.int64) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_pool_patcher.py b/test/registered/kv_canary/test_self_unit_pool_patcher.py new file mode 100644 index 000000000..5f6a35408 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_pool_patcher.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import unittest + +from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES +from sglang.srt.kv_canary.buffer_group import PoolKind +from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import ( + DEFAULT_DEVICE, + make_base_config, + make_mha_pool, +) +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small") + + +class PoolPatcherHelper: + def setUp(self): + self.device = DEFAULT_DEVICE + self.config = make_base_config() + + +class TestAttachCanaryBuffers(PoolPatcherHelper, CustomTestCase): + def test_canary_buffer_group_allocate_full_only(self): + """Verify MHA pools allocate only full canary buffers.""" + pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2) + groups_tuple = attach_canary_buffers( + pool=pool, + config=self.config, + device=self.device, + kv_token_id_vs_position_offset=0, + ) + groups = {g.kind: g for g in groups_tuple} + self.assertEqual(set(groups.keys()), {PoolKind.FULL}) + group = groups[PoolKind.FULL] + self.assertEqual(group.k_head.shape, (16, CANARY_SLOT_BYTES)) + self.assertEqual(group.k_tail.shape, (16, CANARY_SLOT_BYTES)) + self.assertIsNotNone(group.v_head) + self.assertIsNotNone(group.v_tail) + self.assertEqual(group.v_head.shape, (16, CANARY_SLOT_BYTES)) + + +class TestPoolPatcherBufferInfos(PoolPatcherHelper, CustomTestCase): + def test_get_contiguous_buf_infos_inserts_canary_entries(self): + """Verify contiguous buffer metadata includes canary entries after patching.""" + for patched in (False, True): + with self.subTest(patched=patched): + pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2) + ptrs_before, _, _ = pool.get_contiguous_buf_infos() + n_before = len(ptrs_before) + + if patched: + attach_canary_buffers( + pool=pool, + config=self.config, + device=self.device, + kv_token_id_vs_position_offset=0, + ) + ptrs_after, _, _ = pool.get_contiguous_buf_infos() + self.assertEqual(len(ptrs_after), n_before + 4) + else: + ptrs_after, _, _ = pool.get_contiguous_buf_infos() + self.assertEqual(ptrs_after, ptrs_before) + + def test_pd_layout_canary_inserted_correctly(self): + """Verify PD (prefill-decode disaggregation) canary buffers are inserted in layout order.""" + pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2) + k_ptrs_orig = [b.data_ptr() for b in pool.k_buffer] + v_ptrs_orig = [b.data_ptr() for b in pool.v_buffer] + + groups_tuple = attach_canary_buffers( + pool=pool, + config=self.config, + device=self.device, + kv_token_id_vs_position_offset=0, + ) + group = {g.kind: g for g in groups_tuple}[PoolKind.FULL] + ptrs_after, _, _ = pool.get_contiguous_buf_infos() + + canary_k_ptrs = [group.k_head.data_ptr(), group.k_tail.data_ptr()] + canary_v_ptrs = [group.v_head.data_ptr(), group.v_tail.data_ptr()] + self.assertEqual(ptrs_after[0], canary_k_ptrs[0]) + self.assertEqual(ptrs_after[1 : 1 + len(k_ptrs_orig)], k_ptrs_orig) + k_tail_idx = 1 + len(k_ptrs_orig) + self.assertEqual(ptrs_after[k_tail_idx], canary_k_ptrs[1]) + self.assertEqual(ptrs_after[k_tail_idx + 1], canary_v_ptrs[0]) + v_start = k_tail_idx + 2 + self.assertEqual(ptrs_after[v_start : v_start + len(v_ptrs_orig)], v_ptrs_orig) + self.assertEqual(ptrs_after[-1], canary_v_ptrs[1]) + + +class TestCanaryBufferBudget(PoolPatcherHelper, CustomTestCase): + def test_canary_buf_per_token_bytes_within_budget(self): + """Verify canary per-token storage stays below the real KV budget.""" + pool = make_mha_pool(self.device, num_slots=16, dim=64, layer_num=2) + groups_tuple = attach_canary_buffers( + pool=pool, + config=self.config, + device=self.device, + kv_token_id_vs_position_offset=0, + ) + group = {g.kind: g for g in groups_tuple}[PoolKind.FULL] + + slot_stride_bytes = group.k_head.stride(0) * group.k_head.element_size() + self.assertLessEqual(slot_stride_bytes, CANARY_SLOT_BYTES) + + real_kv_per_token_bytes = ( + pool.k_buffer[0].stride(0) * pool.k_buffer[0].element_size() + ) + self.assertLess(slot_stride_bytes, real_kv_per_token_bytes) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_pool_patcher_utils.py b/test/registered/kv_canary/test_self_unit_pool_patcher_utils.py new file mode 100644 index 000000000..cfb1019dc --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_pool_patcher_utils.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import unittest + +from sglang.srt.kv_canary.pool_patcher.utils import wrap_method +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small") + + +class _FakeObj: + def greet(self, name: str) -> str: + return f"hello {name}" + + +class TestPoolPatcherUtils(CustomTestCase): + def test_wrap_method_delegates_to_wrapper(self) -> None: + """Verify wrapped methods delegate through the wrapper.""" + obj = _FakeObj() + + def _with_validation(original, name: str) -> str: + return original(name) + "!" + + wrap_method(obj, "greet", wrapper=_with_validation) + self.assertEqual(obj.greet("world"), "hello world!") + + def test_wrap_method_missing_method_raises_attribute_error(self) -> None: + """Verify wrapping a missing method raises AttributeError.""" + obj = _FakeObj() + with self.assertRaisesRegex(AttributeError, "missing required method"): + wrap_method( + obj, "nonexistent", wrapper=lambda orig, *a, **kw: orig(*a, **kw) + ) + + def test_wrap_method_double_wrap_raises_runtime_error(self) -> None: + """Verify wrapping the same method twice raises RuntimeError.""" + obj = _FakeObj() + wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw)) + with self.assertRaisesRegex(RuntimeError, "already wrapped by kv-canary"): + wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw)) + + def test_wrap_method_preserves_functools_wraps_metadata(self) -> None: + """Verify wrapping preserves method metadata.""" + obj = _FakeObj() + original_name = obj.greet.__name__ + wrap_method(obj, "greet", wrapper=lambda orig, *a, **kw: orig(*a, **kw)) + self.assertEqual(obj.greet.__name__, original_name) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_violation.py b/test/registered/kv_canary/test_self_unit_violation.py new file mode 100644 index 000000000..ac272cd19 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_violation.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.jit_kernel.kv_canary import consts +from sglang.jit_kernel.kv_canary.consts import FailReason +from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag +from sglang.srt.kv_canary.config import CanaryMode +from sglang.srt.kv_canary.runner import violation_reporter as violation_reporter_module +from sglang.srt.kv_canary.runner.violation_reporter import ( + ViolationReporter, + _format_violation, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=5, stage="extra-a", runner_config="1-gpu-small") + + +def _make_row( + *, + kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL, + slot_idx: int = 17, + position: int = 42, + stored_token: int = 111, + expected_token: int = 0, + stored_chain_hash: int = 0, + expected_aux: int = 0, + fail_reason_bits: int = 0, +) -> list[int]: + row = [0] * consts.VIOLATION_FIELDS + row[consts.VIOLATION_FIELD_KERNEL_KIND] = int(kernel_kind) + row[consts.VIOLATION_FIELD_SLOT_IDX] = slot_idx + row[consts.VIOLATION_FIELD_POSITION] = position + row[consts.VIOLATION_FIELD_STORED_TOKEN] = stored_token + row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = expected_token + row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = stored_chain_hash + row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_aux + row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = fail_reason_bits + return row + + +class TestViolationReporter(CustomTestCase): + def test_format_violation_verify_path_labels_each_bit(self) -> None: + """Verify verify-path violations render each fail-reason bit.""" + row = _make_row( + stored_chain_hash=0x1111111111111111, + expected_aux=0x2222222222222222, + fail_reason_bits=int( + FailReason.VERIFY_CHAIN_HASH_MISMATCH + | FailReason.VERIFY_POSITION_MISMATCH + ), + ) + out = _format_violation( + row=row, total=1, ring_overflow=False, step_when_pumped=7 + ) + self.assertEqual( + out, + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash+verify_position " + "slot_idx=17 position=42 stored_token=111 expected_token=0 stored_chain_hash=0x1111111111111111 " + "expected_aux=0x2222222222222222\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: verify_chain_hash verify_position\n" + " stored: token_id=111 position=42 prev_hash=0x1111111111111111\n" + " expected: prev_hash=0x2222222222222222\n" + " total_violations=1 ring_overflow=False step_when_pumped=7", + ) + + def test_format_violation_write_token_mismatch_labels_and_position(self) -> None: + """Verify write-token violations render token and position details.""" + row = _make_row( + position=42, + stored_token=999, + expected_token=888, + stored_chain_hash=0xDEADBEEFCAFEBABE, + expected_aux=43, + fail_reason_bits=int(FailReason.WRITE_TOKEN_MISMATCH), + ) + out = _format_violation( + row=row, total=1, ring_overflow=False, step_when_pumped=0 + ) + self.assertEqual( + out, + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_token slot_idx=17 position=42 " + "stored_token=999 expected_token=888 stored_chain_hash=0xdeadbeefcafebabe " + "expected_aux=0x000000000000002b\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: write_token\n" + " actual: token_id=999 position=42 prev_hash=0xdeadbeefcafebabe\n" + " expected: token_id=888 position=43\n" + " total_violations=1 ring_overflow=False step_when_pumped=0", + ) + + def test_format_violation_write_position_mismatch_uses_expected_aux_as_position( + self, + ) -> None: + """Verify write-position violations render expected_aux as a position.""" + row = _make_row( + position=42, + stored_token=111, + expected_token=111, + expected_aux=99, + fail_reason_bits=int(FailReason.WRITE_POSITION_MISMATCH), + ) + out = _format_violation( + row=row, total=1, ring_overflow=False, step_when_pumped=0 + ) + self.assertEqual( + out, + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_position slot_idx=17 position=42 " + "stored_token=111 expected_token=111 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000063\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: write_position\n" + " actual: token_id=111 position=42 prev_hash=0x0000000000000000\n" + " expected: token_id=111 position=99\n" + " total_violations=1 ring_overflow=False step_when_pumped=0", + ) + + def test_format_violation_combined_write_bits_render_both_labels(self) -> None: + """Verify combined write violation bits render both labels.""" + row = _make_row( + fail_reason_bits=int( + FailReason.WRITE_TOKEN_MISMATCH | FailReason.WRITE_POSITION_MISMATCH + ), + ) + out = _format_violation( + row=row, total=1, ring_overflow=False, step_when_pumped=0 + ) + self.assertEqual( + out, + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=write_token+write_position " + "slot_idx=17 position=42 stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=17, position=42)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: write_token write_position\n" + " actual: token_id=111 position=42 prev_hash=0x0000000000000000\n" + " expected: token_id=0 position=0\n" + " total_violations=1 ring_overflow=False step_when_pumped=0", + ) + + def test_format_violation_unknown_kernel_kind_renders_unknown_label(self) -> None: + """Verify unknown kernel kinds render an unknown label.""" + row = _make_row(fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH)) + row[consts.VIOLATION_FIELD_KERNEL_KIND] = 9999 + out = _format_violation( + row=row, total=1, ring_overflow=False, step_when_pumped=0 + ) + self.assertEqual( + out, + "kv_canary violation: launch_tag=unknown(9999) fail_reason=verify_chain_hash slot_idx=17 position=42 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=unknown(9999), slot_idx=17, position=42)\n" + "canary_kind: unknown(9999)\n" + " fail_reasons: verify_chain_hash\n" + " stored: token_id=111 position=42 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=1 ring_overflow=False step_when_pumped=0", + ) + + +def _make_reporter( + *, + rows: list[list[int]], + write_index: int, + ring_capacity: int, + mode: CanaryMode = CanaryMode.LOG, +) -> ViolationReporter: + ring = torch.zeros(ring_capacity, consts.VIOLATION_FIELDS, dtype=torch.int64) + for i, row in enumerate(rows): + ring[i] = torch.tensor(row, dtype=torch.int64) + violation_log = SimpleNamespace( + violation_ring=ring, + violation_write_index=torch.tensor([write_index], dtype=torch.int32), + ) + device_state = SimpleNamespace(violation_log=violation_log) + config = SimpleNamespace(mode=mode) + return ViolationReporter(config=config, device_state=device_state) + + +class TestLogOrRaiseViolation(CustomTestCase): + def test_log_or_raise_violation_empty_ring_is_noop(self) -> None: + """Empty ring (write_index=0) emits no warning and leaves reporter non-raised.""" + reporter = _make_reporter( + rows=[], write_index=0, ring_capacity=4, mode=CanaryMode.LOG + ) + with patch.object(violation_reporter_module.logger, "warning") as mock_warning: + reporter.log_or_raise_violation(outer_step_counter=0) + mock_warning.assert_not_called() + self.assertFalse(reporter.is_raised) + + def test_log_mode_emits_one_warning_per_violation(self) -> None: + """Log mode with 2 valid rows emits 2 warnings, each a full _format_violation snapshot for that row.""" + rows = [ + _make_row( + slot_idx=11, + position=101, + fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH), + ), + _make_row( + slot_idx=22, + position=202, + fail_reason_bits=int(FailReason.VERIFY_POSITION_MISMATCH), + ), + ] + reporter = _make_reporter( + rows=rows, write_index=2, ring_capacity=4, mode=CanaryMode.LOG + ) + with patch.object(violation_reporter_module.logger, "warning") as mock_warning: + reporter.log_or_raise_violation(outer_step_counter=7) + + self.assertEqual(mock_warning.call_count, 2) + messages: list[str] = [call.args[0] for call in mock_warning.call_args_list] + self.assertEqual( + messages[0], + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash slot_idx=11 position=101 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: verify_chain_hash\n" + " stored: token_id=111 position=101 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=2 ring_overflow=False step_when_pumped=7", + ) + self.assertEqual( + messages[1], + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_position slot_idx=22 position=202 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: verify_position\n" + " stored: token_id=111 position=202 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=2 ring_overflow=False step_when_pumped=7", + ) + self.assertFalse(reporter.is_raised) + + def test_raise_mode_raises_one_error_containing_all_violations(self) -> None: + """Raise mode raises a single RuntimeError whose text is the 2 formatted rows joined with single newlines.""" + rows = [ + _make_row( + slot_idx=11, + position=101, + fail_reason_bits=int(FailReason.VERIFY_CHAIN_HASH_MISMATCH), + ), + _make_row( + slot_idx=22, + position=202, + fail_reason_bits=int(FailReason.VERIFY_POSITION_MISMATCH), + ), + ] + reporter = _make_reporter( + rows=rows, write_index=2, ring_capacity=4, mode=CanaryMode.RAISE + ) + with self.assertRaises(RuntimeError) as ctx: + reporter.log_or_raise_violation(outer_step_counter=5) + + self.assertEqual( + str(ctx.exception), + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_chain_hash slot_idx=11 position=101 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: verify_chain_hash\n" + " stored: token_id=111 position=101 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=2 ring_overflow=False step_when_pumped=5\n" + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=verify_position slot_idx=22 position=202 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: verify_position\n" + " stored: token_id=111 position=202 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=2 ring_overflow=False step_when_pumped=5", + ) + self.assertTrue(reporter.is_raised) + + def test_log_mode_ring_overflow_marks_overflow_in_each_row(self) -> None: + """Log mode with write_index=5 but ring_capacity=2 emits 2 warnings, each a full snapshot with overflow footer.""" + rows = [ + _make_row(slot_idx=11, position=101), + _make_row(slot_idx=22, position=202), + ] + reporter = _make_reporter( + rows=rows, write_index=5, ring_capacity=2, mode=CanaryMode.LOG + ) + with patch.object(violation_reporter_module.logger, "warning") as mock_warning: + reporter.log_or_raise_violation(outer_step_counter=0) + + self.assertEqual(mock_warning.call_count, 2) + messages: list[str] = [call.args[0] for call in mock_warning.call_args_list] + self.assertEqual( + messages[0], + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=none slot_idx=11 position=101 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=11, position=101)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: none\n" + " stored: token_id=111 position=101 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=5 ring_overflow=True step_when_pumped=0", + ) + self.assertEqual( + messages[1], + "kv_canary violation: launch_tag=HEAD_K_FULL fail_reason=none slot_idx=22 position=202 " + "stored_token=111 expected_token=0 stored_chain_hash=0x0000000000000000 " + "expected_aux=0x0000000000000000\n" + "KV cache canary violation detected (kernel_kind=HEAD_K_FULL, slot_idx=22, position=202)\n" + "canary_kind: per_forward_head_k_full\n" + " fail_reasons: none\n" + " stored: token_id=111 position=202 prev_hash=0x0000000000000000\n" + " expected: prev_hash=0x0000000000000000\n" + " total_violations=5 ring_overflow=True step_when_pumped=0", + ) + + +if __name__ == "__main__": + unittest.main()