perf(moe): one pipeline stage for the tiny-M MoE default on MUSA (S5000 4k/128: -18% TPOT on Nemotron 3.5 MTP6)

[yeahdongcn]

What

get_default_config pipelines the MoE K loop 4 stages deep for M <= 32
(num_stages = 4 at BLOCK_K = 128). On S5000 that is a CUDA-shaped default: at the
tiny-M decode shape there is one tile per touched expert with 1-3 useful rows inside it,
and 4 stages reserve ~4x the shared memory per block, which costs resident blocks. This
sets num_stages = 1 for M <= 32 on MUSA only, and logs the resolved value once so
the branch is observable from a serving log.

Found while pricing the decode-MoE padding hypothesis in MUSA-100038: the padding is not
the problem (per-expert tiles are already at their floor), the default config is.

Measured

S5000, daily image vllm:v0.28.0-ph1-5.2.0-torch2.11.0.post1-latest, 4096-in / 128-out,
one request, production graph path, fresh container per arm, 1 warmup + 3 measured
requests, VLLM_LOGGING_LEVEL=INFO, no probes:

config	TPOT median	delta	TTFT	output SHA
shipping default (4 stages)	11.635 ms	-	1648 ms	c214d714...
num_stages=1 only	9.563 ms	-17.80%	1678 ms	c214d714...
BLOCK_SIZE_N=128 only	12.619 ms	+8.46%	2632 ms	c214d714...
this patch (in-tree, no env override)	9.528 ms	-18.11%	1622 ms	c214d714...

• Per-arm spread across the 3 measured requests is 0.07-0.45%, so the delta is ~40x noise.
• Every measured request in every arm produced the same output SHA: byte-equivalent on
  this workload. Acceptance length is the model's and unchanged across arms (4.13/4.13/4.11).
• Activation is asserted rather than inferred: this patch's logger.info_once printed
  MoE default config for tiny M: M=1 num_stages=1 (also M=7/14/21) in the serving log,
  while the unpatched control printed vLLM's own
  WARNING ... Using default MoE config. Performance might be sub-optimal!.
• Isolation: num_stages alone is the whole effect, and BLOCK_SIZE_N=128 alone regresses,
  so the win is not an artifact of the tuned-config path.

Scope and risk

• MUSA only: CUDA keeps 4 stages, ROCm keeps its own num_stages_rocm.
• M > 32 untouched (the test asserts the 512 case still resolves to 3 stages).
• It changes the default for all tiny-M MoE shapes on MUSA. Only Nemotron-3.5-A3B MTP6
  decode was measured here; other MoE models with tiny-M batches should benefit for the
  same reason but were not run.

Tests

tests/test_moe_default_config.py - 4 passed in the patched MUSA container
(python -m pytest tests/test_moe_default_config.py -q). The test is platform-aware, so it
also holds on stacks where this patch is not applied.

Not run

• 4k/1000 arm (only 4k/128 was measured end to end)
• other MoE models at tiny M
• hardware-counter confirmation of the occupancy mechanism: the roofline reading is inferred
  from the timing delta, not from counters.

Ticket: MUSA-100038.

[yeahdongcn]

Family-safety note (in the light of MUSA-100043)

This A/B is immune to the MTP6 output-family instability tracked in MUSA-100043. All 15
measured requests across the five arms (control, m100035, w8bn128, s1only, bn128only)
produced exactly one output SHA, c214d7149093748b, and the in-tree patched arm produced
the same SHA again — every arm served a byte-identical token stream. MTP mean acceptance
length was equal across arms (4.13 / 4.13 / 4.11 / 3.98 / 3.95), so the ~18% is a per-step
cost change, not a change in how many tokens the model accepts.

For contrast: an earlier run on the same image (r7-final-ab.log, MUSA-100037) had its two
arms split across families (control family A, gate arm family B) and reads as +46% if
taken naively. That run is not the basis of PR #241, whose numbers are family-matched.

[yeahdongcn]

Patch-file fix: the series file was malformed, now corrected and re-verified

0166-...patch had been assembled by hand from diff -u output with the --- a/... /
+++ b/... file headers dropped, so git apply rejected it ("patch fragment without header")
on every tree. Since build_apply.py runs strict=True, a merged #242 would have failed the
build instead of silently doing nothing. Fixed by regenerating the diff against this branch's
actual pinned upstream (VLLM_COMMIT=2cf0a6915ce544dc493a0990f2ea38d81601128a, VLLM_TAG=v0.28.0);
the branch head is now 4d456e977.

Independent re-verification (local, reproducible via generated/MUSA-100045/tools/series_dryrun.sh):

• p166.patch APPLIES CLEAN on post-series tree, and clean on the pristine pin file too.
• The patch still sends num_stages = 4 to num_stages = 1 on MUSA only, with the same
  logger.info_once activation line; the test file is unchanged.
