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

.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<usize>` 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<think>\n\n</think>\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.

crates/onnx-genai-engine/src/native_decode/cuda.rs

@@ -121,6 +121,14 @@ pub(crate) struct DecodeCudaState {
     pub(crate) bindings: Vec<DeviceIoBinding>,
     pub(crate) base_binding_count: usize,
     pub(crate) kv_binding_range: std::ops::Range<usize>,
+    /// 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<usize>,
     pub(crate) auxiliary_binding_range: std::ops::Range<usize>,
     pub(crate) input_ids_binding: usize,
     pub(crate) position_ids_binding: Option<usize>,
@@ -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::<i64>(), &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)
     }
 

crates/onnx-genai-engine/src/native_decode/tests.rs

@@ -1010,6 +1010,205 @@ fn binding_addresses(session: &NativeDecodeSession) -> Vec<usize> {
         .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<NativeDecodeSession> {
+    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<Vec<f32>> {
+        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");