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

[Bugfix] Fix full prefill CUDA graph padding and EAGLE capture (#35588)

Browse Source
This commit is contained in:
Aurick Qiao
2026-08-30 21:30:24 -07:00
committed by GitHub
parent 2ea6d17eab
commit 9a9e167179
7 changed files with 171 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py
+67 -5
View File
@@ -3,6 +3,8 @@
The Qwen3-8B test checks end-to-end accuracy with FlashInfer. The smaller
DeepSeek-Coder-V2-Lite test checks that an MLA radix-prefix hit selects the
OSS FA4 cached-prefix graph variant and matches an eager cold request.
The EAGLE3 test checks that target prefills replay FullCG rather than silently
falling back to eager while speculative decoding remains active.
The attention backend is pinned to flashinfer: plain EXTEND under full
CUDA graph requires the backend's init_forward_metadata_out_graph to
@@ -20,6 +22,7 @@ from sglang.srt.utils import get_device_sm, kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.mock_model.utils import run_mock_model_bench_serving
from sglang.test.run_eval import run_eval
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
@@ -30,7 +33,18 @@ from sglang.test.test_utils import (
# OSS FA4 coverage requires Blackwell. The PP test uses two GPUs; the other
# tests use one GPU.
register_cuda_ci(est_time=240, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=300, stage="base-b", runner_config="4-gpu-b200")
def _prefill_graph_count(base_url: str) -> float:
metrics = requests.get(base_url + "/metrics", timeout=30).text
matches = re.findall(
r'^sglang:cuda_graph_passes_total\{[^}]*mode="prefill_cuda_graph"[^}]*\}'
r"\s+([0-9.eE+-]+)$",
metrics,
re.MULTILINE,
)
return sum(map(float, matches), 0.0)
class TestFullCudaGraphPrefill(CustomTestCase):
@@ -93,6 +107,52 @@ class TestFullCudaGraphPipelineParallel(CustomTestCase):
)
class TestFullCudaGraphPrefillWithEagle3(Eagle3Base):
"""EAGLE3 target prefills replay a FullCG captured with full hidden states."""
spec_steps = 3
spec_topk = 1
spec_tokens = 4
mem_fraction_static = 0.6
max_running_requests = 1
chunked_prefill_size = 64
extra_args = (
"--enable-metrics",
"--disable-flashinfer-autotune",
"--cuda-graph-config",
(
'{"decode":{"backend":"full","max_bs":1},'
'"prefill":{"backend":"full","bs":[16],'
'"full_prefill_max_req":1}}'
),
)
def test_eagle_target_prefill_replays_full_cuda_graph(self):
prompt = (
"The capital of France is Paris. Write a concise paragraph about "
"its history, architecture, food, and culture."
)
input_ids = self.tokenizer.encode(prompt)[:16]
self.assertEqual(len(input_ids), 16)
graph_count = _prefill_graph_count(self.base_url)
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_ids,
"sampling_params": {"max_new_tokens": 32, "temperature": 0},
},
timeout=120,
)
response.raise_for_status()
output = response.json()
self.assertTrue(output["output_ids"])
self.assertGreater(output["meta_info"]["spec_verify_ct"], 0)
self.assertGreater(output["meta_info"]["spec_num_proposed_drafts"], 0)
self.assertEqual(_prefill_graph_count(self.base_url), graph_count + 1)
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
class TestFullCudaGraphChunkedPrefix(unittest.TestCase):
"""A radix-cache hit replays the OSS FA4 FullCG prefix variant."""
@@ -117,10 +177,12 @@ class TestFullCudaGraphChunkedPrefix(unittest.TestCase):
"--skip-server-warmup",
"--enable-metrics",
"--cuda-graph-config",
'{"decode":{"backend":"disabled"},'
'"prefill":{"backend":"full","bs":[32],"max_bs":32,'
'"full_prefill_max_req":1,'
'"full_prefill_prefix_chunk_tokens":64}}',
(
'{"decode":{"backend":"disabled"},'
'"prefill":{"backend":"full","bs":[32],"max_bs":32,'
'"full_prefill_max_req":1,'
'"full_prefill_prefix_chunk_tokens":64}}'
),
],
)
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py
+30 -3
View File
@@ -40,8 +40,11 @@ as tests.
import os
import random
import re
import unittest
import requests
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kl_multiturn_utils import (
make_mamba_decode_assert,
@@ -97,7 +100,9 @@ def _random_suffixes(n: int, length: int, seed: int) -> list[list[int]]:
return [[rng.randint(1, 30000) for _ in range(length)] for _ in range(n)]
def _base_args(mamba_strategy: str = "extra_buffer") -> list[str]:
def _base_args(
mamba_strategy: str = "extra_buffer", *, mem_fraction_static: float = 0.6
) -> list[str]:
return [
"--trust-remote-code",
"--attention-backend",
@@ -114,13 +119,24 @@ def _base_args(mamba_strategy: str = "extra_buffer") -> list[str]:
# the static pool leaves ~19 GB for the prefill graphs, the fa4 workspace
# and the chunked-prefill activations, which is what this config needs.
"--mem-fraction-static",
"0.6",
str(mem_fraction_static),
"--mamba-track-interval",
str(TRACK_INTERVAL),
"--enable-deterministic-inference",
]
def _prefill_graph_count(base_url: str) -> float:
metrics = requests.get(base_url + "/metrics", timeout=30).text
matches = re.findall(
r'^sglang:cuda_graph_passes_total\{[^}]*mode="prefill_cuda_graph"[^}]*\}'
r"\s+([0-9.eE+-]+)$",
metrics,
re.MULTILINE,
)
return sum(map(float, matches), 0.0)
class TestUnifiedHybridBitExact(CustomTestCase):
"""Prefill and decode must score a token identically once every kernel on the
path is batch-invariant, so any drift is a stale conv/mamba checkpoint or a
@@ -321,7 +337,9 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
def setUpClass(cls):
cls.model = _MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = _base_args() + [
# Target FullCG and both MTP draft workers retain graph pools, so this
# class needs more dynamic-memory headroom than the non-spec tests.
other_args = _base_args(mem_fraction_static=0.58) + [
"--speculative-algorithm",
"EAGLE",
"--enable-multi-layer-eagle",
@@ -333,6 +351,7 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
"3",
"--chunked-prefill-size",
"16384",
"--enable-metrics",
]
if _MODEL_REVISION:
other_args += ["--revision", _MODEL_REVISION]
@@ -353,6 +372,9 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
terminate_and_kill_process_tree(cls.process, wait_timeout=60)
def _run(self, helper):
server_info = requests.get(self.base_url + "/server_info", timeout=30).json()
self.assertEqual(server_info["cuda_graph_config"]["prefill"]["backend"], "full")
graph_count = _prefill_graph_count(self.base_url)
helper(
self.base_url,
{self.model: {"kl_div": KL_DIV_THRESHOLD}},
@@ -361,6 +383,11 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
max_new_tokens=MAX_NEW_TOKENS,
trust_remote_code=True,
)
self.assertGreater(
_prefill_graph_count(self.base_url),
graph_count,
"MTP target prefill did not replay the configured Full CUDA graph.",
)
def test_logprobs_match(self):
self._run(assert_logprobs_match)
test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py
+4 -5
View File
@@ -1403,13 +1403,12 @@ class TestPrefillNumTokenNonPaddedPostFill(unittest.TestCase):
# pads. local = clamp(1018 - 512, 0, 512).
self.assertEqual(self._fill(attn_tp_rank=1, attn_tp_size=2), 506)
def test_non_gathered_keeps_plain_fb_copy(self):
# Without a gathered buffer there is no attn-TP scatter; the plain FB
# copy must be preserved (post_fill no-op), mirroring the decode
# registry's contract.
def test_non_gathered_uses_raw_token_count(self):
# Full prefill graphs need the live raw boundary even without a
# gathered buffer so model layers can discard the padded bucket tail.
self.assertEqual(
self._fill(attn_tp_rank=0, attn_tp_size=2, require_gathered_buffer=False),
509,
1018,
)
test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py
+39 -12
View File
@@ -152,28 +152,25 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
def test_eagle_target_tc_piecewise_skips_last_mode_capture(self):
eager_runner = object()
# The server-side hidden-state ceiling is a bag leaf.
# The server-side hidden-state ceiling and graph config are bag leaves.
override = get_context().override_server_args(
enable_return_hidden_states=True,
return_hidden_states_mode="last",
cuda_graph_config=SimpleNamespace(
prefill=SimpleNamespace(backend=Backend.TC_PIECEWISE)
),
)
override.install()
self.addCleanup(override.restore)
model_runner = SimpleNamespace(
is_draft_worker=False,
spec_algorithm=SimpleNamespace(is_eagle=lambda: True),
server_args=SimpleNamespace(),
)
with patch.object(
graph_setup,
"check_cuda_graph_backend",
return_value=False,
):
capture = capture_prefill_graph(
model_runner=model_runner,
eager_runner=eager_runner,
)
capture = capture_prefill_graph(
model_runner=model_runner,
eager_runner=eager_runner,
)
self.assertIs(capture.runner, eager_runner)
@@ -234,6 +231,36 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
self.assertIs(static_batch.mm_input_embeds, mm_input_embeds)
def test_eagle_target_full_reaches_graph_construction(self):
override = get_context().override_server_args(
enable_return_hidden_states=True,
return_hidden_states_mode="last",
cuda_graph_config=SimpleNamespace(
prefill=SimpleNamespace(backend=Backend.FULL)
),
)
override.install()
self.addCleanup(override.restore)
model_runner = SimpleNamespace(
is_draft_worker=False,
lora_manager=None,
model=object(),
spec_algorithm=SimpleNamespace(is_eagle=lambda: True),
)
with (
patch.object(
graph_setup,
"resolve_language_model",
side_effect=RuntimeError("reached graph construction"),
),
self.assertRaisesRegex(RuntimeError, "reached graph construction"),
):
capture_prefill_graph(
model_runner=model_runner,
eager_runner=object(),
)
def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
graph_config = SimpleNamespace(
prefill=SimpleNamespace(full_prefill_prefix_chunk_tokens=None, max_bs=8)
@@ -243,7 +270,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
override = get_context().override_server_args(
chunked_prefill_size=16, cuda_graph_config=graph_config
)
published = override.install()
override.install()
self.addCleanup(override.restore)
model_runner = SimpleNamespace(
server_args=SimpleNamespace(),