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/batch_overlap/two_batch_overlap.py
+2 -2
View File
@@ -736,8 +736,8 @@ class TboForwardBatchPreparer:
"forward_mode",
"is_extend_in_batch",
"return_logprob",
"can_run_dp_cuda_graph",
"can_run_dp_breakable_cuda_graph",
"can_run_decode_cuda_graph",
"can_run_dp_prefill_cuda_graph",
"dp_padding_mode",
"global_forward_mode",
"is_prefill_only",
python/sglang/srt/managers/schedule_batch.py
+4 -4
View File
@@ -2161,8 +2161,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# 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
tbo_split_seq_index: Optional[int] = None
# Rank-consistent forward mode for the recv skipper, derived from the MLP
# sync all-gather (the TBO-only `global_forward_mode` is None without TBO).
@@ -3355,8 +3355,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
spec_info=self.spec_info,
global_num_tokens=self.global_num_tokens,
global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_cuda_graph,
can_run_decode_cuda_graph=self.can_run_decode_cuda_graph,
can_run_dp_prefill_cuda_graph=self.can_run_dp_prefill_cuda_graph,
is_extend_in_batch=self.is_extend_in_batch,
is_prefill_only=self.is_prefill_only,
seq_lens_cpu=self.seq_lens_cpu,
python/sglang/srt/managers/scheduler_components/dp_attn.py
+8 -7
View File
@@ -221,8 +221,8 @@ def _update_gather_batch(
batch.global_forward_mode = mlp_sync_info.global_forward_mode
# Check forward mode for cuda graph
batch.can_run_dp_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph
batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph
batch.can_run_decode_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph
batch.can_run_dp_prefill_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph
def prepare_mlp_sync_batch_raw(
@@ -271,14 +271,15 @@ 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
breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
coordinated_prefill = check_cuda_graph_backend(
Phase.PREFILL, Backend.BREAKABLE
) or check_cuda_graph_backend(Phase.PREFILL, Backend.FULL)
prefill_graph_runner = (
model_runner.prefill_cuda_graph_runner if breakable_prefill else None
model_runner.prefill_cuda_graph_runner if coordinated_prefill else None
)
can_run_prefill_cuda_graph = (
local_batch is None
or local_batch.forward_mode.is_idle()
# Breakable Cuda Graph Backend Check.
or (
local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
and (
@@ -287,7 +288,7 @@ def prepare_mlp_sync_batch_raw(
batch_size=local_batch.batch_size(),
num_tokens=local_batch.extend_num_tokens,
input_embeds=local_batch.input_embeds,
replace_embeds=None,
replace_embeds=local_batch.replace_embeds,
prefix_lens=local_batch.prefix_lens,
is_target_verify=local_batch.forward_mode.is_target_verify(),
capture_hidden_mode=None,
@@ -295,7 +296,7 @@ def prepare_mlp_sync_batch_raw(
lora_ineligible=prefill_graph_runner.enable_lora,
)
)
and breakable_prefill
and coordinated_prefill
)
)
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(
python/sglang/srt/speculative/dspark_components/dspark_draft.py
+1 -1
View File
@@ -444,7 +444,7 @@ class DraftBlockProposer:
) -> None:
# The dense DSpark draft still reuses the target batch's graph tier.
# Set graph eligibility before the DP-MoE-only metadata early return.
forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph
forward_batch.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph
if not self._dp_moe_sync or batch.global_num_tokens is None:
return
# Graph bucket selection uses the raw per-rank request counts. Keep
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+3 -1
View File
@@ -318,7 +318,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
)
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
)
return is_bs_supported
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+3 -1
View File
@@ -309,7 +309,9 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
)
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
)
return is_bs_supported
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+3 -1
View File
@@ -237,7 +237,9 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
else cuda_graph_bs <= self.max_bs
)
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
)
return is_bs_supported
def capture_one_shape(
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+1 -1
View File
@@ -485,7 +485,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
self.cuda_graph_runner.execute(forward_batch)
)
else:
forward_batch.can_run_dp_cuda_graph = False
forward_batch.can_run_decode_cuda_graph = False
parent_list, top_scores_index, draft_tokens = self.draft_forward(
forward_batch
)
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+3 -1
View File
@@ -255,7 +255,9 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
)
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
)
return is_bs_supported
test/registered/dp_attn/test_dp_attention_bcg_kl.py
+88 -4
View File
@@ -1,6 +1,11 @@
from __future__ import annotations
import os
import random
import re
import shutil
import tempfile
import time
import unittest
import numpy as np
@@ -23,7 +28,12 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=160, stage="base-b", runner_config="2-gpu-large")
register_cuda_ci(est_time=320, stage="base-b", runner_config="2-gpu-large")
PREFILL_GRAPH_REPLAY_PATTERN = re.compile(r"Prefill batch.*cuda graph: True")
CACHED_PREFIX_GRAPH_REPLAY_PATTERN = re.compile(
r"Prefill batch.*#cached-token: [1-9][0-9]*.*cuda graph: True"
)
# ---------------------------------------------------------------------------
@@ -185,10 +195,11 @@ def _select_attention_backend():
)
class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
class _DPAttentionPrefillCudaGraphKLMixin:
num_samples = 48
max_prompt_tokens = 1024
max_new_tokens = 256
prefill_backend: str
@classmethod
def setUpClass(cls):
@@ -196,6 +207,11 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
cls.model = DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN
cls.base_url = DEFAULT_URL_FOR_TEST
cls.attention_backend = _select_attention_backend()
cls.log_dir = tempfile.mkdtemp(prefix=f"dp_attn_{cls.prefill_backend}_")
cls.stdout_path = os.path.join(cls.log_dir, "server.out")
cls.stderr_path = os.path.join(cls.log_dir, "server.err")
cls.stdout = open(cls.stdout_path, "w")
cls.stderr = open(cls.stderr_path, "w")
cls.process = popen_launch_server(
cls.model,
cls.base_url,
@@ -212,12 +228,15 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
cls.attention_backend,
"--moe-runner-backend",
"triton",
"--cuda-graph-backend-prefill=breakable",
f"--cuda-graph-backend-prefill={cls.prefill_backend}",
"--chunked-prefill-size",
"2048",
"--prefill-max-requests",
"2",
"--mem-fraction-static",
"0.70",
],
return_stdout_stderr=(cls.stdout, cls.stderr),
)
server_info = requests.get(f"{cls.base_url}/server_info", timeout=30).json()
@@ -235,6 +254,36 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
for attr in ("stdout", "stderr"):
output = getattr(cls, attr, None)
if output is not None and not output.closed:
output.close()
if hasattr(cls, "log_dir"):
shutil.rmtree(cls.log_dir, ignore_errors=True)
def _wait_for_prefill_graph_replay(
self,
offsets,
pattern=PREFILL_GRAPH_REPLAY_PATTERN,
case="a lone request",
):
deadline = time.monotonic() + 30
while time.monotonic() < deadline:
chunks = []
for path, offset in zip(
(self.stdout_path, self.stderr_path), offsets, strict=True
):
with open(path, "rb") as log:
log.seek(offset)
chunks.append(log.read().decode(errors="replace"))
logs = "\n".join(chunks)
if pattern.search(logs):
return
time.sleep(0.5)
self.fail(
f"No {self.prefill_backend} prefill CUDA graph replay was logged "
f"for {case}"
)
def test_prefill_and_decode_cache_hit_kl_is_zero(self):
server_info = requests.get(self.base_url + "/server_info", timeout=30).json()
@@ -243,15 +292,24 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
self.assertTrue(server_info["enable_deterministic_inference"])
self.assertEqual(server_info["attention_backend"], self.attention_backend)
self.assertEqual(
server_info["cuda_graph_config"]["prefill"]["backend"], "breakable"
server_info["cuda_graph_config"]["prefill"]["backend"],
self.prefill_backend,
)
print("=== Radix Cache KL Divergence Eval ===")
print(f"Server: {self.base_url} Samples: {self.num_samples}\n")
offsets = [
os.path.getsize(path) for path in (self.stdout_path, self.stderr_path)
]
prefill_kl = test_prefill_cache_hit(
self.base_url, self.input_ids, self.max_new_tokens
)
self._wait_for_prefill_graph_replay(
offsets,
CACHED_PREFIX_GRAPH_REPLAY_PATTERN,
"a cached-prefix request",
)
decode_kl = test_decode_cache_hit(
self.base_url, self.input_ids, self.max_new_tokens
)
@@ -259,6 +317,32 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
self.assertEqual(prefill_kl, 0.0)
self.assertEqual(decode_kl, 0.0)
def test_lone_request_replays_prefill_cuda_graph(self):
_flush_cache(self.base_url)
offsets = [
os.path.getsize(path) for path in (self.stdout_path, self.stderr_path)
]
result = _generate(
self.base_url,
self.input_ids[0],
max_new_tokens=1,
return_logprob=True,
)
self.assertNotIn("error", result)
self._wait_for_prefill_graph_replay(offsets)
class TestDPAttentionBreakablePrefillCudaGraphKL(
_DPAttentionPrefillCudaGraphKLMixin, CustomTestCase
):
prefill_backend = "breakable"
class TestDPAttentionFullPrefillCudaGraphKL(
_DPAttentionPrefillCudaGraphKLMixin, CustomTestCase
):
prefill_backend = "full"
if __name__ == "__main__":
unittest.main()
test/registered/spec/dspark/test_dspark_dp_tier.py
+2 -2
View File
@@ -67,7 +67,7 @@ class TestDraftDpSyncMetadata(CustomTestCase):
batch = SimpleNamespace(
global_num_tokens=[1, 3, 0, 2],
global_num_tokens_for_logprob=[1, 3, 0, 2],
can_run_dp_cuda_graph=True,
can_run_decode_cuda_graph=True,
)
with patch(
@@ -84,7 +84,7 @@ class TestDraftDpSyncMetadata(CustomTestCase):
self.assertEqual(forward_batch.num_token_non_padded.item(), 6)
self.assertEqual(forward_batch.num_token_non_padded.dtype, torch.int32)
self.assertEqual(forward_batch.num_token_non_padded_cpu, 6)
self.assertTrue(forward_batch.can_run_dp_cuda_graph)
self.assertTrue(forward_batch.can_run_decode_cuda_graph)
class TestBusyIdleGraphKeyIdentity(CustomTestCase):
test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py
+2 -2
View File
@@ -59,7 +59,7 @@ class TestMlpSyncPadUnpad(CustomTestCase):
batch = SimpleNamespace(
global_num_tokens=[2, 0, 3],
global_num_tokens_for_logprob=[2, 0, 3],
can_run_dp_cuda_graph=True,
can_run_decode_cuda_graph=True,
)
fb.init_mlp_sync_metadata(batch, torch.device("cpu"))
@@ -71,7 +71,7 @@ class TestMlpSyncPadUnpad(CustomTestCase):
torch.testing.assert_close(
fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6])
)
self.assertTrue(fb.can_run_dp_cuda_graph)
self.assertTrue(fb.can_run_decode_cuda_graph)
def test_draft_input_without_hidden_states_can_be_padded(self):
spec_info = SimpleNamespace(