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

Fix KV-canary workspace accounting after graph capture (#38596)

Browse Source
This commit is contained in:
cctry
2026-09-10 16:10:12 -07:00
committed by GitHub
parent 42bbaac259
commit 203d7e812c
12 changed files with 274 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kv_canary/test_self_unit_capacities.py
+59
View File
@@ -2,7 +2,13 @@ from __future__ import annotations
import unittest
import torch
from sglang.kernels.ops.kv_canary.verify import VerifyPlan
from sglang.kernels.ops.kv_canary.write import WritePlan
from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
from sglang.srt.kv_canary.plan_input import PlanInput
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
CudaGraphConfig,
@@ -88,6 +94,59 @@ class TestComputeLaunchCapacities(CustomTestCase):
with self.assertRaisesRegex(ValueError, "pool_slot_count"):
self._from_args(max_bs=1, max_seq_len=1, max_total_num_tokens=0)
def test_workspace_matches_allocated_tensors(self) -> None:
device = torch.device("cpu")
for slots, requests, entries, groups in ((1, 1, 1, 1), (1024, 8, 128, 3)):
with self.subTest(slots=slots, groups=groups):
capacities = CanaryLaunchCapacities(
per_forward_verify_capacity=3 * slots,
per_forward_write_req_capacity=requests,
per_forward_write_entry_capacity=entries,
)
verify = VerifyPlan.allocate(verify_capacity=3 * slots, device=device)
write = WritePlan.allocate(write_req_capacity=requests, device=device)
expected = ExpectedInputs.allocate(capacity=entries, device=device)
plan = PlanInput.allocate(bs_capacity=requests, device=device)
group_tensors = (
verify.verify_slot_indices,
verify.verify_expected_tokens,
verify.verify_expected_positions,
verify.verify_prev_slot_indices,
verify.verify_num_valid,
verify.enable,
write.write_offsets,
write.write_seed_slot_indices,
write.write_num_valid_reqs,
)
shared_tensors = (
expected.tokens,
expected.positions,
plan.req_pool_indices,
plan.prefix_lens,
plan.extend_seq_lens,
plan.req_to_verify_expected_tokens_valid_lens,
)
actual_bytes = groups * sum(t.nbytes for t in group_tensors) + sum(
t.nbytes for t in shared_tensors
)
self.assertEqual(
capacities.per_forward_workspace_bytes(num_buffer_groups=groups),
actual_bytes,
)
def test_workspace_scales_with_pool_slots(self) -> None:
for groups in (1, 4):
with self.subTest(groups=groups):
small, large = (
self._from_args(max_bs=8, max_seq_len=64, max_total_num_tokens=n)
for n in (1024, 4096)
)
self.assertEqual(
large.per_forward_workspace_bytes(num_buffer_groups=groups)
- small.per_forward_workspace_bytes(num_buffer_groups=groups),
(4096 - 1024) * 96 * groups,
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py
+11 -3
View File
@@ -4,7 +4,7 @@ import dataclasses
import re
import unittest
from types import SimpleNamespace
from unittest.mock import call, patch
from unittest.mock import Mock, call, patch
import torch
from torch import nn
@@ -711,7 +711,9 @@ class _SchedulerWorker:
forward_stream=object(),
prewarm_sampling=lambda: trace.append("prewarm"),
token_to_kv_pool=SimpleNamespace(post_capture_active=post_capture_active),
post_capture_resize_kv_pool=lambda: trace.append("resize"),
post_capture_resize_kv_pool=Mock(
side_effect=lambda *, draft_runners: trace.append("resize")
),
)
def start_startup_weight_load(self):
@@ -750,7 +752,10 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
trace = []
worker = _SchedulerWorker(trace, post_capture_active=True)
draft_worker = (
SimpleNamespace(prewarm_sampling=lambda: trace.append("draft_prewarm"))
SimpleNamespace(
prewarm_sampling=lambda: trace.append("draft_prewarm"),
_draft_model_runners=lambda: (worker.model_runner,),
)
if use_draft_worker
else None
)
@@ -795,6 +800,9 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
):
scheduler.init_model_worker()
worker.model_runner.post_capture_resize_kv_pool.assert_called_once_with(
draft_runners=(worker.model_runner,) if use_draft_worker else ()
)
return trace
def test_serial_path_skips_overlap_hooks(self):
test/registered/unit/model_executor/test_kv_canary_headroom.py
+139
View File
@@ -0,0 +1,139 @@
import unittest
from types import SimpleNamespace
from unittest.mock import Mock, patch
from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
CudaGraphConfig,
PhaseConfig,
)
from sglang.srt.model_executor.model_runner_components import kv_pool_runtime
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, stage="base-a", runner_config="cpu")
_GIB = 1 << 30
def _manager(slots, groups):
manager = CanaryManager.__new__(CanaryManager)
manager._launch_capacities = CanaryLaunchCapacities(
per_forward_verify_capacity=3 * slots,
per_forward_write_req_capacity=128,
per_forward_write_entry_capacity=4096,
)
manager._buffer_groups = (None,) * groups
return manager
class TestCanaryHeadroom(CustomTestCase):
def _resize(self, target, drafts=(), *, graph_borrow=False, eager_gap=False):
config = MemoryPoolConfig(max_total_num_tokens=1024)
pool = Mock(post_capture_backed_bytes=2 * _GIB, dtype="bfloat16")
runner = SimpleNamespace(
token_to_kv_pool=pool,
device="cuda",
gpu_id=0,
pre_model_load_memory=32,
mem_fraction_static=0.875,
max_running_requests=16 if eager_gap else None,
model_config=SimpleNamespace(is_multimodal=False),
sampling_prewarm_result=SimpleNamespace(sampling_headroom_bytes=6 * _GIB),
canary_manager=target,
kv_cache_configurator=Mock(),
max_total_num_tokens=1_000_000,
token_to_kv_pool_allocator=Mock(),
req_to_token_pool=Mock(),
)
runner.kv_cache_configurator.config_from_budget.return_value = config
runner.kv_cache_configurator.resolve_max_num_reqs.return_value = 16
with (
get_context().override_server_args(
disaggregation_mode="null",
cuda_graph_config=CudaGraphConfig(
decode=PhaseConfig(backend=Backend.FULL, max_bs=8)
),
),
patch.object(kv_pool_runtime.torch.cuda, "synchronize"),
patch.object(
kv_pool_runtime,
"get_world_group",
return_value=SimpleNamespace(world_size=1, cpu_group=None),
),
patch.object(kv_pool_runtime, "get_available_gpu_memory", return_value=20),
patch.object(kv_pool_runtime, "mambaish_config", return_value=None),
patch.object(
kv_pool_runtime, "get_device_memory_capacity", return_value=32
),
patch.object(
kv_pool_runtime,
"pre_capture_activation_reserve_mb",
return_value=8 * 1024,
),
patch.object(
kv_pool_runtime, "graph_pool_borrow_enabled", return_value=graph_borrow
),
patch.object(kv_pool_runtime, "mm_runtime_reservation_gb", return_value=1),
):
resize = kv_pool_runtime.compute_post_capture_kv_resize(
runner,
draft_runners=tuple(SimpleNamespace(canary_manager=m) for m in drafts),
)
pool.finalize_backing.assert_called_once_with(config)
runner.token_to_kv_pool_allocator.resize.assert_called_once_with(config)
runner.req_to_token_pool.reset_aux_cache_allocator.assert_called_once_with()
self.assertEqual(resize.max_total_num_tokens, config.max_total_num_tokens)
args, kwargs = runner.kv_cache_configurator.config_from_budget.call_args
self.assertEqual(kwargs, {"cap_tokens": 1_000_000})
return args[0]
def test_canary_off_preserves_every_budget_byte(self):
for graph_borrow, eager_gap, headroom in (
(False, False, 6),
(True, False, 4),
(True, True, 8),
):
with self.subTest(graph_borrow=graph_borrow, eager_gap=eager_gap):
self.assertEqual(
self._resize(
None, (None,), graph_borrow=graph_borrow, eager_gap=eager_gap
),
(20 - headroom - 1 + 2) * _GIB,
)
def test_workspace_is_added_to_other_headroom_at_installed_capacity(self):
manager = _manager(1_000_000, 3)
workspace = manager.per_forward_workspace_bytes()
for graph_borrow, eager_gap in ((False, False), (True, False), (True, True)):
with self.subTest(graph_borrow=graph_borrow, eager_gap=eager_gap):
baseline = self._resize(
None, graph_borrow=graph_borrow, eager_gap=eager_gap
)
self.assertEqual(
self._resize(
manager, graph_borrow=graph_borrow, eager_gap=eager_gap
),
baseline - workspace,
)
def test_sequential_target_and_drafts_reserve_largest_workspace(self):
small, large = _manager(1024, 1), _manager(4096, 4)
for target, drafts in (
(small, (large, None)),
(large, (small,)),
(None, (small, large)),
):
with self.subTest(target=target, drafts=drafts):
self.assertEqual(
self._resize(target, drafts),
self._resize(None) - large.per_forward_workspace_bytes(),
)
if __name__ == "__main__":
unittest.main()