Fix MTP MoE weight loading with NVFP4 target model. (#10758)

2025-09-22 17:21:57 -04:00
parent 7ca1bea63d
commit 9c53dad809
1 changed files with 4 additions and 1 deletions
--- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
@@ -575,7 +575,10 @@ class FusedMoE(torch.nn.Module):
            )

        # Flashinfer assumes w31 format for w13_weight. Same for the scales.
-        if should_use_flashinfer_trtllm_moe():
+        if (
+            should_use_flashinfer_trtllm_moe()
+            and self.quant_method.__class__.__name__ == "ModelOptNvFp4FusedMoEMethod"
+        ):
            shard_id = {"w1": "w3", "w3": "w1", "w2": "w2"}[shard_id]

        WEIGHT_SCALE_SUPPORTED = [e.value for e in FusedMoeWeightScaleSupported]