Add kv_canary PP self-test fixture and SWA divergence coverage (#27410)

This commit is contained in:
fzyzcjy
2026-06-06 09:06:19 +08:00
committed by GitHub
parent b3e4c204fd
commit 21201ef718
8 changed files with 208 additions and 21 deletions
@@ -149,6 +149,13 @@ class SwaDivergenceLog:
return None return None
return cls(**json.loads(last_match.group(1))), last_match.group(0) return cls(**json.loads(last_match.group(1))), last_match.group(0)
@classmethod
def find_all(cls, text: str) -> list[tuple["SwaDivergenceLog", str]]:
return [
(cls(**json.loads(match.group(1))), match.group(0))
for match in _SWA_DIVERGENCE_LINE_RE.finditer(text)
]
def compute_swa_out_of_window_tokens( def compute_swa_out_of_window_tokens(
*, *,
+31 -17
View File
@@ -126,13 +126,17 @@ class CanaryE2EBase(CapturedServerE2EBase):
assert_all_success: bool = True, assert_all_success: bool = True,
max_new_tokens: int = 2048, max_new_tokens: int = 2048,
timeout: float = 240.0, timeout: float = 240.0,
ignore_eos: Optional[bool] = None,
) -> list[dict]: ) -> list[dict]:
"""Fan out n parallel /generate requests; return list of response dicts.""" """Fan out n parallel /generate requests; return list of response dicts."""
if ignore_eos is None:
ignore_eos = self.model_mode == "swa"
results = post_parallel_generate( results = post_parallel_generate(
url=self.base_url + "/generate", url=self.base_url + "/generate",
prompts=self.make_prompts(n), prompts=self.make_prompts(n),
max_new_tokens=max_new_tokens, max_new_tokens=max_new_tokens,
timeout=timeout, timeout=timeout,
ignore_eos=ignore_eos,
) )
if assert_all_success: if assert_all_success:
for result in results: for result in results:
@@ -155,9 +159,9 @@ class CanaryE2EBase(CapturedServerE2EBase):
"""Assert that the SWA path was genuinely exercised. """Assert that the SWA path was genuinely exercised.
Three signals must all hold: Three signals must all hold:
- ``swa_out_of_window_tokens >= 1``: at least one prefix token has been clipped - ``swa_out_of_window_tokens >= 1``: at least one token has slid out of the
out of the sliding window (its SWA mapping is 0). Any prompt longer than the sliding window (its SWA mapping is 0). This only appears once a request decodes
SWA window produces this — proves the SWA window slide actually ran. past the window, so the window evicts — proves the SWA window slide actually ran.
- ``swa_full_idx_divergence >= 1``: SWA pool has actually remapped at least one - ``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 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 workload must drive SWA pool pressure for this to fire — required because the
@@ -165,34 +169,44 @@ class CanaryE2EBase(CapturedServerE2EBase):
traffic, and we must keep it covered. traffic, and we must keep it covered.
- ``verify_swa < verify_full``: SWA verify kernel processed fewer tokens than - ``verify_swa < verify_full``: SWA verify kernel processed fewer tokens than
FULL — proves both kernel groups ran and the window short-circuited SWA. FULL — proves both kernel groups ran and the window short-circuited SWA.
The first two signals are checked as the *peak* across all sampled forwards, not
only the last sample. The divergence reporter snapshots one live forward batch per
interval; under PP it snapshots a single micro-batch, which may hold only in-window
requests even when another micro-batch diverged. "Was the SWA path ever exercised?"
is a max-over-samples question, so a trailing in-window sample must not mask an
earlier diverging one. ``verify_swa``/``verify_full`` are monotonic running totals,
so the lag check reads the last sample.
""" """
last_parsed = None samples: list[tuple[SwaDivergenceLog, str]] = []
last_line: str = ""
for _ in range(max_retries): for _ in range(max_retries):
time.sleep(flush_wait_seconds) time.sleep(flush_wait_seconds)
log_text = self._captured_log_text() samples = SwaDivergenceLog.find_all(self._captured_log_text())
found = SwaDivergenceLog.find_last(log_text) if samples:
if found is not None:
last_parsed, last_line = found
break break
if last_parsed is None: if not samples:
raise AssertionError( raise AssertionError(
"No kv_canary swa_divergence line found in server log after " "No kv_canary swa_divergence line found in server log after "
f"{max_retries} retries (wait={flush_wait_seconds}s each). " f"{max_retries} retries (wait={flush_wait_seconds}s each). "
f"Log tail:\n{self._captured_log_text()[-2000:]}" f"Log tail:\n{self._captured_log_text()[-2000:]}"
) )
if last_parsed.swa_out_of_window_tokens < min_swa_out_of_window_tokens: peak_out_of_window = max(p.swa_out_of_window_tokens for p, _ in samples)
peak_full_idx_divergence = max(p.swa_full_idx_divergence for p, _ in samples)
last_parsed, last_line = samples[-1]
if peak_out_of_window < min_swa_out_of_window_tokens:
raise AssertionError( raise AssertionError(
f"SWA path not exercised: swa_out_of_window_tokens={last_parsed.swa_out_of_window_tokens} " f"SWA path not exercised: peak swa_out_of_window_tokens={peak_out_of_window} "
f"< min={min_swa_out_of_window_tokens}. Line: {last_line}" f"< min={min_swa_out_of_window_tokens} across {len(samples)} samples. "
f"Last line: {last_line}"
) )
if last_parsed.swa_full_idx_divergence < min_swa_full_idx_divergence: if peak_full_idx_divergence < min_swa_full_idx_divergence:
raise AssertionError( raise AssertionError(
f"SWA pool reuse not exercised: swa_full_idx_divergence={last_parsed.swa_full_idx_divergence} " f"SWA pool reuse not exercised: peak swa_full_idx_divergence={peak_full_idx_divergence} "
f"< min={min_swa_full_idx_divergence}. The workload did not drive enough SWA pool pressure " f"< min={min_swa_full_idx_divergence} across {len(samples)} samples. The workload "
f"to force slot remap. Line: {last_line}" f"did not drive enough SWA pool pressure to force slot remap. Last line: {last_line}"
) )
if require_verify_lag and not ( if require_verify_lag and not (
last_parsed.verify_swa < last_parsed.verify_full last_parsed.verify_swa < last_parsed.verify_full
@@ -0,0 +1,25 @@
from __future__ import annotations
from typing import ClassVar
from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
PP_SIZE: int = 2
class CanaryPPFixture(CanaryE2EBase):
model_mode: ClassVar[str] = "swa"
workload_n_batches: ClassVar[int] = 2
@classmethod
def setUpClass(cls) -> None:
cls.extra_server_args = (
"--pp-size",
str(PP_SIZE),
"--disable-cuda-graph",
*SWA_POOL_SERVER_ARGS,
*cls.extra_server_args,
)
super().setUpClass()
+2
View File
@@ -33,6 +33,7 @@ def post_parallel_generate(
prompts: list[str], prompts: list[str],
max_new_tokens: int, max_new_tokens: int,
timeout: float, timeout: float,
ignore_eos: bool = False,
) -> list[dict]: ) -> list[dict]:
def _send(prompt: str) -> dict: def _send(prompt: str) -> dict:
try: try:
@@ -43,6 +44,7 @@ def post_parallel_generate(
"sampling_params": { "sampling_params": {
"max_new_tokens": max_new_tokens, "max_new_tokens": max_new_tokens,
"temperature": 0.0, "temperature": 0.0,
"ignore_eos": ignore_eos,
}, },
}, },
timeout=timeout, timeout=timeout,
@@ -0,0 +1,24 @@
from __future__ import annotations
import unittest
from sglang.srt.kv_canary.config import CanaryMode
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kv_canary.pp_fixture import CanaryPPFixture
register_cuda_ci(est_time=220, stage="extra-a", runner_config="2-gpu-large")
class TestPPBaselineSwa(CanaryPPFixture):
kv_canary_mode = CanaryMode.LOG
def test_no_violation(self) -> None:
for _ in range(self.workload_n_batches):
self.send_parallel_requests()
self.assert_no_violation(wait_seconds=2.0)
self.maybe_assert_swa_divergence_observed()
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,40 @@
from __future__ import annotations
import unittest
from typing import ClassVar
from sglang.srt.kv_canary.config import CanaryMode
from sglang.srt.kv_canary.perturb.config import TargetGroupKind
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kv_canary.pp_fixture import CanaryPPFixture
register_cuda_ci(est_time=220, stage="extra-a", runner_config="2-gpu-large")
class TestPPPerturbSwaSwa(CanaryPPFixture):
kv_canary_mode = CanaryMode.LOG
target_group: ClassVar[TargetGroupKind] = TargetGroupKind.SWA
extra_server_args = ("--kv-canary-real-data", "partial")
@classmethod
def setUpClass(cls) -> None:
cls.extra_env = {
"SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB": "0.1",
"SGLANG_KV_CANARY_PERTURB_TARGET_GROUP": str(cls.target_group),
"SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0",
}
super().setUpClass()
def test_real_kv_used_perturbation_reports_real_kv_hash_violation(self) -> None:
for _ in range(self.workload_n_batches):
self.send_parallel_requests()
self.assert_per_forward_violation_reported(
fail_reason="verify_real_kv_hash",
target_group=self.target_group,
)
self.maybe_assert_swa_divergence_observed()
if __name__ == "__main__":
unittest.main()
@@ -69,17 +69,58 @@ class TestAssertSwaDivergenceObserved(CustomTestCase):
max_retries=1, max_retries=1,
) )
def test_assert_swa_divergence_observed_uses_latest_line(self) -> None: def test_assert_swa_divergence_observed_uses_peak_out_of_window(self) -> None:
log = _GOOD_LINE + "\n" + _LATER_LINE + "\n" diverged = SwaDivergenceLog(
harness, patcher = self._make_harness(log) forward_ct=120,
verify_full=10000,
verify_swa=4200,
swa_full_idx_divergence=512,
swa_out_of_window_tokens=8192,
).format()
trailing_zero = SwaDivergenceLog(
forward_ct=240,
verify_full=20000,
verify_swa=8400,
swa_full_idx_divergence=0,
swa_out_of_window_tokens=0,
).format()
harness, patcher = self._make_harness(diverged + "\n" + trailing_zero + "\n")
with patcher: with patcher:
harness.assert_swa_divergence_observed( harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1000, min_swa_out_of_window_tokens=1,
min_swa_full_idx_divergence=1,
require_verify_lag=True, require_verify_lag=True,
flush_wait_seconds=0.0, flush_wait_seconds=0.0,
max_retries=1, max_retries=1,
) )
def test_assert_swa_divergence_observed_checks_verify_lag_on_latest_line(
self,
) -> None:
lagging = SwaDivergenceLog(
forward_ct=120,
verify_full=10000,
verify_swa=4200,
swa_full_idx_divergence=512,
swa_out_of_window_tokens=8192,
).format()
no_lag = SwaDivergenceLog(
forward_ct=240,
verify_full=20000,
verify_swa=20000,
swa_full_idx_divergence=1024,
swa_out_of_window_tokens=16384,
).format()
harness, patcher = self._make_harness(lagging + "\n" + no_lag + "\n")
with patcher:
with self.assertRaisesRegex(AssertionError, "verify_swa=20000"):
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1,
require_verify_lag=True,
flush_wait_seconds=0.0,
max_retries=1,
)
def test_assert_swa_divergence_observed_raises_when_below_threshold(self) -> None: def test_assert_swa_divergence_observed_raises_when_below_threshold(self) -> None:
zero_mapping_line = SwaDivergenceLog( zero_mapping_line = SwaDivergenceLog(
forward_ct=100, forward_ct=100,
@@ -410,6 +410,40 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
self.assertEqual(parsed.verify_swa, 3) self.assertEqual(parsed.verify_swa, 3)
class TestSwaDivergenceLogFindAll(CustomTestCase):
def test_find_all_returns_every_sample_in_order(self) -> None:
text = "\n".join(
SwaDivergenceLog(
forward_ct=ct,
verify_full=100 * ct,
verify_swa=10 * ct,
swa_full_idx_divergence=ct,
swa_out_of_window_tokens=0,
).format()
for ct in (20, 40, 60)
)
parsed = SwaDivergenceLog.find_all(text)
self.assertEqual([p.forward_ct for p, _ in parsed], [20, 40, 60])
def test_find_all_peak_survives_trailing_zero_sample(self) -> None:
text = "\n".join(
SwaDivergenceLog(
forward_ct=ct,
verify_full=1,
verify_swa=0,
swa_full_idx_divergence=1,
swa_out_of_window_tokens=oow,
).format()
for ct, oow in ((20, 0), (40, 4080), (60, 0))
)
parsed = SwaDivergenceLog.find_all(text)
self.assertEqual(max(p.swa_out_of_window_tokens for p, _ in parsed), 4080)
self.assertEqual(parsed[-1][0].swa_out_of_window_tokens, 0)
def test_find_all_returns_empty_list_when_no_lines(self) -> None:
self.assertEqual(SwaDivergenceLog.find_all("nothing here\n"), [])
class TestCanaryManagerSwaDivergenceWiring(CanaryManagerTestCase): class TestCanaryManagerSwaDivergenceWiring(CanaryManagerTestCase):
def test_swa_divergence_report_is_none_when_env_disabled(self) -> None: def test_swa_divergence_report_is_none_when_env_disabled(self) -> None:
with envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override( with envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override(