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
test/registered/kv_canary/test_self_e2e_baseline.py
+1
View File
@@ -28,6 +28,7 @@ class _BaselineBase(CanaryE2EBase):
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()
class TestBaselineMha(_BaselineBase):
test/registered/kv_canary/test_self_e2e_perturb_real_kv_unused_cache.py
+1
View File
@@ -51,6 +51,7 @@ class _PerturbRealKvUnusedCacheBase(CanaryE2EBase):
target_group=self.target_group,
flush_wait_seconds=5.0,
)
self.maybe_assert_swa_divergence_observed()
class TestPerturbRealKvUnusedCacheMhaFull(_PerturbRealKvUnusedCacheBase):
test/registered/kv_canary/test_self_e2e_perturb_real_kv_used.py
+1
View File
@@ -39,6 +39,7 @@ class _PerturbRealKvUsedBase(CanaryE2EBase):
fail_reason="verify_real_kv_hash",
target_group=self.target_group,
)
self.maybe_assert_swa_divergence_observed()
class TestPerturbRealKvUsedMhaFull(_PerturbRealKvUsedBase):
test/registered/kv_canary/test_self_e2e_perturb_req_to_token.py
+1
View File
@@ -34,6 +34,7 @@ class _PerturbReqToTokenBase(CanaryE2EBase):
for _ in range(self.workload_n_batches):
self.send_parallel_requests()
self.assert_per_forward_violation_reported(fail_reason="verify_chain_hash")
self.maybe_assert_swa_divergence_observed()
class TestPerturbReqToTokenMha(_PerturbReqToTokenBase):
test/registered/kv_canary/test_self_unit_e2e_base.py
+163
View File
@@ -0,0 +1,163 @@
from __future__ import annotations
import unittest
from unittest.mock import patch
from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceLog
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-b-test-cpu")
_GOOD_LINE: str = SwaDivergenceLog(
forward_ct=120,
verify_full=10000,
verify_swa=4200,
swa_full_idx_divergence=512,
swa_out_of_window_tokens=8192,
).format()
_LATER_LINE: str = SwaDivergenceLog(
forward_ct=240,
verify_full=20000,
verify_swa=8400,
swa_full_idx_divergence=1024,
swa_out_of_window_tokens=16384,
).format()
class _DummyHarness(CanaryE2EBase):
model_mode = "swa"
kv_canary_mode = "log"
@classmethod
def setUpClass(cls) -> None:
return
@classmethod
def tearDownClass(cls) -> None:
return
class TestAssertSwaDivergenceObserved(CustomTestCase):
def _make_harness(
self, log_text_or_sequence
) -> tuple[_DummyHarness, "patch._patch[None]"]:
harness = _DummyHarness()
harness._stderr_buf = None
harness._stdout_buf = None
if isinstance(log_text_or_sequence, list):
patcher = patch.object(
_DummyHarness, "_captured_log_text", side_effect=log_text_or_sequence
)
else:
patcher = patch.object(
_DummyHarness,
"_captured_log_text",
return_value=log_text_or_sequence,
)
return harness, patcher
def test_assert_swa_divergence_observed_passes_when_above_threshold(self) -> None:
harness, patcher = self._make_harness(_LATER_LINE + "\n" + _GOOD_LINE + "\n")
with patcher:
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=100,
require_verify_lag=True,
flush_wait_seconds=0.0,
max_retries=1,
)
def test_assert_swa_divergence_observed_uses_latest_line(self) -> None:
log = _GOOD_LINE + "\n" + _LATER_LINE + "\n"
harness, patcher = self._make_harness(log)
with patcher:
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1000,
require_verify_lag=True,
flush_wait_seconds=0.0,
max_retries=1,
)
def test_assert_swa_divergence_observed_raises_when_below_threshold(self) -> None:
zero_mapping_line = SwaDivergenceLog(
forward_ct=100,
verify_full=5000,
verify_swa=2000,
swa_full_idx_divergence=0,
swa_out_of_window_tokens=8192,
).format()
harness, patcher = self._make_harness(zero_mapping_line + "\n")
with patcher:
with self.assertRaisesRegex(AssertionError, "swa_full_idx_divergence=0"):
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1,
require_verify_lag=False,
flush_wait_seconds=0.0,
max_retries=1,
)
def test_assert_swa_divergence_observed_raises_when_no_verify_lag(self) -> None:
equal_verify_line = SwaDivergenceLog(
forward_ct=100,
verify_full=5000,
verify_swa=5000,
swa_full_idx_divergence=200,
swa_out_of_window_tokens=8192,
).format()
harness, patcher = self._make_harness(equal_verify_line + "\n")
with patcher:
with self.assertRaisesRegex(AssertionError, "verify_swa=5000"):
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_retries_until_stats_emitted(self) -> None:
sequence = ["", "", "", _GOOD_LINE + "\n", _GOOD_LINE + "\n"]
harness, patcher = self._make_harness(sequence)
with patcher:
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1,
require_verify_lag=True,
flush_wait_seconds=0.0,
max_retries=5,
)
def test_assert_swa_divergence_observed_raises_when_no_stats_emitted(self) -> None:
harness, patcher = self._make_harness("nothing here\n")
with patcher:
with self.assertRaisesRegex(AssertionError, "No kv_canary swa_divergence"):
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1,
require_verify_lag=True,
flush_wait_seconds=0.0,
max_retries=2,
)
def test_assert_swa_divergence_observed_catches_zero_swa_full_idx_divergence(
self,
) -> None:
zero_divergence_line = SwaDivergenceLog(
forward_ct=200,
verify_full=10000,
verify_swa=2000,
swa_full_idx_divergence=0,
swa_out_of_window_tokens=8192,
).format()
harness, patcher = self._make_harness(zero_divergence_line + "\n")
with patcher:
with self.assertRaisesRegex(AssertionError, "swa_full_idx_divergence=0"):
harness.assert_swa_divergence_observed(
min_swa_full_idx_divergence=1,
require_verify_lag=True,
flush_wait_seconds=0.0,
max_retries=1,
)
if __name__ == "__main__":
unittest.main()
test/registered/kv_canary/test_self_unit_runner_swa_divergence.py
+431
View File
@@ -0,0 +1,431 @@
from __future__ import annotations
import logging
import unittest
from types import SimpleNamespace
import torch
from sglang.jit_kernel.kv_canary.verify import VerifyPlan
from sglang.srt.environ import envs
from sglang.srt.kv_canary.buffer_group import PoolKind
from sglang.srt.kv_canary.runner import swa_divergence as swa_div_module
from sglang.srt.kv_canary.runner.swa_divergence import (
SwaDivergenceLog,
SwaDivergenceReporter,
compute_swa_full_idx_divergence,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kv_canary.fixtures import make_buffer_group
from sglang.test.kv_canary.runner_test_base import CanaryManagerTestCase, make_manager
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
_DEVICE = torch.device("cuda")
_EMPTY_FORWARD_BATCH = SimpleNamespace(
req_pool_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
seq_lens=torch.empty(0, dtype=torch.int64, device=_DEVICE),
)
def _make_verify_plan(value: int) -> VerifyPlan:
plan = VerifyPlan.allocate(verify_capacity=4, device=_DEVICE)
plan.verify_num_valid.copy_(torch.tensor([value], dtype=torch.int32))
return plan
def _make_allocator_stub(mapping: torch.Tensor) -> SimpleNamespace:
return SimpleNamespace(full_to_swa_index_mapping=mapping)
def _make_req_to_token_pool_stub(req_to_token: torch.Tensor) -> SimpleNamespace:
return SimpleNamespace(req_to_token=req_to_token)
def _make_identity_mapping(size: int) -> torch.Tensor:
return torch.arange(size, dtype=torch.int64, device=_DEVICE)
def _make_identity_req_to_token(num_reqs: int, max_seq_len: int) -> torch.Tensor:
base = torch.arange(num_reqs * max_seq_len, dtype=torch.int64, device=_DEVICE)
return base.view(num_reqs, max_seq_len)
def _make_forward_batch(
*, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
) -> SimpleNamespace:
return SimpleNamespace(req_pool_indices=req_pool_indices, seq_lens=seq_lens)
def _parse_swa_divergence_line(line: str) -> SwaDivergenceLog:
parsed = SwaDivergenceLog.parse(line)
if parsed is None:
raise AssertionError(f"line does not match swa_divergence format: {line!r}")
return parsed
def _run_compute(
*,
swa_allocator: SimpleNamespace,
req_to_token_pool: SimpleNamespace,
forward_batch: SimpleNamespace,
) -> int:
count = compute_swa_full_idx_divergence(
swa_allocator=swa_allocator,
req_to_token_pool=req_to_token_pool,
maybe_inaccurate_forward_batch=forward_batch,
)
return int(count.item())
class TestSwaDivergenceReporter(CustomTestCase):
def test_swa_divergence_log_emitted(self) -> None:
d2h_stream = torch.cuda.Stream(device=_DEVICE)
stats = SwaDivergenceReporter(
device=_DEVICE,
d2h_stream=d2h_stream,
interval=10,
swa_allocator=None,
req_to_token_pool=None,
)
# First 3 forwards stay below the interval trigger (1, 2, 3 % 10 != 0) so
# step() just bumps forward_ct and stages nothing.
for forward_idx in range(3):
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(10),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(3),
)
stats.step(
outer_step_counter=forward_idx + 1,
maybe_inaccurate_forward_batch=_EMPTY_FORWARD_BATCH,
)
# 4th forward lands on outer_step_counter=10 = interval, so compute_on_device
# snapshots {forward_ct:4, verify_full:40, verify_swa:12} into the dict and
# the staged future hangs onto it. forward_ct is now 4.
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(10),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(3),
)
stats.step(
outer_step_counter=10, maybe_inaccurate_forward_batch=_EMPTY_FORWARD_BATCH
)
# 5th step drains the previous stage and emits the log; forward_ct is now 5
# but the staged dict still carries the snapshot forward_ct=4 from step 4.
with self.assertLogs(
swa_div_module.logger.name, level=logging.INFO
) as captured:
stats.step(
outer_step_counter=11,
maybe_inaccurate_forward_batch=_EMPTY_FORWARD_BATCH,
)
lines = [
line for line in captured.output if SwaDivergenceLog.parse(line) is not None
]
self.assertEqual(len(lines), 1, lines)
fields = _parse_swa_divergence_line(lines[0])
self.assertEqual(fields.forward_ct, 4)
self.assertEqual(fields.verify_full, 40)
self.assertEqual(fields.verify_swa, 12)
self.assertEqual(fields.swa_full_idx_divergence, 0)
def test_swa_divergence_counts_monotonic_increasing(self) -> None:
d2h_stream = torch.cuda.Stream(device=_DEVICE)
stats = SwaDivergenceReporter(
device=_DEVICE,
d2h_stream=d2h_stream,
interval=10,
swa_allocator=None,
req_to_token_pool=None,
)
snapshots: list[SwaDivergenceLog] = []
def _take_snapshot(stage_step: int, drain_step: int) -> None:
# Stage the dict at the interval-aligned step (no log emitted yet,
# DelayedDeviceHostHandler still has nothing to drain), then call
# step() again at the next counter to drain and emit the log.
stats.step(
outer_step_counter=stage_step,
maybe_inaccurate_forward_batch=_EMPTY_FORWARD_BATCH,
)
with self.assertLogs(
swa_div_module.logger.name, level=logging.INFO
) as captured:
stats.step(
outer_step_counter=drain_step,
maybe_inaccurate_forward_batch=_EMPTY_FORWARD_BATCH,
)
matching = [
line
for line in captured.output
if SwaDivergenceLog.parse(line) is not None
]
self.assertTrue(matching, captured.output)
snapshots.append(_parse_swa_divergence_line(matching[-1]))
for batch in range(3):
for _ in range(5):
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(7),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(2),
)
stage_step = 10 + 20 * batch
_take_snapshot(stage_step=stage_step, drain_step=stage_step + 1)
for idx in range(1, len(snapshots)):
self.assertGreaterEqual(
snapshots[idx].verify_full, snapshots[idx - 1].verify_full
)
self.assertGreaterEqual(
snapshots[idx].verify_swa, snapshots[idx - 1].verify_swa
)
class TestSwaFullIdxDivergenceCompute(CustomTestCase):
def test_compute_returns_zero_when_empty_batch(self) -> None:
mapping = _make_identity_mapping(size=64)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
forward_batch = _make_forward_batch(
req_pool_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
seq_lens=torch.empty(0, dtype=torch.int64, device=_DEVICE),
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=forward_batch,
),
0,
)
def test_compute_returns_zero_when_all_identity(self) -> None:
mapping = _make_identity_mapping(size=64)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=_DEVICE),
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=forward_batch,
),
0,
)
def test_compute_counts_swa_full_idx_divergence_in_live_range(self) -> None:
mapping = _make_identity_mapping(size=64)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
mapping[0] = 50
mapping[1] = 51
mapping[17] = 60
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=_DEVICE),
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=forward_batch,
),
3,
)
def test_compute_ignores_swa_mapping_zero(self) -> None:
# SWATokenToKVPoolAllocator writes 0 into full_to_swa_index_mapping for
# FULL pool slots beyond the sliding window. Those entries are expected,
# not real divergence, so the count must skip them.
mapping = _make_identity_mapping(size=64)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
mapping[3] = 0
mapping[5] = 0
mapping[7] = 42
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=forward_batch,
),
1,
)
def test_compute_ignores_writes_outside_seq_lens(self) -> None:
mapping = _make_identity_mapping(size=128)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=32)
mapping[20] = 99
mapping[28] = 77
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([10], dtype=torch.int64, device=_DEVICE),
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=forward_batch,
),
0,
)
def test_compute_reflects_current_forward_batch(self) -> None:
mapping = _make_identity_mapping(size=64)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
mapping[0] = 41
mapping[1] = 42
mapping[32] = 99
mapping[33] = 100
fb_req0 = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
)
fb_req2 = _make_forward_batch(
req_pool_indices=torch.tensor([2], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=fb_req0,
),
2,
)
self.assertEqual(
_run_compute(
swa_allocator=_make_allocator_stub(mapping),
req_to_token_pool=_make_req_to_token_pool_stub(req_to_token),
forward_batch=fb_req2,
),
2,
)
class TestSwaDivergenceReporterWithCompute(CustomTestCase):
def test_swa_divergence_report_emits_swa_full_idx_divergence_from_compute(
self,
) -> None:
mapping = _make_identity_mapping(size=64)
req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
mapping[0] = 50
mapping[1] = 51
mapping[2] = 52
forward_batch = _make_forward_batch(
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
)
swa_allocator = _make_allocator_stub(mapping)
req_to_token_pool = _make_req_to_token_pool_stub(req_to_token)
d2h_stream = torch.cuda.Stream(device=_DEVICE)
stats = SwaDivergenceReporter(
device=_DEVICE,
d2h_stream=d2h_stream,
interval=10,
swa_allocator=swa_allocator,
req_to_token_pool=req_to_token_pool,
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(11),
)
stats.observe_after_invoke_plan(
group=make_buffer_group(
device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
),
verify_plan=_make_verify_plan(3),
)
# Stage at the interval-aligned step, then drain on the next step so the
# DelayedDeviceHostHandler has a pending future to postprocess.
stats.step(outer_step_counter=10, maybe_inaccurate_forward_batch=forward_batch)
with self.assertLogs(
swa_div_module.logger.name, level=logging.INFO
) as captured:
stats.step(
outer_step_counter=11, maybe_inaccurate_forward_batch=forward_batch
)
matching = [
line for line in captured.output if SwaDivergenceLog.parse(line) is not None
]
self.assertEqual(len(matching), 1, matching)
parsed = SwaDivergenceLog.parse(matching[0])
assert parsed is not None
self.assertEqual(parsed.swa_full_idx_divergence, 3)
self.assertEqual(parsed.verify_full, 11)
self.assertEqual(parsed.verify_swa, 3)
class TestCanaryManagerSwaDivergenceWiring(CanaryManagerTestCase):
def test_swa_divergence_report_is_none_when_env_disabled(self) -> None:
with envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override(
0
), envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.override("full"):
manager = make_manager(device=self.device)
self.assertIsNone(manager._swa_divergence_report)
def test_swa_divergence_report_present_when_env_enabled(self) -> None:
with envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override(
20
), envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.override("full"):
manager = make_manager(device=self.device)
self.assertIsNotNone(manager._swa_divergence_report)
self.assertIsInstance(manager._swa_divergence_report, SwaDivergenceReporter)
if __name__ == "__main__":
unittest.main()