[BCG][GLM5] perf: BCG support and prefill enhancements (#27053)

python/sglang/srt/environ.py

@@ -594,6 +594,7 @@ class Envs:
         False, deprecated_name="SGLANG_NSA_HIP_DISABLE_PRESHUFFLE"
     )
     SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)
+    SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False)
     SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
     SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
 

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -12,16 +12,24 @@ from sglang.jit_kernel.fused_store_index_cache import (
     can_use_dsa_fused_store,
     fused_store_index_k_cache,
 )
+from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import (
     aiter_can_use_preshuffle_paged_mqa,
     is_dsa_enable_prefill_cp,
     is_dsa_prefill_cp_in_seq_split,
+    is_graph_dsa_split_op_surface,
 )
 from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
 from sglang.srt.layers.layernorm import LayerNorm
 from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
 from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     get_tc_piecewise_forward_context,
     is_in_tc_piecewise_cuda_graph,
@@ -39,6 +47,7 @@ from sglang.srt.utils import (
     is_hip,
     is_npu,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 logger = logging.getLogger(__name__)
 
@@ -97,6 +106,16 @@ if TYPE_CHECKING:
 
 
 DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0
+GRAPH_WEIGHTS_PROJ_LORA_ERROR = (
+    "DSA indexer weights_proj LoRA is incompatible with "
+    "piecewise/breakable CUDA graph; remove the explicit "
+    "prefill cuda-graph backend override or drop "
+    "indexer.weights_proj from the LoRA target modules."
+)
+
+
+def _is_in_piecewise_or_breakable_cuda_graph() -> bool:
+    return is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph()
 
 
 def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool:
@@ -128,52 +147,8 @@ def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool:
 
 
 if _is_cuda:
-    from sglang.srt.compilation.compilation_config import register_split_op
-    from sglang.srt.utils.custom_op import register_custom_op
 
-    @register_custom_op(mutates_args=["topk_result"])
+    def _logits_head_gate_graph_fake_impl(
-    @register_split_op()
-    def k_cache_and_topk_result(
-        layer_id: int,
-        key: torch.Tensor,
-        q_fp8: torch.Tensor,
-        weights: torch.Tensor,
-        topk_result: torch.Tensor,
-    ) -> None:
-        assert (
-            _is_cuda
-        ), "Internal error: piecewise CUDA graph is only supported on CUDA"
-        from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
-
-        forward_batch = get_tc_piecewise_forward_context().forward_batch
-        indexer = get_tc_piecewise_forward_context().dsa_indexers[layer_id]
-        metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
-        assert metadata is not None, (
-            "DSA piecewise CUDA graph requires indexer metadata from the DSA "
-            "attention backend"
-        )
-
-        # slice off padding from piecewise CUDA graph
-        extend_num_tokens = forward_batch.extend_num_tokens
-
-        indexer._store_index_k_cache(
-            forward_batch=forward_batch,
-            layer_id=layer_id,
-            key=key[:extend_num_tokens],
-            act_quant=act_quant,
-            out_cache_loc=forward_batch.out_cache_loc[:extend_num_tokens],
-        )
-        indexer._get_topk_ragged(
-            False,
-            forward_batch,
-            layer_id,
-            q_fp8[:extend_num_tokens],
-            weights,
-            metadata,
-            topk_result,
-        )
-
-    def _logits_head_gate_pcg_fake_impl(
         x: torch.Tensor,
         weight: torch.Tensor,
         n_heads_inv_sqrt: float,
@@ -186,8 +161,9 @@ if _is_cuda:
             device=x.device,
         )
 
-    @register_custom_op(fake_impl=_logits_head_gate_pcg_fake_impl)
+    # In-graph (PCG/BCG) head gate for the NON-prefill path
-    def logits_head_gate_pcg(
+    @register_custom_op(fake_impl=_logits_head_gate_graph_fake_impl)
+    def logits_head_gate_graph(
         x: torch.Tensor,
         weight: torch.Tensor,
         n_heads_inv_sqrt: float,
@@ -467,6 +443,15 @@ class Indexer(MultiPlatformOp):
     ):
         return weights.unsqueeze(-1) * q_scale * self.softmax_scale
 
+    def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
+        if (
+            forward_batch.forward_mode.is_extend_without_speculative()
+            and forward_batch.seq_lens_cpu is not None
+        ):
+            max_kv_len = forward_batch.seq_lens_cpu.max().item()
+            return max_kv_len <= self.index_topk
+        return False
+
     def _get_q_k_bf16(
         self,
         q_lora: torch.Tensor,
@@ -986,21 +971,36 @@ class Indexer(MultiPlatformOp):
         enable_dual_stream: bool,
         metadata: BaseIndexerMetadata,
         return_indices: bool = True,
+        *,
+        num_tokens: Optional[int] = None,
+        topk_result: Optional[torch.Tensor] = None,
     ) -> Optional[torch.Tensor]:
+        # Shared by the eager path and the graph DSA split-op dispatch. The two
+        # keyword args carry the graph contract and default to the eager behavior:
+        #   - num_tokens: slice key/out_cache_loc to the unpadded count (the graph
+        #     runs at a static padded shape). None => full (eager) shape.
+        #   - topk_result: pre-allocated padded buffer to fill in place (a downstream
+        #     captured graph reads it at a fixed address). None => return a fresh,
+        #     naturally-sized tensor.
         assert forward_batch.forward_mode.is_extend_without_speculative()
         x_meta = x[0] if isinstance(x, tuple) else x
 
         # Fast path: only compute and store k cache, skip all q and weights ops
         key = self._get_k_bf16(x, positions, enable_dual_stream)
-
+        out_cache_loc = None
-        if not forward_batch.out_cache_loc.is_contiguous():
+        if num_tokens is not None:
+            assert num_tokens <= key.shape[0]
+            assert num_tokens <= forward_batch.out_cache_loc.shape[0]
+            key = key[:num_tokens]
+            out_cache_loc = forward_batch.out_cache_loc[:num_tokens]
+        elif not forward_batch.out_cache_loc.is_contiguous():
             forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
-
         self._store_index_k_cache(
             forward_batch=forward_batch,
             layer_id=layer_id,
             key=key,
             act_quant=act_quant,
+            out_cache_loc=out_cache_loc,
         )
 
         # MHA doesn't need topk_indices
@@ -1016,7 +1016,13 @@ class Indexer(MultiPlatformOp):
             dtype=torch.float32,
             device=x_meta.device,
         )
-        return metadata.topk_transform(dummy_logits, self.index_topk)
+        raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk)
+        if topk_result is not None:
+            # PCG/BCG: fill the valid prefix of the padded static buffer and
+            # leave padded rows at the -1 sentinel.
+            topk_result[: raw_topk_result.shape[0]] = raw_topk_result
+            return None
+        return raw_topk_result
 
     def _get_topk_ragged_with_cp(
         self,
@@ -1029,9 +1035,10 @@ class Indexer(MultiPlatformOp):
         actual_seq_q: int,
         cp_index: List[Tuple[int, int, int]] = None,
     ) -> torch.Tensor:
-        assert (
+        assert not _is_in_piecewise_or_breakable_cuda_graph(), (
-            not is_in_tc_piecewise_cuda_graph()
+            "DSA context parallel (_get_topk_ragged_with_cp) not supported under "
-        ), "DSA context parallel (_get_topk_ragged_with_cp) not supported under piecewise CUDA graph"
+            "piecewise/breakable CUDA graph"
+        )
         if TYPE_CHECKING:
             assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
 
@@ -1178,9 +1185,10 @@ class Indexer(MultiPlatformOp):
         topk: int,
         layer_id: int,
     ) -> Optional[torch.Tensor]:
