Add Opt-In for GLM-5.3 Flash breakable prefill CUDA graphs (#38522)

python/sglang/srt/arg_groups/cuda_graph_hook.py

@@ -280,6 +280,8 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any):
     rules = [
         (
             "KDA hybrid linear attention",
+            # GLM-5.3 Flash supports explicit BCG opt-in, but stays off by
+            # default like other KDA models. Explicit backends skip these rules.
             lambda: uses_kda_attention(model_config_of(server_args).hf_config),
         ),
         # DSV4 is BCG-compatible but introduces heavy memory pressure: the
@@ -399,6 +401,51 @@ def disable_prefill_cuda_graph_for_deepseek_trtllm_mla(server_args: Any):
     )
 
 
+def apply_glm5_chunked_prefill_default(server_args: Any):
+    """Set the opted-in GLM BCG chunk default before memory budgeting."""
+    cfg = resolving_view(server_args)
+    if (
+        get_platform().is_cuda
+        and (Phase.PREFILL, "backend") in server_args._cuda_graph_config_locked
+        and cfg.cuda_graph_config.prefill.backend == Backend.BREAKABLE
+        and cfg.chunked_prefill_size is None
+        and "Glm5NextForConditionalGeneration"
+        in model_config_of(server_args).hf_config.architectures
+    ):
+        declare_resolution(
+            server_args,
+            "_apply_glm5_chunked_prefill_default",
+            chunked_prefill_size=4096,
+        )
+
+
+def apply_glm5_prefill_cuda_graph_policy(server_args: Any):
+    """Set capture sizes for explicitly enabled GLM breakable prefill graphs."""
+    cfg = resolving_view(server_args)
+    if (
+        cfg.cuda_graph_config.prefill.backend != Backend.BREAKABLE
+        or "Glm5NextForConditionalGeneration"
+        not in model_config_of(server_args).hf_config.architectures
+    ):
+        return
+    locked = server_args._cuda_graph_config_locked
+    if any((Phase.PREFILL, key) in locked for key in ("max_bs", "bs")):
+        return
+    # Capacity defaults have already populated buckets. Replace the unlocked
+    # ceiling and its buckets together.
+    declare_resolution(
+        server_args,
+        "_apply_glm5_prefill_cuda_graph_policy",
+        cuda_graph_config=with_phase(
+            cfg.cuda_graph_config,
+            Phase.PREFILL,
+            max_bs=4096,
+            bs=generate_prefill_cuda_graph_batch_sizes(4096),
+        ),
+    )
+    apply_deepep_adjustments(server_args)
+
+
 def apply_deepep_adjustments(server_args: Any):
     """Config adjustments required by the DeepEP a2a backend."""
     cfg = resolving_view(server_args)

python/sglang/srt/arg_groups/pipeline.py

