fix: opt Qwen3-Next out of no-ID chunked prefill

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

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

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<i32> {
         vec![151645] // Qwen3 EOS token
     }

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<f32> {
+    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::<i32>() 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::MlxArray> {
+    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<MlxArray>` 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(&regular, &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"
+    );
+}