[BugFix][EPLB] Split unquant fused MC2 weights (vllm-project#11282)

### What this PR does / why we need it?

Fixes vllm-project#11284.

This PR fixes the unquantized fused MC2 + dynamic EPLB path for MoE
weights.

When `enable_fused_mc2 == 1` and dynamic EPLB is enabled, unquantized
MoE weights were kept as whole NZ tensors while the EPLB adaptor
expected per-expert movable weights. That could make EPLB communication
send an NZ tensor that was not split into a tensor list, which may hang
in `batch_isend_irecv`.

The fix keeps the existing NZ format cast for unquantized fused MC2
weights, but for the dynamic EPLB fused MC2 case it additionally splits
the cast weights into per-expert tensor lists, releases the original
whole tensors, and teaches the unquantized MLP path to consume either
whole tensors or tensor lists. The EPLB adaptor now uses
`w13_weight_list` and `w2_weight_list` for the unquantized fused MC2
case.

Regression coverage was added for:

- splitting unquantized fused MC2 dynamic EPLB weights after loading
- using split weight lists in the fused MC2 unquantized apply path
- accepting split weight lists in `unquant_apply_mlp`
- routing EPLB unquantized fused MC2 weights through `w13_weight_list`
and `w2_weight_list`

### Does this PR introduce _any_ user-facing change?

No. This is a bug fix for an internal NPU MoE/EPLB communication path.

### How was this patch tested?

```bash
bash format.sh ci
```

Passed.

```bash
python -m pytest -q tests/ut/eplb/adaptor/test_vllm_adaptor.py tests/ut/ops/test_fused_moe.py tests/ut/ops/test_moe_mlp.py
```

Passed on A3:

```text
43 passed, 6 skipped, 16 warnings in 2.81s
```

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@1f486d9

Signed-off-by: freyfwt <freytian1996@gmail.com>

Author: freyfwt

tests/ut/eplb/adaptor/test_vllm_adaptor.py

@@ -50,6 +50,8 @@ def test_init_fp16(self, mock_get_config, mock_func):
         self.model.quant_config = None
         adaptor = VllmEplbAdaptor(self.model)
         self.assertEqual(adaptor.expert_weight_key_per_layer[0], (QuantType.NONE, True))
+        self.assertIs(adaptor.expert_param_per_layer[0][0][0], self.mock_layer.w13_weight_list[0])
+        self.assertIs(adaptor.expert_param_per_layer[0][0][1], self.mock_layer.w2_weight_list[0])
 
     @patch("torch.empty_like", return_value=torch.zeros(16, 32))
     @patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")

tests/ut/ops/test_fused_moe.py

@@ -384,6 +384,7 @@ def _build_layer(self, *, has_bias=True, zero_expert_num=0):
     @pytest.mark.parametrize("enable_fused_mc2", [True, False])
     def test_process_weights_after_loading_transposes_and_formats(self, monkeypatch, enable_fused_mc2):
         method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod)
+        method.dynamic_eplb = False
         method._maybe_pad_weight = MagicMock(side_effect=lambda weight: weight)
         layer = self._build_layer()
         original_w13 = layer.w13_weight.detach().clone()
