fix: correct RecurrentGemma prefill state contract

TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.en.md

@@ -0,0 +1,38 @@
+# PR #2300: Correct the RecurrentGemma no-ID prefill contract
+
+**Date**: 2026-10-14
+
+**Status**: Complete
+
+**Risk**: Medium (prefill policy and long-prompt transient memory)
+
+## Summary
+
+RecurrentGemma now declares `supports_chunked_prefill() == false`. A multi-token call without a sequence ID deliberately replaces its model-owned RGLRU state so independent prompts do not inherit history; advertising chunk continuation to no-ID callers could therefore discard earlier chunks.
+
+## Problem
+
+The model inherited the trait default of chunked-prefill support even though `forward_for_sequence` resets fallback state at the start of every multi-token no-ID forward. Current DirectEngine uses a stable sequence ID and avoids that reset, so this change does not claim that its previous state handling truncated prompts.
+
+## Design
+
+- Override the capability while preserving independent no-ID prompt reset and stable-ID continuation.
+- Load a deterministic nonzero one-layer recurrent Griffin model through `GriffinModel::from_weights`.
+- Populate fallback convolution and RGLRU state with a single-token forward, then prove a new multi-token no-ID prompt matches the same prompt on a fresh model.
+- Split a five-token prompt into two and three tokens under one stable sequence ID, then compare the continuation logits with the matching suffix of a one-pass prompt using absolute or relative tolerance below `1e-3`.
+
+## Verification
+
+The root and surgery CUDA `test-fast` build graph passed in 14 minutes 25 seconds. The GPU-locked `recurrent_gemma_` focused suite passed 5 of 5 tests in 0.84 seconds, including the loaded-model contract and existing model-owned state tests. The source also passes `git diff --check`. The acceptance fixture is synthetic and checkpoint-shaped; no external checkpoint run is claimed.
+
+## Tradeoff
+
+The capability cannot express support conditional on a sequence ID. Every honoring caller, including DirectEngine, now receives a whole-prompt prefill plan even though DirectEngine's stable-ID state handling is safe. This can increase peak transient memory for long prompts compared with chunk-bounded prefill.
+
+## Change summary
+
+One Rust model file adds the capability override, a deterministic loaded-model fixture, and one state-contract regression without new dependencies.
+
+## Final integration validation
+
+Integrated main `5f334629`, including MLX `243bdf09`, mlxcelverse, and preceding model fixes, without source conflicts. The rebuilt CUDA `test-fast` binary passes all seven `models::recurrent_gemma` tests, including the loaded nonzero-weight RGLRU fixture. This freshly checks convolution/recurrent state population, no-ID prompt reset, and stable-ID split logits against one-pass execution on the final pin. Formatting and diff checks pass. No external-checkpoint inference result is claimed.

TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.ko.md

