From 5207f074a4781d20e4d56d588831379b14b7c211 Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Sun, 10 May 2026 22:28:49 -0700 Subject: [PATCH] Breakable Cuda Graph Support for bs > 1 (#24662) Co-authored-by: Claude Opus 4.7 (1M context) --- .../breakable_cuda_graph_runner.py | 139 +++++++++++------- 1 file changed, 83 insertions(+), 56 deletions(-) diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py b/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py index da8b5ebde..31206ce3c 100644 --- a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py @@ -114,19 +114,20 @@ class BreakableCudaGraphRunner: self.moe_layers = model_runner.moe_layers self.moe_fusions = model_runner.moe_fusions - with torch.device(self.device): - self.static_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64) - self.static_extend_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64) - self.static_extend_prefix_lens = torch.zeros( - (self.max_bs,), dtype=torch.int64 - ) - self.static_extend_start_loc = torch.zeros( - (self.max_bs,), dtype=torch.int64 - ) - self.static_req_pool_indices = torch.zeros( - (self.max_bs,), dtype=torch.int64 - ) - self.static_orig_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64) + # Resolve the inner transformer-stack module (the same boundary PCG draws + # via patch_model). At replay we monkey-patch this module's forward with + # a closure that replays the captured CUDAGraph and returns the captured + # hidden_states; the outer model.forward then runs logits_processor / + # pooler eagerly with the live (multi-req) forward_batch. + language_model = getattr( + model_runner.model, "language_model", model_runner.model + ) + self.layer_model = ( + language_model.model + if hasattr(language_model, "model") + and hasattr(language_model.model, "layers") + else language_model + ) # Memory pool if get_global_graph_memory_pool() is None: @@ -185,8 +186,21 @@ class BreakableCudaGraphRunner: ) self.buffers.share_buffers() + @torch.no_grad() def _run_forward(self, forward_batch, num_tokens): - """Run model forward with proper context.""" + """Run layer-stack forward with proper context. + + Captures only the inner transformer stack (layer_model). The outer + model.forward's tail (logits_processor / pooler) is intentionally + excluded — it has bs-shaped kernels that would bake batch_size=1 + into the captured graph. + + ``@torch.no_grad`` mirrors the decorator on the outer ``*ForCausalLM.forward`` + (e.g. qwen3.py:507). Calling ``layer_model.forward`` directly skips that + decorator, so we apply it here — without it some MoE @torch.compile + kernels (``torch.sum(out=...)``) fail dynamo with "out= doesn't support + autograd", and mamba state ops can spuriously track gradients. + """ forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len()) set_is_extend_in_batch(False) @@ -198,7 +212,7 @@ class BreakableCudaGraphRunner: self.moe_layers, self.moe_fusions, ): - output = self.model_runner.model.forward( + output = self.layer_model.forward( forward_batch.input_ids, forward_batch.positions, forward_batch, @@ -206,7 +220,13 @@ class BreakableCudaGraphRunner: return output def _build_capture_forward_batch(self, num_tokens): - """Build a ForwardBatch for capture using static buffers for stable addresses.""" + """Build a bs=1 placeholder ForwardBatch for capture. + + bs=1 here is only a placeholder for attention/mamba breaks' metadata + shapes; replay supplies live multi-req metadata via replay_prepare. + Captured kernels run only on the token-major layer stack and are + bs-invariant. + """ from sglang.srt.layers.dp_attention import DpPaddingMode from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -216,12 +236,13 @@ class BreakableCudaGraphRunner: buffers = self.buffers bs = 1 - self.static_seq_lens[:bs].fill_(num_tokens) - self.static_extend_seq_lens[:bs].fill_(num_tokens) - self.static_extend_prefix_lens[:bs].zero_() - self.static_extend_start_loc[:bs].zero_() - self.static_req_pool_indices[:bs].copy_(torch.arange(bs, device=self.device)) - self.static_orig_seq_lens[:bs].fill_(num_tokens) + with torch.device(self.device): + seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64) + extend_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64) + extend_prefix_lens = torch.zeros((bs,), dtype=torch.int64) + extend_start_loc = torch.zeros((bs,), dtype=torch.int64) + req_pool_indices = torch.arange(bs, dtype=torch.int64) + orig_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64) return ForwardBatch( forward_mode=ForwardMode.EXTEND, @@ -230,10 +251,10 @@ class BreakableCudaGraphRunner: input_embeds=( buffers.input_embeds[:num_tokens] if self.is_multimodal else None ), - req_pool_indices=self.static_req_pool_indices[:bs], - seq_lens=self.static_seq_lens[:bs], + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, next_token_logits_buffer=None, - orig_seq_lens=self.static_orig_seq_lens[:bs], + orig_seq_lens=orig_seq_lens, seq_lens_cpu=torch.tensor([num_tokens], device="cpu"), req_to_token_pool=self.model_runner.req_to_token_pool, token_to_kv_pool=self.model_runner.token_to_kv_pool, @@ -251,9 +272,9 @@ class BreakableCudaGraphRunner: encoder_lens=None, return_logprob=False, extend_num_tokens=num_tokens, - extend_seq_lens=self.static_extend_seq_lens[:bs], - extend_prefix_lens=self.static_extend_prefix_lens[:bs], - extend_start_loc=self.static_extend_start_loc[:bs], + extend_seq_lens=extend_seq_lens, + extend_prefix_lens=extend_prefix_lens, + extend_start_loc=extend_start_loc, extend_prefix_lens_cpu=torch.tensor([0], device="cpu"), extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"), extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"), @@ -309,12 +330,7 @@ class BreakableCudaGraphRunner: self.output_buffers[num_tokens] = output def can_run(self, forward_batch: "ForwardBatch"): - # BCG graphs are captured with batch_size=1 (see _build_capture_forward_batch); - # the captured logits-gather / sampler path yields bs=1 outputs. Multi-req - # prefill would silently return wrong-shaped logits, corrupting downstream - # output_ids and breaking the subsequent decode step. Reject here so the - # caller falls back to the eager extend path. - if forward_batch.batch_size > 1: + if forward_batch.forward_mode.is_target_verify(): return False if forward_batch.input_embeds is not None: return False @@ -358,33 +374,44 @@ class BreakableCudaGraphRunner: index = bisect.bisect_left(self.capture_num_tokens, num_tokens) static_num_tokens = self.capture_num_tokens[index] + captured_graph = self.graphs[static_num_tokens] + captured_hidden = self.output_buffers[static_num_tokens] + + # Closure replaces layer_model.forward for the duration of the outer + # model.forward call. Replays the captured CUDAGraph and hands the + # outer forward the captured hidden_states; logits_processor / pooler + # then runs eagerly on top with the live multi-req forward_batch. + def replay_layer_forward(*args, **layer_kwargs): + captured_graph.replay() + return captured_hidden + with enable_breakable_cuda_graph(): static_forward_batch = self.replay_prepare(forward_batch, **kwargs) - bs = forward_batch.batch_size - # Update static buffers used by graph segments (esp. logits processor). - # The graph reads from these addresses — they must have serving-time values. - self.static_seq_lens[:bs].copy_(forward_batch.seq_lens) - self.static_extend_seq_lens[:bs].copy_(forward_batch.extend_seq_lens) - self.static_extend_prefix_lens[:bs].copy_(forward_batch.extend_prefix_lens) - self.static_extend_start_loc[:bs].copy_(forward_batch.extend_start_loc) - self.static_req_pool_indices[:bs].copy_(forward_batch.req_pool_indices) - if forward_batch.orig_seq_lens is not None: - self.static_orig_seq_lens[:bs].copy_(forward_batch.orig_seq_lens) + original_layer_forward = self.layer_model.forward + self.layer_model.forward = replay_layer_forward + try: + self.model_runner.attn_backend.init_forward_metadata(forward_batch) + with set_forward_context( + static_forward_batch, + self.attention_layers, + self.quant_config, + self.moe_layers, + self.moe_fusions, + ): + output = self.model_runner.model.forward( + static_forward_batch.input_ids, + static_forward_batch.positions, + static_forward_batch, + **kwargs, + ) + finally: + self.layer_model.forward = original_layer_forward - # Set forward context and replay - self.model_runner.attn_backend.init_forward_metadata(forward_batch) - with set_forward_context( - static_forward_batch, - self.attention_layers, - self.quant_config, - self.moe_layers, - self.moe_fusions, - ): - self.graphs[static_num_tokens].replay() - - output = self.output_buffers[static_num_tokens] if isinstance(output, LogitsProcessorOutput): + # Slice trailing-padding off hidden_states; next_token_logits is + # bs-shaped from logits_processor (bs <= raw_num_tokens), so the + # slice is a no-op for that field but matches PCG's pattern. return LogitsProcessorOutput( next_token_logits=output.next_token_logits[: self.raw_num_tokens], hidden_states=(