@@ -408,6 +409,36 @@ def test_process_weights_after_loading_transposes_and_formats(self, monkeypatch,
             assert maybe_trans_nz.call_count == 2
             format_cast.assert_not_called()
 
+    def test_process_weights_after_loading_splits_dynamic_eplb_fused_mc2_weights(self, monkeypatch):
+        method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod)
+        method.dynamic_eplb = True
+        method._maybe_pad_weight = MagicMock(side_effect=lambda weight: weight)
+        layer = nn.Module()
+        layer.w13_weight = nn.Parameter(torch.randn(2, 3, 4))
+        layer.w2_weight = nn.Parameter(torch.randn(2, 4, 3))
+        expected_w13 = layer.w13_weight.detach().clone().transpose(1, 2).contiguous()
+        expected_w2 = layer.w2_weight.detach().clone().transpose(1, 2).contiguous()
+        format_cast = MagicMock(side_effect=lambda weight, _: weight)
+        empty_cache = MagicMock()
+
+        mock_ascend_config = MagicMock()
+        mock_ascend_config.enable_fused_mc2 = True
+        monkeypatch.setattr(fused_moe_module, "get_ascend_config", lambda: mock_ascend_config)
+        monkeypatch.setattr(fused_moe_module.torch_npu, "npu_format_cast", format_cast)
+        monkeypatch.setattr(fused_moe_module.torch, "npu", SimpleNamespace(empty_cache=empty_cache), raising=False)
+
+        method.process_weights_after_loading(layer)
+
+        assert "w13_weight" not in layer._parameters
+        assert "w2_weight" not in layer._parameters
+        assert len(layer.w13_weight_list) == 2
+        assert len(layer.w2_weight_list) == 2
+        torch.testing.assert_close(layer.w13_weight_list[0], expected_w13[0])
+        torch.testing.assert_close(layer.w2_weight_list[1], expected_w2[1])
+        assert layer.w13_weight_list[0].untyped_storage().data_ptr() != expected_w13[0].untyped_storage().data_ptr()
+        assert format_cast.call_count == 2
+        empty_cache.assert_called_once()
+
     @pytest.mark.parametrize("moe_comm_type", [MoECommType.MC2, MoECommType.FUSED_MC2])
     def test_apply_builds_fused_experts_input(self, monkeypatch, moe_comm_type):
         method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod)
@@ -459,12 +490,105 @@ def test_apply_builds_fused_experts_input(self, monkeypatch, moe_comm_type):
             assert fused_input.weights.w2[0] is layer.w2_weight
             assert isinstance(fused_input.weights.w1_scale, list)
             assert isinstance(fused_input.weights.w2_scale, list)
+            assert fused_input.weights.w1_scale[0].dtype == torch.int64
+            assert fused_input.weights.w2_scale[0].dtype == torch.int64
+            assert fused_input.weights.w1_scale_bias[0].dtype == torch.float32
+            assert fused_input.weights.w2_scale_bias[0].dtype == torch.float32
         else:
             assert fused_input.weights.w1 is layer.w13_weight
             assert fused_input.weights.w2 is layer.w2_weight
             assert fused_input.weights.w1_scale is None
             assert fused_input.weights.w2_scale is None
 
