[Tianlei Wu (tianleiwu)]

Summary

• Tune tensor-core NVFP4 GEMV dispatch to require multiple waves before selecting wide column tiles.
• Preserve a larger K split for long reductions while avoiding excessive K splitting on Qwen gate/up shapes.
• Add shape-scoped benchmark overrides, boundary tests, and CUDA contrib documentation.

Motivation

On H200, the previous dispatch could select a configuration with too few blocks for Qwen MTP shapes such as N=17408, K=5120. The updated policy keeps KSplit=8 for the longer K=8192 reduction while using KSplit=2 for K=5120, and requires sufficient grid waves before selecting wide column tiles.

This PR is the tiling follow-up to #32128 and contains no duplicate vectorized NVFP4 dequantization changes. It can be rebased/stacked on #32128 after that PR merges.

Validation

• CUDA 13.0 / SM90 build passed.
• git diff --check passed.
• Added Qwen boundary coverage for the selected tensor-core tilings.
• The local build configuration had provider unit tests disabled, so the new gtest was not executed locally.

[Copilot]

Pull request overview

Tunes CUDA NVFP4 tensor-core GEMV dispatch for Qwen MTP workloads.

Changes:

• Adds wave-based column tiling and K-split selection.
• Adds shape-scoped benchmark overrides.
• Adds boundary tests and dispatch documentation.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
matmul_block_scaled_fp4.cu	Updates GEMV tiling policy and overrides.
matmul_block_scaled_fp4_test.cc	Adds Qwen tiling cases.
matmul_block_scaled_fp4.md	Documents policy and environment variables.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Hariharan Seshadri (hariharans29)]

Review — PR #32140: [CUDA] Tune NVFP4 GEMV tiling for Qwen MTP

Small, focused perf-tuning PR (+131/-22 across 5 files) with strong empirical justification. Extends the NVFP4 tensor-core GEMV dispatch policy from a two-arm decision tree (wide-column-if-any-wave, else column-starved) to a three-arm tree that adds a middle "moderate column-tile" branch, plus adds shape-scoped benchmark overrides and boundary tests. Author is a trusted CUDA EP contributor with a clean numeric argument in the doc.

Bug diagnosis — the pre-PR dispatch collapse

Walking the pre-PR policy in matmul_block_scaled_fp4.cu:

if ((col_tiles + 3) / 4 >= sm_count) {           // 1 wave of wide-column blocks
  return {std::min(2, windows), 4};              // wide-column: ColTiles=4, small KSplit
}
// Column-starved: max-out KSplit up to 16-warp block
int k_split = 1;
while (k_split < 16 && k_split * 2 <= windows) k_split <<= 1;
return {k_split, 1};

For Qwen MTP N=17408, K=5120 on H200 (132 SMs): col_tiles = 17408/16 = 1088, wide_col_blocks = 272 ≥ 132 → picks {2, 4}. But 272 wide-column blocks over 132 SMs is only ~2 waves — the wave-quantization loss on the last partial wave dominates, and the smaller-tile config {2, 1} (which gives 4× the block count) would let the GPU hide the tail better. The pre-PR code took the wide-column path too eagerly because "any wave" was the threshold.

Similarly for N=17408, K=8192: wide-column collapses to {2, 4} when the long K reduction (64 windows) would prefer {8, 1} — giving the K reduction more parallelism and keeping every block busy.

The fix — three-arm policy

const int wide_col_blocks = (col_tiles + 3) / 4;
if (wide_col_blocks >= kTargetGridWaves * sm_count &&           // 4 waves for wide-col
    (windows < kLongReductionWindows || wide_col_blocks >= kLongReductionGridWaves * sm_count)) {  // 8 waves for long-K wide-col
  return {std::min(2, windows), 4};
}
if (col_tiles >= kTargetGridWaves * sm_count) {                 // 4 waves for single-col
  const int k_split = windows >= kLongReductionWindows ? 8 : std::min(2, windows);
  return {k_split, 1};
}
// Column-starved fallback, unchanged

Constants:

