perf(moe): tune the Nemotron MTP decode row on S5000 with a checked-in recipe

[yeahdongcn]

Tune the Nemotron MTP decode row on S5000 with a checked-in recipe

Status: draft, pending review. Six commits on top of main: the packaged Nemotron shape map and the
decode/graph work, then the decode table row plus the autotune tooling that derives it (440265c, on top
of 37fba15). Pre-squash history for the older lineup is preserved on
backup/pr115-before-squash-20260919.

What it does

1. The decoded table row. E=128,N=1856,device_name=MTT_S5000.json gave the decode bucket
(M = 1/7/14/21/28 all resolve to nearest key 16) a row carried over from the previous
five-key map: BLOCK_M=32, BLOCK_N=64, BLOCK_K=64, GROUP_SIZE_M=16, num_warps=8, num_stages=2.
It now carries the tiny-M tile this hardware actually wants:

BLOCK_M=16, BLOCK_N=64, BLOCK_K=128, GROUP_SIZE_M=1, SPLIT_K=1, num_warps=4, num_stages=1

Key 512 keeps the config it was measured with - the two buckets used to share a row and are now
independent. Packaged keys stay 16, 49, 80, 113, 144, 511, 512, 513; the serving routes are
M≈115/120 -> 113 and M=512 -> 512, and larger M falls back to the generic heuristic.
Packaged table SHA256: 3e39677d2dbf5c984325f94255dbba3fa5eb2ac331fde3a86a79f75105fc531b.

2. The recipe that produces the row. src/torchada/triton/autotune/fused_moe/ci/shapes.json
records the kernel shape, the bucket and the row itself (pinned_rows, together with the
measurement behind it), so the table is generated rather than typed:

tune_moe.py --config .../ci/shapes.json --materialize --merge-configs

writes it without a GPU, and tests/test_tune_moe_recipe.py asserts the result matches the shipped
file row for row and key for key. --tune --merge-configs reruns the search and keeps the rows a
run did not measure, so a disagreement shows up as a reviewable diff.

3. Why this row is pinned rather than swept. The uniform-routing microbenchmark is not
authoritative for this bucket and not self-consistent: five sweeps produced three distinct winners
and only two of them landed on the row this table ships.

4. Runtime and kernel changes the decode path depends on (squashed earlier work): SPLIT_K
stays in the fused-MoE constexpr map - popping that key changed the generated launch variant and
caused a TPOT regression; the tensor-descriptor capability check now requires the Triton language
API triton.language.make_tensor_descriptor, so Triton 3.2 takes the fallback path while newer
Triton can use descriptors; the resolver regression test pins which M lands on which key.

Measurements

End-to-end on one MTT_S5000 card, same container and image per comparison, canonical Nemotron-3.5
MTP6 protocol (4096-in / 1000-out, one warmup plus five measured repeats, byte-identical output
digests in every arm):

arm	median TPOT
table as shipped (old key 16)	6.6207 ms/token
this row	6.2626 ms/token
this row + vLLM MUSA #241 and #242	5.5649 ms/token
control: #241+#242 with the old row	5.9338 ms/token

TTFT is unchanged: prefill still runs the large-M rows. MTP-depth ladder for the same stack,
median of five measured requests:

path	median TPOT
target-only	17.937391 ms/token
MTP1	11.436920 ms/token
MTP2	9.271596 ms/token
MTP4	7.392256 ms/token
MTP6	6.847434 ms/token

A rerun of the same final wheel and source reproduced MTP6 within 0.117% of the historical manual
result (6.616876 over ten requests vs 6.609144 ms/token) with matching output SHA256. Absolute
numbers are host- and image-specific.

Validation

• tests/test_tune_moe_recipe.py + tests/test_nemotron_moe_config.py: 22 passed on the branch
  artifact.
• Focused config/resolver and descriptor tests: 10 passed.
• Full torchada suite: 509 passed, 19 skipped; the only reported failure is the source-only
  project-version check without project metadata.
• Same-source vLLM MUSA E2E: pinned vLLM 2cf0a6915ce544dc493a0990f2ea38d81601128a, post1/Triton
  3.2 image, TP1 BF16, Triton Mamba, piecewise + full graphs, 4096 input / 1000 output.

Scope and known limits

• The table is keyed by kernel shape (E, N, dtype, block shape, device) plus the Triton version
  directory, not by model. This row affects the M <= 28 bucket of any model that resolves to that
  shape on this device; the other seven rows are untouched.
