minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Breakable Cuda Graph Support for bs > 1 (#24662)

Browse Source 
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuwei An
2026-05-11 13:28:49 +08:00
committed by GitHub
co-authored by Claude Opus 4.7
parent a623ee4cb5
commit 5207f074a4
1 changed files with 83 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/breakable_cuda_graph_runner.py
+83 -56
View File
@@ -114,19 +114,20 @@ class BreakableCudaGraphRunner:
self.moe_layers = model_runner.moe_layers self.moe_layers = model_runner.moe_layers
self.moe_fusions = model_runner.moe_fusions self.moe_fusions = model_runner.moe_fusions
with torch.device(self.device): # Resolve the inner transformer-stack module (the same boundary PCG draws
self.static_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64) # via patch_model). At replay we monkey-patch this module's forward with
self.static_extend_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64) # a closure that replays the captured CUDAGraph and returns the captured
self.static_extend_prefix_lens = torch.zeros( # hidden_states; the outer model.forward then runs logits_processor /
(self.max_bs,), dtype=torch.int64 # pooler eagerly with the live (multi-req) forward_batch.
) language_model = getattr(
self.static_extend_start_loc = torch.zeros( model_runner.model, "language_model", model_runner.model
(self.max_bs,), dtype=torch.int64 )
) self.layer_model = (
self.static_req_pool_indices = torch.zeros( language_model.model
(self.max_bs,), dtype=torch.int64 if hasattr(language_model, "model")
) and hasattr(language_model.model, "layers")
self.static_orig_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64) else language_model
)
# Memory pool # Memory pool
if get_global_graph_memory_pool() is None: if get_global_graph_memory_pool() is None:
@@ -185,8 +186,21 @@ class BreakableCudaGraphRunner:
) )
self.buffers.share_buffers() self.buffers.share_buffers()
@torch.no_grad()
def _run_forward(self, forward_batch, num_tokens): 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 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_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
set_is_extend_in_batch(False) set_is_extend_in_batch(False)
@@ -198,7 +212,7 @@ class BreakableCudaGraphRunner:
self.moe_layers, self.moe_layers,
self.moe_fusions, self.moe_fusions,
): ):
output = self.model_runner.model.forward( output = self.layer_model.forward(
forward_batch.input_ids, forward_batch.input_ids,
forward_batch.positions, forward_batch.positions,
forward_batch, forward_batch,
@@ -206,7 +220,13 @@ class BreakableCudaGraphRunner:
return output return output
def _build_capture_forward_batch(self, num_tokens): 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.layers.dp_attention import DpPaddingMode
from sglang.srt.model_executor.forward_batch_info import ( from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode, CaptureHiddenMode,
@@ -216,12 +236,13 @@ class BreakableCudaGraphRunner:
buffers = self.buffers buffers = self.buffers
bs = 1 bs = 1
self.static_seq_lens[:bs].fill_(num_tokens) with torch.device(self.device):
self.static_extend_seq_lens[:bs].fill_(num_tokens) seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
self.static_extend_prefix_lens[:bs].zero_() extend_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
self.static_extend_start_loc[:bs].zero_() extend_prefix_lens = torch.zeros((bs,), dtype=torch.int64)
self.static_req_pool_indices[:bs].copy_(torch.arange(bs, device=self.device)) extend_start_loc = torch.zeros((bs,), dtype=torch.int64)
self.static_orig_seq_lens[:bs].fill_(num_tokens) req_pool_indices = torch.arange(bs, dtype=torch.int64)
orig_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
return ForwardBatch( return ForwardBatch(
forward_mode=ForwardMode.EXTEND, forward_mode=ForwardMode.EXTEND,
@@ -230,10 +251,10 @@ class BreakableCudaGraphRunner:
input_embeds=( input_embeds=(
buffers.input_embeds[:num_tokens] if self.is_multimodal else None buffers.input_embeds[:num_tokens] if self.is_multimodal else None
), ),
req_pool_indices=self.static_req_pool_indices[:bs], req_pool_indices=req_pool_indices,
seq_lens=self.static_seq_lens[:bs], seq_lens=seq_lens,
next_token_logits_buffer=None, 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"), seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
req_to_token_pool=self.model_runner.req_to_token_pool, req_to_token_pool=self.model_runner.req_to_token_pool,
token_to_kv_pool=self.model_runner.token_to_kv_pool, token_to_kv_pool=self.model_runner.token_to_kv_pool,
@@ -251,9 +272,9 @@ class BreakableCudaGraphRunner:
encoder_lens=None, encoder_lens=None,
return_logprob=False, return_logprob=False,
extend_num_tokens=num_tokens, extend_num_tokens=num_tokens,
extend_seq_lens=self.static_extend_seq_lens[:bs], extend_seq_lens=extend_seq_lens,
extend_prefix_lens=self.static_extend_prefix_lens[:bs], extend_prefix_lens=extend_prefix_lens,
extend_start_loc=self.static_extend_start_loc[:bs], extend_start_loc=extend_start_loc,
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"), extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"), extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_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 self.output_buffers[num_tokens] = output
def can_run(self, forward_batch: "ForwardBatch"): def can_run(self, forward_batch: "ForwardBatch"):
# BCG graphs are captured with batch_size=1 (see _build_capture_forward_batch); if forward_batch.forward_mode.is_target_verify():
# 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:
return False return False
if forward_batch.input_embeds is not None: if forward_batch.input_embeds is not None:
return False return False
@@ -358,33 +374,44 @@ class BreakableCudaGraphRunner:
index = bisect.bisect_left(self.capture_num_tokens, num_tokens) index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
static_num_tokens = self.capture_num_tokens[index] 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(): with enable_breakable_cuda_graph():
static_forward_batch = self.replay_prepare(forward_batch, **kwargs) static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
bs = forward_batch.batch_size
# Update static buffers used by graph segments (esp. logits processor). original_layer_forward = self.layer_model.forward
# The graph reads from these addresses — they must have serving-time values. self.layer_model.forward = replay_layer_forward
self.static_seq_lens[:bs].copy_(forward_batch.seq_lens) try:
self.static_extend_seq_lens[:bs].copy_(forward_batch.extend_seq_lens) self.model_runner.attn_backend.init_forward_metadata(forward_batch)
self.static_extend_prefix_lens[:bs].copy_(forward_batch.extend_prefix_lens) with set_forward_context(
self.static_extend_start_loc[:bs].copy_(forward_batch.extend_start_loc) static_forward_batch,
self.static_req_pool_indices[:bs].copy_(forward_batch.req_pool_indices) self.attention_layers,
if forward_batch.orig_seq_lens is not None: self.quant_config,
self.static_orig_seq_lens[:bs].copy_(forward_batch.orig_seq_lens) 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): 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( return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[: self.raw_num_tokens], next_token_logits=output.next_token_logits[: self.raw_num_tokens],
hidden_states=( hidden_states=(