diff --git a/.squad/decisions/inbox/mary-session-reuse-reset.md b/.squad/decisions/inbox/mary-session-reuse-reset.md new file mode 100644 index 0000000000..90f60cf552 --- /dev/null +++ b/.squad/decisions/inbox/mary-session-reuse-reset.md @@ -0,0 +1,119 @@ +# Decision: native-CUDA session-reuse recurrent-state reset (LinearAttention corruption) + +- **Author:** Mary (Lead, large-model memory-offload workstream) +- **Date:** 2026-07-31T08:11:36+0000 +- **Requested by:** Justin (@justinchuby) +- **Branch:** `squad/native-session-reuse-reset` +- **References:** #384 (native-CUDA large-model decode); discovered during the + 27B weight-offload A/B (see `mary-27b-native-offload-ab.md`). Possible link to + Justin's "repeated / degenerate sentence" report on multi-turn native decode. + +## Scope (decided EARLY, per Step 0) + +**LinearAttention-only — NOT a general session-reuse bug.** + +Step-0 blast-radius test on a plain GQA transformer (KV-only, no conv/recurrent +state), `qwen2.5-0.5b-instruct-cuda`, `profile_native --runs 2`: + +- GQA: **clean / deterministic** — gen#2 byte-identical to gen#1. +- 27B hybrid LinearAttention (`conv_state`/`recurrent_state`): **corrupts** gen#2+. + +So the general KV/position/decode-step reset path is fine for everyone. The bug +is confined to models that carry fixed-size recurrent/conv state. + +## Root cause + +`DecodeCudaState` (`crates/onnx-genai-engine/src/native_decode/cuda.rs`) appends +the fixed recurrent/conv-state device bindings after the growable KV bindings. +Those fixed-state bindings are `memset_zero`'d **once, in `new()`**, and never +again. `reset()` → `rewind(0)` (called at the start of every `generate()`) only +re-zeroed `bindings[0]` (the attention mask) and reset `logical_len`. + +Growable KV is safe on reuse because it is masked and length-tracked (stale +slots are inert). Fixed recurrent state is a **wholesale, unmasked rolling +cache**: on the 2nd generation it inherited generation #1's *terminal* recurrent +state, so decode #2 started from garbage → non-deterministic degenerate output. + +This reproduces both with CUDA-graph capture ON and with `ONNX_GENAI_CUDA_GRAPH=0` +(eager), confirming it is a state-reset bug, not a capture bug. + +## Fix (general, no model-specific special-casing) + +In `cuda.rs`: +- Added `fixed_state_binding_range: Range` to `DecodeCudaState`, + populated in `new()` as `kv_end..fixed_state_end` (empty for pure-KV models). +- In `rewind()`, when `target_len == 0`, re-zero every binding in + `fixed_state_binding_range` via `native_cuda_memset_zero` (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. + +No model names, no architecture switches: any decoder that declares fixed +recurrent/conv state pairs is reset correctly. Pure-KV decoders have an empty +range and are entirely unaffected (verified: GQA path unchanged). + +CPU native was checked and is clean here — its full reset clears the recurrent +state via the `past` map on `rewind(0)` — so the regression test is CUDA-gated. + +## Regression test (non-vacuous) + +`native_cuda_reused_session_rezeros_recurrent_state` +(`crates/onnx-genai-engine/src/native_decode/tests.rs`, `#[cfg(feature="cuda")]`, +gated by `ONNX_GENAI_RUN_CUDA_SMOKE=1`). + +Builds a synthetic recurrent decoder whose **`logits` are a direct function of +the incoming `conv_state`** (`logits = ReduceSum(Cast(conv_state))`), and whose +next state accumulates the decoded token id. Decodes a fixed 3-token sequence, +`reset()`s, decodes it again, and asserts the two logits sequences are identical. +It also asserts the per-step logits strictly grow, proving the state actually +feeds the logits (guards against a vacuous pass). + +**Non-vacuity proof** — with the re-zero disabled: + +``` +gen#1 [0.0, 60.0, 144.0] != gen#2 [252.0, 312.0, 396.0] → FAIL +``` + +With the fix, gen#2 == gen#1 → PASS. (The existing `profile_native --runs 2` +determinism check is the natural end-to-end CI guard on real models.) + +## End-to-end verification (real models, device 4/6, greedy) + +Prompt `"The capital of France is"`, `--tokens 16 --warmups 0 --runs 2`: + +| model | capture | gen#1==gen#2 | output | +|---|---|---|---| +| 27B int4 LinearAttention | ON (captures=2) | ✅ deterministic | `" Paris.\n\n\n\n\n\nThat is correct. Paris is the capital and"` | +| 27B int4 LinearAttention | OFF (`ONNX_GENAI_CUDA_GRAPH=0`) | ✅ deterministic | identical to above | +| qwen2.5-0.5b GQA (regression) | ON | ✅ deterministic | `" Paris. It is the largest city in France and the second largest in the European"` | + +Before the fix the 27B failed: `native greedy decode was not deterministic: +first=[11751,...] rerun=[279,6511,...]`. + +## Exact repro command + +```bash +source .cudaenv.sh +export ONNX_GENAI_ORT_LIB=$ORT_ROOT/lib/libonnxruntime.so.1.27.0 +export ONNX_GENAI_REQUIRE_CUDA=1 CUDA_VISIBLE_DEVICES=4 +./target/release/profile_native \ + --model <27b-int4-native-pkg> --ep cuda --backend native \ + --tokens 16 --warmups 0 --runs 2 --prompt "The capital of France is" +``` + +(27B package: `model.onnx`/`model.onnx.data`/`tokenizer.json` from +`~/mary-models/qwen3.6-27b-int4-cuda/` + `inference_metadata.yaml` carrying the +`io:` block copied from the `qwen36-ortref` sibling, as in the offload A/B.) + +Unit regression: +```bash +ONNX_GENAI_RUN_CUDA_SMOKE=1 CUDA_VISIBLE_DEVICES=6 \ + cargo test --release -p onnx-genai-engine --lib --features "cuda native-backend" \ + native_cuda_reused_session_rezeros_recurrent_state +``` + +## Guardrails honored + +- Did not touch `weight_paging.rs` / `provider.rs` (harry-6, #544 / squad/87-async-pagein). +- Did not switch the main checkout branch; all work in worktree `wt-mary-session-reset` off `origin/main` (fa1afed3). +- This is the sync page-in / native-decode session-state domain only. diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index 713b84b40c..9fd6e620dc 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -121,6 +121,14 @@ pub(crate) struct DecodeCudaState { pub(crate) bindings: Vec, pub(crate) base_binding_count: usize, pub(crate) kv_binding_range: std::ops::Range, + /// Bindings for fixed-size recurrent/conv states (hybrid linear-attention + /// `conv_state`/`recurrent_state`). Unlike growable KV — which is masked and + /// tracked by `logical_len`, so stale slots are inert — these are wholesale + /// rolling caches with no masking. They are zero-initialized once in `new()`; + /// `rewind(0)` must re-zero them so a reused session starts every generation + /// from the declared `init: zeros` state (see `rewind`). Empty for pure-KV + /// decoders. + pub(crate) fixed_state_binding_range: std::ops::Range, pub(crate) auxiliary_binding_range: std::ops::Range, pub(crate) input_ids_binding: usize, pub(crate) position_ids_binding: Option, @@ -1105,6 +1113,7 @@ impl DecodeCudaState { )?; bindings.push(binding); } + let fixed_state_end = bindings.len(); let logits_meta = session .outputs() @@ -1400,6 +1409,7 @@ impl DecodeCudaState { bindings, base_binding_count, kv_binding_range: kv_start..kv_end, + fixed_state_binding_range: kv_end..fixed_state_end, auxiliary_binding_range: auxiliary_start..auxiliary_end, input_ids_binding, position_ids_binding, @@ -1483,6 +1493,25 @@ impl DecodeCudaState { self.bindings[0].write_bytes(target_len * std::mem::size_of::(), &zeros)?; } self.bindings[0].set_logical_shape(vec![1, target_len])?; + if target_len == 0 { + // Fixed-size recurrent/conv states are unmasked rolling caches: a + // reused session would otherwise inherit the previous generation's + // terminal state, corrupting generation #2+. A full reset restores + // the declared `init: zeros`. Speculative recurrent rewind to a + // non-zero length is unsupported (mirrors the CPU path), so only the + // reset boundary re-zeros here. + for index in self.fixed_state_binding_range.clone() { + let binding = &self.bindings[index]; + let bytes = checked_shape_bytes(binding.physical_shape(), binding.dtype) + .with_context(|| { + format!( + "fixed CUDA decoder state re-zero size overflow for binding {index} shape {:?}", + binding.physical_shape() + ) + })?; + native_cuda_memset_zero(binding.device_ptr() as usize, bytes)?; + } + } self.set_logical_len(target_len) } diff --git a/crates/onnx-genai-engine/src/native_decode/tests.rs b/crates/onnx-genai-engine/src/native_decode/tests.rs index 184f4ccb5a..3000fb5540 100644 --- a/crates/onnx-genai-engine/src/native_decode/tests.rs +++ b/crates/onnx-genai-engine/src/native_decode/tests.rs @@ -1010,6 +1010,205 @@ fn binding_addresses(session: &NativeDecodeSession) -> Vec { .collect() } +/// Build a synthetic single-layer CUDA decoder whose `logits` are a direct +/// function of the *incoming* recurrent `conv_state`, and whose next state +/// accumulates the decoded token id. Unlike `build_cuda_decoder_with_fixed_state` +/// (where `conv_state` only passes through `Identity` and never reaches the +/// logits, so a stale state is invisible), here a non-reset recurrent state +/// changes the emitted logits on the next generation — the exact signature of +/// the session-reuse corruption bug. Used by the regression test that a reused +/// `NativeDecodeSession` re-zeros recurrent state on `reset()`. +#[cfg(feature = "cuda")] +fn build_cuda_recurrent_logits_decoder(max_len: usize) -> anyhow::Result { + use prost::Message; + + let mut graph = Graph::new(); + graph.opset_imports.insert(String::new(), 13); + let total = graph.intern_symbol("total"); + let past = graph.intern_symbol("past"); + + let input_ids = + graph.create_named_value("input_ids", DataType::Int64, vec![1.into(), 1.into()]); + let attention_mask = graph.create_named_value( + "attention_mask", + DataType::Int64, + vec![1.into(), total.into()], + ); + let position_ids = + graph.create_named_value("position_ids", DataType::Int64, vec![1.into(), 1.into()]); + let past_key = graph.create_named_value( + "past_key_values.0.key", + DataType::Float32, + vec![1.into(), 1.into(), past.into(), 1.into()], + ); + let past_value = graph.create_named_value( + "past_key_values.0.value", + DataType::Float32, + vec![1.into(), 1.into(), past.into(), 1.into()], + ); + let conv_state = graph.create_named_value( + "past_key_values.0.conv_state", + DataType::Float16, + vec![1.into(), 4.into(), 3.into()], + ); + for input in [ + input_ids, + attention_mask, + position_ids, + past_key, + past_value, + conv_state, + ] { + graph.add_input(input); + } + + // logits = ReduceSum(Cast(conv_state)) — depends on the INCOMING recurrent + // state, so a stale (non-zeroed) conv_state changes the emitted logits. Cast + // to f32 first: cuDNN reduction does not support Float16 inputs. + let conv_f32 = graph.create_named_value( + "conv_state_f32", + DataType::Float32, + vec![1.into(), 4.into(), 3.into()], + ); + insert_op( + &mut graph, + "Cast", + vec![conv_state], + conv_f32, + &[("to", Attribute::Int(DataType::Float32 as i64))], + ); + let logits = graph.create_named_value( + "logits", + DataType::Float32, + vec![1.into(), 1.into(), 1.into()], + ); + insert_op( + &mut graph, + "ReduceSum", + vec![conv_f32], + logits, + &[("keepdims", Attribute::Int(1))], + ); + + // present.0.conv_state = conv_state + Cast(input_ids) — accumulate the token + // id into every recurrent slot so the state grows deterministically per step. + let token_f16 = + graph.create_named_value("token_f16", DataType::Float16, vec![1.into(), 1.into()]); + insert_op( + &mut graph, + "Cast", + vec![input_ids], + token_f16, + &[("to", Attribute::Int(DataType::Float16 as i64))], + ); + let present_conv_state = graph.create_named_value( + "present.0.conv_state", + DataType::Float16, + vec![1.into(), 4.into(), 3.into()], + ); + insert_op( + &mut graph, + "Add", + vec![conv_state, token_f16], + present_conv_state, + &[], + ); + + let present_key = graph.create_named_value( + "present.0.key", + DataType::Float32, + vec![1.into(), 1.into(), past.into(), 1.into()], + ); + insert_op( + &mut graph, + "Cast", + vec![past_key], + present_key, + &[("to", Attribute::Int(DataType::Float32 as i64))], + ); + let present_value = graph.create_named_value( + "present.0.value", + DataType::Float32, + vec![1.into(), 1.into(), past.into(), 1.into()], + ); + insert_op( + &mut graph, + "Cast", + vec![past_value], + present_value, + &[("to", Attribute::Int(DataType::Float32 as i64))], + ); + + for output in [logits, present_key, present_value, present_conv_state] { + graph.add_output(output); + } + + let model = onnx_std::Model::new(graph).to_proto()?.encode_to_vec(); + let session = InferenceSession::builder() + .model_bytes(&model) + .device(DevicePreference::Gpu { index: Some(0) }) + .build() + .context("build recurrent-logits CUDA decoder")?; + let mut io = tiny_decoder_io(); + io.state_pairs = Some(vec![LoopStatePair { + input: "past_key_values.0.conv_state".into(), + output: "present.0.conv_state".into(), + init: Some("zeros".into()), + update: Some("replace".into()), + }]); + NativeDecodeSession::from_session_with_cuda_kv_max_len_and_io(session, Some(max_len), Some(&io)) +} + +/// Regression guard for the native-CUDA session-reuse corruption bug: a reused +/// `NativeDecodeSession` must re-zero fixed recurrent/conv state on `reset()`, +/// so generation #2 starts from the declared `init: zeros` — not generation #1's +/// terminal state. Before the fix, `rewind(0)` left the recurrent binding stale +/// and the second decode sequence produced different logits (non-deterministic +/// degenerate output on hybrid LinearAttention models). This test is +/// non-vacuous: the logits are a direct function of the incoming recurrent +/// state, so a stale state changes the asserted values. +#[cfg(feature = "cuda")] +#[test] +fn native_cuda_reused_session_rezeros_recurrent_state() -> anyhow::Result<()> { + if std::env::var_os("ONNX_GENAI_RUN_CUDA_SMOKE").is_none() { + eprintln!("skipping CUDA smoke; set ONNX_GENAI_RUN_CUDA_SMOKE=1 to run"); + return Ok(()); + } + + const TOKENS: [TokenId; 3] = [5, 7, 9]; + let mut session = build_cuda_recurrent_logits_decoder(16)?; + + let run_once = |session: &mut NativeDecodeSession| -> anyhow::Result> { + session.reset()?; + let mut logits = Vec::new(); + for (position, token) in TOKENS.iter().enumerate() { + let step = session.decode(&[*token], position)?; + logits.push(step[0][0]); + } + Ok(logits) + }; + + let first = run_once(&mut session)?; + // The incoming recurrent state accumulates the decoded token ids, so the + // per-step logits form a strictly growing sequence — proving the state does + // flow into the logits (guards against a vacuous test). + assert_eq!( + first[0], 0.0, + "generation must start from zeroed recurrent state" + ); + assert!( + first[1] > first[0] && first[2] > first[1], + "recurrent state must accumulate into logits: {first:?}" + ); + + let second = run_once(&mut session)?; + assert_eq!( + first, second, + "reused session must re-zero recurrent state on reset(): gen#1 {first:?} != gen#2 {second:?}" + ); + Ok(()) +} + #[cfg(feature = "cuda")] fn input_update_stats(session: &NativeDecodeSession) -> [DeviceBindingTransferStats; 3] { let state = session.cuda.as_ref().expect("CUDA state");