• hidden_size and topk are not part of the table key (upstream naming convention); the recipe
  does record them.
• The search grid emits num_stages=1 only, so the seven rows carrying s2/s3/s4 cannot be
  regenerated by the current grid - a full-table re-tune would be a deliberate re-baseline.
• Measurements come from one card and one image; no cross-SKU transfer is claimed.

[yeahdongcn]

Tiny-M entry of E=128,N=1856,device_name=MTT_S5000.json is stale

The table's nearest-key entry for decode (M=1/14/21 -> key 16) carried
BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_K=64, GROUP_SIZE_M=16, num_warps=8, num_stages=2.
Because a tuned table wins the lookup, that entry made the MUSA tiny-M default unreachable - so
vllm-musa#242 (one pipeline stage for tiny M) had no effect at all in any image that ships this
table. Measured, not inferred:

host / driver	arm	config	4k/128 TPOT	acceptance	output SHA
5.2.0 host	baseline + table	old "16" entry	8.996 ms	n/a	c214d7149093748b
5.2.0 host	#241+#242 + table	old "16" entry	7.99 ms	4.04-4.11	same
5.2.0 host	#241+#242, no table	tiny-M default	7.622 ms	3.90	same
3.3.8 host	#241+#242 + table	new "16" entry	8.473 ms	3.84	same
3.3.8 host	#241+#242, no table	tiny-M default	8.511 ms	4.11	same

Nemotron-3.5 MTP6, 4096-in/128-out, graph path, fresh container per arm, 3 samples per arm, one
output family throughout, so the arms are subtractable. The old entry cost -4.6% versus the
default; with this commit the table matches the default (8.473 vs 8.511 ms, and the no-table arm ran
at higher acceptance, which inflates its TPOT). Larger-M entries are untouched.

The new "16" entry is the tiny-M default configuration verbatim
(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64, BLOCK_SIZE_K=128, GROUP_SIZE_M=1, SPLIT_K=1, num_warps=4, num_stages=1, from get_default_config at M<=32 for bf16), so the table keeps whatever value it
has at larger M without paying for it at decode.

Note for whoever retunes further: get_configs_compute_bound() searches num_stages=1 only, so a
proper re-tune of the tiny-M rows is a separate, welcome follow-up - it may beat the default via
tiles/warps/GROUP_SIZE_M, which this change deliberately does not claim.

Edited: internal identifiers (tracker references, host names, card indices) removed; the measurements are unchanged.

[yeahdongcn]

Follow-up on this table: measured cost, and the right way to fix it

I first pushed a hand-edit of "16" here and then reverted it: that file is a
generated artifact (tune_moe.py -> save_configs), so editing it by hand is the wrong layer -
the next regeneration would silently undo it. Please squash the revert away when merging; the only
thing it must not leave behind is a hand-edited generated file.

What the measurements say (Nemotron-3.5 MTP6, 4096-in/128-out, graph path, one output family
c214d7149093748b in every arm and sample, 3 samples per arm, fresh container per arm):

host / card / driver	arms	4k/128 TPOT
5.2.0 host, card A	#241+#242 + this table (stale "16")	7.99 ms
5.2.0 host, card A	#241+#242 + no table (tiny-M default)	7.622 ms
5.2.0 host, card B	#241+#242 + table whose "16" entry is the tiny-M default	7.542 ms

So the stale entry costs about 4.6% once vllm-musa#242 (one pipeline stage for tiny M) is in,
and it also makes #242 inert: with this table present get_default_config never runs, so the
log line MoE default config for tiny M: never appears. Fixing the entry recovers the loss and can
even do slightly better than the naive default.

The right fix is at the generation layer, one of:

1. Regenerate this table with the decode batch sizes included, e.g.
   --batch-size 1,2,4,8,14,16,21,24,32,42,49. The tuner's default list already contains them, so
   the committed file simply was not produced that way (its keys are 16/49/80/113/144/511/512/513).
   Its candidate space is num_stages=1-only, so a regenerated tiny-M entry cannot reintroduce the
   staging regression. Blocker found while trying: in the shipped image tune_moe.py never gets
   past its workers - torch.ops._moe_C.moe_align_block_size is unregistered unless
   vllm._moe_C is imported explicitly, and the spawned workers lose that registration (patching
   mp.get_context is not enough because mp.Pool binds its own reference). Both were root-caused
   here and are documented in the tracker write-up for those two defects; the tuning harness needs that fixed before regenerating.
2. Or gate the lookup for tiny M: have get_moe_configs fall back to get_default_config when
   M <= 32, so a stale tuned entry can never outrank the platform default. That is a code change
   and therefore regeneration-proof.

