fix(hcu): avoid double router weighting in DeepEP GEMM

tests/runtime_patch/test_deep_gemm_utils.py

@@ -141,15 +141,25 @@ def _load_deep_gemm_apply():
             torch.zeros_like(kwargs["topk_ids"], dtype=torch.int32),
             256,
         ),
-        "deepgemm_unpermute_and_reduce": lambda **_kwargs: None,
+        "deepgemm_unpermute_and_reduce": lambda **kwargs: namespace[
+            "reduce_calls"
+        ].append(kwargs),
+        "topk_weights_for_unpermute": lambda weights, apply: (
+            torch.ones_like(weights) if apply else weights
+        ),
         "m_grouped_fp8_gemm_nt_contiguous": lambda *_args, **_kwargs: None,
+        "m_grouped_w8a8_gemm_nt_contig_asm": lambda *_args, **_kwargs: None,
+        "reduce_calls": [],
     }
     exec(compile(ast.fix_missing_locations(module), source_path, "exec"), namespace)
     return namespace
 
 
 def _run_w8a8_apply(
-    namespace: dict[str, object], *, packed_weights: bool = False
+    namespace: dict[str, object],
+    *,
+    packed_weights: bool = False,
+    apply_router_weight_on_input: bool = False,
 ) -> None:
     logical_n = 16 if packed_weights else 4
     logical_k = 64 if packed_weights else 4