• Caveat observed while re-checking: that script defaults to whatever upstream pin the checked-out
  tree has. Run against the main checkout's pin (v0.24.0) it reports ~94 spurious series
  failures and a spurious failure for perf(moe): parallel FP32 router-gate GEMM for Nemotron 3.5 decode (-8% TPOT, bitwise) #241's patch; at the dev branch's pin the series applies
  with 0 failures. The verdict above is from the correct pin.

[yeahdongcn]

Verified stack measurement (parent session, 2026-09-19)

Host 10.20.34.83/GPU1, daily image v0.28.0-ph1-5.2.0-torch2.11.0.post1-latest, fresh container and
fresh cache per arm, graph path (mode:3 + FULL_AND_PIECEWISE), Nemotron-3.5-A3B MTP6 (6 spec
tokens), 4k prompt, 1 warmup + 3 measurements, exclusivity asserted per arm. All four arms produced
the same output SHA on both benches (4k/128 c214d7149093, 4k/1k ad63e94a5991), so the arms
are subtractable.

arm	4k/128 TPOT	vs base	4k/1k TPOT	vs base
baseline	11.635 ms	-	8.169 ms	-
+#241	10.761 ms	-7.51%	7.531 ms	-7.81%
+#242	9.553 ms	-17.89%	6.639 ms	-18.73%
+#241+#242	8.594 ms	-26.13%	5.967 ms	-26.95%

Super-additive: -26.13% measured vs -24.05% predicted by multiplying the two independent deltas
(arm spread 0.19-0.46%). The two PRs should be evaluated and landed together - judging either
one alone understates the pair (on a second host the same ordering reproduced).

Interaction with the torchada #115 tuned table

With a tuned config for E=128,N=1856,device_name=MTT_S5000 present, get_moe_configs wins and
get_default_config never runs (logs show Using configuration from ... and zero
MoE default config for tiny M: lines), so the table's num_stages=2 (nearest key 16 for M<=7)
overrides the new default 1. Matched comparison (same host, same SHA family, same acceptance):
7.636 -> 7.971 ms, i.e. the table costs 3-4% inside the combined stack. Follow-up: retune the
tiny-M rows against the new default, or gate them. Not a blocker - with the table present the
stack still gains -23.8%.

Other models

Applicability is visible in the serving logs: #242 only affects MoE shapes with no tuned table in
the image (verified active on Nemotron-3.5, and on DeepSeek-V2-Lite-Chat:
MoE default config for tiny M: M=1/2/4 num_stages=1), while Qwen3-30B-A3B (has a table) measured
as a no-op: 13.702 -> 13.704 ms at 4k/128, 13.817 -> 13.798 ms at 4k/1k, output SHA
identical in both arms. #241's guard geometry (128 experts x 2688 hidden) matched only
Nemotron-3.5 in a fleet-wide scan, and its kernel parity is bitwise (26/26 cases, M=1..512, worst
max_abs_diff 0.0, foreign gate geometries declined at runtime). #240 (torchada 0.1.88) could not
be measured: the in-container index only offers 0.1.87.

[yeahdongcn]

Correction to the "with the #115 table" figure posted above

The -23.8% I quoted for the table-present case compared different configurations (levers+table vs
no-levers+no-table), so it mixed the table's presence into the delta. A clean recheck on one host
with one output family (10.20.34.82/GPU1, Nemotron-3.5-A3B MTP6, 4k prompt, graph path, 3
measurements per arm, identical SHA c214d7149093748b in every arm and every sample) gives:

arm	#241+#242	#115 table	4k/128 TPOT	acceptance	activation lines
A	no	yes	8.996 ms	n/a	Using configuration from /tmp/tuned/E=128,N=1856,...
B	yes	yes	7.99 ms	4.04-4.11	table line + router-gate GEMM is serving
C	yes	no	7.622 ms	3.90	MoE default config for tiny M: M=1/14/21 num_stages=1 + gate line

• A -> B = -11.2%: what merging these two PRs buys in an image that ships the tuned table.
• A -> C = -15.3%: the same merge once the table's tiny-M rows are gated or retuned.
• B -> C = -4.6%: the table's own cost (C runs at lower acceptance, which inflates its TPOT,
  so the table costs at least this much).

Mechanism, straight from the logs: with the table present get_default_config is never entered -
arm B has no MoE default config for tiny M: line at all - so #242 contributes nothing while that
table is in place and only #241 acts. Fixing the table is therefore what unlocks #242's ~-18%,
not an optional extra. Both PRs remain worth merging as they stand (-11% today), and #241's -7.5%
measured without any table is consistent with this picture.

Reproducibility: arm B/C agree with the earlier independent sweep on the same host (B3 7.636 ms @4.13
vs C 7.622 ms; B4 7.971 ms @4.13 vs B 7.99 ms).