+    @pytest.mark.parametrize("moe_comm_type", [MoECommType.MC2, MoECommType.FUSED_MC2])
+    def test_apply_uses_weight_lists_when_dynamic_eplb_splits_weights(self, monkeypatch, moe_comm_type):
+        method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod)
+        method.moe = SimpleNamespace(has_bias=False)
+        method.dynamic_eplb = True
+        method.tid2eid = None
+        layer = self._build_layer(has_bias=False)
+        layer.w13_weight_list = [torch.randn(4, 6), torch.randn(4, 6)]
+        layer.w2_weight_list = [torch.randn(3, 4), torch.randn(3, 4)]
+        hidden_states = torch.randn(2, 4, dtype=torch.float16)
+        topk_weights = torch.ones(2, 2, dtype=torch.float32)
+        topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64)
+        moe_comm_method = MagicMock()
+        moe_comm_method.fused_experts.return_value = torch.ones_like(hidden_states)
+        monkeypatch.setattr(
+            fused_moe_module,
+            "_EXTRA_CTX",
+            SimpleNamespace(moe_comm_type=moe_comm_type, moe_comm_method=moe_comm_method),
+        )
+        monkeypatch.setattr(fused_moe_module, "select_experts", MagicMock(return_value=(topk_weights, topk_ids)))
+        monkeypatch.setattr(fused_moe_module, "get_forward_context", MagicMock(return_value=MagicMock(input_ids=None)))
+
+        method.apply(
+            layer=layer,
+            x=hidden_states,
+            use_grouped_topk=False,
+            top_k=2,
+            router_logits=torch.randn(2, 4),
+            renormalize=True,
+            num_experts=4,
+        )
+
+        fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"]
+        assert fused_input.weights.w1 is layer.w13_weight_list
+        assert fused_input.weights.w2 is layer.w2_weight_list
+        if moe_comm_type == MoECommType.FUSED_MC2:
+            assert len(fused_input.weights.w1_scale) == 1
+            assert len(fused_input.weights.w2_scale) == 1
+            assert fused_input.weights.w1_scale[0].dtype == torch.int64
+            assert fused_input.weights.w2_scale[0].dtype == torch.int64
+            assert fused_input.weights.w1_scale[0].numel() == 0
+            assert fused_input.weights.w2_scale[0].numel() == 0
+            assert fused_input.weights.w1_scale_bias[0].dtype == torch.float32
+            assert fused_input.weights.w2_scale_bias[0].dtype == torch.float32
+            assert fused_input.weights.w1_scale_bias[0].numel() == 0
+            assert fused_input.weights.w2_scale_bias[0].numel() == 0
+        else:
+            assert fused_input.weights.w1_scale is None
+            assert fused_input.weights.w2_scale is None
+
+    def test_apply_warns_when_dynamic_eplb_fused_mc2_weights_are_not_split(self, monkeypatch):
+        method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod)
+        method.moe = SimpleNamespace(has_bias=False)
+        method.dynamic_eplb = True
+        method.tid2eid = None
+        layer = self._build_layer(has_bias=False)
+        hidden_states = torch.randn(2, 4, dtype=torch.float16)
+        topk_weights = torch.ones(2, 2, dtype=torch.float32)
+        topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64)
+        moe_comm_method = MagicMock()
+        moe_comm_method.fused_experts.return_value = torch.ones_like(hidden_states)
+        warning_once = MagicMock()
+        monkeypatch.setattr(
+            fused_moe_module,
+            "_EXTRA_CTX",
+            SimpleNamespace(moe_comm_type=MoECommType.FUSED_MC2, moe_comm_method=moe_comm_method),
+        )
+        monkeypatch.setattr(fused_moe_module, "select_experts", MagicMock(return_value=(topk_weights, topk_ids)))
+        monkeypatch.setattr(fused_moe_module, "get_forward_context", MagicMock(return_value=MagicMock(input_ids=None)))
+        monkeypatch.setattr(fused_moe_module.logger, "warning_once", warning_once)
+
+        method.apply(
+            layer=layer,
+            x=hidden_states,
+            use_grouped_topk=False,
+            top_k=2,
+            router_logits=torch.randn(2, 4),
+            renormalize=True,
+            num_experts=4,
+        )
+
+        warning_once.assert_called_once()
+        warning_msg = warning_once.call_args.args[0]
+        assert "dynamic EPLB" in warning_msg
+        assert "not split into tensor lists" in warning_msg
+        fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"]
+        assert fused_input.weights.w1[0] is layer.w13_weight
+        assert fused_input.weights.w2[0] is layer.w2_weight
+
     def test_apply_adds_zero_expert_result_and_force_balances(self, monkeypatch):
         method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod)
         method.moe = SimpleNamespace(has_bias=False)

tests/ut/ops/test_moe_mlp.py

