diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 0786b22ea1e4..21559d1ba57a 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -756,6 +756,7 @@ def filter_batch( "dp_spec_prefill_coordination_applied", "return_logprob", "can_run_decode_cuda_graph", + "can_run_dp_draft_cuda_graph", "can_run_dp_prefill_cuda_graph", "dp_prefill_cuda_graph_max_prefix_len", "dp_padding_mode", diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 361ab1b5d1bd..f4e23f74b6a7 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -201,6 +201,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: if _is_hip: from fast_hadamard_transform import hadamard_transform @@ -1591,6 +1604,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: @@ -1599,10 +1627,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 be2a9fd1d434..84a21fc5c960 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -173,6 +173,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""" @@ -3441,9 +3481,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( @@ -3459,8 +3514,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, @@ -3509,7 +3562,7 @@ def _forward_trtllm( multi_ctas_kv_counter_buffer=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 f09036192177..52c72a0e24b2 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2471,6 +2471,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For DP attention is_extend_in_batch: bool = False can_run_decode_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_prefill_cuda_graph: bool = False dp_prefill_cuda_graph_max_prefix_len: int = 0 tbo_split_seq_index: Optional[int] = None @@ -2524,6 +2527,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 # Full DP token vector retained for Aiter MegaMoE even when the normal MLP # TP gather path stores only this rank's token count. global_spec_verify_tier_num_tokens: Optional[List[int]] = None @@ -3831,6 +3837,7 @@ def copy(self): global_num_tokens=self.global_num_tokens, global_num_tokens_for_logprob=self.global_num_tokens_for_logprob, can_run_decode_cuda_graph=self.can_run_decode_cuda_graph, + can_run_dp_draft_cuda_graph=self.can_run_dp_draft_cuda_graph, can_run_dp_prefill_cuda_graph=self.can_run_dp_prefill_cuda_graph, dp_prefill_cuda_graph_max_prefix_len=self.dp_prefill_cuda_graph_max_prefix_len, is_extend_in_batch=self.is_extend_in_batch, diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 58483c2be190..ea64aea930bf 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -34,6 +34,7 @@ get_memory, get_parallel, get_schedule, + get_spec, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.common import require_mlp_tp_gather @@ -45,6 +46,25 @@ _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 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 + 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 + spec_info = local_batch.spec_info + return spec_info is None or spec_info.cuda_graph_compatible + + def _resolve_elastic_world_dp_size( dp_size: int, *, @@ -92,6 +112,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 @@ -117,6 +138,7 @@ def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor: self.local_forward_mode, int(self.can_run_prefill_cuda_graph), self.prefill_cuda_graph_max_prefix_len, + int(self.can_run_draft_cuda_graph), ], device=device, dtype=dtype, @@ -133,6 +155,7 @@ def _get_fallback_tensor(self, device, dtype=torch.int64) -> torch.Tensor: ForwardMode.IDLE.value, # local_forward_mode 0, # can_run_prefill_cuda_graph 0, # prefill_cuda_graph_max_prefix_len + 1, # can_run_draft_cuda_graph ], device=device, dtype=dtype, @@ -215,6 +238,7 @@ def all_gather( self.is_extend_in_batch = bool(tp0_info_cpu[:, 3].max()) self.can_run_prefill_cuda_graph = bool(tp0_info_cpu[:, 6].min()) self.prefill_cuda_graph_max_prefix_len = int(tp0_info_cpu[:, 7].max()) + self.can_run_draft_cuda_graph = bool(tp0_info_cpu[:, 8].min()) if _ENABLE_METRICS_DP_ATTENTION: self.dp_cooperation_info = DPCooperationInfo.create( tp0_info_cpu[:, 5].tolist() @@ -225,6 +249,7 @@ def _update_gather_batch( batch: ScheduleBatch, mlp_sync_info: MLPSyncBatchInfo, require_mlp_tp_gather: bool, + draft_require_mlp_tp_gather: Optional[bool] = None, skip_global_metadata=False, ): if not require_mlp_tp_gather: @@ -235,6 +260,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 envs.SGLANG_ENABLE_DP_SPEC_PREFILL_COORDINATION.get(): # Fresh counts have not yet been adjusted by the coordination plan. batch.dp_spec_prefill_coordination_applied = False @@ -251,6 +289,7 @@ def _update_gather_batch( # Check forward mode for cuda graph batch.can_run_decode_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_prefill_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph batch.dp_prefill_cuda_graph_max_prefix_len = ( mlp_sync_info.prefill_cuda_graph_max_prefix_len @@ -378,6 +417,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, ): parallel = get_parallel() @@ -413,6 +453,7 @@ def prepare_mlp_sync_batch_raw( can_run_decode_cuda_graph = _local_decode_cuda_graph_vote( local_batch=local_batch, disable_cuda_graph=disable_cuda_graph ) + can_run_draft_cuda_graph = _spec_input_cuda_graph_compatible(local_batch) breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE) full_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.FULL) coordinated_prefill = breakable_prefill or full_prefill @@ -472,6 +513,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, @@ -515,6 +557,7 @@ def prepare_mlp_sync_batch_raw( batch_to_gather, mlp_sync_info, require_mlp_tp_gather, + draft_require_mlp_tp_gather, skip_global_metadata=not metadata_ready, ) @@ -546,12 +589,22 @@ 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 = get_spec().speculative_moe_a2a_backend + if draft_moe_a2a_backend is None: + draft_moe_a2a_backend = get_exec().moe.moe_a2a_backend + draft_require_mlp_tp_gather = require_mlp_tp_gather( + moe_a2a_backend=draft_moe_a2a_backend, + ) + return prepare_mlp_sync_batch_raw( local_batch, model_runner=self.model_runner, get_idle_batch=self.get_idle_batch, disable_cuda_graph=cuda_graph_fully_disabled(), require_mlp_tp_gather=require_mlp_tp_gather(), + 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 cf0382740b13..71dd85d2266d 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -558,6 +558,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): is_extend_in_batch: bool = False dp_spec_prefill_coordination_applied: bool = False can_run_decode_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_prefill_cuda_graph: bool = False dp_prefill_cuda_graph_max_prefix_len: int = 0 global_forward_mode: Optional[ForwardMode] = None @@ -835,16 +837,33 @@ 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.""" self.dp_spec_prefill_coordination_applied = ( batch.dp_spec_prefill_coordination_applied ) - 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 + # A DP spec/prefill coordination plan already wrote this phase's counts. + if ( + is_draft_worker + and batch.draft_global_num_tokens is not None + and not batch.dp_spec_prefill_coordination_applied + ): + 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 and not batch.dp_spec_prefill_coordination_applied @@ -854,15 +873,15 @@ def init_mlp_sync_metadata( 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 pin_memory = is_pin_memory_available(device) self.global_num_tokens_gpu = torch.tensor( @@ -875,6 +894,7 @@ def init_mlp_sync_metadata( pin_memory=pin_memory, ).to(device, non_blocking=True) self.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph + self.can_run_dp_draft_cuda_graph = batch.can_run_dp_draft_cuda_graph @classmethod def init_new( @@ -976,6 +996,7 @@ def init_new( return_logprob=batch.return_logprob, is_extend_in_batch=batch.is_extend_in_batch, can_run_decode_cuda_graph=batch.can_run_decode_cuda_graph, + can_run_dp_draft_cuda_graph=batch.can_run_dp_draft_cuda_graph, can_run_dp_prefill_cuda_graph=batch.can_run_dp_prefill_cuda_graph, dp_prefill_cuda_graph_max_prefix_len=batch.dp_prefill_cuda_graph_max_prefix_len, global_forward_mode=batch.global_forward_mode, @@ -1045,7 +1066,9 @@ def init_new( ).to(device, non_blocking=True) ret.global_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) @@ -1979,6 +2002,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/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index afec8205b221..37335b21cd2d 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -4,6 +4,7 @@ import torch +from sglang.srt.disaggregation.utils import get_dsa_seed_metadata_dim from sglang.srt.layers.attention.dsa.utils import ( should_remap_pd_dsa_seed_to_local_slots, ) @@ -17,6 +18,10 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch +def _requires_dsa_seed_for_cuda_graph(hf_config, topk: int) -> bool: + return topk == 1 and get_dsa_seed_metadata_dim(hf_config) > 0 + + def build_eagle_disagg_draft_input( batch: ScheduleBatch, last_tokens_tensor: torch.Tensor, @@ -87,12 +92,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, + spec.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_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 0147031ccd26..543690c11321 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -340,7 +340,9 @@ def can_run_graph(self, forward_batch: ForwardBatch): if self.require_mlp_sync: is_bs_supported = ( - is_bs_supported and forward_batch.can_run_decode_cuda_graph + is_bs_supported + and forward_batch.can_run_decode_cuda_graph + and forward_batch.can_run_dp_draft_cuda_graph ) return is_bs_supported diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 80ad187ea8cf..f28260586c5a 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -177,6 +177,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) @@ -237,19 +241,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 @@ -259,6 +275,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 af35e649e54c..55c5f902c193 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1114,6 +1114,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: @@ -1283,6 +1284,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): @@ -1600,7 +1604,12 @@ def _forward_dp_spec_prefill_coordination( """Run draft, target, and draft-extend with each rank's local mode.""" rank = get_parallel().attn_dp_rank target_local_only = len(batch.global_num_tokens) == 1 - draft_local_only = target_local_only or self.draft_worker.draft_owns_attention + # The draft MoE A2A backend may gather differently from the target. + draft_local_only = ( + target_local_only + if batch.draft_global_num_tokens is None + else len(batch.draft_global_num_tokens) == 1 + ) or self.draft_worker.draft_owns_attention is_prefill = batch.forward_mode.is_extend() draft_batch = batch if is_prefill: diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index bc0256cc8269..f72644a24d1d 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -382,6 +382,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 diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 2c016becef7b..59572210b982 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3958,11 +3958,16 @@ def with_value(self, new_value: T): self._value = None -def require_mlp_tp_gather(): +def require_mlp_tp_gather(*, 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 + + 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: @@ -3980,9 +3985,9 @@ def require_mlp_tp_gather(): 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 @@ -3992,7 +3997,7 @@ def require_mlp_tp_gather(): # reuse this flag's DP-sync bookkeeping (uniform global_num_tokens + # max-based graph bucket). See #30432 re: the misleading flag name. return True - elif get_moe_a2a_backend().is_mori() and get_bool_env_var( + elif moe_a2a_backend.is_mori() and get_bool_env_var( "SGLANG_MORI_RECV_BOUND", "false" ): # Same bookkeeping, for the same reason. Bounding mori's receive 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 09015e75cb47..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 @@ -97,5 +97,43 @@ 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", + # 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-decode", + "8", + "--max-running-requests", + "8", + "--disable-prefill-cuda-graph", + "--disable-flashinfer-autotune", + ] + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index e85c0a94e6c2..d759a364803a 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -898,6 +898,13 @@ 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( + architectures=["GlmMoeDsaForCausalLM"], + index_share_for_mtp_iteration=True, + index_topk=3, + ) + ), ) # The draft-input shape comes from the spec bag. override = get_context().override_server_args( @@ -911,6 +918,7 @@ def test_decode_input_requires_valid_seed_for_every_request(self): draft_input = build_eagle_disagg_draft_input(batch, 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, @@ -919,6 +927,34 @@ def test_decode_input_requires_valid_seed_for_every_request(self): batch.reqs[1].output_dsa_topk_indices = invalid_seed draft_input = build_eagle_disagg_draft_input(batch, 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()), + ) + 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, 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 = ( @@ -938,6 +974,13 @@ 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( + architectures=["GlmMoeDsaForCausalLM"], + 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), 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..c17cec7c407c --- /dev/null +++ b/test/registered/unit/layers/attention/test_dsa_backend_dp_padding.py @@ -0,0 +1,187 @@ +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.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") + + +class TestDSABackendDPPadding(unittest.TestCase): + def test_occupied_uneven_decode_with_idle_rank_selects_max_len(self): + with ( + patch( + "sglang.srt.layers.moe.utils.get_moe_a2a_backend", + return_value=MoeA2ABackend.NONE, + ), + patch( + "sglang.srt.layers.dp_attention.dp_gather_width", + 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 + ) + backend._pad_trtllm_sparse_page_table = MethodType( + DeepseekSparseAttnBackend._pad_trtllm_sparse_page_table, backend + ) + backend._multi_ctas_kv_counter_for = MethodType( + DeepseekSparseAttnBackend._multi_ctas_kv_counter_for, 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/layers/moe/test_aiter_megamoe_rank_sync.py b/test/registered/unit/layers/moe/test_aiter_megamoe_rank_sync.py index 7db7a01c6ad9..dfb15eee60ef 100644 --- a/test/registered/unit/layers/moe/test_aiter_megamoe_rank_sync.py +++ b/test/registered/unit/layers/moe/test_aiter_megamoe_rank_sync.py @@ -38,6 +38,7 @@ def test_mlp_sync_retains_full_vector_without_tp_gather(): tbo_split_seq_index=None, global_forward_mode=None, can_run_decode_cuda_graph=False, + can_run_draft_cuda_graph=False, can_run_prefill_cuda_graph=False, prefill_cuda_graph_max_prefix_len=0, tp0_info_cpu=torch.tensor( 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 397db24288c7..d8201c343ef3 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 ), diff --git a/test/registered/unit/managers/scheduler_components/test_dp_attn.py b/test/registered/unit/managers/scheduler_components/test_dp_attn.py index 4b348e51db8b..b70882092300 100644 --- a/test/registered/unit/managers/scheduler_components/test_dp_attn.py +++ b/test/registered/unit/managers/scheduler_components/test_dp_attn.py @@ -27,6 +27,7 @@ def test_dp1_skip_preserves_local_tbo_metadata(self): batch = SimpleNamespace( forward_mode=ForwardMode.DECODE, batch_size=lambda: 4, + spec_info=None, ) tbo_preparer = Mock() tbo_preparer.prepare_all_gather.return_value = ( 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 600751cf708c..68449056aa68 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 @@ -18,12 +18,19 @@ from sglang.srt.distributed import parallel_state from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor +from sglang.srt.managers.scheduler_components.dp_attn import MLPSyncBatchInfo from sglang.srt.model_executor.cuda_graph_config import CudaGraphConfig, PhaseConfig 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.runtime_context import get_context, get_flags, get_parallel +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 @@ -90,6 +97,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_decode_cuda_graph=True, + can_run_dp_draft_cuda_graph=False, dp_spec_prefill_coordination_applied=False, ) @@ -103,6 +111,55 @@ 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_decode_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, + prefill_cuda_graph_max_prefix_len=128, + ) + + 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(), 128) + self.assertEqual(local[8].item(), 0) + # Idle/inactive ranks stay permissive; an active incompatible rank wins + # through the all-gathered min reduction. + self.assertEqual(fallback[7].item(), 0) + self.assertEqual(fallback[8].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_decode_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_runs_padding_hook(self): pad_batch = MagicMock() diff --git a/test/registered/unit/spec/test_dp_spec_prefill_coordination.py b/test/registered/unit/spec/test_dp_spec_prefill_coordination.py index f3156d83b0fa..9f7d8958aa18 100644 --- a/test/registered/unit/spec/test_dp_spec_prefill_coordination.py +++ b/test/registered/unit/spec/test_dp_spec_prefill_coordination.py @@ -36,9 +36,11 @@ def make_batch(counts, logprobs): return SimpleNamespace( global_num_tokens=list(counts), global_num_tokens_for_logprob=list(logprobs), + draft_global_num_tokens=None, dp_spec_prefill_coordination_applied=False, is_extend_in_batch=False, can_run_decode_cuda_graph=True, + can_run_dp_draft_cuda_graph=True, can_run_dp_prefill_cuda_graph=True, ) @@ -136,6 +138,7 @@ def test_gather_resets_coordination_and_restores_pure_decode_flags(self): tbo_split_seq_index=None, global_forward_mode=ForwardMode.DECODE, can_run_decode_cuda_graph=True, + can_run_draft_cuda_graph=True, can_run_prefill_cuda_graph=False, prefill_cuda_graph_max_prefix_len=0, )