Add the KV-canary core: data layer, MHA KV-pool patcher, and per-forward runner (#26808)
This commit is contained in:
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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(),
|
||||
)
|
||||
@@ -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],
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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:]}"
|
||||
)
|
||||
Reference in New Issue
Block a user