onnxruntime

e56b1073 - [CUDA] Speed up the NVFP4 QMoE decode GEMV and enable it for MTP verify (#31159)

[CUDA] Speed up the NVFP4 QMoE decode GEMV and enable it for MTP verify (#31159) ### Description Four changes to the fused NVFP4 QMoE decode GEMV. Stacked on #31154 — **review only the top four commits**; the base branch is that PR. 1. **Packed E2M1 dequantize** (`Fp4I2FConverter::decode_quad`) — decode a whole 32-bit weight word (eight codes) per step instead of one code at a time. 2. **`ORT_FP4_GEMV_DEFAULT_TILING`** — env switch to bypass the autotuner and take the default tiling, for A/B and for avoiding autotune cost in short runs. 3. **Cut memory and ALU traffic in the decode GEMV.** 4. **`kMaxProfiledExpandedRows` 8 -> 64** so MTP verify steps stay on the GEMV path. ### Motivation and Context Prior profiling established that this kernel is **ALU-pipeline bound**, not memory- or tiling-bound. On the actual Qwen3.6 decode shapes (`hidden=2048`, `inter=512`, `E=256`, `top_k=8`, bf16, SwiGLU), ncu reported for the FC1 SwiGLU-fused GEMV: > ALU 78.9%, DRAM 7.3%, occupancy 21% (register-limited) That is why the levers here are instruction-count levers. Two things were measured and explicitly **dropped** because of it: smaller `CtaN` tiling (the autotuner still picks `threads64`/`CtaN=8`; `CtaN=4` never wins because the kernel is compute-bound, not occupancy-bound), and halving scale bandwidth by storing combined scales as 1-byte e4m3 (DRAM is only ~7%, so it cannot move the needle). All numbers below: 1x H200 SXM (SM90, 132 SM, ~4.8 TB/s HBM), CUDA 13.0, Qwen3.6-35B-A3B-NVFP4 + MTP `N=3` (verify batch `M=4`). ### 1. Packed E2M1 dequantize `prmt` selects four bytes per instruction, so a 4-element magnitude lookup costs one instruction instead of four. Bit-identical to the per-element path (same magnitude tables, same sign handling). The FP4 GEMV kernel SASS shrinks ~30%, and the two QMoE GEMVs drop: | kernel | before | after | |---|---:|---:| | fc1 (SwiGLU-fused) | 33.2 µs | **26.2 µs** | | fc2 | 30.2 µs | **22.2 µs** | ### 2. Cut memory and ALU traffic — −0.46 ms/step (−5.1%) The scales of the `CtaN` columns a block owns sit `Interleave` elements apart, so for the non-interleaved ColumnMajor layout (`Interleave == 1`) the whole `CtaN`-wide scale vector is contiguous and can be fetched with one wide access instead of `CtaN` scalar ones. This matters far more than the byte count suggests: with a groupwise scale (NVFP4 `GroupSize = 16`) and `StepK = 8`, a warp's 32 lanes cover 16 distinct scale rows that are `n` elements apart, so *every* scale load touches 16 different sectors — `CtaN * 16` sectors, using 2 bytes out of each 32-byte sector. Per-kernel (graph OFF, 40 launches/step each): | kernel | before | after | |---|---:|---:| | `moe_gemv_interleaved_swiglu_kernel` | 0.956 ms/step | **0.678 ms/step** | | `moe_gemv_kernel` | 0.700 ms/step | **0.494 ms/step** | | **family total** | **1.657 ms/step** | **1.174 ms/step** | End-to-end (4 interleaved `.so`-swap reps per arm): * before: 8.973 / 9.009 / 8.992 / 9.017 * after: 8.567 / 8.508 / 8.498 / 8.583 **8.998 -> 8.539 ms/step.** No overlap between the two sets. ### 3. `kMaxProfiledExpandedRows` 8 -> 64 The fused GEMV rejects `expanded_num_rows > kMaxProfiledExpandedRows`. Qwen3.6 is top-8, so single-token decode expands to 8 rows (accepted), but an MTP verify does not: an `(N+1)`-token verify for `num_speculative_tokens = N` expands to `(N+1) * 8` rows, i.e. **up to 64 for N=7**. Those steps fell out of the window and back onto the dequantize + CUTLASS grouped-GEMM path, which re-dequantizes all 256 experts per token. The impact of that fallback is large: with the limit at 8, the 2-token verify (expanded 16) dropped MTP to **~2.4 tok/s**; raising the limit put it at **~30–55 tok/s (12–23x)**. 64 covers the `N=3` shape used today with headroom to `N=7`. ### Tests * `onnxruntime_provider_test` FP4/FP8/QMoE: 18/18 pass. * `onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py`: 22/22 pass, including new multi-token GEMV cases and a `gemv_mode="0"` dequant-fallback companion on the identical shape, so both must match the same exact dequantized reference. ### Methodology note End-to-end deltas are quoted as **ms/step** from a fixed-step measurement, never tok/s: any numerics change alters the generated sequence and therefore the MTP acceptance rate, which swamps the speed delta. Per-kernel durations are taken with CUDA graphs **off** — `nsys --cuda-graph-trace=node` inflates durations ~35% globally and up to 3.8x for large-grid kernels. > [!IMPORTANT] > `Fp4I2FConverter::convert()` gained a `PairInterleaved` template parameter in #31154. The > packed path added here assumes the **plain** nibble order (nibble `j` of the word is logical > element `j`), which is what its `prmt` selectors encode, so it is nested inside > `if constexpr (!PairInterleaved)`. Please check that guard carefully during review — applied > without it, the pair-interleaved SM80 layout would silently decode to the wrong values.