Guard large unaligned Qwen expert gathers

src/engine/metal/qwen_mtplx.rs

@@ -39,6 +39,8 @@ mod issue106;
 mod issue108;
 #[cfg(test)]
 mod issue109;
+#[cfg(test)]
+mod issue110;
 mod linear;
 mod load;
 mod mlp;

src/engine/metal/qwen_mtplx/allocator.rs

@@ -230,6 +230,11 @@ impl Allocator {
         }
     }
 
+    #[cfg(test)]
+    pub(super) fn set_resource_limit_for_test(&self, limit: usize) -> usize {
+        std::mem::replace(&mut self.state.lock().unwrap().resource_limit, limit)
+    }
+
     pub(in super::super) fn stats(&self) -> Stats {
         let state = self.state.lock().unwrap();
         Stats {

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -250,13 +250,14 @@ fn issue103_sorted_gather_receipts() {
     super::configure_sources().unwrap();
     let _context = super::super::gpu::Context::open_qwen(0).unwrap();
     assert!(super::sparse_attention::nax_available());
-    let scope = super::encoder::Scope::new().unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
     let buffer =
         |count: usize, bits: u16| super::scalar_buffer(&bits.to_le_bytes().repeat(count)).unwrap();
     let w = super::scalar_buffer(&vec![0; 64 * 256 * 256 / 8 * 4]).unwrap();
     let scales = buffer(64 * 256 * 8, 0);
     let biases = buffer(64 * 256 * 8, 0x3f80); // all dequantized weights exactly 1
     for tokens in [3276_u32, 3277, 3404, 4095, 4096] {
+        let scope = super::encoder::Scope::new().unwrap();
         let rows = tokens * 10;
         let x = buffer(rows as usize * 256, 0x3f80);
         let lhs = super::scalar_buffer(&(0..rows).flat_map(u32::to_le_bytes).collect::<Vec<_>>())
@@ -295,6 +296,8 @@ fn issue103_sorted_gather_receipts() {
             "{}",
             serde_json::json!({"case":"sorted_gather","tokens":tokens,"routed_rows":rows,"bad_rows":bad.len(),"first_bad":bad.first(),"last_bad":bad.last(),"dispatches":format!("{dispatches:?}")})
         );
+        drop(scope);
+        super::issue110::sentinel(tokens, &streams);
         if tokens == 3276 || tokens == 4096 {
             assert!(bad.is_empty());
         } else {

src/engine/metal/qwen_mtplx/issue104.rs

@@ -12,7 +12,7 @@ fn fingerprint(execution: &Execution, array: &Array) -> serde_json::Value {
         super::super::qwen_mtplx_tests::hex_digest(&values.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>())})
 }
 
-fn state(execution: &Execution, saved: &SessionSnapshot) -> serde_json::Value {
+pub(super) fn state(execution: &Execution, saved: &SessionSnapshot) -> serde_json::Value {
     let states = saved
         .trunk
         .iter()

src/engine/metal/qwen_mtplx/issue110.rs

@@ -0,0 +1,344 @@
+//! Isolated stock sorted-gather guard checks; the fused row-map port is #112.
+use super::array::{Array, Dtype};
+use super::stream::{Device, Streams};
+use super::{linear, ops};
+use crate::engine::metal::qwen_mtplx_tests::{hex_digest, pattern};
+
+fn bytes(a: &Array, streams: &Streams) -> Vec<u8> {
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    let a = ops::contiguous(a, false, gpu).unwrap();
+    ops::evaluate(streams, std::slice::from_ref(&a), gpu, false).unwrap();
+    let mut data = vec![0; a.layout().nbytes()];
+    a.buffer().read(a.offset(), &mut data).unwrap();
+    data
+}
+
+fn unit_gather(rows: i32, streams: &Streams) -> Array {
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    let array = |shape: &[i32], bits: u16| {
+        let count = shape.iter().map(|&n| n as usize).product::<usize>();
+        Array::new(
+            shape,
+            Dtype::BF16,
+            super::scalar_buffer(&bits.to_le_bytes().repeat(count)).unwrap(),
+        )
+        .unwrap()
+    };
+    let x = array(&[rows, 1, 256], 0x3f80);
+    let w = Array::new(
+        &[64, 256, 32],
+        Dtype::U32,
+        super::scalar_buffer(&vec![0; 64 * 256 * 32 * 4]).unwrap(),
+    )
+    .unwrap();
+    let s = array(&[64, 256, 8], 0);
+    let b = array(&[64, 256, 8], 0x3f80);
+    let rhs = Array::new(
+        &[rows],
+        Dtype::U32,
+        super::scalar_buffer(
+            &(0..rows as u32)
+                .flat_map(|r| (r * 64 / rows as u32).to_le_bytes())
+                .collect::<Vec<_>>(),
+        )
+        .unwrap(),
+    )
+    .unwrap();
+    linear::guarded_gather_qmm_bf16(
+        &x,
+        [&w, &s, &b],
+        [None, Some(&rhs)],
+        [4, 32],
+        true,
+        streams,
+        gpu,
+    )
+    .unwrap()
+}
+
+pub(super) fn sentinel(tokens: u32, streams: &Streams) {
+    let rows = (tokens * 10) as i32;
+    let out = unit_gather(rows, streams);
+    assert_eq!(out.layout().shape(), &[rows, 1, 256]);
+    let (data, dispatches) =
+        crate::engine::metal::qwen_mtplx_tests::capture_dispatches(|| bytes(&out, streams));
+    let bad = data
+        .chunks_exact(512)
+        .enumerate()
+        .filter_map(|(i, row)| {
+            row.chunks_exact(2)
+                .any(|b| u16::from_le_bytes(b.try_into().unwrap()) != 0x4380)
+                .then_some(i)
+        })
+        .collect::<Vec<_>>();
+    assert!(bad.is_empty(), "unwritten or incorrect rows: {bad:?}");
+    println!(
+        "{}",
+        serde_json::json!({"case":"guarded_sentinel","tokens":tokens,"rows":rows,"bad_rows":bad.len(),"shape":out.layout().shape(),"dispatches":format!("{dispatches:?}")})
+    );
+}
+
+#[test]
+#[ignore = "installed layer weights and pinned oracle receipts; external memory supervisor required"]
+fn issue110_actual_projection_rows_match_oracle() {
+    super::configure_sources().unwrap();
+    let _context = crate::engine::metal::gpu::Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    assert!(super::sparse_attention::nax_available());
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    let receipts =
+        std::fs::read_to_string(std::env::var_os("DS4_ISSUE110_ORACLE").unwrap()).unwrap();
+    let weights = super::weights::load_shards(&root, &streams).unwrap();
+    let parts = |name: &str| {
+        ["weight", "scales", "biases"].map(|part| {
+            weights[&format!("language_model.model.layers.0.mlp.switch_mlp.{name}.{part}")].clone()
+        })
+    };
+    let gate = parts("gate_proj");
+    let up = parts("up_proj");
+    let gu = std::array::from_fn::<_, 3, _>(|i| {
+        ops::concatenate(&[gate[i].clone(), up[i].clone()], 1, gpu).unwrap()
+    });
+    let down = parts("down_proj");
+    ops::evaluate(&streams, &[gu.to_vec(), down.to_vec()].concat(), gpu, false).unwrap();
+    for line in receipts.lines() {
+        let row: serde_json::Value = serde_json::from_str(line).unwrap();
+        if row["case"] != "projection" {
+            continue;
+        }
+        let rows = row["rows"].as_i64().unwrap() as i32;
+        let k = row["k"].as_i64().unwrap() as i32;
+        let n = row["n"].as_i64().unwrap() as i32;
+        let layer = if row["projection"] == "gate_up" {
+            &gu
+        } else {
+            &down
+        };
+        let x = Array::new(&[rows, 1, k], Dtype::BF16, pattern((rows * k) as u32, 7)).unwrap();
+        let rhs = Array::new(
+            &[rows],
+            Dtype::U32,
+            super::scalar_buffer(
+                &(0..rows as u32)
+                    .flat_map(|r| (r * 512 / rows as u32).to_le_bytes())
+                    .collect::<Vec<_>>(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+        let lhs = row["left"].as_bool().unwrap().then(|| {
+            Array::new(
+                &[rows],
+                Dtype::I32,
+                super::scalar_buffer(
+                    &(0..rows)
+                        .rev()
+                        .flat_map(i32::to_le_bytes)
+                        .collect::<Vec<_>>(),
+                )
+                .unwrap(),
+            )
+            .unwrap()
+        });
+        let out = linear::guarded_gather_qmm_bf16(
+            &x,
+            [&layer[0], &layer[1], &layer[2]],
+            [lhs.as_ref(), Some(&rhs)],
+            [4, 64],
+            true,
+            &streams,
+            gpu,
+        )
+        .unwrap();
+        assert_eq!(out.layout().shape(), &[rows, 1, n]);
+        let (data, dispatches) =
+            crate::engine::metal::qwen_mtplx_tests::capture_dispatches(|| bytes(&out, &streams));
+        assert_eq!(hex_digest(&data), row["sha256"].as_str().unwrap(), "{row}");
+        println!(
+            "{}",
+            serde_json::json!({"case":"projection","projection":row["projection"],"rows":rows,"k":k,"n":n,"left":lhs.is_some(),"exact":true,"sha256":hex_digest(&data),"dispatches":format!("{dispatches:?}")})
+        );
+    }
+}
+
+thread_local! {
+    static PADS: std::cell::RefCell<Option<Vec<serde_json::Value>>> = const { std::cell::RefCell::new(None) };
+}
+
+pub(super) fn record(rows: i32, pad: i32, k: i32, n: i32) {
+    PADS.with_borrow_mut(|slot| {
+        if let Some(receipts) = slot {
+            receipts.push(serde_json::json!({"rows":rows,"pad":pad,"k":k,"n":n}));
+        }
+    });
+}
+
+#[test]
+#[ignore = "full installed model and large partial chunks; serial memory supervisor required"]
+fn issue110_large_cold_and_warm_prefill() {
+    assert_eq!(
+        std::env::var("MTPLX_GDN_BOUNDARY_TAIL_INTERVAL").as_deref(),
+        Ok("0")
+    );
+    use super::session_cache::{DecodeOptions, TurnOptions, TurnProgress};
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE110_DIR").unwrap());
+    std::fs::create_dir_all(&directory).unwrap();
+    let settings = crate::settings::effective_settings(
+        crate::model::ModelChoice::Qwen38FlashNext,
+        &crate::settings::GenerationPreferences {
+            context_tokens: 16384,
+            max_generated_tokens: 3,
+            temperature: Some(0.),
+            ..Default::default()
+        },
+        &crate::settings::RuntimePreferences::default(),
+        root.parent().unwrap(),
+    )
+    .unwrap();
+    let (mut execution, _) = super::execution::Execution::open(&settings.engine).unwrap();
+    let tokens = |len: usize| {
+        (0..len)
+            .map(|i| [9707, 11, 1075, 374, 264, 1092, 13][i % 7])
+            .collect::<Vec<_>>()
+    };
+    for chunk in [2048, 4096] {
+        let sampler = super::sampling::SamplerConfig {
+            temperature: 0.,
+            ..Default::default()
+        };
+        let options = TurnOptions {
+            prefill: super::prefill::Options {
+                chunk_size: Some(chunk),
+                final_logits_only: true,
+                external_cache_only: false,
+                external_emit_logits: true,
+                return_hidden: false,
+            },
+            decode: DecodeOptions::Ar(super::ar::Options {
+                max_tokens: 3,
+                seed: 1,
+                stop_tokens: &[],
+                sync_eval: false,
+                return_hidden: false,
+                capture_final_state: true,
+            }),
+            sampler,
+            draft_sampler: sampler,
+            lazy_kv: false,
+            store_min_suffix: None,
+            fused_max: 512,
+        };
+        let mut prompt = tokens(7374);
+        let mut saved = None;
+        for phase in ["cold", "warm"] {
+            if phase == "warm" && chunk == 4096 {
+                let original = super::issue104::state(&execution, saved.as_ref().unwrap());
+                let cancelled = execution.session_turn(
+                    &prompt,
+                    saved.as_ref(),
+                    &options,
+                    |event| {
+                        if matches!(
+                            event,
+                            TurnProgress::Prefill(super::prefill::Progress::Settled { .. })
+                        ) {
+                            Err("cancelled after padded chunk".into())
+                        } else {
+                            Ok(())
+                        }
+                    },
+                    |_| Ok(()),
+                    |_| Ok(()),
+                );
+                assert!(matches!(cancelled,Err(e) if e=="cancelled after padded chunk"));
+                assert_eq!(
+                    super::issue104::state(&execution, saved.as_ref().unwrap()),
+                    original
+                );
+                println!(
+                    "{}",
+                    serde_json::json!({"case":"cancel","chunk":chunk,"saved_state_unchanged":true})
+                );
+            }
+            PADS.with_borrow_mut(|slot| *slot = Some(Vec::new()));
+            let out = execution
+                .session_turn(
+                    &prompt,
+                    saved.as_ref(),
+                    &options,
+                    |event| {
+                        if let TurnProgress::Prefill(super::prefill::Progress::Settled {
+                            tokens,
+                            total,
+                        }) = event
+                        {
+                            eprintln!("#110 {chunk}/{phase} prefill {tokens}/{total}");
+                        }
+                        Ok(())
+                    },
+                    |_| Ok(()),
+                    |_| Ok(()),
+                )
+                .unwrap();
+            let pads = PADS.with_borrow_mut(Option::take).unwrap();
+            assert_eq!(pads.is_empty(), chunk == 2048);
+            assert!(out.final_snapshot_error.is_none());
+            let snapshot = out.snapshot.unwrap();
+            let state = super::issue104::state(&execution, &snapshot);
+            let record = serde_json::json!({"case":"turn","chunk":chunk,"phase":phase,"cached":out.cached_tokens,"prompt_len":prompt.len(),"state":state,"pads":pads});
+            std::fs::write(
+                directory.join(format!("{chunk}-{phase}.json")),
+                serde_json::to_vec(&record).unwrap(),
+            )
+            .unwrap();
+            println!(
+                "{}",
+                serde_json::json!({"case":"turn","chunk":chunk,"phase":phase,"cached":out.cached_tokens,"prompt_len":prompt.len(),"generated":&snapshot.token_ids[prompt.len()..],"padded_calls":pads.len()})
+            );
+            prompt = snapshot.token_ids.clone();
+            prompt.extend(tokens(3278));
+            saved = Some(snapshot);
+        }
+    }
+}
+
+#[test]
+#[ignore = "Metal allocation refusal and retry; external supervisor required"]
+fn issue110_guard_allocation_refusal_releases_graph_and_retries() {
+    super::configure_sources().unwrap();
+    let _context = crate::engine::metal::gpu::Context::open_qwen(0).unwrap();
+    for limit in [0, 1, 2] {
+        let allocator = super::allocator::Allocator::new().unwrap();
+        let streams = Streams::new(Some(allocator.clone()));
+        let gpu = streams.default_stream(Device::Gpu).unwrap();
+        let out = unit_gather(34040, &streams);
+        allocator.clear_cache();
+        let old = allocator.set_resource_limit_for_test(limit);
+        let result = ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false);
+        assert!(matches!(result, Err(e) if e.contains("resource limit")));
+        let partial_peak = allocator.stats().peak;
+        assert_eq!(partial_peak > 0, limit > 0);
+        drop(out);
+        drop(streams);
+        allocator.set_resource_limit_for_test(old);
+        allocator.clear_cache();
+        assert_eq!(allocator.stats().active, 0);
+        let streams = Streams::new(Some(allocator.clone()));
+        let out = unit_gather(34040, &streams);
+        let data = bytes(&out, &streams);
+        assert!(
+            data.chunks_exact(2)
+                .all(|v| u16::from_le_bytes(v.try_into().unwrap()) == 0x4380)
+        );
+        drop(out);
+        drop(streams);
+        allocator.clear_cache();
+        assert_eq!(allocator.stats().active, 0);
+        println!(
+            "{}",
+            serde_json::json!({"case":"allocation_refusal","limit":limit,"partial_peak":partial_peak,"released":true,"retry_exact":true})
+        );
+    }
+}

src/engine/metal/qwen_mtplx/linear.rs

@@ -297,6 +297,157 @@ fn ensure_row_contiguous_matrix(input: &Array) -> Result<Array, String> {
     Ok(copy)
 }
 
+/// MTPLX moe_sorted_gather._call_pad for our pinned 0.32.2 affine,
+/// transposed primitive. Hardware truth matters here, not the route fallback
+/// switch: the stock MLX kernel does not read that switch.
+fn sorted_gather_pad(
+    x: &[i32],
+    lhs: Option<&[i32]>,
+    rhs: Option<&[i32]>,
+    sorted: bool,
+    nax: bool,
+) -> i32 {
+    let Some(&[rows]) = rhs else { return 0 };
+    if !sorted
+        || !nax
+        || rows <= 32767
+        || rows % 64 == 0
+        || x.len() < 2
+        || x[x.len() - 2] != 1
+        || x[0] != rows
+        || lhs.is_some_and(|shape| shape != [rows])
+    {
+        return 0;
+    }
+    64 - rows % 64
+}
+
+#[test]
+fn sorted_gather_guard_matches_pinned_call_predicates() {
+    for (rows, pad) in [
+        (1, 0),
+        (32767, 0),
+        (32768, 0),
+        (32769, 63),
+        (32770, 62),
+        (34040, 8),
+        (40950, 10),
+        (40960, 0),
+        (65535, 1),
+        (65537, 63),
+        (81910, 10),
+    ] {
+        assert_eq!(
+            sorted_gather_pad(&[rows, 1, 256], None, Some(&[rows]), true, true),
+            pad
+        );
+        assert_eq!(
+            sorted_gather_pad(&[rows, 1, 256], Some(&[rows]), Some(&[rows]), true, true),
+            pad
+        );
+        assert_eq!(
+            sorted_gather_pad(&[rows, 1, 256], None, Some(&[rows]), false, true),
+            0
+        );
+        assert_eq!(
+            sorted_gather_pad(&[rows, 1, 256], None, Some(&[rows]), true, false),
+            0
+        );
+    }
+    for (x, lhs, rhs) in [
+        (vec![34040, 2, 256], None, Some(vec![34040])),
+        (vec![1, 1, 256], None, Some(vec![34040])),
+        (vec![34040, 1, 256], None, None),
+        (vec![34040, 1, 256], None, Some(vec![1, 34040])),
+        (vec![34040, 1, 256], Some(vec![1]), Some(vec![34040])),
+        (vec![34040, 1, 256], Some(vec![1, 34040]), Some(vec![34040])),
+    ] {
+        assert_eq!(
+            sorted_gather_pad(&x, lhs.as_deref(), rhs.as_deref(), true, true),
+            0
+        );
+    }
+}
+
+/// Guard both stock Gate/Up and Down gathers, including future fused-path
+/// fallbacks. Keep the underlying runtime primitive unchanged.
+pub(super) fn guarded_gather_qmm_bf16(
+    x: &Array,
+    weights: [&Array; 3],
+    indices: [Option<&Array>; 2],
+    quant: [u32; 2],
+    sorted: bool,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    use super::{indexing, ops};
+    let shape = x.layout().shape().to_vec();
+    let pad = sorted_gather_pad(
+        &shape,
+        indices[0].map(|a| a.layout()).as_ref().map(|l| l.shape()),
+        indices[1].map(|a| a.layout()).as_ref().map(|l| l.shape()),
+        sorted,
+        super::sparse_attention::nax_available(),
+    );
+    if pad == 0 {
+        return gather_qmm_bf16(x, weights, indices, quant, sorted, streams, stream);
+    }
+    let rows = shape[0];
+    #[cfg(test)]
+    super::issue110::record(
+        rows,
+        pad,
+        *shape.last().unwrap(),
+        weights[0].layout().dim(-2)?,
+    );
+    let end = rows
+        .checked_add(pad)
+        .ok_or("sorted gather padding overflow")?;
+    let mut tail = shape.clone();
+    tail[0] = pad;
+    let x = ops::concatenate(
+        &[x.clone(), ops::zeros(&tail, x.layout().dtype(), stream)?],
+        0,
+        stream,
+    )?;
+    let rhs = indices[1].unwrap();
+    let last = ops::slice(rhs, &[rows - 1], &[rows], &[1], stream)?;
+    let rhs = ops::concatenate(
+        &[rhs.clone(), ops::broadcast_to(&last, &[pad], stream)?],
+        0,
+        stream,
+    )?;
+    let lhs = indices[0]
+        .map(|lhs| {
+            let tail = ops::astype(
+                &indexing::arange_range_i32(rows, end, stream)?,
+                lhs.layout().dtype(),
+                false,
+                stream,
+            )?;
+            ops::concatenate(&[lhs.clone(), tail], 0, stream)
+        })
+        .transpose()?;
+    let out = gather_qmm_bf16(
+        &x,
+        weights,
+        [lhs.as_ref(), Some(&rhs)],
+        quant,
+        sorted,
+        streams,
+        stream,
+    )?;
+    let mut stop = out.layout().shape().to_vec();
+    stop[0] = rows;
+    ops::slice(
+        &out,
+        &vec![0; stop.len()],
+        &stop,
+        &vec![1; stop.len()],
+        stream,
+    )
+}
+
 /// Installed affine/BF16, transposed GatherQMM from ops.cpp. Explicit index
 /// casts use the default stream, as indices_or_default does in the reference.
 pub(super) fn gather_qmm_bf16(

src/engine/metal/qwen_mtplx/moe.rs

@@ -27,7 +27,7 @@ fn gather(
     streams: &Streams,
     stream: Stream,
 ) -> Result<Array, String> {
-    linear::gather_qmm_bf16(
+    linear::guarded_gather_qmm_bf16(
         input,
         [&layer.weight, &layer.scales, &layer.biases],
         [None, Some(ids)],

src/engine/metal/qwen_mtplx/session_checkpoint.rs

@@ -32,6 +32,8 @@ struct Manifest {
     #[serde(default)]
     verify_gate: bool,
     #[serde(default)]
+    sorted_gather_guard: bool,
+    #[serde(default)]
     images: Vec<crate::engine::qwen_vision::Identity>,
     tokens: Vec<i32>,
     lazy_kv: bool,
@@ -119,6 +121,7 @@ impl LiveSession {
             let offset = file.stream_position().map_err(|e| e.to_string())?;
             let manifest = serde_json::to_vec(&Manifest {
                 verify_gate: true,
+                sorted_gather_guard: true,
                 identity,
                 tag,
                 images: self.current_images.clone(),
@@ -193,6 +196,14 @@ impl LiveSession {
                 "Qwen checkpoint uses obsolete MTP attention-gate math; rebuild the cache".into(),
             );
         }
+        // Old files do not record prefill chunk sizes or producer hardware.
+        // At ten routes/token this length can contain an affected chunk;
+        // shorter sessions cannot. Reject before allocating/restoring state.
+        if !manifest.sorted_gather_guard && manifest.tokens.len() >= 3277 {
+            return Err(
+                "Qwen checkpoint predates the sorted-gather row guard; rebuild the cache".into(),
+            );
+        }
         if !manifest.payload["mtp_history_snapshot"].is_null() != identity.mtp {
             return Err("Qwen checkpoint MTP payload differs".into());
         }
@@ -480,6 +491,40 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
             .unwrap()
     );
     assert_eq!(fingerprint(&live), ar_original);
+    let bytes = fs::read(&path).unwrap();
+    let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+    let mut manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
+    manifest
+        .as_object_mut()
+        .unwrap()
+        .remove("sorted_gather_guard");
+    let write_manifest = |manifest: &serde_json::Value| {
+        let mut old = bytes[..offset].to_vec();
+        old.extend(serde_json::to_vec(manifest).unwrap());
+        fs::write(&broken, old).unwrap();
+    };
+    write_manifest(&manifest);
+    assert!(
+        live.load_checkpoint(&broken, ar_identity, &streams, stream, |_| {}, || false)
+            .unwrap()
+    );
+    assert_eq!(fingerprint(&live), ar_original);
+    manifest["tokens"] = serde_json::json!(vec![1; 3277]);
+    manifest["identity"]["context"] = 4096.into();
+    write_manifest(&manifest);
+    let large_identity = Identity {
+        context: 4096,
+        ..ar_identity
+    };
+    assert!(
+        live.load_checkpoint(&broken, large_identity, &streams, stream, |_| {}, || false)
+            .unwrap_err()
+            .contains("sorted-gather row guard")
+    );
+    assert_eq!(fingerprint(&live), ar_original);
+    println!(
+        "Qwen sorted-gather checkpoint policy: old short state exact, possibly affected long state rejected before restore"
+    );
     println!(
         "Qwen checkpoint: exact roundtrip, cancellation/corruption controls, old MTP rejected and old AR restored"
     );