@@ -177,6 +177,8 @@ def run_resolution_pipeline(server_args: Any) -> None:
     # resolution (the declarative registry materializes too late to affect
     # it). Inkling opts into full-graph prefill capture here.
     from sglang.srt.arg_groups.cuda_graph_hook import (
+        apply_glm5_chunked_prefill_default,
+        apply_glm5_prefill_cuda_graph_policy,
         apply_inkling_prefill_cuda_graph_default,
         apply_muse_glimmer_prefill_cuda_graph_max_bs_default,
         disable_prefill_cuda_graph_for_deepseek_trtllm_mla,
@@ -190,6 +192,9 @@ def run_resolution_pipeline(server_args: Any) -> None:
     handle_dwdp(server_args)
 
     handle_cuda_graph_config(server_args)
+    # Requires the parsed backend and explicit-input locks, and must precede
+    # handle_gpu_memory_settings so the chunk size feeds memory budgeting.
+    apply_glm5_chunked_prefill_default(server_args)
 
     # Handle device-specific backends.
     from sglang.srt.arg_groups.platform_hook import (
@@ -261,6 +266,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
     handle_mamba_backend(server_args)
     handle_int8_mamba_checkpoint(server_args)
     handle_linear_attn_backend(server_args)
+    apply_glm5_prefill_cuda_graph_policy(server_args)
     handle_kv4_compatibility(server_args)
     handle_mxfp8_kv_cache_compatibility(server_args)
     run_post_process_pass(server_args, _page_size_default)

python/sglang/srt/configs/model_config.py

@@ -2132,13 +2132,14 @@ multimodal_piecewise_cuda_graph_supported_model_archs = [
 ]
 
 # Multimodal archs whose LM prefill is validated under breakable CUDA graph;
-# embed-carrying batches are rejected at replay (can_run_graph) and run eager.
+# replay eligibility for embed-carrying batches is checked by can_run_graph.
 # The Kimi archs are structurally multimodal -- their configs always carry a
 # vision_config, so is_multimodal is True even for text-only serving -- and the
 # generic multimodal rule disabled prefill CG for them despite the LM prefill
 # capturing cleanly.
 multimodal_breakable_cuda_graph_supported_model_archs = [
     "Cohere2VisionForConditionalGeneration",
+    "Glm5NextForConditionalGeneration",
     "InternS2MobiusForConditionalGeneration",
     "PaddleOCRVLForConditionalGeneration",
     "Qwen3_5ForConditionalGeneration",

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

@@ -39,6 +39,9 @@ from sglang.srt.model_executor.forward_context import (
     get_token_to_kv_pool,
 )
 from sglang.srt.model_executor.runner import get_is_capture_mode
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    is_in_breakable_cuda_graph,
+)
 from sglang.srt.runtime_context import get_device
 
 if TYPE_CHECKING:
@@ -1356,6 +1359,42 @@ class IndexerKPool(MultiPlatformOp):
         forward_batch: ForwardBatch,
         layer_id: int,
         return_indices: bool = True,
+    ) -> Optional[torch.Tensor]:
+        if (
+            is_in_breakable_cuda_graph()
+            and forward_batch.forward_mode.is_extend_without_speculative()
+        ):
+            from sglang.srt.layers.attention.dsa.kpool_prefill_cuda_graph import (
+                bcg_kpool_indexer_prefill_with_output,
+            )
+
+            # K-pool prefill plans contain request-specific tensors and launch
+            # counts. Like the ordinary DSA indexer, execute them eagerly and
+            # bridge the result into a stable buffer for captured attention.
+            output = torch.empty(
+                (
+                    x.shape[0] if return_indices else 0,
+                    self.index_topk + self.index_kpool - 1,
+                ),
+                dtype=torch.int32,
+                device=x.device,
+            )
+            bcg_kpool_indexer_prefill_with_output(
+                self, x, q_lora, positions, output, layer_id
+            )
+            return output if return_indices else None
+        return self._forward_cuda_impl(
+            x, q_lora, positions, forward_batch, layer_id, return_indices
+        )
+
+    def _forward_cuda_impl(
+        self,
+        x: torch.Tensor,
+        q_lora: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        return_indices: bool = True,
     ) -> Optional[torch.Tensor]:
         if is_hip():
             from sglang.kernels.ops.attention.dsa.tilelang_kernel import act_quant
@@ -1372,6 +1411,12 @@ class IndexerKPool(MultiPlatformOp):
             and get_is_capture_mode()
             and q_lora.shape[0] > 0
             and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD
+            # The BCG eager break must finish its indexer work before starting
+            # the next capture segment; keep its projections on one stream.
+            and not (
+                is_in_breakable_cuda_graph()
+                and forward_batch.forward_mode.is_extend_without_speculative()
+            )
         )
 
         # Skip DSA if the attention backend chooses to skip this batch.

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