• kTargetGridWaves = 4 — the wave count required for the wide- and moderate-column paths to pay off.
• kLongReductionGridWaves = 8 — the higher wave count required for wide-column under long-K reductions (because the wider tile costs even more when reductions are long, so we need even more parallelism to justify it).
• kLongReductionWindows = 64 (K = 8192) — the long-reduction threshold. When true, single-column paths use KSplit = 8 instead of min(2, windows).

Verifying the algorithm hits the intended Qwen shapes:

Shape	col_tiles	wide_col_blocks	Old	New
N=17408, K=5120 (Qwen gate/up)	1088	272	{2, 4} (272 ≥ 132)	{2, 1} (272 < 528 and 1088 ≥ 528)
N=17408, K=8192 (Qwen long)	1088	272	{2, 4} (272 ≥ 132)	{8, 1} (272 < 528 and 1088 ≥ 528, long-K)
N=248320, K=2048 (lm_head)	15520	3880	{2, 4} (3880 ≥ 132)	{2, 4} (3880 ≥ 528 and windows=16 < 64) — unchanged ✓
N=512, K=2048 (small gate/up)	32	8	column-starved → {16, 1}	column-starved → {16, 1} — unchanged ✓
N=2048, K=512 (down)	128	32	column-starved → {4, 1}	column-starved → {4, 1} — unchanged ✓

Two rows change (the two Qwen shapes), three rows don't. Measured speedups: 50.7 → 46.6 us (1.09×) and 61.4 → 55.0 us (1.12×). Regression-free on the other four documented shapes. ✓

The wave constants are load-bearing but empirical

kTargetGridWaves = 4, kLongReductionGridWaves = 8, kLongReductionWindows = 64 — these are the tuning knobs. They're stated as constexpr int at function scope with no comment explaining the units or the tuning procedure. The doc §"launcher targets four waves of SMs..." explains the reasoning at the higher level, but the constants themselves are unlabeled magic numbers in code. Suggest adding a one-line comment above them:

// Wave-count thresholds tuned on H100/H200 (SM 90, 132 SMs); H200 measurements in the doc.
constexpr int kTargetGridWaves = 4;
constexpr int kLongReductionGridWaves = 8;
constexpr int kLongReductionWindows = 64;  // K = 8192, the threshold where doubling K-split beats doubling column tiles.

Ideally these would be per-architecture (SM 80 vs 90 vs 100 may have different sweet spots), but the current dispatch only enables the tensor-core path on device_prop.major >= 8 and validation focused on H200, so parameterizing later is fine.

API split — public PickFp4MmaConfig vs internal PickFp4MmaConfigImpl

The refactor extracts the tiling logic to a new header matmul_block_scaled_fp4_tiling.h:

struct Fp4MmaConfig {
  int k_split;
  int col_tiles;
};

Fp4MmaConfig PickFp4MmaConfig(int n, int k, int sm_count);

Clean, no CUDA includes, and importantly lets the host-side test matmul_block_scaled_fp4_test.cc call the tiling algorithm without compiling CUDA. ✓

The launcher calls the internal PickFp4MmaConfigImpl directly:

const Fp4MmaConfig cfg = PickFp4MmaConfigImpl(n, k, device_prop.multiProcessorCount);

Rationale is presumably to avoid the CUDA-version-check overhead in the outer wrapper, but the wrapper's overhead is a single #if — the compiler folds it away. The Impl suffix is a bit ad-hoc. Minor style suggestion: put the implementation into a detail namespace inside the anonymous namespace instead of using an Impl suffix, and have the launcher call detail::PickFp4MmaConfig symmetrically with the public one. Non-blocking.

Environment-variable overrides — well-scoped, but ORT_ENFORCE runs per-call

static int const k_split = ParseEnvironmentVariableWithDefault<int>("ORT_FP4_GEMV_KSPLIT", 0);
...
ORT_ENFORCE(k_split == 0 || k_split == 1 || k_split == 2 || k_split == 4 ||
            k_split == 8 || k_split == 16,
            "ORT_FP4_GEMV_KSPLIT must be 0, 1, 2, 4, 8, or 16.");

