test(engine): gemma4-e2b dual head_dim (256+512) e2e golden lock — closes #1438 512 arm end-to-end

[justinchuby]

Summary

Closes the last mile of #1438 (MAX_HEAD_DIM 256→512): proves the head_dim=512 full-attention arm runs end-to-end on a real dual-head-size model, not just unit parity.

Gemma-3n / gemma4-e2b interleaves 28 sliding layers at head_dim=256 with 7 full-attention layers at head_dim=512 (KV layers 4, 9, 14, 19, 24, 29, 34). This adds a native-CUDA greedy-decode coherence golden lock that exercises all 35 layers every step, including the 512 layers on the fused split-K decode kernel raised in #1438.

• crates/onnx-genai-engine/tests/gemma4_e2b_head_dim_512_decode_lock.rs — the lock.
• crates/onnx-genai-engine/tests/common/decode_lock.rs — assert_native_matches_golden_eager helper.

Result

• Coherent e2e (Y): "Hello" → "! How can I help you today?" (+ <end_of_turn>/<eos>).
• Byte-identical (Y): native CUDA == ORT CPU-EP greedy for the same export, 16/16 tokens.
• Reproducible: PASSES from a clean detached worktree checkout against a persisted, self-contained export (GPU-pinned, --test-threads=1, ~33 s).

Export (persisted, self-contained)

GEMMA4_E2B_512_DIR=/home/justinchu/gemma4-e2b-it-text-cuda

Every on-disk gemma4 export was broken (dangling symlink target / speculative draft / multimodal-mandatory audio+vision). I composed a standalone text-only single graph (input_ids+attention_mask+past KV → logits+present KV) by fusing the official pipeline's embedding (produces inputs_embeds+routed per_layer_inputs) into the decoder, baking image_features/audio_features to empty constants, with real weight copies in-dir and a <bos>-prepending tokenizer patch (Gemma degenerates without BOS). No Mobius change was needed.

RULES.md §2 (general head size)

Per-layer head sizes are read structurally by the engine KV bridge from each present.N.key shape (kv_bridge.rs::layer_configs_from_key_outputs) — never from a model name or a fixed value. No value-keyed 256/512 branch; a future 3-distinct-head-size model needs only a new export + golden.

Capture caveat (documented; not the 512 head)

The lock runs with CUDA-graph capture OFF. The composed graph's merged present-KV sequence axis is an opaque symbol the prefill workspace planner can't yet upper-bound for capture — and it trips on a head_dim=256 sliding layer (present.13.key), not a 512 layer. Eager runs the identical fused kernels, so the 512 fast-path coherence is fully locked. Root cause + minimal-fix directions are in .squad/decisions/inbox/deckard-gemma4-512-e2e.md (cleanest: a non-merged native Gemma-3n text export).

Do not self-merge — coordinator validation requested.

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