-        assert (
+        assert not _is_in_piecewise_or_breakable_cuda_graph(), (
-            not is_in_tc_piecewise_cuda_graph()
+            "DSA forward_indexer (non-CUDA loop path) not supported under "
-        ), "DSA forward_indexer (non-CUDA loop path) not supported under piecewise CUDA graph"
+            "piecewise/breakable CUDA graph"
+        )
         if not _is_npu:
             from sglang.srt.layers.attention.dsa.tilelang_kernel import fp8_index
 
@@ -1370,10 +1378,15 @@ class Indexer(MultiPlatformOp):
         # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries.
         x_meta = x[0] if isinstance(x, tuple) else x
 
-        # In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_tc_piecewise_forward_context() to
+        in_piecewise_or_breakable_cuda_graph = (
-        # prevent Dynamo from guarding on forward_metadata identity (which changes each
+            _is_in_piecewise_or_breakable_cuda_graph()
-        # replay when init_forward_metadata creates a new ForwardMetadata object).
+        )
-        if not is_in_tc_piecewise_cuda_graph():
+
+        # In piecewise/breakable CUDA graph mode, metadata is fetched inside
+        # custom ops via get_tc_piecewise_forward_context() to prevent Dynamo
+        # from guarding on forward_metadata identity, which changes each replay
+        # when init_forward_metadata creates a new ForwardMetadata object.
+        if not in_piecewise_or_breakable_cuda_graph:
             metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
             if metadata is None:
                 return None
@@ -1390,13 +1403,10 @@ class Indexer(MultiPlatformOp):
         # Determine if should skip topk based on sequence length
         # We can only skip the logits computation if cuda graph is not involved
         skip_logits_computation = False