The env vars are read once (static const), but ORT_ENFORCE runs on every ApplyFp4MmaConfigOverrides call. Cheap comparisons, and the per-inference dispatch cost is already dominated by other work, but the idiomatic pattern is to move validation into a lambda IIFE that runs during static initialization:

static const int k_split = []() {
  int v = ParseEnvironmentVariableWithDefault<int>("ORT_FP4_GEMV_KSPLIT", 0);
  ORT_ENFORCE(v == 0 || v == 1 || v == 2 || v == 4 || v == 8 || v == 16,
              "ORT_FP4_GEMV_KSPLIT must be 0, 1, 2, 4, 8, or 16.");
  return v;
}();

That way validation happens exactly once, on first entry. Micro-nit, non-blocking.

Kernel-shape constraint enforcement — good defensive check

ORT_ENFORCE(config.col_tiles == 1 || config.k_split <= 2,
            "ORT_FP4_GEMV_COL_TILES=4 supports only KSplit 1 or 2.");

Placed after applying overrides so that a ORT_FP4_GEMV_KSPLIT=8 ORT_FP4_GEMV_COL_TILES=4 combination fails clearly instead of running an unsupported kernel shape. Correct and well-placed. ✓

ORT_FP4_GEMV_MATCH_N / ORT_FP4_GEMV_MATCH_K shape-scoping — nice detail

The MATCH_N/MATCH_K env vars let benchmark overrides apply only to a specific shape:

if ((match_n != 0 && n != match_n) || (match_k != 0 && k != match_k)) {
  return config;  // shape doesn't match; use default
}

This is genuinely useful for benchmarking campaigns where you want to override one shape without affecting others in the same model run. Rarely seen in kernel dispatch code; good addition.

Test coverage — hits the three policy branches

The new GemvTensorCoreQwenTilingBoundaries test validates four cases:

{64 * sm_count, 5120, 2, 1},       // moderate-col, ordinary K
{64 * sm_count, 8192, 8, 1},       // moderate-col, long K
{512 * sm_count, 512, 2, 4},       // wide-col, 8 waves
{64 * sm_count - 16, 2048, 16, 1}, // column-starved

I re-walked each against the new algorithm:

• Case 1: col_tiles = 4·sm, wide_col_blocks = sm (1 wave). Fails wide-col (sm < 4·sm). Passes moderate-col (4·sm ≥ 4·sm). K=5120 → windows=40, ordinary. Returns {min(2,40), 1} = {2, 1}. ✓
• Case 2: Same shape as Case 1 but K=8192 → windows=64, long. Moderate-col long branch: k_split = 8. Returns {8, 1}. ✓
• Case 3: col_tiles = 32·sm, wide_col_blocks = 8·sm (8 waves). Passes wide-col (8·sm ≥ 4·sm and windows=4 < 64 short OR 8·sm ≥ 8·sm for long — actually 4 < 64 so short-K branch is entered). Returns {min(2,4), 4} = {2, 4}. ✓
• Case 4: col_tiles = 4·sm - 1, wide_col_blocks = sm (still 1 wave). Fails wide-col. Fails moderate-col (4·sm - 1 < 4·sm). Column-starved. windows=16 → k_split doubles to 16. Returns {16, 1}. ✓

Boundary gap: no test at exactly wide_col_blocks == 4·sm under long-K, which is the case that hits the second condition (wide_col_blocks < 8·sm → falls through wide-col even though it passes the first condition). E.g., n = 16·(4·sm) = 8448, k = 8192. Not blocking (Case 2 exercises the moderate-col long-K path), but adding one more entry:

{16 * kTargetGridWaves * sm_count, 8192, 8, 1},  // exactly at 4-wave wide-col threshold, long K -> falls through to moderate

would sharpen the wide-col-long-K interaction. Optional nit.

Documentation update — clean policy table

The added table in matmul_block_scaled_fp4.md is exactly the right level of abstraction:

