perf(cuda): prefetch-pipeline the block-32 scales-fp16 split-K int4 GEMV (+4.4% q38 decode)

[justinchuby]

What

Adds prefetch-pipelined _pf siblings of the block-32 scales-fp16 split-K int4 decode GEMV (both the asymmetric zero-point entry matmul_nbits_gemv_f16_scales_f16_zp_splitk — q38's dominant decode kernel — and the symmetric entry) and routes to them by default. ONNX_GENAI_ZP_SPLITK_PREFETCH=0 restores the plain split-K entry for A/B.

Why

The int4 M=1 GEMV (MatMulNBits) is the confirmed #1 q38 decode cost (~25% of captured decode). The split-K kernel already fills idle SMs and adds per-column memory-level parallelism (K_SPLIT warps/column), but each lane still walks its weight-word chain with a single 32-bit global load in flight — a load -> accumulate -> load dependency that stalls on the Long-Scoreboard weight-load latency the single-warp _pipe kernel already diagnosed (ncu on the sibling path: ~75% of warp cycles stalled on the one in-flight load; DRAM ~24% of an H200's 4.8 TB/s; the dequant math is not the bottleneck).

The _pf variant reuses the proven _pipe transform on the split-K loop: each lane keeps PF=3 weight-word loads in flight via a register-resident shift register (manual rotation, no dynamic-indexed array -> no local spill), adding per-lane MLP on top of split-K's per-column MLP.

Bit-identity

The lane->nibble mapping, accumulate_int4x8_f16_zp calls, per-split accumulation order and the K_SPLIT shared-memory reduction are unchanged — only the timing of the weight loads changes. Output is bit-identical to the plain split-K entry:

• New unit test scales_f16_zp_splitk_prefetch_is_bit_identical_to_splitk (asserts 0/N fp16 bit mismatches on grid-starved split-K shapes; also asserts the dims actually select split-K so the check isn't vacuous). All 41 matmul_nbits tests pass.
• End-to-end greedy token IDs are identical ON vs OFF on both models.

Capture-safety

Registers + launch-time shared memory only; no alloc/sync. Verified captures>0, fallbacks=0 on q38 and mary.

Generality

The swap keys off the already-shape-gated split-K entry (use_scales_f16_zp_splitk / use_f16_symmetric_splitk), which derive from k/n/block_size/live SM count. No hardcoded head counts, dims, or block sizes.

A/B (mandatory env block)

• GPU: NVIDIA H200, ordinal 2, idle (GPU pinned via CUDA_VISIBLE_DEVICES)
• Batch: 1   Decode window: 128-tok steady, 1 warmup, decode_skip=8   Stat: median of 3 steady_median runs (each itself median-of-3)
• Build: --features bench-native,cuda, ORT_ROOT=.ort-cuda-1.28/root
• Base: origin/main 70db81192 (merged up to b77e3d3fc); branch: squad/int4-gemv-m1

Model	prefetch OFF	prefetch ON	Δ	capture
qwen38-27b-int4-cuda (asymmetric int4 block-32, bf16 io) — target	59.77 tok/s	62.42 tok/s	+4.43%	captures=4, fallbacks=0
qwen3.6-27b-int4-cuda (symmetric f16) — control	coherent, byte-identical tokens ON vs OFF			captures=4, fallbacks=0

(Pre-merge A/B on the same box measured +4.85%; both are well above the +2% gate.)

Profiling disclosure

ncu is blocked on this box (RmProfilingAdminOnly=1 for non-admin), so attribution uses op-count + roofline + captured A/B timing rather than per-kernel counters. The bottleneck diagnosis (Long-Scoreboard weight-load latency) is quoted from the pre-existing _pipe kernel's committed ncu analysis on the sibling single-warp path.

Files

• crates/onnx-runtime-ep-cuda/src/kernels/matmul_nbits.rs — new _pf device template + two extern entries, use_scales_f16_zp_splitk_pf gate, dispatch swap, bf16_direct_capable update, run_symmetric_block_raw prefetch toggle, new parity test.

Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.70%. Comparing base (0faa635) to head (b0c28d1).
⚠️ Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1584      +/-   ##
==========================================
- Coverage   80.77%   80.70%   -0.07%     
==========================================
  Files         381      379       -2     
  Lines      173645   170738    -2907     
  Branches   173645   170738    -2907     
==========================================
- Hits       140262   137799    -2463     
+ Misses      28492    28060     -432     
+ Partials     4891     4879      -12     

Flag	Coverage Δ
cli-ort-linux	82.60% <ø> (ø)
cli-ort-windows	82.10% <ø> (ø)
mlas	?
offline	80.64% <ø> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 6 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Rebased onto #1585 (deep-split GQA k/v) + re-measured

origin/main advanced to 0faa63547 (#1585), which refactored the split-K template to <bool HasZp, int K_SPLIT> and added a K_SPLIT=8 deep-split entry (matmul_nbits_gemv_f16_scales_f16_zp_splitk8) for the grid-starved GQA k/v projection. 3-way merged (no cherry-pick); head is now b0c28d135.

Reconciliation (the deep-split path also wants prefetch):

• Resolved the bf16_direct_capable conflict as a union.
• Parametrized my prefetch template on K_SPLIT too (matmul_nbits_gemv_f16_scales_f16_splitk_pf_tpl<bool HasZp, int K_SPLIT>) and added a matmul_nbits_gemv_f16_scales_f16_zp_splitk8_pf entry (<true, 8>). The deep-split path is still single-32-bit-load-in-flight per lane after the grid is deepened, so the PF register-shift-register is an additional latency-hiding win there, not a duplicate of perf(cuda): fix the GEMV probe that was measuring its own dispatch, then deepen split-K for GQA k/v #1585's grid fix.
• Extended the dispatch swap + bf16_direct_capable to route _splitk8 → _splitk8_pf under the same ONNX_GENAI_ZP_SPLITK_PREFETCH gate. Now the whole K_SPLIT∈{2,8} split-K family gets prefetch.

Re-measured A/B on the post-#1585 base (so #1585's deep-split is in both arms; this isolates the prefetch):

	GPU	prefetch OFF	prefetch ON	Δ	capture
q38 (asymmetric int4 block-32, bf16)	H200 ord 2, idle	59.76 tok/s	62.49 tok/s	+4.57%	captures=4, fallbacks=0
mary (symmetric f16) control	H200 ord 2, idle	coherent, byte-identical tokens ON vs OFF			captures=4, fallbacks=0

Env: batch=1, 128-tok steady, 1 warmup, median-of-4 steady_median runs; features bench-native,cuda, ORT_ROOT=.ort-cuda-1.28/root; base 0faa63547 vs branch head b0c28d135. q38 greedy token IDs identical ON vs OFF (byte-identical, stronger than the relaxed bar). ncu blocked (RmProfilingAdminOnly=1) → op-count + roofline + captured A/B.

Tests: all 43 matmul_nbits tests pass (incl. the prefetch parity test and #1585's zp_split_factor_deepens...). Still ≥2% → merge-ready.