Breakable Cuda Graph Support for bs > 1 (#24662)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuwei An
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
a623ee4cb5
commit
5207f074a4
@@ -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=(
|
||||
|
||||
Reference in New Issue
Block a user