@@ -5,7 +5,7 @@
 import torch
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 
-from vllm_ascend.ops.fused_moe.moe_mlp import cumsum_group_list, unified_apply_mlp
+from vllm_ascend.ops.fused_moe.moe_mlp import cumsum_group_list, unified_apply_mlp, unquant_apply_mlp
 from vllm_ascend.ops.fused_moe.moe_runtime_args import (
     MoEMlpComputeInput,
     MoEQuantParams,
@@ -64,6 +64,40 @@ def test_w4a8_per_channel_gmm_swiglu_flag(self):
 
 
 class TestUnifiedApplyMlpRequest(unittest.TestCase):
+    def test_unquant_apply_mlp_wraps_tensor_weights_for_grouped_matmul(self):
+        hidden_states = torch.randn(2, 8)
+        gate_up_out = torch.randn(2, 16)
+        expected = torch.randn(2, 8)
+        w1 = torch.randn(2, 8, 16)
+        w2 = torch.randn(2, 8, 8)
+
+        with (
+            patch(
+                "vllm_ascend.ops.fused_moe.moe_mlp.torch_npu.npu_grouped_matmul",
+                side_effect=[[gate_up_out], [expected]],
+                create=True,
+            ) as mock_grouped_matmul,
+            patch(
+                "vllm_ascend.ops.fused_moe.moe_mlp.torch_npu.npu_swiglu",
+                return_value=gate_up_out,
+                create=True,
+            ),
+        ):
+            output, _ = unquant_apply_mlp(
+                hidden_states=hidden_states,
+                w1=w1,
+                w2=w2,
+                group_list=torch.tensor([1, 1]),
+                need_trans=True,
+            )
+
+        self.assertTrue(output is expected)
+        first_call, second_call = mock_grouped_matmul.call_args_list
+        self.assertEqual(len(first_call.kwargs["weight"]), 1)
+        self.assertEqual(len(second_call.kwargs["weight"]), 1)
+        self.assertEqual(first_call.kwargs["weight"][0].shape, torch.Size([2, 16, 8]))
+        self.assertEqual(second_call.kwargs["weight"][0].shape, torch.Size([2, 8, 8]))
+
     def test_request_unquant_path(self):
         hidden_states = torch.randn(2, 8)
         expected = torch.randn(2, 8)

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -27,7 +27,7 @@
 
 EPLB_EXPERT_WEIGHT_NAMES = {
     (QuantType.NONE, False): ("w13_weight", "w2_weight"),
-    (QuantType.NONE, True): ("w13_weight", "w2_weight"),
+    (QuantType.NONE, True): ("w13_weight_list", "w2_weight_list"),
     (QuantType.W8A8, False): (
         "w13_weight_list",
         "w2_weight_list",

vllm_ascend/ops/fused_moe/fused_moe.py

@@ -131,9 +131,16 @@ def process_weights_after_loading(self, layer):
         # ND format (or other formats), remove this specific 'if' check and the forced
         # npu_format_cast. At that point, the operator should be able to handle weights
         # in their native format without explicit casting here.
-        if get_ascend_config().enable_fused_mc2:
+        enable_fused_mc2 = get_ascend_config().enable_fused_mc2
+        if enable_fused_mc2:
             layer.w13_weight.data = torch_npu.npu_format_cast(layer.w13_weight.data, ACL_FORMAT_FRACTAL_NZ)
             layer.w2_weight.data = torch_npu.npu_format_cast(layer.w2_weight.data, ACL_FORMAT_FRACTAL_NZ)
+            if enable_fused_mc2 == 1 and self.dynamic_eplb:
+                layer.w13_weight_list = [weight.clone() for weight in layer.w13_weight.data.unbind(dim=0)]
+                layer.w2_weight_list = [weight.clone() for weight in layer.w2_weight.data.unbind(dim=0)]
+                del layer.w13_weight
+                del layer.w2_weight
+                torch.npu.empty_cache()
         else:
             layer.w13_weight.data = maybe_trans_nz(layer.w13_weight.data)
             layer.w2_weight.data = maybe_trans_nz(layer.w2_weight.data)
@@ -231,17 +238,25 @@ def apply(
         # (due to signature constraints), we are forced to use a placeholder empty tensor.
         # This TODO tracks the requirement to update the C++ operator to accept Optional[Tensor]
         # or None for scales in non-quantized scenarios.
+        w13_weight_list = getattr(layer, "w13_weight_list", None)
+        w2_weight_list = getattr(layer, "w2_weight_list", None)
+        has_split_weight_lists = isinstance(w13_weight_list, list) and isinstance(w2_weight_list, list)
         if _EXTRA_CTX.moe_comm_type == MoECommType.FUSED_MC2:
-            w1 = [layer.w13_weight]
+            if self.dynamic_eplb and not has_split_weight_lists:
+                logger.warning_once(
+                    "FUSED_MC2 is enabled with dynamic EPLB, but unquantized MoE weights are not split into "
+                    "tensor lists. This may cause accuracy issues or communication hangs."
+                )
+            w1 = w13_weight_list if isinstance(w13_weight_list, list) else [layer.w13_weight]
+            w2 = w2_weight_list if isinstance(w2_weight_list, list) else [layer.w2_weight]
             w1_scale = [torch.tensor([], dtype=torch.int64)]
-            w2 = [layer.w2_weight]
             w2_scale = [torch.tensor([], dtype=torch.int64)]
             w1_scale_bias = [torch.tensor([], dtype=torch.float32)]
             w2_scale_bias = [torch.tensor([], dtype=torch.float32)]
         else:
-            w1 = layer.w13_weight
+            w1 = w13_weight_list if isinstance(w13_weight_list, list) else layer.w13_weight
             w1_scale = None
-            w2 = layer.w2_weight
+            w2 = w2_weight_list if isinstance(w2_weight_list, list) else layer.w2_weight
             w2_scale = None
             w1_scale_bias = None
             w2_scale_bias = None