Support piecewise cuda graph for MLA (#11812)

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

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

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(

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
 

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
         ]

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"):
+                if hasattr(layer, "self_attn"):
-                    self.attention_layers.append(layer.self_attn.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 (
-            if self.piecewise_cuda_graph_runner.can_run(forward_batch):
+            self.piecewise_cuda_graph_runner is not None
-                return self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
+            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
+_in_piecewise_cuda_graph = False
-is_capture_mode = False
 
 
-def get_is_capture_mode():
+def is_in_piecewise_cuda_graph():
-    return is_capture_mode
+    return _in_piecewise_cuda_graph
 
 
 @contextmanager
-def model_capture_mode():
+def enable_piecewise_cuda_graph():
-    global is_capture_mode
+    global _in_piecewise_cuda_graph
-    is_capture_mode = True
+    _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(
+        with enable_piecewise_cuda_graph():
-            self.model_runner.model.model, self.compile_config.compiler
+            with patch_model(
-        ) as patched_model:
+                self.model_runner.model.model, self.compile_config.compiler
-            install_torch_compiled(
+            ) as patched_model:
-                patched_model,
+                install_torch_compiled(
-                fullgraph=True,
+                    patched_model,
-                dynamic_arg_dims=None,
+                    fullgraph=True,
-                compile_config=self.compile_config,
+                    dynamic_arg_dims=None,
-                graph_pool=get_global_graph_memory_pool(),
+                    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 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:
+        with enable_piecewise_cuda_graph():
-            old_ca_disable = self.model_runner.tp_group.ca_comm.disabled
+            if self.model_runner.tp_group.ca_comm is not None:
-            self.model_runner.tp_group.ca_comm.disabled = True
+                old_ca_disable = self.model_runner.tp_group.ca_comm.disabled
-        static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
+                self.model_runner.tp_group.ca_comm.disabled = True
-        # Replay
+            self.model_runner.attn_backend.init_forward_metadata(forward_batch)
-        with set_forward_context(static_forward_batch, self.attention_layers):
+            static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
-            with set_compiled(True):
+            # Replay
-                output = self.model_runner.model.forward(
+            with set_forward_context(static_forward_batch, self.attention_layers):
-                    static_forward_batch.input_ids,
+                with set_compiled(True):
-                    static_forward_batch.positions,
+                    output = self.model_runner.model.forward(
-                    static_forward_batch,
+                        static_forward_batch.input_ids,
-                    **kwargs,
+                        static_forward_batch.positions,
-                )
+                        static_forward_batch,
-            if isinstance(output, LogitsProcessorOutput):
+                        **kwargs,
-                return LogitsProcessorOutput(
+                    )
-                    next_token_logits=output.next_token_logits[: self.raw_num_tokens],
+                if isinstance(output, LogitsProcessorOutput):
-                    hidden_states=(
+                    return LogitsProcessorOutput(
-                        output.hidden_states[: self.raw_num_tokens]
+                        next_token_logits=output.next_token_logits[
-                        if output.hidden_states is not None
+                            : self.raw_num_tokens
-                        else None
+                        ],
-                    ),
+                        hidden_states=(
-                )
+                            output.hidden_states[: self.raw_num_tokens]
-            else:
+                            if output.hidden_states is not None
-                assert isinstance(output, PPProxyTensors)
+                            else None
-                # TODO(Yuwei): support PP Support
+                        ),
-                raise NotImplementedError(
+                    )
-                    "PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet."
+                else:
-                )
+                    assert isinstance(output, PPProxyTensors)
-        if self.model_runner.tp_group.ca_comm is not None:
+                    # TODO(Yuwei): support PP Support
-            self.model_runner.tp_group.ca_comm.disabled = old_ca_disable
+                    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

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(
+            if is_in_piecewise_cuda_graph():
-                (attn_output.shape[0], self.num_local_heads * self.v_head_dim),
+                # torch dynamo requires out= op was called where output tensor was non-contiguous
-                dtype=attn_output.dtype,
+                attn_bmm_output = (
-                device=attn_output.device,
+                    torch.bmm(attn_output.transpose(0, 1), self.w_vc)
-            )
+                    .transpose(0, 1)
-            torch.bmm(
+                    .flatten(1, 2)
-                attn_output.transpose(0, 1),
+                )
-                self.w_vc,
+            else:
-                out=attn_bmm_output.view(
+                attn_bmm_output = torch.empty(
-                    -1, self.num_local_heads, self.v_head_dim
+                    (attn_output.shape[0], self.num_local_heads * self.v_head_dim),
-                ).transpose(0, 1),
+                    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

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),

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()