fix(native-cuda): re-zero recurrent/conv state on session reuse (LinearAttention corruption)

[justinchuby]

Summary

Native CUDA decode corrupted generation #2+ within a reused NativeDecodeSession on hybrid LinearAttention models (dense conv_state/recurrent_state). The first generate() was correct; subsequent generations on the same session returned non-deterministic degenerate garbage. Discovered during the 27B weight-offload A/B (see #384).

Repro (before): profile_native --runs 2 on the 27B →
native greedy decode was not deterministic: first=[11751,...] rerun=[279,6511,...].

Scope (Step-0 blast-radius test — decided early)

LinearAttention-only, NOT a general session-reuse bug.

• Plain GQA transformer (KV-only, no recurrent state), qwen2.5-0.5b-instruct-cuda, --runs 2: clean / deterministic.
• 27B hybrid LinearAttention: corrupts gen#2+.

So the general KV / position / decode-step reset path is fine for all multi-turn/server/REPL usage. Only models carrying fixed-size recurrent/conv state were affected.

Root cause

DecodeCudaState zero-initializes the fixed recurrent/conv-state device bindings once, in new(). reset() → rewind(0) (start of every generate()) only re-zeroed the attention-mask binding and reset logical_len. Growable KV is masked + length-tracked (stale slots inert); fixed recurrent state is an unmasked rolling cache, so gen#2 inherited gen#1's terminal recurrent state → garbage. Reproduces with CUDA-graph capture ON and with ONNX_GENAI_CUDA_GRAPH=0 (eager) → a state-reset bug, not a capture bug.

Fix (general — no model-specific special-casing)

• Track the fixed-state bindings as fixed_state_binding_range on DecodeCudaState (empty for pure-KV models).
• In rewind(), when target_len == 0, re-zero those bindings (same zero-init the constructor applies; state_pairs declare init: zeros). Only at the reset boundary — speculative recurrent rewind to a non-zero length is intentionally unsupported, mirroring the CPU path.

Pure-KV decoders have an empty range and are entirely unaffected (GQA path verified unchanged). CPU native was checked and is clean (its rewind(0) clears recurrent state via the past map), so the regression test is CUDA-gated.

Regression test (non-vacuous)

