Add a periodic full-radix-tree KV-canary sweep (#26812)
This commit is contained in:
fzyzcjy
@@ -11,7 +11,7 @@ register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _BaselineBase(CanaryE2EBase):
|
||||
"""No perturb, kv-canary=log. Server should run clean with no canary
|
||||
"""No perturb, kv-canary=log, sweep off. Server should run clean with no canary
|
||||
violations and every request must come back 200."""
|
||||
|
||||
kv_canary_mode = CanaryMode.LOG
|
||||
|
||||
@@ -63,6 +63,28 @@ class TestSelfUnitEndpoint(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.device = DEFAULT_DEVICE
|
||||
|
||||
def test_launch_sweep_only_calls_verify(self):
|
||||
"""Verify sweep launch invokes only the verify kernel."""
|
||||
calls: list[str] = []
|
||||
with patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_verify_kernel",
|
||||
lambda **kwargs: calls.append("verify"),
|
||||
), patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_write_kernel",
|
||||
lambda **kwargs: calls.append("write"),
|
||||
):
|
||||
ep = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_K_FULL
|
||||
)
|
||||
args = _make_kernel_args(self.device)
|
||||
ep.launch_sweep(
|
||||
verify_plan=args.verify_plan,
|
||||
violation_log=args.violation_log,
|
||||
)
|
||||
self.assertEqual(calls, ["verify"])
|
||||
|
||||
def test_launch_per_forward_passes_kernel_kind(self):
|
||||
"""Verify per-forward launch passes the endpoint kernel kind."""
|
||||
captured: list[tuple[str, CanaryLaunchTag]] = []
|
||||
@@ -109,33 +131,19 @@ class TestSelfUnitEndpoint(CustomTestCase):
|
||||
):
|
||||
shared_log = ViolationLog.allocate(ring_capacity=2, device=self.device)
|
||||
ep_a = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_K_FULL
|
||||
)
|
||||
ep_b = _make_endpoint(
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.HEAD_V_FULL
|
||||
device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_V_FULL
|
||||
)
|
||||
|
||||
args = _make_kernel_args(self.device)
|
||||
ep_a.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
plan = VerifyPlan.allocate(verify_capacity=1, device=self.device)
|
||||
ep_a.launch_sweep(
|
||||
verify_plan=plan,
|
||||
violation_log=shared_log,
|
||||
)
|
||||
ep_b.launch_per_forward(
|
||||
verify_plan=args.verify_plan,
|
||||
write_plan=args.write_plan,
|
||||
input_ids=args.input_ids,
|
||||
positions=args.positions,
|
||||
out_cache_loc=args.out_cache_loc,
|
||||
enable_write_input_assert=args.enable_write_input_assert,
|
||||
enable_verify_token_assert=args.enable_verify_token_assert,
|
||||
expected_inputs=args.expected_inputs,
|
||||
ep_b.launch_sweep(
|
||||
verify_plan=plan,
|
||||
violation_log=shared_log,
|
||||
)
|
||||
self.assertEqual(captured_rings[0], captured_rings[1])
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary
|
||||
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_radix_cache
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestSelfUnitRadixWalker(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.device = DEFAULT_DEVICE
|
||||
|
||||
def test_single_node_chain_positions_increase(self):
|
||||
"""Verify a single radix chain emits increasing positions."""
|
||||
chain = [10, 20, 30, 40]
|
||||
cache = make_radix_cache([[], chain], device=self.device)
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache)
|
||||
self.assertEqual(result.slot_indices.tolist(), chain)
|
||||
self.assertEqual(result.positions.tolist(), [0, 1, 2, 3])
|
||||
self.assertEqual(result.prev_slot_indices.tolist(), [-1, 10, 20, 30])
|
||||
|
||||
def test_child_node_first_slot_prev_is_parent_last(self):
|
||||
"""Verify child chains link their first slot to the parent tail."""
|
||||
parent = [7, 8]
|
||||
child = [9, 10]
|
||||
cache = make_radix_cache([[], parent, child], device=self.device)
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache)
|
||||
self.assertEqual(result.slot_indices.tolist(), parent + child)
|
||||
self.assertEqual(result.prev_slot_indices.tolist()[len(parent)], parent[-1])
|
||||
|
||||
def test_root_child_first_slot_prev_minus_one(self):
|
||||
"""Verify root child chains use -1 as the initial previous slot."""
|
||||
cache = make_radix_cache([[], [42, 43]], device=self.device)
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache)
|
||||
self.assertEqual(int(result.prev_slot_indices[0]), -1)
|
||||
|
||||
def test_position_equals_depth_from_root(self):
|
||||
"""Verify emitted positions match depth from the radix root."""
|
||||
cache = make_radix_cache([[], [1, 2], [3], [4, 5]], device=self.device)
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache)
|
||||
self.assertEqual(result.positions.tolist(), [0, 1, 2, 3, 4])
|
||||
self.assertEqual(result.slot_indices.tolist(), [1, 2, 3, 4, 5])
|
||||
|
||||
def test_walk_includes_locked_nodes_by_default(self):
|
||||
"""Verify radix walking includes locked nodes by default."""
|
||||
cache = make_radix_cache([[], [1, 2], [3, 4]], device=self.device)
|
||||
locked_node = next(iter(cache.root_node.children.values()))
|
||||
locked_node.lock_ref = 1
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache)
|
||||
self.assertEqual(result.slot_indices.tolist(), [1, 2, 3, 4])
|
||||
|
||||
def test_walk_unlocked_only_skips_locked(self):
|
||||
"""Verify unlocked-only radix walking skips locked nodes."""
|
||||
cache = make_radix_cache([[], [1, 2], [3, 4]], device=self.device)
|
||||
locked_node = next(iter(cache.root_node.children.values()))
|
||||
locked_node.lock_ref = 1
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True)
|
||||
self.assertEqual(result.slot_indices.tolist(), [3, 4])
|
||||
|
||||
def test_walk_unlocked_only_uses_swa_full_lock_ref(self):
|
||||
"""Verify SWA radix walking honors full-pool lock references."""
|
||||
cache = SWARadixCache.__new__(SWARadixCache)
|
||||
cache.device = self.device
|
||||
cache.page_size = 1
|
||||
cache.disable = False
|
||||
|
||||
root = TreeNode()
|
||||
root.value = torch.tensor([], dtype=torch.int32, device=self.device)
|
||||
cache.root_node = root
|
||||
|
||||
locked_child = TreeNode()
|
||||
locked_child.value = torch.tensor([1, 2], dtype=torch.int32, device=self.device)
|
||||
locked_child.parent = root
|
||||
locked_child.full_lock_ref = 1
|
||||
root.children[locked_child.id] = locked_child
|
||||
|
||||
unlocked_child = TreeNode()
|
||||
unlocked_child.value = torch.tensor(
|
||||
[3, 4], dtype=torch.int32, device=self.device
|
||||
)
|
||||
unlocked_child.parent = root
|
||||
root.children[unlocked_child.id] = unlocked_child
|
||||
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True)
|
||||
self.assertEqual(result.slot_indices.tolist(), [3, 4])
|
||||
|
||||
def test_swa_resident_only_skips_tombstoned_nodes(self):
|
||||
"""Verify SWA radix walking skips nodes whose SWA storage was evicted."""
|
||||
cache = SWARadixCache.__new__(SWARadixCache)
|
||||
cache.device = self.device
|
||||
cache.page_size = 1
|
||||
cache.disable = False
|
||||
|
||||
root = TreeNode()
|
||||
root.value = torch.tensor([], dtype=torch.int32, device=self.device)
|
||||
cache.root_node = root
|
||||
|
||||
tombstoned_child = TreeNode()
|
||||
tombstoned_child.value = torch.tensor(
|
||||
[1, 2], dtype=torch.int32, device=self.device
|
||||
)
|
||||
tombstoned_child.parent = root
|
||||
tombstoned_child.swa_tombstone = True
|
||||
root.children[tombstoned_child.id] = tombstoned_child
|
||||
|
||||
resident_child = TreeNode()
|
||||
resident_child.value = torch.tensor(
|
||||
[3, 4], dtype=torch.int32, device=self.device
|
||||
)
|
||||
resident_child.parent = root
|
||||
root.children[resident_child.id] = resident_child
|
||||
|
||||
result = walk_radix_cache_for_canary(
|
||||
radix_cache=cache,
|
||||
swa_resident_only=True,
|
||||
)
|
||||
self.assertEqual(result.slot_indices.tolist(), [3, 4])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from sglang.srt.kv_canary import endpoint as endpoint_module
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import (
|
||||
make_forward_batch,
|
||||
make_radix_cache,
|
||||
make_req_to_token_pool,
|
||||
)
|
||||
from sglang.test.kv_canary.runner_test_base import (
|
||||
CanaryManagerTestCase,
|
||||
make_config,
|
||||
make_manager,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _run_one_cycle(manager, forward_batch) -> None:
|
||||
with manager.with_ops_outside_graph(
|
||||
single_forward_indices=[0],
|
||||
maybe_inaccurate_forward_batch=forward_batch,
|
||||
):
|
||||
with manager.with_active_single_forward_manager(0):
|
||||
pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch)
|
||||
manager.post_ops_maybe_inside_graph(forward_batch, pre_ops_output)
|
||||
|
||||
|
||||
class TestSelfUnitManagerSweep(CanaryManagerTestCase):
|
||||
def test_sweep_every_n_cadence(self) -> None:
|
||||
"""Verify sweep execution follows the configured step cadence."""
|
||||
config = make_config(sweep_interval=4)
|
||||
manager = make_manager(device=self.device, config=config)
|
||||
forward_batch = make_forward_batch(self.device)
|
||||
|
||||
sweep_calls: list[int] = []
|
||||
real_maybe = manager._sweep_orchestrator.maybe_run_sweep
|
||||
|
||||
def _spy() -> None:
|
||||
before = manager._sweep_orchestrator._last_sweep_step
|
||||
real_maybe()
|
||||
if manager._sweep_orchestrator._last_sweep_step != before:
|
||||
sweep_calls.append(manager._outer_step_counter)
|
||||
|
||||
with patch.object(manager._sweep_orchestrator, "maybe_run_sweep", _spy):
|
||||
for _ in range(12):
|
||||
_run_one_cycle(manager, forward_batch)
|
||||
self.assertEqual(sweep_calls, [0, 4, 8])
|
||||
|
||||
def test_sweep_path_launches_sweep_kernels(self) -> None:
|
||||
"""Verify sweep paths launch sweep verify kernels."""
|
||||
config = make_config(sweep_interval=1)
|
||||
manager = make_manager(device=self.device, config=config)
|
||||
forward_batch = make_forward_batch(self.device)
|
||||
manager._single_forward_managers[0].pre_ops_outside_graph(
|
||||
maybe_inaccurate_forward_batch=forward_batch
|
||||
)
|
||||
with manager.with_active_single_forward_manager(0):
|
||||
manager.pre_ops_maybe_inside_graph(forward_batch)
|
||||
|
||||
cache = make_radix_cache([[], [10, 11, 12]], device=self.device)
|
||||
cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
manager.attach_radix_cache(cache)
|
||||
|
||||
sweep_kernel_kinds: list[str] = []
|
||||
with patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_verify_kernel",
|
||||
lambda **kwargs: sweep_kernel_kinds.append(
|
||||
kwargs["context"].kernel_kind.name
|
||||
),
|
||||
):
|
||||
manager._sweep_orchestrator.maybe_run_sweep()
|
||||
self.assertTrue(any("SWEEP" in kind for kind in sweep_kernel_kinds))
|
||||
|
||||
def test_sweep_allocates_verify_plan_from_walker_output(self) -> None:
|
||||
"""Verify sweep planning sizes the verify plan from walker output."""
|
||||
manager = make_manager(device=self.device)
|
||||
cache = make_radix_cache([[], [10, 11], [12, 13, 14]], device=self.device)
|
||||
cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
manager.attach_radix_cache(cache)
|
||||
|
||||
valid_counts: list[int] = []
|
||||
with patch.object(
|
||||
endpoint_module,
|
||||
"launch_canary_verify_kernel",
|
||||
lambda **kwargs: valid_counts.append(
|
||||
int(kwargs["plan"].verify_num_valid.item())
|
||||
),
|
||||
):
|
||||
manager._sweep_orchestrator.maybe_run_sweep()
|
||||
self.assertTrue(all(count == 5 for count in valid_counts))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.sweep_plan_builder import build_verify_plan_radix_sweep
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import (
|
||||
DEFAULT_DEVICE,
|
||||
make_radix_cache,
|
||||
make_req_to_token_pool,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestSelfUnitSweepPlanBuilder(CustomTestCase):
|
||||
def setUp(self) -> None:
|
||||
self.device = DEFAULT_DEVICE
|
||||
|
||||
def test_build_verify_plan_radix_sweep(self) -> None:
|
||||
"""Verify radix sweep verify plans include cached slot chains."""
|
||||
empty_cache = make_radix_cache([[]], device=self.device)
|
||||
empty_cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
empty_out = build_verify_plan_radix_sweep(
|
||||
radix_cache=empty_cache,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
self.assertEqual(int(empty_out.verify_num_valid.item()), 0)
|
||||
|
||||
cache = make_radix_cache([[], [100, 101, 102]], device=self.device)
|
||||
cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
out = build_verify_plan_radix_sweep(
|
||||
radix_cache=cache,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
self.assertEqual(int(out.verify_num_valid.item()), 3)
|
||||
self.assertEqual(out.verify_slot_indices.dtype, torch.int64)
|
||||
self.assertEqual(out.verify_expected_positions.dtype, torch.int64)
|
||||
self.assertEqual(out.verify_prev_slot_indices.dtype, torch.int64)
|
||||
self.assertEqual(out.verify_slot_indices[:3].tolist(), [100, 101, 102])
|
||||
self.assertEqual(out.verify_expected_positions[:3].tolist(), [0, 1, 2])
|
||||
self.assertEqual(out.verify_prev_slot_indices[:3].tolist(), [-1, 100, 101])
|
||||
|
||||
def test_radix_held_slot_still_swept(self) -> None:
|
||||
"""Verify held radix slots are still included in sweep plans."""
|
||||
cache = make_radix_cache([[], [42, 43, 44]], device=self.device)
|
||||
cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
out = build_verify_plan_radix_sweep(
|
||||
radix_cache=cache,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
num_valid = int(out.verify_num_valid.item())
|
||||
self.assertEqual(num_valid, 3)
|
||||
self.assertEqual(
|
||||
set(out.verify_slot_indices[:num_valid].tolist()), {42, 43, 44}
|
||||
)
|
||||
|
||||
def test_truly_free_slot_not_swept(self) -> None:
|
||||
"""Verify free radix slots are excluded from sweep plans."""
|
||||
empty_cache = make_radix_cache([[]], device=self.device)
|
||||
empty_cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
out = build_verify_plan_radix_sweep(
|
||||
radix_cache=empty_cache,
|
||||
swa_window_size=0,
|
||||
full_to_swa_index_mapping=None,
|
||||
)
|
||||
self.assertEqual(int(out.verify_num_valid.item()), 0)
|
||||
|
||||
def test_swa_translate_preserves_evicted_as_padding_sentinel(self) -> None:
|
||||
"""Evicted (LUT=0) slots stay in the plan as the padding sentinel; the kernel does the skipping."""
|
||||
cache = make_radix_cache([[], [100, 101, 102]], device=self.device)
|
||||
cache.req_to_token_pool = make_req_to_token_pool(self.device)
|
||||
|
||||
lut = torch.zeros(200, dtype=torch.int64, device=self.device)
|
||||
lut[100] = 500
|
||||
lut[101] = 0
|
||||
lut[102] = 502
|
||||
|
||||
out = build_verify_plan_radix_sweep(
|
||||
radix_cache=cache,
|
||||
swa_window_size=128,
|
||||
full_to_swa_index_mapping=lut,
|
||||
)
|
||||
num_valid = int(out.verify_num_valid.item())
|
||||
self.assertEqual(num_valid, 3)
|
||||
self.assertEqual(out.verify_slot_indices[:num_valid].tolist(), [500, 0, 502])
|
||||
self.assertEqual(
|
||||
out.verify_prev_slot_indices[:num_valid].tolist(), [-1, 500, 0]
|
||||
)
|
||||
self.assertEqual(out.verify_expected_positions[:num_valid].tolist(), [0, 1, 2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user