[Bugfix] Support non-gated MoE in online quantization and Marlin MoE tile padding

tests/kernels/quantization/test_marlin_tile_padding.py

@@ -186,6 +186,23 @@ def test_marlin_moe_pad_helpers_shapes():
     assert bias_shards[..., N:].abs().sum() == 0
 
 
+@pytest.mark.parametrize("num_shards", [1, 2])
+def test_moe_pad_shard_rows(num_shards):
+    """_moe_pad_shard_rows must pad each w13 shard independently: 2 gate/up
+    shards for gated MoE, a single up shard for non-gated MoE."""
+    from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+        _moe_pad_shard_rows,
+    )
+
+    E, N, padded_N, K = 2, 96, 128, 16
+    w13 = torch.arange(E * num_shards * N * K).reshape(E, num_shards * N, K).float()
+    padded = _moe_pad_shard_rows(w13, N, padded_N, num_shards)
+    assert padded.shape == (E, num_shards * padded_N, K)
+    shards = padded.view(E, num_shards, padded_N, K)
+    assert torch.equal(shards[:, :, :N], w13.view(E, num_shards, N, K))
+    assert shards[:, :, N:].abs().sum() == 0
+
+
 def _gpu_marlin_unsupported() -> bool:
     return not (
         current_platform.is_cuda() and current_platform.has_device_capability(80)
@@ -702,14 +719,17 @@ def make_layer(hidden, intermediate):
     _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
     reason="FP8 Marlin is not supported on this GPU type.",
 )
+@pytest.mark.parametrize("gated", [True, False])
 @pytest.mark.parametrize("quant", ["channel", "tensor"])
 @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
-def test_fp8_marlin_moe_padded_round_trip(shape, quant):
+def test_fp8_marlin_moe_padded_round_trip(shape, quant, gated):
     """FP8 weight-only MoE: pad a tile-misaligned intermediate and check the
-    real prepare + fused_marlin_moe against the dequantized reference."""
+    real prepare + fused_marlin_moe against the dequantized reference. Covers
+    gated (w13 = gate/up shards) and non-gated (w13 = up only) layouts."""
     from tests.kernels.utils import torch_experts
     from vllm.config import VllmConfig, set_current_vllm_config
     from vllm.model_executor.layers.fused_moe import fused_topk
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
     from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
         fused_marlin_moe,
     )
@@ -726,6 +746,8 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant):
     fp8 = torch.float8_e4m3fn
     dtype = torch.bfloat16
     device = torch.device("cuda")
+    w13_up_dim = 2 * n if gated else n
+    activation = MoEActivation.SILU if gated else MoEActivation.RELU2_NO_MUL
     padded_n = marlin_moe_padded_intermediate(n, -1)
     assert padded_n != n
 
@@ -737,8 +759,9 @@ def q(w):  # (out, in) -> fp8 weight, scale, dequant reference
         s = s.reshape(1) if quant == "tensor" else s.squeeze(1)
         return wq, s, ref
 
-    a = torch.randn((m, k), device=device, dtype=dtype) / 10
-    w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
+    torch.manual_seed(0)
+    a = torch.randn((m, k), device=device, dtype=dtype)
+    w1 = torch.randn((e, w13_up_dim, k), device=device, dtype=dtype) / k**0.5
     w2 = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
     w13_q, w13_s, w1_ref = zip(*(q(w1[i]) for i in range(e)))
     w2_q, w2_s, w2_ref = zip(*(q(w2[i]) for i in range(e)))
@@ -770,6 +793,7 @@ def q(w):  # (out, in) -> fp8 weight, scale, dequant reference
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
+        activation=activation,
         is_k_full=True,
         workspace=layer.workspace,
     )
@@ -781,6 +805,7 @@ def q(w):  # (out, in) -> fp8 weight, scale, dequant reference
             topk_weight=topk_weights,
             topk_ids=topk_ids,
             global_num_experts=e,
+            activation=activation,
         )
     torch.testing.assert_close(out, ref, atol=8e-2, rtol=0)
 
@@ -789,13 +814,16 @@ def q(w):  # (out, in) -> fp8 weight, scale, dequant reference
     _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
     reason="FP8 Marlin is not supported on this GPU type.",
 )
+@pytest.mark.parametrize("gated", [True, False])
 @pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
-def test_mxfp8_marlin_moe_padded_round_trip(shape):
+def test_mxfp8_marlin_moe_padded_round_trip(shape, gated):
     """MXFP8 weight-only MoE round-trip at a tile-misaligned intermediate, with