-        if (
+        if not in_piecewise_or_breakable_cuda_graph:
-            not is_in_tc_piecewise_cuda_graph()
+            skip_logits_computation = self._should_skip_logits_computation(
-            and forward_batch.forward_mode.is_extend_without_speculative()
+                forward_batch
-        ):
+            )
-            if forward_batch.seq_lens_cpu is not None:
-                max_kv_len = forward_batch.seq_lens_cpu.max().item()
-                skip_logits_computation = max_kv_len <= self.index_topk
 
         # Optimization: fast path when skipping topk computation
         if skip_logits_computation and (not self.dsa_enable_prefill_cp):
@@ -1417,6 +1427,43 @@ class Indexer(MultiPlatformOp):
         # wrapper owns base+delta and no LoRA kernel runs under torch.compile
         weights_proj_lora = getattr(self.weights_proj, "set_lora", False)
 
+        if (
+            is_graph_dsa_split_op_surface(forward_batch)
+            and not self.dsa_enable_prefill_cp
+        ):
+            # Default path for non-CP prefill under PCG/BCG: run the whole indexer
+            # (q/k proj, head gate, k-cache store, topk) as a single eager split op
+            # instead of capturing it piecemeal in the graph.
+            if weights_proj_lora:
+                raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR)
+            if return_indices:
+                topk_result = torch.full(
+                    (x.shape[0], self.index_topk),
+                    -1,
+                    device=x.device,
+                    dtype=torch.int32,
+                )
+            else:
+                topk_result = torch.empty(
+                    (0, self.index_topk), device=x.device, dtype=torch.int32
+                )
+            graph_dispatch_fn = (
+                bcg_dsa_indexer_prefill_split
+                if is_in_breakable_cuda_graph()
+                else pcg_dsa_indexer_prefill_split
+            )
+            graph_dispatch_fn(
+                layer_id=layer_id,
+                x=x,
+                q_lora=q_lora,
+                positions=positions,
+                topk_result=topk_result,
+            )
+            result = _broadcast_indexer_topk_from_rank0(
+                topk_result if return_indices else None
+            )
+            return maybe_capture_indexer_topk(layer_id, result)
+
         if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
             current_stream = torch.cuda.current_stream()
             self.alt_stream.wait_stream(current_stream)
@@ -1455,7 +1502,7 @@ class Indexer(MultiPlatformOp):
                         act_quant=act_quant,
                     )
                 current_stream.wait_stream(self.alt_stream)
-            elif not is_in_tc_piecewise_cuda_graph():
+            elif not in_piecewise_or_breakable_cuda_graph:
                 q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
                 self._store_index_k_cache(
                     forward_batch=forward_batch,
@@ -1464,8 +1511,10 @@ class Indexer(MultiPlatformOp):
                     act_quant=act_quant,
                 )
             else:
-                # piecewise CUDA graph need to split graph on store_k_cache and mqa_logits,
+                # Graph paths not handled by the full DSA indexer split op
-                # so delay store_k_cache after weights proj.
+                # still need q_fp8 for paged topk and q_scale for
+                # logits_head_gate_graph. K-cache storage is handled by the
+                # full graph split path when prefill requires it.
                 q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
 
             # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from
@@ -1509,13 +1558,10 @@ class Indexer(MultiPlatformOp):
             else:
                 x_for_gate = x
 
-            if is_in_tc_piecewise_cuda_graph():
+            if in_piecewise_or_breakable_cuda_graph:
                 if weights_proj_lora:
-                    raise RuntimeError(
+                    raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR)
-                        "DSA indexer weights_proj LoRA is incompatible with TC piecewise CUDA graph; remove the explicit"
+                weights = logits_head_gate_graph(
-                        " prefill cuda-graph backend override or drop indexer.weights_proj from the LoRA target modules."
-                    )
-                weights = logits_head_gate_pcg(
                     x_for_gate,
                     self.weights_proj.weight,
                     self.n_heads**-0.5,
@@ -1529,9 +1575,10 @@ class Indexer(MultiPlatformOp):
                 weights = self._get_logits_head_gate(x_for_gate, q_scale)
 
         if _is_cuda or _is_hip:
-            # In piecewise CUDA graph, any access to seq_lens_cpu creates a Dynamo shape guard.
+            # In piecewise/breakable CUDA graph, any access to seq_lens_cpu
-            # Piecewise CUDA graph never has empty batches.
+            # creates a Dynamo shape guard. These graph modes never have empty
-            if not is_in_tc_piecewise_cuda_graph():
+            # batches.
+            if not in_piecewise_or_breakable_cuda_graph:
                 assert forward_batch.seq_lens_cpu is not None
                 if len(forward_batch.seq_lens_cpu) == 0:
                     # this seems b/c max-pad, no worries?
@@ -1602,28 +1649,14 @@ class Indexer(MultiPlatformOp):
                     topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0)
                     topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
                     return maybe_capture_indexer_topk(layer_id, topk_result)
