diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index c567a6d6f..6f8a102fa 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -217,8 +217,12 @@ class PiecewiseCudaGraphRunner: self.capture_forward_mode = ForwardMode.EXTEND self.capture_hidden_mode = CaptureHiddenMode.NULL - # If returning hidden states is enabled, set initial capture hidden mode to full to avoid double-capture on startup - if model_runner.server_args.enable_return_hidden_states: + # If returning hidden states is enabled, or if speculative prefill needs + # aux hidden states (DFLASH), capture the FULL variant up front. + if ( + model_runner.server_args.enable_return_hidden_states + or model_runner.spec_algorithm.is_dflash() + ): self.capture_hidden_mode = CaptureHiddenMode.FULL self.max_num_tokens = ( @@ -416,7 +420,7 @@ class PiecewiseCudaGraphRunner: mrope_positions=mrope_positions, spec_algorithm=None, spec_info=None, - capture_hidden_mode=CaptureHiddenMode.NULL, + capture_hidden_mode=self.capture_hidden_mode, num_token_non_padded=None, num_token_non_padded_cpu=num_tokens, global_forward_mode=ForwardMode.EXTEND, @@ -587,7 +591,7 @@ class PiecewiseCudaGraphRunner: mrope_positions=mrope_positions, spec_algorithm=None, spec_info=None, - capture_hidden_mode=CaptureHiddenMode.NULL, + capture_hidden_mode=self.capture_hidden_mode, num_token_non_padded=None, num_token_non_padded_cpu=num_tokens, global_forward_mode=ForwardMode.EXTEND, diff --git a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_dflash.py b/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_dflash.py new file mode 100644 index 000000000..3095aacca --- /dev/null +++ b/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_dflash.py @@ -0,0 +1,47 @@ +"""Test piecewise CUDA graph coexisting with speculative decoding (DFLASH). + +PCG handles prefill/extend path while DFlash needs target aux hidden states +from prefill to materialize draft KV cache. This verifies PCG captures that +path with the DFlash hidden-state variant enabled. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.server_fixtures.pcg_spec_fixture import PCGSpecBase +from sglang.test.test_utils import ( + DEFAULT_DRAFT_MODEL_DFLASH, + DEFAULT_TARGET_MODEL_DFLASH, + CustomTestCase, +) + +register_cuda_ci(est_time=531, stage="base-b", runner_config="1-gpu-small") + + +class TestPCGWithDFlash(PCGSpecBase, CustomTestCase): + """PCG + DFLASH on Llama-3.1-8B-Instruct.""" + + model = DEFAULT_TARGET_MODEL_DFLASH + server_args = [ + "--trust-remote-code", + "--attention-backend", + "flashinfer", + "--enforce-piecewise-cuda-graph", + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + DEFAULT_DRAFT_MODEL_DFLASH, + "--page-size", + "1", + "--max-running-requests", + "64", + "--cuda-graph-bs", + *[str(i) for i in range(1, 65)], + ] + server_env = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"} + accuracy_threshold = 0.75 + speedup_threshold = 2.8 + + +if __name__ == "__main__": + unittest.main()