[MoE] MegaMoE A4W4: remove dispatch-pointer scalarization waterfalls and retune bs=512 dispatch CUs

[jhinpan]

Motivation

While benchmarking this branch's A4W4 path on 8×MI355X I found that two shapes miss
the golden latencies committed in _MEGA_PERF_BASELINE — and that they are exactly
the two shapes _MEGA_BENCH_PARAMS gates for a4w4:

shape	committed golden	branch head (972e1a94)
v4_pro:a4w4:512	0.5350 ms	0.5602 ms	+4.7%
v4_pro:a4w4:1024	0.7050 ms	0.7451 ms	+5.7%

This is currently invisible to CI: prepare-mlir fails in the Build LLVM once step,
so test and the multi-GPU jobs are all SKIPPED and the a4w4 perf gate never runs.

This PR recovers part of that and documents the rest. It does not close the whole
gap — see "What this does not fix" below.

Technical Details

1. Dispatch-table pointers no longer force scalarization waterfalls.

dp() reads a 64-bit pointer out of the dispatch slot table at a compile-time constant
index, and rdisp is wave-uniform, so every lane gets the same value. The vector load
still leaves it in VGPRs, so codegen wraps every buffer access built from those
pointers in a waterfall loop:

v_readfirstlane_b32 s39, v13
v_cmp_eq_u64_e32    vcc, s[36:37], v[10:11]
v_cmp_eq_u64_e64    s[6:7], s[38:39], v[12:13]
s_and_saveexec_b64  s[6:7], s[6:7]
buffer_store_dword  v14, off, s[36:39], 0
s_xor_b64           exec, exec, s[6:7]
s_cbranch_execnz    .LBB0_21

Reading the two halves with buffer_load(..., is_scalar=True) lands them in SGPRs, so
the descriptor is uniform by construction. Final ISA for the Stage1 kernel:

	before	after
total instructions	4275	3351 (-21.6%)
v_readfirstlane_b32	421	131
v_cmp_eq_u64_e64	174	42
s_and_saveexec_b64	134	63
s_cbranch_execnz	112	41
s_buffer_load_dwordx2	0	18

VGPR (108), SGPR (112), LDS (79872) and grid geometry are unchanged.

2. num_dispatch_cu retuned for the FP4 bucket-512 override.

Sweeping the knob on a4w4 bs=512 gives a clear minimum at 64, not the 160 currently
set (prequant_e2e, ms): 24→0.5787, 32→0.5723, 48→0.5605, 64→0.5523, 88→0.5577,
104→0.5594, 128→0.5600, 160→0.5608, 192→0.5696, 224→0.5731.

Note _select_bounded_stage1 already uses dispatch_cu=64 for bucket 512 on the
dynamic-MTPR path, so this brings the FP4 override in line with it.

grid_mult, use_tile_resource, external_grouping / external_counting and
persist_cu were swept too and are all within noise at these shapes.

Test Plan

• 8×MI355X, v4_pro, --mega-only, accuracy against the FP32 oracle for a4w4
  (bs 64/512/1024) and a8w4 (bs 512/4096).
• Latency measured as interleaved A/B against branch head at single-shape launch
  granularity, medians over 5 rounds, on a node verified idle over a 24 s sampling
  window (all 8 GPUs at 0% across 12 samples).
• Full a4w4 shape sweep, bs 1 … 32768.
• Stage1 ISA and rocprofv3 kernel/counter collection on rank 0.

Test Result

Interleaved A/B, a4w4, prequant_e2e (ms), median of 5 rounds:

shape	branch head	this PR	delta	rounds won
bs=512	0.5602	0.5531	-1.3%	4/5
bs=1024	0.7451	0.7391	-0.8%	5/5, disjoint distributions

Per round — bs=512 base 0.5585 0.5649 0.5635 0.5602 0.5592, PR
0.5511 0.5531 0.5507 0.5612 0.5542; bs=1024 base 0.7451 0.7449 0.7460 0.7432 0.7457,
PR 0.7408 0.7401 0.7382 0.7375 0.7391 (base minimum 0.7432 > PR maximum 0.7408).

