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

This commit is contained in:
fzyzcjy
2026-05-31 09:54:24 +08:00
committed by GitHub
parent 736ad1f32a
commit 11391b2a1c
46 changed files with 3756 additions and 0 deletions
+6
View File
@@ -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
@@ -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
+114
View File
@@ -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,
)
+49
View File
@@ -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(),
)
+156
View File
@@ -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)
@@ -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],
)
+135
View File
@@ -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}"
)
@@ -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,)
@@ -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
@@ -0,0 +1,77 @@
from __future__ import annotations
from typing import Any, Callable, List, Tuple
import torch
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
BufInfoTriple = Tuple[List[int], List[int], List[int]]
def patch_buf_info_method(
pool: object,
*,
method_name: str,
group: CanaryBufferGroup,
has_v_half: bool,
page_size: int,
) -> None:
"""Wrap ``pool.<method_name>()`` so its (ptrs, lens, item_lens) triple is spliced with K/V
head and tail entries from ``group``."""
def _with_splice(original: Callable, *args: Any, **kwargs: Any) -> BufInfoTriple:
ptrs, lens, item_lens = original(*args, **kwargs)
return splice_kv_buf_info(
ptrs=ptrs,
lens=lens,
item_lens=item_lens,
group=group,
has_v_half=has_v_half,
page_size=page_size,
)
wrap_method(pool, method_name, wrapper=_with_splice)
def splice_kv_buf_info(
*,
ptrs: List[int],
lens: List[int],
item_lens: List[int],
group: CanaryBufferGroup,
has_v_half: bool,
page_size: int,
) -> BufInfoTriple:
entries = list(zip(ptrs, lens, item_lens))
k_head = _entry_triple(group.k_head, page_size=page_size)
k_tail = _entry_triple(group.k_tail, page_size=page_size)
if not has_v_half:
out = [k_head, *entries, k_tail]
else:
assert group.v_head is not None and group.v_tail is not None
v_head = _entry_triple(group.v_head, page_size=page_size)
v_tail = _entry_triple(group.v_tail, page_size=page_size)
if len(entries) % 2 != 0:
raise RuntimeError(
f"kv-canary: K/V split adapter expects even-length buf_info list, got {len(entries)}"
)
mid = len(entries) // 2
out = [k_head, *entries[:mid], k_tail, v_head, *entries[mid:], v_tail]
return _untranspose_entries(out)
def _entry_triple(buf: torch.Tensor, *, page_size: int) -> Tuple[int, int, int]:
return (
buf.data_ptr(),
buf.nbytes,
buf[0].nbytes * page_size,
)
def _untranspose_entries(entries: List[Tuple[int, int, int]]) -> BufInfoTriple:
out_ptrs, out_lens, out_item_lens = (list(col) for col in zip(*entries))
return out_ptrs, out_lens, out_item_lens
@@ -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)
@@ -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)
@@ -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
@@ -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,
)
@@ -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
)
@@ -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()
@@ -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]
@@ -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])
@@ -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)
@@ -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,
)
+126
View File
@@ -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,
)
@@ -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
+13
View File
@@ -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,
+141
View File
@@ -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)
]
+191
View File
@@ -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)
@@ -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",
),
}
@@ -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)
+55
View File
@@ -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))
@@ -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))
@@ -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:]}"
)