Grid condition	Reduction	Configuration
4-col grid covers 4 waves	ordinary or wide enough long	ColTiles=4, KSplit=min(2, K/128)
4-col grid < 4 waves, 1-col grid ≥ 4 waves	K/128 < 64	ColTiles=1, KSplit=min(2, K/128)
4-col grid < 8 waves, 1-col grid ≥ 4 waves	K/128 ≥ 64	ColTiles=1, KSplit=8
1-col grid < 4 waves	any	ColTiles=1, up to 16 K-split warps

Correctly captures the dispatch behaviour and lets a future maintainer look up "why did this shape get this config?" without reading the C++. Docs also cover the four new env vars. ✓

One minor doc claim to verify

"the launcher targets four waves of SMs before using one-column blocks"

This describes the wide → moderate transition, but "before using one-column blocks" reads slightly ambiguously — it could be parsed as "prefers wide-col until wave-count drops below four", which is what it means. Minor wordsmithing suggestion:

"the launcher requires at least four waves of SMs before selecting the wide-column tile, and eight waves when the reduction is long (K/128 ≥ 64)"

Non-blocking.

Test file wiring

The tiling header is included inside a #if CUDA_VERSION >= 12080 block in the test file, matching the runtime guard on the tensor-core path. That means GemvTensorCoreQwenTilingBoundaries is CUDA-12.8-only. Correct — on older CUDA the tensor-core dispatch is inactive and testing the selector would compare against a stub {1,1} return. ✓

CI status

87 checks total on tip 86435da. Latest snapshot showed 73/87 OK — down from 84/87 on the intermediate commit 66cb98e. That's a 14-check drop, which is bigger than the usual bot-scaffolding noise pattern (~3 non-green). Worth investigating: could be flaky Windows GPU CUDA legs from a rebase, or the last commit test(cuda): verify FP4 GEMV tiling policy may have introduced a header-hygiene issue in the test file. Author note says "The local build configuration had provider unit tests disabled, so the new gtest was not executed locally" — so the CI is the first place it runs. Please verify the failing checks before merge and confirm they're the tiling test itself (in which case, fix), the include-order lint (which is mechanical), or unrelated infra flakes.

The CI trend across commits is:

• e622e00 (tiling logic): 80/87
• 66cb98e (Windows CUDA runtime fix): 84/87
• 86435da (test): 73/87

The test commit dropped the count more than the tiling-logic commit did. Suggests the test itself may have added CI cost or has an issue.

Stacking note

Author notes this PR is a follow-up to #32128 and "can be rebased/stacked on #32128 after that PR merges." I don't see #32128 as merged yet — worth checking the base branch situation before merge. If #32128 is still open, this needs to rebase and its +131/-22 count could shift as the vectorized dequant paths in #32128 might overlap.

Summary of items

Blocking:

1. Verify the 73/87 CI slip on 86435da isn't a real test failure in GemvTensorCoreQwenTilingBoundaries. Author's Validation section explicitly notes local test runner was disabled, so CI is the first exercise.
2. Confirm [CUDA] Vectorize the NVFP4 weight dequantization for prefill #32128 stacking situation before merge.

Non-blocking but suggested:

3. Add a one-line comment explaining that the wave constants kTargetGridWaves = 4, kLongReductionGridWaves = 8, kLongReductionWindows = 64 are empirically tuned on H100/H200 (not fundamental invariants).
4. Move env-var validation into a lambda IIFE so ORT_ENFORCE runs during static init instead of per-call.
5. Add one boundary test at the wide_col_blocks == 4·sm long-K interaction edge.
6. Consider renaming PickFp4MmaConfigImpl → detail::PickFp4MmaConfig for cleaner symmetry with the public wrapper. Style-only.
7. Doc wording nit on "four waves of SMs before using one-column blocks".

Recommendation

Approve — pending confirmation of the CI slip and the #32128 stacking situation. The tuning is well-diagnosed, the algorithm change is precisely bounded (only two shape classes shift; four documented shapes are unchanged), the new selector API extraction enables clean host-side testing, and the doc + env-var overrides make future tuning campaigns straightforward. The magic constants deserve one line of comment but are legitimate empirical tuning knobs, not accidental. Good example of how to add a shape-dispatch performance fix without regressing sibling shapes.