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

[Memory] Size the CUDA graph pool from warmup measurements and fix graph-pool borrowing (#36911)

Browse Source 
Co-authored-by: cctry <cctry@fb.com>
This commit is contained in:
cctry
2026-09-01 09:32:38 -07:00
committed by GitHub
co-authored by cctry
parent c34f378342
commit 9a05b470fa
17 changed files with 681 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/spec/dflash/test_dflash.py
+35 -7
View File
@@ -1,4 +1,5 @@
import unittest
from contextlib import ExitStack
import openai
@@ -22,7 +23,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=420, stage="base-b", runner_config="1-gpu-small")
register_cuda_ci(est_time=500, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=420, stage="stage-b", runner_config="1-gpu-small-amd")
@@ -45,6 +46,8 @@ class TestDFlashServerBase(
draft_model = DEFAULT_DRAFT_MODEL_DFLASH
gsm8k_accuracy_thres = 0.75
gsm8k_accept_length_thres = 2.8
# (env, value) pairs applied around the server launch.
extra_env_overrides: tuple = ()
@classmethod
def setUpClass(cls):
@@ -71,12 +74,15 @@ class TestDFlashServerBase(
if cls.disable_overlap:
launch_args.append("--disable-overlap-schedule")
launch_args.extend(cls.other_launch_args)
with (
envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.override(cls.overlap_plan_stream),
envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1),
envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True),
envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.override(True),
):
with ExitStack() as stack:
for env, value in (
(envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM, cls.overlap_plan_stream),
(envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),
(envs.SGLANG_ENABLE_ASYNC_ASSERT, True),
(envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN, True),
*cls.extra_env_overrides,
):
stack.enter_context(env.override(value))
cls.process = popen_launch_server(
cls.model,
cls.base_url,
@@ -173,5 +179,27 @@ class TestDFlashServerOverlapPlanStream(TestDFlashServerOverlap):
overlap_plan_stream = True
@unittest.skipIf(
is_hip(),
"borrowing is CUDA-only and ROCm has no DFLASH sampling-verify kernel",
)
class TestDFlashServerGraphPoolBorrow(TestDFlashServerBase):
"""Verify probabilities served out of idle CUDA graph storage. Only the
non-greedy accept path borrows, hence the sampled GSM8K below."""
extra_env_overrides = (
(envs.SGLANG_ENABLE_GRAPH_POOL_BORROW, 1),
(envs.SGLANG_ENABLE_GRAPH_POOL_PRECARVE, 1),
)
disable_overlap = False
gsm8k_temperature = 0.6
gsm8k_top_p = 0.95
# Measured over 4x200 examples per arm: score 0.7525 (sd 0.021) and accept
# length 2.80, borrowing on or off. Corruption collapses accept length
# toward 1.0; the accuracy floor is loose to absorb other hardware.
gsm8k_accuracy_thres = 0.60
gsm8k_accept_length_thres = 2.0
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py
+66 -6
View File
@@ -674,6 +674,9 @@ class _SchedulerWorker:
def __init__(self, trace, *, post_capture_active=False):
self._trace = trace
self.model_runner = SimpleNamespace(
device="cuda",
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"),
)
@@ -687,7 +690,7 @@ class _SchedulerWorker:
class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
@staticmethod
def _scheduler(worker, trace, *, mode):
def _scheduler(worker, trace, *, mode, draft_worker=None):
from sglang.srt.managers.scheduler import Scheduler
scheduler = Scheduler.__new__(Scheduler)
@@ -696,17 +699,38 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
)
scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker)
scheduler.maybe_init_draft_worker = lambda: setattr(
scheduler, "draft_worker", None
scheduler, "draft_worker", draft_worker
)
scheduler.init_memory_pools = lambda: trace.append("memory_pool")
scheduler.init_all_attention_backends = lambda: trace.append("attention")
scheduler.init_all_cuda_graphs = lambda: trace.append("capture")
return scheduler
def _run_startup(self, mode):
def _run_startup(self, mode, *, use_draft_worker=False):
trace = []
worker = _SchedulerWorker(trace, post_capture_active=True)
scheduler = self._scheduler(worker, trace, mode=mode)
draft_worker = (
SimpleNamespace(prewarm_sampling=lambda: trace.append("draft_prewarm"))
if use_draft_worker
else None
)
scheduler = self._scheduler(
worker,
trace,
mode=mode,
draft_worker=draft_worker,
)
class StreamContext:
def __enter__(self):
trace.append("stream_enter")
def __exit__(self, *_args):
trace.append("stream_exit")
def stream_context(stream):
self.assertIs(stream, worker.model_runner.forward_stream)
return StreamContext()
def stop_after_startup():
raise RuntimeError("stop after startup")
@@ -723,6 +747,10 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
)
),
),
patch(
"sglang.srt.managers.scheduler.torch.get_device_module",
return_value=SimpleNamespace(stream=stream_context),
),
self.assertRaisesRegex(RuntimeError, "stop after startup"),
):
scheduler.init_model_worker()
@@ -732,13 +760,45 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
def test_serial_path_skips_overlap_hooks(self):
self.assertEqual(
self._run_startup("serial"),
["memory_pool", "attention", "capture", "resize"],
[
"memory_pool",
"attention",
"capture",
"stream_enter",
"prewarm",
"stream_exit",
"resize",
],
)
def test_overlap_starts_before_capture_and_finalizes_after(self):
self.assertEqual(
self._run_startup("overlap"),
["start", "memory_pool", "attention", "capture", "resize", "finalize"],
[
"start",
"memory_pool",
"attention",
"capture",
"stream_enter",
"prewarm",
"stream_exit",
"resize",
"finalize",
],
)
def test_draft_worker_prewarm_uses_target_forward_stream(self):
self.assertEqual(
self._run_startup("serial", use_draft_worker=True),
[
"memory_pool",
"attention",
"capture",
"stream_enter",
"draft_prewarm",
"stream_exit",
"resize",
],
)
test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py
+3
View File
@@ -56,6 +56,9 @@ def _make_backend(runner):
backend._outputs = {}
backend._pool = None
backend._capture_stream = None
backend._precarve = SimpleNamespace(
measure=contextlib.nullcontext, mint=mock.Mock()
)
backend._memory_saver_adapter = None
backend._cuda_graph_runner = runner
backend._device_module = runner.device_module
test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py
+194 -33
View File
@@ -12,7 +12,8 @@ from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
FullCudaGraphBackend,
)
from sglang.srt.model_executor.runner_utils import pool
from sglang.srt.speculative import eagle_utils
from sglang.srt.speculative import dflash_utils, dflash_worker_v2, eagle_utils
from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -43,6 +44,25 @@ class TestGraphPoolBorrow(CustomTestCase):
pool._borrow_extents_total = 0
pool._largest_logged_graph_pool_borrow = 0
def test_mixed_segment_runs_exclude_live_blocks(self):
"""A mixed segment's free runs are borrowable, but a returned run must
never overlap the live block — overlap would silently corrupt
graph-owned data instead of raising an OOM."""
snapshot = [
{
"allocated_size": 4096,
"total_size": 3 * 4096,
"blocks": [
{"state": "inactive", "address": 0x1000, "size": 4096},
{"state": "active_allocated", "address": 0x2000, "size": 4096},
{"state": "inactive", "address": 0x3000, "size": 4096},
],
}
]
with patch.object(pool.torch.cuda, "memory_snapshot", return_value=snapshot):
runs = pool.find_free_graph_pool_runs((0, 1))
self.assertEqual(sorted(runs), [(0x1000, 4096), (0x3000, 4096)])
def test_graph_replay_fails_during_active_pool_borrow(self):
graph = Mock()
backend = object.__new__(FullCudaGraphBackend)
@@ -74,6 +94,24 @@ class TestGraphPoolBorrow(CustomTestCase):
graph.replay.assert_not_called()
def test_high_cursor_keeps_reusable_cached_segments(self):
stub = MagicMock(cursor_bytes=600, freed_bytes=0)
mem_pool = MagicMock()
with (
patch.object(pool, "graph_pool_borrow_enabled", return_value=True),
patch.object(pool, "_borrow_stub", stub),
patch.object(pool, "_borrow_mem_pool", mem_pool),
patch.object(pool, "_borrow_extents_total", 1000),
patch.object(pool, "_teardown_borrow_pool") as teardown,
patch.object(pool.torch, "empty"),
patch.object(pool.torch.cuda, "use_mem_pool"),
):
with pool.borrow_graph_pool(user="test"):
pass
teardown.assert_not_called()
def test_external_graph_storage_can_disable_borrowing(self):
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
@@ -84,29 +122,8 @@ class TestGraphPoolBorrow(CustomTestCase):
pool.disable_graph_pool_borrow("graph storage is externally managed")
self.assertFalse(pool.graph_pool_borrow_enabled())
def test_eagle_non_greedy_probabilities_use_borrow_scope(self):
state = {"active": False, "users": []}
@contextmanager
def tracking_borrow(*, user):
self.assertFalse(state["active"])
state["active"] = True
state["users"].append(user)
try:
yield
finally:
state["active"] = False
import torch.nn.functional as F
real_softmax = F.softmax
def checked_softmax(*args, **kwargs):
self.assertTrue(state["active"])
return real_softmax(*args, **kwargs)
def test_eagle_non_greedy_probabilities_do_not_borrow_graph_pool(self):
def fake_sampling(**kwargs):
self.assertTrue(state["active"])
kwargs["predicts"].fill_(3)
kwargs["accept_index"].fill_(0)
kwargs["accept_token_num"].fill_(1)
@@ -146,9 +163,8 @@ class TestGraphPoolBorrow(CustomTestCase):
tp_group = SimpleNamespace(world_size=1)
with (
patch.object(eagle_utils, "borrow_graph_pool", tracking_borrow),
patch.object(pool, "borrow_graph_pool") as borrow_graph_pool,
patch.object(eagle_utils, "get_spec", return_value=spec_config),
patch("torch.nn.functional.softmax", side_effect=checked_softmax),
patch(
"sglang.srt.layers.dp_attention.is_dp_attention_enabled",
return_value=False,
@@ -163,8 +179,7 @@ class TestGraphPoolBorrow(CustomTestCase):
verify_input, batch, logits_output
)
self.assertEqual(state["users"], ["EAGLE probability borrow"])
self.assertFalse(state["active"])
borrow_graph_pool.assert_not_called()
self.assertTrue(torch.equal(predict, torch.full_like(predict, 3)))
self.assertTrue(torch.equal(accept_lens, torch.full_like(accept_lens, 2)))
self.assertTrue(torch.equal(accept_index, torch.zeros_like(accept_index)))
@@ -176,8 +191,9 @@ class TestGraphPoolBorrow(CustomTestCase):
graph = torch.cuda.CUDAGraph()
x = torch.zeros(8, device="cuda")
stream = torch.cuda.Stream()
with torch.cuda.stream(stream), torch.cuda.graph(
graph, pool=handle, stream=stream
with (
torch.cuda.stream(stream),
torch.cuda.graph(graph, pool=handle, stream=stream),
):
# Two capture-only transients become disjoint free graph-pool runs.
transient_a = torch.empty(48 << 20, dtype=torch.uint8, device="cuda")
@@ -237,8 +253,9 @@ class TestGraphPoolBorrow(CustomTestCase):
graph = torch.cuda.CUDAGraph()
x = torch.zeros(8, device="cuda")
stream = torch.cuda.Stream()
with torch.cuda.stream(stream), torch.cuda.graph(
graph, pool=handle, stream=stream
with (
torch.cuda.stream(stream),
torch.cuda.graph(graph, pool=handle, stream=stream),
):
transient = torch.empty(64 << 20, dtype=torch.uint8, device="cuda")
y = x + 1
@@ -267,6 +284,40 @@ class TestGraphPoolBorrow(CustomTestCase):
del graph, y
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
def test_oversized_borrow_raises_oom_then_regular_allocation_succeeds(self):
handle = torch.cuda.graph_pool_handle()
graph = torch.cuda.CUDAGraph()
x = torch.zeros(8, device="cuda")
stream = torch.cuda.Stream()
with (
torch.cuda.stream(stream),
torch.cuda.graph(graph, pool=handle, stream=stream),
):
transient = torch.empty(64 << 20, dtype=torch.uint8, device="cuda")
y = x + 1
del transient
torch.cuda.synchronize()
runs = pool.find_free_graph_pool_runs(handle)
address, run_bytes = next(run for run in runs if run[1] >= 16 << 20)
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=None),
):
pool.set_graph_pool_borrow_runs([(address, 8 << 20)])
with self.assertRaises(torch.OutOfMemoryError):
with pool.borrow_graph_pool(user="undersized-test"):
torch.empty(16 << 20, dtype=torch.uint8, device="cuda")
pool.disable_graph_pool_borrow("undersized test pool")
regular = torch.empty(16 << 20, dtype=torch.uint8, device="cuda")
self.assertEqual(regular.nbytes, 16 << 20)
del regular
self.assertGreaterEqual(run_bytes, 16 << 20)
del graph, y
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
def test_cross_stream_borrow_frees_resolve_before_pointer_reuse(self):
"""Deferred record_stream frees must not collide on the next borrow."""
@@ -274,8 +325,9 @@ class TestGraphPoolBorrow(CustomTestCase):
graph = torch.cuda.CUDAGraph()
x = torch.zeros(8, device="cuda")
stream = torch.cuda.Stream()
with torch.cuda.stream(stream), torch.cuda.graph(
graph, pool=handle, stream=stream
with (
torch.cuda.stream(stream),
torch.cuda.graph(graph, pool=handle, stream=stream),
):
transient = torch.empty(128 << 20, dtype=torch.uint8, device="cuda")
y = x + 1
@@ -301,6 +353,115 @@ class TestGraphPoolBorrow(CustomTestCase):
del graph, y
def test_dflash_verify_output_buffers_predate_the_borrow_scope(self):
"""The chain verify buffers outlive the step, so creating them inside
the borrow scope would let the next replay overwrite the accept
length instead of raising."""
events = []
@contextmanager
def recording_borrow(user):
events.append(f"borrow:{user}")
yield
events.append("release")
real_buffers = dflash_utils._get_or_create_chain_verify_buffers
def recording_buffers(**kwargs):
events.append("buffers")
return real_buffers(**kwargs)
def fake_sampling(**kwargs):
kwargs["predicts"].fill_(3)
kwargs["accept_index"].fill_(0)
kwargs["accept_token_num"].fill_(1)
sampling_info = SimpleNamespace(
temperatures=torch.ones((1, 1)),
top_ks=torch.ones(1, dtype=torch.int32),
top_ps=torch.ones(1),
need_top_k_sampling=False,
need_top_p_sampling=False,
)
with (
patch.object(dflash_utils, "borrow_graph_pool", recording_borrow),
patch.object(
dflash_utils,
"_get_or_create_chain_verify_buffers",
recording_buffers,
),
patch.object(dflash_utils, "_DFLASH_SAMPLING_VERIFY_AVAILABLE", True),
patch.object(
dflash_utils,
"tree_speculative_sampling_target_only",
fake_sampling,
),
):
correct_len, bonus = (
dflash_utils.compute_dflash_sampling_correct_drafts_and_bonus(
candidates=torch.zeros((1, 2), dtype=torch.int64),
next_token_logits=torch.randn((2, 8)),
sampling_info=sampling_info,
threshold_single=1.0,
threshold_acc=1.0,
)
)
self.assertEqual(
events, ["buffers", "borrow:DFLASH verify probabilities", "release"]
)
self.assertTrue(torch.equal(correct_len, torch.ones_like(correct_len)))
self.assertTrue(torch.equal(bonus, torch.full_like(bonus, 3)))
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
def test_dflash_prewarm_falls_back_when_the_rehearsal_exhausts_the_pool(self):
"""A rehearsal too large for the pool must retire borrowing and
re-measure, rather than crash startup or leave KV sizing without the
headroom it now has to reserve."""
worker = object.__new__(DFlashWorkerV2)
worker.block_size = 4
worker.device = "cuda"
worker._target_worker = SimpleNamespace(
model_runner=SimpleNamespace(
max_running_requests=2,
max_decode_logits_rows=lambda: 8,
sampling_prewarm_result=None,
),
model_config=SimpleNamespace(vocab_size=32),
)
worker.model_runner = worker._target_worker.model_runner
calls = []
def rehearse(**kwargs):
calls.append(pool.graph_pool_borrow_enabled())
if len(calls) == 2:
raise torch.OutOfMemoryError("rehearsal too large")
return torch.zeros(2), torch.zeros(2)
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=(1, 2)),
patch.object(
dflash_worker_v2,
"compute_dflash_sampling_correct_drafts_and_bonus",
rehearse,
),
):
self.assertTrue(pool.graph_pool_borrow_enabled())
result = worker.prewarm_sampling()
self.assertFalse(pool.graph_pool_borrow_enabled())
# Warm pass outside the pool, borrowed pass that OOMs, retry after the
# fallback retires borrowing.
self.assertEqual(calls, [False, True, False])
# 2 rows x 4 draft tokens x 32 vocab x 4 bytes.
self.assertEqual(result.sampling_input_bytes, 2 * 4 * 32 * 4)
self.assertGreaterEqual(
result.sampling_headroom_bytes, result.sampling_input_bytes
)
self.assertIs(worker.model_runner.sampling_prewarm_result, result)
if __name__ == "__main__":
unittest.main()