diff --git a/TECHNICAL_REPORTS/2295-qwen3-next-chunked-prefill-20261014.en.md b/TECHNICAL_REPORTS/2295-qwen3-next-chunked-prefill-20261014.en.md new file mode 100644 index 000000000..1ed9d5f13 --- /dev/null +++ b/TECHNICAL_REPORTS/2295-qwen3-next-chunked-prefill-20261014.en.md @@ -0,0 +1,43 @@ +# Technical report: PR #2295 — Qwen3-Next chunked-prefill opt-out + +**Date**: 2026-10-14 + +**Status**: Completed + +**Risk level**: Medium + +## Executive summary + +Qwen3-Next now advertises `supports_chunked_prefill() == false`. This prevents cache-less `LanguageModel` callers from splitting a prompt across no-sequence-ID forwards that each reset the model-owned GatedDelta and KV state. A one-layer GatedDelta model loaded through the production weight constructor verifies the capability and both state paths. + +## Problem statement + +`Qwen3NextModel::forward_for_sequence` treats every multi-token forward without a sequence ID as a new prompt. That reset isolates independent cache-less prompts, but a no-ID chunking caller would discard the state built by earlier chunks. The inherited trait default advertised chunking support despite this contract. + +Current DirectEngine prefill passes one stable sequence ID through every piece, so its state handling retains the per-sequence state. `PrefillCaps` consumes this model-wide capability for every token-ID prefill plan, however, so returning false disables chunking for DirectEngine as well as cache-less callers without changing the server state lifecycle. + +## Technical decision + +The model opts out instead of changing `ModelOwnedSequenceState`. An explicit no-ID continuation mechanism would need a cross-family lifecycle signal that still resets between independent prompts. The local opt-out preserves that isolation, but the global capability makes DirectEngine and cache-less callers replace per-chunk transient memory bounds with one graph over a long prompt. + +## Change summary + +- Added the Qwen3-Next `supports_chunked_prefill` override with the no-ID state rationale and DirectEngine distinction. +- Added deterministic nonzero weights for a real one-layer GatedDelta fixture loaded by `Qwen3NextModel::from_weights`. +- Forced two no-ID prompt forwards and verified the second prompt resets its cache offset to 3; forced two forwards with one stable `SequenceId` and verified continuation reaches offset 5. + +## Validation + +The CUDA `mlxcel` and `mlxcel-surgery` library test graph builds successfully. The focused `qwen3_next_loaded_model_chunking_contract_is_safe` test passes 1/1 under the GPU lock in 0.99 seconds. Rust formatting and diff checks pass, and independent implementation review found no correctness issue. + +All required repository contracts pass, including 123 Python tests and 49 WebUI contract tests. An initial timing-sensitive GPU guard sample missed a short foreign process during concurrent native compilation; the isolated guard retry and broad contract rerun passed without code changes. + +## Follow-up actions + +- Parent integration validates the combined issue set after all authorized follow-up units land. + +## References + +- Issue #2285 +- Issue #1749 and PR #2293 +- Issue #672 and PR #1741 diff --git a/TECHNICAL_REPORTS/2295-qwen3-next-chunked-prefill-20261014.ko.md b/TECHNICAL_REPORTS/2295-qwen3-next-chunked-prefill-20261014.ko.md new file mode 100644 index 000000000..93f4345c4 --- /dev/null +++ b/TECHNICAL_REPORTS/2295-qwen3-next-chunked-prefill-20261014.ko.md @@ -0,0 +1,43 @@ +# 기술 보고서: PR #2295 — Qwen3-Next chunked-prefill opt-out + +**작성일**: 2026-10-14 + +**상태**: 완료 + +**위험도**: 중간 + +## 요약 + +Qwen3-Next가 이제 `supports_chunked_prefill() == false`를 명시한다. 이에 따라 cache-less `LanguageModel` 호출자가 prompt를 sequence ID 없는 여러 forward로 나눠 각 호출마다 model-owned GatedDelta 및 KV state를 초기화하는 경로를 선택하지 않는다. production weight constructor로 불러온 실제 1-layer GatedDelta 모델이 capability와 두 state 경로를 검증한다. + +## 1. 문제 정의 + +`Qwen3NextModel::forward_for_sequence`는 sequence ID가 없는 multi-token forward를 매번 새 prompt로 해석한다. 이 초기화는 독립적인 cache-less prompt를 격리하지만, no-ID chunking 호출자는 앞 chunk가 만든 state를 다음 chunk에서 버리게 된다. 기존 trait 기본값은 이 계약과 달리 chunking 지원을 선언했다. + +현재 DirectEngine prefill은 모든 piece에 동일한 sequence ID를 전달하므로 state 처리 관점에서 per-sequence state를 안전하게 유지한다. 그러나 `PrefillCaps`는 이 model-wide capability를 모든 token-ID prefill plan에 적용하므로 false를 반환하면 server state lifecycle은 바꾸지 않으면서 DirectEngine과 cache-less caller 모두의 chunking을 비활성화한다. + +## 3. 기술적 선택과 그 이유 + +`ModelOwnedSequenceState` 계약을 바꾸지 않고 모델이 opt-out하도록 선택했다. 명시적인 no-ID continuation mechanism은 독립 prompt 사이에서 계속 초기화하면서도 continuation을 구분하는 여러 family 공통 lifecycle signal이 필요하다. 이 opt-out은 격리를 보존하지만 model-wide capability이므로 DirectEngine과 cache-less caller 모두 긴 prompt에서 chunk별 transient memory bound 대신 전체 prompt를 하나의 graph로 처리한다. + +## 7. 변경 요약 + +- no-ID state 동작의 근거와 DirectEngine 경로의 차이를 설명하는 Qwen3-Next `supports_chunked_prefill` override를 추가했다. +- `Qwen3NextModel::from_weights`로 실제 1-layer GatedDelta fixture를 불러올 수 있도록 deterministic nonzero weight를 추가했다. +- no-ID prompt forward 두 번을 강제 평가해 두 번째 prompt의 cache offset이 3으로 초기화되는지 검증하고, 동일한 `SequenceId`의 forward 두 번은 offset 5까지 이어지는지 검증했다. + +## 검증 + +CUDA `mlxcel` 및 `mlxcel-surgery` library test graph가 성공적으로 빌드됐다. GPU lock 아래 focused `qwen3_next_loaded_model_chunking_contract_is_safe` test가 0.99초에 1/1 통과했다. Rust formatting과 diff check가 통과했으며 독립 implementation review에서 correctness 문제가 없었다. + +필수 repository contract가 모두 통과했으며 Python test 123개와 WebUI contract test 49개를 포함한다. 최초 timing-sensitive GPU guard sample은 native compile이 동시에 진행되는 동안 짧은 foreign process를 놓쳤지만, 코드 변경 없이 isolated guard retry와 broad contract rerun이 통과했다. + +## 8. 후속 조치 + +- 부모 통합 단계에서 승인된 모든 후속 unit이 반영된 뒤 결합된 issue set을 검증한다. + +## 참고 자료 + +- Issue #2285 +- Issue #1749 및 PR #2293 +- Issue #672 및 PR #1741 diff --git a/src/models/qwen3_next.rs b/src/models/qwen3_next.rs index 62d236c9e..28c99bcc3 100644 --- a/src/models/qwen3_next.rs +++ b/src/models/qwen3_next.rs @@ -1871,6 +1871,19 @@ impl LanguageModel for Qwen3NextModel { false } + /// A multi-token forward without a sequence id starts an independent + /// prompt and resets the model-owned GatedDelta/KV state. Splitting one + /// no-ID prompt across multiple forwards would therefore discard the + /// earlier chunks. DirectEngine carries a stable sequence id across its + /// prefill pieces, so its state handling is safe, but `PrefillCaps` applies + /// this model-wide capability to every token-ID prefill plan. The opt-out + /// therefore makes DirectEngine and cache-less callers prefill this family + /// in one whole-prompt graph, trading chunk-bounded transient memory for + /// correctness across the trait contract. + fn supports_chunked_prefill(&self) -> bool { + false + } + fn eos_token_ids(&self) -> Vec { vec![151645] // Qwen3 EOS token } diff --git a/src/models/qwen3_next_tests.rs b/src/models/qwen3_next_tests.rs index 36889ccd2..426ca9150 100644 --- a/src/models/qwen3_next_tests.rs +++ b/src/models/qwen3_next_tests.rs @@ -16,9 +16,11 @@ //! Qwen3-Next `SwitchLinear` loader applies before it stores them on a //! quantized expert plane (issue #958). +use mlxcel_core::cache::SequenceId; +use mlxcel_core::generate::LanguageModel; use mlxcel_core::weights::WeightMap; -use super::{Qwen3NextConfig, SwitchLinear}; +use super::{Qwen3NextCache, Qwen3NextConfig, Qwen3NextModel, SwitchLinear}; use crate::models::switch_layers::{HOSTILE_QUANT_PARAMS, insert_stacked_quantized_expert_plane}; /// Honest 4-bit expert geometry: `packed_in * 32 == bits * num_groups * @@ -62,6 +64,93 @@ fn config_with(group_size: i32, bits: i32) -> Qwen3NextConfig { .expect("test config must parse") } +fn recurrent_config() -> Qwen3NextConfig { + serde_json::from_value(serde_json::json!({ + "model_type": "qwen3_next", + "hidden_size": 8, + "vocab_size": 32, + "num_hidden_layers": 1, + "intermediate_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "linear_num_value_heads": 2, + "linear_num_key_heads": 1, + "linear_key_head_dim": 4, + "linear_value_head_dim": 4, + "linear_conv_kernel_dim": 4, + "full_attention_interval": 4, + "num_experts": 0, + "num_experts_per_tok": 1, + "decoder_sparse_step": 1, + "moe_intermediate_size": 8, + "shared_expert_intermediate_size": 8, + "tie_word_embeddings": true, + })) + .expect("test config must parse") +} + +fn deterministic_values(seed: u32, count: usize, scale: f32) -> Vec { + let mut state = seed; + (0..count) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + let unit = (state >> 8) as f32 / (1u32 << 24) as f32; + let value = (unit * 2.0 - 1.0) * scale; + if value.abs() < 1e-5 { + scale / 17.0 + } else { + value + } + }) + .collect() +} + +fn recurrent_weights() -> WeightMap { + let mut weights = WeightMap::new(); + let mut seed = 2285u32; + let mut tensor = |name: &str, shape: &[i32], scale: f32| { + seed += 1; + let count = shape.iter().product::() as usize; + weights.insert( + name.to_string(), + mlxcel_core::from_slice_f32(&deterministic_values(seed, count, scale), shape), + ); + }; + + tensor("model.embed_tokens.weight", &[32, 8], 0.2); + tensor("model.layers.0.linear_attn.conv1d.weight", &[16, 4, 1], 0.2); + tensor( + "model.layers.0.linear_attn.in_proj_qkvz.weight", + &[24, 8], + 0.2, + ); + tensor("model.layers.0.linear_attn.in_proj_ba.weight", &[4, 8], 0.2); + tensor("model.layers.0.linear_attn.dt_bias", &[2], 0.1); + tensor("model.layers.0.linear_attn.A_log", &[2], 0.1); + tensor("model.layers.0.linear_attn.norm.weight", &[4], 0.2); + tensor("model.layers.0.linear_attn.out_proj.weight", &[8, 8], 0.2); + tensor("model.layers.0.mlp.gate_proj.weight", &[8, 8], 0.2); + tensor("model.layers.0.mlp.up_proj.weight", &[8, 8], 0.2); + tensor("model.layers.0.mlp.down_proj.weight", &[8, 8], 0.2); + tensor("model.layers.0.input_layernorm.weight", &[8], 0.2); + tensor("model.layers.0.post_attention_layernorm.weight", &[8], 0.2); + tensor("model.norm.weight", &[8], 0.2); + + weights +} + +fn tokens(values: &[i32]) -> mlxcel_core::UniquePtr { + mlxcel_core::from_slice_i32(values, &[1, values.len() as i32]) +} + +fn linear_offset(caches: &[Qwen3NextCache]) -> i32 { + match caches.first().expect("one-layer fixture has one cache") { + Qwen3NextCache::Linear(cache) => cache.offset, + Qwen3NextCache::Attention(_) => panic!("layer 0 must be GatedDeltaNet"), + } +} + /// The pair this loader stores is handed straight to `gather_qmm`, which /// crosses the cxx bridge as `UniquePtr` rather than `Result`. A C++ /// throw there is an uncatchable `std::terminate`, @@ -116,3 +205,56 @@ fn qwen3_next_switch_linear_rejects_quantization_params_that_would_abort_gather_ SwitchLinear::from_weights(®ular, &config_with(0, 0), PREFIX) .expect("a bf16 expert plane must not be gated on quantization params"); } + +#[test] +fn qwen3_next_loaded_model_chunking_contract_is_safe() { + let _guard = crate::models::embedding_test_support::mlx_test_guard(); + let model = Qwen3NextModel::from_weights(&recurrent_weights(), &recurrent_config()) + .expect("one-layer GatedDeltaNet model must load"); + let language_model: &dyn LanguageModel = &model; + assert!( + !language_model.supports_chunked_prefill(), + "no-ID prompt chunks cannot retain the state built by earlier chunks" + ); + + let mut compatibility_caches = LanguageModel::make_caches(&model); + let first_prompt = + LanguageModel::forward(&model, &tokens(&[1, 2]), &mut compatibility_caches, None); + mlxcel_core::eval(&first_prompt); + let second_prompt = + LanguageModel::forward(&model, &tokens(&[3, 4, 5]), &mut compatibility_caches, None); + mlxcel_core::eval(&second_prompt); + assert_eq!( + model + .sequence_state + .with_sequence_state(None, |caches| linear_offset(caches)), + 3, + "a second no-ID prompt must start from fresh model-owned state" + ); + + let seq_id = SequenceId::from_raw(2285); + LanguageModel::prepare_sequence_state(&model, seq_id); + let first_chunk = LanguageModel::forward_with_sequence_id( + &model, + &tokens(&[6, 7]), + Some(seq_id), + &mut compatibility_caches, + None, + ); + mlxcel_core::eval(&first_chunk); + let second_chunk = LanguageModel::forward_with_sequence_id( + &model, + &tokens(&[8, 9, 10]), + Some(seq_id), + &mut compatibility_caches, + None, + ); + mlxcel_core::eval(&second_chunk); + assert_eq!( + model + .sequence_state + .with_sequence_state_ref(seq_id, linear_offset), + Some(5), + "stable sequence IDs must retain state across prefill chunks" + ); +}