From 9a9e1671794eb803aabe2867be85ac0d83834e59 Mon Sep 17 00:00:00 2001 From: Aurick Qiao Date: Sun, 30 Aug 2026 21:30:24 -0700 Subject: [PATCH] [Bugfix] Fix full prefill CUDA graph padding and EAGLE capture (#35588) --- .../cuda_graph_buffer_registry.py | 19 +++-- .../cuda_graph_setup.py | 8 +-- .../runner/prefill_cuda_graph_runner.py | 21 ++++-- .../test_full_cuda_graph_prefill.py | 72 +++++++++++++++++-- ..._unified_radix_cache_kl_hybrid_bitexact.py | 33 ++++++++- .../test_cuda_graph_buffer_registry.py | 9 ++- .../test_prefill_cuda_graph_runner.py | 51 +++++++++---- 7 files changed, 171 insertions(+), 42 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 484d3bd4e..c07654b69 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -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( diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 2b46f86af..adc24df07 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -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 " diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index d7787a4c1..59e8568c7 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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() diff --git a/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py b/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py index 9001da960..2529f940b 100644 --- a/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py +++ b/test/registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py @@ -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}}' + ), ], ) diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py index 36cea24ed..7f4ef8b40 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py @@ -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) diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 51b173dcd..1cf9d42a7 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py @@ -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, ) diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index a26da03ba..4806aaa4c 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py @@ -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(),