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

[puririshi98]

Purpose

Non-gated MoE models — where w13 holds only an up projection instead of fused
gate+up shards — currently break in two places on the online-quantization +
Marlin path. The concrete motivation is the NemotronH architecture
(mlp_hidden_act="relu2", is_act_and_mul=False) used by the public
Nemotron-3 family:

• nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
• nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16
• nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16

MarlinExperts already supports non-gated MoE (_supports_no_act_and_mul() -> True, RELU2_NO_MUL, w13_num_shards = 2 if activation.is_gated else 1), but
the weight-prep side does not:

Bug 1 — online-quant allocation.
OnlineMoEMethodBase.create_weights unconditionally allocates w13_weight
(and w13_bias) at 2 * intermediate_size_per_partition. For non-gated
models the kernel-format conversion then packs a 2N-wide w13 while the
Marlin MoE gemm computes size_n = N, tripping
STD_TORCH_CHECK(size_n == actual_size_n) in marlin_moe_wna16. This is the
shared allocation base inherited by all online MoE methods (fp8
per-tensor/per-block/ptpc, mxfp8, int8), so one fix covers all of them; it
mirrors UnquantizedFusedMoEMethod.create_weights (w13_up_dim). The fix also
makes Fp8PerBlockOnlineMoEMethod._zero_padding shard-aware: it
previously assumed two gate/up shards, so for a non-gated model whose
intermediate size is rounded up to the 128 quant block, uninitialized pad rows
were never zeroed and would contaminate the shared per-block scales.

