[CUDA] Block-scaled MatMul decode GEMV: tensor cores, packed FP4 decode and M-tiling

docs/contrib_ops/cuda/matmul_block_scaled_fp4.md

@@ -107,7 +107,7 @@ quantization, CUTLASS setup, and underutilized tensor-core GEMM work.
 - `block_size == 16`,
 - `K % 32 == 0`.
 
-Each warp computes one output element `Y[row, col]`. A lane consumes 32 K
+Each warp computes one output column `col`. A lane consumes 32 K
 elements per iteration, which is exactly two 16-element scale blocks. The kernel
 loads:
 
@@ -118,6 +118,88 @@ loads:
 The per-block scales are folded into the partial sums and `weight_scale_2` is
 applied once after the warp reduction. Optional bias is fused in lane 0.
 
+### Row tiling
+
+A warp produces `RowsPerBlock` rows of `Y` at once (1, 2 or 4). The packed weight
+load and the E2M1 decode are shared by all rows in the tile, which matters for
+speculative decoding / MTP verify where `M = N_spec + 1 > 1`. This trades grid
+parallelism for reuse, because `M` no longer contributes to `gridDim.y`, so it is
+only enabled when the column grid `ceil(N / 8)` covers at least one full wave of
+SMs on its own. Measured on H200 (132 SMs, `M = 4`, FP16):
+
+| Shape | N | column blocks | `RowsPerBlock = 4` vs `1` |
+| --- | --- | --- | --- |
+| `lm_head` | 248320 | 31040 | 615.8 -> 537.7 us (1.15x) |
+| shared `down_proj` | 2048 | 256 | 4.30 -> 3.33 us (1.29x) |
+| shared `gate_up_proj` | 512 | 64 | 3.47 -> 4.27 us (0.81x) - gated off |
+
+Per-row fp32 accumulation order does not depend on `RowsPerBlock`, so results are
+bit-identical across tilings. Set `ORT_FP4_GEMV_ROW_TILING=0` to force
+`RowsPerBlock == 1`.
+
+Note that the tensor-core sub-path below takes precedence on SM80+ whenever
+`K % 128 == 0`, which covers most production shapes. Row tiling is therefore what
+actually runs on pre-SM80 devices or when `K` is not a multiple of 128.
+
+### Tensor-core sub-path (SM80+)
+
+On SM80 and newer, when `K % 128 == 0` and `M <= 8`, the warp reduction above is
+replaced by `mma.m16n8k16`, so a warp produces **16 output columns** at once
+instead of one. Set `ORT_FP4_GEMV_MMA=0` to fall back to the scalar path.
+
+The scalar path re-reads the whole `A` tile once per output column: at
+`RowsPerBlock = 4` a warp pulls 16 KB of activation for 1 KB of weight, a 16:1
+amplification that dominates `lm_head`. Producing 16 columns per warp cuts those
+re-reads by 16x.
+
+The fragments are free because the **weight goes in the mma `A` slot** and the
+**activation in the mma `B` slot**: `B` is `[N, K]` row-major, which is exactly the
+`A`-row-major fragment, and `A` is `[M, K]` row-major, which is exactly the
+`B`-col-major fragment. No transpose, no `ldmatrix`, no shared-memory staging. The
+mma `M` extent becomes the column count (16) and the mma `N` extent becomes `M`.
+
+A lane does not own a whole dot product: with `g = lane >> 2` and `t = lane & 3`,
+it supplies the weights of output columns `g` and `g + 8` (the `A` fragment) and
+activation row `g` (the `B` fragment), and the accumulator it receives back covers
+output rows `2t` and `2t + 1` of those two columns. Loads are therefore keyed off
+`g` and stores off `t`; the kernel source carries the full fragment table, and the
+`GemvTensorCoreLaneOwnership*` tests probe the mapping with a one-hot activation.
+
+The K axis is then permuted so the four k-slots a lane needs are contiguous in
+memory, which is legal because K is a reduction axis and the same permutation is
+applied to both operands. A window is 128 K elements = 64 packed bytes; lane
+`(g, t)` owns elements `[32t, 32t + 32)`, i.e. one `uint4` of weight and one
+`uint4` of activation, spanning exactly two 16-element scale blocks.
+
+Because the mma sums across the four `t` lanes, which hold *different* scale
+blocks, the accumulator cannot be flushed per block. The E4M3 scale is instead
+folded into the decoded weight before the mma. That is exact in both FP16 and
+BF16: E2M1 magnitudes carry 2 significand bits and E4M3 scales carry 4, so the
+product needs at most 6, inside FP16's 11 and BF16's 8; the range is safe too
+(max `6 * 448 = 2688`, min `0.5 * 2^-9 = 2^-10`).
+
+`KSplit` warps per block take a strided share of the K windows and reduce through
+shared memory. Without it, 16 columns per warp yields 16x fewer warps than the
+scalar path and the small MLP shapes lose more to idle SMs than they gain. The
+launcher picks `ColTiles = 4, KSplit = 2` when the column grid alone still covers
+a wave of SMs, and otherwise `ColTiles = 1` with as many `KSplit` warps as there
+are K windows.
+
+Measured on H200 (132 SMs, `M = 4`, FP16), scalar -> tensor core:
+
+| Shape | N | K | scalar | tensor core | speedup |
+| --- | --- | --- | --- | --- | --- |
+| `lm_head` | 248320 | 2048 | 537.6 us | 108.3 us | **4.96x** |
+| shared `gate_up_proj` | 512 | 2048 | 3.48 us | 2.99 us | 1.17x |
+| shared `down_proj` | 2048 | 512 | 3.32 us | 2.52 us | 1.32x |
+
+Over a Qwen3.6 NVFP4 MTP decode step (121 FP4 GEMV launches) this is
+0.949 -> 0.448 ms/step, and 9.54 -> 8.99 ms/step end to end.
+
+The fp32 accumulation order differs from the scalar path, so this path is not
+bit-identical to it. Max relative error against an fp64 reference is unchanged
+(2.4e-04 .. 3.5e-04, i.e. NVFP4 quantization noise, not accumulation noise).
+
 This kernel reads the original unswizzled `[N, K / 16]` scale layout. Experiments
 with the native SM120 swizzled scale layout for GEMV were slower; see
 [matmul_block_scaled_fp4_experiments.md](matmul_block_scaled_fp4_experiments.md).