@@ -0,0 +1,38 @@
+# PR #2300: RecurrentGemma no-ID prefill 계약 수정
+
+**작성일**: 2026-10-14
+
+**상태**: 완료
+
+**위험도**: 중간 (prefill 정책 및 긴 prompt의 transient memory)
+
+## 요약
+
+RecurrentGemma가 이제 `supports_chunked_prefill() == false`를 선언합니다. Sequence ID가 없는 다중 토큰 호출은 독립적인 prompt 사이에 history가 섞이지 않도록 model-owned RGLRU state를 의도적으로 교체하므로, no-ID 호출자에게 chunk continuation을 제공하면 앞선 chunk를 잃을 수 있습니다.
+
+## 문제
+
+`forward_for_sequence`가 모든 다중 토큰 no-ID forward 시작 시 fallback state를 초기화하는데도 이 모델은 trait의 chunked-prefill 지원 기본값을 상속했습니다. 현재 DirectEngine은 안정적인 sequence ID를 사용해 이 초기화를 피하므로, 이번 변경은 이전 DirectEngine 상태 처리가 prompt를 잘랐다고 주장하지 않습니다.
+
+## 설계
+
+- 독립 no-ID prompt 초기화와 stable-ID continuation을 보존하면서 capability를 override합니다.
+- 0이 아닌 결정적 one-layer recurrent Griffin 모델을 `GriffinModel::from_weights`로 불러옵니다.
+- Single-token forward로 fallback convolution 및 RGLRU state를 채운 뒤 새 multi-token no-ID prompt가 fresh model의 동일 prompt와 일치하는지 검증합니다.
+- 하나의 stable sequence ID로 5-token prompt를 2개와 3개로 나누고 continuation logits를 one-pass prompt의 동일 suffix와 비교합니다. 절대 또는 상대 허용 오차는 `1e-3` 미만입니다.
+
+## 검증
+
+Root 및 surgery CUDA `test-fast` build graph가 14분 25초에 통과했습니다. GPU lock 아래에서 실행한 `recurrent_gemma_` focused suite도 loaded-model contract 및 기존 model-owned state 테스트를 포함해 5개 중 5개가 0.84초에 통과했습니다. 소스는 `git diff --check`도 통과했습니다. 승인 fixture는 synthetic checkpoint shape이며 외부 checkpoint 실행을 주장하지 않습니다.
+
+## 트레이드오프
+
+Capability는 sequence ID 유무에 따른 지원 여부를 표현할 수 없습니다. DirectEngine의 stable-ID 상태 처리는 안전하지만 DirectEngine을 포함해 capability를 따르는 모든 호출자가 whole-prompt prefill plan을 사용합니다. 따라서 긴 prompt에서 chunk-bounded prefill보다 peak transient memory가 늘 수 있습니다.
+
+## 변경 요약
+
+Rust 모델 파일 하나에 capability override, 결정적 loaded-model fixture, state-contract 회귀 테스트 하나를 새 의존성 없이 추가합니다.
+
+## 최종 통합 검증
+
+MLX `243bdf09`, mlxcelverse, 앞서 머지한 모델 수정을 포함한 main `5f334629`를 소스 충돌 없이 통합했다. 다시 빌드한 CUDA `test-fast` binary에서 `models::recurrent_gemma` 테스트 7개가 모두 통과했다. 비영 가중치를 적재한 RGLRU fixture를 포함하며, 최종 pin에서 convolution/recurrent 상태 생성, ID 없는 prompt 초기화, 같은 ID의 분할 logits와 한 번에 실행한 결과의 일치를 새로 검증했다. Formatting과 diff 검사도 통과했다. 외부 체크포인트 추론 결과를 주장하지 않는다.

src/models/recurrent_gemma.rs

@@ -1080,6 +1080,24 @@ impl LanguageModel for GriffinModel {
         self.config.num_hidden_layers
     }
 
+    /// Prefill through the cache-less [`LanguageModel::forward`] entry point
+    /// runs in one pass for this family.
+    ///
+    /// [`Self::forward_for_sequence`] treats a multi-token forward carrying no
+    /// sequence id as a new prompt and resets the model-owned recurrent and KV
+    /// state. That keeps unrelated no-ID prompts isolated, but it also means a
+    /// no-ID caller splitting one prompt across forwards would discard every
+    /// earlier chunk. DirectEngine carries a stable sequence id and therefore
+    /// avoids this reset, but the trait capability also covers cache-less
+    /// callers and cannot express support conditional on an id being present.
+    ///
+    /// Opting out makes every prefill-plan caller, including DirectEngine,
+    /// process the full prompt in one graph. This preserves the no-ID contract
+    /// at the cost of the transient-memory bound that chunked prefill provides.
+    fn supports_chunked_prefill(&self) -> bool {
+        false
+    }
+
     fn supports_padded_prefill(&self) -> bool {
         false // Padding tokens corrupt RGLRU recurrent state
     }
