diff --git a/tests/runtime_patch/test_deep_gemm_utils.py b/tests/runtime_patch/test_deep_gemm_utils.py index fb4bda80..2628f99f 100644 --- a/tests/runtime_patch/test_deep_gemm_utils.py +++ b/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)) diff --git a/vllm_hcu/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm_hcu/model_executor/layers/fused_moe/deep_gemm_utils.py index 1f2037f0..30be0c27 100644 --- a/vllm_hcu/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/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, diff --git a/vllm_hcu/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm_hcu/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index bf3a5886..4052fbc6 100644 --- a/vllm_hcu/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/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, diff --git a/vllm_hcu/model_executor/layers/fused_moe/experts/dpsk_v4_deep_gemm_moe.py b/vllm_hcu/model_executor/layers/fused_moe/experts/dpsk_v4_deep_gemm_moe.py index a1220bad..7661b71a 100644 --- a/vllm_hcu/model_executor/layers/fused_moe/experts/dpsk_v4_deep_gemm_moe.py +++ b/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, diff --git a/vllm_hcu/model_executor/layers/quantization/slimquant_w4a8_deepgemm_runtime.py b/vllm_hcu/model_executor/layers/quantization/slimquant_w4a8_deepgemm_runtime.py index f2693984..0911e28f 100644 --- a/vllm_hcu/model_executor/layers/quantization/slimquant_w4a8_deepgemm_runtime.py +++ b/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"