[CUDA] Coalesce QMoE MXFP4/NVFP4 weight dequantization

[tianleiwu]

Description

Speeds up MXFP4/NVFP4 QMoE weight dequantization (LaunchQMoEDequantizeFp4Weights / LaunchQMoEDequantizeNvfp4Weights) by fixing the memory access pattern.

The existing scalar kernels map one thread to one output element with k varying fastest. Packed weights are stored [E, K, N/2] (n-packed) while the output is [E, N, K], so consecutive lanes read packed bytes packed_n apart — a separate 32-byte sector per lane for half a byte of payload, plus a 64-bit div/mod per element. Measured at ~5% of HBM peak on a 40-layer NVFP4 MoE prefill (3.8 ms per launch, 307 ms per forward).

This PR adds QMoEDequantizeFp4WeightsVecKernel, which tiles the output instead:

• a block covers 64 consecutive rows x 64 consecutive k; threadIdx.x selects the k group, threadIdx.y the row;
• each thread emits 8 values as one 16-byte store, so 8 lanes of the same row issue one 128-byte contiguous store (a warp covers 4 rows in 4 requests instead of 32);
• kTileN = 64 rows is exactly 32 packed bytes, so a block fully consumes every packed sector it touches;
• the block-scale byte and per-expert global scale are read once per 8 values;
• the index decomposition is pure grid arithmetic — no integer division.

The mapping matters more than the vector width. An earlier revision that gave each thread 16 (then 32) consecutive k of a single row reported ~95% of Max Bandwidth at only ~26% DRAM throughput in Nsight Compute: bound by memory requests, not DRAM. Widening the per-lane store leaves requests-per-byte unchanged, and measurement confirmed no improvement from 16 -> 32 values per thread.

One template covers both codecs: kE4M3Scale selects the NVFP4 scale codec (Float8E4M3FN, block 16) over the MXFP4 one (Float8E8M0, block 32).

Also in this PR:

• DecodeFp4E2M1 is made branch-free by assembling the float bits directly. The previous runtime-indexed local table compiles to a constant-bank load that the hardware replays once per distinct address in a warp, and neighbouring weights rarely share a code.
• DecodeFloat8E4M3FN is moved earlier in the file (no behavior change) so the shared kernel can use it.
• Added the missing cudaGetLastError check after the scalar dequantize launches.

Correctness

The tiled kernel reproduces the scalar kernels' indexing exactly (weight nibble selection, block-scale index, output index); the scalar kernels remain the fallback for any shape the tiled mapping cannot cover (odd n, k not a multiple of 64, or grid dimensions over 65535). Every shape produced by the QMoE quantizers takes the new path, so the existing test_qmoe_fp4_cuda.py and test_qmoe_nvfp4_cuda.py end-to-end tests cover it.

Motivation and Context

Prefill of NVFP4/MXFP4 MoE models spends a significant fraction of time in weight dequantization; this removes it as a bottleneck.

[Copilot]

Pull request overview

This PR optimizes CUDA QMoE MXFP4/NVFP4 weight dequantization by adding a coalescing-friendly, tiled/vectorized kernel path while keeping the existing scalar kernels as a fallback.

Changes:

• Added QMoEDequantizeFp4WeightsVecKernel and gating (QMoEDequantizeFp4VecApplies) to use a tiled output mapping for better memory coalescing.
• Updated FP4 and NVFP4 dequantize launchers to use the vector kernel when applicable and added cudaGetLastError() checks after the scalar launches.
• Made DecodeFp4E2M1 branch-free and moved DecodeFloat8E4M3FN earlier for shared use.