@@ -1123,10 +1141,190 @@ impl LanguageModel for GriffinModel {
 
 #[cfg(test)]
 mod tests {
-    use super::{GriffinConfig, GriffinLayerCache, make_griffin_caches};
+    use super::{GriffinConfig, GriffinLayerCache, GriffinModel, make_griffin_caches};
     use crate::models::model_owned::ModelOwnedSequenceState;
+    use mlxcel_core::MlxArray;
     use mlxcel_core::cache::SequenceId;
     use mlxcel_core::dtype;
+    use mlxcel_core::generate::LanguageModel;
+    use mlxcel_core::weights::WeightMap;
+
+    fn recurrent_gemma_test_config() -> GriffinConfig {
+        serde_json::from_value(serde_json::json!({
+            "model_type": "recurrent_gemma",
+            "hidden_size": 8,
+            "intermediate_size": 16,
+            "num_hidden_layers": 1,
+            "num_attention_heads": 2,
+            "num_key_value_heads": 1,
+            "vocab_size": 32,
+            "block_types": ["recurrent"],
+            "conv1d_width": 4
+        }))
+        .expect("the one-layer RecurrentGemma config must parse")
+    }
+
+    fn recurrent_gemma_test_weights() -> WeightMap {
+        let mut weights = WeightMap::new();
+        let mut seed = 0x2288_u32;
+        let mut insert = |name: &str, shape: &[i32]| {
+            let n = shape.iter().product::<i32>() as usize;
+            let values: Vec<f32> = (0..n)
+                .map(|_| {
+                    seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
+                    ((seed % 97) + 1) as f32 / 970.0
+                })
+                .collect();
+            weights.insert(
+                name.to_string(),
+                mlxcel_core::from_slice_f32(&values, shape),
+            );
+        };
+
+        insert("model.embed_tokens.weight", &[32, 8]);
+        insert("model.layers.0.temporal_block.linear_y.weight", &[8, 8]);
+        insert("model.layers.0.temporal_block.linear_x.weight", &[8, 8]);
+        insert("model.layers.0.temporal_block.linear_out.weight", &[8, 8]);
+        insert("model.layers.0.temporal_block.conv_1d.weight", &[8, 4, 1]);
+        insert("model.layers.0.temporal_block.conv_1d.bias", &[8]);
+        insert("model.layers.0.temporal_block.rg_lru.recurrent_param", &[8]);
+        insert(
+            "model.layers.0.temporal_block.rg_lru.input_gate_weight",
+            &[2, 4, 4],
+        );
+        insert(
+            "model.layers.0.temporal_block.rg_lru.input_gate_bias",
+            &[2, 4],
+        );
+        insert(
+            "model.layers.0.temporal_block.rg_lru.recurrent_gate_weight",
+            &[2, 4, 4],
+        );
+        insert(
+            "model.layers.0.temporal_block.rg_lru.recurrent_gate_bias",
+            &[2, 4],
+        );
+        insert("model.layers.0.mlp_block.gate_proj.weight", &[16, 8]);
+        insert("model.layers.0.mlp_block.up_proj.weight", &[16, 8]);
+        insert("model.layers.0.mlp_block.down_proj.weight", &[8, 16]);
+        insert("model.layers.0.temporal_pre_norm.weight", &[8]);
+        insert("model.layers.0.channel_pre_norm.weight", &[8]);
+        insert("model.final_norm.weight", &[8]);
+        weights
+    }
+
+    fn read_all_f32(array: &MlxArray) -> Vec<f32> {
+        mlxcel_core::eval(array);
+        let flat = mlxcel_core::reshape(array, &[-1]);
+        let len = mlxcel_core::array_shape(&flat)[0];
+        (0..len)
+            .map(|i| mlxcel_core::item_f32(&mlxcel_core::slice(&flat, &[i], &[i + 1])))
+            .collect()
+    }
+
+    fn assert_close(actual: &MlxArray, expected: &MlxArray, context: &str) {
+        let actual = read_all_f32(actual);
+        let expected = read_all_f32(expected);
+        assert_eq!(actual.len(), expected.len(), "{context}: element count");
+        for (index, (actual, expected)) in actual.iter().zip(&expected).enumerate() {
+            let abs = (actual - expected).abs();
+            let rel = abs / actual.abs().max(expected.abs()).max(1e-6);
+            assert!(
+                abs < 1e-3 || rel < 1e-3,
+                "{context}: element {index} differs: actual={actual}, expected={expected}, abs={abs}, rel={rel}"
+            );
+        }
+    }
+
+    fn loaded_recurrent_gemma() -> GriffinModel {
+        GriffinModel::from_weights(
+            recurrent_gemma_test_config(),
+            recurrent_gemma_test_weights(),
+        )
+        .expect("the one-layer RecurrentGemma model must load")
+    }
+
+    #[test]
+    fn recurrent_gemma_loaded_model_chunking_contract_is_safe() {
+        let reset_model = loaded_recurrent_gemma();
+        let language_model: &dyn LanguageModel = &reset_model;
+        assert!(
+            !language_model.supports_chunked_prefill(),
+            "a no-ID chunked prefill would discard the state built by every chunk but the last"
+        );
+
+        let mut unused_caches = language_model.make_caches();
+        let contaminant = mlxcel_core::from_slice_i32(&[7], &[1, 1]);
+        let contaminated = language_model.forward(&contaminant, &mut unused_caches, None);
+        mlxcel_core::eval(&contaminated);
+        reset_model
+            .sequence_state
+            .with_sequence_state(None, |caches| match &caches[0] {
+                GriffinLayerCache::Recurrent(cache) => {
+                    assert!(
+                        cache.conv_cache.is_some(),
+                        "the contaminant must populate conv state"
+                    );
+                    assert!(
+                        cache.state_cache.is_some(),
+                        "the contaminant must populate RGLRU state"
+                    );
+                }
+                GriffinLayerCache::Attention(_) => panic!("layer 0 must be recurrent"),
+            });
+
+        let target = mlxcel_core::from_slice_i32(&[1, 2, 3], &[1, 3]);
+        let reset_logits = language_model.forward(&target, &mut unused_caches, None);
+        let fresh_model = loaded_recurrent_gemma();
+        let mut fresh_caches = LanguageModel::make_caches(&fresh_model);
+        let fresh_logits = LanguageModel::forward(&fresh_model, &target, &mut fresh_caches, None);
+        assert_close(
+            &reset_logits,
+            &fresh_logits,
+            "a multi-token no-ID prompt must start from fresh state",
+        );
+
+        let split_model = loaded_recurrent_gemma();
+        let split_id = SequenceId::from_raw(2288);
+        LanguageModel::prepare_sequence_state(&split_model, split_id);
+        let mut split_caches = LanguageModel::make_caches(&split_model);
+        let first = mlxcel_core::from_slice_i32(&[1, 2], &[1, 2]);
+        let second = mlxcel_core::from_slice_i32(&[3, 4, 5], &[1, 3]);
+        let first_logits = LanguageModel::forward_with_sequence_id(
+            &split_model,
+            &first,
+            Some(split_id),
+            &mut split_caches,
+            None,
+        );
+        mlxcel_core::eval(&first_logits);
+        let split_logits = LanguageModel::forward_with_sequence_id(
+            &split_model,
+            &second,
+            Some(split_id),
+            &mut split_caches,
+            None,
+        );
+
+        let whole_model = loaded_recurrent_gemma();
+        let whole_id = SequenceId::from_raw(2289);
+        LanguageModel::prepare_sequence_state(&whole_model, whole_id);
+        let mut whole_caches = LanguageModel::make_caches(&whole_model);
+        let whole = mlxcel_core::from_slice_i32(&[1, 2, 3, 4, 5], &[1, 5]);
+        let whole_logits = LanguageModel::forward_with_sequence_id(
+            &whole_model,
+            &whole,
+            Some(whole_id),
+            &mut whole_caches,
+            None,
+        );
+        let expected_tail = mlxcel_core::slice(&whole_logits, &[0, 2, 0], &[1, 5, 32]);
+        assert_close(
+            &split_logits,
+            &expected_tail,
+            "stable-ID continuation must match one-pass recurrent state",
+        );
+    }
 
     /// The conv used to be called with the PyTorch `[B, C, L]` layout, which
     /// MLX rejects because it reads the channel count off the last axis. That