diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index bb1b8f624d0c..20d03d8cc7f7 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 5b82f536a5eb..c8bfe3a8f8f0 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 1c6c29dc805d..69c31855f0cc 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/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 diff --git a/python/sglang/srt/multimodal/dsv41/vl_routing.py b/python/sglang/srt/multimodal/dsv41/vl_routing.py index eaedc57f2706..b03978d84e19 100644 --- a/python/sglang/srt/multimodal/dsv41/vl_routing.py +++ b/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 diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py index 6bb7736a82e3..5965ea59ba38 100644 --- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py +++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py @@ -14,7 +14,9 @@ maybe_stub_sgl_kernel() from sglang.srt.layers.moe import MoeA2ABackend, MoeRunnerBackend, mega_moe +from sglang.srt.layers.moe import topk as topk_module from sglang.srt.layers.moe.fused_moe_triton import layer as fused_moe_layer_module +from sglang.srt.layers.moe.topk import TopKConfig from sglang.srt.layers.moe.utils import draft_model_build_scope from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod from sglang.srt.runtime_context import get_context, get_exec, get_flags, get_parallel @@ -224,7 +226,9 @@ def test_mxf4_pre_dispatch_uses_typed_api(self): swiglu_limit=None, ), experts=experts, - gate=MagicMock(return_value=torch.empty((1, 8))), + gate=MagicMock( + return_value=torch.empty((1, 8)), e_score_correction_bias_vl=None + ), topk=MagicMock(return_value=topk_output), is_hash=False, num_fused_shared_experts=0, @@ -342,15 +346,143 @@ def test_run_mega_routed_experts_generic_entry(self): self.assertIs(mega_call.args[2], experts.mega_l2_weights) self.assertEqual(mega_call.kwargs.get("activation_clamp"), 7.0) - def test_shape_check_rejects_unaligned_intermediate(self): - # Qwen3-30B-A3B: intermediate 768 leaves a 24-byte scale row. - with self.assertRaisesRegex(ValueError, "multiples of 512"): - mega_moe.check_mega_moe_shapes(2048, 768, "fp8xfp4") + def test_v41_megamoe_selects_image_bias_on_local_rows(self): + """Image rows must retain their router bias after attention-TP sharding. + + Equal logits make the two biases select disjoint experts. A numerical + expert boundary exposes accidentally routing image rows as text, and + padded rows must contribute nothing. + """ + for fused_shared in (0, 1): + for scale_in_topk in (False, True): + with self.subTest( + fused_shared=fused_shared, scale_in_topk=scale_in_topk + ): + moe = self._v41_moe(fused_shared, scale_in_topk) + ids = torch.tensor([7, 99, 99]) + hidden = torch.ones((3, 4)) + batch = SimpleNamespace( + moe_num_token_non_padded=lambda: torch.tensor( + 2, dtype=torch.int32 + ) + ) + + def run_experts(_experts, x, topk_ids, weights, **kwargs): + # Routed experts multiply x by their physical ID + 1; + # the rank-1 shared slot (ID 5) is an identity expert. + factors = (topk_ids + 1).float() + factors[topk_ids == 5] = 1.0 + self.assertTrue((topk_ids[2] == -1).all()) + return x * ( + (factors * weights).sum(-1, keepdim=True) + * kwargs["routed_scaling_factor"] + ) + + with ( + get_flags().moe.override(a2a_backend=MoeA2ABackend.MEGAMOE), + get_parallel().override(moe_ep_size=2, moe_ep_rank=1), + patch.object( + mega_moe.ExpertLocationDispatchInfo, + "init_new", + return_value=None, + ), + patch.object(mega_moe, "run_mega_routed_experts", run_experts), + # The fused padded-row fill requires Triton/GPU. + patch( + "sglang.srt.multimodal.dsv41.vl_routing.is_cuda", + return_value=False, + ), + patch.object(topk_module, "_is_cuda", False), + patch.object( + topk_module, "_can_fuse_padded_region", return_value=False + ), + ): + out = mega_moe._run_mega_routed( + moe, hidden, batch, ids, num_tokens=3 + ) + text_out = mega_moe._run_mega_routed( + moe, hidden, batch, None, num_tokens=3 + ) + # Text selects logical 0/1; image selects logical 2/3. With + # a shared slot, rank-1 routed IDs shift by one (2/3 -> 3/4). + expected = torch.tensor( + [3.0 + fused_shared, 7.0 + 3.0 * fused_shared, 0.0] + )[:, None].expand_as(hidden) + torch.testing.assert_close(out, expected) + text_expected = torch.tensor( + [3.0 + fused_shared, 3.0 + fused_shared, 0.0] + )[:, None].expand_as(hidden) + torch.testing.assert_close(text_out, text_expected) + + def test_v41_megamoe_empty_rank_still_participates(self): + """An idle DP rank must reach A2A combine without running its router.""" + moe = self._v41_moe(0, False) + hidden = torch.empty((0, 4)) + + def run_experts(_experts, x, ids, weights, **kwargs): + self.assertIsNone(ids) + self.assertIsNone(weights) + # The collective boundary owns the result even on an idle rank. + return x.new_empty((0, kwargs["hidden_size"])) + + with patch.object(mega_moe, "run_mega_routed_experts", run_experts): + out = mega_moe._run_mega_routed( + moe, hidden, None, torch.empty(0, dtype=torch.long), num_tokens=0 + ) + self.assertEqual(out.shape, (0, 4)) + moe.gate.assert_not_called() + + @staticmethod + def _v41_moe(fused_shared, scale_in_topk): + return SimpleNamespace( + config=SimpleNamespace( + hidden_size=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + image_token_id=99, + ), + experts=SimpleNamespace( + should_fuse_routed_scaling_factor_in_topk=scale_in_topk, + moe_runner_config=SimpleNamespace(swiglu_limit=None), + ), + gate=MagicMock( + return_value=torch.zeros((3, 4)), + e_score_correction_bias=torch.tensor([4.0, 3.0, 0.0, 0.0]), + e_score_correction_bias_vl=torch.tensor([0.0, 0.0, 3.0, 4.0]), + ), + topk=SimpleNamespace( + topk_config=TopKConfig( + top_k=2 + fused_shared, + num_fused_shared_experts=fused_shared, + renormalize=True, + routed_scaling_factor=2.0, + apply_routed_scaling_factor_on_output=scale_in_topk, + ) + ), + is_hash=False, + num_fused_shared_experts=fused_shared, + layer_id=0, + routed_scaling_factor=2.0, + mega_shared_l1_weights=None, + mega_shared_l2_weights=None, + ) + + def test_shape_check_uses_pinned_fp8_fp4_alignment(self): + # Packed SF groups need 128 elements, not a 16-byte token SF row. + mega_moe.check_mega_moe_shapes(2048, 768, "fp8xfp4") + mega_moe.check_mega_moe_shapes(5120, 2304, "fp8xfp4") + mega_moe.check_mega_moe_shapes(512, 128, "fp8xfp4") + with self.assertRaisesRegex(ValueError, "intermediate_size.*multiple of 128"): + mega_moe.check_mega_moe_shapes(2048, 736, "fp8xfp4") + with self.assertRaisesRegex(ValueError, "hidden_size.*multiple of 512"): + mega_moe.check_mega_moe_shapes(2016, 768, "fp8xfp4") + with self.assertRaisesRegex(ValueError, "hidden_size.*multiple of 512"): + mega_moe.check_mega_moe_shapes(256, 128, "fp8xfp4") mega_moe.check_mega_moe_shapes(4096, 1536, "fp8xfp4") mega_moe.check_mega_moe_shapes(4096, 1024, "mxf4xmxf4") # 768 is a multiple of 256, so the NVFP4 (g16) rule accepts it. mega_moe.check_mega_moe_shapes(2048, 768, "nvfp4xnvfp4") - with self.assertRaisesRegex(ValueError, "multiples of 256"): + with self.assertRaisesRegex(ValueError, "intermediate_size.*multiple of 256"): mega_moe.check_mega_moe_shapes(2048, 384, "nvfp4xnvfp4") def test_mxf4_l1_uses_packed_gate_up_interleave(self): diff --git a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py index fd47b414d9f8..a93bffd2a692 100644 --- a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py +++ b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py @@ -6,10 +6,12 @@ from torch import nn from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config +from sglang.srt.layers.moe import MoeA2ABackend from sglang.srt.layers.moe.utils import ( install_shared_experts_fusion_decision, is_shared_experts_fusion_disabled, ) +from sglang.srt.models import deepseek_v4 from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM from sglang.srt.models.deepseek_v4_dspark import DeepseekV4ForCausalLMDSpark from sglang.srt.runtime_context import get_context, get_exec, get_flags, get_parallel @@ -79,6 +81,68 @@ def test_disables_shared_fusion_without_enforce(self): self.assertTrue(is_shared_experts_fusion_disabled()) self.assertFalse(get_exec().moe.disable_shared_experts_fusion) + def test_v41_vision_admits_cuda_megamoe_with_dp_attention(self): + """A vision-capable checkpoint must initialize for text or image serving + with DP attention and CUDA MegaMoE. CP, PP and other A2A paths still lack + vision routing support and must fail before constructing the tower. + """ + config = DeepSeekV4Config( + architectures=["DeepseekV4ForCausalLM"], + model_type="deepseek_v41", + vision_n_layers=1, + hidden_size=4, + tie_word_embeddings=True, + quantization_config={}, + rope_scaling={}, + compress_ratios=[], + kv_source_layer_ids=[], + index_source_layer_ids=[], + engram_layer_ids=[], + engram_num_embeddings=[], + ) + backbone = nn.Module() + backbone.embed_tokens = nn.Identity() + backbone.start_layer = 0 + backbone.end_layer = 0 + cases = ( + (MoeA2ABackend.MEGAMOE, True, 1, 1, True), + (MoeA2ABackend.NONE, True, 1, 1, True), + (MoeA2ABackend.MEGAMOE, False, 1, 1, False), + (MoeA2ABackend.DEEPEP, True, 1, 1, False), + (MoeA2ABackend.MEGAMOE, True, 2, 1, False), + (MoeA2ABackend.MEGAMOE, True, 1, 2, False), + ) + for backend, cuda, cp, pp, admitted in cases: + with self.subTest(backend=backend, cuda=cuda, cp=cp, pp=pp): + with ( + get_parallel().override( + tp_size=4, + attn_dp_size=2, + attn_tp_size=2 // cp, + attn_cp_size=cp, + moe_ep_size=4, + pp_group=SimpleNamespace(world_size=pp, is_last_rank=True), + ), + get_flags().moe.override( + a2a_backend=backend, disable_shared_experts_fusion=True + ), + patch.object(deepseek_v4, "_is_cuda", cuda), + patch.object(deepseek_v4, "ViT", return_value=nn.Identity()), + patch.object(deepseek_v4, "Aligner", return_value=nn.Identity()), + patch.object(deepseek_v4, "DeepseekV4Model", return_value=backbone), + patch.object( + deepseek_v4, "LogitsProcessor", return_value=nn.Identity() + ), + patch.object(deepseek_v4, "get_attn_tp_context"), + ): + if admitted: + model = DeepseekV4ForCausalLM(config) + self.assertIsInstance(model.vision, nn.Module) + self.assertEqual(model.image_start.shape, (4,)) + else: + with self.assertRaisesRegex(ValueError, "V4.1 vision"): + DeepseekV4ForCausalLM(config) + def test_enables_shared_fusion_when_enforced(self): self._publish(enforce=True) self.assertIsNone( diff --git a/test/registered/unit/spec/test_dspark_dp_spec_prefill_coordination.py b/test/registered/unit/spec/test_dspark_dp_spec_prefill_coordination.py index 80e90e635fef..72d36ad5e42e 100644 --- a/test/registered/unit/spec/test_dspark_dp_spec_prefill_coordination.py +++ b/test/registered/unit/spec/test_dspark_dp_spec_prefill_coordination.py @@ -6,6 +6,7 @@ import torch from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.models.deepseek_v4_dspark import DSparkAttention from sglang.srt.speculative.dp_spec_prefill_coordination import ( DPSpecPrefillCoordinationPlan, ) @@ -20,6 +21,23 @@ class TestDSparkDPSpecPrefillCoordination(CustomTestCase): + def test_v4_idle_draft_attention_returns_empty_rows_without_kv_work(self): + """A peer prefill has no local draft KV rows or attention metadata. + + Its empty HC output can have arbitrary strides; sending it through + MXFP8 projection fails before the rank can join the peer's MoE pass. + """ + attention = object.__new__(DSparkAttention) + hidden = torch.empty_strided((0, 5120), (1, 1), dtype=torch.bfloat16) + output = attention.forward( + torch.empty(0, dtype=torch.int64), + hidden, + SimpleNamespace(forward_mode=ForwardMode.IDLE), + ) + self.assertEqual(output.shape, hidden.shape) + self.assertEqual(output.dtype, hidden.dtype) + self.assertEqual(output.device, hidden.device) + def test_draft_counts_are_not_scaled_twice(self): for width in (6, 7): for rank, local_tokens in ((0, 0), (1, 3 * width), (2, 0)):