Support piecewise cuda graph for MLA (#11812)

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,

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