[1/N][DSv41] Enable DP attention with MegaMoE

python/sglang/srt/layers/moe/mega_moe.py

@@ -33,6 +33,7 @@
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.models.deepseek_common.utils import _device_sm
+from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
 from sglang.srt.utils import is_hip, is_sm100_supported
 
 if TYPE_CHECKING:
@@ -97,14 +98,19 @@ def _configure_mega_moe_deep_gemm_num_sms(deep_gemm):
 
 
 def check_mega_moe_shapes(hidden: int, intermediate: int, mma_type: str) -> None:
-    # DeepGEMM keeps one scale row per token and needs 16-byte TMA alignment
-    # on it (layout/mega_moe.cuh), so both dims must be multiples of 16 * group.
+    # SGLang's pre-dispatch requires 16-byte input SF rows (hidden % 512).
+    # The pinned DeepGEMM FP8/FP4 path copies intermediate scales separately
+    # from token TMA loads, so intermediate % 128 suffices, including V4.1
+    # Flash's 2304-wide intermediate dimension.
+    # Keep the existing restrictions on the other MMA paths.
     scale_group = 16 if mma_type == "nvfp4xnvfp4" else 32
-    align = 16 * scale_group
-    if hidden % align != 0 or intermediate % align != 0:
+    align = 128 if mma_type == "fp8xfp4" else 16 * scale_group
+    hidden_align = 16 * scale_group
+    if hidden % hidden_align != 0 or intermediate % align != 0:
         raise ValueError(
-            f"DeepGEMM MegaMoE ({mma_type}) needs hidden_size and "
-            f"moe_intermediate_size to be multiples of {align}; got "
+            f"DeepGEMM MegaMoE ({mma_type}) needs hidden_size to be a multiple "
+            f"of {hidden_align} and moe_intermediate_size to be a multiple of "
+            f"{align}; got "
             f"hidden_size={hidden}, moe_intermediate_size={intermediate}. "
             "Use another --moe-a2a-backend for this model."
         )
