diff --git a/docs/cookbook/autoregressive/GLM/GLM-5.3-Flash.mdx b/docs/cookbook/autoregressive/GLM/GLM-5.3-Flash.mdx index 69966c10c..8b220890a 100644 --- a/docs/cookbook/autoregressive/GLM/GLM-5.3-Flash.mdx +++ b/docs/cookbook/autoregressive/GLM/GLM-5.3-Flash.mdx @@ -27,7 +27,9 @@ Choose your hardware, then choose the operating point that matches your workload Every listed hardware platform exposes both strategies. A **Verified** badge means that exact hardware and command were tested. **Final Verification In Progress** means the recipe runs and is queued for measurement on the final weights. **Not Verified** means the command is a supported starting point that still needs workload validation. A choice is disabled only when the underlying runtime combination is known to be unsupported. -The recommended selection is only a starting point. The same panel also lets you override the KV/DSA pairing, multimodal feature transport, and HiCache tiers. Changing an option that was not part of the measured command changes the badge to **Not Verified** without hiding the option. +The recommended selection is only a starting point. The same panel also lets you override the KV/DSA pairing, multimodal feature transport, Breakable Cuda Graph, and HiCache tiers. Changing an option that was not part of the measured command changes the badge to **Not Verified** without hiding the option. + +**Breakable Cuda Graph** defaults to **Off**. Select **On** to add `--cuda-graph-backend-prefill breakable` to the generated command. This requires a build that includes [PR #38522](https://github.com/sgl-project/sglang/pull/38522). import { Deployment } from "/src/snippets/_deployment.jsx"; import { config } from "/src/snippets/configs/zai-org/glm-5.3-flash.jsx"; diff --git a/docs/src/snippets/configs/zai-org/glm-5.3-flash.jsx b/docs/src/snippets/configs/zai-org/glm-5.3-flash.jsx index ffcbd4385..0c6da4f62 100644 --- a/docs/src/snippets/configs/zai-org/glm-5.3-flash.jsx +++ b/docs/src/snippets/configs/zai-org/glm-5.3-flash.jsx @@ -33,6 +33,7 @@ export const config = { s.kvDsaPair === pairing && s.mmTransport === "auto" && s.hicache === "off" && + s.bcg === "off" && s.dcp === "off" ); }, @@ -68,6 +69,20 @@ export const config = { }, ], }, + { + id: "bcg", + title: "Breakable Cuda Graph", + default: "off", + options: [ + { id: "off", label: "Off" }, + { + id: "on", + label: "On", + flags: ["--cuda-graph-backend-prefill breakable"], + hints: ["Enables breakable prefill CUDA graphs; requires a build with PR #38522."], + }, + ], + }, { id: "mmTransport", title: "VLM Transport", @@ -331,6 +346,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["bf16-tilelang", "fp8-trtllm"].includes(s.kvDsaPair) && s.mmTransport === "auto" && s.hicache === "off" && @@ -362,6 +378,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["bf16-tilelang", "fp8-trtllm"].includes(s.kvDsaPair) && s.mmTransport === "auto" && s.hicache === "off" && @@ -393,6 +410,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["bf16-tilelang", "fp8-trtllm"].includes(s.kvDsaPair) && s.mmTransport === "auto" && s.hicache === "off" && @@ -426,6 +444,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["bf16-tilelang", "fp8-trtllm"].includes(s.kvDsaPair) && s.mmTransport === "auto" && s.hicache === "off" && @@ -594,6 +613,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : s.mmTransport === "auto" && s.hicache === "off" ? "verified" : "unverified", @@ -623,6 +643,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["off", "l2"].includes(s.hicache) ? "verified" : "unverified", env: [], flags: [ @@ -645,6 +666,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : s.mmTransport === "auto" && s.hicache === "off" ? "verified" : "unverified", @@ -674,6 +696,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["off", "l2"].includes(s.hicache) ? "verified" : "unverified", env: [], flags: [ @@ -694,7 +717,8 @@ sgl-eval run gsm8k \\ match: { hw: "b200", strategy: "low-latency", quant: "fp8" }, nnodes: 1, verified: true, - verificationStatus: (s) => (s.hicache === "off" ? "verified" : "unverified"), + verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : (s.hicache === "off" ? "verified" : "unverified"), env: [], flags: [ "--model-path {{MODEL_NAME}}", @@ -720,6 +744,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["off", "l2"].includes(s.hicache) ? "verified" : "unverified", env: [], flags: [ @@ -740,7 +765,8 @@ sgl-eval run gsm8k \\ match: { hw: "b300", strategy: "low-latency", quant: "fp8" }, nnodes: 1, verified: true, - verificationStatus: (s) => (s.hicache === "off" ? "verified" : "unverified"), + verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : (s.hicache === "off" ? "verified" : "unverified"), env: [], flags: [ "--model-path {{MODEL_NAME}}", @@ -766,6 +792,7 @@ sgl-eval run gsm8k \\ nnodes: 1, verified: true, verificationStatus: (s) => + s.bcg !== "off" ? "unverified" : ["off", "l2"].includes(s.hicache) ? "verified" : "unverified", env: [], flags: [ diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py index bb0f22772..44b526828 100644 --- a/python/sglang/srt/arg_groups/cuda_graph_hook.py +++ b/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) diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 28440f244..96c996c8e 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/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) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index b589e250e..d01c5d318 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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", diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py index 4d24d61ff..45a4f30df 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py +++ b/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. diff --git a/python/sglang/srt/layers/attention/dsa/kpool_prefill_cuda_graph.py b/python/sglang/srt/layers/attention/dsa/kpool_prefill_cuda_graph.py new file mode 100644 index 000000000..8fb92d09f --- /dev/null +++ b/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) diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 01d829de0..6c6218820 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/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) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 7448c0ea4..49d24ab38 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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(