Add the KV-canary perturb modes and PD-disaggregation e2e tests (#26819)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.pd_fixture import CanaryPDFixture
|
||||
|
||||
register_cuda_ci(est_time=180, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
|
||||
class TestPDBaselineMha(CanaryPDFixture):
|
||||
model_mode = "mha"
|
||||
|
||||
def test_clean_pd_run_produces_no_canary_violation_on_either_side(self) -> None:
|
||||
self.send_parallel_short_requests(n=4)
|
||||
self.assert_no_violation(side="prefill")
|
||||
self.assert_no_violation(side="decode")
|
||||
|
||||
|
||||
class TestPDBaselineSwa(CanaryPDFixture):
|
||||
model_mode = "swa"
|
||||
|
||||
def test_clean_pd_run_produces_no_canary_violation_on_either_side(self) -> None:
|
||||
self.send_parallel_short_requests(n=4)
|
||||
self.assert_no_violation(side="prefill")
|
||||
self.assert_no_violation(side="decode")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,74 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from typing import ClassVar
|
||||
|
||||
from sglang.srt.kv_canary.perturb.config import TargetGroupKind
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.pd_fixture import CanaryPDFixture
|
||||
|
||||
register_cuda_ci(est_time=180, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
|
||||
class _PDPerturbBase(CanaryPDFixture):
|
||||
target_group: ClassVar[TargetGroupKind]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
if cls is _PDPerturbBase:
|
||||
raise unittest.SkipTest(
|
||||
"abstract base; concrete subclasses set model_mode + target_group"
|
||||
)
|
||||
cls.extra_prefill_env = {
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_POST_FORWARD_PROB": "1.0",
|
||||
"SGLANG_KV_CANARY_PERTURB_TARGET_GROUP": str(cls.target_group),
|
||||
"SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0",
|
||||
}
|
||||
cls.extra_decode_env = {
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_POST_FORWARD_PROB": "0",
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB": "0",
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB": "0",
|
||||
"SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB": "0",
|
||||
}
|
||||
super().setUpClass()
|
||||
|
||||
def test_p_side_perturb_surfaces_real_kv_hash_violation_on_decode_side(
|
||||
self,
|
||||
) -> None:
|
||||
# send_parallel_short_requests defaults to max_new_tokens=100 so D-side runs
|
||||
# decode forwards that exercise canary verify on the transferred prefix.
|
||||
self.send_parallel_short_requests(n=4)
|
||||
# D-side: first decode forward re-verifies the transferred prefix slots,
|
||||
# so the flip MUST surface as real_kv_hash violation.
|
||||
self.assert_per_forward_violation_reported(
|
||||
fail_reason="verify_real_kv_hash",
|
||||
target_group=self.target_group,
|
||||
side="decode",
|
||||
flush_wait_seconds=4.0,
|
||||
)
|
||||
# P-side: flip happens post-TAIL of the prefill forward, and PD prefill
|
||||
# does not run another forward on P that would verify the perturbed slot,
|
||||
# so P MUST stay silent (no false-positive violations) for this perturb
|
||||
# point. If a future canary feature adds post-prefill verify on P, this
|
||||
# assert will start failing and should be upgraded to assert the
|
||||
# violation on P-side too.
|
||||
self.assert_no_violation(side="prefill", wait_seconds=0.5)
|
||||
|
||||
|
||||
class TestPDPerturbMhaFull(_PDPerturbBase):
|
||||
model_mode = "mha"
|
||||
target_group = TargetGroupKind.FULL
|
||||
|
||||
|
||||
class TestPDPerturbSwaFull(_PDPerturbBase):
|
||||
model_mode = "swa"
|
||||
target_group = TargetGroupKind.FULL
|
||||
|
||||
|
||||
class TestPDPerturbSwaSwa(_PDPerturbBase):
|
||||
model_mode = "swa"
|
||||
target_group = TargetGroupKind.SWA
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
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.e2e_base import CanaryE2EBase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestPerturbRaiseMha(CanaryE2EBase):
|
||||
model_mode = "mha"
|
||||
kv_canary_mode = CanaryMode.RAISE
|
||||
extra_server_args = ("--kv-canary-real-data", "partial", "--skip-server-warmup")
|
||||
extra_env = {
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB": "0.1",
|
||||
"SGLANG_KV_CANARY_PERTURB_TARGET_GROUP": "full",
|
||||
"SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0",
|
||||
}
|
||||
|
||||
def test_real_kv_used_perturbation_raises_in_raise_mode(self) -> None:
|
||||
"""Verify raise mode surfaces real KV perturbation as a logged violation."""
|
||||
try:
|
||||
self.send_parallel_requests(
|
||||
n=4,
|
||||
assert_all_success=False,
|
||||
timeout=30.0,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
self.assert_per_forward_violation_reported(
|
||||
fail_reason="verify_real_kv_hash",
|
||||
target_group=TargetGroupKind.FULL,
|
||||
flush_wait_seconds=3.0,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,80 @@
|
||||
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.consts import SWA_POOL_SERVER_ARGS
|
||||
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _PerturbRealKvUnusedCacheBase(CanaryE2EBase):
|
||||
kv_canary_mode = CanaryMode.LOG
|
||||
extra_server_args = (
|
||||
"--kv-canary-real-data",
|
||||
"partial",
|
||||
"--kv-canary-sweep-interval",
|
||||
"4",
|
||||
)
|
||||
use_unique_prompts = True
|
||||
|
||||
target_group: ClassVar[TargetGroupKind]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
if cls is _PerturbRealKvUnusedCacheBase:
|
||||
raise unittest.SkipTest(
|
||||
"abstract base; concrete subclasses set model_mode + target_group"
|
||||
)
|
||||
cls.extra_env = {
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_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_unused_cache_perturbation_reports_sweep_real_kv_hash_violation(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify cached unused KV perturbation is caught by sweep verification."""
|
||||
# Step 1: first batch builds radix entries that will become orphans once finished.
|
||||
self.send_parallel_requests(n=8)
|
||||
# Step 2: second batch drives more forward passes so the sweep cadence fires
|
||||
# while the orphan slots are still cached.
|
||||
self.send_parallel_requests(n=8)
|
||||
self.assert_sweep_violation_reported(
|
||||
fail_reason="verify_real_kv_hash",
|
||||
target_group=self.target_group,
|
||||
flush_wait_seconds=5.0,
|
||||
)
|
||||
|
||||
|
||||
class TestPerturbRealKvUnusedCacheMhaFull(_PerturbRealKvUnusedCacheBase):
|
||||
model_mode = "mha"
|
||||
target_group = TargetGroupKind.FULL
|
||||
|
||||
|
||||
class TestPerturbRealKvUnusedCacheSwaFull(_PerturbRealKvUnusedCacheBase):
|
||||
model_mode = "swa"
|
||||
target_group = TargetGroupKind.FULL
|
||||
extra_server_args = (
|
||||
*_PerturbRealKvUnusedCacheBase.extra_server_args,
|
||||
*SWA_POOL_SERVER_ARGS,
|
||||
)
|
||||
|
||||
|
||||
class TestPerturbRealKvUnusedCacheSwaSwa(_PerturbRealKvUnusedCacheBase):
|
||||
model_mode = "swa"
|
||||
target_group = TargetGroupKind.SWA
|
||||
extra_server_args = (
|
||||
*_PerturbRealKvUnusedCacheBase.extra_server_args,
|
||||
*SWA_POOL_SERVER_ARGS,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,68 @@
|
||||
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.consts import SWA_POOL_SERVER_ARGS
|
||||
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _PerturbRealKvUsedBase(CanaryE2EBase):
|
||||
kv_canary_mode = CanaryMode.LOG
|
||||
extra_server_args = ("--kv-canary-real-data", "partial")
|
||||
|
||||
target_group: ClassVar[TargetGroupKind]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
if cls is _PerturbRealKvUsedBase:
|
||||
raise unittest.SkipTest(
|
||||
"abstract base; concrete subclasses set model_mode + target_group"
|
||||
)
|
||||
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:
|
||||
"""Verify active real KV perturbation reports a real KV hash violation."""
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class TestPerturbRealKvUsedMhaFull(_PerturbRealKvUsedBase):
|
||||
model_mode = "mha"
|
||||
target_group = TargetGroupKind.FULL
|
||||
|
||||
|
||||
class TestPerturbRealKvUsedSwaFull(_PerturbRealKvUsedBase):
|
||||
model_mode = "swa"
|
||||
target_group = TargetGroupKind.FULL
|
||||
extra_server_args = (
|
||||
*_PerturbRealKvUsedBase.extra_server_args,
|
||||
*SWA_POOL_SERVER_ARGS,
|
||||
)
|
||||
|
||||
|
||||
class TestPerturbRealKvUsedSwaSwa(_PerturbRealKvUsedBase):
|
||||
model_mode = "swa"
|
||||
target_group = TargetGroupKind.SWA
|
||||
extra_server_args = (
|
||||
*_PerturbRealKvUsedBase.extra_server_args,
|
||||
*SWA_POOL_SERVER_ARGS,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,49 @@
|
||||
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.consts import SWA_POOL_SERVER_ARGS
|
||||
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _PerturbReqToTokenBase(CanaryE2EBase):
|
||||
kv_canary_mode = CanaryMode.LOG
|
||||
extra_env = {
|
||||
"SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB": "0.1",
|
||||
"SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0",
|
||||
# req_to_token perturbation deliberately corrupts the slot mapping
|
||||
# by design, which the scheduler's on-idle invariant checker reports
|
||||
# as a pool memory leak (perturbed slot is freed, original slot
|
||||
# still looks busy). That's expected for this test; disable strict
|
||||
# mode so the leak warning doesn't crash the scheduler.
|
||||
"SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "0",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
if cls is _PerturbReqToTokenBase:
|
||||
raise unittest.SkipTest("abstract base; concrete subclasses set model_mode")
|
||||
super().setUpClass()
|
||||
|
||||
def test_req_to_token_perturbation_reports_chain_hash_violation(self) -> None:
|
||||
"""Verify req_to_token perturbation reports a chain hash violation."""
|
||||
for _ in range(self.workload_n_batches):
|
||||
self.send_parallel_requests()
|
||||
self.assert_per_forward_violation_reported(fail_reason="verify_chain_hash")
|
||||
|
||||
|
||||
class TestPerturbReqToTokenMha(_PerturbReqToTokenBase):
|
||||
model_mode = "mha"
|
||||
|
||||
|
||||
class TestPerturbReqToTokenSwa(_PerturbReqToTokenBase):
|
||||
model_mode = "swa"
|
||||
extra_server_args = SWA_POOL_SERVER_ARGS
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -2,28 +2,44 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from typing import cast
|
||||
from typing import TYPE_CHECKING, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify import RealKvSource
|
||||
from sglang.srt.kv_canary.buffer_group import PoolKind
|
||||
from sglang.srt.kv_canary.perturb import (
|
||||
real_kv_post_forward,
|
||||
)
|
||||
from sglang.srt.kv_canary.perturb import (
|
||||
real_kv_unused_cache as real_kv_unused_cache_module,
|
||||
)
|
||||
from sglang.srt.kv_canary.perturb.config import (
|
||||
PerturbConfig,
|
||||
TargetGroupKind,
|
||||
_parse_target_group_kind,
|
||||
)
|
||||
from sglang.srt.kv_canary.perturb.manager import PerturbManager
|
||||
from sglang.srt.kv_canary.perturb.slot_picker import collect_active_slots
|
||||
from sglang.srt.kv_canary.perturb.utils import (
|
||||
WarmupGate,
|
||||
flip_first_byte_in_source,
|
||||
pick_target_group,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import (
|
||||
DEFAULT_DEVICE,
|
||||
make_buffer_group,
|
||||
make_forward_batch,
|
||||
make_radix_cache,
|
||||
make_req_to_token_pool,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
|
||||
|
||||
register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
@@ -56,11 +72,20 @@ class TestParseTargetGroupKind(CustomTestCase):
|
||||
):
|
||||
_parse_target_group_kind(raw)
|
||||
|
||||
def test_from_env_allows_missing_target(
|
||||
def test_from_env_allows_missing_target_when_real_kv_perturb_is_disabled(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify normal canary startup does not require a perturb target group."""
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"SGLANG_KV_CANARY_PERTURB_REQ_TO_TOKEN_PROB": "0",
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB": "0",
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB": "0",
|
||||
"SGLANG_KV_CANARY_PERTURB_REAL_KV_POST_FORWARD_PROB": "0",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
os.environ.pop("SGLANG_KV_CANARY_PERTURB_TARGET_GROUP", None)
|
||||
config = PerturbConfig.from_env()
|
||||
|
||||
@@ -111,7 +136,204 @@ class TestPickTargetGroup(CustomTestCase):
|
||||
self.assertIsNone(group)
|
||||
|
||||
|
||||
class TestPerturbWarmupAndUtils(CustomTestCase):
|
||||
class TestPerturbManager(CustomTestCase):
|
||||
def test_perturb_manager_perturb_post_forward_dispatches_real_kv_post_forward(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify perturb_post_forward() routes only to the post_forward dispatch."""
|
||||
device = DEFAULT_DEVICE
|
||||
manager = PerturbManager(
|
||||
config=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=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
),
|
||||
req_to_token_pool=make_req_to_token_pool(device, max_reqs=4, max_seq_len=8),
|
||||
buffer_groups=(),
|
||||
outer_step_counter_getter=lambda: 10,
|
||||
)
|
||||
forward_batch = make_forward_batch(device, bs=1, seq_lens_list=(1,))
|
||||
calls: list[str] = []
|
||||
|
||||
with patch.object(
|
||||
manager,
|
||||
"perturb_real_kv_post_forward",
|
||||
lambda batch: calls.append("real_kv_post_forward"),
|
||||
), patch.object(
|
||||
manager,
|
||||
"perturb_req_to_token",
|
||||
lambda batch: calls.append("req_to_token"),
|
||||
), patch.object(
|
||||
manager,
|
||||
"perturb_real_kv_used",
|
||||
lambda batch: calls.append("real_kv_used"),
|
||||
), patch.object(
|
||||
manager,
|
||||
"perturb_real_kv_unused_cache",
|
||||
lambda batch: calls.append("real_kv_unused_cache"),
|
||||
):
|
||||
manager.perturb_post_forward(maybe_inaccurate_forward_batch=forward_batch)
|
||||
|
||||
self.assertEqual(calls, ["real_kv_post_forward"])
|
||||
|
||||
|
||||
class TestRealKvPostForwardPerturb(CustomTestCase):
|
||||
def test_real_kv_post_forward_flips_a_byte_in_out_cache_loc_slot(self) -> None:
|
||||
"""Verify post-forward perturbation flips one real-KV byte and leaves canary buffers untouched."""
|
||||
device = DEFAULT_DEVICE
|
||||
group = make_buffer_group(kind=PoolKind.FULL, has_real_kv=True)
|
||||
source = group.real_kv_sources_k[0]
|
||||
config = 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=1.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
)
|
||||
warmup_gate = WarmupGate(config=config, outer_step_counter_getter=lambda: 10)
|
||||
|
||||
forward_batch = make_forward_batch(device, bs=1, seq_lens_list=(1,))
|
||||
forward_batch.out_cache_loc = torch.tensor(
|
||||
[2], dtype=torch.int32, device=device
|
||||
)
|
||||
forward_batch.num_token_non_padded_cpu = 1
|
||||
|
||||
head_snapshot = group.k_head.clone()
|
||||
v_head_snapshot = group.v_head.clone()
|
||||
k_tail_snapshot = group.k_tail.clone()
|
||||
v_tail_snapshot = group.v_tail.clone()
|
||||
source_snapshot = source.tensor.clone()
|
||||
|
||||
with patch.object(torch, "rand", return_value=torch.tensor(0.0)):
|
||||
real_kv_post_forward.run(
|
||||
maybe_inaccurate_forward_batch=forward_batch,
|
||||
config=config,
|
||||
buffer_groups=(group,),
|
||||
warmup_gate=warmup_gate,
|
||||
)
|
||||
|
||||
diff = source.tensor != source_snapshot
|
||||
self.assertEqual(int(diff.sum().item()), 1)
|
||||
self.assertTrue(bool(diff[2, 0].item()))
|
||||
self.assertEqual(int(source.tensor[2, 0].item()), 0 ^ 0xFF)
|
||||
self.assertTrue(torch.equal(group.k_head, head_snapshot))
|
||||
self.assertTrue(torch.equal(group.v_head, v_head_snapshot))
|
||||
self.assertTrue(torch.equal(group.k_tail, k_tail_snapshot))
|
||||
self.assertTrue(torch.equal(group.v_tail, v_tail_snapshot))
|
||||
|
||||
|
||||
class TestReqToTokenPerturb(CustomTestCase):
|
||||
def test_req_to_token_perturb_uses_live_slot_as_replacement(self) -> None:
|
||||
"""Verify req_to_token perturbation replaces a slot with another live slot."""
|
||||
device = DEFAULT_DEVICE
|
||||
pool = make_req_to_token_pool(device, max_reqs=4, max_seq_len=8)
|
||||
pool.req_to_token[1, :3] = torch.tensor(
|
||||
[11, 22, 33], dtype=torch.int32, device=device
|
||||
)
|
||||
pool.req_to_token[2, :3] = torch.tensor(
|
||||
[44, 55, 66], dtype=torch.int32, device=device
|
||||
)
|
||||
manager = PerturbManager(
|
||||
config=PerturbConfig(
|
||||
req_to_token_prob=1.0,
|
||||
real_kv_used_prob=0.0,
|
||||
real_kv_unused_cache_prob=0.0,
|
||||
real_kv_post_forward_prob=0.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
),
|
||||
req_to_token_pool=pool,
|
||||
buffer_groups=(),
|
||||
outer_step_counter_getter=lambda: 10,
|
||||
)
|
||||
forward_batch = make_forward_batch(device, bs=2, seq_lens_list=(3, 3))
|
||||
forward_batch.out_cache_loc = torch.tensor(
|
||||
[11], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
snapshot = pool.req_to_token.clone()
|
||||
with patch.object(torch, "rand", return_value=torch.tensor(0.0)):
|
||||
manager.perturb_req_to_token(forward_batch)
|
||||
|
||||
diff = pool.req_to_token != snapshot
|
||||
self.assertEqual(int(diff.sum().item()), 1)
|
||||
rows, cols = torch.nonzero(diff, as_tuple=True)
|
||||
row, col = int(rows[0].item()), int(cols[0].item())
|
||||
original = int(snapshot[row, col].item())
|
||||
replacement = int(pool.req_to_token[row, col].item())
|
||||
live_slots = {11, 22, 33, 44, 55, 66}
|
||||
self.assertIn(original, live_slots)
|
||||
self.assertIn(replacement, live_slots)
|
||||
self.assertNotEqual(replacement, original)
|
||||
self.assertFalse(bool(diff[1, 0].item()))
|
||||
|
||||
def test_collect_active_slots_ignores_padded_out_cache_loc(self) -> None:
|
||||
"""Verify out_cache_loc padding does not exclude a live slot."""
|
||||
device = DEFAULT_DEVICE
|
||||
pool = make_req_to_token_pool(device, max_reqs=4, max_seq_len=8)
|
||||
pool.req_to_token[1, :2] = torch.tensor(
|
||||
[0, 7], dtype=torch.int32, device=device
|
||||
)
|
||||
forward_batch = make_forward_batch(device, bs=1, seq_lens_list=(2,))
|
||||
forward_batch.out_cache_loc = torch.tensor(
|
||||
[7, 0, 0], dtype=torch.int32, device=device
|
||||
)
|
||||
forward_batch.num_token_non_padded_cpu = 1
|
||||
|
||||
targets = collect_active_slots(
|
||||
maybe_inaccurate_forward_batch=forward_batch,
|
||||
req_to_token_pool=pool,
|
||||
)
|
||||
|
||||
self.assertEqual([target.value for target in targets], [0])
|
||||
|
||||
|
||||
class TestRealKvUsedPerturb(CustomTestCase):
|
||||
def test_real_kv_used_flips_first_real_kv_byte_for_active_full_slot(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify real_kv_used flips only the first real KV byte for an active FULL slot."""
|
||||
device = DEFAULT_DEVICE
|
||||
pool = make_req_to_token_pool(device, max_reqs=4, max_seq_len=8)
|
||||
pool.req_to_token.fill_(-1)
|
||||
pool.req_to_token[1, 0] = 2
|
||||
group = make_buffer_group(kind=PoolKind.FULL, has_real_kv=True)
|
||||
source = group.real_kv_sources_k[0]
|
||||
source.tensor.copy_(
|
||||
torch.arange(source.tensor.numel(), dtype=torch.uint8).view_as(
|
||||
source.tensor
|
||||
)
|
||||
)
|
||||
manager = PerturbManager(
|
||||
config=PerturbConfig(
|
||||
req_to_token_prob=0.0,
|
||||
real_kv_used_prob=1.0,
|
||||
real_kv_unused_cache_prob=0.0,
|
||||
real_kv_post_forward_prob=0.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
),
|
||||
req_to_token_pool=pool,
|
||||
buffer_groups=(group,),
|
||||
outer_step_counter_getter=lambda: 10,
|
||||
)
|
||||
forward_batch = make_forward_batch(device, bs=1, seq_lens_list=(1,))
|
||||
forward_batch.out_cache_loc = torch.tensor(
|
||||
[99], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
snapshot = source.tensor.clone()
|
||||
with patch.object(torch, "rand", return_value=torch.tensor(0.0)):
|
||||
manager.perturb_real_kv_used(forward_batch)
|
||||
|
||||
expected = snapshot.clone()
|
||||
expected[2, 0] = int(snapshot[2, 0].item()) ^ 0xFF
|
||||
self.assertTrue(torch.equal(source.tensor, expected))
|
||||
|
||||
def test_flip_first_byte_in_source_maps_swa_logical_slot_through_lut(
|
||||
self,
|
||||
) -> None:
|
||||
@@ -137,6 +359,161 @@ class TestPerturbWarmupAndUtils(CustomTestCase):
|
||||
expected[1, 16] = int(snapshot[1, 16].item()) ^ 0xFF
|
||||
self.assertTrue(torch.equal(source.tensor, expected))
|
||||
|
||||
def test_warmup_gate_prevents_perturbation_when_probabilities_are_one(self) -> None:
|
||||
"""Verify warmup prevents all perturbations even when every probability is one."""
|
||||
device = DEFAULT_DEVICE
|
||||
pool = make_req_to_token_pool(device, max_reqs=4, max_seq_len=8)
|
||||
pool.req_to_token.fill_(-1)
|
||||
pool.req_to_token[1, 0] = 2
|
||||
group = make_buffer_group(kind=PoolKind.FULL, has_real_kv=True)
|
||||
source = group.real_kv_sources_k[0]
|
||||
source.tensor.copy_(
|
||||
torch.arange(source.tensor.numel(), dtype=torch.uint8).view_as(
|
||||
source.tensor
|
||||
)
|
||||
)
|
||||
manager = PerturbManager(
|
||||
config=PerturbConfig(
|
||||
req_to_token_prob=1.0,
|
||||
real_kv_used_prob=1.0,
|
||||
real_kv_unused_cache_prob=1.0,
|
||||
real_kv_post_forward_prob=0.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=20,
|
||||
),
|
||||
req_to_token_pool=pool,
|
||||
buffer_groups=(group,),
|
||||
outer_step_counter_getter=lambda: 10,
|
||||
)
|
||||
manager.attach_radix_cache(cast("BasePrefixCache", object()))
|
||||
forward_batch = make_forward_batch(device, bs=1, seq_lens_list=(1,))
|
||||
forward_batch.out_cache_loc = torch.tensor(
|
||||
[99], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
pool_snapshot = pool.req_to_token.clone()
|
||||
source_snapshot = source.tensor.clone()
|
||||
with patch.object(torch, "rand", return_value=torch.tensor(0.0)), patch.object(
|
||||
real_kv_unused_cache_module,
|
||||
"_pick_sweep_slot_for_group",
|
||||
return_value=3,
|
||||
):
|
||||
manager.perturb(maybe_inaccurate_forward_batch=forward_batch)
|
||||
|
||||
self.assertTrue(torch.equal(pool.req_to_token, pool_snapshot))
|
||||
self.assertTrue(torch.equal(source.tensor, source_snapshot))
|
||||
|
||||
|
||||
class TestRealKvUnusedCachePerturb(CustomTestCase):
|
||||
def test_real_kv_unused_cache_flips_first_real_kv_byte_for_orphan_slot(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify real_kv_unused_cache flips only the first real KV byte for an orphan slot."""
|
||||
device = DEFAULT_DEVICE
|
||||
pool = make_req_to_token_pool(device, max_reqs=4, max_seq_len=8)
|
||||
group = make_buffer_group(kind=PoolKind.FULL, has_real_kv=True)
|
||||
source = group.real_kv_sources_k[0]
|
||||
source.tensor.copy_(
|
||||
torch.arange(source.tensor.numel(), dtype=torch.uint8).view_as(
|
||||
source.tensor
|
||||
)
|
||||
)
|
||||
manager = PerturbManager(
|
||||
config=PerturbConfig(
|
||||
req_to_token_prob=0.0,
|
||||
real_kv_used_prob=0.0,
|
||||
real_kv_unused_cache_prob=1.0,
|
||||
real_kv_post_forward_prob=0.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
),
|
||||
req_to_token_pool=pool,
|
||||
buffer_groups=(group,),
|
||||
outer_step_counter_getter=lambda: 10,
|
||||
sweep_interval=1,
|
||||
)
|
||||
manager.attach_radix_cache(make_radix_cache([[], [3]], device=device))
|
||||
|
||||
snapshot = source.tensor.clone()
|
||||
with patch.object(torch, "rand", return_value=torch.tensor(0.0)), patch.object(
|
||||
torch,
|
||||
"randint",
|
||||
return_value=torch.tensor(0),
|
||||
):
|
||||
manager.perturb_real_kv_unused_cache(None)
|
||||
|
||||
expected = snapshot.clone()
|
||||
expected[3, 0] = int(snapshot[3, 0].item()) ^ 0xFF
|
||||
self.assertTrue(torch.equal(source.tensor, expected))
|
||||
|
||||
def test_pick_sweep_slot_for_group_skips_locked_radix_nodes(self) -> None:
|
||||
"""Verify unused-cache perturbation chooses only unlocked radix-cache slots."""
|
||||
device = DEFAULT_DEVICE
|
||||
group = make_buffer_group(kind=PoolKind.FULL, has_real_kv=True)
|
||||
cache = make_radix_cache([[], [1, 2], [3]], device=device)
|
||||
locked_node = next(iter(cache.root_node.children.values()))
|
||||
locked_node.lock_ref = 1
|
||||
|
||||
with patch.object(torch, "randint", return_value=torch.tensor(0)):
|
||||
slot = real_kv_unused_cache_module._pick_sweep_slot_for_group(
|
||||
radix_cache=cache,
|
||||
group=group,
|
||||
swa_window_size=0,
|
||||
)
|
||||
|
||||
self.assertEqual(slot, 3)
|
||||
|
||||
def test_pick_sweep_slot_for_group_translates_swa_slots(self) -> None:
|
||||
"""Verify unused-cache SWA perturbation translates full slots to physical SWA slots."""
|
||||
device = DEFAULT_DEVICE
|
||||
lut = torch.tensor([-1, 2], dtype=torch.int64, device=device)
|
||||
group = make_buffer_group(
|
||||
kind=PoolKind.SWA, has_real_kv=True, swa_index_lut=lut
|
||||
)
|
||||
cache = make_radix_cache([[], [1]], device=device)
|
||||
|
||||
with patch.object(torch, "randint", return_value=torch.tensor(0)):
|
||||
slot = real_kv_unused_cache_module._pick_sweep_slot_for_group(
|
||||
radix_cache=cache,
|
||||
group=group,
|
||||
swa_window_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(slot, 2)
|
||||
|
||||
def test_real_kv_unused_cache_skips_without_radix_cache_when_forward_batch_is_none(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify unused-cache perturbation accepts no forward batch but skips without radix_cache."""
|
||||
device = DEFAULT_DEVICE
|
||||
group = make_buffer_group(kind=PoolKind.FULL, has_real_kv=True)
|
||||
source = group.real_kv_sources_k[0]
|
||||
source.tensor.copy_(
|
||||
torch.arange(source.tensor.numel(), dtype=torch.uint8).view_as(
|
||||
source.tensor
|
||||
)
|
||||
)
|
||||
manager = PerturbManager(
|
||||
config=PerturbConfig(
|
||||
req_to_token_prob=0.0,
|
||||
real_kv_used_prob=0.0,
|
||||
real_kv_unused_cache_prob=1.0,
|
||||
real_kv_post_forward_prob=0.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
),
|
||||
req_to_token_pool=make_req_to_token_pool(device, max_reqs=4, max_seq_len=8),
|
||||
buffer_groups=(group,),
|
||||
outer_step_counter_getter=lambda: 10,
|
||||
sweep_interval=1,
|
||||
)
|
||||
|
||||
snapshot = source.tensor.clone()
|
||||
with patch.object(torch, "rand", return_value=torch.tensor(0.0)):
|
||||
manager.perturb_real_kv_unused_cache(None)
|
||||
|
||||
self.assertTrue(torch.equal(source.tensor, snapshot))
|
||||
|
||||
|
||||
class TestPerturbUtils(CustomTestCase):
|
||||
def test_flip_first_byte_in_physical_swa_slot_does_not_translate_twice(
|
||||
|
||||
@@ -128,5 +128,44 @@ class TestPdTransferCanaryClean(_MockModelPDBase, unittest.TestCase):
|
||||
self.assert_no_canary_violation()
|
||||
|
||||
|
||||
class TestPdTransferChecksumFullRealData(_MockModelPDBase, unittest.TestCase):
|
||||
"""--kv-canary-real-data=all + sweep every step, no perturb, no violation."""
|
||||
|
||||
extra_prefill_args: ClassVar[List[str]] = mock_model_server_args(
|
||||
"--skip-server-warmup",
|
||||
"--kv-canary-real-data",
|
||||
"all",
|
||||
"--kv-canary-sweep-interval",
|
||||
"1",
|
||||
)
|
||||
extra_decode_args: ClassVar[List[str]] = mock_model_server_args(
|
||||
"--skip-server-warmup",
|
||||
"--kv-canary-real-data",
|
||||
"all",
|
||||
"--kv-canary-sweep-interval",
|
||||
"1",
|
||||
"--disaggregation-decode-enable-radix-cache",
|
||||
)
|
||||
|
||||
def test_pd_transfer_checksum_full_real_data(self) -> None:
|
||||
# Step 1: drive traffic through the PD path with full real-KV hashing.
|
||||
results = _send_parallel_requests(
|
||||
self.lb_url,
|
||||
n=_NUM_PROMPTS,
|
||||
max_new_tokens=_OUTPUT_LEN,
|
||||
timeout=240.0,
|
||||
max_workers=_NUM_PROMPTS,
|
||||
)
|
||||
|
||||
# Step 2: all requests must succeed.
|
||||
for result in results:
|
||||
self.assertEqual(result.get("status_code"), 200, result)
|
||||
|
||||
# Step 3: servers must stay healthy.
|
||||
self.assertIsNone(self.process_prefill.poll(), "Prefill server died")
|
||||
self.assertIsNone(self.process_decode.poll(), "Decode server died")
|
||||
self.assert_no_canary_violation()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
import logging
|
||||
import unittest
|
||||
from unittest.mock import Mock
|
||||
|
||||
from sglang.srt.kv_canary.perturb import real_kv_used
|
||||
from sglang.srt.kv_canary.perturb.config import PerturbConfig, TargetGroupKind
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestCanaryPerturb(CustomTestCase):
|
||||
def test_real_kv_used_logs_when_target_group_has_no_real_kv_sources(self) -> None:
|
||||
"""Verify real KV used perturbation logs when the target group has no sources."""
|
||||
config = PerturbConfig(
|
||||
req_to_token_prob=0.0,
|
||||
real_kv_used_prob=1.0,
|
||||
real_kv_unused_cache_prob=0.0,
|
||||
real_kv_post_forward_prob=0.0,
|
||||
target_group_kind=TargetGroupKind.FULL,
|
||||
warmup_steps=0,
|
||||
)
|
||||
warmup_gate = Mock()
|
||||
warmup_gate.is_in_warmup.return_value = False
|
||||
|
||||
# Empty buffer_groups means pick_target_group returns None, so run() takes
|
||||
# the early-return branch before any slot is picked.
|
||||
with self.assertLogs(real_kv_used.logger.name, level=logging.INFO) as logs:
|
||||
real_kv_used.run(
|
||||
maybe_inaccurate_forward_batch=Mock(),
|
||||
config=config,
|
||||
req_to_token_pool=Mock(),
|
||||
buffer_groups=(),
|
||||
swa_window_size=0,
|
||||
warmup_gate=warmup_gate,
|
||||
)
|
||||
|
||||
self.assertIn(
|
||||
"kv_canary perturb real_kv_used: skipped because no target group with "
|
||||
"real_kv_sources_k matched target_group_kind=full",
|
||||
"\n".join(logs.output),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user