@@ -256,19 +262,31 @@ def _run_mega_routed(
     if num_tokens > 0:
         router_logits = moe.gate(hidden_states, forward_batch=forward_batch)
         topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
-        topk_output = moe.topk(
-            hidden_states,
-            router_logits,
-            num_token_non_padded=(
-                forward_batch.moe_num_token_non_padded()
-                if forward_batch is not None
-                else None
-            ),
-            expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
-                layer_id=moe.layer_id,
-            ),
-            **topk_kwargs,
+        num_token_non_padded = (
+            forward_batch.moe_num_token_non_padded()
+            if forward_batch is not None
+            else None
         )
+        dispatch_info = ExpertLocationDispatchInfo.init_new(layer_id=moe.layer_id)
+        if (
+            not moe.is_hash
+            and getattr(moe.gate, "e_score_correction_bias_vl", None) is not None
+        ):
+            topk_output = vision_topk(
+                moe,
+                router_logits,
+                input_ids_global,
+                num_token_non_padded=num_token_non_padded,
+                expert_location_dispatch_info=dispatch_info,
+            )
+        else:
+            topk_output = moe.topk(
+                hidden_states,
+                router_logits,
+                num_token_non_padded=num_token_non_padded,
+                expert_location_dispatch_info=dispatch_info,
+                **topk_kwargs,
+            )
         topk_ids = topk_output.topk_ids
         topk_weights = topk_output.topk_weights
     else:

python/sglang/srt/models/deepseek_v4.py

@@ -4455,10 +4455,14 @@ def __init__(
             if (
                 get_parallel().attn_cp_size != 1
                 or get_parallel().pp_group.world_size != 1
-                or not get_moe_a2a_backend().is_none()
+                or not (
+                    get_moe_a2a_backend().is_none()
+                    or (_is_cuda and get_moe_a2a_backend().is_megamoe())
+                )
             ):
                 raise ValueError(
-                    "V4.1 vision currently supports TP/EP/DP without CP, PP or MoE A2A"
+                    "V4.1 vision supports TP/EP/DP without CP or PP; "
+                    "MoE A2A requires CUDA MegaMoE"
                 )
 
             args = SimpleNamespace(**vars(config), dim=config.hidden_size)

python/sglang/srt/models/deepseek_v4_dspark.py

@@ -257,7 +257,9 @@ def forward(
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
-        if _is_npu and forward_batch.forward_mode.is_idle():
+        # Idle ranks have no draft KV rows or attention metadata, but still
+        # participate in the shared MoE forward later in the draft stage.
+        if forward_batch.forward_mode.is_idle():
             return torch.zeros_like(hidden_states)
 
         from sglang.srt.model_executor.forward_context import get_attn_backend

python/sglang/srt/multimodal/dsv41/vl_routing.py

@@ -1,14 +1,19 @@
 import torch
 import torch.nn.functional as F
 
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.moe.topk import (
     _RENORMALIZE_SUM_EPSILON,
     StandardTopKOutput,
     StandardTopKOutputPacked,
     _mask_topk_ids_padded_region,
+    _post_process_topk_ids,
     _zero_topk_weights_padded_region,
 )
-from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots
+from sglang.srt.layers.moe.utils import (
+    get_moe_a2a_backend,
+    has_per_rank_fused_shared_slots,
+)
 from sglang.srt.utils import is_cuda
 
 
@@ -20,10 +25,17 @@ def _scale_fused_shared_weights(weights, num_fused_shared_experts, scaling_facto
     return weights
 
 
-def vision_topk(moe, logits, input_ids, num_token_non_padded=None):
+def vision_topk(
+    moe,
+    logits,
+    input_ids,
+    num_token_non_padded=None,
+    expert_location_dispatch_info=None,
+):
     config = moe.topk.topk_config
     num_fused_shared_experts = config.num_fused_shared_experts
-    if num_fused_shared_experts:
+    use_mega_moe = get_moe_a2a_backend().is_megamoe()
+    if num_fused_shared_experts and not use_mega_moe:
         # This path bypasses _post_process_topk_ids, which appends the per-rank slots.
         assert not has_per_rank_fused_shared_slots(num_fused_shared_experts), (
             "VL routing does not support per-rank fused shared slots"
@@ -37,6 +49,7 @@ def vision_topk(moe, logits, input_ids, num_token_non_padded=None):
         packed_topk = None
         if (
             num_fused_shared_experts == 0
+            and not use_mega_moe
             and get_moe_runner_backend().is_flashinfer_mxfp4()
         ):
             packed_topk = torch.empty(
@@ -60,10 +73,14 @@ def vision_topk(moe, logits, input_ids, num_token_non_padded=None):
             packed_out=packed_topk,
             sqrtsoftplus_log1p=True,
         )
-        weights = _scale_fused_shared_weights(
+        weights, indices = _finish_vision_topk(
+            moe,
+            logits,
             weights,
-            num_fused_shared_experts,
-            config.fused_shared_experts_scaling_factor,
+            indices,
+            num_token_non_padded,
+            expert_location_dispatch_info,
+            padded_rows_masked=True,
         )
         if packed_topk is not None:
             return StandardTopKOutputPacked(weights, indices, logits, packed_topk)
@@ -96,11 +113,63 @@ def vision_topk(moe, logits, input_ids, num_token_non_padded=None):
         weights = weights / (routed_sum + _RENORMALIZE_SUM_EPSILON)
     if config.apply_routed_scaling_factor_on_output:
         weights = weights * config.routed_scaling_factor
-    weights = _scale_fused_shared_weights(
-        weights, num_fused_shared_experts, config.fused_shared_experts_scaling_factor
-    )
     weights, indices = weights.float(), indices.int()
-    if num_token_non_padded is not None:
+    weights, indices = _finish_vision_topk(
+        moe,
+        logits,
+        weights,
+        indices,
+        num_token_non_padded,
+        expert_location_dispatch_info,
+    )
+    return StandardTopKOutput(weights, indices, logits)
+
+
+def _finish_vision_topk(
+    moe,
+    logits,
+    weights,
+    indices,
+    num_token_non_padded,
+    expert_location_dispatch_info,
+    padded_rows_masked=False,
+):
+    config = moe.topk.topk_config
+    use_mega_moe = get_moe_a2a_backend().is_megamoe()
+    if use_mega_moe:
+        # Use the same EPLB mapping, recording and home-rank shared slots as
+        # ordinary top-k before handing physical expert IDs to MegaMoE.
+        indices, weights, recorder_ids = _post_process_topk_ids(
+            indices,
+            weights,
+            config,
+            logits,
+            moe.layer_id,
+            num_token_non_padded=num_token_non_padded,
+            expert_location_dispatch_info=expert_location_dispatch_info,
+            padded_rows_masked=padded_rows_masked,
+        )
+        if (
+            config.num_fused_shared_experts
+            and config.apply_routed_scaling_factor_on_output
+        ):
+            # MegaMoE applies no post-expert scale when it is already in top-k.
+            weights[:, -config.num_fused_shared_experts :] = 1.0
+        if recorder_ids is not None:
+            get_global_expert_distribution_recorder().on_select_experts(
+                topk_ids=recorder_ids
+            )
+    else:
+        weights = _scale_fused_shared_weights(
+            weights,
+            config.num_fused_shared_experts,
+            config.fused_shared_experts_scaling_factor,
+        )
+    # Shared-slot remapping can overwrite padded weights; mask last so idle
+    # and partially padded DP ranks never dispatch a phantom shared expert.
+    if num_token_non_padded is not None and (
+        not padded_rows_masked or (use_mega_moe and config.num_fused_shared_experts)
+    ):
         _mask_topk_ids_padded_region(indices, num_token_non_padded)
         _zero_topk_weights_padded_region(weights, num_token_non_padded)
-    return StandardTopKOutput(weights, indices, logits)
+    return weights, indices