From db24d34603d3797ece082f0b8f7e73e16c65f22a Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Mon, 10 Nov 2025 09:13:48 +0800 Subject: [PATCH] Support piecewise cuda graph for MLA (#11812) --- python/sglang/srt/compilation/compile.py | 3 + .../compilation/piecewise_context_manager.py | 3 +- .../attention/flashinfer_mla_backend.py | 5 + python/sglang/srt/layers/radix_attention.py | 26 +++- python/sglang/srt/layers/rotary_embedding.py | 1 - .../sglang/srt/model_executor/model_runner.py | 23 ++- .../piecewise_cuda_graph_runner.py | 141 ++++++++++-------- python/sglang/srt/models/deepseek_v2.py | 41 +++-- test/srt/run_suite.py | 2 +- test/srt/test_piecewise_cuda_graph.py | 40 +++++ 10 files changed, 196 insertions(+), 89 deletions(-) diff --git a/python/sglang/srt/compilation/compile.py b/python/sglang/srt/compilation/compile.py index a77f5aee7..e0ec9a1f3 100644 --- a/python/sglang/srt/compilation/compile.py +++ b/python/sglang/srt/compilation/compile.py @@ -112,6 +112,9 @@ def _infer_dynamic_arg_dims_from_annotations(forward_fn): for a in getattr(ann, "__args__", []) ): dyn[name] = 0 + elif ann == "torch.Tensor" or ann == "Optional[torch.Tensor]": + # For future import annotations (e.g. from __future__ import annotations), the annotation is a string + dyn[name] = 0 if not dyn: raise ValueError("No dynamic dims inferred; pass dynamic_arg_dims explicitly.") return dyn diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index 38d17a6df..55884c0d7 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/python/sglang/srt/compilation/piecewise_context_manager.py @@ -30,11 +30,10 @@ def get_forward_context() -> Optional[ForwardContext]: @contextmanager def set_forward_context(forward_batch: ForwardBatch, attention_layers: List[Any]): global _forward_context - prev_forward_context = _forward_context _forward_context = ForwardContext() _forward_context.set_forward_batch(forward_batch) _forward_context.set_attention_layers(attention_layers) try: yield finally: - _forward_context = prev_forward_context + _forward_context = None diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 7a4376379..348b3d3f5 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -22,6 +22,9 @@ from sglang.srt.layers.attention.flashinfer_backend import ( ) from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( + is_in_piecewise_cuda_graph, +) from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import ( @@ -322,6 +325,8 @@ class FlashInferMLAAttnBackend(AttentionBackend): use_ragged = ( not get_global_server_args().flashinfer_mla_disable_ragged and extend_no_prefix + # Piecewise cuda graph should use paged prefill to be compatible with prefix cache + and not is_in_piecewise_cuda_graph() ) self.indices_updater_prefill.update( diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 79a324302..4d22f1887 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -110,9 +110,12 @@ class RadixAttention(nn.Module): k = k.view(-1, self.tp_k_head_num, self.v_head_dim) if forward_batch.forward_mode.is_extend() and get_forward_context() is not None: - output = torch.empty_like(q) + if self.qk_head_dim != self.v_head_dim: + output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim)) + else: + output = torch.empty_like(q) torch.ops.sglang.unified_attention_with_output( - q, k, v, output, save_kv_cache, self.layer_id + q, k, v, output, save_kv_cache, self.layer_id, **kwargs ) return output else: @@ -134,13 +137,26 @@ def unified_attention_with_output( output: torch.Tensor, save_kv_cache: bool, layer_id: int, + *, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, ) -> None: context = get_forward_context() forward_batch = context.forward_batch attention_layers = context.attention_layers attention_layer = attention_layers[layer_id] + + kwargs = {} + if q_rope is not None: + kwargs["q_rope"] = q_rope + if k_rope is not None: + kwargs["k_rope"] = k_rope + if sinks is not None: + kwargs["sinks"] = sinks + ret = forward_batch.attn_backend.forward( - query, key, value, attention_layer, forward_batch, save_kv_cache + query, key, value, attention_layer, forward_batch, save_kv_cache, **kwargs ) assert ( output.numel() == ret.numel() @@ -157,6 +173,10 @@ def unified_attention_with_output_fake( output: torch.Tensor, save_kv_cache: bool, layer_id: int, + *, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, ) -> None: return diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index a2e8e60b2..51981da81 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -823,7 +823,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): query_pass = query[..., self.rotary_dim :] key_pass = key[..., self.rotary_dim :] - self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device) cos_sin = self.cos_sin_cache[ torch.add(positions, offsets) if offsets is not None else positions ] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6e3a2ef6c..8433b0247 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -337,8 +337,13 @@ class ModelRunner: ): self.attention_layers = [] for layer in self.model.model.layers: - if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "attn"): - self.attention_layers.append(layer.self_attn.attn) + if hasattr(layer, "self_attn"): + if hasattr(layer.self_attn, "attn"): + self.attention_layers.append(layer.self_attn.attn) + elif hasattr(layer.self_attn, "attn_mqa"): + # For DeepSeek model + self.attention_layers.append(layer.self_attn.attn_mqa) + if len(self.attention_layers) < self.model_config.num_hidden_layers: # TODO(yuwei): support Non-Standard GQA log_info_on_rank0( @@ -2075,9 +2080,6 @@ class ModelRunner: skip_attn_backend_init: bool = False, pp_proxy_tensors=None, ) -> LogitsProcessorOutput: - if not skip_attn_backend_init: - self.attn_backend.init_forward_metadata(forward_batch) - kwargs = {} if self.support_pp: kwargs["pp_proxy_tensors"] = pp_proxy_tensors @@ -2086,9 +2088,14 @@ class ModelRunner: if not self.is_generation: kwargs["get_embedding"] = True - if self.piecewise_cuda_graph_runner is not None: - if self.piecewise_cuda_graph_runner.can_run(forward_batch): - return self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs) + if ( + self.piecewise_cuda_graph_runner is not None + and self.piecewise_cuda_graph_runner.can_run(forward_batch) + ): + return self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs) + + if not skip_attn_backend_init: + self.attn_backend.init_forward_metadata(forward_batch) return self.model.forward( forward_batch.input_ids, 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 736242b01..2a8b4ae21 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -55,22 +55,21 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner -# Detect whether the current forward pass is in capture mode -is_capture_mode = False +_in_piecewise_cuda_graph = False -def get_is_capture_mode(): - return is_capture_mode +def is_in_piecewise_cuda_graph(): + return _in_piecewise_cuda_graph @contextmanager -def model_capture_mode(): - global is_capture_mode - is_capture_mode = True +def enable_piecewise_cuda_graph(): + global _in_piecewise_cuda_graph + _in_piecewise_cuda_graph = True yield - is_capture_mode = False + _in_piecewise_cuda_graph = False @contextmanager @@ -125,6 +124,15 @@ def set_global_graph_memory_pool(val): global_graph_memory_pool = val +def set_torch_compile_config(): + import torch._dynamo.config + + # Resolve torch._dynamo.exc.FailOnRecompileLimitHit + torch._dynamo.config.accumulated_cache_size_limit = 1024 + if hasattr(torch._dynamo.config, "cache_size_limit"): + torch._dynamo.config.cache_size_limit = 1024 + + class PiecewiseCudaGraphRunner: """A PiecewiseCudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile.""" @@ -142,6 +150,8 @@ class PiecewiseCudaGraphRunner: self.attn_tp_size = get_attention_tp_size() self.attn_tp_rank = get_attention_tp_rank() + set_torch_compile_config() + assert ( self.model_runner.server_args.piecewise_cuda_graph_tokens is not None ), "piecewise_cuda_graph_tokens is not set" @@ -166,7 +176,6 @@ class PiecewiseCudaGraphRunner: if model_runner.server_args.enable_return_hidden_states: self.capture_hidden_mode = CaptureHiddenMode.FULL - # Attention backend self.max_num_tokens = max(self.capture_num_tokens) # Graph inputs @@ -185,29 +194,29 @@ class PiecewiseCudaGraphRunner: # Set graph pool id globally to be able to use symmetric memory set_graph_pool_id(get_global_graph_memory_pool()) - with patch_model( - self.model_runner.model.model, self.compile_config.compiler - ) as patched_model: - install_torch_compiled( - patched_model, - fullgraph=True, - dynamic_arg_dims=None, - compile_config=self.compile_config, - graph_pool=get_global_graph_memory_pool(), - ) - - with set_compiled(True): - self.warmup_and_capture() - - # Capture - try: - with model_capture_mode(): - self.capture() - except RuntimeError as e: - raise Exception( - f"Capture cuda graph failed: {e}\n{PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG}" + with enable_piecewise_cuda_graph(): + with patch_model( + self.model_runner.model.model, self.compile_config.compiler + ) as patched_model: + install_torch_compiled( + patched_model, + fullgraph=True, + dynamic_arg_dims=None, + compile_config=self.compile_config, + graph_pool=get_global_graph_memory_pool(), ) + with set_compiled(True): + self.warmup_and_capture() + + # Capture + try: + self.capture() + except RuntimeError as e: + raise Exception( + f"Capture cuda graph failed: {e}\n{PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG}" + ) + self.raw_num_tokens = 0 def warmup_and_capture(self): @@ -225,7 +234,9 @@ class PiecewiseCudaGraphRunner: req_to_token_pool=self.model_runner.req_to_token_pool, token_to_kv_pool=self.model_runner.token_to_kv_pool, attn_backend=self.model_runner.attn_backend, - out_cache_loc=torch.randint(0, 100, (num_tokens,), device=self.device), + out_cache_loc=torch.zeros( + (num_tokens,), device=self.device, dtype=self._cache_loc_dtype() + ), seq_lens_sum=num_tokens, encoder_lens=None, return_logprob=False, @@ -378,6 +389,8 @@ class PiecewiseCudaGraphRunner: if lora_ids is not None: self.model_runner.lora_manager.prepare_lora_batch(forward_batch) + self.model_runner.attn_backend.init_forward_metadata(forward_batch) + # Run and capture def run_once(): # Clean intermediate result cache for DP attention @@ -401,7 +414,7 @@ class PiecewiseCudaGraphRunner: ) return - for _ in range(2): + for _ in range(3): self.device_module.synchronize() self.model_runner.tp_group.barrier() run_once() @@ -483,36 +496,40 @@ class PiecewiseCudaGraphRunner: forward_batch: ForwardBatch, **kwargs, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: - if self.model_runner.tp_group.ca_comm is not None: - old_ca_disable = self.model_runner.tp_group.ca_comm.disabled - self.model_runner.tp_group.ca_comm.disabled = True - static_forward_batch = self.replay_prepare(forward_batch, **kwargs) - # Replay - with set_forward_context(static_forward_batch, self.attention_layers): - with set_compiled(True): - output = self.model_runner.model.forward( - static_forward_batch.input_ids, - static_forward_batch.positions, - static_forward_batch, - **kwargs, - ) - if isinstance(output, LogitsProcessorOutput): - return LogitsProcessorOutput( - next_token_logits=output.next_token_logits[: self.raw_num_tokens], - hidden_states=( - output.hidden_states[: self.raw_num_tokens] - if output.hidden_states is not None - else None - ), - ) - else: - assert isinstance(output, PPProxyTensors) - # TODO(Yuwei): support PP Support - raise NotImplementedError( - "PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet." - ) - if self.model_runner.tp_group.ca_comm is not None: - self.model_runner.tp_group.ca_comm.disabled = old_ca_disable + with enable_piecewise_cuda_graph(): + if self.model_runner.tp_group.ca_comm is not None: + old_ca_disable = self.model_runner.tp_group.ca_comm.disabled + self.model_runner.tp_group.ca_comm.disabled = True + self.model_runner.attn_backend.init_forward_metadata(forward_batch) + static_forward_batch = self.replay_prepare(forward_batch, **kwargs) + # Replay + with set_forward_context(static_forward_batch, self.attention_layers): + with set_compiled(True): + output = self.model_runner.model.forward( + static_forward_batch.input_ids, + static_forward_batch.positions, + static_forward_batch, + **kwargs, + ) + if isinstance(output, LogitsProcessorOutput): + return LogitsProcessorOutput( + next_token_logits=output.next_token_logits[ + : self.raw_num_tokens + ], + hidden_states=( + output.hidden_states[: self.raw_num_tokens] + if output.hidden_states is not None + else None + ), + ) + else: + assert isinstance(output, PPProxyTensors) + # TODO(Yuwei): support PP Support + raise NotImplementedError( + "PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet." + ) + if self.model_runner.tp_group.ca_comm is not None: + self.model_runner.tp_group.ca_comm.disabled = old_ca_disable def get_spec_info(self, num_tokens: int): spec_info = None diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 28dd58483..105bbf5d7 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -112,6 +112,9 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( + is_in_piecewise_cuda_graph, +) from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args @@ -320,6 +323,9 @@ def _support_mha_one_shot(attn: DeepseekV2AttentionMLA, forward_batch, backend_n def _handle_attention_backend( attn: DeepseekV2AttentionMLA, forward_batch, backend_name ): + if is_in_piecewise_cuda_graph(): + return AttnForwardMethod.MLA + sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch) disable_ragged = ( backend_name in ["flashinfer", "flashmla"] @@ -427,6 +433,9 @@ def handle_attention_nsa(attn, forward_batch): def handle_attention_triton(attn, forward_batch): + if is_in_piecewise_cuda_graph(): + return AttnForwardMethod.MLA + # when deterministic inference is enabled, use MLA if get_global_server_args().enable_deterministic_inference: return _dispatch_mla_subtype(attn, forward_batch) @@ -1841,18 +1850,26 @@ class DeepseekV2AttentionMLA(nn.Module): ) attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) else: - attn_bmm_output = torch.empty( - (attn_output.shape[0], self.num_local_heads * self.v_head_dim), - dtype=attn_output.dtype, - device=attn_output.device, - ) - torch.bmm( - attn_output.transpose(0, 1), - self.w_vc, - out=attn_bmm_output.view( - -1, self.num_local_heads, self.v_head_dim - ).transpose(0, 1), - ) + if is_in_piecewise_cuda_graph(): + # torch dynamo requires out= op was called where output tensor was non-contiguous + attn_bmm_output = ( + torch.bmm(attn_output.transpose(0, 1), self.w_vc) + .transpose(0, 1) + .flatten(1, 2) + ) + else: + attn_bmm_output = torch.empty( + (attn_output.shape[0], self.num_local_heads * self.v_head_dim), + dtype=attn_output.dtype, + device=attn_output.device, + ) + torch.bmm( + attn_output.transpose(0, 1), + self.w_vc, + out=attn_bmm_output.view( + -1, self.num_local_heads, self.v_head_dim + ).transpose(0, 1), + ) output, _ = self.o_proj(attn_bmm_output) return output diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index befce0872..345a718c3 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -103,7 +103,7 @@ suites = { TestFile("test_original_logprobs.py", 41), TestFile("test_page_size.py", 60), TestFile("test_penalty.py", 82), - TestFile("test_piecewise_cuda_graph.py", 180), + TestFile("test_piecewise_cuda_graph.py", 300), TestFile("test_priority_scheduling.py", 130), TestFile("test_pytorch_sampling_backend.py", 66), TestFile("test_radix_attention.py", 105), diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index d50c67453..51b92e93f 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/test/srt/test_piecewise_cuda_graph.py @@ -1,9 +1,11 @@ import unittest from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_MODEL_NAME_FOR_TEST_MLA, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, @@ -92,5 +94,43 @@ class TestPiecewiseCudaGraphQwen3MoE(CustomTestCase): self.assertGreaterEqual(metrics["score"], 0.90) +class TestPiecewiseCudaGraphDeepSeek(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--enable-piecewise-cuda-graph", + "--piecewise-cuda-graph-compiler", + "eager", + "--piecewise-cuda-graph-max-tokens", + "4096", # should less than max_context_len + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(metrics) + + self.assertGreater(metrics["accuracy"], 0.62) + + if __name__ == "__main__": unittest.main()