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..926f7f72d --- /dev/null +++ b/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. 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..8e99627e6 --- /dev/null +++ b/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 검사도 통과했다. 외부 체크포인트 추론 결과를 주장하지 않는다. 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