-                elif is_in_tc_piecewise_cuda_graph():
-                    assert (
-                        not enable_dual_stream
-                    ), "Internal error: piecewise CUDA graph should not be enabled with dual stream"
-
-                    if not _uses_dsa_attention_backend(forward_batch):
-                        return None
-
-                    topk_result = torch.full(
-                        (q_fp8.shape[0], self.index_topk),
-                        -1,
-                        device=q_fp8.device,
-                        dtype=torch.int32,
-                    )
-                    k_cache_and_topk_result(
-                        layer_id=layer_id,
-                        key=key,
-                        q_fp8=q_fp8,
-                        weights=weights,
-                        topk_result=topk_result,
-                    )
                 else:
+                    # In-graph (PCG/BCG) non-CP prefill is handled earlier by the
+                    # graph DSA split-op dispatch, so only the eager path reaches
+                    # here.
+                    assert not in_piecewise_or_breakable_cuda_graph, (
+                        "Internal error: in-graph DSA prefill must go through the "
+                        "graph DSA split-op dispatch"
+                    )
                     topk_result = self._get_topk_ragged(
                         enable_dual_stream,
                         forward_batch,
@@ -1974,6 +2007,88 @@ class Indexer(MultiPlatformOp):
         return topk_indices_prev[0], topk_indices_next[0]
 
 
+@register_custom_op(mutates_args=["topk_result"])
+@register_split_op()
+def pcg_dsa_indexer_prefill_split(
+    layer_id: int,
+    x: torch.Tensor,
+    q_lora: torch.Tensor,
+    positions: torch.Tensor,
+    topk_result: torch.Tensor,
+) -> None:
+    # Default in-graph indexer path for non-CP prefill: runs the whole indexer
+    # (q/k proj, head gate, k-cache store, topk) as one eager split op. PCG calls
+    # this as a split op; BCG uses the explicit eager wrapper below.
+    #
+    # Output contract (differs from the eager `forward` path): a split op returns
+    # None, so results are delivered only by mutating `topk_result` in place. The
+    # call site pre-allocates it at a static, padded shape and a downstream
+    # captured graph reads it at a fixed address; eager code instead allocates
+    # and returns a fresh, naturally-sized tensor each call.
+    assert _is_cuda, "Internal error: DSA graph dispatch is only supported on CUDA"
+    from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
+
+    forward_context = get_tc_piecewise_forward_context()
+    forward_batch = forward_context.forward_batch
+    indexer = forward_context.dsa_indexers[layer_id]
+    metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
+
+    extend_num_tokens = forward_batch.extend_num_tokens
+    # Empty buffer encodes return_indices=False for graph dispatch.
+    return_indices = topk_result.numel() != 0
+    k_only = not return_indices or (
+        indexer._should_skip_logits_computation(forward_batch)
+        and not indexer.dsa_enable_prefill_cp
+    )
+    if k_only:
+        indexer._forward_cuda_k_only(
+            x,
+            positions,
+            forward_batch,
+            layer_id,
+            act_quant,
+            enable_dual_stream=False,
+            metadata=metadata,
+            return_indices=return_indices,
+            num_tokens=extend_num_tokens,
+            topk_result=topk_result,
+        )
+        return
+
+    query, key = indexer._get_q_k_bf16(
+        q_lora,
+        x,
+        positions,
+        enable_dual_stream=False,
+        forward_batch=forward_batch,
+    )
+    q_fp8, q_scale = act_quant(query, indexer.block_size, indexer.scale_fmt)
+    # Reuse the compiled head-gate util shared with the eager path.
+    weights = indexer._get_logits_head_gate(x, q_scale)
+    # Store K cache + ragged top-k, sliced to the unpadded count and writing into
+    # the static padded topk_result buffer (the graph contract). Mirrors the eager
+    # path's store + _get_topk_ragged.
+    indexer._store_index_k_cache(
+        forward_batch=forward_batch,
+        layer_id=layer_id,
+        key=key[:extend_num_tokens],
+        act_quant=act_quant,
+        out_cache_loc=forward_batch.out_cache_loc[:extend_num_tokens],
+    )
+    indexer._get_topk_ragged(
+        False,
+        forward_batch,
+        layer_id,
+        q_fp8[:extend_num_tokens],
+        weights,
+        metadata,
+        topk_result,
+    )
+
+
+bcg_dsa_indexer_prefill_split = eager_on_graph(True)(pcg_dsa_indexer_prefill_split)
+
+
 def scattered_to_tp_attn_full(
     hidden_states: torch.Tensor,
     forward_batch,

python/sglang/srt/layers/attention/dsa/utils.py

@@ -9,9 +9,15 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
 )
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    is_in_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
-from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip
 from sglang.srt.utils.common import ceil_align, ceil_div
 
 
@@ -80,6 +86,19 @@ def is_dsa_prefill_cp_round_robin_split():
     )
 
 
