Evaluate native IU4/IU8 WMMA for grouped-INT MoE prefill on gfx1151

[randomvariable]

Problem

RDNA3.5 (gfx1151) WMMA provides native IU4 and IU8 integer matrix inputs (ISA §7.9, Table 33), but no path in this fork uses them. Grouped-INT4 MoE currently runs through TritonWNA16Experts, which dequantises packed weights to bf16 and issues a bf16 tl.dot. The integer matrix hardware is idle.

This does not affect decode: at batch 1 with top-k routing, each expert GEMM has M=1, so the workload is GEMV — bound by weight streaming and dispatch count, not MAC throughput. Padding M=1 into a 16-wide WMMA tile wastes 15/16 of the matrix unit regardless of input type.

It does affect prefill, where per-expert M is large enough for matrix cores to matter, and long-context prompts make prefill the dominant cost.

Constraint

IU4/IU8 WMMA requires both operands to be the same integer type, accumulating int32. There is no mixed int4 x bf16 shape. W4A16 / W8A16 checkpoints therefore cannot use it directly — dequantise-then-bf16-dot is the correct lowering for those.

Exploiting integer WMMA requires an activation-quantised scheme (W8A8-int8, or W4A8 with upconversion of weights to IU8 before the MMA).

Scope

1. Measure prefill kernel time share for a grouped-INT4 MoE model on gfx1151 with rocprofv3 --kernel-trace, to establish how much of prefill is MoE GEMM versus attention and routing.
2. Determine whether an int8-activation MoE path is reachable in-tree for gfx1151, and which checkpoint schemes would qualify.
3. If viable, prototype an IU8 WMMA expert GEMM for the prefill (large-M) branch only, keeping the existing bf16 dequant path for decode.

Acceptance criteria

• Kernel-trace evidence quantifying MoE GEMM share of prefill time before any kernel work.
• Any prototype must show a measured prefill-time reduction of at least 25% on a saturated grid at a fixed prompt length, with unchanged output token IDs against the bf16 path under greedy decoding.
• Decode throughput must not regress.
• Tile dimensions stay multiples of 16 and per-workgroup LDS stays within 64 kB (see docs/contributing/kernel_targets/rdna35.md).

Notes

Related but distinct: tuned Triton WNA16 configs for gfx1151 do not exist (vllm/model_executor/layers/fused_moe/configs/ has no gfx1151 entries), which is a separate and cheaper optimisation applying to both prefill and decode.