[ROCm][MoE] W4A16 MoE routing-distribution benchmark suite for the gfx11 prefill GEMM

[roberteg16]

Summary

Adds benchmarks/kernels/moe_w4a16_bench/, a small benchmark suite that maps where the gfx11 W4A16 MoE prefill GEMM (torch.ops._rocm_C.moe_gemm_w4a16, the rdna_moe_gemm kernel) wins or loses against the Triton reference (fused_moe_kernel_gptq_awq via invoke_fused_moe_kernel_hybrid_triton) as a function of the token→expert routing distribution. MoE prefill performance is routing-dependent: a custom WMMA kernel's only lever over Triton is keeping a hot expert's weights cache-resident across its blocks, so the win is entirely a function of routing skew and per-block fill. This suite makes that surface measurable and reproducible across synthetic routing archetypes (balanced → heavily skewed) and a sweep of token counts, on identical inputs with Triton as the correctness oracle. It is benchmark-only (no kernel, library, or runtime changes) and adds nothing to the build or import graph of vLLM proper.

What this adds

• benchmarks/kernels/moe_w4a16_bench/gen_distributions.py: a numpy-only routing-distribution generator. It exposes an Archetype enum (balanced, uniform, zipf1, zipf2, hot16) and an importable synth_topk(arch, T) that builds topk_ids with a controlled per-expert popularity (gumbel-top-k sampling), so the benchmark constructs routing in-process with no on-disk dataset step. Run as a script it prints a GPU-free moe_align padding/fill space-map (per-expert load stats, dead-expert count, WMMA-useful %, and skinny-block fraction at block_m 16 and 32), useful for choosing which shapes are worth GPU time. It deliberately imports neither torch nor vLLM so it runs on any machine.
• benchmarks/kernels/moe_w4a16_bench/bench_moe_gemm.py: the GPU benchmark. For each (archetype, T) it times the rdna_moe_gemm op against the Triton reference on the same tensors and shared moe_align@32 layout, reporting Triton µs, kernel µs, speedup %, achieved kernel TFLOP/s, and the relative error vs Triton. It covers both MoE prefill GEMMs via --gemm 1,2 (gemm1 = up/gate, top_k=8, K=2048 N=1024; gemm2 = down, top_k=1, K=512 N=2048), at the single production block_m=32 the op supports. Timing uses triton.testing.do_bench (L2 flushed between iterations) so each launch is measured cold, matching the production pattern where the same GEMM does not run back-to-back. The achieved TFLOP/s is computed over the real padded work (valid_blocks * block_m, i.e. num_tokens_post_padded), not just the useful routed rows, so it reflects the hardware throughput the kernel actually delivers rather than an inflated useful-only figure.

The suite reuses existing vLLM building blocks rather than reimplementing them: pack_int4_exllama_shuffle for the ExLlama INT4 weight layout, moe_align_block_size for the block layout, and invoke_fused_moe_kernel_hybrid_triton for the Triton reference; only the routing generator and the harness are new.

Why this is useful

Real trained routers are skewed (zipf-like), and the rdna_moe_gemm kernel wins more the more skewed the routing is, while losing on the synthetic perfectly-flat balanced case a real router never produces. Reviewers and future tuners can use this suite to confirm the kernel sits in its win regime for realistic routing, find the token-count crossover where it starts winning for a given archetype, and correlate the win with the per-block fill (use%) that the space-map predicts without a GPU. Example crossover finding on gfx1151: for zipf1 gemm1 under the cold-cache methodology the kernel turns positive around T≈1024 (block fill ~63%) and grows to ~+31% by T=4096.