+# Structural surface where the graph DSA split-op dispatch (DSA indexer) and the
+# MLA BMM-into-attention fusion apply: a non-speculative extend (prefill) running
+# inside a piecewise/breakable CUDA graph. Both fusions are now on by default on
+# this surface (no feature flag); each adds its own extra carve-outs at its call
+# site (e.g. the indexer also excludes DSA prefill context parallelism).
+def is_graph_dsa_split_op_surface(forward_batch: "ForwardBatch") -> bool:
+    return (
+        is_cuda()
+        and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
+        and forward_batch.forward_mode.is_extend_without_speculative()
+    )
+
+
 def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
     if not forward_batch.forward_mode.is_context_parallel_extend():
         return False

python/sglang/srt/layers/attention/dsa_backend.py

@@ -2471,14 +2471,18 @@ class DeepseekSparseAttnBackend(
         """
         Decide all attention prefill dispatch strategies for this batch.
         """
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+            is_in_breakable_cuda_graph,
+        )
         from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
             is_in_tc_piecewise_cuda_graph,
         )
         from sglang.srt.utils import get_device_sm, is_blackwell
 
         # Decide MHA vs MLA
-        if is_in_tc_piecewise_cuda_graph():
+        if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
-            # Can't branch on seq_lens_cpu in PCG, force mha off to guarantee correctness.
+            # Can't branch on seq_lens_cpu in graph replay, force MHA off to
+            # guarantee correctness.
             self.use_mha = False
         elif (
             forward_batch and forward_batch.forward_mode.is_extend_without_speculative()

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -1,18 +1,24 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional
 
 import torch
 
+from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
-from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
+from sglang.srt.layers.attention.dsa.utils import (
+    dsa_use_prefill_cp,
+    is_graph_dsa_split_op_surface,
+)
 from sglang.srt.layers.communicator import get_attn_tp_context
 from sglang.srt.layers.quantization.fp8_kernel import (
     fp8_dtype,
     per_tensor_quant_mla_fp8,
     per_token_group_quant_mla_deep_gemm_masked_fp8,
 )
+from sglang.srt.layers.radix_attention import unified_attention_with_output
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.lora.deepseek_mla_correction import (
     apply_q_correction as apply_kv_b_lora_q_correction,
@@ -28,6 +34,12 @@ from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
     get_token_to_kv_pool,
 )
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
@@ -48,17 +60,25 @@ from sglang.srt.state_capturer.indexer_topk import (
     maybe_capture_indexer_topk,
 )
 from sglang.srt.utils import BumpAllocator
+from sglang.srt.utils.custom_op import register_custom_op
 
 _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
 
 if TYPE_CHECKING:
     from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
 
+
+@dataclass(frozen=True)
+class MlaBmmFusionPlan:
+    q_nope_t: torch.Tensor
+    q_nope_out_buf: torch.Tensor
+    q_nope_out_view: torch.Tensor
+    attn_output_buf: torch.Tensor
+
+
 if _is_cuda:
     from sgl_kernel import bmm_fp8 as _raw_bmm_fp8
 
-    from sglang.srt.utils.custom_op import register_custom_op
-
     # TODO(yuwei): remove this wrapper after sgl-kernel registers its own fake/meta impl
     # Wrap bmm_fp8 as a custom op so torch.compile does not trace into
     # torch.cuda.current_blas_handle() (which returns a non-Tensor).
@@ -140,6 +160,63 @@ class DeepseekMLAForwardMixin:
             get_global_server_args().flashinfer_mla_disable_ragged
         )
 