Against the committed goldens: bs=512 goes from +4.7% to +3.4%, bs=1024 from
+5.7% to +4.8%. Both gate shapes are back under the +5% threshold on the median,
where bs=1024 was previously over it.

A8W4 (bs=4096, interleaved, 3 rounds): within noise, no regression.

Accuracy (unchanged from branch head): a4w4 relL2 2.292e-01 / 2.256e-01 /
2.269e-01 at bs 64/512/1024 (floor ~0.28); a8w4 relL2 4.666e-02 / 5.352e-02 at
bs 512/4096 (floor ~0.1). All 8 ranks pass.

What this does not fix

The remaining +3.4% / +4.8% over golden is a real regression that entered with
0369a804, the first merge of main into this branch, and the only commit in that
merge touching kernels/mega_moe/ is #985 (2d65dea1, "Optimize megamoe
performance"). Ladder, prequant_e2e (ms):

state	bs=512	bs=1024
golden, committed at a62685de	0.5350	0.7050
ee9fb4fc (last A4W4 commit before the merges)	0.5348	0.7272
0369a804 (first merge of main)	0.5621	0.7414
9baf8cec, 972e1a94 (later merges)	0.5668, 0.5671	0.7419, 0.7416

ee9fb4fc reproduces the golden at bs=512 to within 0.04%, so the measurement
environment is comparable to the one the goldens were captured on.

This looks like an A8W4-tuned change that A4W4 does not benefit from, which is
consistent with #985 having been validated on A8W4 while A4W4 lived on this branch.
The following were each ruled out by measurement, and are recorded so nobody
repeats them:

• The LLVM bump ([LLVM][FlyDSL] Bump up LLVM and Adapt FlyDSL to upstream LLVM/MLIR API changes #945) / compiler version. Two independent checks: reverting only
  kernels/common/buffer_ops.py and running 0369a804 on the older wheel reproduces
  the regressed numbers; and reverting all of kernels/mega_moe/ while running on the
  newer wheel fully recovers them.
• mega_moe_config.py. Keeping Optimize megamoe performance #985's config and reverting only the four kernel
  files recovers to 0.5353 / 0.7306.
• Stage2. 176.62 → 177.16 µs, flat.
• Instruction count / issue bound. This PR removes 21.6% of Stage1's instructions
  and buys 0.8–1.3%.
• Occupancy, registers, LDS, rank skew. All identical; per-rank mean/max spread is
  0.0004 vs 0.0005 ms.
• Shader work. In the slower build, SQ_BUSY_CYCLES -1.8%, SQ_WAIT_ANY -2.6%,
  SQ_WAIT_INST_ANY -5.0%.

A build that does strictly less shader work, at the same occupancy, with no rank skew,
yet takes 7.7% longer points at intra-kernel workgroup scheduling / critical path
rather than at any single hot instruction sequence. Resolving that needs per-wave (ATT)
data or knowledge of the intended payload-overlap schedule.

@GwilliamHu — #985's dispatch geometry and skewed-expert scheduling changes are the
remaining suspect; could you take a look at whether the payload-overlap schedule
assumes the A8W4 payload size? A FP4 row is 896 i32 against A8W4's 1792.

@Yaowu-Xiong — the num_dispatch_cu change lands in your FP4 override; happy to drop
it from this PR if you would rather retune that table as a whole.

[jhinpan]

Closing this. The measured effect is too small to justify review bandwidth: bs=512 is -1.3% but the base and patched distributions overlap (4/5 rounds, p≈0.19), so it is not separable from run-to-run noise; bs=1024 is a clean -0.8% (5/5, disjoint distributions) but 0.8% against a 5% regression is not worth a PR.

The diagnosis is the useful part and it now lives on #972: #972 (comment) — the regression is traced to #985, with six hypotheses ruled out by measurement and the remaining evidence pointing at intra-kernel workgroup scheduling rather than any hot instruction sequence.

The branch megamoe-a4w4-dispatch-perf stays up if anyone wants the scalar-load change (Stage1 4275 → 3351 static instructions) or the num_dispatch_cu = 64 value.