From ff5578eb4e1d68d842303e3ab9ce4c86038c1c26 Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Thu, 27 Aug 2026 12:59:56 -0700 Subject: [PATCH] [1/N][Mix] Mixed Chunk Prefill Base (#36288) Co-authored-by: Claude Fable 5 --- python/sglang/srt/managers/schedule_batch.py | 30 ++++ python/sglang/srt/managers/scheduler.py | 9 + .../batch_result_processor.py | 15 +- .../managers/scheduler_components/dp_attn.py | 169 ++++++++++++++---- .../srt/model_executor/runner/__init__.py | 3 +- .../srt/model_executor/runner/eager_runner.py | 15 +- .../runner/prefill_cuda_graph_runner.py | 9 +- 7 files changed, 209 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index b16d3f8ec..40a3c4d1b 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2839,6 +2839,36 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ) self.is_prefill_only = False + def convert_decode_to_extend(self): + """View every decode request as a 1-token extend with its context as + prefix, making this a plain extend (prefill) batch. Called after the + DP mlp-sync when a peer rank runs extend, so this rank replays the + extend graphs; requires prepare_for_decode(); token count unchanged. + """ + bs = self.batch_size() + self.forward_mode = ForwardMode.EXTEND + # Stale residue from this object's life as a prefill batch; the + # extend-merge path would exclude (silently drop) that req otherwise. + self.chunked_req = None + # Also stale residue; None keeps the prefill result path from + # re-reporting old prefill stats for what is decode work. + self.prefill_stats = None + for req in self.reqs: + req._refresh_fill_ids() + full_len = len(req.full_untruncated_fill_ids) + req.set_extend_range(full_len - 1, full_len) + + # Same one-step output_ids delay handling as mix_with_running. + delta = 0 if self.enable_overlap else -1 + self.prefix_lens = [ + len(r.origin_input_ids) + len(r.output_ids) + delta for r in self.reqs + ] + self.extend_lens = [1] * bs + self.extend_num_tokens = bs + self.extend_logprob_start_lens = [0] * bs + self.decoding_reqs = self.reqs + self.is_prefill_only = False + def new_tokens_required_next_decode( self, selected_indices: Optional[List[int]] = None ): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index edc52550b..38f13e95a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3270,6 +3270,15 @@ class Scheduler( ret = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch( ret, need_sync=need_mlp_sync ) + # Decode->extend conversion keeps a heterogeneous dp step replayable. + converted = self.dp_attn_adapter.maybe_convert_decode_to_extend(ret) + if converted is running_batch and converted.forward_mode.is_extend(): + # The converted batch re-enters via the last_batch extend-merge + # next iteration; empty running_batch or it merges with itself. + running_batch = ScheduleBatch( + reqs=[], batch_is_full=running_batch.batch_is_full + ) + ret = converted self._arm_prefill_decode_interval(ret) # Handle ngram embedding diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 8c7c66239..9ba9160c6 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -432,12 +432,15 @@ class SchedulerBatchResultProcessor: ) can_run_cuda_graph = result.can_run_cuda_graph - self.metrics_reporter.report_prefill_stats( - batch=batch, - prefill_stats=batch.prefill_stats, - can_run_cuda_graph=can_run_cuda_graph, - dp_cooperation_info=batch.dp_cooperation_info, - ) + # None on decode->extend converted batches; they are decode work and + # have no prefill stats to report. + if batch.prefill_stats is not None: + self.metrics_reporter.report_prefill_stats( + batch=batch, + prefill_stats=batch.prefill_stats, + can_run_cuda_graph=can_run_cuda_graph, + dp_cooperation_info=batch.dp_cooperation_info, + ) def _convert_embeddings(self, *, result: EmbeddingBatchResult) -> list: is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set() diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 738ff67d9..8ac9da6f7 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -10,13 +10,16 @@ from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.environ import envs +from sglang.srt.layers.cp.utils import get_cp_strategy from sglang.srt.layers.dp_attention import world_dp_gather_enabled +from sglang.srt.layers.moe.utils import get_moe_a2a_backend from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler_components.recv_skipper import ( SchedulerRecvSkipper, ) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache +from sglang.srt.mem_cache.kv_cache_builder import uses_ssm_state from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.model_executor.cuda_graph_config import ( Backend, @@ -25,8 +28,14 @@ from sglang.srt.model_executor.cuda_graph_config import ( cuda_graph_fully_disabled, ) from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.model_executor.runner import PrefillCudaGraphRunner from sglang.srt.observability.metrics_collector import DPCooperationInfo -from sglang.srt.runtime_context import get_parallel, get_schedule +from sglang.srt.runtime_context import ( + get_exec, + get_memory, + get_parallel, + get_schedule, +) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.common import require_mlp_tp_gather @@ -224,6 +233,91 @@ def _update_gather_batch( batch.can_run_dp_prefill_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph +def _local_decode_cuda_graph_vote( + *, + local_batch: Optional[ScheduleBatch], + disable_cuda_graph: bool, +) -> bool: + """This rank's vote for the decode graph (min-reduced across dp ranks).""" + if disable_cuda_graph: + return False + return ( + local_batch is None + or local_batch.forward_mode.is_decode_or_idle() + or local_batch.forward_mode.is_prebuilt() + ) + + +def _local_prefill_cuda_graph_vote( + *, + local_batch: Optional[ScheduleBatch], + prefill_graph_runner, + coordinated_prefill: bool, + breakable_prefill: bool, + spec_algorithm: SpeculativeAlgorithm, + model_config, +) -> bool: + """This rank's vote for the prefill graph (min-reduced across dp + ranks). Extend/mixed batches vote their own replayability; a decode + batch eligible for the decode->extend conversion votes as its 1-token- + extend view, so the vote and the post-sync conversion always agree.""" + if local_batch is None or local_batch.forward_mode.is_idle(): + return True + if not coordinated_prefill: + return False + + mode = local_batch.forward_mode + if mode in (ForwardMode.EXTEND, ForwardMode.MIXED): + num_tokens = local_batch.extend_num_tokens + input_embeds = local_batch.input_embeds + replace_embeds = local_batch.replace_embeds + prefix_lens = local_batch.prefix_lens + return_logprob = local_batch.return_logprob + elif ( + mode.is_decode() + # Conversion replays the breakable graphs only; full's fixed + # request-slot geometry does not cover converted decode tails. + and breakable_prefill + # decode->extend conversion eligibility; needs the captured-graph + # prefill runner, not the eager fallback. + and isinstance(prefill_graph_runner, PrefillCudaGraphRunner) + and spec_algorithm.is_none() + and not local_batch.return_logprob + # Grammar FSMs advance through the decode result path only. + and not local_batch.has_grammar + # Small-bucket BCG replays amplify the a2a EP logits drift (#30898) + # into an accuracy loss. + and get_moe_a2a_backend().is_none() + # The converted view lacks prepare_for_extend's mamba-track fills. + and not uses_ssm_state(model_config) + # HiSparse decode has its own batch lifecycle and host-offloaded KV. + and not get_memory().enable_hisparse + and not get_exec().overlap.enable_two_batch_overlap + and get_cp_strategy() is None + ): + num_tokens = local_batch.batch_size() + input_embeds = None + replace_embeds = None + prefix_lens = None + return_logprob = False + else: + return False + + if prefill_graph_runner is None: + return True + return prefill_graph_runner.can_replay_locally( + batch_size=local_batch.batch_size(), + num_tokens=num_tokens, + input_embeds=input_embeds, + replace_embeds=replace_embeds, + prefix_lens=prefix_lens, + is_target_verify=mode.is_target_verify(), + capture_hidden_mode=None, + return_logprob=return_logprob, + lora_ineligible=prefill_graph_runner.enable_lora, + ) + + def prepare_mlp_sync_batch_raw( local_batch: ScheduleBatch, model_runner: ModelRunner, @@ -265,38 +359,23 @@ def prepare_mlp_sync_batch_raw( ) skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get() - can_run_decode_cuda_graph = ( - local_batch is None - or local_batch.forward_mode.is_decode_or_idle() - or local_batch.forward_mode.is_prebuilt() - ) and not disable_cuda_graph - coordinated_prefill = check_cuda_graph_backend( - Phase.PREFILL, Backend.BREAKABLE - ) or check_cuda_graph_backend(Phase.PREFILL, Backend.FULL) + can_run_decode_cuda_graph = _local_decode_cuda_graph_vote( + local_batch=local_batch, disable_cuda_graph=disable_cuda_graph + ) + breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE) + coordinated_prefill = breakable_prefill or check_cuda_graph_backend( + Phase.PREFILL, Backend.FULL + ) prefill_graph_runner = ( 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() - or ( - local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED) - and ( - prefill_graph_runner is None - or prefill_graph_runner.can_replay_locally( - batch_size=local_batch.batch_size(), - num_tokens=local_batch.extend_num_tokens, - input_embeds=local_batch.input_embeds, - 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, - return_logprob=local_batch.return_logprob, - lora_ineligible=prefill_graph_runner.enable_lora, - ) - ) - and coordinated_prefill - ) + can_run_prefill_cuda_graph = _local_prefill_cuda_graph_vote( + local_batch=local_batch, + prefill_graph_runner=prefill_graph_runner, + coordinated_prefill=coordinated_prefill, + breakable_prefill=breakable_prefill, + spec_algorithm=model_runner.spec_algorithm, + model_config=model_runner.model_config, ) is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False @@ -438,6 +517,36 @@ class SchedulerDPAttnAdapter: batch = self.prepare_mlp_sync_batch(batch) return batch + def maybe_convert_decode_to_extend( + self, batch: Optional[ScheduleBatch] + ) -> Optional[ScheduleBatch]: + """After the mlp-sync gather: convert an eligible decode batch to the + extend view when a peer rank runs extend this step, so the step stays + mode-homogeneous and every rank replays the extend graphs instead of + all falling to eager.""" + if batch is None or not batch.forward_mode.is_decode(): + return batch + # Global triggers from the gather. This rank's own eligibility (spec/ + # TBO/CP/logprob/replayability) is folded into the min-reduced vote: + # if it failed, can_run_dp_prefill_cuda_graph is already False. + if not batch.is_extend_in_batch: + return batch + if not batch.can_run_dp_prefill_cuda_graph: + # The step is eager everywhere; eager decode beats eager mixed. + return batch + global_tokens = batch.global_num_tokens + if ( + global_tokens is not None + and len(global_tokens) > 1 + and min(global_tokens) == 0 + ): + # An idle rank makes the prefill runner reject replay for every + # rank (_has_inactive_dp_rank); converting would only trade eager + # decode for eager mixed. + return batch + batch.convert_decode_to_extend() + return batch + def get_idle_batch(self) -> ScheduleBatch: idle_batch = ScheduleBatch.init_new( [], diff --git a/python/sglang/srt/model_executor/runner/__init__.py b/python/sglang/srt/model_executor/runner/__init__.py index fd03b2752..9aa65fbf5 100644 --- a/python/sglang/srt/model_executor/runner/__init__.py +++ b/python/sglang/srt/model_executor/runner/__init__.py @@ -12,7 +12,8 @@ Public API: - BaseCudaGraphRunner — abstract cuda-graph base; bucket padding + capture-loop scaffolding on top of BaseRunner. - DecodeCudaGraphRunner — concrete decode-phase runner. - - PrefillCudaGraphRunner — concrete prefill-phase runner. + - PrefillCudaGraphRunner — concrete prefill-phase runner (extend family; + MIXED batches replay the EXTEND-captured graphs). - EagerRunner — no-cuda-graph runner; runs model.forward live (the eager dual of the cuda-graph runners), mode-dispatched over decode + extend + idle. diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index dcc23b730..1e6a37c07 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -27,6 +27,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.cp.utils import ( cp_gather_after_forward, cp_shard_model_inputs, + get_cp_strategy, is_cp_v2_active, prepare_cp_forward, ) @@ -37,7 +38,11 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import ( from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( create_chunked_prefix_cache_kv_indices, ) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + ForwardMode, + PPProxyTensors, +) from sglang.srt.model_executor.forward_context import ( ForwardContext, forward_context, @@ -59,7 +64,7 @@ from sglang.srt.runtime_context import ( max_prefill_buffer_tokens, max_speculative_num_draft_tokens, ) -from sglang.srt.utils import is_hip +from sglang.srt.utils import is_hip, is_npu from sglang.srt.utils.common import ( ceil_align, get_eager_max_batch_size, @@ -208,6 +213,12 @@ class EagerRunner(BaseRunner): self, forward_batch: ForwardBatch, pp_proxy_tensors=None, **kwargs ) -> Any: mode = forward_batch.forward_mode + if mode.is_mixed() and not is_npu() and get_cp_strategy() is None: + # A mixed batch is extend-shaped (decode tails are 1-token + # extends); run it as EXTEND. NPU keeps MIXED for its dedicated + # kernel; CP keeps it to skip the zigzag split. + forward_batch.forward_mode = ForwardMode.EXTEND + mode = ForwardMode.EXTEND if mode.is_decode(): return self._execute_decode(forward_batch, pp_proxy_tensors) if mode.is_idle(): diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index ad83af86c..0eb3ff600 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -262,6 +262,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): """ def __init__(self, model_runner: ModelRunner): + if get_schedule().enable_mixed_chunk: + backend = get_exec().graph.cuda_graph_config.prefill.backend + assert backend == Backend.BREAKABLE, ( + "Mixed chunk prefill requires the breakable prefill CUDA " + f"graph backend; got '{backend}'." + ) super().__init__(model_runner) # --- model flags ---------------------------------------------- self.quant_config = getattr(model_runner.model, "quant_config", None) @@ -1527,8 +1533,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): else forward_batch.num_token_non_padded ) - # Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1) - # doesn't fail on MIXED=3. + # MIXED replays the EXTEND-captured graphs. pcg_forward_mode = ( ForwardMode.EXTEND if forward_batch.forward_mode == ForwardMode.MIXED