diff --git a/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.en.md b/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.en.md new file mode 100644 index 000000000..44dc362fb --- /dev/null +++ b/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.en.md @@ -0,0 +1,38 @@ +# PR #2299: Correct the Bailing MoE Linear no-ID prefill contract + +**Date**: 2026-10-14 + +**Status**: Complete + +**Risk**: Medium (prefill policy and long-prompt transient memory) + +## Summary + +Bailing MoE Linear now declares `supports_chunked_prefill() == false`. A multi-token call without a sequence ID deliberately replaces its model-owned recurrent and KV 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 the existing deterministic nonzero six-layer fixture through `BailingMoeLinearModel::from_weights`; it contains four recurrent linear-attention layers and two global-attention layers. +- Force no-ID two-token and three-token forwards, then verify layer 0 ends at offset 3 rather than 5. +- Force prepared stable-ID two-token and three-token forwards, then verify the same recurrent layer continues to offset 5. + +## Verification + +The root and surgery CUDA `test-fast` build graph passed in 14 minutes 20 seconds. The GPU-locked `bailing_moe_linear_loaded_model_chunking_contract_is_safe` regression passed 1 of 1 tests in 0.78 seconds. The source also passes `git diff --check` and independent implementation review. 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 + +The model adds one capability override, factors the existing deterministic nonzero fixture fill into a helper, and adds one loaded-model state-contract regression without new dependencies. + +## Final integration validation + +Integrated main `b159c89d`, including MLX `243bdf09`, mlxcelverse, and the preceding inference fixes, without source conflicts. The rebuilt CUDA `test-fast` binary passes all 36 `models::bailing_moe_linear` tests, including the loaded six-layer request-state regression. This freshly checks no-ID prompt reset and stable-ID continuation on the final pin. Formatting and diff checks pass. The model uses synthetic loaded weights; no external-checkpoint inference result is claimed. diff --git a/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.ko.md b/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.ko.md new file mode 100644 index 000000000..e68db3070 --- /dev/null +++ b/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.ko.md @@ -0,0 +1,38 @@ +# PR #2299: Bailing MoE Linear no-ID prefill 계약 수정 + +**작성일**: 2026-10-14 + +**상태**: 완료 + +**위험도**: 중간 (prefill 정책 및 긴 prompt의 transient memory) + +## 요약 + +Bailing MoE Linear가 이제 `supports_chunked_prefill() == false`를 선언합니다. Sequence ID가 없는 다중 토큰 호출은 독립적인 prompt 사이에 history가 섞이지 않도록 model-owned recurrent 및 KV 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이 아닌 결정적 six-layer fixture를 `BailingMoeLinearModel::from_weights`로 불러옵니다. 이 fixture는 recurrent linear-attention layer 네 개와 global-attention layer 두 개를 포함합니다. +- no-ID 2-token 및 3-token forward를 모두 평가한 뒤 layer 0 offset이 5가 아닌 3인지 확인합니다. +- 준비된 stable ID로 2-token 및 3-token forward를 모두 평가한 뒤 같은 recurrent layer가 offset 5까지 이어지는지 확인합니다. + +## 검증 + +Root 및 surgery CUDA `test-fast` build graph가 14분 20초에 통과했습니다. GPU lock 아래에서 실행한 `bailing_moe_linear_loaded_model_chunking_contract_is_safe` 회귀 테스트도 1개 중 1개가 0.78초에 통과했습니다. 소스는 `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가 늘 수 있습니다. + +## 변경 요약 + +모델에 capability override 하나를 추가하고 기존의 0이 아닌 결정적 fixture fill을 helper로 분리하며 새 의존성 없이 loaded-model state-contract 회귀 테스트 하나를 추가합니다. + +## 최종 통합 검증 + +MLX `243bdf09`, mlxcelverse, 앞서 머지한 추론 수정을 포함한 main `b159c89d`를 소스 충돌 없이 통합했다. 다시 빌드한 CUDA `test-fast` binary에서 `models::bailing_moe_linear` 테스트 36개가 모두 통과했다. 가중치를 적재한 6-layer 요청 상태 회귀 테스트를 포함하며, 최종 pin에서 ID 없는 prompt의 초기화와 같은 ID의 연속 처리를 새로 검증했다. Formatting과 diff 검사도 통과했다. 합성 가중치를 적재한 모델을 사용했으며 외부 체크포인트 추론을 주장하지 않는다. diff --git a/src/models/bailing_moe_linear.rs b/src/models/bailing_moe_linear.rs index 5ec91673e..3956cbd37 100644 --- a/src/models/bailing_moe_linear.rs +++ b/src/models/bailing_moe_linear.rs @@ -2726,6 +2726,24 @@ impl LanguageModel for BailingMoeLinearModel { false } + /// 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 + } + /// Opt out of NA tile-aligned padded prefill (issue #1755). /// /// The `BailingLinearCache::Linear` layers carry a `LinearAttentionCache`, diff --git a/src/models/bailing_moe_linear_tests.rs b/src/models/bailing_moe_linear_tests.rs index c358f8ab5..ba6fd2aa8 100644 --- a/src/models/bailing_moe_linear_tests.rs +++ b/src/models/bailing_moe_linear_tests.rs @@ -46,6 +46,7 @@ use super::{ Quantization, TokenIdField, alibi_slopes, chunked_prefill_enabled_from, gla_chunked, gla_sequential, gla_step, layer_decay, validate_weights, }; +use mlxcel_core::cache::SequenceId; use mlxcel_core::generate::LanguageModel; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; @@ -980,6 +981,28 @@ fn synthetic_weights(args: &ModelArgs) -> WeightMap { weights } +fn filled_weights(args: &ModelArgs) -> WeightMap { + let mut weights = synthetic_weights(args); + // An all-zero stack produces finite logits without exercising the decay, + // router, or gate. Fill every tensor deterministically so forward-based + // contract tests cover the real six-layer computation. + let mut keys: Vec = weights.keys().cloned().collect(); + // `WeightMap` is a `HashMap`, so sorting keeps the advancing seed stable + // across processes (issue #1265). + keys.sort(); + let mut seed = 0x5EED_1234u32; + for key in keys { + let shape = mlxcel_core::array_shape(weights.get(&key).expect("key just listed")); + let n: i32 = shape.iter().product(); + seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + weights.insert( + key, + mlxcel_core::from_slice_f32(&noise(n as usize, seed), &shape), + ); + } + weights +} + /// A shrunk Ring-mini: 6 layers in groups of 3 (so layers 2 and 5 are global), /// 4 experts, small dims. fn small_args() -> ModelArgs { @@ -1087,25 +1110,7 @@ fn a_dense_prefix_layer_is_validated_as_a_plain_mlp() { #[test] fn a_synthetic_model_builds_and_produces_finite_logits() { let args = small_args(); - let mut weights = synthetic_weights(&args); - // Replace the lazy zeros with real values: an all-zero stack produces - // technically-finite logits from a graph that never exercises the decay, - // the router or the gate. - let mut keys: Vec = weights.keys().cloned().collect(); - // Sorted because `WeightMap` is a `HashMap`: its iteration order is - // randomized per process, and the seed below advances once per key, so an - // unsorted walk builds a different random model on every run (issue #1265). - keys.sort(); - let mut seed = 0x5EED_1234u32; - for key in keys { - let shape = mlxcel_core::array_shape(weights.get(&key).expect("key just listed")); - let n: i32 = shape.iter().product(); - seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); - weights.insert( - key, - mlxcel_core::from_slice_f32(&noise(n as usize, seed), &shape), - ); - } + let weights = filled_weights(&args); let model = BailingMoeLinearModel::from_weights(&weights, &args).expect("the model builds"); assert_eq!(model.num_layers(), args.num_hidden_layers); @@ -1138,6 +1143,60 @@ fn a_synthetic_model_builds_and_produces_finite_logits() { } } +#[test] +fn bailing_moe_linear_loaded_model_chunking_contract_is_safe() { + let args = small_args(); + let weights = filled_weights(&args); + let model = BailingMoeLinearModel::from_weights(&weights, &args).expect("the model builds"); + let language_model: &dyn LanguageModel = &model; + assert!( + !language_model.supports_chunked_prefill(), + "a no-ID chunked prefill would discard the state built by every chunk but the last" + ); + + let two_tokens = mlxcel_core::from_slice_i32(&[1, 2], &[1, 2]); + let three_tokens = mlxcel_core::from_slice_i32(&[3, 4, 5], &[1, 3]); + let mut unused_caches = language_model.make_caches(); + + let first_prompt = language_model.forward(&two_tokens, &mut unused_caches, None); + mlxcel_core::eval(&first_prompt); + let second_prompt = language_model.forward(&three_tokens, &mut unused_caches, None); + mlxcel_core::eval(&second_prompt); + + let fallback_offset = + model + .sequence_state + .with_sequence_state(None, |caches| match &caches[0] { + BailingLinearCache::Linear(cache) => cache.offset, + BailingLinearCache::Attention(_) => panic!("layer 0 must use linear recurrence"), + }); + assert_eq!( + fallback_offset, 3, + "a second no-ID prompt must replace, rather than inherit, fallback state" + ); + + let seq = SequenceId::from_raw(2287); + language_model.prepare_sequence_state(seq); + let first_chunk = + language_model.forward_with_sequence_id(&two_tokens, Some(seq), &mut unused_caches, None); + mlxcel_core::eval(&first_chunk); + let second_chunk = + language_model.forward_with_sequence_id(&three_tokens, Some(seq), &mut unused_caches, None); + mlxcel_core::eval(&second_chunk); + + let continued_offset = model + .sequence_state + .with_sequence_state_ref(seq, |caches| match &caches[0] { + BailingLinearCache::Linear(cache) => cache.offset, + BailingLinearCache::Attention(_) => panic!("layer 0 must use linear recurrence"), + }) + .expect("the prepared sequence state must remain available"); + assert_eq!( + continued_offset, 5, + "consecutive chunks carrying one sequence id must retain their state" + ); +} + #[test] fn the_cache_flavour_follows_the_layer_schedule() { let args = small_args();