Parallelize Kn int4 dequant in prefill load path + profiling (#959)

[justinchuby]

Summary

Native CPU-only (mlas-off) load path for int4 no-zero-point models spent its first-prefill time in a single-threaded int4→f32 dequantization (the Kn/[k,n] transpose consumed by the dense prefill GEMM), while the sibling Nk (decode) path was already rayon-parallelized. On this box that serial phase was ~⅔ of load wall and pinned the whole load to ~1.27 cores busy — the bulk of the native-vs-ORT load-time gap in #959.

This PR:

1. Instrumentation (ONNX_GENAI_PROFILE_MM=1): env-gated per-phase timing of the one-time weight dequant/prepack (dequant-kn, dequant-nk, mlas-*). Zero-cost / byte-identical when unset.
2. Fix — parallelize dequant-kn: each Kn depth-row [depth*n .. (depth+1)*n) is contiguous, so it parallelizes with par_chunks_mut(n) over depth-rows. Byte-identical to the serial scatter (same dequantize_nbits_value, same indices; grouped path stays serial). No unsafe.

Measured (qwen0.5b-q4, CPU-only build, warm cache, best-of-3, --max-new-tokens 1 = load only)

	wall	cores	dequant-kn
before	7.6 s	1.27	~5.06 s
after	~4.8 s	~3.9	~2.1 s

Approaching ORT's 3.5 s / ~4.5 cores — the load gap was threading, not a costlier algorithm. Greedy output unchanged ("The capital of France is Paris.").

Scope / what this does NOT fix

Load only. The 14B ~102 s/token decode cliff is a separate symptom of the same f32-materialization root cause: the non-mlas decode caches each weight as f32 (weight_nk), an 8× expansion (6.995 GB int4 → ~56 GB f32 > this 63.8 GB box → paging). That needs a footprint fix (blockwise int4 kernels / mlas), not threading — noted on #959, not in this PR.

Gates

• cargo fmt -p onnx-genai-engine -p onnx-runtime-ep-cpu — ok
• cargo clippy -p onnx-runtime-ep-cpu -- -D warnings — ok
• cargo test -p onnx-runtime-ep-cpu — 1065 passed; 0 failed
• cargo test -p onnx-genai-engine --lib — 382 passed; 0 failed

Refs #959.

[codecov]

Codecov Report

❌ Patch coverage is 79.10448% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.99%. Comparing base (a66f0dc) to head (ec4ee6a).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
...es/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs	79.10%	12 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #970      +/-   ##
==========================================
- Coverage   79.00%   78.99%   -0.02%     
==========================================
  Files         356      356              
  Lines      150158   150211      +53     
  Branches   150158   150211      +53     
==========================================
+ Hits       118639   118661      +22     
- Misses      26868    26897      +29     
- Partials     4651     4653       +2     

Flag	Coverage Δ
cli-ort-linux	83.52% <ø> (ø)
cli-ort-windows	83.11% <ø> (ø)
mlas	80.96% <ø> (+0.04%)	⬆️
offline	78.79% <79.10%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
...es/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs	74.28% <79.10%> (-0.26%)	⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	39.69 µs	80.32 µs	+102.4%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	520.93 µs	1.01 ms	+94.5%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	407.89 µs	718.84 µs	+76.2%
🔴	matmul/small_generic_f32_threads=8/1x256x256	34.31 µs	60.19 µs	+75.4%
🔴	matmul/medium_generic_f16_threads=8/32x512x512	30.23 µs	52.33 µs	+73.1%
🔴	gather/large_bf16_threads=1-internal/131072	11.80 µs	20.21 µs	+71.2%
🔴	matmul/small_generic_f16_threads=8/1x256x256	30.07 µs	47.41 µs	+57.6%
🔴	matmul/medium_generic_f32_threads=8/32x512x512	1.06 ms	1.62 ms	+52.1%
🔴	sampling_latency/top_p_per_token	378.69 µs	535.10 µs	+41.3%
🔴	matmul/small_generic_f16_threads=1/1x256x256	31.20 µs	43.61 µs	+39.8%
🔴	gather/large_f32_threads=1-internal/131072	31.99 µs	43.55 µs	+36.1%
🔴	matmul/large_generic_f16_threads=8/32x1024x1024	87.80 µs	117.97 µs	+34.4%
🔴	matmul/small_generic_bf16_threads=1/1x256x256	31.70 µs	42.28 µs	+33.4%
🔴	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	78.48 µs	104.10 µs	+32.6%
🔴	matmul/medium_generic_bf16_threads=1/32x512x512	527.63 µs	696.86 µs	+32.1%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	378.47 µs	498.64 µs	+31.8%
🔴	matmul/small_generic_bf16_threads=8/1x256x256	31.99 µs	41.63 µs	+30.2%
⚠️	sampling_latency/top_k_per_token	52.29 µs	67.42 µs	+28.9%
⚠️	grammar_masking/llguidance_compute_mask/32	76.67 µs	98.35 µs	+28.3%
⚠️	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.67 ms	7.26 ms	+28.0%
⚠️	matmul/small_generic_f32_threads=1/1x256x256	37.59 µs	47.46 µs	+26.2%
⚠️	tokenization/encode_tokens_per_second	374.18 µs	457.08 µs	+22.2%
⚠️	matmul/large_generic_f32_threads=1/32x1024x1024	9.37 ms	11.42 ms	+21.9%
⚠️	sampling_latency/min_p_per_token	207.75 µs	253.05 µs	+21.8%
⚠️	kv_cache/alloc_dealloc_pages	38.74 µs	47.17 µs	+21.7%
⚠️	qwen3_sampling_processors/top_p_fast_after_top_k	517.57 µs	625.04 µs	+20.8%
⚠️	gather/large_f16_threads=1-internal/131072	13.06 µs	15.47 µs	+18.4%
⚠️	matmul/medium_generic_f16_threads=1/32x512x512	30.93 µs	36.47 µs	+17.9%
⚠️	gather/medium_f32_threads=1-internal/32768	4.51 µs	5.32 µs	+17.9%
⚠️	logit_processing/seven_processor_chain_per_step	324.51 µs	381.33 µs	+17.5%
⚠️	qwen3_sampling_processors/top_k_top_p_fast	656.84 µs	763.36 µs	+16.2%
⚠️	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	51.95 µs	59.81 µs	+15.1%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.11 ms	2.39 ms	+13.5%
✅	qwen3_sampling_processors/top_k_partial_selection	138.60 µs	156.52 µs	+12.9%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.97 ms	2.21 ms	+12.2%
✅	sampling_latency/greedy_per_token	3.23 µs	3.62 µs	+12.1%
✅	tokenization/decode_tokens_per_second	6.21 ms	6.94 ms	+11.8%
✅	gather/medium_bf16_threads=1-internal/32768	2.48 µs	2.77 µs	+11.5%
✅	gather/medium_f16_threads=1-internal/32768	2.69 µs	2.98 µs	+10.8%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	1.34 ms	1.48 ms	+10.7%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	88.67 µs	97.61 µs	+10.1%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.32 ms	2.45 ms	+5.9%
✅	gather/small_f32_threads=1-internal/4096	769.7 ns	813.5 ns	+5.7%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.60 ms	3.79 ms	+5.2%
✅	gather/small_f16_threads=1-internal/4096	549.4 ns	535.0 ns	-2.6%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	5.45 ms	5.26 ms	-3.4%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.54 µs	14.98 µs	-3.6%
✅	reduce_mean/medium_f32_threads=1-internal/65536	261.20 µs	243.08 µs	-6.9%
✅	gather/small_bf16_threads=1-internal/4096	640.3 ns	589.1 ns	-8.0%
✅	add/large_f32_threads=1-internal/4194304	45.96 ms	40.32 ms	-12.3%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.16 ms	1.02 ms	-12.4%
✅	add/medium_bf16_threads=1-internal/262144	3.08 ms	2.67 ms	-13.3%
🟢	add/large_f16_threads=1-internal/4194304	49.24 ms	40.88 ms	-17.0%
🟢	add/medium_f16_threads=1-internal/262144	3.50 ms	2.62 ms	-25.2%
🟢	add/small_f16_threads=1-internal/1024	17.40 µs	12.96 µs	-25.5%
🟢	add/large_bf16_threads=1-internal/4194304	56.63 ms	41.85 ms	-26.1%
🟢	add/medium_f32_threads=1-internal/262144	3.49 ms	2.55 ms	-27.1%
🟢	add/small_bf16_threads=1-internal/1024	17.59 µs	12.82 µs	-27.1%
🟢	add/small_f32_threads=1-internal/1024	280.6 ns	196.7 ns	-29.9%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 5.46 4.88 6.19 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)