Until one of them lands, this table is not a regression - it just keeps #242 inactive - so it should
not block either PR.

Edited: internal identifiers (tracker references, host names, card indices) removed; the measurements are unchanged.

[yeahdongcn]

Correction: the stale "16" entry is probably not the mechanism

My two comments above attributed a 4.6% difference to the stale "16" entry and recommended
regenerating the tiny-M rows. That attribution is withdrawn for three reasons grounded in
the protocol spec:

1. Bucket routing. The steady-state decode batch in the reference workload is M ~ 115/120, which
   resolves to key 113, not 16. The "16" entry is only touched by the first/tiny buckets, so it
   cannot explain a steady-state difference.
2. Protocol mismatch. My arms used 4096-in/128-out; the reference numbers (6.847 / 6.6091
   ms/token) use 1000 output tokens. Short runs mix in transient behaviour, so the two are not
   directly comparable.
3. Runtime mismatch. Most of my arms used the image's stock torchada 0.1.83, not the
   candidate that fixes the SPLIT_K launch-variant issue (~1 ms/token class). That confound is
   larger than the effect I was attributing.

What still stands from my measurements: on one host, one image, one output family, the eight-key map
and the no-table default differ by ~4.6% at 128-out on stock torchada (7.99 vs 7.622 ms) - an
observation, not an explanation, and it needs re-taking with the candidate runtime at 1000 output
tokens before anyone acts on it. The SPLIT_K validation fix in this PR is the mechanism that
the protocol spec already identified, and this table's entries are only meaningful together with it.

Edited: internal identifiers (tracker references, host names, card indices) removed; the measurements are unchanged.

[yeahdongcn]

tiny-M row corrected in place (head 440265c), and re-measured on the branch artifact.

Why this row

Key 16 is the bucket every MTP6 verify shape lands on (real decode calls were logged at M = 1/7/14/21/28, all resolving to key 16) and it dominates the step — but it is the one row this campaign never derived: 37fba15 repackaged the tuner's five-key map {1,4,6,8,16} into the current eight keys and carried key 16 over verbatim (only SPLIT_K was added). The shipped stage grid already contains the config we want: get_configs_compute_bound() iterates num_stages = 1 only, and its grid includes (BLOCK_M, BLOCK_N) = (16, 64) with BLOCK_K 32/64/128, num_warps 4/8/16 and GROUP_SIZE_M 1/16/32/64. This is also the only table in the tree without an s1 row (99 of 100 sweeps carry one).

Measurements

one MTT_S5000 card, image v0.28.0-ph1-5.2.0-torch2.11.0.post1-20260917, canonical Nemotron-3.5 MTP6 protocol (4096-in / 1000-out, 1 warmup + 5 measured, one fresh container per arm, byte-identical output SHA ad63e94a59912806… in every arm):

configuration	TPOT (ms)	TTFT (ms)	acceptance
table as shipped	6.6207	1516.5	6.04
this row	6.2626	1520.5	6.05
this row + vllm-musa #241 + #242	5.5649	1513.1	6.07
control: same stack, previous row	5.9338	1514.2	5.98
control: same stack, key 16 removed (tiny M falls to key 49, stages 4)	8.5499	1542.7	6.02

The last row is the mechanism check: with the table otherwise intact and only the tiny-M key missing, TPOT lands exactly on the untuned number, so the whole effect is this bucket's staging and nothing else in the map.

Branch artifact re-verified

The "this row" numbers above swapped the row into the installed table; this run uses the file as committed:

• B1R = overlay table (pinned a0b192a4…) + #241 + #242 → 5.9282 ms, TTFT 1513.5
• B3 = this branch's table (3e39677d…) + #241 + #242 → 5.5659 ms, TTFT 1513.2
• A/B = −6.11%, TTFT unchanged (+0.02%), same output SHA

Notes

• TTFT does not move because the large-M rows still serve prefill; only the decode bucket changes. Key 512 keeps the config it was measured with — it used to be the same row as key 16, so the two buckets are now independent, and the routing test pins which M lands on which key.
• tests/test_nemotron_moe_config.py: 9 passed in the container.
• Context for reviewers: with this table present, vllm-musa #242's get_default_config branch is never reached, so the decode win in the current stack comes from this row, not from that patch. The row we are landing is the config #242 selects for M ≤ 32.

Full nine-arm record and the shadowing mechanism: the internal campaign write-up.

