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

Enable breakable prefill CUDA graph for DP attention (#30898)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-07-12 17:10:04 -07:00
committed by GitHub
parent c616d5a55e
commit b94ac87e0c
9 changed files with 478 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/dp_attn.py
+13 -5
View File
@@ -14,7 +14,12 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
cuda_graph_fully_disabled,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.metrics_collector import DPCooperationInfo
from sglang.srt.server_args import ServerArgs
@@ -184,11 +189,14 @@ def prepare_mlp_sync_batch_raw(
or local_batch.forward_mode.is_decode_or_idle()
or local_batch.forward_mode.is_prebuilt()
) and not disable_cuda_graph
# Idle/None ranks are permissive (like can_cuda_graph): the all-gather
# min()-reduces this across DP ranks, so a prefill batch with idle ranks
# still resolves to True (idle ranks become a padded dummy extend).
can_run_breakable_cuda_graph = (
local_batch is not None
and local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
and not disable_cuda_graph
)
local_batch is None
or local_batch.forward_mode.is_idle()
or local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
) and check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
if local_batch is not None:
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+10
View File
@@ -788,6 +788,7 @@ def build_prefill_registry(
hidden_size: int = 0,
embed_dtype: Optional[torch.dtype] = None,
enable_mamba_track: bool = False,
enable_num_token_non_padded: bool = False,
register_input_embeds: bool = True,
share_pool: bool = True,
source: Optional[Any] = None,
@@ -876,6 +877,15 @@ def build_prefill_registry(
slots.append(GraphSlot("mamba_track_indices", _bs, torch.int64, axis="bs"))
slots.append(GraphSlot("mamba_track_mask", _bs, torch.bool, axis="bs"))
slots.append(GraphSlot("mamba_track_seqlens", _bs, torch.int32, axis="bs"))
if enable_num_token_non_padded:
slots.append(
GraphSlot(
"num_token_non_padded",
lambda _bs2, _mt: (1,),
torch.int32,
axis="none",
)
)
for slot in slots:
bind = None
python/sglang/srt/model_executor/forward_batch_info.py
+39 -1
View File
@@ -1177,6 +1177,26 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
self.is_extend_in_batch, global_num_tokens
)
# Prefill breakable CUDA graph requires every DP rank to run the SAME
# captured shape. Under SUM_LEN each rank pads to its own local token
# count and can select a different capture bucket, so the in-graph DP
# collectives (all_gather / reduce_scatter) mismatch across ranks and
# corrupt the output. Force MAX_LEN so every rank pads to the global
# max and picks the same bucket (mirrors the decode cuda graph
# contract, which always runs MAX_LEN).
#
# Only force MAX_LEN when the batch fits a captured breakable prefill
# graph; larger prefills fall back to eager and keep the
# memory-efficient SUM_LEN. global_num_tokens is identical across ranks
# (all-gathered), so the decision is consistent cluster-wide.
prefill_cg = model_runner.server_args.cuda_graph_config.prefill
if (
self.can_run_dp_breakable_cuda_graph
and self.is_extend_in_batch
and prefill_cg.bs
and max(global_num_tokens) <= max(prefill_cg.bs)
):
dp_padding_mode = DpPaddingMode.MAX_LEN
self.dp_padding_mode = dp_padding_mode
if dp_padding_mode.is_max_len():
@@ -1233,7 +1253,13 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
elif self.is_extend_in_batch and dp_padding_mode.is_max_len():
self._original_forward_mode = self.forward_mode
self.forward_mode = ForwardMode.EXTEND
if hybrid_ssm:
# Fabricate a single dummy request covering num_tokens for an
# empty (idle) rank. Hybrid-SSM families always take this path;
# non-hybrid ranks reach it once MAX_LEN is forced for the
# prefill breakable CUDA graph (idle + prefill), which needs
# every DP rank to run the same captured shape. The `else`
# branch handles decode rows padded to a 1-token extend.
if hybrid_ssm or self.seq_lens.shape[0] == 0:
dev = self.seq_lens.device
assert (
self.seq_lens.shape[0] == 0
@@ -1251,6 +1277,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
self.seq_lens = torch.tensor(
[num_tokens], dtype=self.seq_lens.dtype, device=dev
)
# orig_seq_lens is not padded by _pad_inputs_to_size, so
# fabricate it to match the dummy request (the breakable
# prefill CUDA graph runner reads it).
self.orig_seq_lens = torch.tensor(
[num_tokens], dtype=self.orig_seq_lens.dtype, device=dev
)
self.seq_lens_sum = int(num_tokens)
if self.seq_lens_cpu is not None:
self.seq_lens_cpu = torch.tensor(
@@ -1260,6 +1292,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
self.extend_seq_lens_cpu = [int(num_tokens)]
self.extend_logprob_start_lens_cpu = [0]
bs = self.batch_size = 1
# Count the dummy tokens as real, else MoE topk/all-to-all
# treats this rank as empty and starves later layers.
# (num_token_non_padded is None unless moe_ep_size > 1.)
if self.num_token_non_padded is not None:
self.num_token_non_padded.fill_(num_tokens)
self.num_token_non_padded_cpu = num_tokens
else:
self.extend_num_tokens = bs
self.extend_seq_lens = torch.full_like(self.seq_lens, 1)
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+70 -8
View File
@@ -61,6 +61,8 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardMode,
PPProxyTensors,
compute_local_num_token_non_padded,
enable_num_token_non_padded,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
@@ -98,6 +100,7 @@ from sglang.srt.utils import (
is_hip,
is_npu,
require_attn_tp_gather,
require_gathered_buffer,
require_mlp_tp_gather,
)
@@ -229,6 +232,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
hidden_size=self.model_runner.model_config.hidden_size,
embed_dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
enable_num_token_non_padded=enable_num_token_non_padded(),
source=self.buffers,
)
@@ -381,6 +385,39 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.model_runner.model_config.vocab_size, rows=rows
)
def _prefill_logits_buffer_rows(self, forward_batch: ForwardBatch) -> int:
if not forward_batch.return_logprob:
return forward_batch.batch_size
if not isinstance(self.backend, BreakableCudaGraphBackend):
return forward_batch.batch_size
global_num_tokens = forward_batch.global_num_tokens_for_logprob_cpu
if global_num_tokens is not None:
dp_rank = get_parallel().attn_dp_rank
return int(global_num_tokens[dp_rank if len(global_num_tokens) > 1 else 0])
return sum(
max(int(seq_len) - int(start_len), 1)
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
)
)
def _capture_num_token_non_padded(self, num_tokens: int) -> Optional[torch.Tensor]:
if not self.buffer_registry.has_slot("num_token_non_padded"):
return None
buf = self.buffer_registry.get_slot("num_token_non_padded").buffer
buf.fill_(num_tokens)
if require_gathered_buffer(self.model_runner.server_args):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=buf,
num_tokens_per_dp=num_tokens,
)
buf.copy_(local)
return buf
_aiter_chip_info_cached = False
@classmethod
@@ -592,7 +629,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
):
return False
num_tokens = len(forward_batch.input_ids)
if forward_batch.return_logprob:
if forward_batch.return_logprob and not isinstance(
self.backend, BreakableCudaGraphBackend
):
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
@@ -631,7 +670,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
Returns ``(forward_batch, attn_backend)`` to mirror decode's
capture_prepare signature.
"""
buffers = self.buffers
bs = self._capture_req_slots
# Slot 0 carries num_tokens; slots 1..bs-1 are zero-length sentinels.
lens_cpu = [num_tokens] + [0] * (bs - 1)
@@ -748,7 +786,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# FULL aux hidden states) captures with the right mode.
# Ported from main #27468.
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=None,
num_token_non_padded=self._capture_num_token_non_padded(num_tokens),
num_token_non_padded_cpu=num_tokens,
global_forward_mode=ForwardMode.EXTEND,
lora_ids=None,
@@ -829,7 +867,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""Pad, populate static buffers, and build the static_forward_batch
the model code reads during replay.
"""
buffers = self.buffers
num_tokens = len(forward_batch.input_ids)
static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
self.raw_num_tokens = num_tokens
@@ -881,6 +918,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
and forward_batch.mrope_positions is not None
else None
)
num_token_non_padded = (
_slot("num_token_non_padded")
if registry.has_slot("num_token_non_padded")
else forward_batch.num_token_non_padded
)
# Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1)
# doesn't fail on MIXED=3.
@@ -902,7 +944,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
input_embeds=input_embeds,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
next_token_logits_buffer=self._next_token_logits_buffer(bs),
next_token_logits_buffer=self._next_token_logits_buffer(
self._prefill_logits_buffer_rows(forward_batch)
),
orig_seq_lens=forward_batch.orig_seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
out_cache_loc=out_cache_loc,
@@ -911,25 +955,34 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
encoder_lens=forward_batch.encoder_lens,
return_logprob=False,
return_logprob=(
forward_batch.return_logprob
if isinstance(self.backend, BreakableCudaGraphBackend)
else False
),
is_prefill_only=forward_batch.is_prefill_only,
extend_seq_lens=forward_batch.extend_seq_lens,
extend_prefix_lens=forward_batch.extend_prefix_lens,
extend_start_loc=forward_batch.extend_start_loc,
extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
top_logprobs_nums=forward_batch.top_logprobs_nums,
token_ids_logprobs=forward_batch.token_ids_logprobs,
multi_item_delimiter_indices=forward_batch.multi_item_delimiter_indices,
extend_num_tokens=forward_batch.extend_num_tokens,
extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
positions=positions,
global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
dp_padding_mode=forward_batch.dp_padding_mode,
global_dp_buffer_len=forward_batch.global_dp_buffer_len,
mrope_positions=mrope_positions,
spec_algorithm=forward_batch.spec_algorithm,
spec_info=forward_batch.spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=forward_batch.num_token_non_padded,
num_token_non_padded=num_token_non_padded,
num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
global_forward_mode=pcg_global_forward_mode,
lora_ids=forward_batch.lora_ids,
@@ -1096,12 +1149,21 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.raw_bs if self._is_full_backend else self.raw_num_tokens
)
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[:logits_rows],
next_token_logits=(
output.next_token_logits[:logits_rows]
if output.next_token_logits is not None
else None
),
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
input_token_logprobs=output.input_token_logprobs,
input_top_logprobs_val=output.input_top_logprobs_val,
input_top_logprobs_idx=output.input_top_logprobs_idx,
input_token_ids_logprobs_val=output.input_token_ids_logprobs_val,
input_token_ids_logprobs_idx=output.input_token_ids_logprobs_idx,
mm_input_embeds=mm_input_embeds,
)
elif isinstance(output, EmbeddingPoolerOutput):
python/sglang/srt/model_executor/runner_utils/buffers.py
+3 -1
View File
@@ -62,7 +62,6 @@ def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -
@dataclass
class DecodeInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
input_embeds: torch.Tensor
req_pool_indices: torch.Tensor
@@ -328,6 +327,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
class PrefillInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
out_cache_loc: torch.Tensor
num_token_non_padded: torch.Tensor
mamba_track_indices: Optional[torch.Tensor]
mamba_track_mask: Optional[torch.Tensor]
mamba_track_seqlens: Optional[torch.Tensor]
@@ -351,6 +351,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
with torch.device(device):
input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
mamba_track_indices = (
torch.zeros((max_bs,), dtype=torch.int64)
if enable_mamba_track
@@ -376,6 +377,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
return cls(
input_ids=input_ids,
out_cache_loc=out_cache_loc,
num_token_non_padded=num_token_non_padded,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
python/sglang/srt/server_args.py
+33 -5
View File
@@ -219,9 +219,15 @@ ATTENTION_BACKEND_CHOICES = [
"intel_xpu",
]
DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton", "ascend"]
DETERMINISTIC_ATTENTION_BACKEND_CHOICES = [
"ascend",
"fa3",
"fa4",
"flashinfer",
"triton",
]
RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton", "ascend"]
RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"]
DISAGG_TRANSFER_BACKEND_CHOICES = [
"mooncake",
@@ -5256,11 +5262,32 @@ class ServerArgs:
if self._resolved().enable_dp_attention:
self.schedule_conservativeness = self.schedule_conservativeness * 0.3
assert self.tp_size % self.dp_size == 0
original_chunked_prefill_size = self.chunked_prefill_size
self.chunked_prefill_size = self.chunked_prefill_size // self.dp_size
logger.warning(
f"DP attention is enabled. The chunked prefill size is adjusted to {self.chunked_prefill_size} to avoid MoE kernel issues. "
f"DP attention is enabled. chunked prefill size is adjusted "
f"from {original_chunked_prefill_size} to {self.chunked_prefill_size}."
)
# The prefill CUDA graph max_bs was derived from the pre-DP-division
# chunked_prefill_size in _handle_gpu_memory_settings (which runs
# before this handler). Re-clamp it (and the captured shape list) to
# the per-DP-rank chunked_prefill_size so breakable CUDA graph
# capture never exceeds the MoE all-to-all's max_num_tokens budget,
# which is also sized from the DP-adjusted chunked_prefill_size.
prefill_cfg = self.cuda_graph_config.prefill
if (
prefill_cfg.backend != Backend.DISABLED
and prefill_cfg.max_bs is not None
and prefill_cfg.max_bs > self.chunked_prefill_size
and (Phase.PREFILL, "max_bs") not in self._cuda_graph_config_locked
):
prefill_cfg.max_bs = self.chunked_prefill_size
if (Phase.PREFILL, "bs") not in self._cuda_graph_config_locked:
prefill_cfg.bs = self._generate_prefill_cuda_graph_batch_sizes(
prefill_cfg.max_bs
)
# The dp-lm-head validation moved to the resolution pipeline
# (arg_groups/overrides.py: _dp_lm_head_validation), invoked here at
# its legacy slot.
@@ -6190,9 +6217,10 @@ class ServerArgs:
attention_backend = resolved_view(self).attention_backend
if is_deepseek_model:
if attention_backend not in ["fa3", "triton"]:
deepseek_deterministic_attention_backends = ["fa3", "triton"]
if attention_backend not in deepseek_deterministic_attention_backends:
raise ValueError(
f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
f"Currently only {deepseek_deterministic_attention_backends} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
)
if attention_backend not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND: