From 4abf816e917bf3b6f1d607f8c8a6419f6d009e95 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Thu, 23 Jul 2026 19:19:59 +0800 Subject: [PATCH 01/11] Synchronize EAGLE DSA graph fallback across DP ranks Root cause: A missing GLM-5.2 DSA seed disabled draft CUDA graph replay only on the local active rank, after DP graph eligibility had already been reduced. Idle peers could stay on the graph path and desynchronize later collectives. Fix: Carry a graph-compatibility bit with EAGLE draft inputs and fold it into the existing DP graph-admission reduction. Keep idle inputs neutral and preserve the bit through merge and overlap state. Validation: - pre-commit run --all-files --show-diff-on-failure - Python compileall for the changed modules - 4-GPU GLM-5.2 EAGLE PD warmup A/B: baseline hung; the fix completed all DP-rank warmups and reached server ready --- .../managers/scheduler_components/dp_attn.py | 21 +++++++++++++++ .../srt/speculative/eagle_disaggregation.py | 16 +++++++++++ python/sglang/srt/speculative/eagle_info.py | 27 ++++++++++++++++--- .../sglang/srt/speculative/eagle_worker_v2.py | 4 +++ 4 files changed, 64 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index af318261161a..d76fe3ec589b 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -39,6 +39,24 @@ _ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get() +def _spec_input_cuda_graph_compatible(local_batch: Optional[ScheduleBatch]) -> bool: + """Return the local spec-input graph admission bit. + + None/idle/prebuilt inputs stay permissive so an active rank with complete + runtime state can still use graphs. Any active incompatible input is + min-reduced by ``MLPSyncBatchInfo`` and forces every DP rank eager. + """ + if ( + local_batch is None + or local_batch.forward_mode.is_idle() + or local_batch.forward_mode.is_prebuilt() + ): + return True + return getattr( + getattr(local_batch, "spec_info", None), "cuda_graph_compatible", True + ) + + def _resolve_elastic_world_dp_size( dp_size: int, *, @@ -263,6 +281,9 @@ 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 + can_run_decode_cuda_graph = ( + can_run_decode_cuda_graph and _spec_input_cuda_graph_compatible(local_batch) + ) breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE) prefill_graph_runner = ( model_runner.prefill_cuda_graph_runner if breakable_prefill else None diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index 0c7b84c1c5aa..8a79dcde1a75 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -17,6 +17,14 @@ from sglang.srt.server_args import ServerArgs +def _requires_dsa_seed_for_cuda_graph(hf_config, topk: int) -> bool: + return ( + topk == 1 + and getattr(hf_config, "index_share_for_mtp_iteration", False) + and getattr(hf_config, "index_topk", None) is not None + ) + + def build_eagle_disagg_draft_input( batch: ScheduleBatch, server_args: ServerArgs, @@ -85,12 +93,20 @@ def build_eagle_disagg_draft_input( if torch.any(torch.all(dsa_topk_indices < 0, dim=1)).item(): dsa_topk_indices = None + requires_dsa_seed_for_cuda_graph = _requires_dsa_seed_for_cuda_graph( + batch.model_config.hf_config, + server_args.speculative_eagle_topk, + ) + spec_info = EagleDraftInput( topk_p=topk_p, topk_index=topk_index, hidden_states=hidden_states, bonus_tokens=last_tokens_tensor, dsa_topk_indices=dsa_topk_indices, + cuda_graph_compatible=not ( + requires_dsa_seed_for_cuda_graph and dsa_topk_indices is None + ), ) spec_info.capture_hidden_mode = CaptureHiddenMode.LAST diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index a927efc8bf48..a04f9fecf19c 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -175,6 +175,10 @@ class EagleDraftInput(SpecInput): future_indices: Optional[torch.Tensor] = None future_dsa_topk_indices_available: bool = False + # False when this input is missing model-required state for graph replay. + # DP attention synchronizes the bit across ranks before the worker runs. + cuda_graph_compatible: bool = True + def __post_init__(self): super().__init__(SpecInputType.EAGLE_DRAFT) @@ -227,19 +231,31 @@ def filter_batch( def merge_batch(self, spec_info: "EagleDraftInput"): if self.future_indices is not None: assert spec_info.future_indices is not None + self_was_empty = len(self.future_indices) == 0 + spec_info_is_empty = len(spec_info.future_indices) == 0 + if self_was_empty: + self.cuda_graph_compatible = spec_info.cuda_graph_compatible + self.future_dsa_topk_indices_available = ( + spec_info.future_dsa_topk_indices_available + ) + elif not spec_info_is_empty: + self.cuda_graph_compatible = ( + self.cuda_graph_compatible and spec_info.cuda_graph_compatible + ) + self.future_dsa_topk_indices_available = ( + self.future_dsa_topk_indices_available + and spec_info.future_dsa_topk_indices_available + ) self.future_indices = torch.cat( [self.future_indices, spec_info.future_indices] ) - self.future_dsa_topk_indices_available = ( - self.future_dsa_topk_indices_available - and spec_info.future_dsa_topk_indices_available - ) return # Detect idle stub by `topk_index` length (idle inputs have # shape[0] == 0 across all fields). Don't use `hidden_states is None`: # for STANDALONE all non-idle inputs also have None hidden_states. if len(self.topk_index) == 0: + self.cuda_graph_compatible = spec_info.cuda_graph_compatible self.hidden_states = spec_info.hidden_states self.bonus_tokens = spec_info.bonus_tokens self.topk_p = spec_info.topk_p @@ -249,6 +265,9 @@ def merge_batch(self, spec_info: "EagleDraftInput"): return if len(spec_info.topk_index) == 0: return + self.cuda_graph_compatible = ( + self.cuda_graph_compatible and spec_info.cuda_graph_compatible + ) if self.hidden_states is not None and spec_info.hidden_states is not None: self.hidden_states = torch.cat( [self.hidden_states, spec_info.hidden_states], axis=0 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 0b15eb5f11da..73a335e0314f 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -833,6 +833,7 @@ def _draft_extend_for_prefill( num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, dsa_topk_indices=prefill_dsa_topk, + cuda_graph_compatible=not (seed_from_extend and prefill_dsa_topk is None), ) def _get_dsa_extend_topk_buf(self, num_tokens: int) -> torch.Tensor: @@ -998,6 +999,9 @@ def _draft_extend_for_decode( next_draft_input.draft_probs = ret_draft_probs if self.seed_dsa_topk_from_draft_extend: next_draft_input.dsa_topk_indices = dsa_seed_topk_indices + next_draft_input.cuda_graph_compatible = not ( + self.seed_dsa_topk_from_draft_extend and dsa_seed_topk_indices is None + ) class EAGLEWorkerV2(BaseSpecWorker): From edb4b065269059e87c58959f02e9d88fc3801388 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Sun, 26 Jul 2026 02:11:52 +0800 Subject: [PATCH 02/11] Preserve draft DP token counts across MoE backends --- python/sglang/srt/managers/schedule_batch.py | 3 ++ .../managers/scheduler_components/dp_attn.py | 32 +++++++++++++++- .../srt/model_executor/forward_batch_info.py | 38 ++++++++++++++----- python/sglang/srt/utils/common.py | 13 +++++-- 4 files changed, 72 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c273912b037d..c65d47728326 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2158,6 +2158,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For DP attention global_num_tokens: Optional[List[int]] = None global_num_tokens_for_logprob: Optional[List[int]] = None + # The draft model can use a different MoE A2A backend than the target. + draft_global_num_tokens: Optional[List[int]] = None + draft_global_num_tokens_for_logprob: Optional[List[int]] = None global_spec_verify_tier_num_tokens: Optional[List[int]] = None # === Compound crossing to ForwardBatch (carry their own device tensors) === diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index d76fe3ec589b..e9c20756fd80 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -214,6 +214,7 @@ def _update_gather_batch( batch: ScheduleBatch, mlp_sync_info: MLPSyncBatchInfo, require_mlp_tp_gather: bool, + draft_require_mlp_tp_gather: Optional[bool], skip_all_gather=False, ): # TODO: handle the case when moe_dense_tp_size != 1 @@ -225,6 +226,19 @@ def _update_gather_batch( batch.global_num_tokens_for_logprob = ( mlp_sync_info.global_num_tokens_for_logprob ) + # Reuse the same all-gather result for a draft model whose A2A backend + # requires a different local/full token-count representation. + if draft_require_mlp_tp_gather is not None: + if draft_require_mlp_tp_gather: + batch.draft_global_num_tokens = mlp_sync_info.global_num_tokens + batch.draft_global_num_tokens_for_logprob = ( + mlp_sync_info.global_num_tokens_for_logprob + ) + else: + batch.draft_global_num_tokens = [mlp_sync_info.num_tokens] + batch.draft_global_num_tokens_for_logprob = [ + mlp_sync_info.num_tokens_for_logprob + ] if not skip_all_gather: batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index @@ -247,6 +261,7 @@ def prepare_mlp_sync_batch_raw( require_mlp_tp_gather: bool, disable_overlap_schedule: bool, offload_tags: set[str], + draft_require_mlp_tp_gather: Optional[bool] = None, dwdp: bool = False, ): # Check if other DP workers have running batches @@ -386,7 +401,11 @@ def prepare_mlp_sync_batch_raw( if batch_to_gather is not None: _update_gather_batch( - batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather + batch_to_gather, + mlp_sync_info, + require_mlp_tp_gather, + draft_require_mlp_tp_gather, + skip_all_gather, ) # Set on `local_batch`, not `batch_to_gather`: for PREBUILT batches the @@ -420,6 +439,16 @@ class SchedulerDPAttnAdapter: get_require_mlp_sync: Callable[[], bool] def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch): + draft_require_mlp_tp_gather = None + if self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone(): + draft_moe_a2a_backend = self.server_args.speculative_moe_a2a_backend + if draft_moe_a2a_backend is None: + draft_moe_a2a_backend = self.server_args.moe_a2a_backend + draft_require_mlp_tp_gather = require_mlp_tp_gather( + self.server_args, + moe_a2a_backend=draft_moe_a2a_backend, + ) + return prepare_mlp_sync_batch_raw( local_batch, model_runner=self.model_runner, @@ -430,6 +459,7 @@ def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch): get_idle_batch=self.get_idle_batch, disable_cuda_graph=cuda_graph_fully_disabled(), require_mlp_tp_gather=require_mlp_tp_gather(self.server_args), + draft_require_mlp_tp_gather=draft_require_mlp_tp_gather, disable_overlap_schedule=get_schedule().disable_overlap_schedule, offload_tags=self.offload_tags, dwdp=get_parallel().dwdp_size > 1, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 2d01ef258b2e..726795df1753 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -703,28 +703,40 @@ def apply_deprecated_skip_attn_backend_init( self.mark_forward_metadata_ready() def init_mlp_sync_metadata( - self, batch: ScheduleBatch, device: Union[str, torch.device] + self, + batch: ScheduleBatch, + device: Union[str, torch.device], + *, + is_draft_worker: bool = False, ) -> None: """Populate per-rank token counts for DP-attention MLP synchronization.""" - if batch.global_num_tokens is None: + global_num_tokens_source = batch.global_num_tokens + global_num_tokens_for_logprob_source = batch.global_num_tokens_for_logprob + if is_draft_worker and batch.draft_global_num_tokens is not None: + global_num_tokens_source = batch.draft_global_num_tokens + global_num_tokens_for_logprob_source = ( + batch.draft_global_num_tokens_for_logprob + ) + + if global_num_tokens_source is None: return - assert batch.global_num_tokens_for_logprob is not None + assert global_num_tokens_for_logprob_source is not None if self.spec_info is not None: from sglang.srt.speculative.spec_info import spec_scale_global_num_tokens global_num_tokens, global_num_tokens_for_logprob = ( spec_scale_global_num_tokens( self.spec_info, - batch.global_num_tokens, - batch.global_num_tokens_for_logprob, + global_num_tokens_source, + global_num_tokens_for_logprob_source, ) ) else: - global_num_tokens = batch.global_num_tokens - global_num_tokens_for_logprob = batch.global_num_tokens_for_logprob + global_num_tokens = global_num_tokens_source + global_num_tokens_for_logprob = global_num_tokens_for_logprob_source - self.original_global_num_tokens_cpu = batch.global_num_tokens + self.original_global_num_tokens_cpu = global_num_tokens_source self.global_num_tokens_cpu = global_num_tokens self.global_num_tokens_gpu = torch.tensor( global_num_tokens, dtype=torch.int64 @@ -874,7 +886,9 @@ def init_new( ) ret.num_token_non_padded_cpu = num_tokens - ret.init_mlp_sync_metadata(batch, device) + ret.init_mlp_sync_metadata( + batch, device, is_draft_worker=model_runner.is_draft_worker + ) if ret.forward_mode.is_idle(): ret.positions = torch.empty((0,), dtype=torch.int64, device=device) @@ -1591,6 +1605,12 @@ def post_forward_mlp_sync_batch(self, logits_output: LogitsProcessorOutput): ] logits_output.hidden_states = logits_output.hidden_states[:bs] elif self.forward_mode.is_extend() or self.forward_mode.is_idle(): + if self.forward_mode.is_idle(): + self.positions = self.positions[:bs] + self.seq_lens = self.seq_lens[:bs] + self.req_pool_indices = self.req_pool_indices[:bs] + if self.seq_lens_cpu is not None: + self.seq_lens_cpu = self.seq_lens_cpu[:bs] if logits_output.next_token_logits is not None: logits_output.next_token_logits = logits_output.next_token_logits[ :bs diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 749a7e96a649..ff52c2dc9f2e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3570,13 +3570,18 @@ def with_value(self, new_value: T): self._value = None -def require_mlp_tp_gather(server_args: ServerArgs): +def require_mlp_tp_gather(server_args: ServerArgs, *, moe_a2a_backend=None): """ Check if the input of MLP is obtained by all-gather rather than all-reduce. This only happens when each MLP TP group contains multiple attention DP groups. """ - from sglang.srt.layers.moe.utils import get_moe_a2a_backend + from sglang.srt.layers.moe.utils import MoeA2ABackend, get_moe_a2a_backend from sglang.srt.runtime_context import get_exec, get_parallel + if moe_a2a_backend is None: + moe_a2a_backend = get_moe_a2a_backend() + elif not isinstance(moe_a2a_backend, MoeA2ABackend): + moe_a2a_backend = MoeA2ABackend(moe_a2a_backend) + # elastic-EP scale-up rewrites dp_size on the published config if get_parallel().enable_dp_attention: assert get_parallel().dp_size > 1, "dp_size must be greater than 1" @@ -3593,9 +3598,9 @@ def require_mlp_tp_gather(server_args: ServerArgs): return True elif not get_parallel().enable_dp_lm_head: return True - elif get_moe_a2a_backend().is_none(): + elif moe_a2a_backend.is_none(): return True - elif get_moe_a2a_backend().is_flashinfer(): + elif moe_a2a_backend.is_flashinfer(): # FlashInfer MoE A2A needs a rank-invariant, DP-synchronized per-rank # token count: MoeAlltoAll uses fixed-geometry buffers and the decode # cuda-graph bucket must be identical across EP ranks, otherwise ranks From 486de181f32bb7a23d4c3aebd4f9127fb1f8ebb4 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:15:00 +0800 Subject: [PATCH 03/11] Cover DSA graph fallback in disagg draft input test build_eagle_disagg_draft_input now reads batch.model_config.hf_config to decide whether a missing DSA seed forces the draft off the CUDA graph, so the fake batch in the wire test needs a model_config. Give it one and pin both branches: index-sharing models fall back to eager when any request lacks a seed, non-sharing models keep the captured graph regardless. --- .../test_disaggregation_wire.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 0b4639661d4b..ed0e2a0ba743 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -181,6 +181,12 @@ def test_decode_input_requires_valid_seed_for_every_request(self): reqs=[self._make_req(seed) for seed in seeds], device="cpu", enable_overlap=False, + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + index_share_for_mtp_iteration=True, + index_topk=3, + ) + ), ) server_args = SimpleNamespace( speculative_eagle_topk=1, @@ -194,6 +200,7 @@ def test_decode_input_requires_valid_seed_for_every_request(self): batch, server_args, last_tokens, None ) self.assertTrue(torch.equal(draft_input.dsa_topk_indices, torch.stack(seeds))) + self.assertTrue(draft_input.cuda_graph_compatible) for invalid_seed in ( None, @@ -204,6 +211,32 @@ def test_decode_input_requires_valid_seed_for_every_request(self): batch, server_args, last_tokens, None ) self.assertIsNone(draft_input.dsa_topk_indices) + # A model that shares the DSA index across MTP iterations cannot run + # the captured graph without a seed, so the draft must fall back to + # eager on every DP rank -- not just the ranks missing the seed. + self.assertFalse(draft_input.cuda_graph_compatible) + + def test_decode_input_stays_graph_compatible_without_dsa_index_sharing(self): + """Models that don't share the DSA index keep the captured graph even + when no seed is transferred; only index_share_for_mtp_iteration models + need the seed.""" + batch = SimpleNamespace( + reqs=[self._make_req(None)], + device="cpu", + enable_overlap=False, + model_config=SimpleNamespace(hf_config=SimpleNamespace()), + ) + server_args = SimpleNamespace( + speculative_eagle_topk=1, + speculative_num_steps=5, + enable_multi_layer_eagle=False, + ) + + draft_input = build_eagle_disagg_draft_input( + batch, server_args, torch.tensor([11], dtype=torch.int64), None + ) + self.assertIsNone(draft_input.dsa_topk_indices) + self.assertTrue(draft_input.cuda_graph_compatible) def test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots(self): wire_positions = ( From d44420bf417b587e73544fd7d131775815615dc2 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:04:30 +0800 Subject: [PATCH 04/11] Avoid dynamic speculative graph compatibility lookup Root cause: The DP-attention graph admission check used nested getattr calls because only EagleDraftInput declared the compatibility flag, even though ScheduleBatch and SpecInput already provide stable typed interfaces. Fix: Add a permissive cuda_graph_compatible default to SpecInput and read ScheduleBatch.spec_info directly. EagleDraftInput keeps its dataclass field so model-specific inputs can still disable graph replay. Validation: - pre-commit run --files python/sglang/srt/speculative/spec_info.py python/sglang/srt/managers/scheduler_components/dp_attn.py --show-diff-on-failure - python3 -m compileall -q on both changed modules - git diff --check --- python/sglang/srt/managers/scheduler_components/dp_attn.py | 5 ++--- python/sglang/srt/speculative/spec_info.py | 4 ++++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index e9c20756fd80..7549bce40e3d 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -52,9 +52,8 @@ def _spec_input_cuda_graph_compatible(local_batch: Optional[ScheduleBatch]) -> b or local_batch.forward_mode.is_prebuilt() ): return True - return getattr( - getattr(local_batch, "spec_info", None), "cuda_graph_compatible", True - ) + spec_info = local_batch.spec_info + return spec_info is None or spec_info.cuda_graph_compatible def _resolve_elastic_world_dp_size( diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 7d19a18d333c..b034b7bbf75b 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -328,6 +328,10 @@ class SpecInputType(IntEnum): class SpecInput(ABC): + # Whether this input carries all state required for CUDA graph replay. + # Algorithm-specific dataclass subclasses can override this per input. + cuda_graph_compatible: bool = True + # Per-request verify lengths for the ragged-verify graphs (see # sglang.srt.speculative.ragged_verify); verify inputs of algorithms with # supports_ragged_verify() override it per step. Must stay a class-level From d6504a6e80292d85e8379869c4de5c747d02cb52 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:20:43 +0800 Subject: [PATCH 05/11] Fix fused TopK disaggregation test model fixture Root cause: The fused TopK disaggregation test builds a lightweight batch mock without model_config, but graph compatibility now reads the model HF config for every disaggregated draft input. Fix: Provide the GLM IndexShare model configuration expected by the tested fused TopK path. Validation: - uvx pre-commit run --all-files --show-diff-on-failure - python3 -m compileall -q test/registered/unit/disaggregation/test_disaggregation_wire.py --- .../unit/disaggregation/test_disaggregation_wire.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index ed0e2a0ba743..53965bbc1428 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -256,6 +256,12 @@ def test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots(self): reqs=[self._make_req(seed) for seed in wire_positions], device="cpu", enable_overlap=False, + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + index_share_for_mtp_iteration=True, + index_topk=3, + ) + ), req_pool_indices=torch.tensor([3, 1], dtype=torch.int64), req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), seq_lens=torch.tensor([4, 4], dtype=torch.int32), From 247cbab431ce75681288f701dadd5925e47d7d7b Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:18:30 +0800 Subject: [PATCH 06/11] Fix EAGLE eager DSA padding across DP ranks Root cause: Eager draft fallback can pad activations after DSA metadata is planned, leaving query and top-k rows misaligned with page tables. The draft-state compatibility vote also incorrectly affected generic graph phases. Fix: Separate draft-only DP graph eligibility, trim TRT-LLM DSA inputs to real metadata rows and restore output padding, short-circuit eager idle indexing, and add mixed-MoE PD plus uneven DP regression coverage. Validation: - uvx pre-commit run --all-files --show-diff-on-failure - GLM-5.2-NVFP4 TP4/DP4 EAGLE GSM8K: 500/500 requests completed, score 0.934, average accept length 4.1541; no DSA/CUDA/rank hang --- .../srt/batch_overlap/two_batch_overlap.py | 1 + .../srt/layers/attention/dsa/dsa_indexer.py | 32 +++- .../srt/layers/attention/dsa_backend.py | 63 +++++- python/sglang/srt/managers/schedule_batch.py | 4 + .../managers/scheduler_components/dp_attn.py | 16 +- .../srt/model_executor/forward_batch_info.py | 4 + .../srt/speculative/eagle_disaggregation.py | 2 +- .../eagle_draft_cuda_graph_runner.py | 6 +- .../test_dsa_glm52_pd_mtp_cp_layersplit.py | 35 ++++ .../attention/test_dsa_backend_dp_padding.py | 181 ++++++++++++++++++ .../layers/attention/test_dsa_indexer_idle.py | 70 +++++++ .../model_executor/test_mlp_sync_pad_unpad.py | 54 ++++++ 12 files changed, 452 insertions(+), 16 deletions(-) create mode 100644 test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py create mode 100644 test/registered/unit/layers/attention/test_dsa_indexer_idle.py diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index f26955de4963..80871b2803fa 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -737,6 +737,7 @@ def filter_batch( "is_extend_in_batch", "return_logprob", "can_run_dp_cuda_graph", + "can_run_dp_draft_cuda_graph", "can_run_dp_breakable_cuda_graph", "dp_padding_mode", "global_forward_mode", diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index eaa5cc6044bd..bd553455c5bb 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -181,6 +181,19 @@ def _broadcast_indexer_topk_from_rank0( return topk_indices +def _make_eager_idle_topk_result( + x: torch.Tensor, index_topk: int, return_indices: bool +) -> Optional[torch.Tensor]: + if not return_indices: + return None + return torch.full( + (x.shape[0], index_topk), + -1, + dtype=torch.int32, + device=x.device, + ) + + def rotate_activation(x: torch.Tensor) -> torch.Tensor: # from sgl_kernel import hadamard_transform if _is_hip: @@ -1562,6 +1575,21 @@ def forward_cuda( layer_id: int, return_indices: bool = True, ) -> Optional[torch.Tensor]: + # A padded eager IDLE rank has no real requests and therefore no valid + # DSA page-table rows to index. Return invalid rows with the physical DP + # shape so later MLP/EP collectives still agree across ranks. + x_meta = x[0] if isinstance(x, tuple) else x + if ( + _is_cuda + and forward_batch.forward_mode.is_idle() + and not get_is_capture_mode() + ): + topk_result = _make_eager_idle_topk_result( + x_meta, self.index_topk, return_indices + ) + topk_result = _broadcast_indexer_topk_from_rank0(topk_result) + return maybe_capture_indexer_topk(layer_id, topk_result) + if _is_hip: from sglang.kernels.ops.attention.dsa.tilelang_kernel import act_quant elif not _is_npu: @@ -1570,10 +1598,6 @@ def forward_cuda( if TYPE_CHECKING: assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) - # When upstream uses fused FP8 RMSNorm+quant, activations may be passed as - # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries. - x_meta = x[0] if isinstance(x, tuple) else x - in_piecewise_or_breakable_cuda_graph = ( _is_in_piecewise_or_breakable_cuda_graph() ) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 9078dbac619a..bd5afc8ef93e 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -167,6 +167,46 @@ def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tens return seqlens_32.contiguous().view(-1, 1) +def _trim_trtllm_decode_dp_padding( + q_all: torch.Tensor, + topk_indices: Optional[torch.Tensor], + real_batch_size: int, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], int]: + """Align eager decode inputs with metadata planned before DP padding.""" + physical_batch_size = q_all.shape[0] + assert real_batch_size <= physical_batch_size, ( + f"DSA metadata batch size ({real_batch_size}) exceeds q batch size " + f"({physical_batch_size})" + ) + if topk_indices is not None: + assert real_batch_size <= topk_indices.shape[0], ( + f"DSA metadata batch size ({real_batch_size}) exceeds topk batch size " + f"({topk_indices.shape[0]})" + ) + + num_padding_rows = physical_batch_size - real_batch_size + if num_padding_rows == 0: + return q_all, topk_indices, 0 + + return ( + q_all[:real_batch_size], + topk_indices[:real_batch_size] if topk_indices is not None else None, + num_padding_rows, + ) + + +def _restore_trtllm_decode_dp_padding( + output: torch.Tensor, num_padding_rows: int +) -> torch.Tensor: + """Restore the physical DP shape required by downstream MLP collectives.""" + if num_padding_rows == 0: + return output + return torch.cat( + [output, output.new_zeros((num_padding_rows, *output.shape[1:]))], + dim=0, + ) + + @dataclass(frozen=True) class DSAFlashMLAMetadata: """Metadata only needed by FlashMLA""" @@ -3193,9 +3233,24 @@ def _forward_trtllm( else: q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim) + # Eager DP attention can pad q beyond metadata that was deliberately + # planned on the real draft batch. Pad top-k to the physical q shape, + # then run decode attention only on metadata-backed rows. The output is + # restored below before downstream MLP/EP collectives. + if (self.use_fused_topk or not is_prefill) and topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) + + num_decode_padding_rows = 0 + if not is_prefill: + q_all, topk_indices, num_decode_padding_rows = ( + _trim_trtllm_decode_dp_padding( + q_all, + topk_indices, + metadata.cache_seqlens_int32.shape[0], + ) + ) + if self.use_fused_topk: - if topk_indices is not None: - topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) page_table_1 = self._get_fused_topk_page_table(topk_indices) elif is_prefill: page_table_1 = transform_index_page_table_prefill( @@ -3211,8 +3266,6 @@ def _forward_trtllm( cu_seqlens_q=metadata.cu_seqlens_q, ) else: - if topk_indices is not None: - topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) page_table_1 = transform_index_page_table_decode( page_table=metadata.page_table_1, topk_indices=topk_indices, @@ -3270,7 +3323,7 @@ def _forward_trtllm( multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, ) - return out + return _restore_trtllm_decode_dp_padding(out, num_decode_padding_rows) def _pad_topk_indices( self, topk_indices: torch.Tensor, num_tokens: int diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 21aced7091c5..9012ce8a48d1 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2133,6 +2133,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For DP attention is_extend_in_batch: bool = False can_run_dp_cuda_graph: bool = False + # Rank-consistent EAGLE draft replay gate. Keep it separate so missing + # draft-only state does not disable target verification or draft extend. + can_run_dp_draft_cuda_graph: bool = False can_run_dp_breakable_cuda_graph: bool = False tbo_split_seq_index: Optional[int] = None # Rank-consistent forward mode for the recv skipper, derived from the MLP @@ -3324,6 +3327,7 @@ def copy(self): 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_draft_cuda_graph=self.can_run_dp_draft_cuda_graph, can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_cuda_graph, is_extend_in_batch=self.is_extend_in_batch, is_prefill_only=self.is_prefill_only, diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 1e0de44a1235..356315c5665e 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -39,8 +39,10 @@ _ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get() -def _spec_input_cuda_graph_compatible(local_batch: Optional[ScheduleBatch]) -> bool: - """Return the local spec-input graph admission bit. +def _spec_input_cuda_graph_compatible( + local_batch: Optional[ScheduleBatch], +) -> bool: + """Return the local speculative-draft graph admission bit. None/idle/prebuilt inputs stay permissive so an active rank with complete runtime state can still use graphs. Any active incompatible input is @@ -102,6 +104,7 @@ class MLPSyncBatchInfo: num_tokens: int num_tokens_for_logprob: int can_run_decode_cuda_graph: bool + can_run_draft_cuda_graph: bool can_run_prefill_cuda_graph: bool is_extend_in_batch: bool local_can_run_tbo: bool @@ -125,6 +128,7 @@ def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor: int(self.local_can_run_tbo), self.local_forward_mode, int(self.can_run_prefill_cuda_graph), + int(self.can_run_draft_cuda_graph), ], device=device, dtype=dtype, @@ -140,6 +144,7 @@ def _get_fallback_tensor(self, device, dtype=torch.int64) -> torch.Tensor: 1, # local_can_run_tbo ForwardMode.IDLE.value, # local_forward_mode 0, # can_run_prefill_cuda_graph + 1, # can_run_draft_cuda_graph ], device=device, dtype=dtype, @@ -211,6 +216,7 @@ def all_gather( self.can_run_decode_cuda_graph = bool(tp0_info_cpu[:, 2].min()) self.is_extend_in_batch = bool(tp0_info_cpu[:, 3].max()) self.can_run_prefill_cuda_graph = bool(tp0_info_cpu[:, 6].min()) + self.can_run_draft_cuda_graph = bool(tp0_info_cpu[:, 7].min()) if _ENABLE_METRICS_DP_ATTENTION: self.dp_cooperation_info = DPCooperationInfo.create( tp0_info_cpu[:, 5].tolist() @@ -253,6 +259,7 @@ def _update_gather_batch( # Check forward mode for cuda graph batch.can_run_dp_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph + batch.can_run_dp_draft_cuda_graph = mlp_sync_info.can_run_draft_cuda_graph batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph @@ -303,9 +310,7 @@ 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 - can_run_decode_cuda_graph = ( - can_run_decode_cuda_graph and _spec_input_cuda_graph_compatible(local_batch) - ) + can_run_draft_cuda_graph = _spec_input_cuda_graph_compatible(local_batch) breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE) prefill_graph_runner = ( model_runner.prefill_cuda_graph_runner if breakable_prefill else None @@ -372,6 +377,7 @@ def prepare_mlp_sync_batch_raw( num_tokens=num_tokens, num_tokens_for_logprob=num_tokens_for_logprob, can_run_decode_cuda_graph=can_run_decode_cuda_graph, + can_run_draft_cuda_graph=can_run_draft_cuda_graph, can_run_prefill_cuda_graph=can_run_prefill_cuda_graph, is_extend_in_batch=is_extend_in_batch, local_can_run_tbo=local_can_run_tbo, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 15846031caa0..cf5578e0013c 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -443,6 +443,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For DP attention is_extend_in_batch: bool = False can_run_dp_cuda_graph: bool = False + # Draft-only companion to the generic DP decode graph vote. + can_run_dp_draft_cuda_graph: bool = False can_run_dp_breakable_cuda_graph: bool = False global_forward_mode: Optional[ForwardMode] = None @@ -713,6 +715,7 @@ def init_mlp_sync_metadata( 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_dp_draft_cuda_graph = batch.can_run_dp_draft_cuda_graph @classmethod def init_new( @@ -798,6 +801,7 @@ def init_new( 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_draft_cuda_graph=batch.can_run_dp_draft_cuda_graph, can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_cuda_graph, global_forward_mode=batch.global_forward_mode, is_prefill_only=batch.is_prefill_only, diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index 49f88511bc71..55007f9dd2bc 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -97,7 +97,7 @@ def build_eagle_disagg_draft_input( requires_dsa_seed_for_cuda_graph = _requires_dsa_seed_for_cuda_graph( batch.model_config.hf_config, - server_args.speculative_eagle_topk, + spec.speculative_eagle_topk, ) spec_info = EagleDraftInput( diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index a55d18a48821..2aec1df05825 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -314,7 +314,11 @@ def can_run_graph(self, forward_batch: ForwardBatch): ) 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_dp_cuda_graph + and forward_batch.can_run_dp_draft_cuda_graph + ) return is_bs_supported diff --git a/test/registered/models_e2e/test_dsa_glm52_pd_mtp_cp_layersplit.py b/test/registered/models_e2e/test_dsa_glm52_pd_mtp_cp_layersplit.py index 5bf33d972aec..2716d1b5704c 100644 --- a/test/registered/models_e2e/test_dsa_glm52_pd_mtp_cp_layersplit.py +++ b/test/registered/models_e2e/test_dsa_glm52_pd_mtp_cp_layersplit.py @@ -91,5 +91,40 @@ def setUpClass(cls): cls.launch_all() +class TestGLM52DSACacheLayerSplitMixedMoE(TestGLM52DSACacheLayerSplit): + """Cover target DeepEP with a TP-only EAGLE draft on PD decode.""" + + # This variant is a topology/control-flow smoke rather than a second full + # accuracy run. Keep enough uneven concurrent work to exercise DP padding. + gsm8k_accuracy_thres = 0.8 + gsm8k_num_questions = 32 + gsm8k_num_threads = 32 + + extra_decode_args = TestGLM52DSACacheLayerSplit.extra_decode_args + [ + "--dp-size", + "4", + "--enable-dp-attention", + "--enable-dp-lm-head", + "--ep-size", + "4", + "--moe-dense-tp-size", + "1", + "--moe-a2a-backend", + "deepep", + "--deepep-mode", + "low_latency", + "--speculative-moe-a2a-backend", + "none", + "--speculative-moe-runner-backend", + "triton", + "--cuda-graph-max-bs", + "8", + "--max-running-requests", + "8", + "--disable-prefill-cuda-graph", + "--disable-flashinfer-autotune", + ] + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py new file mode 100644 index 000000000000..78b628d9d8b1 --- /dev/null +++ b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py @@ -0,0 +1,181 @@ +import unittest +from types import MethodType, ModuleType, SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.attention.dsa_backend import ( + DeepseekSparseAttnBackend, + _restore_trtllm_decode_dp_padding, + _trim_trtllm_decode_dp_padding, +) +from sglang.srt.layers.dp_attention import DpPaddingMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestDSABackendDPPadding(unittest.TestCase): + def test_occupied_uneven_decode_with_idle_rank_selects_max_len(self): + moe_backend = SimpleNamespace(is_pplx=lambda: False) + with ( + patch( + "sglang.srt.layers.moe.utils.get_moe_a2a_backend", + return_value=moe_backend, + ), + patch( + "sglang.srt.layers.dp_attention.get_attention_dp_size", + return_value=4, + ), + ): + mode = DpPaddingMode.get_dp_padding_mode( + is_extend_in_batch=False, + global_num_tokens=[8, 7, 6, 0], + ) + self.assertEqual(mode, DpPaddingMode.MAX_LEN) + + def test_trim_and_restore_active_eager_decode_padding(self): + q_all = torch.arange(8 * 2 * 3).view(8, 2, 3) + topk_indices = torch.arange(8 * 4, dtype=torch.int32).view(8, 4) + + real_q, real_topk, num_padding_rows = _trim_trtllm_decode_dp_padding( + q_all, + topk_indices, + real_batch_size=6, + ) + + self.assertTrue(torch.equal(real_q, q_all[:6])) + self.assertTrue(torch.equal(real_topk, topk_indices[:6])) + self.assertEqual(num_padding_rows, 2) + + real_output = torch.ones((6, 1, 2, 5), dtype=torch.bfloat16) + output = _restore_trtllm_decode_dp_padding(real_output, num_padding_rows) + self.assertEqual(output.shape, (8, 1, 2, 5)) + self.assertTrue(torch.equal(output[:6], real_output)) + self.assertTrue(torch.all(output[6:] == 0)) + + def test_no_padding_preserves_existing_tensors(self): + q_all = torch.empty((2, 2, 3)) + topk_indices = torch.empty((2, 8), dtype=torch.int32) + + real_q, real_topk, num_padding_rows = _trim_trtllm_decode_dp_padding( + q_all, + topk_indices, + real_batch_size=2, + ) + + self.assertIs(real_q, q_all) + self.assertIs(real_topk, topk_indices) + self.assertEqual(num_padding_rows, 0) + self.assertIs( + _restore_trtllm_decode_dp_padding(real_q, num_padding_rows), + real_q, + ) + + def test_rejects_metadata_larger_than_physical_batch(self): + with self.assertRaisesRegex( + AssertionError, "metadata batch size \\(3\\) exceeds q batch size \\(2\\)" + ): + _trim_trtllm_decode_dp_padding( + torch.empty((2, 2, 3)), + torch.empty((2, 8), dtype=torch.int32), + real_batch_size=3, + ) + + def test_trtllm_decode_runs_real_rows_then_restores_physical_batch(self): + metadata = SimpleNamespace( + cache_seqlens_int32=torch.tensor([8, 12], dtype=torch.int32), + page_table_1=torch.zeros((2, 12), dtype=torch.int32), + max_seq_len_k=12, + ) + backend = SimpleNamespace( + forward_metadata=metadata, + kv_cache_dtype=torch.bfloat16, + token_to_kv_pool=SimpleNamespace( + get_key_buffer=lambda _layer_id: torch.zeros((24, 3)) + ), + real_page_size=1, + kv_cache_dim=3, + use_fused_topk=False, + qk_nope_head_dim=2, + kv_lora_rank=2, + qk_rope_head_dim=1, + workspace_buffer=None, + dsa_index_topk=2, + _multi_ctas_kv_counter_buffer=None, + device="cpu", + num_q_heads=2, + ) + backend._pad_topk_indices = MethodType( + DeepseekSparseAttnBackend._pad_topk_indices, backend + ) + + layer = SimpleNamespace( + layer_id=0, + tp_q_head_num=2, + head_dim=3, + k_scale_float=None, + scaling=1.0, + ) + forward_batch = SimpleNamespace() + q = torch.arange(4 * 2 * 3, dtype=torch.float32).view(4, 2, 3) + topk_indices = torch.arange(4 * 2, dtype=torch.int32).view(4, 2) + + flashinfer = ModuleType("flashinfer") + flashinfer_decode = ModuleType("flashinfer.decode") + captured = {} + + def fake_decode(**kwargs): + captured.update(kwargs) + return torch.ones((2, 1, 2, 2), dtype=torch.bfloat16) + + flashinfer_decode.trtllm_batch_decode_with_kv_cache_mla = fake_decode + flashinfer.decode = flashinfer_decode + + def fake_transform(*, page_table, topk_indices, page_size): + self.assertEqual(page_table.shape[0], 2) + self.assertEqual(topk_indices.shape[0], 2) + self.assertEqual(page_size, 1) + return torch.zeros((2, 2), dtype=torch.int32) + + with ( + patch.dict( + "sys.modules", + { + "flashinfer": flashinfer, + "flashinfer.decode": flashinfer_decode, + }, + ), + patch( + "sglang.srt.layers.attention.dsa_backend.transform_index_page_table_decode", + side_effect=fake_transform, + ), + patch( + "sglang.srt.layers.attention.dsa_backend.dsa_use_prefill_cp", + return_value=False, + ), + patch( + "sglang.srt.layers.attention.dsa_backend.grow_multi_ctas_kv_counter_buffer_if_needed", + return_value=None, + ), + ): + output = DeepseekSparseAttnBackend._forward_trtllm( + backend, + q=q, + k=torch.empty((4, 1, 2)), + v=torch.empty((4, 1, 2)), + layer=layer, + forward_batch=forward_batch, + seq_lens=metadata.cache_seqlens_int32, + save_kv_cache=False, + topk_indices=topk_indices, + ) + + self.assertEqual(captured["query"].shape, (2, 1, 2, 3)) + self.assertEqual(output.shape, (4, 1, 2, 2)) + self.assertTrue(torch.all(output[:2] == 1)) + self.assertTrue(torch.all(output[2:] == 0)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/attention/test_dsa_indexer_idle.py b/test/registered/unit/layers/attention/test_dsa_indexer_idle.py new file mode 100644 index 000000000000..3af04bf07cd0 --- /dev/null +++ b/test/registered/unit/layers/attention/test_dsa_indexer_idle.py @@ -0,0 +1,70 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.attention.dsa.dsa_indexer import ( + Indexer, + _make_eager_idle_topk_result, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + +_INDEXER = "sglang.srt.layers.attention.dsa.dsa_indexer" + + +class TestDSAIndexerIdle(unittest.TestCase): + def test_builds_invalid_topk_rows_for_padded_idle_tokens(self): + result = _make_eager_idle_topk_result( + torch.empty((8, 16)), + index_topk=4, + return_indices=True, + ) + + self.assertEqual(result.shape, (8, 4)) + self.assertEqual(result.dtype, torch.int32) + self.assertTrue(torch.all(result == -1)) + + def test_returns_none_when_indices_are_not_consumed(self): + self.assertIsNone( + _make_eager_idle_topk_result( + torch.empty((8, 16)), + index_topk=4, + return_indices=False, + ) + ) + + def test_eager_idle_short_circuits_before_paged_mqa(self): + indexer = SimpleNamespace(index_topk=4) + batch = SimpleNamespace(forward_mode=ForwardMode.IDLE) + + with ( + patch(f"{_INDEXER}._is_cuda", True), + patch(f"{_INDEXER}.get_is_capture_mode", return_value=False), + patch( + f"{_INDEXER}._broadcast_indexer_topk_from_rank0", + side_effect=lambda result: result, + ), + patch( + f"{_INDEXER}.maybe_capture_indexer_topk", + side_effect=lambda _layer_id, result: result, + ), + ): + result = Indexer.forward_cuda( + indexer, + x=torch.empty((8, 16)), + q_lora=torch.empty((8, 16)), + positions=torch.empty((8,), dtype=torch.int64), + forward_batch=batch, + layer_id=3, + ) + + self.assertEqual(result.shape, (8, 4)) + self.assertTrue(torch.all(result == -1)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py index 55cf84243088..cb693e5a2482 100644 --- a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py +++ b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py @@ -15,10 +15,17 @@ import torch +from sglang.srt.managers.scheduler_components.dp_attn import MLPSyncBatchInfo from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( DecodeCudaGraphRunner, ) +from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( + EAGLEDraftCudaGraphRunner, +) +from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( + EAGLEDraftExtendCudaGraphRunner, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -60,6 +67,7 @@ def test_init_mlp_sync_metadata_scales_speculative_request_width(self): global_num_tokens=[2, 0, 3], global_num_tokens_for_logprob=[2, 0, 3], can_run_dp_cuda_graph=True, + can_run_dp_draft_cuda_graph=False, ) fb.init_mlp_sync_metadata(batch, torch.device("cpu")) @@ -72,6 +80,52 @@ def test_init_mlp_sync_metadata_scales_speculative_request_width(self): fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6]) ) self.assertTrue(fb.can_run_dp_cuda_graph) + self.assertFalse(fb.can_run_dp_draft_cuda_graph) + + def test_draft_graph_gate_has_an_independent_dp_vote(self): + sync_info = MLPSyncBatchInfo( + dp_size=1, + tp_size=1, + cp_size=1, + num_tokens=1, + num_tokens_for_logprob=1, + can_run_decode_cuda_graph=True, + can_run_draft_cuda_graph=False, + can_run_prefill_cuda_graph=False, + is_extend_in_batch=False, + local_can_run_tbo=True, + local_forward_mode=ForwardMode.DECODE.value, + ) + + local = sync_info._get_local_tensor(device="cpu") + fallback = sync_info._get_fallback_tensor(device="cpu") + + self.assertEqual(local[2].item(), 1) + self.assertEqual(local[7].item(), 0) + # Idle/inactive ranks stay permissive; an active incompatible rank wins + # through the all-gathered min reduction. + self.assertEqual(fallback[7].item(), 1) + + def test_draft_only_gate_does_not_disable_draft_extend_graph(self): + draft_runner = object.__new__(EAGLEDraftCudaGraphRunner) + draft_extend_runner = object.__new__(EAGLEDraftExtendCudaGraphRunner) + for runner in (draft_runner, draft_extend_runner): + runner.require_mlp_tp_gather = False + runner.require_mlp_sync = True + runner.disable_padding = False + runner.captured_req_width = 1 + runner.max_bs = 8 + + forward_batch = SimpleNamespace( + spec_info=SimpleNamespace(num_tokens_per_req=1), + batch_size=1, + seq_lens=torch.ones(1), + can_run_dp_cuda_graph=True, + can_run_dp_draft_cuda_graph=False, + ) + + self.assertFalse(draft_runner.can_run_graph(forward_batch)) + self.assertTrue(draft_extend_runner.can_run_graph(forward_batch)) def test_draft_input_without_hidden_states_can_be_padded(self): spec_info = SimpleNamespace( From 2d42a6dc83a0d06d729f9d769bc42452d7f15a02 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:48:38 +0800 Subject: [PATCH 07/11] Fix EAGLE disaggregation test model fixtures Root cause: The DSA seed-transfer fixtures did not identify their mocked HF configs as DSA models after the model-config validation was tightened. A second test also retained the old four-argument draft-input builder call after the builder moved to runtime spec state. Fix: Mark the DSA fixtures with the GLM DSA architecture and update the non-index-sharing case to install runtime spec overrides and use the current builder signature. Validation: - uvx pre-commit run --all-files --show-diff-on-failure - python3 -m compileall -q test/registered/unit/disaggregation/test_disaggregation_wire.py - git diff --check - no added getattr calls --- .../unit/disaggregation/test_disaggregation_wire.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index b9e49d1844ea..d5c410c6fb9d 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -383,6 +383,7 @@ def test_decode_input_requires_valid_seed_for_every_request(self): enable_overlap=False, model_config=SimpleNamespace( hf_config=SimpleNamespace( + architectures=["GlmMoeDsaForCausalLM"], index_share_for_mtp_iteration=True, index_topk=3, ) @@ -424,14 +425,16 @@ def test_decode_input_stays_graph_compatible_without_dsa_index_sharing(self): enable_overlap=False, model_config=SimpleNamespace(hf_config=SimpleNamespace()), ) - server_args = SimpleNamespace( + override = get_context().override_server_args( speculative_eagle_topk=1, speculative_num_steps=5, enable_multi_layer_eagle=False, ) + override.install() + self.addCleanup(override.restore) draft_input = build_eagle_disagg_draft_input( - batch, server_args, torch.tensor([11], dtype=torch.int64), None + batch, torch.tensor([11], dtype=torch.int64), None ) self.assertIsNone(draft_input.dsa_topk_indices) self.assertTrue(draft_input.cuda_graph_compatible) @@ -456,6 +459,7 @@ def test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots(self): enable_overlap=False, model_config=SimpleNamespace( hf_config=SimpleNamespace( + architectures=["GlmMoeDsaForCausalLM"], index_share_for_mtp_iteration=True, index_topk=3, ) From 04c51a7ff476e7779fe6ee39638ed4c154e86763 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:03:54 +0800 Subject: [PATCH 08/11] Use complete MoE backend enum in DSA padding test Root cause: The test backend stub only implemented is_pplx, while current DP padding selection also queries is_deepep_v2. Fix: Use MoeA2ABackend.NONE so the test exercises the production backend interface. Validation: Verified the production enum predicates and MAX_LEN selection on CPU. pre-commit run --all-files --show-diff-on-failure passed. --- .../unit/layers/attention/test_dsa_backend_dp_padding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py index 78b628d9d8b1..efdef1557ecd 100644 --- a/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py +++ b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py @@ -10,6 +10,7 @@ _trim_trtllm_decode_dp_padding, ) from sglang.srt.layers.dp_attention import DpPaddingMode +from sglang.srt.layers.moe.utils import MoeA2ABackend from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=2, suite="base-a-test-cpu") @@ -17,11 +18,10 @@ class TestDSABackendDPPadding(unittest.TestCase): def test_occupied_uneven_decode_with_idle_rank_selects_max_len(self): - moe_backend = SimpleNamespace(is_pplx=lambda: False) with ( patch( "sglang.srt.layers.moe.utils.get_moe_a2a_backend", - return_value=moe_backend, + return_value=MoeA2ABackend.NONE, ), patch( "sglang.srt.layers.dp_attention.get_attention_dp_size", From 4532117bcee1f5b3532daa9b469831d6c7139837 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:54:28 +0800 Subject: [PATCH 09/11] Bind TRT-LLM sparse page-table helper in DSA test Root cause: The lightweight test backend invoked the production TRT-LLM forward path without binding its sparse page-table padding helper. Fix: Bind DeepseekSparseAttnBackend._pad_trtllm_sparse_page_table alongside the existing top-k helper. Validation: Verified the production helper binding and 3-to-4 sparse page-table padding on CPU. pre-commit run --all-files --show-diff-on-failure passed. --- .../unit/layers/attention/test_dsa_backend_dp_padding.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py index efdef1557ecd..f85468583070 100644 --- a/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py +++ b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py @@ -109,6 +109,9 @@ def test_trtllm_decode_runs_real_rows_then_restores_physical_batch(self): backend._pad_topk_indices = MethodType( DeepseekSparseAttnBackend._pad_topk_indices, backend ) + backend._pad_trtllm_sparse_page_table = MethodType( + DeepseekSparseAttnBackend._pad_trtllm_sparse_page_table, backend + ) layer = SimpleNamespace( layer_id=0, From 5fbc285d89ab63b78ac6dc0b23e951f3e229e87a Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Mon, 28 Sep 2026 11:19:24 +0800 Subject: [PATCH 10/11] Add is_draft_worker to mock ModelRunner in mamba prefill track test ForwardBatch.init_new now reads model_runner.is_draft_worker. --- test/registered/unit/layers/test_mamba_prefill_track_metadata.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/layers/test_mamba_prefill_track_metadata.py b/test/registered/unit/layers/test_mamba_prefill_track_metadata.py index dee5d7d5113a..304db6131b50 100644 --- a/test/registered/unit/layers/test_mamba_prefill_track_metadata.py +++ b/test/registered/unit/layers/test_mamba_prefill_track_metadata.py @@ -221,6 +221,7 @@ def test_forward_snapshot_and_padding_do_not_mutate_scheduler_lists(self): ) runner = SimpleNamespace( device="cpu", + is_draft_worker=False, model_config=SimpleNamespace( requires_mm_token_modalities=False, model_is_mrope=False ), From eb2f4242fa457dacf724ef96db9bafbc4b36eb28 Mon Sep 17 00:00:00 2001 From: weireweire <20922698+weireweire@users.noreply.github.com> Date: Mon, 28 Sep 2026 14:22:20 +0800 Subject: [PATCH 11/11] Fix MixedMoE e2e decode args for current server args --cuda-graph-max-bs is now ambiguous (split into -decode/-prefill), and NVFP4 MoE with DeepEP needs the flashinfer_cutedsl runner: the auto flashinfer_trtllm runner has no DeepEP fused func. --- .../e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/registered/e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py b/test/registered/e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py index 39528f85dd73..37611054b128 100644 --- a/test/registered/e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py +++ b/test/registered/e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py @@ -119,11 +119,14 @@ class TestGLM52DSACacheLayerSplitMixedMoE(TestGLM52DSACacheLayerSplit): "deepep", "--deepep-mode", "low_latency", + # NVFP4 MoE registers a DeepEP fused func only for the CuteDSL runner. + "--moe-runner-backend", + "flashinfer_cutedsl", "--speculative-moe-a2a-backend", "none", "--speculative-moe-runner-backend", "triton", - "--cuda-graph-max-bs", + "--cuda-graph-max-bs-decode", "8", "--max-running-requests", "8",