[justinchuby]

Summary

• keep the native pipeline decoder alive across sequential requests
• retain CUDA KV in its original device dtype, including BF16
• rewind directly to the shared token prefix and prefill only the new suffix
• honor model.runtime_configurable.chunked_prefill.chunk_size to bound multi-token prefill activation peaks
• avoid paged F32 host mirroring and GPU/CPU KV round-trips
• preserve the underlying CUDA error unless generation actually failed because a required attachment input was missing

Validation

• cargo build --release -p onnx-genai-cli --features native-cuda
• cargo test -p onnx-genai-cli only_input_not_found_errors_are_classified_as_missing_required_input
• pipeline unit tests: 68 passed
• Muse Glimmer 30B BF16 KV integration (with fix(cli): load native pipelines without ORT preflight #1141): second turn reused 48/73 prompt tokens (66%); TTFT dropped from 6.8s to 1.5s; output remained correct
• long Muse multi-turn run crossed 256, 512, 1024, 2048, and 4096 KV buckets and continued generating through the former 441-token prefill OOM point; the harness stopped it only at its 300-second timeout

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.39%. Comparing base (266a6fe) to head (47a63ee).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1152      +/-   ##
==========================================
- Coverage   80.58%   80.39%   -0.20%     
==========================================
  Files         370      359      -11     
  Lines      162350   157243    -5107     
  Branches   162350   157243    -5107     
==========================================
- Hits       130827   126412    -4415     
+ Misses      26786    26270     -516     
+ Partials     4737     4561     -176     

Flag	Coverage Δ
cli-ort-linux	?
cli-ort-windows	?
mlas	85.52% <ø> (-0.39%)	⬇️
offline	80.29% <ø> (-0.08%)	⬇️

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

Reviewed the diff and pushed one fix (48afe093). Merging.

The bug I found

The retained decoder and the context describing it were stored under different conditions:

• decoder: use_native_decoder && paged_session.is_none()
• context: retains_kv && inputs_digest.is_some()

These diverge. digest_request_identity returns None whenever absorb_value cannot canonicalize an input, and use_kv() can be false. Such a request overwrote the decoder KV while leaving self.retained describing an earlier request. The next request derives its rewind target from retained, so it would rewind to a prefix length computed against a context the KV no longer holds — reusing one request's KV as another's prefix.

This fails silently. No crash, no error, just plausible text conditioned on the wrong context. That is the worst failure mode for a caching change, and it is exactly the class of bug a TTFT measurement cannot see.

Fixed by storing the decoder inside the same branch that refreshes the context, so the two cannot diverge. When the context is not refreshed the decoder is dropped, costing only the reuse.

What I verified

• cargo test -p onnx-genai-engine --features native-backend --lib pipeline — 71 passed, 0 failed with the fix applied.
• cargo check -p onnx-genai-engine --features native-backend clean.
• The rewind_kv -> 0 decline branch is currently unreachable: NativePipelineDecoder::rewind_kv always returns Ok(true), and every other implementor reports current_kv_len() == None, so target == current == 0 and the call is never made. Defensive, not dead-wrong — but note that if a future backend ever returns false, that path maps to "reuse nothing" without having rewound, which would be the same silent-corruption shape as the bug above. Worth an assertion when a second implementor appears.
• take() on entry is the right choice: any early return drops the decoder rather than leaving a stale one, so the failure direction is "lose the speedup", not "corrupt the context".

What I could not verify

The Muse Glimmer 30B BF16 numbers (48/73 tokens reused, TTFT 6.8s -> 1.5s). This box has no CUDA runtime. I accept them as reported and am explicit that they are unreproduced here; the correctness review above is from the code, which is where the risk actually was.

Follow-up (not blocking)

This PR adds a second prefix-reuse implementation: the native path computes shared/current/reusable/must_reset/target inline, in parallel with reusable_prefix_len used by ORT. Same policy, two copies, and they already disagree in detail — the native copy has no use_kv gate, the ORT copy has no shared.min(current) clamp. The decode loop itself is properly shared (one run_decode_loop, one PipelineDecodeLoopBackend), so this is the one place the backends are treated differently. I am unifying it now behind a single policy plus a small KV-rewind seam, so the backends stay interchangeable.

[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
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.40 ms	6.12 ms	+13.3%
✅	sampling_latency/min_p_per_token	193.19 µs	218.64 µs	+13.2%
✅	gather/small_bf16_threads=1-internal/4096	442.0 ns	498.9 ns	+12.9%
✅	kv_cache/alloc_dealloc_pages	37.22 µs	41.92 µs	+12.6%
✅	qwen3_sampling_processors/top_k_partial_selection	132.93 µs	148.86 µs	+12.0%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.31 ms	3.67 ms	+10.8%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	506.36 µs	541.30 µs	+6.9%
✅	grammar_masking/llguidance_compute_mask/32	72.55 µs	76.65 µs	+5.6%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.04 ms	2.14 ms	+4.8%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	438.44 µs	456.30 µs	+4.1%
✅	qwen3_sampling_processors/top_k_top_p_fast	670.90 µs	694.55 µs	+3.5%
✅	sampling_latency/top_k_per_token	48.83 µs	50.38 µs	+3.2%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	153.99 µs	158.38 µs	+2.9%
✅	tokenization/encode_tokens_per_second	346.38 µs	355.04 µs	+2.5%
✅	sampling_latency/greedy_per_token	2.97 µs	2.98 µs	+0.4%
✅	matmul/small_generic_bf16_threads=1/1x256x256	29.23 µs	29.08 µs	-0.5%
✅	logit_processing/seven_processor_chain_per_step	335.95 µs	332.72 µs	-1.0%
✅	gather/medium_bf16_threads=1-internal/32768	2.25 µs	2.21 µs	-1.7%
✅	gather/small_f16_threads=1-internal/4096	448.1 ns	439.8 ns	-1.8%
✅	reduce_mean/medium_f32_threads=1-internal/65536	230.17 µs	225.78 µs	-1.9%
✅	gather/medium_f16_threads=1-internal/32768	2.27 µs	2.21 µs	-2.7%
✅	tokenization/decode_tokens_per_second	5.79 ms	5.62 ms	-3.0%
✅	reduce_mean/large_f32_threads=1-internal/262144	943.84 µs	914.62 µs	-3.1%
✅	gather/medium_f32_threads=1-internal/32768	3.51 µs	3.40 µs	-3.1%
✅	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	41.23 µs	39.33 µs	-4.6%
✅	add/small_bf16_threads=1-internal/1024	442.0 ns	419.8 ns	-5.0%
✅	add/large_f16_threads=1-internal/4194304	1.67 ms	1.59 ms	-5.2%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	9.15 ms	8.62 ms	-5.8%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.95 µs	13.96 µs	-6.7%
✅	sampling_latency/top_p_per_token	382.81 µs	355.85 µs	-7.0%
✅	gather/large_f16_threads=1-internal/131072	11.66 µs	10.84 µs	-7.1%
✅	add/medium_bf16_threads=1-internal/262144	104.97 µs	95.81 µs	-8.7%
✅	add/small_f16_threads=1-internal/1024	458.3 ns	410.1 ns	-10.5%
✅	gather/large_f32_threads=1-internal/131072	24.55 µs	21.91 µs	-10.7%
✅	add/medium_f16_threads=1-internal/262144	110.84 µs	98.24 µs	-11.4%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	88.75 µs	78.00 µs	-12.1%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	4.04 ms	3.53 ms	-12.7%
✅	add/large_bf16_threads=1-internal/4194304	1.76 ms	1.53 ms	-13.2%
🟢	gather/small_f32_threads=1-internal/4096	736.6 ns	624.1 ns	-15.3%
🟢	matmul/medium_generic_bf16_threads=1/32x512x512	580.04 µs	487.99 µs	-15.9%
🟢	add/medium_f32_threads=1-internal/262144	26.63 µs	22.23 µs	-16.5%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	89.04 µs	73.70 µs	-17.2%
🟢	add/small_f32_threads=1-internal/1024	235.8 ns	194.3 ns	-17.6%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	50.34 µs	41.43 µs	-17.7%
🟢	matmul/large_generic_bf16_threads=1/32x1024x1024	2.22 ms	1.80 ms	-18.6%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	36.32 µs	29.30 µs	-19.3%
🟢	matmul/small_generic_f16_threads=1/1x256x256	34.96 µs	28.19 µs	-19.4%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	1.55 ms	1.24 ms	-20.2%
🟢	matmul/small_generic_f32_threads=1/1x256x256	44.08 µs	34.98 µs	-20.6%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	2.83 ms	2.13 ms	-24.7%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	37.49 µs	27.92 µs	-25.5%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	37.29 µs	27.71 µs	-25.7%
🟢	gather/large_bf16_threads=1-internal/131072	14.52 µs	10.68 µs	-26.5%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	489.08 µs	356.50 µs	-27.1%
🟢	add/large_f32_threads=1-internal/4194304	828.48 µs	573.56 µs	-30.8%
🟢	matmul/small_generic_f16_threads=8/1x256x256	41.31 µs	28.10 µs	-32.0%
🟢	matmul/small_generic_f32_threads=8/1x256x256	53.41 µs	32.50 µs	-39.2%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	860.01 µs	509.24 µs	-40.8%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	2.03 ms	883.11 µs	-56.5%

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.40 4.59 5.25 }

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)