From ec1c32c8a7f8a0226bc62c2a9d4e49badcd3bff9 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 00:50:52 +0900 Subject: [PATCH 1/2] fix(models): correct RecurrentGemma prefill state contract Disable chunked prefill because no-ID multi-token forwards reset model-owned RGLRU state while preserving stable-ID continuation. Add a loaded one-layer regression that compares reset prompts with a fresh model and stable-ID split execution with one-pass logits. Focused CUDA validation is running under the orchestrator. Refs #2288 --- src/models/recurrent_gemma.rs | 200 +++++++++++++++++++++++++++++++++- 1 file changed, 199 insertions(+), 1 deletion(-) diff --git a/src/models/recurrent_gemma.rs b/src/models/recurrent_gemma.rs index 7b571143a..678277c86 100644 --- a/src/models/recurrent_gemma.rs +++ b/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::() as usize; + let values: Vec = (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 { + 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 From 66e9d40a1573951e9216f541c284be599ba82450 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 01:05:13 +0900 Subject: [PATCH 2/2] docs: record RecurrentGemma prefill contract fix --- ...rrent-gemma-chunked-prefill-20261014.en.md | 34 +++++++++++++++++++ ...rrent-gemma-chunked-prefill-20261014.ko.md | 34 +++++++++++++++++++ 2 files changed, 68 insertions(+) create mode 100644 TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.en.md create mode 100644 TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.ko.md diff --git a/TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.en.md b/TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.en.md new file mode 100644 index 000000000..7bd44ea68 --- /dev/null +++ b/TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.en.md @@ -0,0 +1,34 @@ +# 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. diff --git a/TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.ko.md b/TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.ko.md new file mode 100644 index 000000000..b43d58aac --- /dev/null +++ b/TECHNICAL_REPORTS/2300-recurrent-gemma-chunked-prefill-20261014.ko.md @@ -0,0 +1,34 @@ +# 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 회귀 테스트 하나를 새 의존성 없이 추가합니다.