[ROCm][Bugfix] Stop declaring the MXFP4 MoE intermediate padding to AITER

[PerryZhang01]

Purpose

The AMD CI job dsv41-mono-decode-numerics-amd goes red intermittently with

AssertionError: (0, 1, (973, 973), 'residual', nan)

(e.g. https://buildkite.com/vllm/ci/builds/94971). The NaN comes from AITER, on
vLLM's ordinary MoE path, not from the mono kernels the test compares against.

AITER's a8w4 MoE (FP8 activations x MXFP4 weights) allocates several scratch
buffers with torch.empty, sized for the padded sorted length, and reads them
where the producing kernel never wrote. In FP8/E8M0 a garbage 0xFF byte is
NaN, so a whole token row of the output comes back NaN. It only happens with a
non-zero intermediate_pad, at decode-sized batches, after larger calls have
run in the same process — which is why it shows up as a flaky decode failure.
Root cause, a standalone repro and a fix are filed upstream:

• [Issue]: a8w4 fused_moe writes NaN at small token counts on gfx950 (uninitialized padded scratch) ROCm/aiter#6565
• [FlyDSL] [Bugfix] Initialize padded MoE scratch to stop NaN at small token counts (gfx950 a8w4) ROCm/aiter#6568

This change

Stop declaring intermediate_pad to AITER on the MXFP4 path. It is redundant
there, and dropping it keeps the call off the buggy AITER path.

Why it is redundant — the padded region is zero at the point AITER sees it, and
stays zero through every step between allocation and the kernel call:

step	padded region	where
allocation	torch.zeros for w13/w2, their scales and w13_bias/w2_bias, at the padded intermediate size	Mxfp4MoEMethod.create_weights
checkpoint load	the destination is narrowed to the loaded shard before the copy, so the padded rows are never written — this is deliberate, the helper is named for it and its docstring calls out "padding ... on the shard (intermediate) dimension (e.g. for MXFP4 GEMM)"	RoutedExperts._narrow_expert_data_for_padding, used by _load_w13 / _load_w2 right before expert_data.copy_()
weight shuffle	pure view + permute (asserts N % 16 == 0, K % 32 == 0, no re-padding), so zeros stay zeros	aiter.ops.shuffle.shuffle_weight
scale shuffle	re-pads the scale to 256 rows / 8 cols, but fills the pad with 0x7F (E8M0 = 1.0) before copying, so it is never garbage	aiter.utility.fp4_utils.e8m0_shuffle

So for the padded columns the kernel computes

• stage 1: x @ 0 = 0 (gate and up are both zero there, and the bias is zero,
  so the activation output is exactly silu(0) * 0 = 0)
• stage 2: 0 @ 0 = 0

Exact zeros also mean the padded columns cannot disturb the stage-2 activation
quantization: amax over an MX group is unaffected by zeros, whether or not a
group straddles real and padded columns.

Two things are deliberately left alone:

• hidden_pad — the activation columns it covers are not guaranteed to be
  zero, so dropping it would not be equivalent.
• SITU (activation == MoEActivation.SITU, the a16w4 FlyDSL port used by
  Kimi K3) — that kernel consumes the pad. The surrounding code already skips
  the rounding for it ("SITU's A16W4 FlyDSL kernel pads per gate/up half; pass
  through unrounded"), and on that path AITER keys its own scratch
  initialization on the value: _alloc = torch.zeros if inter_dim_pad > 0 else torch.empty in aiter/ops/flydsl/moe_kernels.py. Zeroing the declaration
  there would remove an initialization that is currently happening, which is
  the opposite of what this PR is for.

This is a workaround and should be reverted once the AITER fix is merged and
pinned. The surrounding code already carries a similar TODO about AITER padding
behaviour.

Test Result

MI355X (gfx950), ROCm 7.2.3, stock AITER 0.1.25 — i.e. without the
AITER-side fix — running tests/models/test_dsv41_mono_layers.py unmodified,
each run covering TP2 and TP4:

build	runs	cells	result
exactly this patch	4	8	all passed
same logic, earlier iterations of the comment/wording	10	20	all passed

14 runs, 28 cells, no NaN. Before the change the same test failed intermittently
— locally about 2 runs in 6 — so 14 clean runs is about 0.3% likely if the
change did nothing.

Overhead: none measurable. Timing the same shape with and without the
declaration (median over 6 rounds, per fused_moe_2stages call):

M	declared	not declared
2	0.139 ms	0.141 ms
4	0.156 ms	0.158 ms
1439	0.534 ms	0.536 ms
16384	2.010 ms	2.010 ms

All within run-to-run noise (~4%).

Scope and limits

• Validated on DeepSeek-V4.1-Flash at TP4 (model_dim=5120, per-rank
  intermediate 576 padded to 640, E=384, topk=6). Other MXFP4 models on this
  path (e.g. gpt-oss with INTERLEAVE) are covered by the same zero-fill
  argument but were not run here, so reviewers may prefer a narrower gate.
• The zero-fill chain above is established from the code, with the citations in
  the table. It is not confirmed by a run against a real MXFP4 checkpoint that
  actually has padding: the test that exercises this path uses a random-weight
  loader which also zero-fills the padding, and every MXFP4 checkpoint we have
  locally has moe_intermediate_size = 2048, which is already aligned per rank
  (pad = 0). DeepSeek-V4.1-Flash (2304) is the one that pads, and the local copy
  of it is FP8 rather than MXFP4.
• Models whose per-rank intermediate size is already 128-aligned get
  intermediate_pad == 0 and are unaffected either way.
• The VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 opt-in is a sub-case of
  use_mxfp4_w4a16, so it is covered by this change as well. The same
  zero-fill argument applies to it, but it runs different kernels and was not
  exercised here.

[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.

[wangyicong52]

Thanks for the fix! Meet the same problem when testing CI: buildkite #95188.