@@ -196,6 +278,8 @@ GEMM and the original scale tensor for the other paths.
 | Variable | Default | Meaning |
 |----------|---------|---------|
 | `ORT_MATMUL_BLOCK_SCALED_FP4_NATIVE_SM120` | `0` | Enables the opt-in native SM120 NVFP4 x NVFP4 GEMM path when the shape and device guards pass. |
+| `ORT_FP4_GEMV_MMA` | `1` | Set to `0` to disable the decode GEMV tensor-core sub-path (`mma.m16n8k16`, SM80+, `K % 128 == 0`) and use the scalar warp-reduction path. |
+| `ORT_FP4_GEMV_ROW_TILING` | `1` | Set to `0` to force `RowsPerBlock == 1` in the scalar decode GEMV. |
 
 The default remains the existing weight-only semantics: decode GEMV for small
 `M`, otherwise dequantize `B` and call cuBLAS.

docs/contrib_ops/cuda/matmul_block_scaled_fp8.md

@@ -129,6 +129,42 @@ This path avoids a materialized dequant buffer and runs on all supported CUDA
 architectures because it uses regular FP8 conversion and warp-shuffle reduction,
 not architecture-specific block-scaled tensor cores.
 
+### 4.1 Tensor-Core Sub-Path (SM80+)
+
+When the device is SM80 or newer and, in addition to the predicate above,
+
+- `K % 64 == 0` and `K >= 256`,
+- `block_size % 64 == 0`,
+- `M <= 8` (the mma "N" extent),
+
+the GEMV runs `MatMulBlockScaledFp8MmaGemvKernel`, which replaces the FP32 FMA
+dot products with `mma.m16n8k16` (FP32 accumulate). The FMA kernel is ALU bound
+for `M > 1` because it re-widens `A` and `B` to FP32 for every row/column pair;
+the tensor-core kernel cuts instructions per weight byte roughly 10x.
+
+The weight is fed to the mma **A** operand and the activation to the **B**
+operand, so both fragments match the tensors' natural row-major layouts. The mma
+"M" extent is therefore the output column count (16 columns per warp) and the mma
+"N" extent is `M` (up to 8 rows). A lane does not own a whole dot product: with
+`g = lane >> 2` and `t = lane & 3`, it supplies the weights of output columns
+`g` and `g + 8` (the A fragment) and activation row `g` (the B fragment), and the
+accumulator it receives back covers output rows `2t` and `2t + 1` of those two
+columns. Loads are therefore keyed off `g` and stores off `t`; the kernel source
+carries the full fragment table, and the `GemvTensorCoreLaneOwnership*` tests
+probe the mapping with a one-hot activation. Fragment loads are made fully
+coalesced by permuting the K axis - K is a reduction axis, so any permutation
+applied to both operands leaves the result unchanged - which lets each lane load
+one contiguous `uint4` of weight bytes per 64-element K window and feed four mma
+instructions with it. `KSplit` warps per block take a strided share of the K
+windows and are reduced through shared memory, which restores the memory-level
+parallelism lost by giving each warp 16 columns instead of 1-4.
+
+Every individual product is exact (E4M3 to FP16/BF16 is lossless) and the mma
+accumulates in FP32 just like the FMA path, but the summation order differs, so
+the result is **not** bit-identical to the FMA kernel. Note also that the mma
+path is preferred over the FMA kernel by default on SM80+, including at `M == 1`.
+Set `ORT_FP8_GEMV_MMA=0` to fall back to the FMA kernel.
+
 ---
 
 ## 5. Default Path - Dequantize + cuBLAS

docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md

@@ -18,8 +18,9 @@ Related documentation:
 3. [Prefill Bottleneck - Weight Dequantization](#3-prefill-bottleneck---weight-dequantization)
 4. [Optimization - Vectorized Dequantization Kernel](#4-optimization---vectorized-dequantization-kernel)
 5. [Decode GEMV - Memory-Level Parallelism](#5-decode-gemv---memory-level-parallelism)
-6. [Benchmark Commands](#6-benchmark-commands)
-7. [Lessons](#7-lessons)
+6. [Decode GEMV - Tensor Cores](#6-decode-gemv---tensor-cores)
+7. [Benchmark Commands](#7-benchmark-commands)
+8. [Lessons](#8-lessons)
 
 ---
 
@@ -234,6 +235,11 @@ Below `N = 4096` the wider tiles leave too few warps to fill the GPU, and for
 `M > 1` the extra live registers (accumulators plus pre-issued loads) cost more
 than the added parallelism returns. Both measured slower, hence the guards.
 
+> Superseded for `M > 1` by section 6.4, which re-tunes `ColsPerWarp` / `Unroll`
+> for the speculative-decode tiles. Because `Unroll` changes the K chunk each
+> lane accumulates first, the `M > 1` dispatch there is not bit-identical to
+> `<RowsPerWarp, 1, 1>` (last-ulp only; the accumulation is still FP32).
+
 ### 5.4 Results (H200, `M = 1`, CUDA graph, us, includes 0.68 us node overhead)
 
 | Shape (N x K) | cuBLAS FP16 | GEMV before | GEMV after | vs before | vs cuBLAS |
@@ -253,7 +259,132 @@ preferred on speed alone.
 
 ---
 
-## 6. Benchmark Commands
+## 6. Decode GEMV - Tensor Cores
+
+Section 5 tuned memory-level parallelism at `M = 1`. At `M = 4` - the width of a
+speculative-decode / MTP verify forward - the kernel is limited by something
+else. With `RowsPerWarp = 4` a lane executes roughly 240 instructions per 32
+weight bytes, only 128 of which are the FMAs that do useful work, and effective
+bandwidth falls from about 2.35 TB/s at `M = 1` to about 1.25 TB/s at `M = 4`.
+More ILP cannot fix that; the dot products have to leave the FMA pipe.
+
+### 6.1 Design
+
+`MatMulBlockScaledFp8MmaGemvKernel` uses `mma.m16n8k16` with FP32 accumulation.
+The operand assignment is the key decision:
+
+| mma operand | fed from | why it fits |
+|---|---|---|
+| `A[16, 16]` row-major | weight `[16 output cols][16 k]` | `B` is `[N, K]` row-major |
+| `B[16, 8]` col-major | activation `[16 k][8 rows]` | `A` is `[M, K]` row-major |
+| `D[16, 8]` | `y[16 output cols][8 rows]` | |
+
+So the mma "M" extent is the output column count and the mma "N" extent is `M`.
+At `M = 4` half the mma N lanes are idle, which is irrelevant: the kernel is
+bound by weight traffic and instruction issue, and both improve about 10x per
+weight byte.
+
+The naive fragment load is badly coalesced - a lane needs bytes
+`{2t, 2t+1, 2t+8, 2t+9}` of a row, which spreads a warp across 16 rows x 16
+bytes and over-fetches every 32-byte sector 2x. The fix is to **permute the K
+axis**. K is a reduction axis, so any permutation applied to *both* operands
+leaves the result unchanged. Inside a 64-element K window the permutation used is
+
+```
+mma k-slot (of step j)  ->  actual k
+  2t,   2t+1                 16t + 4j,     16t + 4j + 1
+  2t+8, 2t+9                 16t + 4j + 2, 16t + 4j + 3
+```
+
+so lane `(g = lane >> 2, t = lane & 3)` loads one contiguous `uint4` of weight
+bytes `[16t, 16t + 16)` and the matching 32 activation bytes, four lanes cover 64
+contiguous bytes of one weight row, and that single `uint4` feeds all four mma
+steps.
+
+16 columns per warp gives about 8x fewer warps than the FMA kernel, which alone
+costs more in lost memory-level parallelism than the instruction saving is worth.
+`KSplit` warps per block therefore take a strided share of the K windows and are
+reduced through shared memory at the end. `KSplit = 8` for `N >= 8192` (the
+column count already fills the grid) and 16 otherwise.
+
+Preconditions: SM80+, `K % 64 == 0`, `K >= 256`, `block_size % 64 == 0`, `M <= 8`.
+Otherwise the FMA kernel runs unchanged. `ORT_FP8_GEMV_MMA=0` forces the FMA
+kernel for A/B testing in a single binary.
+
+### 6.2 Accuracy
+
+Not bit-identical to the FMA kernel (different summation order), but not less
+accurate either. E4M3 to FP16 is lossless, E4M3 to BF16 is lossless, FP16 x FP16
+products are exact in FP32, and the mma accumulates in FP32 exactly as the FMA
+path does. Scored against an FP64 CPU reference on the shapes below, the maximum
+error is *identical* for the two kernels (2-4e-4, i.e. pure FP16 output
+rounding).
+
+### 6.3 Results (H200, us, standalone, `fp8_gemv_m4_bench.cu`)
+
+| Shape (N x K) | M | cuBLAS FP16 | FMA kernel | mma kernel | vs FMA | vs cuBLAS |
+|---|---|---|---|---|---|---|
+| 8192 x 2048 | 1 | 11.0 | 6.3 | **5.1** | 1.23x | 2.16x |
+| 4096 x 2048 | 1 | 8.4 | 4.8 | **4.0** | 1.20x | 2.09x |
+| 2048 x 4096 | 1 | 9.0 | 5.1 | **4.2** | 1.21x | 2.13x |
+| 512 x 2048 | 1 | 6.7 | 3.3 | **3.1** | 1.06x | 2.12x |
+| 8192 x 2048 | 4 | 11.0 | 9.8 | **5.2** | 1.87x | 2.10x |
+| 4096 x 2048 | 4 | 8.5 | 6.9 | **4.1** | 1.69x | 2.09x |
+| 2048 x 4096 | 4 | 8.5 | 8.0 | **4.4** | 1.82x | 1.92x |
+| 512 x 2048 | 4 | 6.8 | 4.7 | **3.3** | 1.43x | 2.08x |
+| 8192 x 2048 | 8 | 10.9 | 17.5 | **5.7** | 3.09x | 1.93x |
+| 2048 x 4096 | 8 | 8.5 | 13.0 | **4.5** | 2.90x | 1.90x |
+
+The mma kernel is faster at every measured `M`, so it is preferred whenever its
+preconditions hold rather than only for `M > 1`. Note also that the FMA kernel
+crosses over and loses to cuBLAS at `M = 8`, while the mma kernel stays about 1.9x
+ahead.
+
+End to end on a 40-layer Qwen3.6-35B-A3B NVFP4 MTP decode (130 FP8 matmul nodes
+per step, `M = 4`), CUDA graphs on:
+
+| | FMA kernel | mma kernel |
+|---|---|---|
+| FP8 GEMV kernel time | 1.052 ms/step | **0.713 ms/step** |
+| total kernel time | 7.368 ms/step | **7.021 ms/step** |
+| wall | 9.80 ms/step | **9.54 ms/step** |
+
+No other kernel family moved. This optimization is kept.
+
+### 6.4 FMA fallback re-tune for `M > 1`
+
+The FMA kernel still runs when the tensor-core preconditions do not hold (pre-SM80,
+`K < 256`, `K % 64 != 0` or `block_size % 64 != 0`), so the `M > 1` tiles were
+re-tuned there as well. Widening `A` to FP32 is now hoisted out of the column loop
+(one widening per row instead of one per row/column pair), which makes `ColsPerWarp`
+profitable at `M > 1` for a second reason beyond memory-level parallelism:
+
+| Condition | Config |
+|---|---|
+| `2 <= M <= 2, N >= 8192` | `<2, 4, 1>` |
+| `2 <= M <= 2, N >= 2048` | `<2, 2, 1>` |
+| `2 <= M <= 2` otherwise | `<2, 1, 2>` |
+| `3 <= M <= 4, N >= 4096` | `<4, 4, 1>` |
+| `3 <= M <= 4, N >= 2048` | `<4, 2, 2>` |
+| `3 <= M <= 4` otherwise | `<4, 1, 2>` |
+| `M > 4` | `<8, 1, 1>` |
+
+Measured on H200 (us, `M = 4`, versus the previous `<R, 1, 1>` and cuBLAS FP16):
+
+| Shape (N x K) | cuBLAS | `<4, 1, 1>` | tuned |
+|---|---|---|---|
+| 8192 x 2048 | 10.9 | 13.7 | **9.7** (`<4, 4, 1>`) |
+| 4096 x 2048 | 8.3 | 8.1 | **6.9** (`<4, 4, 1>`) |
+| 2048 x 4096 | 8.3 | 9.0 | **7.9** (`<4, 2, 2>`) |
+| 512 x 2048 | 7.3 | 5.0 | **4.6** (`<4, 1, 2>`) |
+
+The hoisting itself is bit-identical (the per-lane `fmaf` sequence is unchanged),
+but a different `Unroll` changes which K chunk a lane accumulates first, so the
+re-tuned dispatch is a last-ulp change relative to `<RowsPerWarp, 1, 1>`.
+
+---
+
+## 7. Benchmark Commands
 
 The commands below use `ORT_REPO` and `ORT_BUILD` so they can be copied without
 editing developer-specific paths. Set them once:
@@ -307,7 +438,7 @@ CUDA_VISIBLE_DEVICES=0 "$ORT_BUILD/onnxruntime_provider_test" \
 
 ---
 
-## 7. Lessons
+## 8. Lessons
 
 - The prefill path is memory bound on weight dequantization, not on the GEMM;
   latency there scales with `N*K` and is independent of `M`.

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc

@@ -196,6 +196,7 @@ Status MatMulBlockQuantizedFp4Weight::ComputeImpl(OpKernelContext* context) cons
         k_i,
         SafeInt<int>(block_size_),
         std::is_same<T, BFloat16>::value,
+        GetDeviceProp(),
         Stream(context));
   }