-    unit e8m0 scales so the reference is the exact fp8 dequant."""
+    unit e8m0 scales so the reference is the exact fp8 dequant. Covers gated
+    (w13 = gate/up shards) and non-gated (w13 = up only) layouts."""
     from tests.kernels.utils import torch_experts
     from vllm.config import VllmConfig, set_current_vllm_config
     from vllm.model_executor.layers.fused_moe import fused_topk
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
     from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
         fused_marlin_moe,
     )
@@ -812,16 +840,19 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
     fp8 = torch.float8_e4m3fn
     dtype = torch.bfloat16
     device = torch.device("cuda")
+    w13_up_dim = 2 * n if gated else n
+    activation = MoEActivation.SILU if gated else MoEActivation.RELU2_NO_MUL
     padded_n = marlin_moe_padded_intermediate(n, gs)
     assert padded_n != n
 
-    a = torch.randn((m, k), device=device, dtype=dtype) / 10
-    w13_weight = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
+    torch.manual_seed(0)
+    a = torch.randn((m, k), device=device, dtype=dtype)
+    w13_weight = torch.randn((e, w13_up_dim, k), device=device, dtype=dtype) / k**0.5
     w2_weight = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
     w13_weight = w13_weight.clamp(-448, 448).to(fp8)
     w2_weight = w2_weight.clamp(-448, 448).to(fp8)
     w13_scale = torch.full(
-        (e, 2 * n, k // gs), e8m0_one, dtype=torch.uint8, device=device
+        (e, w13_up_dim, k // gs), e8m0_one, dtype=torch.uint8, device=device
     )
     w2_scale = torch.full((e, k, n // gs), e8m0_one, dtype=torch.uint8, device=device)
 
@@ -848,6 +879,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
+        activation=activation,
         is_k_full=True,
         workspace=layer.workspace,
     )
@@ -859,5 +891,6 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
             topk_weight=topk_weights,
             topk_ids=topk_ids,
             global_num_experts=e,
+            activation=activation,
         )
     torch.testing.assert_close(out, ref, atol=8e-2, rtol=0)

tests/quantization/test_online.py

@@ -176,3 +176,73 @@ def test_online_quant_load_format_dummy(
     ) as llm:
         outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
         print(outputs[0][1])
+
+
+@pytest.mark.parametrize("is_act_and_mul", [True, False])
+@pytest.mark.parametrize("has_bias", [True, False])
+def test_online_moe_create_weights_w13_dim(is_act_and_mul, has_bias):
+    """w13 holds gate+up (2N) for gated MoE and up only (N) for non-gated MoE
+    (is_act_and_mul=False, e.g. NemotronH relu2)."""
+    from types import SimpleNamespace
+
+    from vllm.model_executor.layers.quantization.online.moe_base import (
+        OnlineMoEMethodBase,
+    )
+
+    class _Method(OnlineMoEMethodBase):
+        def __init__(self, moe):
+            self.moe = moe
+
+        def process_weights_after_loading(self, layer):
+            pass
+
+        def get_fused_moe_quant_config(self, layer):
+            return None
+
+    e, n, k = 4, 96, 64
+    moe = SimpleNamespace(is_act_and_mul=is_act_and_mul, has_bias=has_bias)
+    layer = torch.nn.Module()
+    _Method(moe).create_weights(
+        layer=layer,
+        num_experts=e,
+        hidden_size=k,
+        intermediate_size_per_partition=n,
+        params_dtype=torch.bfloat16,
+    )
+    w13_up_dim = 2 * n if is_act_and_mul else n
+    assert layer.w13_weight.shape == (e, w13_up_dim, k)
+    assert layer.w2_weight.shape == (e, k, n)
+    if has_bias:
+        assert layer.w13_bias.shape == (e, w13_up_dim)
+        assert layer.w2_bias.shape == (e, k)
+
+
+@pytest.mark.parametrize("is_act_and_mul", [True, False])
+def test_fp8_per_block_zero_padding(is_act_and_mul):
+    """_zero_padding must zero the roundup-padded rows of each w13 shard: two
+    gate/up shards for gated MoE, a single up shard for non-gated MoE."""
+    from types import SimpleNamespace
+
+    e, n, k = 2, 96, 256  # unpadded sizes; k is already block-aligned
+    n_pad = 128  # rounded up to the 128 quant block
+    num_shards = 2 if is_act_and_mul else 1
+    layer = SimpleNamespace(
+        moe_config=SimpleNamespace(
+            hidden_dim_unpadded=k, intermediate_size_per_partition_unpadded=n
+        ),
+        w13_weight=torch.ones(e, num_shards * n_pad, k),
+        w2_weight=torch.ones(e, k, n_pad),
+        w13_bias=torch.ones(e, num_shards * n_pad),
+        w2_bias=torch.ones(e, k),
+    )
+    method = SimpleNamespace(moe=SimpleNamespace(is_act_and_mul=is_act_and_mul))
+    Fp8PerBlockOnlineMoEMethod._zero_padding(method, layer)
+
+    w13 = layer.w13_weight.view(e, num_shards, n_pad, k)
+    assert (w13[:, :, :n] == 1).all()
+    assert (w13[:, :, n:] == 0).all()
+    bias = layer.w13_bias.view(e, num_shards, n_pad)
+    assert (bias[..., :n] == 1).all()
+    assert (bias[..., n:] == 0).all()
+    assert (layer.w2_weight[:, :, :n] == 1).all()
+    assert (layer.w2_weight[:, :, n:] == 0).all()

vllm/model_executor/layers/quantization/online/fp8.py

@@ -582,12 +582,14 @@ def _zero_padding(self, layer: Module) -> None:
         hidden_size = layer.moe_config.hidden_dim_unpadded
         intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded
 
-        w13_half_size = layer.w13_weight.shape[1] // 2
-        if w13_half_size > intermediate_size:
-            layer.w13_weight[:, intermediate_size:w13_half_size, :] = 0
-            layer.w13_weight[
-                :, w13_half_size + intermediate_size : 2 * w13_half_size, :
-            ] = 0
+        # w13 holds two gate/up shards for gated MoE, a single up shard for
+        # non-gated MoE (is_act_and_mul=False).
+        num_w13_shards = 2 if self.moe.is_act_and_mul else 1
+        w13_shard_size = layer.w13_weight.shape[1] // num_w13_shards
+        if w13_shard_size > intermediate_size:
+            for s in range(num_w13_shards):
+                start = s * w13_shard_size + intermediate_size
+                layer.w13_weight[:, start : (s + 1) * w13_shard_size, :] = 0
         if layer.w13_weight.shape[2] > hidden_size:
             layer.w13_weight[:, :, hidden_size:] = 0
 
@@ -597,12 +599,11 @@ def _zero_padding(self, layer: Module) -> None:
             layer.w2_weight[:, :, intermediate_size:] = 0
 
         if getattr(layer, "w13_bias", None) is not None:
-            w13_bias_half_size = layer.w13_bias.shape[1] // 2
-            if w13_bias_half_size > intermediate_size:
-                layer.w13_bias[:, intermediate_size:w13_bias_half_size] = 0
-                layer.w13_bias[
-                    :, w13_bias_half_size + intermediate_size : 2 * w13_bias_half_size
-                ] = 0
+            w13_bias_shard_size = layer.w13_bias.shape[1] // num_w13_shards
+            if w13_bias_shard_size > intermediate_size:
+                for s in range(num_w13_shards):
+                    start = s * w13_bias_shard_size + intermediate_size
+                    layer.w13_bias[:, start : (s + 1) * w13_bias_shard_size] = 0
 
         if (
             getattr(layer, "w2_bias", None) is not None

vllm/model_executor/layers/quantization/online/moe_base.py

@@ -37,11 +37,19 @@ def create_weights(
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None
 
+        # Non-gated MoE (is_act_and_mul=False) has no gate_proj, so w13 holds
+        # only the up projection.
+        w13_up_dim = (
+            2 * intermediate_size_per_partition
+            if self.moe.is_act_and_mul
+            else intermediate_size_per_partition
+        )
+
         # Fused gate_up_proj (column parallel) — full precision on meta device
         w13_weight = torch.nn.Parameter(
             torch.empty(
                 num_experts,
-                2 * intermediate_size_per_partition,
+                w13_up_dim,
                 hidden_size,
                 device="meta",
                 dtype=params_dtype,
@@ -70,7 +78,7 @@ def create_weights(
             w13_bias = torch.nn.Parameter(
                 torch.zeros(
                     num_experts,
-                    2 * intermediate_size_per_partition,
+                    w13_up_dim,
                     device="meta",
                     dtype=layer.orig_dtype,
                 ),

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

@@ -217,16 +217,20 @@ def prepare_fp8_layer_for_marlin(
         replace_parameter(layer, "bias", bias)
 
 
-def _moe_pad_shard_rows(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor:
-    """Zero-pad each gate/up shard of a ``(E, 2 * n, ...)`` tensor to padded_n
-    rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing."""
+def _moe_pad_shard_rows(
+    x: torch.Tensor, n: int, padded_n: int, num_shards: int
+) -> torch.Tensor:
+    """Zero-pad each w13 shard of a ``(E, num_shards * n, ...)`` tensor to
+    padded_n rows. FP8 zero decodes to 0.0, so the padded rows contribute
+    nothing. Gated MoE stores two gate/up shards; non-gated MoE
+    (``is_act_and_mul=False``) stores a single up shard."""
     if padded_n == n:
         return x
     e = x.size(0)
     rest = x.shape[2:]
-    x = x.view(e, 2, n, *rest)
+    x = x.view(e, num_shards, n, *rest)
     x = torch.nn.functional.pad(x, (0, 0) * len(rest) + (0, padded_n - n))
-    return x.reshape(e, 2 * padded_n, *rest)
+    return x.reshape(e, num_shards * padded_n, *rest)
 
 
 def _moe_pad_last(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor:
@@ -265,6 +269,10 @@ def prepare_fp8_moe_layer_for_marlin(
     k = layer.hidden_size
     n = layer.intermediate_size_per_partition
     w13_n = w13_weight.size(1)
+    # w13 holds two gate/up shards for gated MoE, a single up shard for
+    # non-gated MoE (is_act_and_mul=False).
+    num_w13_shards = w13_n // n
+    assert w13_n == num_w13_shards * n and num_w13_shards in (1, 2)
     weight_block_size = getattr(layer, "weight_block_size", None)
     group_size = -1 if weight_block_size is None else weight_block_size[1]
 
@@ -273,7 +281,7 @@ def prepare_fp8_moe_layer_for_marlin(
     # are padded to match below (the padded values are irrelevant).
     padded_n = marlin_moe_padded_intermediate(n, group_size)
     if padded_n != n:
-        w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n)
+        w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n, num_w13_shards)
         w2_weight = _moe_pad_last(w2_weight, n, padded_n)
 
     # WORKSPACE
@@ -345,10 +353,10 @@ def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
         if padded_n != n:
             if "w13" in name:
                 g = scales.size(1)
-                scales = scales.view(e, g, 2, n)
+                scales = scales.view(e, g, num_w13_shards, n)
                 scales = torch.nn.functional.pad(scales, (0, padded_n - n))
-                scales = scales.reshape(e, g, 2 * padded_n)
-                size_n = 2 * padded_n
+                scales = scales.reshape(e, g, num_w13_shards * padded_n)
+                size_n = num_w13_shards * padded_n
             else:
                 if group_size > 0:
                     pad_groups = (padded_n - n) // group_size
@@ -522,9 +530,10 @@ def prepare_mxfp8_moe_layer_for_marlin(
 
     Args:
         layer: MoE layer (used to read params_dtype and attach workspace).
-        w13: [E, 2*N, K] float8_e4m3fn weights.
+        w13: [E, S*N, K] float8_e4m3fn weights, where S is 2 for gated MoE
+            (gate/up shards) and 1 for non-gated MoE (is_act_and_mul=False).
         w2:  [E, K, N] float8_e4m3fn weights.
-        w13_scale: [E, 2*N, K//32] uint8 e8m0 scales.
+        w13_scale: [E, S*N, K//32] uint8 e8m0 scales.
         w2_scale:  [E, K, N//32] uint8 e8m0 scales.
 
     Returns:
@@ -534,12 +543,14 @@ def prepare_mxfp8_moe_layer_for_marlin(
     e = w13.shape[0]
     k = w13.shape[2]
     n = w2.shape[2]
+    num_w13_shards = w13.shape[1] // n
+    assert w13.shape[1] == num_w13_shards * n and num_w13_shards in (1, 2)
 
     # Pad a tile-misaligned intermediate size to a valid Marlin thread tile.
     padded_n = marlin_moe_padded_intermediate(n, group_size)
     if padded_n != n:
-        w13 = _moe_pad_shard_rows(w13, n, padded_n)
-        w13_scale = _moe_pad_shard_rows(w13_scale, n, padded_n)
+        w13 = _moe_pad_shard_rows(w13, n, padded_n, num_w13_shards)
+        w13_scale = _moe_pad_shard_rows(w13_scale, n, padded_n, num_w13_shards)
         w2 = _moe_pad_last(w2, n, padded_n)
         w2_scale = _moe_pad_last(w2_scale, n // group_size, padded_n // group_size)
         n = padded_n