+    def _can_fuse_bmm_into_attention(
+        self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch
+    ) -> bool:
+        # Shared activation surface with the DSA indexer graph dispatch
+        # (in piecewise/breakable graph + non-speculative extend). Like the indexer
+        # dispatch, this fusion is on by default on that surface.
+        if not is_graph_dsa_split_op_surface(forward_batch):
+            return False
+        if not self.use_dsa:
+            return False
+        if self.use_deep_gemm_bmm or _is_hip:
+            return False
+        if is_kv_b_lora_active(self):
+            return False
+        # The isolated 1-kernel graph is the bf16 fallback BMM. The fp8 and
+        # DeepGEMM branches already use different fused paths.
+        if self.w_kc.dtype == torch.float8_e4m3fn:
+            return False
+        if self.current_attention_backend not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
+            return False
+        return True
+
+    def _split_q_nope_pe(
+        self: DeepseekV2AttentionMLA,
+        q: torch.Tensor,
+        latent_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
+        return q_nope, q_pe, k_pe
+
+    def _make_mla_bmm_fusion_plan(
+        self: DeepseekV2AttentionMLA,
+        q: torch.Tensor,
+        q_nope: torch.Tensor,
+    ) -> MlaBmmFusionPlan:
+        q_nope_out_buf = q.new_empty(
+            (
+                self.num_local_heads,
+                q.shape[0],
+                self.kv_lora_rank,
+            )
+        )
+        q_nope_out_view = q_nope_out_buf.transpose(0, 1)
+        attn_output_buf = q.new_empty(
+            (
+                q.shape[0],
+                self.num_local_heads * self.kv_lora_rank,
+            )
+        )
+        return MlaBmmFusionPlan(
+            q_nope_t=q_nope.transpose(0, 1),
+            q_nope_out_buf=q_nope_out_buf,
+            q_nope_out_view=q_nope_out_view,
+            attn_output_buf=attn_output_buf,
+        )
+
     def forward_absorb_prepare(
         self: DeepseekV2AttentionMLA,
         positions: torch.Tensor,
@@ -151,8 +228,16 @@ class DeepseekMLAForwardMixin:
     ):
         from sglang.srt.model_executor.runner import get_is_capture_mode
 
+        fuse_bmm_attention = (
+            self.q_lora_rank is not None
+            and self._can_fuse_bmm_into_attention(forward_batch)
+        )
         q_lora = None
         topk_indices = None
+        q_nope = None
+        q_pe = None
+        k_pe = None
+        fusion_plan: Optional[MlaBmmFusionPlan] = None
         if self.q_lora_rank is not None:
             q, latent_cache = (
                 get_attn_tp_context()
@@ -276,6 +361,13 @@ class DeepseekMLAForwardMixin:
             else:
                 k_nope = k_nope.unsqueeze(1)
                 q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+
+                # Hoist these above the DSA indexer split op so the indexer
+                # and the composite bmm+attention split op are adjacent in FX.
+                if fuse_bmm_attention:
+                    q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache)
+                    fusion_plan = self._make_mla_bmm_fusion_plan(q, q_nope)
+
                 if q_lora is not None:
                     # See the skip_topk note above: shared layers have no
                     # indexer weights, so this gate must not fall back to
@@ -302,10 +394,15 @@ class DeepseekMLAForwardMixin:
             k_nope = latent_cache[..., : self.kv_lora_rank]
             k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1)
 
-        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        if q_nope is None:
-        k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
+            q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache)
 
         _kvb_q = None
+        if fusion_plan is not None:
+            # The composite split op fills q_nope_out_buf and attention reads
+            # this transposed alias directly.
+            q_nope_out = fusion_plan.q_nope_out_view
+        else:
             if _SGLANG_EXPERIMENTAL_LORA_OPTI:
                 # Fork the kv_b q-correction A-step onto the LoRA side stream to overlap the bmm.
                 from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import (
@@ -321,7 +418,9 @@ class DeepseekMLAForwardMixin:
                     masked_m,
                     expected_m,
                     aligned_m,
-            ) = per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1))
+                ) = per_token_group_quant_mla_deep_gemm_masked_fp8(
+                    q_nope.transpose(0, 1)
+                )
                 q_nope_out = q_nope.new_empty(
                     (self.num_local_heads, aligned_m, self.kv_lora_rank)
                 )
@@ -352,8 +451,11 @@ class DeepseekMLAForwardMixin:
                         q_nope_out,
                     )
                 else:
-                if (_use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn) or (
+                    if (
-                    get_is_capture_mode() and self.w_kc.dtype == torch.float8_e4m3fnuz
+                        _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn
+                    ) or (
+                        get_is_capture_mode()
+                        and self.w_kc.dtype == torch.float8_e4m3fnuz
                     ):
                         # fp8 Triton kernel: always on gfx950,
                         # cudagraph-only on gfx942 (hides launch overhead)
@@ -391,7 +493,11 @@ class DeepseekMLAForwardMixin:
                         ),
                     )
                     q_nope_out = bmm_fp8(
-                    q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
+                        q_nope_val,
+                        self.w_kc,
+                        q_nope_scale,
+                        self.w_scale,
+                        torch.bfloat16,
                     )
             else:
                 q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
@@ -433,6 +539,7 @@ class DeepseekMLAForwardMixin:
             positions,
             topk_indices,
             llama_4_scaling,
+            fusion_plan,
         )
 
     def forward_absorb_core(
@@ -446,6 +553,7 @@ class DeepseekMLAForwardMixin:
         positions,
         topk_indices,
         llama_4_scaling,
+        fusion_plan: Optional[MlaBmmFusionPlan] = None,
     ):
         save_kv_cache = True
 
@@ -527,6 +635,30 @@ class DeepseekMLAForwardMixin:
                         "is_neox": self.rotary_emb.is_neox_style,
                         "llama_4_scaling": llama_4_scaling,
                     }
