minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+12 -7
View File
@@ -815,7 +815,7 @@ def build_prefill_registry(
source: Optional[Any] = None,
) -> CudaGraphBufferRegistry:
"""Registry mirroring the **token-axis** FB-shared buffers for the
piecewise / breakable (prefill) cuda-graph runners.
piecewise / breakable / full (prefill) cuda-graph runners.
``register_input_embeds`` (default ``True``) registers the multimodal
``input_embeds`` slot; the eager extend path passes ``False`` so it is
@@ -910,13 +910,18 @@ def build_prefill_registry(
# blank real tokens whenever raw < bucket. Recompute the local
# count against the padded bucket from the batch's un-adjusted
# global count, mirroring the decode registry's post_fill.
if require_gathered_buffer and not enable_prefill_cp:
buf.fill_(
compute_local_num_token_non_padded_cpu(
global_num_token_non_padded=fb.num_token_non_padded_cpu,
num_tokens_per_dp=ctx.padded_num_tokens,
if require_gathered_buffer:
if not enable_prefill_cp:
buf.fill_(
compute_local_num_token_non_padded_cpu(
global_num_token_non_padded=fb.num_token_non_padded_cpu,
num_tokens_per_dp=ctx.padded_num_tokens,
)
)
)
else:
# Non-gathered FullCG still needs the live boundary rather
# than a stale/absent ForwardBatch tensor.
buf.fill_(ctx.raw_num_tokens)
slots.append(
GraphSlot(
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+4 -4
View File
@@ -322,14 +322,14 @@ def capture_prefill_graph(
# Skip prefill CG for EAGLE target on tc_piecewise when the fixed server
# capture ceiling is below FULL. EAGLE target prefill requests FULL, so a
# NULL or LAST graph is dead; capturing it can perturb FP4/TRTLLM-MoE
# state and corrupt decode replay (see #28386 and #28870). BCG captures
# FULL for EAGLE target in PrefillCudaGraphRunner.__init__, so it does not
# need this skip.
# state and corrupt decode replay (see #28386 and #28870). BCG and FullCG
# capture FULL for EAGLE targets in PrefillCudaGraphRunner.__init__, so
# they do not need this skip.
if (
model_runner.spec_algorithm.is_eagle()
and not model_runner.is_draft_worker
and get_server_return_hidden_states_mode() < CaptureHiddenMode.FULL
and not check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
and check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
):
logger.info(
"Disable prefill CUDA graph for EAGLE target on tc_piecewise "
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+15 -6
View File
@@ -296,16 +296,20 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.capture_forward_mode = ForwardMode.EXTEND
# Hidden-state capture mode cases:
# - Breakable EAGLE draft: LAST.
# - Breakable EAGLE target: FULL.
# - EAGLE target: FULL.
# - Return-hidden-states or DFLASH: FULL.
# - Otherwise: NULL.
is_breakable_eagle = (
is_eagle = model_runner.spec_algorithm.is_eagle()
is_breakable_eagle_draft = (
self.prefill_backend_name == Backend.BREAKABLE
and model_runner.spec_algorithm.is_eagle()
and is_eagle
and model_runner.is_draft_worker
)
if is_breakable_eagle and model_runner.is_draft_worker:
if is_breakable_eagle_draft:
self.capture_hidden_mode = CaptureHiddenMode.LAST
elif is_breakable_eagle or model_runner.spec_algorithm.is_dflash_family():
elif (is_eagle and not model_runner.is_draft_worker) or (
model_runner.spec_algorithm.is_dflash_family()
):
self.capture_hidden_mode = CaptureHiddenMode.FULL
else:
self.capture_hidden_mode = self.return_hidden_states_mode
@@ -350,7 +354,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
hidden_size=input_embeds_hidden_size,
embed_dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
enable_num_token_non_padded=enable_num_token_non_padded(),
# FullCG always pads to a capture bucket. Models that mask padded
# hidden rows need the live boundary even without expert parallelism.
enable_num_token_non_padded=(
enable_num_token_non_padded()
or self.prefill_backend_name == Backend.FULL
),
require_gathered_buffer=require_gathered_buffer(),
enable_prefill_cp=(
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
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(),