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

[Feature] Coordinate FullCG prefill across DP-attention ranks (#35640)

Browse Source 
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
This commit is contained in:
Aurick Qiao
2026-08-26 02:16:02 -07:00
committed by GitHub
co-authored by Yuwei An
parent 2511743bd7
commit 58ecbba0bd
15 changed files with 155 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/forward_batch_info.py
+11 -12
View File
@@ -446,8 +446,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# For DP attention
is_extend_in_batch: bool = False
can_run_dp_cuda_graph: bool = False
can_run_dp_breakable_cuda_graph: bool = False
can_run_decode_cuda_graph: bool = False
can_run_dp_prefill_cuda_graph: bool = False
global_forward_mode: Optional[ForwardMode] = None
# For two-batch overlap
@@ -703,7 +703,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
self.global_num_tokens_for_logprob_gpu = torch.tensor(
global_num_tokens_for_logprob, dtype=torch.int64
).to(device, non_blocking=True)
self.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph
self.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph
@classmethod
def init_new(
@@ -788,8 +788,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# Scalar config / flags
return_logprob=batch.return_logprob,
is_extend_in_batch=batch.is_extend_in_batch,
can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph,
can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_cuda_graph,
can_run_decode_cuda_graph=batch.can_run_decode_cuda_graph,
can_run_dp_prefill_cuda_graph=batch.can_run_dp_prefill_cuda_graph,
global_forward_mode=batch.global_forward_mode,
is_prefill_only=batch.is_prefill_only,
spec_algorithm=batch.spec_algorithm,
@@ -1312,21 +1312,20 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
):
# Joined ranks require real token counts instead of MAX_LEN padding.
dp_padding_mode = DpPaddingMode.SUM_LEN
# 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
# Prefill CUDA graphs require 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. This mismatches the
# rank-coupled communication geometry: DP gather/combine uses
# all_gather_into_tensor / reduce_scatter_tensor, while MoE backends may
# use A2A dispatch/combine. Force MAX_LEN so every rank pads to the global
# max and picks the same bucket.
#
# 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.
# 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 = get_exec().graph.cuda_graph_config.prefill
if (
self.can_run_dp_breakable_cuda_graph
self.can_run_dp_prefill_cuda_graph
and self.is_extend_in_batch
and prefill_cg.bs
and max(global_num_tokens) <= max(prefill_cg.bs)
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+4 -2
View File
@@ -694,7 +694,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
)
if self.require_mlp_sync:
is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
is_bs_supported = (
is_bs_supported and forward_batch.can_run_decode_cuda_graph
)
# NOTE: cuda graph cannot handle mixed batch (encoder_len = 0)
# If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph
@@ -735,7 +737,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
] and forward_batch.batch_size <= self._ragged_capture_slots(admission_tokens)
is_dp_supported = (
forward_batch.can_run_dp_cuda_graph if self.require_mlp_sync else True
forward_batch.can_run_decode_cuda_graph if self.require_mlp_sync else True
)
is_encoder_lens_supported = (
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+20 -16
View File
@@ -40,6 +40,7 @@ from __future__ import annotations
import copy
import inspect
import logging
from collections.abc import Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
@@ -826,19 +827,29 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
return prefix_chunk_len, prefix_chunk_len * capture_req_slots
def _select_prefix_capture_chunks(
self, forward_batch: ForwardBatch
self, prefix_lens: Sequence[int]
) -> Optional[int]:
"""Smallest captured variant covering the batch's max prefix, or None."""
max_prefix_len = max(
int(length) for length in forward_batch.extend_prefix_lens_cpu
)
max_prefix_len = max(int(length) for length in prefix_lens)
real_n = _ceil_div(max_prefix_len, self._prefix_chunk_len)
return next((n for n in self._prefix_capture_variants if n >= real_n), None)
def _has_uncapturable_chunked_prefix(
self, prefix_lens: Sequence[int] | None
) -> bool:
return (
self._capture_chunked_prefix
and prefix_lens is not None
and any(int(length) > 0 for length in prefix_lens)
and self._select_prefix_capture_chunks(prefix_lens) is None
)
def _shape_key(self, num_tokens: int, forward_batch: ForwardBatch) -> ShapeKey:
variant = None
if self._capture_chunked_prefix and self._has_prefix_hit(forward_batch):
captured_n = self._select_prefix_capture_chunks(forward_batch)
captured_n = self._select_prefix_capture_chunks(
forward_batch.extend_prefix_lens_cpu
)
assert captured_n is not None, "prefix batch has no captured FullCG variant"
variant = _chunked_prefix_variant(captured_n)
return ShapeKey(size=num_tokens, variant_label=variant)
@@ -1055,7 +1066,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
capture_hidden_mode,
return_logprob: bool,
lora_ineligible: bool = False,
chunked_prefix_uncapturable: bool = False,
) -> bool:
"""Rank-local replay eligibility: the single source of truth for
``can_run_graph`` (ForwardBatch, forward time) and the dp mlp-sync
@@ -1084,10 +1094,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
and any(prefix_lens)
):
return False
# FullCG's chunked-prefix topology covers a bounded prefix. The flag
# gating it is FULL-backend-only, so this is inert for the breakable
# vote path.
if chunked_prefix_uncapturable:
# FullCG's chunked-prefix topology covers a bounded prefix. Its capture
# flag is FullCG-only, so this is inert for the BreakableCG vote path.
if self._has_uncapturable_chunked_prefix(prefix_lens):
return False
# tc_piecewise captures with ForwardMode.EXTEND and spec_info=None.
if is_target_verify:
@@ -1115,7 +1124,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# (min-reduced votes; also requires every rank to hold tokens).
if (
forward_batch.global_num_tokens_cpu is not None
and not forward_batch.can_run_dp_breakable_cuda_graph
and not forward_batch.can_run_dp_prefill_cuda_graph
):
return False
@@ -1141,11 +1150,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
forward_batch
)
),
chunked_prefix_uncapturable=(
self._capture_chunked_prefix
and self._has_prefix_hit(forward_batch)
and self._select_prefix_capture_chunks(forward_batch) is None
),
):
return False
if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(