[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_cached_dense_repeated_call/1x1024x1024	75.30 µs	127.30 µs	+69.0%
🔴	gather/small_f32_threads=1-internal/4096	687.2 ns	1.15 µs	+67.2%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	374.12 µs	577.93 µs	+54.5%
🔴	matmul/medium_generic_f16_threads=8/32x512x512	32.93 µs	49.68 µs	+50.9%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	65.11 µs	96.17 µs	+47.7%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	790.36 µs	1.14 ms	+44.4%
⚠️	matmul/small_generic_f16_threads=8/1x256x256	32.47 µs	40.54 µs	+24.9%
⚠️	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	103.14 µs	125.27 µs	+21.4%
⚠️	matmul/medium_generic_f16_threads=1/32x512x512	32.91 µs	39.95 µs	+21.4%
⚠️	add/large_f32_threads=1-internal/4194304	613.90 µs	735.35 µs	+19.8%
⚠️	matmul/small_generic_bf16_threads=8/1x256x256	33.27 µs	39.81 µs	+19.6%
⚠️	matmul/large_generic_bf16_threads=8/32x1024x1024	1.61 ms	1.93 ms	+19.4%
⚠️	gather/small_f16_threads=1-internal/4096	454.7 ns	542.7 ns	+19.4%
⚠️	reduce_mean/large_f32_threads=1-internal/262144	1.01 ms	1.17 ms	+15.6%
✅	kv_cache/alloc_dealloc_pages	37.02 µs	42.23 µs	+14.1%
✅	matmul/small_generic_f16_threads=1/1x256x256	29.38 µs	33.46 µs	+13.9%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.71 ms	6.44 ms	+12.9%
✅	qwen3_sampling_processors/top_k_partial_selection	150.66 µs	169.35 µs	+12.4%
✅	grammar_masking/llguidance_compute_mask/32	74.38 µs	82.79 µs	+11.3%
✅	add/large_bf16_threads=1-internal/4194304	1.76 ms	1.95 ms	+10.9%
✅	gather/large_f16_threads=1-internal/131072	13.12 µs	14.47 µs	+10.3%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	9.60 ms	10.53 ms	+9.7%
✅	logit_processing/seven_processor_chain_per_step	319.30 µs	344.14 µs	+7.8%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.43 ms	2.60 ms	+7.2%
✅	reduce_mean/medium_f32_threads=1-internal/65536	252.88 µs	269.87 µs	+6.7%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	549.38 µs	582.60 µs	+6.0%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.50 ms	3.70 ms	+5.7%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.77 µs	16.45 µs	+4.3%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	535.20 µs	553.59 µs	+3.4%
✅	gather/small_bf16_threads=1-internal/4096	494.7 ns	511.6 ns	+3.4%
✅	qwen3_sampling_processors/top_k_top_p_fast	665.50 µs	687.75 µs	+3.3%
✅	add/medium_f16_threads=1-internal/262144	111.87 µs	112.64 µs	+0.7%
✅	add/medium_bf16_threads=1-internal/262144	107.78 µs	108.52 µs	+0.7%
✅	sampling_latency/min_p_per_token	217.94 µs	215.84 µs	-1.0%
✅	matmul/small_generic_bf16_threads=1/1x256x256	35.35 µs	34.72 µs	-1.8%
✅	add/small_bf16_threads=1-internal/1024	488.7 ns	476.5 ns	-2.5%
✅	tokenization/encode_tokens_per_second	422.21 µs	406.11 µs	-3.8%
✅	add/large_f16_threads=1-internal/4194304	1.89 ms	1.81 ms	-4.0%
✅	matmul/small_generic_f32_threads=1/1x256x256	40.66 µs	38.31 µs	-5.8%
✅	gather/large_bf16_threads=1-internal/131072	14.20 µs	13.23 µs	-6.8%
✅	sampling_latency/top_p_per_token	431.94 µs	396.88 µs	-8.1%
✅	sampling_latency/greedy_per_token	3.79 µs	3.48 µs	-8.2%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.39 ms	2.19 ms	-8.3%
✅	add/small_f16_threads=1-internal/1024	529.3 ns	479.3 ns	-9.5%
✅	tokenization/decode_tokens_per_second	7.29 ms	6.38 ms	-12.5%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.61 ms	2.23 ms	-14.8%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	102.98 µs	87.54 µs	-15.0%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	488.66 µs	412.07 µs	-15.7%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.22 ms	999.23 µs	-18.0%
🟢	sampling_latency/top_k_per_token	68.00 µs	55.51 µs	-18.4%
🟢	gather/medium_f32_threads=1-internal/32768	5.76 µs	4.70 µs	-18.4%
🟢	matmul/small_generic_f32_threads=8/1x256x256	46.10 µs	37.40 µs	-18.9%
🟢	add/small_f32_threads=1-internal/1024	259.5 ns	206.1 ns	-20.6%
🟢	add/medium_f32_threads=1-internal/262144	32.22 µs	24.18 µs	-25.0%
🟢	gather/medium_f16_threads=1-internal/32768	3.57 µs	2.51 µs	-29.6%
🟢	gather/large_f32_threads=1-internal/131072	36.61 µs	25.01 µs	-31.7%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	5.97 ms	4.05 ms	-32.1%
🟢	matmul/large_generic_f16_threads=8/32x1024x1024	150.19 µs	93.33 µs	-37.9%
🟢	gather/medium_bf16_threads=1-internal/32768	4.11 µs	2.53 µs	-38.3%

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.48 3.28 5.09 }

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)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.19%. Comparing base (354be8f) to head (6358b6d).
⚠️ Report is 101 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1442       +/-   ##
===========================================
+ Coverage   80.12%   82.19%    +2.06%     
===========================================
  Files         364       12      -352     
  Lines      160709     5471   -155238     
  Branches   160709     5471   -155238     
===========================================
- Hits       128775     4497   -124278     
+ Misses      27283      775    -26508     
+ Partials     4651      199     -4452     

Flag	Coverage Δ
cli-ort-windows	82.19% <ø> (?)
mlas	?
offline	?

Flags with carried forward coverage won't be shown. Click here to find out more.
see 376 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.