Fix Humming non-gated MoE (#49096)

tests/kernels/moe/test_moe.py

@@ -1165,6 +1165,102 @@ def test_fused_marlin_moe_non_gated(
     torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)
 
 
+@pytest.mark.parametrize(
+    "activation",
+    [
+        MoEActivation.SILU,
+        MoEActivation.RELU2_NO_MUL,
+    ],
+    ids=["gated", "non_gated"],
+)
+def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
+    pytest.importorskip("humming")
+    from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
+        HummingIndexedExperts,
+    )
+    from vllm.model_executor.layers.quantization.utils import humming_utils
+    from vllm.utils import humming
+
+    top_k, num_experts = 6, 12
+    hidden_size, intermediate_size = 2688, 1856
+    gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size
+    num_w13_stacks = 2 if activation.is_gated else 1
+    moe_config = make_dummy_moe_config(
+        num_experts=num_experts,
+        experts_per_token=top_k,
+        hidden_dim=hidden_size,
+        intermediate_size=intermediate_size,
+        activation=activation,
+    )
+
+    layer = torch.nn.Module()
+    layer.moe_config = moe_config
+    layer.params_dtype = torch.bfloat16
+
+    weight_schema = humming.ModeloptNvfp4WeightSchema()
+    for sublayer_name, shape_n, shape_k, stack_size in (
+        ("w13", gate_up_size, hidden_size, num_w13_stacks),
+        ("w2", hidden_size, intermediate_size, 1),
+    ):
+        tensor_attrs = weight_schema.get_tensors_attrs(
+            shape_n=shape_n,
+            shape_k=shape_k,
+            param_dtype=layer.params_dtype,
+            num_experts=num_experts,
+            stack_size=stack_size,
+        )
+        for tensor_name, attrs in tensor_attrs.items():
+            layer.register_parameter(
+                f"{sublayer_name}_{tensor_name}",
+                Parameter(
+                    torch.ones(
+                        attrs["shape"],
+                        dtype=attrs["dtype"],
+                        device="cuda",
+                    ),
+                    requires_grad=False,
+                ),
+            )
+
+    humming_utils.convert_to_humming_moe_kernel_format(
+        layer,
+        weight_schema=weight_schema,
+        input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16),
+    )
+
+    w13_meta, w2_meta = (layer.humming_metas[name] for name in ("w13", "w2"))
+    for meta in (w13_meta, w2_meta):
+        assert meta.a_dtype == humming.dtypes.bfloat16
+        assert meta.b_dtype == humming.dtypes.float4e2m1
+
+    assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size
+    assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size
+
+    layer.local_num_experts = layer.global_num_experts = num_experts
+    layer.hidden_size = hidden_size
+    layer.intermediate_size_per_partition = intermediate_size
+    quant_config = humming_utils.get_humming_moe_quant_config(layer)
+    experts = HummingIndexedExperts(
+        layer,
+        moe_config,
+        quant_config,
+    )
+
+    buffer_metas, _ = experts.get_buffer_metas(
+        M=1,
+        topk=top_k,
+        activation=moe_config.activation,
+    )
+    assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size
+    assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size
+    assert experts.moe_problem_size(
+        a1=torch.empty(1, hidden_size),
+        w1=torch.empty(num_experts, 1),
+        w2=torch.empty(num_experts, 1),
+        topk_ids=torch.empty(1, top_k, dtype=torch.long),
+    ) == (num_experts, 1, intermediate_size, hidden_size, top_k)
+
+
 @pytest.mark.parametrize("ep_size", [1, 2])
 def test_moe_align_block_size_opcheck(ep_size):
     num_experts = 4

vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

@@ -289,7 +289,14 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
             assert a1.size(0) == num_experts
             num_tokens = a1.size(1)
 
-        return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k
+        return (
+            meta1.num_experts,
+            num_tokens,
+            # Logical intermediate width for both gated and non-gated activations
+            self.layer.intermediate_size_per_partition,
+            meta1.shape_k,
+            top_k,
+        )
 
     def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
         from vllm.utils.humming import GemmType as HummingGemmType
@@ -327,7 +334,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
             real_shape_m = M * topk
             output_shape = (M, K)
 
-        down_input_size = N if activation.is_gated else (N * 2)
+        gate_up_size = N * (2 if activation.is_gated else 1)
+        down_input_size = N
         a_dtype = self.layer.humming_metas["w13"].a_dtype
         c_dtype = self.layer.humming_metas["w13"].c_dtype
         num_bits = a_dtype.num_bits
@@ -347,7 +355,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
                 "dtype": torch_dtype_map[a_dtype],
             },
             "gate_up_output": {
-                "shape": (real_shape_m, N * 2),
+                "shape": (real_shape_m, gate_up_size),
                 "dtype": torch_dtype_map[c_dtype],
             },
             "activation_output": {

vllm/model_executor/layers/quantization/utils/humming_utils.py

@@ -788,7 +788,7 @@ def _convert_sublayer_to_humming(
 
     shape_k_stacks = [shape_k]
     shape_n_stacks = [shape_n]
-    if sublayer_name == "w13":
+    if sublayer_name == "w13" and layer.moe_config.activation.is_gated:
         shape_n_stacks = [shape_n // 2] * 2
 
     converted_weight_schema, converted_tensors = weight_schema.convert_humming(
@@ -991,15 +991,15 @@ def convert_to_humming_moe_kernel_format(
     # Build sublayer configs from layer properties if not provided
     if sublayer_configs is None:
         is_gated = layer.moe_config.activation.is_gated
+        intermediate_size = layer.moe_config.intermediate_size_per_partition
         sublayer_configs = {
             "w13": {
-                "shape_n": layer.moe_config.intermediate_size_per_partition * 2,
+                "shape_n": intermediate_size * (2 if is_gated else 1),
                 "shape_k": layer.moe_config.hidden_dim,
             },
             "w2": {
                 "shape_n": layer.moe_config.hidden_dim,
-                "shape_k": layer.moe_config.intermediate_size_per_partition
-                * (1 if is_gated else 2),
+                "shape_k": intermediate_size,
             },
         }