+                if fusion_plan is not None:
+                    bmm_attention_fn = (
+                        bcg_mla_bmm_then_unified_attention
+                        if is_in_breakable_cuda_graph()
+                        else mla_bmm_then_unified_attention
+                    )
+                    bmm_attention_fn(
+                        fusion_plan.q_nope_t,
+                        self.w_kc,
+                        fusion_plan.q_nope_out_buf,
+                        q_nope_out,
+                        k_nope,
+                        fusion_plan.attn_output_buf,
+                        save_kv_cache,
+                        self.layer_id,
+                        q_pe,
+                        k_pe,
+                        cos_sin_cache=extra_args.get("cos_sin_cache"),
+                        is_neox=extra_args.get("is_neox"),
+                        llama_4_scaling=extra_args.get("llama_4_scaling"),
+                        topk_indices=topk_indices,
+                    )
+                    attn_output = fusion_plan.attn_output_buf
+                else:
                     attn_output = self.attn_mqa(
                         q_nope_out,
                         k_nope,
@@ -807,3 +939,54 @@ class DeepseekMLAForwardMixin:
             and self.current_attention_backend
             not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS
         )
+
+
+# Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output`
+# into one eager split op under both PCG and BCG. Without this, the bf16
+# fallback BMM is captured alone in its own single-kernel CUDA graph submodule,
+# paying per-submodule host overhead with no fusion benefit.
+#
+# `q_nope_out_view` aliases `q_nope_out_buf` (transposed). The op writes
+# `q_nope_out_buf` via `torch.bmm(..., out=...)` and then reads through
+# `q_nope_out_view`, so the alias's storage is mutated too. Declare it in
+# `mutates_args` to keep the schema honest.
+@register_custom_op(
+    mutates_args=["q_nope_out_buf", "q_nope_out_view", "attn_output_buf"]
+)
+@register_split_op()
+def mla_bmm_then_unified_attention(
+    q_nope_t: torch.Tensor,
+    w_kc: torch.Tensor,
+    q_nope_out_buf: torch.Tensor,
+    q_nope_out_view: torch.Tensor,
+    k_nope: torch.Tensor,
+    attn_output_buf: torch.Tensor,
+    save_kv_cache: bool,
+    layer_id: int,
+    q_pe: torch.Tensor,
+    k_pe: torch.Tensor,
+    cos_sin_cache: Optional[torch.Tensor] = None,
+    is_neox: Optional[bool] = None,
+    llama_4_scaling: Optional[torch.Tensor] = None,
+    topk_indices: Optional[torch.Tensor] = None,
+) -> None:
+    torch.bmm(q_nope_t, w_kc, out=q_nope_out_buf)
+    unified_attention_with_output(
+        q_nope_out_view,
+        k_nope,
+        k_nope,
+        attn_output_buf,
+        save_kv_cache,
+        layer_id,
+        q_rope=q_pe,
+        k_rope=k_pe,
+        cos_sin_cache=cos_sin_cache,
+        is_neox=is_neox,
+        llama_4_scaling=llama_4_scaling,
+        topk_indices=topk_indices,
+    )
+
+
+bcg_mla_bmm_then_unified_attention = eager_on_graph(True)(
+    mla_bmm_then_unified_attention
+)

python/sglang/srt/models/deepseek_v2.py

@@ -94,7 +94,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import (
     CombineInput,
     DispatchOutput,
 )
-from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
+from sglang.srt.layers.moe.topk import BypassedTopKOutput, TopK, TopKOutputFormat
 from sglang.srt.layers.moe.utils import (
     RoutingMethodType,
     filter_moe_weight_param_global_expert,
@@ -135,6 +135,13 @@ from sglang.srt.model_executor.cuda_graph_config import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.model_executor.runner import get_is_capture_mode
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.models.deepseek_common.attention_backend_handler import (
     AttentionBackendRegistry,
 )
@@ -207,6 +214,10 @@ else:
 
 logger = logging.getLogger(__name__)
 
+_enable_pcg_dsv2_dual_stream = (
+    _is_cuda and envs.SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM.get()
+)
+
 
 class DeepseekV2MLP(nn.Module):
     def __init__(
@@ -821,6 +832,26 @@ class DeepseekV2MoE(nn.Module):
             )
         ]
 
+    def _can_dual_stream_graph(
+        self, hidden_states: torch.Tensor, server_args=None
+    ) -> bool:
+        if server_args is None:
+            server_args = get_global_server_args()
+        return (
+            _enable_pcg_dsv2_dual_stream
+            and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
+            and get_moe_runner_backend().is_flashinfer_trtllm()
+            and self.alt_stream is not None
+            and self.num_fused_shared_experts == 0
+            and hidden_states.shape[0] > 0
+            and hasattr(self, "shared_experts")
+            and getattr(self.experts, "use_flashinfer_trtllm_moe", False)
+            and not self._enable_a2a_moe
+            and not self._fuse_shared_experts_inside_sbo
+            and not getattr(self, "is_hash", False)
+            and not server_args.enable_eplb
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -843,16 +874,24 @@ class DeepseekV2MoE(nn.Module):
             )
 
         if not self._enable_a2a_moe:
-            if (
+            server_args = get_global_server_args()
+            if self._can_dual_stream_graph(hidden_states, server_args):
+                return dsv2_flashinfer_moe_dual_stream_graph(
+                    hidden_states,
+                    self.layer_id,
+                    should_allreduce_fusion,
+                    use_reduce_scatter,
+                )
+            elif (
                 self.alt_stream is not None
                 and self.num_fused_shared_experts == 0
                 and hidden_states.shape[0] > 0
                 and get_is_capture_mode()
                 and not (
-                    get_global_server_args().enable_torch_compile
+                    server_args.enable_torch_compile
                     and hidden_states.shape[0]
-                    <= get_global_server_args().torch_compile_max_bs
+                    <= server_args.torch_compile_max_bs
-                    * (get_global_server_args().speculative_num_draft_tokens or 1)
+                    * (server_args.speculative_num_draft_tokens or 1)
                 )
             ):
                 return self.forward_normal_dual_stream(
@@ -886,6 +925,8 @@ class DeepseekV2MoE(nn.Module):
         gemm_output_zero_allocator: BumpAllocator = None,
         input_ids: Optional[torch.Tensor] = None,
         input_ids_global: Optional[torch.Tensor] = None,
+        *,
+        use_flashinfer_trtllm_bypass: bool = False,
     ) -> torch.Tensor:
         current_stream = torch.cuda.current_stream()
         self.alt_stream.wait_stream(current_stream)
@@ -901,6 +942,13 @@ class DeepseekV2MoE(nn.Module):
         with torch.cuda.stream(self.alt_stream):
             # router_logits: (num_tokens, n_experts)
             router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
+            if use_flashinfer_trtllm_bypass:
+                topk_output = BypassedTopKOutput(
+                    hidden_states=hidden_states,
+                    router_logits=router_logits,
+                    topk_config=self.topk.topk_config,
+                )
+            else:
                 topk_kwargs = (
                     {"input_ids": input_ids_global}
                     if getattr(self, "is_hash", False)
@@ -922,6 +970,10 @@ class DeepseekV2MoE(nn.Module):
                 final_hidden_states = self.experts.forward_deferred_finalize(
                     hidden_states, topk_output
                 )
+            elif use_flashinfer_trtllm_bypass:
+                final_hidden_states = self.experts.forward_impl(
+                    hidden_states, topk_output
+                )
             else:
                 final_hidden_states = self.experts(hidden_states, topk_output)
             if (
@@ -2864,4 +2916,25 @@ def flashinfer_dsv3_router_gemm(
     )
 
 
+@register_custom_op(out_shape="hidden_states")
+def dsv2_flashinfer_moe_dual_stream_graph(
+    hidden_states: torch.Tensor,
+    layer_id: int,
+    should_allreduce_fusion: bool,
+    use_reduce_scatter: bool,
+) -> torch.Tensor:
+    forward_context = get_tc_piecewise_forward_context()
+    assert forward_context is not None
+    assert forward_context.moe_fusions is not None
+
+    moe_fusion = forward_context.moe_fusions[layer_id]
+    assert moe_fusion is not None
+    return moe_fusion.forward_normal_dual_stream(
+        hidden_states,
+        should_allreduce_fusion=should_allreduce_fusion,
+        use_reduce_scatter=use_reduce_scatter,
+        use_flashinfer_trtllm_bypass=True,
+    )
+
+
 EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]

test/registered/cuda_graph/piecewise/test_pcg_glm5_fp8_tp8.py

@@ -0,0 +1,75 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=900, stage="base-c", runner_config="8-gpu-h200")
+
+GLM5_FP8_MODEL = "zai-org/GLM-5-FP8"
+
+
+class TestBCGGlm5Fp8TP8(CustomTestCase):
+    """Breakable CUDA graph prefill on GLM-5-FP8 (DSA model, TP=8, H200)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = GLM5_FP8_MODEL
+        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=[
+                "--tp-size",
+                "8",
+                "--trust-remote-code",
+                "--reasoning-parser",
+                "glm45",
+                "--tool-call-parser",
+                "glm47",
+                "--mem-fraction-static",
+                "0.8",
+                "--disable-flashinfer-autotune",
+                "--cuda-graph-backend-prefill=breakable",
+                # Small chunks => many prefill iterations, each <= the 2048
+                # capture max, so every prefill batch replays the BCG graph and
+                # exercises the DSA split-op / dual-stream / MLA-fusion paths.
+                "--chunked-prefill-size",
+                "512",
+                "--model-loader-extra-config",
+                '{"enable_multithread_load": true, "num_threads": 64}',
+            ],
+            env={
+                "SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM": "1",
+            },
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=200,
+            num_threads=200,
+            max_tokens=4096,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.92)
+
+
+if __name__ == "__main__":
+    unittest.main()