Edited: internal identifiers (tracker references, host names, card indices) removed; the measurements are unchanged.

[yeahdongcn]

History note (force-push, no content change). The last three commits - 005e793 (hand-edit) -> 8eb5303 (revert) -> 9a6562b (the fix) - are now squashed into c612d30, which is the only change in that range that survives. The new commit's tree is byte-identical to 9a6562b's (1b9f00dc…), so every number in the evidence comment above still applies unchanged, and the PR diff is the same 15 files.

[yeahdongcn]

The table row is now derived from a checked-in recipe, not typed by hand

Follow-up to the row itself: the key-16 row was placed from an end-to-end measurement, and nothing in the repo could produce it. That is fixed in the same commit (440265c), and the upshot is that the table can be regenerated from source.

What was added

Path	What it does
src/torchada/triton/autotune/fused_moe/ci/shapes.json	Recipe for the Nemotron shape: E=128, hidden=2688, N=1856, topk=6, non-gated relu2_no_mul, bf16, TP1/EP1, bucket set [16], and the row itself under pinned_rows together with the measurement behind it.
tune_moe.py --materialize	Writes the recipe's pinned_rows into their table — no GPU, no sweep — and --merge-configs keeps every row the recipe does not pin. A pinned row without a source is rejected.
tune_moe.py --tune	Recipe entries may now pin batch_sizes or describe a kernel shape without a model checkout, so a table bucket can be swept from a recipe alone.
A skipped-candidate fix	A candidate that raises inside a tuning worker is logged and skipped. Before this, any unhandled exception killed the worker and the parent blocked forever on a result that never arrived (I hit exactly that with an unsupported config key: no output, no error, just a hung process). If every worker still dies, the run now aborts with a message instead of hanging.

Reproduce:

python src/torchada/triton/autotune/fused_moe/tune_moe.py \
    --config src/torchada/triton/autotune/fused_moe/ci/shapes.json \
    --materialize --merge-configs
python -m pytest tests/test_tune_moe_recipe.py

The test seeds a copy of the shipped table, materializes the recipe into it, and asserts the result has the shipped file's content row for row and key for key.

Why the row is pinned instead of swept

The uniform-routing microbenchmark is not authoritative for this bucket, and it is not even self-consistent. Five sweeps of M = 16 on one MTT_S5000 card (image v0.28.0-ph1-5.2.0-torch2.11.0.post1-20260917, driver 5.2.0-server) produced three distinct winners:

sweep	winner for bucket 16
full grid, clean dir	BM64/BN128/BK64/GSM1/w8/s1
full grid, repeat (--merge-configs)	a different row from the run above
tiny-M family	BM16/BN64/BK64/GSM1/w4/s1
tiny-M family (repeat)	BM16/BN64/BK128/GSM1/w4/s1
tiny-M family + one unsupported candidate	BM16/BN64/BK128/GSM1/w4/s1

Only the last two land on the row this table ships, which is also the one that wins end to end (5.5659 ms vs 6.6207 ms TPOT, the figure in the commit message). The grid's own answer for this bucket therefore cannot be the source of truth; the recipe records the measured row and the measurement.

--merge-configs is what keeps the rest of the table intact while re-deriving one bucket: the other seven rows carry num_stages 2–4 from an earlier sweep, while this grid iterates num_stages = 1 only, so a full-table re-run would replace measured rows with unmeasured ones. Verified: merging the recipe into the shipped table reproduced the file's content exactly and kept all eight rows.

Notes

• The shipped table is hand-formatted (one row per line) while save_configs writes the indented style the other 99 tables use, so an in-place refresh reformats the file without changing its content. The reproduction test compares content, not bytes.
• One deliberate non-change: I first tried giving every grid candidate SPLIT_K, so a swept row would be schema-identical to this table's rows. That is wrong — torchada's fused_moe_kernel does not declare SPLIT_K, so Triton rejects the whole config (KeyError: 'Keyword argument SPLIT_K was specified but unrecognised'), and a sweeper hangs rather than reporting it. The key stays out of the sweep; the pinned row keeps it because that is how the shipped table is written and the vLLM path sets it itself.

Evidence: m9_pr115_retune.sh (derivation + merge + comparison), m9_pr115_retune_check.sh (materialize, sweep, resilience probe, tests), log pr115-retune.log / pr115-check.log; 22 passed for tests/test_tune_moe_recipe.py tests/test_nemotron_moe_config.py on the final tree.

Edited: internal identifiers (tracker references, host names, card indices) removed; the measurements are unchanged.