perf(cpu-ep): packed-panel 6x16 microkernel for gemm_bt MoE prefill

[justinchuby]

What

A packed-panel 6x16 microkernel for gemm_bt, the kernel that computes
C = A·Btᵀ directly on the [out][in]-major MoE expert weights (#1241). It is
entered only from the single-threaded driver on prefill-shaped row extents;
decode, small-token MoE and every multi-threaded shape keep the existing
tile_4x2 path unchanged.

Why

MoE expert GEMMs were the largest remaining single-thread loss on the
production-native matrix — §31.5/§31.8 of the CPU-EP ledger put MoE prefill at
1.22–1.56x ORT at t=512 and attributed it to "microkernel efficiency
against MLAS's packed panels and wider register tiles". That attribution is
right, and it is fixable without handing anything to ORT's CPU EP.

tile_4x2 keeps k-partials in the SIMD lanes. Two consequences: every
output cell ends in a horizontal reduction, and the b side of the inner loop
is strided by k — six loads per eight FMAs.

How

	tile_4x2 (unchanged, still the default)	micro_6x16 (new)
lanes hold	k-partials	c columns
horizontal reductions	one per output cell	none
per k step	6 loads / 8 FMAs = 1.33	2 contiguous loads + 6 broadcasts / 12 FMAs = 1.5
b access	strided by k	two sequential cache lines

• pack_panel transposes one column panel of bt into micro-panel-major order
  (packed[s*k*PNR + p*PNR + t]). This is the only transpose in the path, it is
  bounded by the panel rather than the whole expert bank, and every row band in
  the sweep reuses it.
• block_packed drives column panel → row band → micro-panel. Both remainders
  (rows % 6, n % 16) fall through to the existing full-k dot, which is
  the same code the unpacked path already uses for its edges.
• Single-threaded gemm_bt hands the whole row extent to one block when the
  packed kernel is selected, so the pack is amortised once rather than once per
  64-row MAX_MC block.

The gate is structural, not tuned

PACK_MIN_ROWS = PMR * PACK_MIN_BANDS = 72, deliberately above MAX_MC (64),
pinned by a const assertion:

const _: () = assert!(
    PACK_MIN_ROWS > super::MAX_MC,
    "the packed kernel must stay out of the multi-threaded driver's row \
     blocks, which are capped at MAX_MC and would each re-pack the panel"
);

The multi-threaded driver slices m into 2*threads blocks capped at MAX_MC,
and each block re-packs the same panel. Keeping the bar above that cap makes
"the panel is packed once per sweep" a property of the code rather than of the
shape, and means multi-threaded behaviour is byte-identical to main for
every shape. The assertion stops a future MAX_MC bump from silently undoing
that.

Why I will not make a multi-threaded claim

I tried to. It does not survive its own control. A 2-thread A/B of two binaries
that I traced to be on the identical code path (every row block rows=56..64,
packed=false) still came out ~40% apart on the median. Multi-thread
measurement on this shared 32-core box is not usable at this effect size, in
either direction, so this PR takes the position that costs nothing: MT is left
exactly as it was. Sharing one pack across threads (or making
gemm_bt_col_parallel the packed path, since each of its tasks already owns all
m rows) is the natural follow-up, and it needs a quiet machine.

Evidence

Both binaries built --no-default-features --features bench-native,cuda-13000,
differing only in matmul.rs. The arm is proved by the linker, not the flag:
nm -C <bin> | grep -ci mlas = 0 on both. Arms interleaved innermost in one
driver process, arm order alternating per rep, thread-matched
--native-threads 1 --ort-intra-threads 1. Metric is native_min / ort_min
(lower is better; <1.000 beats ORT).

Single thread, t=512 — the shapes that reach the new kernel

Final gate, shipped binary (6 reps):

model	before	after	Δ min	Δ median	parity
moe_phi35moe_h2048_i6400_e4_t512	1.356	1.009	−25.61%	−16.50%	PASS
moe_qwen3moe_h2048_i768_e16_t512	1.297	1.101	−15.12%	−15.92%	PASS
moe_mixtral_h1024_i3584_e8_t512	1.226	1.173	−4.27%	−2.85%	PASS

Four independent passes, Δ on the min ratio:

model	pass 1	pass 2	pass 3	pass 4 (shipped)
qwen3moe t512	−16.7%	−23.4%	−19.6%	−15.1%
phi35moe t512	−16.9%	−24.9%	−27.6%	−25.6%
mixtral t512	−1.3%	−6.6%	−3.0%	−4.3%

Phi-3.5-MoE prefill crosses 1.000 in the shipped pass and Qwen3-MoE lands at
1.10, from a 1.30–1.54 loss. Mixtral moves least, and I believe that is real
rather than measurement: its expert bank is ~352 MB of fp32, so that shape is
closer to DRAM-bound than microkernel-bound and a better inner loop has less to
win.

t=1 / t=32 — unchanged by construction, and verified

Traced with a temporary instrumented build: every gemm_bt_block at t=32 has
rows in 14–19, far below the gate, so packed=false and the executed code
is the unpacked path. Their run-to-run spread (±10% on min, ±3% on median, sign
flipping between passes) is therefore a direct read of this box's noise
floor. At t=512 the traced extents are 74–365 rows, all packed.

Correctness

• cargo test -p onnx-runtime-ep-cpu --lib → 1434 passed, 0 failed, 17 ignored.
• New test gemm_bt_packed_kernel_matches_naive_single_threaded: 10 shapes
  against naive_bt on a 1-thread pool, covering both remainders, k tails, the
  shortest admitted k, and multi-column-panel sweeps. It carries a
  non-vacuity assertion — if a policy change stops these shapes from reaching
  the packed kernel, the test fails loudly rather than silently measuring the old
  path.
• Falsified twice: swapping the two stored halves of the micro-tile, and
  corrupting the pack layout to dst[t*k+p]. Each mutation fails only the new
  test; the other five gemm_bt tests stay green.
• Clippy clean, cargo fmt --check clean, cargo check --features mlas compiles.
• Parity PASS in every benchmark cell above.

Review

Reviewed by Opus (claude-opus-4.8), verdict APPROVE-WITH-NITS, no blocking
findings. It independently re-ran the packed test with 7 additional adversarial
shapes (the nc-floor branch at k>8192, single-micro-panel wide sweeps,
whole-m widening at m=1000, maximal row/column remainders) under
AddressSanitizer — all clean — and verified the panel-index bounds, the
jend-j0 ≡ 0 (mod PNR) contract, remainder exhaustiveness/disjointness, and that
no consumer of gemm_bt depends on its accumulation order. Its one substantive
nit — that the original rows >= 48 gate did not keep large-m multi-threaded
blocks out of the packed path, contrary to the comment — is what produced the
structural > MAX_MC gate and the const assertion above.

Scope

• No ORT CPU fallback, no deferral — assigned == executed is unchanged.
• Multi-threaded blocking policy untouched; MT is on the identical code path.
• tile_4x2, block, dot, gemm_bt_block_scalar and the mlas arm untouched.

[github-actions]

✅ Benchmarks — No Regression

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
✅	matmul/large_generic_f32_threads=1/32x1024x1024	8.61 ms	9.43 ms	+9.6%
✅	gather/large_f32_threads=1-internal/131072	23.15 µs	25.29 µs	+9.3%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.33 ms	5.73 ms	+7.4%
✅	gather/small_f32_threads=1-internal/4096	625.4 ns	656.3 ns	+4.9%
✅	matmul/small_generic_f16_threads=8/1x256x256	26.92 µs	28.21 µs	+4.8%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	360.49 µs	377.03 µs	+4.6%
✅	matmul/small_generic_f16_threads=1/1x256x256	28.14 µs	29.27 µs	+4.0%
✅	gather/large_f16_threads=1-internal/131072	11.83 µs	12.22 µs	+3.2%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	494.22 µs	506.54 µs	+2.5%
✅	gather/large_bf16_threads=1-internal/131072	10.39 µs	10.65 µs	+2.5%
✅	sampling_latency/top_p_per_token	350.49 µs	358.77 µs	+2.4%
✅	kv_cache/alloc_dealloc_pages	36.43 µs	36.95 µs	+1.4%
✅	matmul/small_generic_bf16_threads=1/1x256x256	28.10 µs	28.49 µs	+1.4%
✅	tokenization/decode_tokens_per_second	5.73 ms	5.80 ms	+1.2%
✅	matmul/medium_generic_f16_threads=8/32x512x512	27.80 µs	28.04 µs	+0.8%
✅	add/small_bf16_threads=1-internal/1024	409.7 ns	413.2 ns	+0.8%
✅	sampling_latency/min_p_per_token	191.81 µs	193.17 µs	+0.7%
✅	gather/medium_f16_threads=1-internal/32768	2.23 µs	2.24 µs	+0.4%
✅	matmul/small_generic_f32_threads=1/1x256x256	34.15 µs	34.28 µs	+0.4%
✅	gather/small_bf16_threads=1-internal/4096	441.0 ns	441.5 ns	+0.1%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	487.93 µs	488.09 µs	+0.0%
✅	reduce_mean/large_f32_threads=1-internal/262144	911.09 µs	911.13 µs	+0.0%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	3.68 ms	3.67 ms	-0.1%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.15 ms	2.15 ms	-0.3%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	73.45 µs	73.24 µs	-0.3%
✅	gather/medium_bf16_threads=1-internal/32768	2.21 µs	2.20 µs	-0.5%
✅	sampling_latency/greedy_per_token	3.01 µs	2.99 µs	-0.5%
✅	grammar_masking/llguidance_compute_mask/32	69.47 µs	69.04 µs	-0.6%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	1.98 ms	1.97 ms	-0.7%
✅	logit_processing/seven_processor_chain_per_step	298.52 µs	296.37 µs	-0.7%
✅	add/medium_bf16_threads=1-internal/262144	97.77 µs	96.73 µs	-1.1%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	79.16 µs	78.20 µs	-1.2%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.31 ms	3.27 ms	-1.2%
✅	add/large_f16_threads=1-internal/4194304	1.56 ms	1.54 ms	-1.3%
✅	qwen3_sampling_processors/top_k_partial_selection	135.40 µs	133.02 µs	-1.8%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.08 µs	13.82 µs	-1.8%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.91 ms	1.88 ms	-1.9%
✅	reduce_mean/medium_f32_threads=1-internal/65536	230.99 µs	226.03 µs	-2.1%
✅	gather/small_f16_threads=1-internal/4096	448.6 ns	438.1 ns	-2.4%
✅	sampling_latency/top_k_per_token	53.75 µs	52.40 µs	-2.5%
✅	add/large_f32_threads=1-internal/4194304	603.63 µs	588.14 µs	-2.6%
✅	tokenization/encode_tokens_per_second	358.63 µs	348.07 µs	-2.9%
✅	matmul/small_generic_f32_threads=8/1x256x256	33.38 µs	32.22 µs	-3.5%
✅	matmul/small_generic_bf16_threads=8/1x256x256	30.39 µs	29.15 µs	-4.1%
✅	add/small_f32_threads=1-internal/1024	186.2 ns	178.5 ns	-4.1%
✅	matmul/medium_generic_f16_threads=1/32x512x512	28.75 µs	27.45 µs	-4.6%
✅	add/large_bf16_threads=1-internal/4194304	1.63 ms	1.53 ms	-6.2%
✅	add/small_f16_threads=1-internal/1024	457.8 ns	427.0 ns	-6.7%
✅	add/medium_f32_threads=1-internal/262144	24.49 µs	22.83 µs	-6.8%
✅	qwen3_sampling_processors/top_k_top_p_fast	657.75 µs	609.61 µs	-7.3%
✅	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	651.14 µs	559.17 µs	-14.1%
✅	gather/medium_f32_threads=1-internal/32768	3.97 µs	3.40 µs	-14.4%
✅	add/medium_f16_threads=1-internal/262144	113.32 µs	96.94 µs	-14.5%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.18 ms	897.91 µs	-23.7%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	1.79 ms	1.34 ms	-25.0%
🟢	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	491.07 µs	361.86 µs	-26.3%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	72.73 µs	44.95 µs	-38.2%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	66.97 µs	38.84 µs	-42.0%
🟢	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	103.11 µs	50.67 µs	-50.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: { 3.25 3.17 5.10 }

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)