Add Optimized MoE Routing Path

[amd-wsung102]

Summary

This PR is an extension of this merged PR: #582

MoE routing fusion previously only paid off for T ≤ 16, where a single kernel does softmax + top-k gating and counting-sort in LDS. Above that, moe_softmax_sort_flydsl fell back to the unfused two-kernel chain, and the earlier serial multi-tile fusion experiment was actually slower than the fallback (e.g. ~20.4 µs vs ~17.6 µs kernel time at T=64 on DeepSeek-R1) because it serialized 16-token tiles inside one workgroup.

This PR adds a dedicated wide routing path for T = 17–64 on gfx95 that beats the unfused baseline on kernel, CUDA-graph, and eager time. The T ≤ 16 single-kernel path is unchanged, and any unsupported shape/arch still uses the existing fallback.

Performance improvement: 1.16x - 1.2x speedup for graph mode.

Optimizations Applied

• Two-stage design instead of one serial workgroup. Rather than replaying 16-token tiles in a single block, routing is split into two launches under one host dispatch:
  • Parallel gating — the softmax/top-k gating kernel keeps its token blocks spread across CUs and writes compact per-assignment expert-ID records plus top-k weights.
  • Expert-parallel sort — one workgroup per expert builds a block-local histogram via LDS atomics, does a padded prefix sum, and scatters directly to unique output slots.
• Compact T×topk records replace the dense T×E LDS mesh and the O(T×E) clear/count/scatter scans, dropping sorting work to O(T×topk).
• LDS atomic ranking (atomicrmw add, workgroup scope) gives each winning assignment a unique per-expert rank, so scatter needs no serialized counting.
• Preserved contracts: expert masking, weight renormalization, sentinel padding, sorted_expert_ids local-index remapping under EP, and the parallel moe_buf zeroing blocks all behave identically to the existing paths.
• Caching: stream-local scratch allocation and a compiled-dispatch cache keep host overhead off the hot path.

Validation

Correctness (gfx950):

PYTHONPATH=./ pytest -q tests/kernels/test_moe_sorting.py -k 'moe_softmax_sort'
34 passed

Coverage added: tile boundaries T = 17, 31, 32, 33, 48, 63, 64 for both supported shapes, expert masking, renormalization disabled, dispatch-label assertions, and fallback correctness above T = 64.

Latencies are in microseconds; speedup = unfused ÷ fused (higher favors the new path). Representative points from the full T = 17–64 sweep:

Test Result - DeepSeek-R1 — E=256, topk=8, model_dim=7168, bf16

T	unfused_eager	fused_eager	unfused_graph	fused_graph	unfused_kernel	fused_kernel	eager speedup	graph speedup	kernel speedup
17	80.4	24.1	17.3	14.5	17.2	11.3	3.34×	1.19×	1.53×
24	80.7	23.9	17.3	15	17.4	11.3	3.38×	1.16×	1.54×
32	80.8	24.1	17.3	14.5	17.5	11.4	3.35×	1.19×	1.53×
48	81.9	24.3	17.6	14.7	17.3	11.3	3.37×	1.20×	1.53×
64	84.5	24.3	17.6	15.1	17.6	11.4	3.48×	1.17×	1.54×

Test Result - GPT-OSS 120B — E=128, topk=4, model_dim=2880, bf16

T	unfused_eager	fused_eager	unfused_graph	fused_graph	unfused_kernel	fused_kernel	eager speedup	graph speedup	kernel speedup
17	81.2	22.8	15.6	13.3	15	9.4	3.56×	1.18×	1.60×
24	79.9	20.8	15.6	13	14.1	8.7	3.85×	1.20×	1.61×
32	80.8	21.3	15.6	13	14.1	8.9	3.80×	1.20×	1.59×
48	81.6	20.6	15.8	13.2	14.3	9	3.95×	1.20×	1.59×
64	80.4	20.8	15.8	13.3	14.4	9.1	3.86×	1.19×	1.59×

Test Result - MiniMax M3 428B — E=128, topk=4, model_dim=6144, bf16

T	unfused_eager	fused_eager	unfused_graph	fused_graph	unfused_kernel	fused_kernel	eager speedup	graph speedup	kernel speedup
17	80.5	22.4	15.7	13.2	14.5	9.2	3.59×	1.19×	1.57×
24	82.1	22.9	15.7	13.1	15	9.5	3.59×	1.20×	1.57×
32	81.4	22.1	15.7	13.5	14.7	9.6	3.69×	1.16×	1.52×
48	82	22.8	15.8	13.2	14.9	9.3	3.59×	1.19×	1.61×
64	84.3	22.6	15.7	13.2	15	10	3.72×	1.19×	1.50×

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.