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

Add a sliding-window-attention divergence reporter for the KV-canary (#26820)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-31 09:59:28 +08:00
committed by GitHub
parent ae9db7ff4b
commit 7dd19ae3d8
13 changed files with 937 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -763,6 +763,7 @@ class Envs:
SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB = EnvFloat(0.0)
SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE = EnvBool(False)
SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT = EnvBool(False)
SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL = EnvInt(0)
SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False)
python/sglang/srt/kv_canary/api.py
+6
View File
@@ -11,6 +11,7 @@ from sglang.srt.kv_canary.perturb.config import PerturbConfig
from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
if TYPE_CHECKING:
@@ -48,6 +49,10 @@ def install_canary(
device=device,
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
)
allocator = model_runner.token_to_kv_pool_allocator
swa_allocator = (
allocator if isinstance(allocator, SWATokenToKVPoolAllocator) else None
)
launch_capacities = CanaryLaunchCapacities.from_args(
server_args=model_runner.server_args,
req_to_token_pool_size=model_runner.req_to_token_pool.size,
@@ -65,6 +70,7 @@ def install_canary(
launch_capacities=launch_capacities,
swa_window_size=swa_window_size,
token_oracle_manager=token_oracle_manager,
swa_allocator=swa_allocator,
speculative_num_steps=speculative_num_steps,
is_eagle_draft_decode=model_runner.is_draft_worker,
)
python/sglang/srt/kv_canary/runner/canary_manager.py
+26
View File
@@ -18,6 +18,7 @@ from sglang.srt.kv_canary.endpoint import (
)
from sglang.srt.kv_canary.perturb.config import PerturbConfig
from sglang.srt.kv_canary.perturb.manager import PerturbManager
from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceReporter
from sglang.srt.kv_canary.runner.sweep import SweepOrchestrator
from sglang.srt.kv_canary.runner.violation_manager import ViolationManager
from sglang.srt.kv_canary.single_forward_manager.manager import (
@@ -30,6 +31,7 @@ from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
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.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
@@ -47,12 +49,14 @@ class CanaryManager:
launch_capacities: CanaryLaunchCapacities,
swa_window_size: int = 0,
token_oracle_manager: Optional[TokenOracleManager] = None,
swa_allocator: Optional["SWATokenToKVPoolAllocator"] = None,
speculative_num_steps: int = 1,
is_eagle_draft_decode: bool = False,
) -> None:
self.config = config
self._req_to_token_pool = req_to_token_pool
self._swa_window_size = swa_window_size
self._swa_allocator: Optional["SWATokenToKVPoolAllocator"] = swa_allocator
self._outer_step_counter: int = 0
self._active_single_forward_manager_index: Optional[int] = None
@@ -85,6 +89,22 @@ class CanaryManager:
self._d2h_stream: torch.cuda.Stream = torch.cuda.Stream(device=device)
swa_divergence_interval = (
envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.get()
)
if swa_divergence_interval > 0:
self._swa_divergence_report: Optional[SwaDivergenceReporter] = (
SwaDivergenceReporter(
device=device,
d2h_stream=self._d2h_stream,
interval=swa_divergence_interval,
swa_allocator=self._swa_allocator,
req_to_token_pool=self._req_to_token_pool,
)
)
else:
self._swa_divergence_report = None
self._violation_manager = ViolationManager(
config=config,
device_state=self._device_state,
@@ -122,6 +142,7 @@ class CanaryManager:
per_forward_write_entry_capacity=launch_capacities.per_forward_write_entry_capacity,
d2h_stream=self._d2h_stream,
token_oracle_manager=token_oracle_manager,
swa_divergence_report=self._swa_divergence_report,
is_eagle_draft_decode=is_eagle_draft_decode,
)
for _ in range(num_sfms)
@@ -212,6 +233,11 @@ class CanaryManager:
self._sweep_orchestrator.maybe_run_sweep()
self._outer_step_counter += 1
self._violation_manager.step()
if self._swa_divergence_report is not None:
self._swa_divergence_report.step(
outer_step_counter=self._outer_step_counter,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
)
def mark_init_finished(self) -> None:
for single_forward_manager in self._single_forward_managers:
python/sglang/srt/kv_canary/runner/swa_divergence.py
+202
View File
@@ -0,0 +1,202 @@
from __future__ import annotations
import json
import logging
import re
from dataclasses import asdict, dataclass
from typing import TYPE_CHECKING, Any, Optional
import torch
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
_SWA_DIVERGENCE_LOG_PREFIX: str = "kv_canary_swa_divergence="
_SWA_DIVERGENCE_LINE_RE = re.compile(re.escape(_SWA_DIVERGENCE_LOG_PREFIX) + r"(\S+)")
_FULL_IDX = 0
_SWA_IDX = 1
class SwaDivergenceReporter:
def __init__(
self,
*,
device: torch.device,
d2h_stream: torch.cuda.Stream,
interval: int,
swa_allocator: Optional["SWATokenToKVPoolAllocator"] = None,
req_to_token_pool: Optional["ReqToTokenPool"] = None,
) -> None:
self._interval = interval
self._swa_allocator = swa_allocator
self._req_to_token_pool = req_to_token_pool
self._forward_ct: int = 0
# Per-group running total of verify entries (shape ``[2]``, int32).
self.verify_total_count_device: torch.Tensor = torch.zeros(
2, dtype=torch.int32, device=device
)
self._handler = DelayedDeviceHostHandler(d2h_stream=d2h_stream)
def observe_after_invoke_plan(
self, *, group: CanaryBufferGroup, verify_plan: VerifyPlan
) -> None:
idx = _FULL_IDX if group.kind is PoolKind.FULL else _SWA_IDX
# verify_num_valid is shape [1]; slice to a length-1 view so the in-place add
# has matching ranks (else torch refuses the broadcast into shape []).
self.verify_total_count_device[idx : idx + 1].add_(verify_plan.verify_num_valid)
def step(
self,
*,
outer_step_counter: int,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
) -> None:
self._forward_ct += 1
self._handler.step(
compute_on_device=lambda: self._compute_on_device(
outer_step_counter=outer_step_counter,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
),
postprocess_on_host=self._postprocess_on_host,
)
def _compute_on_device(
self,
*,
outer_step_counter: int,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"],
) -> Optional[dict[str, Any]]:
if outer_step_counter == 0 or outer_step_counter % self._interval != 0:
return None
result: dict[str, Any] = {
"forward_ct": self._forward_ct,
"verify_total_count": self.verify_total_count_device,
}
if (
self._swa_allocator is not None
and maybe_inaccurate_forward_batch is not None
):
result["swa_full_idx_divergence"] = compute_swa_full_idx_divergence(
swa_allocator=self._swa_allocator,
req_to_token_pool=self._req_to_token_pool,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
)
result["swa_out_of_window_tokens"] = compute_swa_out_of_window_tokens(
swa_allocator=self._swa_allocator,
req_to_token_pool=self._req_to_token_pool,
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
)
return result
def _postprocess_on_host(self, host_data: dict[str, Any]) -> None:
verify_totals = host_data["verify_total_count"].tolist()
swa_full_idx_divergence = (
int(x.item())
if (x := host_data.get("swa_full_idx_divergence")) is not None
else 0
)
swa_out_of_window_tokens = (
int(x.item())
if (x := host_data.get("swa_out_of_window_tokens")) is not None
else 0
)
logger.info(
SwaDivergenceLog(
forward_ct=host_data["forward_ct"],
verify_full=int(verify_totals[_FULL_IDX]),
verify_swa=int(verify_totals[_SWA_IDX]),
swa_full_idx_divergence=swa_full_idx_divergence,
swa_out_of_window_tokens=swa_out_of_window_tokens,
).format()
)
@dataclass(frozen=True, slots=True, kw_only=True)
class SwaDivergenceLog:
forward_ct: int
verify_full: int
verify_swa: int
swa_full_idx_divergence: int
swa_out_of_window_tokens: int = 0
def format(self) -> str:
return _SWA_DIVERGENCE_LOG_PREFIX + json.dumps(
asdict(self), separators=(",", ":")
)
@classmethod
def parse(cls, line: str) -> Optional["SwaDivergenceLog"]:
match = _SWA_DIVERGENCE_LINE_RE.search(line)
if match is None:
return None
return cls(**json.loads(match.group(1)))
@classmethod
def find_last(cls, text: str) -> Optional[tuple["SwaDivergenceLog", str]]:
last_match: Optional[re.Match] = None
for match in _SWA_DIVERGENCE_LINE_RE.finditer(text):
last_match = match
if last_match is None:
return None
return cls(**json.loads(last_match.group(1))), last_match.group(0)
def compute_swa_out_of_window_tokens(
*,
swa_allocator: "SWATokenToKVPoolAllocator",
req_to_token_pool: "ReqToTokenPool",
maybe_inaccurate_forward_batch: "ForwardBatch",
) -> torch.Tensor:
"""Count tokens in the live req_to_token range whose SWA mapping is 0 (out-of-window)."""
full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping
device = full_to_swa_index_mapping.device
req_pool_indices = maybe_inaccurate_forward_batch.req_pool_indices
seq_lens = maybe_inaccurate_forward_batch.seq_lens
if req_pool_indices.numel() == 0:
return torch.zeros(1, dtype=torch.int32, device=device)
req_to_token = req_to_token_pool.req_to_token
rows = req_to_token[req_pool_indices]
positions = torch.arange(rows.shape[1], device=rows.device)
mask = positions[None, :] < seq_lens[:, None]
swa_indices = full_to_swa_index_mapping[rows]
return ((swa_indices == 0) & mask).sum().to(torch.int32).view(1)
def compute_swa_full_idx_divergence(
*,
swa_allocator: "SWATokenToKVPoolAllocator",
req_to_token_pool: "ReqToTokenPool",
maybe_inaccurate_forward_batch: "ForwardBatch",
) -> torch.Tensor:
"""Count non-identity (full, swa) index pairs in the live req_to_token range."""
full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping
device = full_to_swa_index_mapping.device
req_pool_indices = maybe_inaccurate_forward_batch.req_pool_indices
seq_lens = maybe_inaccurate_forward_batch.seq_lens
if req_pool_indices.numel() == 0:
return torch.zeros(1, dtype=torch.int32, device=device)
req_to_token = req_to_token_pool.req_to_token
rows = req_to_token[req_pool_indices]
positions = torch.arange(rows.shape[1], device=rows.device)
mask = positions[None, :] < seq_lens[:, None]
swa_indices = full_to_swa_index_mapping[rows]
# FULL pool slots beyond the sliding window have their SWA mapping written
# to 0 (see SWATokenToKVPoolAllocator.alloc_extend); skip those so they
# don't get counted as divergence.
return (
((swa_indices != rows) & mask & (swa_indices != 0))
.sum()
.to(torch.int32)
.view(1)
)
python/sglang/srt/kv_canary/single_forward_manager/data.py
+15
View File
@@ -11,6 +11,7 @@ class PostOpsInsideGraphOutputBuffer:
kernel_run_counters: torch.Tensor
slot_run_counters: torch.Tensor
violation_write_index: torch.Tensor
swa_verify_total_count: torch.Tensor | None
@classmethod
def allocate(
@@ -18,6 +19,7 @@ class PostOpsInsideGraphOutputBuffer:
*,
num_kernel_tags: int,
num_slot_tags: int,
swa_verify_total_count_shape: tuple[int, ...] | None,
device: torch.device,
) -> "PostOpsInsideGraphOutputBuffer":
return cls(
@@ -29,6 +31,13 @@ class PostOpsInsideGraphOutputBuffer:
num_slot_tags, dtype=torch.int64, device=device
),
violation_write_index=torch.zeros(1, dtype=torch.int32, device=device),
swa_verify_total_count=(
None
if swa_verify_total_count_shape is None
else torch.zeros(
swa_verify_total_count_shape, dtype=torch.int32, device=device
)
),
)
def copy_from(
@@ -38,8 +47,14 @@ class PostOpsInsideGraphOutputBuffer:
kernel_run_counters: torch.Tensor,
slot_run_counters: torch.Tensor,
violation_write_index: torch.Tensor,
swa_verify_total_count: torch.Tensor | None,
) -> None:
self.verify_plan_enable.copy_(verify_plan_enable)
self.kernel_run_counters.copy_(kernel_run_counters)
self.slot_run_counters.copy_(slot_run_counters)
self.violation_write_index.copy_(violation_write_index)
assert (self.swa_verify_total_count is not None) == (
swa_verify_total_count is not None
)
if self.swa_verify_total_count is not None:
self.swa_verify_total_count.copy_(swa_verify_total_count)
python/sglang/srt/kv_canary/single_forward_manager/manager.py
+20
View File
@@ -21,6 +21,7 @@ from sglang.srt.kv_canary.runner.kernel_launcher import (
invoke_plan,
launch_endpoints_per_forward,
)
from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceReporter
from sglang.srt.kv_canary.single_forward_manager.data import (
PostOpsInsideGraphOutputBuffer,
)
@@ -72,6 +73,7 @@ class SingleForwardManager:
per_forward_write_entry_capacity: int,
d2h_stream: torch.cuda.Stream,
token_oracle_manager: Optional[TokenOracleManager],
swa_divergence_report: Optional[SwaDivergenceReporter],
is_eagle_draft_decode: bool,
) -> None:
self._config = config
@@ -83,6 +85,9 @@ class SingleForwardManager:
self._swa_window_size = swa_window_size
self._d2h_stream = d2h_stream
self._token_oracle_manager: Optional[TokenOracleManager] = token_oracle_manager
self._swa_divergence_report: Optional[SwaDivergenceReporter] = (
swa_divergence_report
)
self._is_eagle_draft_decode: bool = is_eagle_draft_decode
self._write_req_capacity = per_forward_write_req_capacity
@@ -101,6 +106,11 @@ class SingleForwardManager:
self._output_buffer = PostOpsInsideGraphOutputBuffer.allocate(
num_kernel_tags=int(device_state.kernel_run_counters.shape[0]),
num_slot_tags=int(device_state.slot_run_counters.shape[0]),
swa_verify_total_count_shape=(
None
if swa_divergence_report is None
else tuple(swa_divergence_report.verify_total_count_device.shape)
),
device=device,
)
@@ -201,6 +211,11 @@ class SingleForwardManager:
swa_window_size=self._swa_window_size,
req_to_verify_expected_tokens=self._device_state.req_to_verify_expected_tokens,
)
if self._swa_divergence_report is not None:
self._swa_divergence_report.observe_after_invoke_plan(
group=group,
verify_plan=verify_plan,
)
launch_endpoints_per_forward(
endpoints=self._endpoints,
group=group,
@@ -261,6 +276,11 @@ class SingleForwardManager:
kernel_run_counters=self._device_state.kernel_run_counters,
slot_run_counters=self._device_state.slot_run_counters,
violation_write_index=self._device_state.violation_log.violation_write_index,
swa_verify_total_count=(
None
if self._swa_divergence_report is None
else self._swa_divergence_report.verify_total_count_device
),
)
def post_ops_outside_graph(self) -> None:
python/sglang/test/kv_canary/e2e_base.py
+69
View File
@@ -3,9 +3,11 @@ from __future__ import annotations
import io
import os
import string
import time
from typing import ClassVar, Literal, Optional
from sglang.srt.kv_canary.config import CanaryMode
from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceLog
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
@@ -81,6 +83,9 @@ class CanaryE2EBase(CapturedServerE2EBase):
server_env.setdefault("SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT", "1")
server_env.update(cls.extra_env)
if cls.model_mode == "swa":
server_env.setdefault(
"SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL", "20"
)
# 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
@@ -134,6 +139,70 @@ class CanaryE2EBase(CapturedServerE2EBase):
self.assertEqual(result.get("status_code"), 200, result)
return results
def maybe_assert_swa_divergence_observed(self) -> None:
if self.model_mode == "swa":
self.assert_swa_divergence_observed()
def assert_swa_divergence_observed(
self,
*,
min_swa_out_of_window_tokens: int = 1,
min_swa_full_idx_divergence: int = 1,
require_verify_lag: bool = True,
flush_wait_seconds: float = 3.0,
max_retries: int = 10,
) -> None:
"""Assert that the SWA path was genuinely exercised.
Three signals must all hold:
- ``swa_out_of_window_tokens >= 1``: at least one prefix token has been clipped
out of the sliding window (its SWA mapping is 0). Any prompt longer than the
SWA window produces this — proves the SWA window slide actually ran.
- ``swa_full_idx_divergence >= 1``: SWA pool has actually remapped at least one
slot to a non-identity index (i.e. real slot reuse / eviction occurred). The
workload must drive SWA pool pressure for this to fire — required because the
"pool reuse" path is the one production hits under sustained long-context
traffic, and we must keep it covered.
- ``verify_swa < verify_full``: SWA verify kernel processed fewer tokens than
FULL — proves both kernel groups ran and the window short-circuited SWA.
"""
last_parsed = None
last_line: str = ""
for _ in range(max_retries):
time.sleep(flush_wait_seconds)
log_text = self._captured_log_text()
found = SwaDivergenceLog.find_last(log_text)
if found is not None:
last_parsed, last_line = found
break
if last_parsed is None:
raise AssertionError(
"No kv_canary swa_divergence line found in server log after "
f"{max_retries} retries (wait={flush_wait_seconds}s each). "
f"Log tail:\n{self._captured_log_text()[-2000:]}"
)
if last_parsed.swa_out_of_window_tokens < min_swa_out_of_window_tokens:
raise AssertionError(
f"SWA path not exercised: swa_out_of_window_tokens={last_parsed.swa_out_of_window_tokens} "
f"< min={min_swa_out_of_window_tokens}. Line: {last_line}"
)
if last_parsed.swa_full_idx_divergence < min_swa_full_idx_divergence:
raise AssertionError(
f"SWA pool reuse not exercised: swa_full_idx_divergence={last_parsed.swa_full_idx_divergence} "
f"< min={min_swa_full_idx_divergence}. The workload did not drive enough SWA pool pressure "
f"to force slot remap. Line: {last_line}"
)
if require_verify_lag and not (
last_parsed.verify_swa < last_parsed.verify_full
):
raise AssertionError(
f"SWA path not exercised: verify_swa={last_parsed.verify_swa} "
f"not strictly less than verify_full={last_parsed.verify_full}. "
f"Line: {last_line}"
)
def _make_unique_prompts(n: int) -> list[str]:
if n > len(_UNIQUE_PROMPT_FIRST_CHARS):