minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add the KV-canary perturb modes and PD-disaggregation e2e tests (#26819)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-31 09:59:09 +08:00
committed by GitHub
parent 6be4b32d8d
commit ae9db7ff4b
20 changed files with 1551 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+4
View File
@@ -754,7 +754,11 @@ class Envs:
# ===================================================================
SGLANG_KV_CANARY_RING_CAPACITY = EnvInt(1024)
SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT = EnvBool(False)
SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB = EnvFloat(0.0)
SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS = EnvInt(50)
SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB = EnvFloat(0.0)
SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB = EnvFloat(0.0)
SGLANG_KV_CANARY_PERTURB_REAL_KV_POST_FORWARD_PROB = EnvFloat(0.0)
SGLANG_KV_CANARY_PERTURB_TARGET_GROUP = EnvStr(None)
SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB = EnvFloat(0.0)
SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE = EnvBool(False)
python/sglang/srt/kv_canary/perturb/config.py
+27
View File
@@ -17,14 +17,32 @@ class TargetGroupKind(IntEnum):
@dataclass(frozen=True, slots=True, kw_only=True)
class PerturbConfig:
req_to_token_prob: float
real_kv_used_prob: float
real_kv_unused_cache_prob: float
real_kv_post_forward_prob: float
target_group_kind: TargetGroupKind | None
warmup_steps: int
@classmethod
def from_env(cls) -> "PerturbConfig":
real_kv_used_prob = envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB.get()
real_kv_unused_cache_prob = (
envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB.get()
)
real_kv_post_forward_prob = (
envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_POST_FORWARD_PROB.get()
)
return cls(
req_to_token_prob=envs.SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB.get(),
real_kv_used_prob=real_kv_used_prob,
real_kv_unused_cache_prob=real_kv_unused_cache_prob,
real_kv_post_forward_prob=real_kv_post_forward_prob,
target_group_kind=_parse_target_group_kind_from_env(
raw=envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.get(),
real_kv_used_prob=real_kv_used_prob,
real_kv_unused_cache_prob=real_kv_unused_cache_prob,
real_kv_post_forward_prob=real_kv_post_forward_prob,
),
warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(),
)
@@ -33,9 +51,18 @@ class PerturbConfig:
def _parse_target_group_kind_from_env(
*,
raw: str | None,
real_kv_used_prob: float,
real_kv_unused_cache_prob: float,
real_kv_post_forward_prob: float,
) -> TargetGroupKind | None:
if raw is not None and raw.strip():
return _parse_target_group_kind(raw)
if (
real_kv_used_prob > 0.0
or real_kv_unused_cache_prob > 0.0
or real_kv_post_forward_prob > 0.0
):
return _parse_target_group_kind(raw)
return None
python/sglang/srt/kv_canary/perturb/manager.py
+65 -1
View File
@@ -4,11 +4,18 @@ from collections.abc import Callable
from typing import TYPE_CHECKING, Optional
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
from sglang.srt.kv_canary.perturb import (
real_kv_post_forward,
real_kv_unused_cache,
real_kv_used,
req_to_token,
)
from sglang.srt.kv_canary.perturb.config import PerturbConfig
from sglang.srt.kv_canary.perturb.utils import WarmupGate
if TYPE_CHECKING:
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -17,12 +24,14 @@ class PerturbManager:
self,
*,
config: PerturbConfig,
req_to_token_pool: "ReqToTokenPool",
buffer_groups: tuple[CanaryBufferGroup, ...],
outer_step_counter_getter: Callable[[], int],
swa_window_size: int = 0,
sweep_interval: int = 0,
) -> None:
self._config = config
self._req_to_token_pool = req_to_token_pool
self._buffer_groups = buffer_groups
self._outer_step_counter_getter = outer_step_counter_getter
self._swa_window_size = swa_window_size
@@ -40,4 +49,59 @@ class PerturbManager:
*,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
) -> None:
pass
self.perturb_req_to_token(maybe_inaccurate_forward_batch)
self.perturb_real_kv_used(maybe_inaccurate_forward_batch)
self.perturb_real_kv_unused_cache(maybe_inaccurate_forward_batch)
def perturb_post_forward(
self,
*,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
) -> None:
self.perturb_real_kv_post_forward(maybe_inaccurate_forward_batch)
def perturb_req_to_token(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"]
) -> None:
req_to_token.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
config=self._config,
req_to_token_pool=self._req_to_token_pool,
warmup_gate=self._warmup_gate,
)
def perturb_real_kv_used(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"]
) -> None:
real_kv_used.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
config=self._config,
req_to_token_pool=self._req_to_token_pool,
buffer_groups=self._buffer_groups,
swa_window_size=self._swa_window_size,
warmup_gate=self._warmup_gate,
)
def perturb_real_kv_unused_cache(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"]
) -> None:
real_kv_unused_cache.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
config=self._config,
buffer_groups=self._buffer_groups,
radix_cache=self._radix_cache,
swa_window_size=self._swa_window_size,
sweep_interval=self._sweep_interval,
outer_step_counter=self._outer_step_counter_getter(),
warmup_gate=self._warmup_gate,
)
def perturb_real_kv_post_forward(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"]
) -> None:
real_kv_post_forward.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
config=self._config,
buffer_groups=self._buffer_groups,
warmup_gate=self._warmup_gate,
)
python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py
+78
View File
@@ -0,0 +1,78 @@
"""Pick a random real-KV source byte derived from out_cache_loc and flip it
in-place AFTER the TAIL kernel has captured its canary hash.
The slot id is taken from maybe_inaccurate_forward_batch.out_cache_loc and used
only as a lookup into the target group's real-KV source buffer; the actual flip
happens inside that buffer. The flip is a PyTorch indexed write on the current
CUDA stream; because TAIL is launched on the same stream, stream ordering
guarantees it happens-after TAIL's canary write.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
from sglang.srt.kv_canary.perturb.config import (
PerturbConfig,
require_target_group_kind,
)
from sglang.srt.kv_canary.perturb.slot_picker import pick_out_cache_loc_slot
from sglang.srt.kv_canary.perturb.utils import (
WarmupGate,
flip_random_source_byte_and_log,
pick_target_group,
should_run_perturbation,
)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
def run(
*,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
config: PerturbConfig,
buffer_groups: tuple[CanaryBufferGroup, ...],
warmup_gate: WarmupGate,
) -> None:
if not should_run_perturbation(
perturb_name="real_kv_post_forward",
probability=config.real_kv_post_forward_prob,
warmup_gate=warmup_gate,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
):
return
slot = pick_out_cache_loc_slot(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch
)
if slot is None:
logger.info(
"kv_canary perturb real_kv_post_forward: skipped because maybe_inaccurate_forward_batch.out_cache_loc "
"had no valid slot"
)
return
group = pick_target_group(
buffer_groups=buffer_groups,
target_kind=require_target_group_kind(
target_group_kind=config.target_group_kind,
perturb_name="real_kv_post_forward",
),
)
if group is None:
logger.info(
"kv_canary perturb real_kv_post_forward: skipped because no target group matched "
"target_group_kind=%s slot=%d",
config.target_group_kind,
slot,
)
return
flip_random_source_byte_and_log(
perturb_name="real_kv_post_forward",
group=group,
slot_idx=slot,
)
python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py
+162
View File
@@ -0,0 +1,162 @@
"""Flip the first byte of a radix-cached but currently-unused (orphan) slot.
Detection should come from sweep (per-forward verify won't even look at this
slot). Designed to surface bugs where cached KV is silently corrupted and
sleeps until much later when a prefix happens to match.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.perturb.config import (
PerturbConfig,
require_target_group_kind,
)
from sglang.srt.kv_canary.perturb.utils import (
WarmupGate,
flip_first_byte_in_source,
pick_target_group,
should_run_perturbation,
)
from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary
if TYPE_CHECKING:
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
def run(
*,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
config: PerturbConfig,
buffer_groups: tuple[CanaryBufferGroup, ...],
radix_cache: Optional["BasePrefixCache"],
swa_window_size: int,
sweep_interval: int,
outer_step_counter: int,
warmup_gate: WarmupGate,
) -> None:
if sweep_interval <= 0 or outer_step_counter % sweep_interval != 0:
return
if not should_run_perturbation(
perturb_name="real_kv_unused_cache",
probability=config.real_kv_unused_cache_prob,
warmup_gate=warmup_gate,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
require_forward_batch=False,
):
return
group = pick_target_group(
buffer_groups=buffer_groups,
target_kind=require_target_group_kind(
target_group_kind=config.target_group_kind,
perturb_name="real_kv_unused_cache",
),
)
if group is None:
logger.info(
"kv_canary perturb real_kv_unused_cache: skipped because no target group with "
"real_kv_sources_k matched target_group_kind=%s",
config.target_group_kind,
)
return
slot = _pick_sweep_slot_for_group(
radix_cache=radix_cache,
group=group,
swa_window_size=swa_window_size,
)
if slot is None:
logger.info(
"kv_canary perturb real_kv_unused_cache: skipped because no orphan sweep slot "
"was found for group=%s",
group.kind.name,
)
return
source_pick = int(torch.randint(0, len(group.real_kv_sources_k), (1,)).item())
source = group.real_kv_sources_k[source_pick]
flip_result = flip_first_byte_in_source(
group=group,
source=source,
slot_idx=slot,
slot_is_physical=True,
)
if flip_result is None:
logger.info(
"kv_canary perturb real_kv_unused_cache: skipped because slot=%d could not be mapped "
"into group=%s source_idx=%d",
slot,
group.kind.name,
source_pick,
)
return
row, col, original_byte = flip_result
logger.info(
"kv_canary perturb real_kv_unused_cache: group=%s source_idx=%d slot=%d row=%d col=%d "
"original_byte=0x%02X new_byte=0x%02X",
group.kind.name,
source_pick,
slot,
row,
col,
original_byte,
original_byte ^ 0xFF,
)
def _pick_sweep_slot_for_group(
*,
radix_cache: Optional["BasePrefixCache"],
group: CanaryBufferGroup,
swa_window_size: int,
) -> Optional[int]:
if radix_cache is None:
return None
walk_result = walk_radix_cache_for_canary(
radix_cache=radix_cache,
unlocked_only=True,
swa_resident_only=group.kind is PoolKind.SWA,
)
slots = [
int(raw_slot)
for raw_slot in walk_result.slot_indices.detach().to("cpu").tolist()
if int(raw_slot) >= 0
]
if group.kind is PoolKind.SWA:
slots = _translate_full_slots_to_swa_slots(
slots=slots,
full_to_swa_index_mapping=group.swa_index_lut,
)
if not slots:
return None
pick = int(torch.randint(0, len(slots), (1,)).item())
return slots[pick]
def _translate_full_slots_to_swa_slots(
*,
slots: list[int],
full_to_swa_index_mapping: Optional[torch.Tensor],
) -> list[int]:
if full_to_swa_index_mapping is None:
return []
lut = full_to_swa_index_mapping.detach().to("cpu").to(torch.int64)
translated: list[int] = []
for slot in slots:
if slot >= int(lut.shape[0]):
continue
physical_slot = int(lut[slot].item())
if physical_slot >= 0:
translated.append(physical_slot)
return translated
python/sglang/srt/kv_canary/perturb/real_kv_used.py
+134
View File
@@ -0,0 +1,134 @@
"""Flip the first byte of a slot currently being used by an active req.
Detection should come from per-forward verify (HEAD/TAIL kernel), NOT from
sweep. Designed to surface CUDA-graph-idle-class bugs where production reads
a slot whose KV byte was silently overwritten.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.perturb.config import (
PerturbConfig,
require_target_group_kind,
)
from sglang.srt.kv_canary.perturb.slot_picker import (
ReqToTokenEntry,
collect_active_slots,
)
from sglang.srt.kv_canary.perturb.utils import (
WarmupGate,
flip_first_byte_in_source,
pick_target_group,
should_run_perturbation,
)
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
def run(
*,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
config: PerturbConfig,
req_to_token_pool: "ReqToTokenPool",
buffer_groups: tuple[CanaryBufferGroup, ...],
swa_window_size: int,
warmup_gate: WarmupGate,
) -> None:
if not should_run_perturbation(
perturb_name="real_kv_used",
probability=config.real_kv_used_prob,
warmup_gate=warmup_gate,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
):
return
group = pick_target_group(
buffer_groups=buffer_groups,
target_kind=require_target_group_kind(
target_group_kind=config.target_group_kind,
perturb_name="real_kv_used",
),
)
if group is None:
logger.info(
"kv_canary perturb real_kv_used: skipped because no target group with real_kv_sources_k "
"matched target_group_kind=%s",
config.target_group_kind,
)
return
target = _pick_active_slot_for_group(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
req_to_token_pool=req_to_token_pool,
group=group,
swa_window_size=swa_window_size,
)
if target is None:
logger.info(
"kv_canary perturb real_kv_used: skipped because no active slot was found "
"for group=%s",
group.kind.name,
)
return
source_pick = int(torch.randint(0, len(group.real_kv_sources_k), (1,)).item())
source = group.real_kv_sources_k[source_pick]
flip_result = flip_first_byte_in_source(
group=group, source=source, slot_idx=target.value
)
if flip_result is None:
logger.info(
"kv_canary perturb real_kv_used: skipped because slot=%d could not be mapped into "
"group=%s source_idx=%d",
target.value,
group.kind.name,
source_pick,
)
return
row, col, original_byte = flip_result
logger.info(
"kv_canary perturb real_kv_used: group=%s source_idx=%d slot=%d row=%d col=%d "
"original_byte=0x%02X new_byte=0x%02X",
group.kind.name,
source_pick,
target.value,
row,
col,
original_byte,
original_byte ^ 0xFF,
)
def _pick_active_slot_for_group(
*,
maybe_inaccurate_forward_batch: "ForwardBatch",
req_to_token_pool: "ReqToTokenPool",
group: CanaryBufferGroup,
swa_window_size: int,
) -> Optional[ReqToTokenEntry]:
candidates = collect_active_slots(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
req_to_token_pool=req_to_token_pool,
exclude_out_cache_loc=True,
)
if group.kind is PoolKind.SWA:
candidates = [
entry
for entry in candidates
if entry.position >= max(0, entry.seq_len - swa_window_size)
]
if not candidates:
return None
pick = int(torch.randint(0, len(candidates), (1,)).item())
return candidates[pick]
python/sglang/srt/kv_canary/perturb/req_to_token.py
+75
View File
@@ -0,0 +1,75 @@
"""Flip the req_to_token pointer of a currently-active req.
The hook picks a random (req_pool_idx, position, value) from active reqs,
filtering out entries whose value is 0 (so slot 0 is excluded), and overwrites
req_to_token[req_pool_idx, position] with another active req's slot id.
KV bytes are not touched.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.kv_canary.perturb.config import PerturbConfig
from sglang.srt.kv_canary.perturb.slot_picker import collect_active_slots
from sglang.srt.kv_canary.perturb.utils import WarmupGate, should_run_perturbation
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
def run(
*,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
config: PerturbConfig,
req_to_token_pool: "ReqToTokenPool",
warmup_gate: WarmupGate,
) -> None:
if not should_run_perturbation(
perturb_name="req_to_token",
probability=config.req_to_token_prob,
warmup_gate=warmup_gate,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
):
return
entries = collect_active_slots(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
req_to_token_pool=req_to_token_pool,
exclude_out_cache_loc=True,
)
entries = [entry for entry in entries if entry.value >= 1]
if not entries:
logger.info(
"kv_canary perturb req_to_token: skipped because no active nonzero slots were found"
)
return
pick = int(torch.randint(0, len(entries), (1,)).item())
target = entries[pick]
replacement_values = [item.value for item in entries if item.value != target.value]
if not replacement_values:
logger.info(
"kv_canary perturb req_to_token: skipped because no replacement slot differs from "
"original_slot=%d",
target.value,
)
return
replacement_pick = int(torch.randint(0, len(replacement_values), (1,)).item())
new_value = replacement_values[replacement_pick]
req_to_token = req_to_token_pool.req_to_token
logger.info(
"kv_canary perturb req_to_token: req_pool_idx=%d position=%d original_slot=%d new_slot=%d",
target.req_pool_idx,
target.position,
target.value,
new_value,
)
req_to_token[target.req_pool_idx, target.position] = new_value
python/sglang/srt/kv_canary/perturb/slot_picker.py
+100
View File
@@ -0,0 +1,100 @@
from __future__ import annotations
import random
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
import torch
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@dataclass(frozen=True, slots=True, kw_only=True)
class ReqToTokenEntry:
req_pool_idx: int
position: int
value: int
seq_len: int = 0
def collect_active_slots(
*,
maybe_inaccurate_forward_batch: "ForwardBatch",
req_to_token_pool: "ReqToTokenPool",
exclude_out_cache_loc: bool = True,
) -> list[ReqToTokenEntry]:
"""Collect every (req_pool_idx, position, value) triple for currently-active reqs.
Excludes slots in ``maybe_inaccurate_forward_batch.out_cache_loc`` when ``exclude_out_cache_loc=True``
so a slot the current forward is about to write isn't picked (write race).
"""
req_pool_indices = maybe_inaccurate_forward_batch.req_pool_indices
seq_lens = maybe_inaccurate_forward_batch.seq_lens
if req_pool_indices is None or seq_lens is None:
return []
req_to_token = req_to_token_pool.req_to_token
if not isinstance(req_to_token, torch.Tensor) or req_to_token.numel() == 0:
return []
excluded: set[int] = set()
if exclude_out_cache_loc:
out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc
if out_cache_loc is not None:
valid_num_tokens = maybe_inaccurate_forward_batch.num_token_non_padded_cpu
if valid_num_tokens is None:
valid_num_tokens = int(out_cache_loc.shape[0])
excluded = set(
int(x)
for x in out_cache_loc[:valid_num_tokens].detach().to("cpu").tolist()
)
req_pool_indices_list = req_pool_indices.detach().to("cpu").tolist()
seq_lens_list = seq_lens.detach().to("cpu").tolist()
rows, cols = int(req_to_token.shape[0]), int(req_to_token.shape[1])
candidates: list[ReqToTokenEntry] = []
for req_pool_idx, seq_len in zip(req_pool_indices_list, seq_lens_list):
req_pool_idx_int = int(req_pool_idx)
seq_len_int = int(seq_len)
if req_pool_idx_int < 0 or req_pool_idx_int >= rows:
continue
upper = min(seq_len_int, cols)
if upper <= 0:
continue
row_values = req_to_token[req_pool_idx_int, :upper].detach().to("cpu").tolist()
candidates.extend(
ReqToTokenEntry(
req_pool_idx=req_pool_idx_int,
position=pos,
value=value,
seq_len=seq_len_int,
)
for pos, raw_value in enumerate(row_values)
if (value := int(raw_value)) >= 0 and value not in excluded
)
return candidates
def pick_out_cache_loc_slot(
*, maybe_inaccurate_forward_batch: "ForwardBatch"
) -> Optional[int]:
out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc
if out_cache_loc is None:
return None
total = int(out_cache_loc.shape[0])
if total <= 0:
return None
valid_num_tokens = maybe_inaccurate_forward_batch.num_token_non_padded_cpu
if valid_num_tokens is None:
valid_num_tokens = total
valid_num_tokens = int(valid_num_tokens)
if valid_num_tokens <= 0:
return None
pick = random.randrange(valid_num_tokens)
slot = int(out_cache_loc[pick].item())
if slot < 0:
return None
return slot
python/sglang/srt/kv_canary/runner/canary_manager.py
+4
View File
@@ -101,6 +101,7 @@ class CanaryManager:
)
self._perturb_manager = PerturbManager(
config=perturb_config,
req_to_token_pool=req_to_token_pool,
buffer_groups=self._buffer_groups,
outer_step_counter_getter=self._get_outer_step_counter,
swa_window_size=self._swa_window_size,
@@ -205,6 +206,9 @@ class CanaryManager:
) -> None:
for idx in single_forward_indices:
self._single_forward_managers[idx].post_ops_outside_graph()
self._perturb_manager.perturb_post_forward(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch
)
self._sweep_orchestrator.maybe_run_sweep()
self._outer_step_counter += 1
self._violation_manager.step()
python/sglang/test/kv_canary/pd_fixture.py
+89
View File
@@ -0,0 +1,89 @@
from __future__ import annotations
from typing import ClassVar, Literal, Optional
from sglang.srt.kv_canary.config import CanaryMode
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.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
_SHORT_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 8).strip()
class CanaryPDFixture(CanaryViolationAssertMixin, PDDisaggregationServerBase):
capture_per_side_logs = True
model_mode: ClassVar[Literal["mha", "swa", "dsv4"]]
kv_canary_mode: ClassVar[CanaryMode] = CanaryMode.LOG
extra_server_args: ClassVar[tuple[str, ...]] = (
"--kv-canary-real-data",
"partial",
"--skip-server-warmup",
)
_cfg: ClassVar[Optional[_ModeConfig]] = None
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
cls._cfg = _MODE_CONFIGS[cls.model_mode]
cls.model = cls._cfg.model_path
canary_args = build_canary_server_args(
kv_canary_mode=cls.kv_canary_mode,
mode_cfg=cls._cfg,
extra_server_args=cls.extra_server_args,
)
cls.extra_prefill_args = list(canary_args)
cls.extra_decode_args = list(canary_args)
if cls.model_mode == "swa":
# SWA mode uses google/gemma-4-E2B-it, whose forward does a
# ``positions += 1`` in-place. canary's WRITE/VERIFY require
# forward_batch.positions to stay 0-indexed, so flip the gemma
# path to out-of-place shift for these tests.
cls.extra_prefill_env = {
**cls.extra_prefill_env,
"SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION": "1",
}
cls.extra_decode_env = {
**cls.extra_decode_env,
"SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION": "1",
}
cls.launch_all()
def send_parallel_short_requests(
self,
n: int,
*,
assert_all_success: bool = True,
max_new_tokens: int = 100,
timeout: float = 60.0,
) -> list[dict]:
results = post_parallel_generate(
url=self.lb_url + "/generate",
prompts=[_SHORT_PROMPT_BODY] * 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 _captured_log_text(
self, side: Optional[Literal["prefill", "decode"]] = None
) -> str:
if side == "prefill":
stdout_buf = type(self)._prefill_stdout_buf
stderr_buf = type(self)._prefill_stderr_buf
elif side == "decode":
stdout_buf = type(self)._decode_stdout_buf
stderr_buf = type(self)._decode_stderr_buf
else:
raise ValueError(f"Unsupported side={side!r}")
stdout_text = stdout_buf.getvalue() if stdout_buf is not None else ""
stderr_text = stderr_buf.getvalue() if stderr_buf is not None else ""
return stdout_text + stderr_text
python/sglang/test/kv_canary/runner_test_base.py
+4
View File
@@ -54,6 +54,10 @@ def make_perturb_config() -> PerturbConfig:
"""Build a PerturbConfig with every probability pinned to 0 so the
perturb hooks do nothing during unit tests."""
return PerturbConfig(
req_to_token_prob=0.0,
real_kv_used_prob=0.0,
real_kv_unused_cache_prob=0.0,
real_kv_post_forward_prob=0.0,
target_group_kind=None,
warmup_steps=0,
)