@@ -187,7 +197,32 @@ def _run_w8a8_apply(
         workspace13=torch.empty(32, dtype=torch.int8),
         workspace2=torch.empty(32),
         expert_tokens_meta=None,
-        apply_router_weight_on_input=False,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+    )
+
+
+def test_w8a8_apply_uses_uniform_reduce_weights_after_input_weighting(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _install_categorized_lightop(
+        monkeypatch,
+        activation_name="fuse_silu_mul_quant",
+        activation_kernel=lambda tensor, **kwargs: (
+            kwargs["output"],
+            torch.ones((tensor.shape[0], 1)),
+        ),
+        gemm_name="m_grouped_w8a8_gemm_nt_contig_asm",
+        gemm_kernel=lambda *_args: None,
+    )
+    namespace = _load_deep_gemm_apply()
+    namespace["current_platform"] = SimpleNamespace(is_rocm=lambda: True)
+    _run_w8a8_apply(namespace, apply_router_weight_on_input=True)
+
+    reduce_calls = namespace["reduce_calls"]
+    assert len(reduce_calls) == 1
+    torch.testing.assert_close(
+        reduce_calls[0]["topk_weights"],
+        torch.ones((2, 1)),
     )
 
 
@@ -446,8 +481,10 @@ def pack(_weight: torch.Tensor) -> torch.Tensor:
     assert pack_calls == 0
 
 
+@pytest.mark.parametrize("apply_router_weight_on_input", [False, True])
 def test_w4a8_contiguous_runs_two_hipc_gemms_with_expert_map_and_scales(
     monkeypatch: pytest.MonkeyPatch,
+    apply_router_weight_on_input: bool,
 ):
     """Both GEMM stages must retain scale and expert-map ownership."""
 
@@ -470,7 +507,7 @@ def test_w4a8_contiguous_runs_two_hipc_gemms_with_expert_map_and_scales(
     hidden_states = torch.arange(8, dtype=torch.int8).reshape(2, 4)
     input_scale = torch.ones((2, 1), dtype=torch.float32)
     topk_ids = torch.tensor([[0], [1]], dtype=torch.int32)
-    topk_weights = torch.ones((2, 1), dtype=torch.float32)
+    topk_weights = torch.tensor([[0.25], [0.75]], dtype=torch.float32)
     expert_map = torch.tensor([1, 0], dtype=torch.int32)
     m_indices = torch.tensor([0, 1], dtype=torch.int64)
     inv_perm = torch.tensor([1, 0], dtype=torch.int32)
@@ -540,7 +577,7 @@ def unpermute(**kwargs: object) -> None:
         workspace13=torch.empty(32, dtype=torch.int8),
         workspace2=torch.empty(32, dtype=torch.bfloat16),
         expert_tokens_meta=expert_tokens_meta,
-        apply_router_weight_on_input=False,
+        apply_router_weight_on_input=apply_router_weight_on_input,
     )
 
     assert len(gemm_calls) == 2
@@ -551,6 +588,10 @@ def unpermute(**kwargs: object) -> None:
     assert permute_call["expert_map"] is expert_map
     assert permute_call["expert_tokens_meta"] is expert_tokens_meta
     assert reduce_call["expert_map"] is expert_map
+    expected_weights = (
+        torch.ones_like(topk_weights) if apply_router_weight_on_input else topk_weights
+    )
+    torch.testing.assert_close(reduce_call["topk_weights"], expected_weights)
     torch.testing.assert_close(output, torch.full_like(output, 2))
 
 

vllm_hcu/model_executor/layers/fused_moe/deep_gemm_utils.py

@@ -637,6 +637,20 @@ def deepgemm_moe_permute(
     return aq_out, aq_scale_out, expert_ids, inv_perm, align_used
 
 
+def topk_weights_for_unpermute(
+    topk_weights: torch.Tensor,
+    apply_router_weight_on_input: bool,
+) -> torch.Tensor:
+    """Return reduction weights for an input-weighted expert computation.
+
+    The modular MoE contract applies router weights either before dispatch or
+    during unpermute/reduce, but never at both stages.
+    """
+    if apply_router_weight_on_input:
+        return torch.ones_like(topk_weights)
+    return topk_weights
+
+
 def deepgemm_unpermute_and_reduce(
     a: torch.Tensor,  # Grouped gemm output
     topk_ids: torch.Tensor,

vllm_hcu/model_executor/layers/fused_moe/experts/deep_gemm_moe.py

@@ -20,6 +20,9 @@
     deepgemm_moe_permute,
     deepgemm_unpermute_and_reduce,
 )
+from vllm_hcu.model_executor.layers.fused_moe.deep_gemm_utils import (
+    topk_weights_for_unpermute,
+)
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
@@ -495,8 +498,10 @@ def apply(
                     **gemm_kwargs,
                 )
 
-        if apply_router_weight_on_input:
-            topk_weights = torch.ones_like(topk_weights)
+        topk_weights = topk_weights_for_unpermute(
+            topk_weights,
+            apply_router_weight_on_input,
+        )
 
         deepgemm_unpermute_and_reduce(
             a=mm2_out,
@@ -719,8 +724,10 @@ def apply(
                 recipe_b=(1, self._WEIGHT_BLOCK_K),
             )
 
-        if apply_router_weight_on_input:
-            topk_weights = torch.ones_like(topk_weights)
+        topk_weights = topk_weights_for_unpermute(
+            topk_weights,
+            apply_router_weight_on_input,
+        )
 
         deepgemm_unpermute_and_reduce(
             a=mm2_out,

vllm_hcu/model_executor/layers/fused_moe/experts/dpsk_v4_deep_gemm_moe.py

@@ -27,6 +27,9 @@
     deepgemm_moe_permute,
     deepgemm_unpermute_and_reduce,
 )
+from vllm_hcu.model_executor.layers.fused_moe.deep_gemm_utils import (
+    topk_weights_for_unpermute,
+)
 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceDelegate,
@@ -341,7 +344,6 @@ def _apply_deepgemm_ht(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        del apply_router_weight_on_input
         if activation != MoEActivation.SILU:
             raise NotImplementedError(
                 "DeepEP DeepGEMM HT path currently mirrors SGLang's silu-only "
@@ -435,6 +437,10 @@ def _apply_deepgemm_ht(
             m_indices,
         )
 
+        topk_weights = topk_weights_for_unpermute(
+            topk_weights,
+            apply_router_weight_on_input,
+        )
         deepgemm_unpermute_and_reduce(
             a=down_output,
             topk_ids=topk_ids,

vllm_hcu/model_executor/layers/quantization/slimquant_w4a8_deepgemm_runtime.py

@@ -20,6 +20,9 @@
     deepgemm_moe_permute,
     deepgemm_unpermute_and_reduce,
 )
+from vllm_hcu.model_executor.layers.fused_moe.deep_gemm_utils import (
+    topk_weights_for_unpermute,
+)
 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
@@ -236,9 +239,13 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ) -> None:
-        del global_num_experts, a2_scale, apply_router_weight_on_input
+        del global_num_experts, a2_scale
         if hidden_states.size(0) == 0:
             return
+        topk_weights = topk_weights_for_unpermute(
+            topk_weights,
+            apply_router_weight_on_input,
+        )
         if activation != MoEActivation.SILU:
             raise NotImplementedError(
                 "SlimQuant W4A8 DeepGEMM supports only SiLU activation"