@@ -0,0 +1,60 @@
+"""Breakable prefill bridge for the request-dependent pooled-key indexer."""
+
+import torch
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+)
+
+
+def _kpool_indexer_prefill_with_output(
+    indexer,
+    x: torch.Tensor,
+    q_lora: torch.Tensor,
+    positions: torch.Tensor,
+    output: torch.Tensor,
+    layer_id: int,
+) -> None:
+    # Metadata, write counts and cache destinations change between requests.
+    # Resolve the live batch inside the eager break, never from capture args.
+    forward_batch = get_tc_piecewise_forward_context().forward_batch
+    n = forward_batch.extend_num_tokens
+    if n is None or not 0 <= n <= x.shape[0]:
+        raise ValueError(f"Invalid pooled-indexer prefill token count: {n}")
+    if n > q_lora.shape[0] or n > positions.shape[0]:
+        raise ValueError("Pooled-indexer prefill inputs have inconsistent rows")
+    return_indices = output.shape[0] != 0
+    result = indexer._forward_cuda_impl(
+        x=x[:n],
+        q_lora=q_lora[:n],
+        positions=positions[:n],
+        forward_batch=forward_batch,
+        layer_id=layer_id,
+        return_indices=return_indices,
+    )
+    if not return_indices:
+        return
+    if result is None or result.shape != (n, output.shape[1]):
+        raise ValueError("Pooled-indexer prefill returned an unexpected top-k shape")
+    # The following captured attention segment reads this stable padded buffer.
+    output[:n].copy_(result)
+    output[n:].fill_(-1)
+
+
+def _kpool_indexer_prefill_capture_stub(
+    indexer,
+    x: torch.Tensor,
+    q_lora: torch.Tensor,
+    positions: torch.Tensor,
+    output: torch.Tensor,
+    layer_id: int,
+) -> None:
+    output.fill_(-1)
+
+
+bcg_kpool_indexer_prefill_with_output = eager_on_graph(
+    True, capture_stub=_kpool_indexer_prefill_capture_stub
+)(_kpool_indexer_prefill_with_output)

python/sglang/srt/layers/radix_linear_attention.py

@@ -171,8 +171,8 @@ def _linear_attention_with_output_impl(
             layer=attention_layer,
             forward_batch=forward_batch,
             mixed_qkv=mixed_qkv[:real_num_tokens],
-            a=a[:real_num_tokens],
-            b=b[:real_num_tokens],
+            a=a.narrow(0 if a.ndim == 2 else 1, 0, real_num_tokens),
+            b=b.narrow(0 if b.ndim == 2 else 1, 0, real_num_tokens),
             linear_attn_output=logical_output,
         )
     finally:
@@ -229,6 +229,16 @@ def unified_linear_attention_with_output(
     )
 
 
-bcg_unified_linear_attention_with_output = eager_on_graph(True)(
-    unified_linear_attention_with_output
-)
+def _linear_attention_capture_stub(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    output: torch.Tensor,
+    layer_id: int,
+) -> None:
+    output.zero_()
+
+
+bcg_unified_linear_attention_with_output = eager_on_graph(
+    True, capture_stub=_linear_attention_capture_stub
+)(unified_linear_attention_with_output)

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -1509,7 +1509,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
 
         # Main's monolithic BCG runner never invokes
         # on_after_cuda_graph_warmup between warmup iterations — the BCG
-        # contract is to keep warmup state untouched and let
+        # contract is to keep warmup metadata untouched and let
         # init_forward_metadata_in_graph (recorded inside the captured
         # forward) do any raw->full upgrade. cg-refactor's runner_backend
         # abstraction exposes a post_warmup_hook for backends that need
@@ -1519,6 +1519,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # corrupt warmup iter 2's metadata read.
         if isinstance(self.backend, BreakableCudaGraphBackend):
             post_warmup_hook = None
+            req_pool = self.model_runner.req_to_token_pool
+            mamba_pool = getattr(req_pool, "mamba_pool", None)
+            if mamba_pool is not None and not prefix_num_chunks:
+                capture_state_indices = req_pool.translate_mamba_indices(
+                    req_pool.get_mamba_indices(forward_batch.req_pool_indices)
+                ).unique()
+
+                def post_warmup_hook():
+                    mamba_pool.clear_slots(capture_state_indices)
+
+                post_warmup_hook()
         else:
             post_warmup_hook = getattr(attn_backend, "on_after_cuda_graph_warmup", None)
         self.backend.capture_one(