[justinchuby]

Replace KvLayout enum with a KV-cache stride descriptor

Refs #783, #782, #787. Contract change plus a kernel-indexing generalization, not a layout migration. Head-major BNSH stays the default and the behavior for all existing metadata. Do not merge; do not close any issue.

Why a descriptor, not a third enum variant

cuMemMap maps whole granule-aligned VA windows onto whole physical granules, so layout controls the VMM residency floor (BNSH: layers x 2 x kv_heads; BSNH seq-major #782: layers x 2; token-major #787: 1 per sequence, a measured 768x reduction). But the next step is not another axis permutation:

1. Token-major collapses the per-(layer, side) buffer boundary both existing variants preserve, so it is not expressible as another enum axis-permutation.
2. Layout preference is per-EP / per-platform, not global. ORT GQA is BNSH-only on every dispatch path (group_query_attention_helper.h hardcodes Q_K_V_BNSH); Intel reportedly prefers BNHS (coherent, because Level Zero exposes ~64 KiB pages where the layout penalty is ~32x smaller than our 2 MiB granule). An enum would need a variant per (platform, preference).
3. We JIT with NVRTC, so strides are known at module-build time and each layout compiles to a specialized variant — MODULE_KEY versioning already does exactly this. Generality is close to free on a JIT backend.

The descriptor

Absolute element strides are a serving-time property (they depend on the cache_capacity a runtime picks), so the descriptor stores each axis stride symbolically — as the unordered set of runtime dims it multiplies. Empty factor list = unit stride.

axis	head-major BNSH	seq-major BSNH
batch	kv_heads·seq·head_dim	seq·kv_heads·head_dim
head	seq·head_dim	head_dim
seq	head_dim	kv_heads·head_dim
head_dim	1	1

head_dim is unit (contiguous) for every honored layout — the fp16 read vectorizes it as half2, the fused write addresses it as dst + d.

• Metadata (onnx-genai-metadata): kv_layout: Option<KvCacheLayout>, an untagged enum of a readable Named shorthand (serde aliases head_major_bnsh / seq_major_bsnh) or an Explicit KvStrideDescriptor. Absent still means head-major exactly as before. Committed JSON schema regenerated; schema-sync test green. A raw stride tuple is unreadable, so the common cases stay nameable and existing string-valued metadata keeps working.
• CUDA EP (onnx-runtime-ep-cuda): KvCacheStrides generates the CUDA #define prelude (GQA_KV_BASE/GQA_KV_STRIDE for the read, GQA_KV_DST for the write) that specializes each kernel per descriptor. Static per-layout module-cache keys (decode v8, prep v3); head-major reuses the shared prep module via an #ifndef GQA_KV_DST baked default, so only seq-major prepends a specialized prelude under a distinct key — avoiding double-compiling the always-f32 metadata/reference prep kernels.

How both existing layouts map

The two named layouts expand to exactly the strides the old enum variants implied (asserted by *_strides_match_legacy_enum tests, so the migration cannot silently change a layout). The kv_layout node attribute (0=BNSH, 1=BSNH) still selects them on the wire via KvCacheStrides::from_attribute; the .with_kv_layout(0/1) test entry point is preserved as a thin adapter.

How token-major would map (not implemented here)

KvStrideDescriptor carries reservation_offset_elements and reservation_seq_slots: token-major stores every layer's tokens in one reservation and hands each (layer, side) a sub-view whose seq stride is taken over the reservation's total token count, starting at a non-zero offset. That is expressible today but rejected by require_converted_path_support (as is any non-unit head_dim or unnamed descriptor) — the hard gate errors rather than mis-indexing.

The hard gate is preserved

Head-major is honored by every reader and writer. Any non-head-major layout is gated to the one converted fp16 single-token decode pair; anything the converted path cannot honor errors, never mis-indexes.

Follow-up increment

Make bindings actual views into one reservation (wire up reservation_* end to end), then convert the remaining BNSH kernels (prefill flash read) so seq-major can be enabled end to end, then token-major and #777 page-level prefix sharing. All explicitly out of scope here.

Validation (all local, green)

• cargo fmt -p onnx-genai-metadata -p onnx-runtime-ep-cuda -- --check
• cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings (exact CI gate, no --all-targets)
• cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings
• metadata: 33 fixtures + schema-sync + 5 descriptor tests
• ep-cuda: 6 kv_stride unit tests (named constructors reproduce legacy strides; reservation-view and non-unit-head_dim descriptors rejected)
• GPU: gqa_seqmajor_parity_gpu (BNSH byte-identical to BSNH, capture on==off, CPU oracle holds) + group_query_attention_gpu (25 head-major regression tests), single-threaded

Note on an unrelated fix

Added the missing memory_strategy_plan field to a native-engine test fixture in runtime.rs — a pre-existing merge artifact between #752 (added the field) and #758 (added a sibling fixture without it) that broke the onnx-genai-engine all-targets clippy gate on main. Independent of this change; included so the required gate passes.

[codecov]

Codecov Report

❌ Patch coverage is 97.87234% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.02%. Comparing base (7246e63) to head (268911c).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-genai-metadata/src/schema/model_io.rs	97.87%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #792      +/-   ##
==========================================
+ Coverage   79.41%   80.02%   +0.61%     
==========================================
  Files         331      331              
  Lines      137690   137782      +92     
  Branches   137690   137782      +92     
==========================================
+ Hits       109346   110266     +920     
+ Misses      23996    23162     -834     
- Partials     4348     4354       +6     

Flag	Coverage Δ
cli-ort-linux	83.50% <ø> (ø)
cli-ort-windows	83.09% <ø> (ø)
mlas	81.29% <ø> (+0.36%)	⬆️
offline	79.86% <97.87%> (+0.64%)	⬆️

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

Files with missing lines	Coverage Δ
crates/onnx-genai-metadata/src/schema/model_io.rs	96.29% <97.87%> (+40.04%)	⬆️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[justinchuby]

Coordinator pre-merge validation on the merged branch (Windows, RTX box), after
merging current main (no conflicts):

• cargo fmt -p onnx-runtime-ep-cuda -p onnx-genai-metadata -p onnx-genai-engine -- --check clean
• cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings clean (exact CI gate, no --all-targets)
• cargo test -p onnx-runtime-ep-cuda --lib kv_stride: 6 passed
• cargo test -p onnx-genai-metadata: 33 fixtures + schema-sync passed
• gqa_seqmajor_parity_gpu: passed on device
• group_query_attention_gpu head-major regression: 25 passed, 0 failed

I also independently checked the "pre-existing broken gate" this PR reports
fixing forward: cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings is clean on current main, so that
breakage has already been resolved upstream. The fix-forward here is harmless
either way, but it is no longer load-bearing.

Why this is the right shape

Three unit tests are what make this safe rather than merely working:

• named_layouts_are_honored — the named constructors reproduce exactly the
  strides the old enum variants implied, so replacing the enum cannot silently
  change a layout. That was my main risk with this change.
• reservation_view_is_rejected and non_unit_head_dim_is_rejected — the hard
  gate holds in descriptor form. Anything the converted path cannot honor
  errors rather than mis-indexing, which is the property that lets seq-major
  exist in the tree without being able to half-work.
• decode_prelude_matches_layout_formulae / prep_prelude_matches_layout_formulae
  — the generated NVRTC preludes are pinned to the layout formulae, so a stride
  bug shows up as a failing test rather than as wrong logits.

And the GPU parity test still asserts what matters, including the independent
CPU oracle — without it, a shared indexing error across both layouts would
still produce "byte-identical" agreement and pass.

Why a descriptor rather than a third enum variant

Recorded so this is not re-litigated:

1. Token-major collapses the per-(layer, side) buffer boundary that both
   existing variants preserve, so it is not expressible as another axis
   permutation — the descriptor's reservation-view fields are what will carry it.
2. Layout preference is per-EP and per-platform, not global. ORT GQA is
   BNSH-only on every dispatch path; Intel reportedly prefers BNHS, which is
   coherent because Level Zero exposes ~64 KiB pages where the layout penalty is
   ~32x smaller than our measured 2 MiB granule. An enum would need a variant per
   (platform, preference) combination.
3. We JIT with NVRTC, so strides are compile-time constants per module —
   generality costs nothing at runtime. That is a real advantage of a JIT backend
   over an AOT one, and this PR uses it: static per-layout module keys (decode v8,
   prep v3), no runtime branch, and head-major reuses the shared prep module via
   #ifndef GQA_KV_DST.

Default behavior is unchanged: absent metadata still means head-major, and
seq-major stays hard-gated to the one converted fp16 path.

Notably, the descriptor expressed everything cleanly — no ModelIoSpec
limitation surfaced, which was the main thing that could have forced a redesign.

Merging.

[github-actions]

⚠️ Benchmark Change 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
⚠️	gather/large_f32_threads=1-internal/131072	30.95 µs	39.40 µs	+27.3%
⚠️	matmul/small_generic_f16_threads=8/1x256x256	31.90 µs	39.95 µs	+25.2%
⚠️	reduce_mean/large_f32_threads=1-internal/262144	1.03 ms	1.29 ms	+25.2%
⚠️	matmul/small_generic_f16_threads=1/1x256x256	30.24 µs	37.13 µs	+22.8%
⚠️	grammar_masking/llguidance_compute_mask/32	81.22 µs	96.05 µs	+18.3%
⚠️	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.19 ms	1.40 ms	+17.9%
✅	matmul/medium_generic_f32_threads=8/32x512x512	1.54 ms	1.76 ms	+14.3%
✅	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	95.57 µs	108.19 µs	+13.2%
✅	qwen3_sampling_processors/top_k_partial_selection	142.39 µs	158.67 µs	+11.4%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	530.94 µs	590.00 µs	+11.1%
✅	matmul/small_generic_bf16_threads=1/1x256x256	37.88 µs	40.84 µs	+7.8%
✅	logit_processing/seven_processor_chain_per_step	352.46 µs	377.55 µs	+7.1%
✅	reduce_mean/medium_f32_threads=1-internal/65536	253.29 µs	269.92 µs	+6.6%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	691.69 µs	736.78 µs	+6.5%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	566.93 µs	599.96 µs	+5.8%
✅	qwen3_sampling_processors/top_k_top_p_fast	637.63 µs	674.72 µs	+5.8%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.44 ms	3.62 ms	+5.4%
✅	matmul/medium_generic_f16_threads=1/32x512x512	35.71 µs	37.58 µs	+5.2%
✅	sampling_latency/greedy_per_token	3.51 µs	3.69 µs	+5.2%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.27 µs	14.82 µs	+3.9%
✅	matmul/small_generic_f32_threads=8/1x256x256	44.74 µs	46.32 µs	+3.5%
✅	gather/small_f32_threads=1-internal/4096	837.7 ns	865.3 ns	+3.3%
✅	matmul/small_generic_bf16_threads=8/1x256x256	37.89 µs	38.92 µs	+2.7%
✅	matmul/medium_generic_f16_threads=8/32x512x512	42.74 µs	43.77 µs	+2.4%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	6.71 ms	6.83 ms	+1.9%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	80.02 µs	80.70 µs	+0.8%
✅	add/small_bf16_threads=1-internal/1024	14.30 µs	14.23 µs	-0.5%
✅	kv_cache/alloc_dealloc_pages	40.01 µs	39.74 µs	-0.7%
✅	sampling_latency/top_k_per_token	55.03 µs	52.52 µs	-4.6%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.10 ms	5.81 ms	-4.8%
✅	add/medium_f16_threads=1-internal/262144	2.76 ms	2.60 ms	-5.8%
✅	tokenization/encode_tokens_per_second	431.06 µs	400.37 µs	-7.1%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.20 ms	2.04 ms	-7.3%
✅	gather/medium_f16_threads=1-internal/32768	2.68 µs	2.48 µs	-7.4%
✅	gather/medium_bf16_threads=1-internal/32768	2.54 µs	2.34 µs	-7.9%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.73 ms	9.87 ms	-8.0%
✅	add/small_f16_threads=1-internal/1024	16.44 µs	14.88 µs	-9.5%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.29 ms	2.07 ms	-9.7%
✅	add/large_bf16_threads=1-internal/4194304	44.76 ms	40.36 ms	-9.8%
✅	sampling_latency/top_p_per_token	429.40 µs	386.24 µs	-10.1%
✅	sampling_latency/min_p_per_token	232.75 µs	206.55 µs	-11.3%
✅	add/medium_f32_threads=1-internal/262144	2.78 ms	2.46 ms	-11.3%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	105.14 µs	92.57 µs	-12.0%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	102.93 µs	90.40 µs	-12.2%
✅	add/large_f16_threads=1-internal/4194304	44.40 ms	38.77 ms	-12.7%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.90 ms	2.48 ms	-14.5%
✅	tokenization/decode_tokens_per_second	7.52 ms	6.41 ms	-14.7%
🟢	gather/medium_f32_threads=1-internal/32768	4.86 µs	4.02 µs	-17.2%
🟢	gather/large_bf16_threads=1-internal/131072	17.79 µs	14.64 µs	-17.7%
🟢	add/medium_bf16_threads=1-internal/262144	3.75 ms	3.04 ms	-18.9%
🟢	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	549.19 µs	418.61 µs	-23.8%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.33 ms	1.77 ms	-24.3%
🟢	add/large_f32_threads=1-internal/4194304	50.19 ms	37.91 ms	-24.5%
🟢	gather/small_f16_threads=1-internal/4096	662.1 ns	489.5 ns	-26.1%
🟢	gather/large_f16_threads=1-internal/131072	18.98 µs	13.80 µs	-27.3%
🟢	matmul/small_generic_f32_threads=1/1x256x256	55.10 µs	38.65 µs	-29.9%
🟢	gather/small_bf16_threads=1-internal/4096	772.5 ns	527.6 ns	-31.7%
🟢	add/small_f32_threads=1-internal/1024	339.2 ns	205.7 ns	-39.4%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	138.17 µs	72.35 µs	-47.6%

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: { 6.53 4.64 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)