Bug 2 — Marlin MoE tile padding is gated-only.
The MoE thread-tile padding added in #45703 hardcodes the gated layout:
_moe_pad_shard_rows views (E, 2, n, ...) and the fp8 scale-pad branch
views (e, g, 2, n). A non-gated model with a tile-misaligned rank-local
intermediate size fails at the view instead of being padded. Concrete case:
Nemotron-3-Super-120B-A12B has moe_intermediate_size = 2688; at TP=4 the
rank-local shard is 672, which is not a multiple of the 64-column Marlin
thread tile (min_thread_n). The fix derives the shard count from the tensor
shapes (w13_n // n: 2 gated, 1 non-gated) and threads it through the weight
and scale padding in prepare_fp8_moe_layer_for_marlin and
prepare_mxfp8_moe_layer_for_marlin. Zero-padding is mathematically exact for
non-gated activations with act(0) = 0 (relu²): the padded w13 output
channels are zero, and the matching zero w2 input columns drop out.

No behavior change for existing gated models: num_w13_shards resolves
to 2 and reproduces the prior layout exactly; the allocation and zero-padding
changes only affect is_act_and_mul=False configurations, which previously
crashed (or, for per-block zero-padding, silently corrupted scales).

Why this does not duplicate existing PRs

• [Bugfix] Fix FP8 MoE double memory allocation for non-gated models #44498 applies the same w13_up_dim fix to Fp8MoEMethod in
  quantization/fp8.py (checkpoint-fp8 path); its second hunk targets an
  online create_weights that no longer exists after the online/moe_base.py
  refactor. This PR fixes the shared online allocation base that all online
  methods inherit from — complementary, not overlapping (no common lines).
• [Bugfix] Pad Marlin FP8 MoE weight dims to tile alignment under TP > 1 #36807 is an earlier tile-padding approach that predates the merged
  [Kernel] Consolidate Marlin thread-tile padding across all dense Marlin paths #45295 (dense) / [Kernel] Extend Marlin thread-tile padding to MoE (WNA16 + FP8/MXFP8) #45703 (MoE) mechanism; this PR extends the merged
  mechanism to non-gated layouts rather than introducing a parallel one.
• [MoE] Allow is_act_and_mul=False (non-gated MoE) on TPU #46795 enables non-gated MoE config on TPU only (fused_moe/config.py).

Changes

1. vllm/model_executor/layers/quantization/online/moe_base.py — size
   w13_weight/w13_bias by self.moe.is_act_and_mul.
2. vllm/model_executor/layers/quantization/online/fp8.py — shard-aware
   Fp8PerBlockOnlineMoEMethod._zero_padding.
3. vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py —
   _moe_pad_shard_rows(..., num_shards) (required parameter);
   shard-count-aware weight and scale padding in both MoE prep functions.
4. Tests:
   • tests/quantization/test_online.py — CPU unit tests
     test_online_moe_create_weights_w13_dim (gated × bias) and
     test_fp8_per_block_zero_padding (gated/non-gated).
   • tests/kernels/quantization/test_marlin_tile_padding.py — CPU shape test
     test_moe_pad_shard_rows[1|2]; gated=[True, False] parametrization of
     test_fp8_marlin_moe_padded_round_trip and
     test_mxfp8_marlin_moe_padded_round_trip (non-gated uses
     MoEActivation.RELU2_NO_MUL in both the kernel call and the
     torch_experts reference); round-trip inputs are no longer downscaled so
     the tolerance is discriminative, and the round-trips are seeded
     (torch.manual_seed(0)) so the unscaled inputs stay deterministic across
     runs.

Known related issue left out of scope: the int-WNA16 oracle pad helpers
(_pad_w13_shard_cols/_pad_w13_bias in fused_moe/oracle/int_wna16.py)
are still gated-only, so a non-gated WNA16 model with a tile-misaligned
intermediate fails loudly at the analogous view; same fix pattern applies
and can follow up.

Test Plan

# GPU: end-to-end online-quant MoE (weight creation -> loading -> post-load
# processing -> kernel run), gated and non-gated
pytest tests/quantization/test_online.py -v -k end_to_end

# GPU: full Marlin tile-padding regression (fp8-Marlin-capable GPU, sm >= 80)
pytest tests/kernels/quantization/test_marlin_tile_padding.py -v

Test Result

All runs on GB200 (aarch64, CUDA 13.0), editable install of this branch with
the precompiled wheel pinned to the main commit this branch merges
(1ff942965), so the compiled ops match the source.

$ pytest tests/quantization/test_online.py -q -k end_to_end
2 passed, 6 deselected in 25.63s

$ pytest tests/kernels/quantization/test_marlin_tile_padding.py -q
103 passed, 13 skipped in 30.42s

# canary: a Marlin MoE test this PR does not touch
$ pytest tests/kernels/quantization/test_marlin_tile_padding.py -q \
    -k test_gptq_marlin_moe_padded_round_trip
4 passed in 45.71s

Negative controls. The end-to-end test was confirmed to fail when either
fix is reverted, and only for the non-gated parametrization:

Reverted fix	Result
create_weights w13 allocation back to unconditional 2 * intermediate	[False] fails, [True] passes
_zero_padding back to assuming 2 shards	[False] fails, [True] passes

The test uses a non-block-aligned intermediate size (n=96, rounded up to the
128 quant block) and leaves the roundup pad rows non-zero, so the kernel output
only matches the bf16 reference if the pad rows are zeroed — this is what makes
the second control fail rather than silently pass.

The test pins VLLM_TEST_FORCE_FP8_MARLIN=1. Marlin is the fp8 MoE backend
that supports non-gated MoE; the Blackwell default
trtllm_fp8_block_scale_moe hardcodes the gated layout and rejects a
single-shard w13 (gemm1_weights_scale.size(1) == intermediate_size_factor * intermediate_size / 128, 1 vs 2). That is a separate gap, not addressed here.

The 13 skips are pre-existing shape/group-divisibility guards in
test_gptq_marlin_padded_round_trip / test_marlin_moe_padded_intermediate
("group straddles the boundary", "group must divide the rank-local K") and
skip identically on upstream main; none are in tests added or modified by
this PR.

Additional validation: a functionally equivalent backport of both fixes to
v0.20.0 was validated on GB200:

• End-to-end online-MXFP8 + Marlin generation for all three Nemotron-3 models
  (Nano-30B-A3B, Super-120B-A12B at TP=4 exercising the 672->704 pad,
  Ultra-550B-A55B at TP=8/PP=2).
• GEMM-level parity for the padded non-gated case: fused Marlin MoE output vs
  a bf16 reference reaches the same agreement as an already-tile-aligned
  control (cosine similarity 0.999180 padded vs 0.999174 control, i.e. the
  padding contributes no error beyond MXFP8 quantization noise). Since the
  change only enables previously-crashing configurations and gated layouts are
  byte-identical, this parity evidence stands in for a model-quality eval.

AI assistance disclosure

This PR was developed with AI assistance (Claude Code), including an
adversarial multi-agent review pass. Every changed line has been reviewed by
the human submitter, who ran the tests and can defend the change end-to-end.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[puririshi98]

@TomerBN-Nvidia do you know who can merge this?

[TomerBN-Nvidia]

@TomerBN-Nvidia do you know who can merge this?

@pavanimajety can you help with this?

[pavanimajety]

A few nits, thanks for the PR!

[pavanimajety]

nit: I think this test is unnecessary. it might be more helpful to add one test that tests the whole online_moe weight creation, post weight loading and kernel run.

[puririshi98]

Agreed, replaced both narrow tests with a single end-to-end test_online_moe_fp8_per_block_end_to_end that covers weight creation, weight loading, post-loading processing and a kernel run (parametrized gated / non-gated).

Two details worth calling out:

• It uses a non-block-aligned intermediate size (n=96, rounded up to 128) and leaves the roundup pad rows non-zero, so the kernel output only matches the bf16 reference if _zero_padding actually zeroes them — that turns the old shape/zero-padding assertions into a behavioural guard.
• It pins VLLM_TEST_FORCE_FP8_MARLIN=1. Marlin is the fp8 MoE backend that supports non-gated MoE; on Blackwell the default trtllm_fp8_block_scale_moe hardcodes the gated layout and rejects a single-shard w13 (gemm1_weights_scale.size(1) == intermediate_size_factor * intermediate_size / 128 (1 vs 2)), which is a separate gap outside this PR.

Verified on GB200: 2 passed. Negative controls confirm it catches both bugs — reverting the w13 allocation fix, or making _zero_padding assume 2 shards again, fails only the non-gated case.

[pavanimajety]

nit: Combine with the above loop

[puririshi98]

Done — the weight and bias pad-zeroing now share one loop:

num_w13_shards = 2 if self.moe.is_act_and_mul else 1
for w13 in (layer.w13_weight, getattr(layer, "w13_bias", None)):
    if w13 is None:
        continue
    shard_size = w13.shape[1] // num_w13_shards
    if shard_size > intermediate_size:
        for s in range(num_w13_shards):
            start = s * shard_size + intermediate_size
            w13[:, start : (s + 1) * shard_size] = 0

Semantics are unchanged: for the 3-D weight, w13[:, a:b] is equivalent to the previous w13[:, a:b, :].

[pavanimajety]

are there other places where _moe_pad_shard_rows is utilized? Do we need to give a default number here for num_shards?

[puririshi98]

_moe_pad_shard_rows is module-private and has exactly three call sites, all in marlin_utils_fp8.py, and all now pass num_w13_shards explicitly:

• prepare_fp8_moe_layer_for_marlin — w13_weight
• prepare_mxfp8_moe_layer_for_marlin — w13 and w13_scale

I would prefer not to give num_shards a default. A default of 2 would silently reinstate the gated-only assumption this PR is fixing: a non-gated layer would then be mis-padded rather than failing loudly at the view. Keeping it required means any future caller has to state the layout. Happy to add one if you feel strongly.

[aoshen02]

Hi, thanks for the pr. I've solve the problem partially in #51125 and list you as an author. Hope it works for you, lmk if any other problems.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @puririshi98.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork