From d3a99fdcd1de2be250c78faa53fc3af341464a5e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 00:50:44 +0900 Subject: [PATCH 1/2] fix(models): correct Bailing linear prefill state contract Disable chunked prefill because no-ID multi-token forwards reset model-owned recurrent and KV state while preserving stable-ID continuation. Add a loaded six-layer regression for independent fallback reset and per-sequence continuation. The source passed independent review; focused CUDA validation is running under the orchestrator. Refs #2287 --- src/models/bailing_moe_linear.rs | 18 +++++ src/models/bailing_moe_linear_tests.rs | 97 +++++++++++++++++++++----- 2 files changed, 96 insertions(+), 19 deletions(-) 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(); From 5931280cdf90d88b11308d00c44a892755453484 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 01:00:24 +0900 Subject: [PATCH 2/2] docs: record Bailing linear prefill contract fix --- ...-moe-linear-chunked-prefill-20261014.en.md | 34 +++++++++++++++++++ ...-moe-linear-chunked-prefill-20261014.ko.md | 34 +++++++++++++++++++ 2 files changed, 68 insertions(+) create mode 100644 TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.en.md create mode 100644 TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.ko.md 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..9a89b400b --- /dev/null +++ b/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.en.md @@ -0,0 +1,34 @@ +# 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. 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..44f949997 --- /dev/null +++ b/TECHNICAL_REPORTS/2299-bailing-moe-linear-chunked-prefill-20261014.ko.md @@ -0,0 +1,34 @@ +# 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 회귀 테스트 하나를 추가합니다.