native_cuda_reused_session_rezeros_recurrent_state (#[cfg(feature="cuda")], gated by ONNX_GENAI_RUN_CUDA_SMOKE=1): a synthetic recurrent decoder whose logits are a direct function of the incoming conv_state, decoded twice across a reset(), asserting gen#1 == gen#2 (and that the per-step logits grow, proving state feeds logits).

Non-vacuity proof (re-zero disabled): gen#1 [0,60,144] != gen#2 [252,312,396] → FAIL. With the fix → PASS.

The existing profile_native --runs 2 determinism check is the end-to-end CI guard on real models.

End-to-end verification (greedy, "The capital of France is", --runs 2)

model	capture	gen#1==gen#2
27B int4 LinearAttention	ON (captures=2)	✅ " Paris.\n\n<think>\n\n</think>\n\nThat is correct. Paris is the capital and"
27B int4 LinearAttention	OFF (ONNX_GENAI_CUDA_GRAPH=0)	✅ identical
qwen2.5-0.5b GQA (regression)	ON	✅ unchanged

Guardrails

• Did not touch weight_paging.rs / provider.rs (harry-6, feat(cuda-offload): async fence-ordered weight page-in overlap (#87 first increment) #544 / squad/87-async-pagein).
• Did not switch the main checkout; worktree off origin/main (fa1afed).
• cargo fmt --all applied.

Refs #384. Possibly related to the repeated/degenerate-sentence report on multi-turn native decode.

⚠️ This will get independent review before merge.

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.16%. Comparing base (fa1afed) to head (aab80e4).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #554      +/-   ##
==========================================
+ Coverage   80.59%   81.16%   +0.57%     
==========================================
  Files         315      315              
  Lines      123446   123446              
  Branches   123446   123446              
==========================================
+ Hits        99489   100199     +710     
+ Misses      19907    19191     -716     
- Partials     4050     4056       +6     

Flag	Coverage Δ
cli-ort-linux	83.27% <ø> (ø)
cli-ort-windows	82.67% <ø> (-0.11%)	⬇️
mlas	77.91% <ø> (ø)
offline	81.12% <ø> (+0.60%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 5 files 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 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
⚠️	sampling_latency/top_p_per_token	905.80 µs	1.08 ms	+19.1%
⚠️	grammar_masking/llguidance_compute_mask/32	68.61 µs	79.54 µs	+15.9%
✅	logit_processing/seven_processor_chain_per_step	1.10 ms	1.25 ms	+13.8%
✅	sampling_latency/top_k_per_token	438.38 µs	494.72 µs	+12.9%
✅	gather/large_f32_threads=1-internal/131072	23.20 µs	25.99 µs	+12.0%
✅	sampling_latency/min_p_per_token	320.19 µs	356.79 µs	+11.4%
✅	kv_cache/alloc_dealloc_pages	35.07 µs	38.94 µs	+11.0%
✅	matmul/medium_generic_f16_threads=1/32x512x512	27.83 µs	29.09 µs	+4.5%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	1.24 ms	1.28 ms	+3.2%
✅	add/large_f32_threads=1-internal/4194304	37.60 ms	38.44 ms	+2.2%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	8.56 ms	8.65 ms	+1.1%
✅	gather/medium_f32_threads=1-internal/32768	3.40 µs	3.43 µs	+1.0%
✅	matmul/medium_generic_f16_threads=8/32x512x512	27.71 µs	27.89 µs	+0.6%
✅	matmul/small_generic_f32_threads=1/1x256x256	33.72 µs	33.93 µs	+0.6%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.14 ms	2.16 ms	+0.6%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.81 ms	1.82 ms	+0.6%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	485.07 µs	487.54 µs	+0.5%
✅	matmul/small_generic_f32_threads=8/1x256x256	32.06 µs	32.17 µs	+0.3%
✅	add/medium_f16_threads=1-internal/262144	2.41 ms	2.42 ms	+0.2%
✅	matmul/medium_generic_f32_threads=8/32x512x512	885.93 µs	887.57 µs	+0.2%
✅	gather/small_f32_threads=1-internal/4096	591.3 ns	591.9 ns	+0.1%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	72.86 µs	72.90 µs	+0.1%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	358.09 µs	358.20 µs	+0.0%
✅	matmul/small_generic_bf16_threads=8/1x256x256	29.72 µs	29.71 µs	-0.0%
✅	add/medium_f32_threads=1-internal/262144	2.35 ms	2.34 ms	-0.1%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	77.93 µs	77.87 µs	-0.1%
✅	matmul/small_generic_f16_threads=1/1x256x256	27.71 µs	27.69 µs	-0.1%
✅	tokenization/decode_tokens_per_second	5.73 ms	5.73 ms	-0.1%
✅	add/small_f16_threads=1-internal/1024	12.08 µs	12.05 µs	-0.2%
✅	add/medium_bf16_threads=1-internal/262144	2.43 ms	2.42 ms	-0.3%
✅	add/small_bf16_threads=1-internal/1024	11.89 µs	11.86 µs	-0.3%
✅	gather/small_f16_threads=1-internal/4096	436.8 ns	435.2 ns	-0.4%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.11 µs	14.06 µs	-0.4%
✅	sampling_latency/greedy_per_token	3.03 µs	3.01 µs	-0.5%
✅	add/large_bf16_threads=1-internal/4194304	39.01 ms	38.80 ms	-0.5%
✅	matmul/small_generic_bf16_threads=1/1x256x256	29.14 µs	28.91 µs	-0.8%
✅	gather/medium_f16_threads=1-internal/32768	2.16 µs	2.14 µs	-0.9%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	3.60 ms	3.56 ms	-1.2%
✅	gather/small_bf16_threads=1-internal/4096	439.3 ns	433.7 ns	-1.3%
✅	add/small_f32_threads=1-internal/1024	183.8 ns	180.7 ns	-1.7%
✅	reduce_mean/large_f32_threads=1-internal/262144	924.98 µs	908.36 µs	-1.8%
✅	gather/large_bf16_threads=1-internal/131072	9.55 µs	9.37 µs	-1.8%
✅	matmul/small_generic_f16_threads=8/1x256x256	28.94 µs	28.38 µs	-1.9%
✅	tokenization/encode_tokens_per_second	355.72 µs	348.63 µs	-2.0%
✅	add/large_f16_threads=1-internal/4194304	39.14 ms	38.34 ms	-2.0%
✅	reduce_mean/medium_f32_threads=1-internal/65536	236.17 µs	227.80 µs	-3.5%
✅	gather/medium_bf16_threads=1-internal/32768	2.26 µs	2.13 µs	-5.4%
✅	gather/large_f16_threads=1-internal/131072	10.48 µs	9.34 µs	-10.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.4.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 5.56 4.10 5.05 }

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)

[justinchuby]

VERDICT: APPROVE

Independent review by Harry (I did NOT author this; Mary did). Every claim in the PR description was re-verified from scratch in a scratch worktree (wt-harry-554) on an H200, device 0. All of Mary's claims hold.

Code review (fix commit aab80e4 — only 3 files: cuda.rs, tests.rs, decision note)

General, no special-casing (VERIFIED). The fix adds fixed_state_binding_range: Range to DecodeCudaState. In new(), bindings are laid out [0]=mask, then KV (kv_start..kv_end), then fixed recurrent/conv state (kv_end..fixed_state_end) — the sort_by_key(contains(fixed_state_inputs)) puts fixed-state pairs last, so the range covers exactly the conv_state/recurrent_state bindings and nothing else (not the mask, not KV). No model names, no architecture switches (cuda.rs:1116, 1412). For pure-KV decoders fixed_state_inputs is empty so kv_end == fixed_state_end -> empty range -> the re-zero loop is a no-op. Common path provably unchanged.

Right time (VERIFIED). rewind() re-zeros the range only at target_len == 0 (cuda.rs:1496-1514). reset() -> rewind(0) is called at the start of every generate_with_callback() (mod.rs:356). Mid-generation never calls rewind(0); non-zero rewind (incremental/speculative) intentionally leaves recurrent state intact, mirroring the CPU path (backend.rs:86-105: past.clear() at 0, recurrent left intact otherwise). The re-zero uses the same native_cuda_memset_zero + checked_shape_bytes(physical_shape, dtype) as the constructor's zero-init (cuda.rs:1106-1113) — symmetric. First-generation short-circuit (backend.rs:58, current_len==0) is safe because the constructor already zero-inits.

Run evidence (all commands via profile_native, which reuses ONE session across runs and bail!s on any cross-run mismatch — profile_native.rs:744-757)

27B int4 hybrid LinearAttention (overlay /home/justinchu/mary-model-overlays/qwen36-ortref, conv_state+recurrent_state state_pairs, 64 layers), prompt "The capital of France is", --tokens 16 --warmups 0 --runs 2:

• capture ON (captures=2 replays=26): DETERMINISTIC -> " Paris.\n\n\n\n\n\nThat is correct. Paris is the capital and"
• capture OFF (ONNX_GENAI_CUDA_GRAPH=0, captures=0): DETERMINISTIC -> identical output.
• BEFORE (re-zero disabled, rebuilt): FAILS -> "native greedy decode was not deterministic: first=[11751, ...], rerun=[279, 6511, 314, 9338, 369, 279, 6511, 314, 9338, ...]" — a degenerate repeating loop (the reported symptom). Matches Mary's claimed rerun=[279,6511,...] exactly.

So it reproduces both capture ON and OFF -> state-reset bug, not a capture bug. Confirmed.

Regression test native_cuda_reused_session_rezeros_recurrent_state (ONNX_GENAI_RUN_CUDA_SMOKE=1, feature cuda):

• WITH fix: PASS.
• NON-VACUITY — re-zero disabled: FAIL -> "gen#1 [0.0, 60.0, 144.0] != gen#2 [252.0, 312.0, 396.0]". Exactly the values Mary cited. Restored -> PASS again.
• Test is genuinely non-vacuous: logits = ReduceSum(Cast(conv_state)) so a stale state changes emitted logits; it asserts first[0]==0.0 and strictly-growing per-step logits (proving state feeds logits) before asserting gen#1==gen#2.

GQA qwen2.5-0.5b int4 (KV-only common path), --runs 2: DETERMINISTIC -> " Paris. It is the largest city in France and the second largest in the European". No regression on the common path.

Gates

• cargo fmt --all --check: clean (exit 0).
• Touched crate onnx-genai-engine builds under cuda,native-backend.
• weight_paging.rs / provider.rs: NOT touched by the fix commit (verified via git show --name-only). Guardrail honored.

Conclusion

Root cause, scope (LinearAttention-only / empty range for pure-KV), fix generality, reset-boundary timing, CPU parity, the non-vacuous regression test, and end-to-end 27B determinism (capture ON and OFF) are all independently confirmed. Approving for merge.

— Harry (independent reviewer)