From 3810f531a89f050bdceb5597bf6550bc7f4f5100 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Fri, 18 Sep 2026 19:13:10 +0800
Subject: [PATCH] Support V4.1 decode vision MegaMoE

(cherry picked from commit 98aa79972d6d31f30a3c5bc83b9d7e3009ae638a)
---
 python/sglang/srt/layers/moe/mega_moe.py | 39 ++++++++++++++++--------
 python/sglang/srt/models/deepseek_v4.py  | 16 ++++++++--
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py
index 783ac819e..8cbeffc38 100644
--- a/python/sglang/srt/layers/moe/mega_moe.py
+++ b/python/sglang/srt/layers/moe/mega_moe.py
@@ -216,20 +216,33 @@ 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.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.num_token_non_padded if forward_batch is not None else None
         )
+        if isinstance(
+            getattr(moe.gate, "e_score_correction_bias_vl", None), torch.Tensor
+        ):
+            # V4.1 uses a different correction bias for image-token rows. The
+            # MegaMoE transport consumes the same routed ids/weights as TopK.
+            from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
+
+            topk_output = vision_topk(
+                moe,
+                router_logits,
+                input_ids_global,
+                num_token_non_padded=num_token_non_padded,
+            )
+        else:
+            topk_kwargs = {"input_ids": input_ids_global} if moe.is_hash else {}
+            topk_output = moe.topk(
+                hidden_states,
+                router_logits,
+                num_token_non_padded=num_token_non_padded,
+                expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
+                    layer_id=moe.layer_id,
+                ),
+                **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 69ae59dcb..3a352c484 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -182,6 +182,7 @@ from sglang.srt.multimodal.deepseek_v41_image_processing import (
 )
 from sglang.srt.runtime_context import (
     get_device,
+    get_disagg,
     get_exec,
     get_forward,
     get_parallel,
@@ -4858,6 +4859,13 @@ class DeepseekV4Model(nn.Module):
         return hidden_states, pre_hc_head
 
 
+def _v41_vision_a2a_supported() -> bool:
+    backend = get_moe_a2a_backend()
+    return backend.is_none() or (
+        backend.is_megamoe() and get_disagg().disaggregation_mode == "decode"
+    )
+
+
 class DeepseekV4ForCausalLM(nn.Module):
     supports_cuda_vmm_feature_transport = True
 
@@ -4889,11 +4897,13 @@ class DeepseekV4ForCausalLM(nn.Module):
             and not getattr(config, "language_model_only", False)
         ):
             if (
-                get_parallel().pp_group.world_size != 1
-                or not get_moe_a2a_backend().is_none()
+                get_parallel().attn_cp_size != 1
+                or get_pp_group().world_size != 1
+                or not _v41_vision_a2a_supported()
             ):
                 raise ValueError(
-                    "V4.1 vision supports TP/EP/DP and prefill CP without PP or MoE A2A"
+                    "V4.1 vision supports TP/EP/DP without CP or PP; "
+                    "MoE A2A is supported only with MegaMoE on a PD decode node"
                 )
 
             args = SimpleNamespace(**vars(config), dim=config.hidden_size)