Port Qwen row-map SwiGLU expert gathers

src/engine/metal/qwen_mtplx.rs

@@ -41,11 +41,14 @@ mod issue108;
 mod issue109;
 #[cfg(test)]
 mod issue110;
+#[cfg(test)]
+mod issue112_bench;
 mod linear;
 mod load;
 mod mlp;
 mod model;
 mod moe;
+mod moe_sorted;
 mod mtp;
 mod ngram;
 mod ngram_sidecar;
@@ -1200,7 +1203,7 @@ fn gather_bf16_broadcast(
     let lhs_strides = lhs_strides.map(|d| d as i64);
     let rhs_strides = rhs_strides.map(|d| d as i64);
 
-    gather_bf16_layout(
+    gather_layout(
         [
             tensor(0, w).input(elements[0]),
             tensor(1, s).input(elements[1]),
@@ -1210,6 +1213,7 @@ fn gather_bf16_broadcast(
             tensor(5, rhs).input(elements[5]),
             tensor(6, y).output(elements[6]),
         ],
+        array::Dtype::BF16,
         [m, n, k, experts, routes, routes],
         [bits, group],
         sorted,
@@ -1232,8 +1236,10 @@ fn gather_bf16_broadcast(
 
 /// Shared original GatherQMM dispatch; array batch strides are passed verbatim,
 /// while the two broadcast index layouts have jointly collapsed dimensions.
-fn gather_bf16_layout(
+#[allow(clippy::too_many_arguments)] // Original GatherQMM ABI and array layouts.
+fn gather_layout(
     bindings: [Binding<'_>; 7],
+    dtype: array::Dtype,
     [m, n, k, experts, routes, sorted_rows]: [u32; 6],
     [bits, group]: [u32; 2],
     sorted: bool,
@@ -1248,7 +1254,12 @@ fn gather_bf16_layout(
         rhs_strides,
     ]: [&[i64]; 6],
 ) -> Result<(), String> {
-    let suffix = format!("bfloat16_t_gs_{group}_b_{bits}");
+    let ty = match dtype {
+        array::Dtype::BF16 => "bfloat16_t",
+        array::Dtype::F16 => "float16_t",
+        _ => return Err("GatherQMM dtype unsupported".into()),
+    };
+    let suffix = format!("{ty}_gs_{group}_b_{bits}");
     if m == 1 && routes >= 16 && sorted && routes / experts >= 4 {
         let (bm, bn, bk, wm, wn) = if nax {
             (

src/engine/metal/qwen_mtplx/binding.rs

@@ -438,6 +438,7 @@ impl Bank {
         policy: model::CompilePolicy,
         fused_moe_decode: bool,
         fused_moe_verify: bool,
+        sorted_gather: bool,
     ) -> Result<model::TextModel<'_>, String> {
         let mut bind = Bind {
             bank: self,
@@ -445,6 +446,7 @@ impl Bank {
             order: Vec::new(),
             fused_moe_decode,
             fused_moe_verify,
+            sorted_gather,
         };
         let embedding = bind.linear("language_model.model.embed_tokens")?;
         let layers = self
@@ -503,6 +505,7 @@ struct Bind<'a> {
     order: Vec<String>,
     fused_moe_decode: bool,
     fused_moe_verify: bool,
+    sorted_gather: bool,
 }
 
 #[test]
@@ -566,7 +569,12 @@ fn mtplx_installed_module_binding_matches_reference() {
             "no eager parameter evaluation"
         );
         let mut model = bank
-            .bind(model::CompilePolicy::from_values(None, None), true, true)
+            .bind(
+                model::CompilePolicy::from_values(None, None),
+                true,
+                true,
+                true,
+            )
             .unwrap();
         assert_eq!(case["mask"], order["mask"]);
         assert_eq!(case["mtp"], order["mtp"]);
@@ -687,7 +695,12 @@ fn mtplx_installed_module_binding_matches_reference() {
                 .all(|(_, a)| !a.has_data())
         );
         let mut owned = BoundModel::try_new(bank, |bank| {
-            bank.bind(model::CompilePolicy::from_values(None, None), true, true)
+            bank.bind(
+                model::CompilePolicy::from_values(None, None),
+                true,
+                true,
+                true,
+            )
         })
         .unwrap();
         let module_address = std::ptr::from_ref(owned.borrow_dependent()).addr();
@@ -705,19 +718,29 @@ fn mtplx_installed_module_binding_matches_reference() {
         let extra = bank.parameters["language_model.model.layers.0.linear_attn.A_log"].clone();
         bank.parameters.insert("unexpected.weight".into(), extra);
         assert!(
-            bank.bind(model::CompilePolicy::from_values(None, None), true, true)
-                .err()
-                .unwrap()
-                .contains("unbound checkpoint parameters")
+            bank.bind(
+                model::CompilePolicy::from_values(None, None),
+                true,
+                true,
+                true
+            )
+            .err()
+            .unwrap()
+            .contains("unbound checkpoint parameters")
         );
         bank.parameters.remove("unexpected.weight");
         bank.parameters
             .remove("language_model.model.layers.0.linear_attn.norm.weight");
         assert!(
-            bank.bind(model::CompilePolicy::from_values(None, None), true, true)
-                .err()
-                .unwrap()
-                .contains("missing parameter")
+            bank.bind(
+                model::CompilePolicy::from_values(None, None),
+                true,
+                true,
+                true
+            )
+            .err()
+            .unwrap()
+            .contains("missing parameter")
         );
         eprintln!(
             "binding metadata, module recipes, all-parameter coverage, MTP norms and SSD table lifecycle exact"
@@ -938,6 +961,7 @@ impl<'a> Bind<'a> {
             normalize: c.norm_topk_prob,
             fused_decode: self.fused_moe_decode,
             fused_verify: self.fused_moe_verify,
+            sorted_gather: self.sorted_gather,
         };
         let attn_hyper = self.hyper(&format!("{p}.attn_hyper_connection"), true)?;
         let mlp_hyper = self.hyper(&format!("{p}.mlp_hyper_connection"), true)?;

src/engine/metal/qwen_mtplx/compile_graph.rs

@@ -565,10 +565,10 @@ fn equivalent_operation(a: &Operation, b: &Operation, x: &Array, y: &Array) -> b
         | (AffineQuantize { bits: a, group: b }, AffineQuantize { bits: x, group: y })
         | (AffineDequantize { bits: a, group: b }, AffineDequantize { bits: x, group: y })
         | (
-            GatherQmmBf16 {
+            GatherQmm {
                 bits: a, group: b, ..
             },
-            GatherQmmBf16 {
+            GatherQmm {
                 bits: x, group: y, ..
             },
         ) => (a, b) == (x, y),

src/engine/metal/qwen_mtplx/decoder.rs

@@ -299,6 +299,7 @@ fn mtplx_decoder_connected_attention_and_moe_advance_caches() {
             shared_gate: quantized(1, 2560),
             topk: 1,
             normalize: true,
+            sorted_gather: true,
             fused_decode: false,
             fused_verify: false,
         }

src/engine/metal/qwen_mtplx/execution.rs

@@ -313,6 +313,7 @@ impl Execution {
                 ),
                 flag("MTPLX_FUSED_MOE_DECODE"),
                 flag("MTPLX_FUSED_MOE_VERIFY"),
+                super::policy::sorted_gather_enabled(get),
             )
         })?;
         if !lazy {

src/engine/metal/qwen_mtplx/issue110.rs

@@ -44,7 +44,7 @@ fn unit_gather(rows: i32, streams: &Streams) -> Array {
         .unwrap(),
     )
     .unwrap();
-    linear::guarded_gather_qmm_bf16(
+    linear::guarded_gather_qmm(
         &x,
         [&w, &s, &b],
         [None, Some(&rhs)],
@@ -141,7 +141,7 @@ fn issue110_actual_projection_rows_match_oracle() {
             )
             .unwrap()
         });
-        let out = linear::guarded_gather_qmm_bf16(
+        let out = linear::guarded_gather_qmm(
             &x,
             [&layer[0], &layer[1], &layer[2]],
             [lhs.as_ref(), Some(&rhs)],
@@ -176,7 +176,7 @@ pub(super) fn record(rows: i32, pad: i32, k: i32, n: i32) {
 
 #[test]
 #[ignore = "full installed model and large partial chunks; serial memory supervisor required"]
-fn issue110_large_cold_and_warm_prefill() {
+pub(super) fn issue110_large_cold_and_warm_prefill() {
     assert_eq!(
         std::env::var("MTPLX_GDN_BOUNDARY_TAIL_INTERVAL").as_deref(),
         Ok("0")

src/engine/metal/qwen_mtplx/issue112_bench.rs

@@ -0,0 +1,180 @@
+//! Paired, serial installed-model prefill timings; state reads follow timing.
+use super::execution::Execution;
+use super::sampling::SamplerConfig;
+use super::session_cache::{DecodeOptions, TurnProgress};
+use std::cell::Cell;
+use std::time::Instant;
+
+#[test]
+#[ignore = "full installed model; serial memory supervisor and explicit output directory"]
+fn issue112_paired_prefill() {
+    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_ISSUE112_DIR").unwrap());
+    std::fs::create_dir_all(&directory).unwrap();
+    let repeats = std::env::var("DS4_ISSUE112_REPEATS")
+        .ok()
+        .map(|n| n.parse::<usize>().unwrap())
+        .unwrap_or(3);
+    let settings = crate::settings::effective_settings(
+        crate::model::ModelChoice::Qwen38FlashNext,
+        &crate::settings::GenerationPreferences {
+            context_tokens: 16384,
+            max_generated_tokens: 1,
+            temperature: Some(0.),
+            ..Default::default()
+        },
+        &crate::settings::RuntimePreferences::default(),
+        root.parent().unwrap(),
+    )
+    .unwrap();
+    let (mut execution, mut options) = Execution::open(&settings.engine).unwrap();
+    let sampler = SamplerConfig {
+        temperature: 0.,
+        ..Default::default()
+    };
+    options.sampler = sampler;
+    options.draft_sampler = sampler;
+    options.prefill.return_hidden = false;
+    options.decode = DecodeOptions::Ar(super::ar::Options {
+        max_tokens: 1,
+        seed: 1,
+        stop_tokens: &[],
+        sync_eval: false,
+        return_hidden: false,
+        capture_final_state: true,
+    });
+    options.store_min_suffix = None;
+    let tokens = |len: usize| {
+        (0..len)
+            .map(|i| [9707, 11, 1075, 374, 264, 1092, 13][i % 7])
+            .collect::<Vec<_>>()
+    };
+    for chunk in [2048, 4096] {
+        options.prefill.chunk_size = Some(chunk);
+        for rep in 0..repeats {
+            let mut prompt = tokens(4096);
+            let mut saved = None;
+            for phase in ["cold", "warm"] {
+                let start = Instant::now();
+                let ttft = Cell::new(None);
+                let out = execution
+                    .session_turn(
+                        &prompt,
+                        saved.as_ref(),
+                        &options,
+                        |event| {
+                            if let TurnProgress::Prefill(super::prefill::Progress::Settled {
+                                tokens,
+                                total,
+                            }) = event
+                            {
+                                eprintln!("#112 {chunk}/{rep}/{phase} {tokens}/{total}");
+                            }
+                            Ok(())
+                        },
+                        |_| {
+                            if ttft.get().is_none() {
+                                ttft.set(Some(start.elapsed().as_secs_f64()));
+                            }
+                            Ok(())
+                        },
+                        |_| Ok(()),
+                    )
+                    .unwrap();
+                let elapsed = start.elapsed().as_secs_f64();
+                assert!(out.final_snapshot_error.is_none());
+                let snapshot = out.snapshot.unwrap();
+                let state = super::issue104::state(&execution, &snapshot);
+                let hidden_state=snapshot.hidden.as_ref().map(|a| {
+                    let values=execution.values_f32(a).unwrap();
+                    assert!(values.iter().all(|v|v.is_finite()));
+                    serde_json::json!({"shape":a.layout().shape(),"sha256":super::super::qwen_mtplx_tests::hex_digest(&values.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>())})
+                });
+                assert!(hidden_state.is_none());
+                let metrics = serde_json::json!({"chunk":chunk,"rep":rep,"phase":phase,"prompt":prompt.len(),"cached":out.cached_tokens,"ttft_seconds":ttft.get().unwrap(),"prefill_seconds":out.prefill_seconds,"prefill_tokens_per_second":(prompt.len()-out.cached_tokens) as f64/out.prefill_seconds,"elapsed_seconds":elapsed});
+                let record =
+                    serde_json::json!({"metrics":metrics,"state":state,"hidden":hidden_state});
+                std::fs::write(
+                    directory.join(format!("{chunk}-{rep}-{phase}.json")),
+                    serde_json::to_vec(&record).unwrap(),
+                )
+                .unwrap();
+                println!(
+                    "{}",
+                    serde_json::json!({"case":"paired_prefill","metrics":metrics})
+                );
+                prompt = snapshot.token_ids.clone();
+                prompt.extend(tokens(2048));
+                saved = Some(snapshot);
+            }
+        }
+    }
+}
+
+#[test]
+#[ignore = "full installed model hidden-state prefill; serial supervisor required"]
+fn issue112_hidden_prefill() {
+    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_ISSUE112_DIR").unwrap());
+    std::fs::create_dir_all(&directory).unwrap();
+    let mut settings = crate::settings::effective_settings(
+        crate::model::ModelChoice::Qwen38FlashNext,
+        &crate::settings::GenerationPreferences {
+            context_tokens: 16384,
+            max_generated_tokens: 1,
+            ..Default::default()
+        },
+        &crate::settings::RuntimePreferences::default(),
+        root.parent().unwrap(),
+    )
+    .unwrap();
+    settings.engine.speculative.glm_mtp = true;
+    let (mut execution, mut options) = Execution::open(&settings.engine).unwrap();
+    options.decode = DecodeOptions::Ar(super::ar::Options {
+        max_tokens: 1,
+        seed: 1,
+        stop_tokens: &[],
+        sync_eval: false,
+        return_hidden: true,
+        capture_final_state: true,
+    });
+    options.prefill.return_hidden = true;
+    options.store_min_suffix = None;
+    let tokens = |len: usize| {
+        (0..len)
+            .map(|i| [9707, 11, 1075, 374, 264, 1092, 13][i % 7])
+            .collect::<Vec<_>>()
+    };
+    for chunk in [2048, 4096] {
+        options.prefill.chunk_size = Some(chunk);
+        let mut saved = None;
+        let mut prompt = tokens(4096);
+        for phase in ["cold", "warm"] {
+            let snapshot = execution
+                .prefill_session(&prompt, saved.as_ref(), &options, |event| {
+                    if let super::prefill::Progress::Settled { tokens, total } = event {
+                        eprintln!("#112 hidden {chunk}/{phase} {tokens}/{total}");
+                    }
+                    Ok(())
+                })
+                .unwrap();
+            let state = super::issue104::state(&execution, &snapshot);
+            let hidden = snapshot.hidden.as_ref().unwrap();
+            let values = execution.values_f32(hidden).unwrap();
+            assert!(values.iter().all(|v| v.is_finite()));
+            let hidden = serde_json::json!({"shape":hidden.layout().shape(),"sha256":super::super::qwen_mtplx_tests::hex_digest(&values.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>())});
+            let record = serde_json::json!({"case":"hidden_prefill","chunk":chunk,"phase":phase,"cached":saved.as_ref().map_or(0,|s|s.token_ids.len()),"state":state,"hidden":hidden});
+            std::fs::write(
+                directory.join(format!("{chunk}-{phase}.json")),
+                serde_json::to_vec(&record).unwrap(),
+            )
+            .unwrap();
+            println!(
+                "{}",
+                serde_json::json!({"case":"hidden_prefill","chunk":chunk,"phase":phase,"hidden":hidden})
+            );
+            prompt.extend(tokens(2048));
+            saved = Some(snapshot);
+        }
+    }
+}

src/engine/metal/qwen_mtplx/linear.rs

@@ -371,7 +371,7 @@ fn sorted_gather_guard_matches_pinned_call_predicates() {
 
 /// 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(
+pub(super) fn guarded_gather_qmm(
     x: &Array,
     weights: [&Array; 3],
     indices: [Option<&Array>; 2],
@@ -390,7 +390,7 @@ pub(super) fn guarded_gather_qmm_bf16(
         super::sparse_attention::nax_available(),
     );
     if pad == 0 {
-        return gather_qmm_bf16(x, weights, indices, quant, sorted, streams, stream);
+        return gather_qmm(x, weights, indices, quant, sorted, streams, stream);
     }
     let rows = shape[0];
     #[cfg(test)]
@@ -428,7 +428,7 @@ pub(super) fn guarded_gather_qmm_bf16(
             ops::concatenate(&[lhs.clone(), tail], 0, stream)
         })
         .transpose()?;
-    let out = gather_qmm_bf16(
+    let out = gather_qmm(
         &x,
         weights,
         [lhs.as_ref(), Some(&rhs)],
@@ -448,9 +448,9 @@ pub(super) fn guarded_gather_qmm_bf16(
     )
 }
 
-/// Installed affine/BF16, transposed GatherQMM from ops.cpp. Explicit index
+/// Affine BF16/F16 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(
+pub(super) fn gather_qmm(
     x: &Array,
     [w, scales, biases]: [&Array; 3],
     indices: [Option<&Array>; 2],
@@ -466,17 +466,19 @@ pub(super) fn gather_qmm_bf16(
     let wl = w.layout();
     let sl = scales.layout();
     if stream.device() != Device::Gpu
-        || xl.dtype() != Dtype::BF16
+        || !matches!(xl.dtype(), Dtype::BF16 | Dtype::F16)
         || wl.dtype() != Dtype::U32
-        || sl.dtype() != Dtype::BF16
-        || biases.layout().dtype() != Dtype::BF16
+        || sl.dtype() != xl.dtype()
+        || biases.layout().dtype() != xl.dtype()
         || xl.shape().len() < 2
         || wl.shape().len() < 2
         || wl.shape().len() != sl.shape().len()
         || !matches!(bits, 4 | 8)
-        || !matches!(group, 32 | 64)
+        || !matches!(group, 32 | 64 | 128)
     {
-        return Err("GatherQMM requires BF16 GPU matrices and U32 affine expert weights".into());
+        return Err(
+            "GatherQMM requires BF16/F16 GPU matrices and U32 affine expert weights".into(),
+        );
     }
     let wrank = wl.shape().len();
     let k = *xl.shape().last().unwrap() as u32;
@@ -491,7 +493,9 @@ pub(super) fn gather_qmm_bf16(
     {
         return Err("GatherQMM input/weight/group shapes do not match".into());
     }
-    affine_plan([m, n, k], bits, group, true)?;
+    if [m, n, k].iter().any(|&d| d == 0 || d > i32::MAX as u32) || !k.is_multiple_of(group) {
+        return Err("invalid GatherQMM shape".into());
+    }
     let index = |a: &Array, indices: Option<&Array>| -> Result<Array, String> {
         if let Some(indices) = indices {
             if !matches!(
@@ -530,7 +534,7 @@ pub(super) fn gather_qmm_bf16(
     )?);
     let mut shape = inputs[4].layout().shape().to_vec();
     shape.extend([m as i32, n as i32]);
-    let layout = Layout::new(&shape, Dtype::BF16)?;
+    let layout = Layout::new(&shape, xl.dtype())?;
     if layout.size() > i32::MAX as usize {
         return Err("GatherQMM output size overflow".into());
     }
@@ -538,7 +542,7 @@ pub(super) fn gather_qmm_bf16(
         stream,
         &inputs,
         layout,
-        Operation::GatherQmmBf16 {
+        Operation::GatherQmm {
             bits,
             group,
             right_sorted: sorted && indices[0].is_none(),
@@ -575,7 +579,7 @@ pub(super) fn gather_evaluate(
         if prepared[0].layout().size() / k as usize != prepared[5].layout().size() {
             let mut shape = prepared[5].layout().shape().to_vec();
             shape.extend([m as i32, k as i32]);
-            let broadcast = Array::unallocated(&shape, Dtype::BF16)?;
+            let broadcast = Array::unallocated(&shape, prepared[0].layout().dtype())?;
             views::broadcast(&prepared[0], &broadcast)?;
             prepared[0] = broadcast;
         }
@@ -596,7 +600,7 @@ pub(super) fn gather_evaluate(
         views::collapse_joint(ll.shape(), &[ll.strides(), rl.strides()]);
     let arrays = [w, s, b, x, lhs, rhs, output];
     let buffers = arrays.map(Array::buffer);
-    super::gather_bf16_layout(
+    super::gather_layout(
         std::array::from_fn(|i| {
             let binding = super::tensor(i as u32, &buffers[i]).at_byte_offset(arrays[i].offset());
             if i == 6 {
@@ -605,6 +609,7 @@ pub(super) fn gather_evaluate(
                 binding.input(arrays[i].data_size())
             }
         }),
+        xl.dtype(),
         [m, n, k, experts, batches, sorted_rows],
         [bits, group],
         right_sorted,
@@ -1058,7 +1063,7 @@ fn mtplx_graph_gather_qmm_matches_reference() {
                     .unwrap();
                 }
                 let (output, construction) = capture_dispatches(|| {
-                    gather_qmm_bf16(
+                    gather_qmm(
                         &x,
                         [&w, &s, &b],
                         [left.then_some(&lhs), right.then_some(&rhs)],

src/engine/metal/qwen_mtplx/model.rs

@@ -721,7 +721,7 @@ impl TextModel<'_> {
                         O::RmsNorm(_) => "RMSNorm",
                         O::Softmax(_) => "Softmax",
                         O::QuantizedLinearBf16 { .. } => "QuantizedMatmul",
-                        O::GatherQmmBf16 { .. } => "GatherQMM",
+                        O::GatherQmm { .. } => "GatherQMM",
                         O::DenseMatmul => "Matmul",
                         O::Sum(_) => "Sum",
                         O::Sort {

src/engine/metal/qwen_mtplx/moe.rs

@@ -27,7 +27,7 @@ fn gather(
     streams: &Streams,
     stream: Stream,
 ) -> Result<Array, String> {
-    linear::guarded_gather_qmm_bf16(
+    linear::guarded_gather_qmm(
         input,
         [&layer.weight, &layer.scales, &layer.biases],
         [None, Some(ids)],
@@ -55,13 +55,19 @@ impl FeedForward {
         &self,
         input: &Array,
         indices: &Array,
+        sorted_gather: bool,
         streams: &Streams,
         stream: Stream,
     ) -> Result<Array, String> {
-        let mut x = ops::expand_dims(input, &[-2, -3], stream)?;
         let do_sort = indices.layout().size() >= 64;
+        let mut x = if do_sort {
+            input.clone()
+        } else {
+            ops::expand_dims(input, &[-2, -3], stream)?
+        };
         let mut ids = indices.clone();
         let mut inverse = None;
+        let mut row_map = None;
         if do_sort {
             let size =
                 i32::try_from(indices.layout().size()).map_err(|_| "MoE index size overflow")?;
@@ -76,33 +82,46 @@ impl FeedForward {
             let k = *input.layout().shape().last().ok_or("MoE input rank")?;
             let flat_rows = i32::try_from(input.layout().size() / k as usize)
                 .map_err(|_| "MoE row count overflow")?;
-            x = indexing::take(
-                &ops::reshape(&x, &[flat_rows, 1, k], stream)?,
-                &rows,
-                0,
-                stream,
-            )?;
+            x = ops::reshape(&x, &[flat_rows, 1, k], stream)?;
+            row_map = Some(rows);
             ids = indexing::take(&ids, &order, 0, stream)?;
         }
-        let hidden = match &self.input {
-            GateUp::Fused(gu) => {
-                let gu = gather(&x, gu, &ids, do_sort, streams, stream)?;
-                let n = *gu.layout().shape().last().ok_or("MoE gate/up rank")?;
-                if n % 2 != 0 {
-                    return Err("MoE fused gate/up width must be even".into());
+        let hidden = if let Some(map) = row_map {
+            let (gate, up) = match &self.input {
+                GateUp::Fused(gu) => (gu, None),
+                GateUp::Separate { gate, up } => (gate, Some(up)),
+            };
+            super::moe_sorted::swiglu_rows(
+                &x,
+                &map,
+                gate,
+                up,
+                &ids,
+                sorted_gather,
+                streams,
+                stream,
+            )?
+        } else {
+            match &self.input {
+                GateUp::Fused(gu) => {
+                    let gu = gather(&x, gu, &ids, do_sort, streams, stream)?;
+                    let n = *gu.layout().shape().last().ok_or("MoE gate/up rank")?;
+                    if n % 2 != 0 {
+                        return Err("MoE fused gate/up width must be even".into());
+                    }
+                    let pair = ops::split(&gu, &[n / 2], -1, stream)?;
+                    binary(
+                        &mlp::silu(&pair[0], stream)?,
+                        &pair[1],
+                        Binary::Multiply,
+                        stream,
+                    )?
+                }
+                GateUp::Separate { gate, up } => {
+                    let up = gather(&x, up, &ids, do_sort, streams, stream)?;
+                    let gate = gather(&x, gate, &ids, do_sort, streams, stream)?;
+                    mlp::swiglu(&gate, &up, stream)?
                 }
-                let pair = ops::split(&gu, &[n / 2], -1, stream)?;
-                binary(
-                    &mlp::silu(&pair[0], stream)?,
-                    &pair[1],
-                    Binary::Multiply,
-                    stream,
-                )?
-            }
-            GateUp::Separate { gate, up } => {
-                let up = gather(&x, up, &ids, do_sort, streams, stream)?;
-                let gate = gather(&x, gate, &ids, do_sort, streams, stream)?;
-                mlp::swiglu(&gate, &up, stream)?
             }
         };
         let mut output = gather(&hidden, &self.down, &ids, do_sort, streams, stream)?;
@@ -123,6 +142,7 @@ pub(super) struct SparseMoe {
     pub(super) shared_gate: QuantizedLinear,
     pub(super) topk: i32,
     pub(super) normalize: bool,
+    pub(super) sorted_gather: bool,
     pub(super) fused_decode: bool,
     pub(super) fused_verify: bool,
 }
@@ -199,7 +219,9 @@ impl SparseMoe {
                 stream,
             )?
         } else {
-            let switched = self.experts.switched(input, &indices, streams, stream)?;
+            let switched =
+                self.experts
+                    .switched(input, &indices, self.sorted_gather, streams, stream)?;
             let scores = ops::expand_dims(&scores, &[-1], stream)?;
             reduce::sum(
                 &binary(&switched, &scores, Binary::Multiply, stream)?,
@@ -413,6 +435,7 @@ fn mtplx_graph_sparse_moe_matches_reference() {
                 shared_gate: mlp::fixture_linear(1, k, 4, 64, 52),
                 topk: 10,
                 normalize: row["normalize"].as_bool().unwrap(),
+                sorted_gather: true,
                 fused_decode: row["fused"].as_bool().unwrap_or(false),
                 fused_verify: row["fused"].as_bool().unwrap_or(false),
             };

src/engine/metal/qwen_mtplx/moe_sorted.rs

@@ -0,0 +1,839 @@
+//! MTPLX 9882703 expert row-map gather, including whole-first-call canaries.
+use super::array::{Array, Dtype, Layout};
+use super::binary::{Binary, binary, floor_divide_integer};
+use super::mlp::{self, QuantizedLinear};
+use super::stream::{Stream, Streams};
+use super::{indexing, linear, ops};
+use std::collections::BTreeMap;
+use std::sync::{LazyLock, Mutex};
+
+#[derive(Default, Clone, Copy, Debug)]
+pub(super) struct Stats {
+    pub(super) calls: u64,
+    pub(super) fallbacks: u64,
+    pub(super) canaries: u64,
+    pub(super) failures: u64,
+}
+static CANARY: LazyLock<Mutex<(BTreeMap<String, bool>, Stats)>> =
+    LazyLock::new(|| Mutex::new((BTreeMap::new(), Stats::default())));
+
+#[cfg(test)]
+pub(super) fn stats() -> Stats {
+    CANARY.lock().unwrap().1
+}
+
+#[cfg(test)]
+static FAILURE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
+
+#[derive(Clone, Copy, Debug)]
+pub(super) struct Kernel {
+    dtype: Dtype,
+    bits: u32,
+    group: u32,
+    k: i32,
+    n: i32,
+    experts: i32,
+    swiglu: bool,
+    separate: bool,
+}
+
+impl Kernel {
+    fn width(self) -> i32 {
+        if self.swiglu && !self.separate {
+            self.n / 2
+        } else {
+            self.n
+        }
+    }
+    fn name(self) -> String {
+        format!(
+            "kernel_qwen_mtplx_sorted_{}_g{}_b{}_k{}_n{}_e{}_s{}_u{}",
+            self.dtype.kernel_name(),
+            self.group,
+            self.bits,
+            self.k,
+            self.n,
+            self.experts,
+            u8::from(self.swiglu),
+            u8::from(self.separate)
+        )
+    }
+    fn source(self) -> String {
+        let ty = if self.dtype == Dtype::BF16 {
+            "bfloat16_t"
+        } else {
+            "float16_t"
+        };
+        let params = [
+            ("x", ty),
+            ("w", "uint"),
+            ("scales", ty),
+            ("biases", ty),
+            ("wu", "uint"),
+            ("scales_u", ty),
+            ("biases_u", ty),
+            ("tile_start", "int"),
+            ("row_start", "int"),
+            ("row_map", "uint"),
+        ]
+        .iter()
+        .enumerate()
+        .map(|(i, (name, ty))| format!("const device {ty}* {name} [[buffer({i})]]"))
+        .collect::<Vec<_>>()
+        .join(",\n");
+        format!(
+            "{}\nusing T={ty};\nconstant int GS={}, BITS={}, KD={}, ND={}, UD={}, UOFF={}, OUTC={}, ED={};\nconstant bool SWIGLU={};\nkernel void {}({params}, device T* y [[buffer(10)]], uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], uint simdgroup_index_in_threadgroup [[simdgroup_index_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) {{\n{}\n}}",
+            include_str!("../../../../metal/mtplx-moe-sorted-header.metal"),
+            self.group,
+            self.bits,
+            self.k,
+            self.n,
+            self.n,
+            if self.swiglu && !self.separate {
+                self.n / 2
+            } else {
+                0
+            },
+            self.width(),
+            self.experts,
+            self.swiglu,
+            self.name(),
+            include_str!("../../../../metal/mtplx-moe-sorted-body.metal")
+        )
+    }
+    pub(super) fn evaluate(self, inputs: &[Array], outputs: &[Array]) -> Result<(), String> {
+        super::compile_custom(&self.name(), || {
+            #[cfg(test)]
+            if FAILURE.load(std::sync::atomic::Ordering::Relaxed) == 1 {
+                return Ok("injected invalid Metal source".into());
+            }
+            Ok(self.source())
+        })?;
+        let rows = outputs[0].layout().shape()[0] as u32;
+        let cols = if self.swiglu {
+            self.width() * 2
+        } else {
+            self.n
+        } as u32;
+        super::gdn::custom_evaluate(inputs, outputs, |bindings| {
+            super::dispatch_geometry(
+                &self.name(),
+                bindings,
+                &[],
+                [
+                    (cols / 64) * 128,
+                    rows.div_ceil(64) + self.experts as u32,
+                    1,
+                ],
+                [128, 1, 1],
+                true,
+            )
+        })
+    }
+}
+
+// Applies mirrors the affine tensor/layout predicates. QuantizedLinear always
+// carries biases and affine mode; separate packs must also share quantization.
+pub(super) fn applies(
+    tokens: &Array,
+    map: &Array,
+    gate: &QuantizedLinear,
+    up: Option<&QuantizedLinear>,
+    ids: &Array,
+    enabled: bool,
+    nax: bool,
+) -> bool {
+    if let Some(up) = up
+        && (gate.bits != up.bits
+            || gate.group != up.group
+            || [&gate.weight, &gate.scales, &gate.biases]
+                .iter()
+                .zip([&up.weight, &up.scales, &up.biases])
+                .any(|(a, b)| {
+                    a.layout().shape() != b.layout().shape()
+                        || a.layout().dtype() != b.layout().dtype()
+                }))
+    {
+        return false;
+    }
+    let x = tokens.layout();
+    let w = gate.weight.layout();
+    let s = gate.scales.layout();
+    let b = gate.biases.layout();
+    let r = map.layout();
+    let i = ids.layout();
+    enabled
+        && nax
+        && matches!(gate.bits, 4 | 8)
+        && matches!(gate.group, 32 | 64 | 128)
+        && matches!(x.dtype(), Dtype::BF16 | Dtype::F16)
+        && x.shape().len() == 3
+        && x.shape()[1] == 1
+        && s.dtype() == x.dtype()
+        && b.dtype() == x.dtype()
+        && i.shape().len() == 1
+        && r.shape() == i.shape()
+        && i.size() >= 4096
+        && r.dtype() == Dtype::U32
+        && w.shape().len() == 3
+        && w.dtype() == Dtype::U32
+        && w.shape()[1] % 64 == 0
+        && (w.shape()[2] * 32 / gate.bits as i32) % 64 == 0
+        && x.shape()[2] == w.shape()[2] * 32 / gate.bits as i32
+        && x.shape()[2] % gate.group as i32 == 0
+}
+
+fn scalar(n: i32) -> Result<Array, String> {
+    Array::new(&[], Dtype::I32, super::scalar_buffer(&n.to_le_bytes())?)
+}
+fn schedule(ids: &Array, experts: i32, stream: Stream) -> Result<(Array, Array), String> {
+    if !matches!(
+        ids.layout().dtype(),
+        Dtype::U8
+            | Dtype::U16
+            | Dtype::U32
+            | Dtype::U64
+            | Dtype::I8
+            | Dtype::I16
+            | Dtype::I32
+            | Dtype::I64
+    ) {
+        return Err("MoE expert indices must be integral".into());
+    }
+    let ids = ops::contiguous(&ops::astype(ids, Dtype::U32, false, stream)?, false, stream)?;
+    let zero = ops::zeros(&[experts], Dtype::I32, stream)?;
+    let counts = Array::make_operation_with_inputs(
+        stream,
+        &[zero, ids.clone()],
+        Layout::new(&[experts], Dtype::I32)?,
+        ops::Operation::MoeSortedCounts,
+    )?;
+    let row_start = ops::concatenate(
+        &[
+            ops::zeros(&[1], Dtype::I32, stream)?,
+            indexing::cumsum_last_i32(&counts, stream)?,
+        ],
+        0,
+        stream,
+    )?;
+    let tiles = floor_divide_integer(
+        &binary(&counts, &scalar(63)?, Binary::Add, stream)?,
+        &scalar(64)?,
+        stream,
+    )?;
+    let tile_start = ops::concatenate(
+        &[
+            ops::zeros(&[1], Dtype::I32, stream)?,
+            indexing::cumsum_last_i32(&tiles, stream)?,
+        ],
+        0,
+        stream,
+    )?;
+    Ok((tile_start, row_start))
+}
+
+pub(super) fn counts_evaluate(inputs: &[Array], output: &Array) -> Result<(), String> {
+    // Same zeros.at[ids].add(1) schedule, with integer atomic sums on the GPU.
+    // The pinned library does not export this integer scatter specialization.
+    const NAME: &str = "kernel_qwen_mtplx_sorted_counts";
+    super::compile_custom(NAME, || {
+        Ok("#include <metal_stdlib>\nusing namespace metal;\nkernel void kernel_qwen_mtplx_sorted_counts(const device uint* ids [[buffer(0)]], device atomic_int* counts [[buffer(1)]], constant uint& size [[buffer(2)]],uint i [[thread_position_in_grid]]) { if(i<size) atomic_fetch_add_explicit(counts+ids[i],1,memory_order_relaxed); }".into())
+    })?;
+    super::views::copy_gpu(&inputs[0], output, true)?;
+    let ids = &inputs[1];
+    let n = ids.layout().size() as u32;
+    super::dispatch_geometry(
+        NAME,
+        &[
+            super::tensor(0, &ids.buffer())
+                .at_byte_offset(ids.offset())
+                .input(ids.data_size()),
+            super::tensor(1, &output.buffer()).output(output.data_size()),
+            super::bytes(2, &n),
+        ],
+        &[],
+        [n, 1, 1],
+        [n.min(1024), 1, 1],
+        true,
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+fn launch(
+    tokens: &Array,
+    map: &Array,
+    gate: &QuantizedLinear,
+    up: Option<&QuantizedLinear>,
+    ids: &Array,
+    kernel: Kernel,
+    stream: Stream,
+) -> Result<Array, String> {
+    let (tiles, rows) = schedule(ids, kernel.experts, stream)?;
+    let u = up.unwrap_or(gate);
+    let input = ops::reshape(tokens, &[-1, kernel.k], stream)?;
+    let y = Array::make_operation_with_inputs(
+        stream,
+        &[
+            input,
+            gate.weight.clone(),
+            gate.scales.clone(),
+            gate.biases.clone(),
+            u.weight.clone(),
+            u.scales.clone(),
+            u.biases.clone(),
+            tiles,
+            rows,
+            map.clone(),
+        ],
+        Layout::new(&[ids.layout().shape()[0], kernel.width()], kernel.dtype)?,
+        ops::Operation::MoeSortedGather(kernel),
+    )?;
+    ops::reshape(&y, &[ids.layout().shape()[0], 1, kernel.width()], stream)
+}
+
+fn gather(
+    x: &Array,
+    q: &QuantizedLinear,
+    ids: &Array,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    linear::guarded_gather_qmm(
+        x,
+        [&q.weight, &q.scales, &q.biases],
+        [None, Some(ids)],
+        [q.bits, q.group],
+        true,
+        streams,
+        stream,
+    )
+}
+
+fn activation(gate: &Array, up: &Array, separate: bool, stream: Stream) -> Result<Array, String> {
+    if gate.layout().dtype() == Dtype::BF16 {
+        return if separate {
+            mlp::swiglu(gate, up, stream)
+        } else {
+            binary(&mlp::silu(gate, stream)?, up, Binary::Multiply, stream)
+        };
+    }
+    let graph = super::compiled::Replay::trace_compiled(&[gate.clone(), up.clone()], |a| {
+        let sig = super::unary::unary(&a[0], super::unary::Unary::Sigmoid, stream)?;
+        let silu = binary(&a[0], &sig, Binary::Multiply, stream)?;
+        Ok(vec![if separate {
+            binary(&silu, &a[1], Binary::Multiply, stream)?
+        } else {
+            silu
+        }])
+    })?;
+    let result = graph.run(&[gate.clone(), up.clone()])?.remove(0);
+    if separate {
+        Ok(result)
+    } else {
+        binary(&result, up, Binary::Multiply, stream)
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn stock(
+    tokens: &Array,
+    map: &Array,
+    gate: &QuantizedLinear,
+    up: Option<&QuantizedLinear>,
+    ids: &Array,
+    swiglu: bool,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    let x = indexing::take(tokens, map, 0, stream)?;
+    let y = gather(&x, gate, ids, streams, stream)?;
+    if !swiglu {
+        return Ok(y);
+    }
+    if let Some(up) = up {
+        activation(&y, &gather(&x, up, ids, streams, stream)?, true, stream)
+    } else {
+        let pair = ops::split(&y, &[y.layout().dim(-1)? / 2], -1, stream)?;
+        activation(&pair[0], &pair[1], false, stream)
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn checked(
+    tokens: &Array,
+    map: &Array,
+    gate: &QuantizedLinear,
+    up: Option<&QuantizedLinear>,
+    ids: &Array,
+    swiglu: bool,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Option<Array>, String> {
+    let kernel = Kernel {
+        dtype: tokens.layout().dtype(),
+        bits: gate.bits,
+        group: gate.group,
+        k: tokens.layout().dim(-1)?,
+        n: gate.weight.layout().dim(-2)?,
+        experts: gate.weight.layout().shape()[0],
+        swiglu,
+        separate: up.is_some(),
+    };
+    let key = format!(
+        "{}_{}",
+        kernel.name(),
+        if ids.layout().size() >= 64 * kernel.experts as usize {
+            "wide"
+        } else {
+            "narrow"
+        }
+    );
+    let mut state = CANARY.lock().map_err(|_| "MoE canary lock poisoned")?;
+    if state.0.get(&key) == Some(&false) {
+        state.1.fallbacks += 1;
+        return Ok(None);
+    }
+    if state.0.get(&key) == Some(&true) {
+        state.1.calls += 1;
+        drop(state);
+        return launch(tokens, map, gate, up, ids, kernel, stream).map(Some);
+    }
+    state.1.canaries += 1;
+    let probe: Result<Array, String> = (|| {
+        let ours = launch(tokens, map, gate, up, ids, kernel, stream)?;
+        let reference = stock(tokens, map, gate, up, ids, swiglu, streams, stream)?;
+        ops::evaluate(streams, &[ours.clone(), reference.clone()], stream, false)?;
+        #[cfg(test)]
+        if FAILURE.load(std::sync::atomic::Ordering::Relaxed) == 2 {
+            let mut bits = [0u8; 2];
+            ours.buffer().read(ours.offset(), &mut bits)?;
+            bits[1] ^= 0x80;
+            ours.buffer().write(ours.offset(), &bits)?;
+        }
+        #[cfg(test)]
+        if FAILURE.load(std::sync::atomic::Ordering::Relaxed) == 3 {
+            ours.buffer().write(ours.offset(), &[0, 128])?;
+            reference.buffer().write(reference.offset(), &[0, 0])?;
+        }
+        let equal = binary(
+            &ops::view(&ours, Dtype::U16, stream)?,
+            &ops::view(&reference, Dtype::U16, stream)?,
+            Binary::Equal,
+            stream,
+        )?;
+        let different = super::unary::unary(&equal, super::unary::Unary::LogicalNot, stream)?;
+        let mismatch = super::reduce::any(&different, &[0, 1, 2], false, stream)?;
+        ops::evaluate(streams, std::slice::from_ref(&mismatch), stream, false)?;
+        let mut failed = [0u8];
+        mismatch.buffer().read(mismatch.offset(), &mut failed)?;
+        if failed[0] != 0 {
+            return Err(
+                "first whole-call output differs bit for bit from the stock sorted gather".into(),
+            );
+        }
+        Ok(ours)
+    })();
+    match probe {
+        Ok(ours) => {
+            state.0.insert(key.clone(), true);
+            state.1.calls += 1;
+            eprintln!(
+                "[moe-sorted-gather] tensor-unit kernel on for {key}: first call equals stock bit for bit"
+            );
+            Ok(Some(ours))
+        }
+        Err(reason) => {
+            state.0.insert(key.clone(), false);
+            state.1.failures += 1;
+            state.1.fallbacks += 1;
+            eprintln!(
+                "[moe-sorted-gather] tensor-unit kernel off for {key}: {reason}; stock sorted gather runs instead (canary_failures={})",
+                state.1.failures
+            );
+            Ok(None)
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+pub(super) fn swiglu_rows(
+    tokens: &Array,
+    map: &Array,
+    gate: &QuantizedLinear,
+    up: Option<&QuantizedLinear>,
+    ids: &Array,
+    enabled: bool,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    let eligible = applies(
+        tokens,
+        map,
+        gate,
+        up,
+        ids,
+        enabled,
+        super::sparse_attention::nax_available(),
+    );
+    if eligible && let Some(y) = checked(tokens, map, gate, up, ids, true, streams, stream)? {
+        return Ok(y);
+    }
+    // The fused-weight wrapper retries its plain row-map gather if the SwiGLU
+    // canary failed. Separate SwitchGLU instead returns to its original chain.
+    if eligible
+        && up.is_none()
+        && let Some(gu) = checked(tokens, map, gate, None, ids, false, streams, stream)?
+    {
+        let pair = ops::split(&gu, &[gu.layout().dim(-1)? / 2], -1, stream)?;
+        return activation(&pair[0], &pair[1], false, stream);
+    }
+    if let Some(up) = up {
+        let x = indexing::take(tokens, map, 0, stream)?;
+        let up = gather(&x, up, ids, streams, stream)?;
+        let gate = gather(&x, gate, ids, streams, stream)?;
+        return activation(&gate, &up, true, stream);
+    }
+    stock(tokens, map, gate, None, ids, true, streams, stream)
+}
+
+#[test]
+#[ignore = "Metal synthetic all supported row-map variants; external supervisor"]
+fn issue112_all_kernel_variants() {
+    use super::super::{
+        gpu::Context,
+        qwen_mtplx_tests::{hex_digest, pattern},
+    };
+    super::configure_sources().unwrap();
+    let _context = Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(super::stream::Device::Gpu).unwrap();
+    let e = 128;
+    let n = 256;
+    let k = 256;
+    let tokens = 512;
+    let read = |a: &Array| {
+        ops::evaluate(&streams, std::slice::from_ref(a), gpu, false).unwrap();
+        let mut v = vec![0; a.layout().nbytes()];
+        a.buffer().read(a.offset(), &mut v).unwrap();
+        v
+    };
+    let u32s = |v: Vec<u32>| {
+        Array::new(
+            &[v.len() as i32],
+            Dtype::U32,
+            super::scalar_buffer(&v.into_iter().flat_map(u32::to_le_bytes).collect::<Vec<_>>())
+                .unwrap(),
+        )
+        .unwrap()
+    };
+    for dtype in [Dtype::BF16, Dtype::F16] {
+        for bits in [4, 8] {
+            for group in [32, 64, 128] {
+                let pack = |salt| {
+                    let mut q = mlp::fixture_linear(e * n, k, bits, group, salt);
+                    for a in [&mut q.weight, &mut q.scales, &mut q.biases] {
+                        let width = a.layout().dim(-1).unwrap();
+                        *a = ops::reshape(a, &[e as i32, n as i32, width], gpu).unwrap();
+                        if a.layout().dtype() == Dtype::BF16 {
+                            *a = ops::astype(a, dtype, false, gpu).unwrap();
+                        }
+                    }
+                    q
+                };
+                let gate = pack(7);
+                let up = pack(9);
+                let x = ops::astype(
+                    &Array::new(
+                        &[tokens, 1, k as i32],
+                        Dtype::BF16,
+                        pattern(tokens as u32 * k, 8),
+                    )
+                    .unwrap(),
+                    dtype,
+                    false,
+                    gpu,
+                )
+                .unwrap();
+                for separate in [false, true] {
+                    for rows in [4096, 8193, 32769] {
+                        eprintln!(
+                            "row-map {dtype:?} b{bits}/g{group} separate={separate} rows={rows}"
+                        );
+                        let ids = u32s(
+                            (0..rows)
+                                .map(|r| (r * (e - 3) / rows) / 2 * 2 + 1)
+                                .collect(),
+                        );
+                        let map = u32s((0..rows).map(|r| (r * 13) % tokens as u32).collect());
+                        let ids = if rows == 32769 {
+                            ops::astype(&ids, Dtype::I64, false, gpu).unwrap()
+                        } else {
+                            ids
+                        };
+                        let map = if rows == 8193 {
+                            let repeated = ops::repeat(&map, 2, 0, gpu).unwrap();
+                            ops::slice(&repeated, &[0], &[(rows * 2) as i32], &[2], gpu).unwrap()
+                        } else {
+                            map
+                        };
+                        let up = separate.then_some(&up);
+                        assert!(applies(&x, &map, &gate, up, &ids, true, true));
+                        assert!(!applies(&x, &map, &gate, up, &ids, false, true));
+                        assert!(!applies(&x, &map, &gate, up, &ids, true, false));
+                        let before = stats();
+                        let y =
+                            swiglu_rows(&x, &map, &gate, up, &ids, true, &streams, gpu).unwrap();
+                        let got = read(&y);
+                        let reference =
+                            read(&stock(&x, &map, &gate, up, &ids, true, &streams, gpu).unwrap());
+                        assert_eq!(got, reference);
+                        let after = stats();
+                        assert_eq!(after.failures, before.failures, "{after:?}");
+                        assert_eq!(after.calls, before.calls + 1);
+                        println!(
+                            "{}",
+                            serde_json::json!({"case":"synthetic","dtype":format!("{dtype:?}"),"bits":bits,"group":group,"separate":separate,"rows":rows,"sha256":hex_digest(&got),"canaries":after.canaries,"calls":after.calls})
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[test]
+#[ignore = "Metal canary rejection and sticky fallback; external supervisor"]
+fn issue112_canary_failures() {
+    use super::super::{gpu::Context, qwen_mtplx_tests::pattern};
+    super::configure_sources().unwrap();
+    let _context = Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(super::stream::Device::Gpu).unwrap();
+    let rows = 4096;
+    let e = 16;
+    let n = 128;
+    let k = 256;
+    let u32s = |v: Vec<u32>| {
+        Array::new(
+            &[v.len() as i32],
+            Dtype::U32,
+            super::scalar_buffer(&v.into_iter().flat_map(u32::to_le_bytes).collect::<Vec<_>>())
+                .unwrap(),
+        )
+        .unwrap()
+    };
+    let x = Array::new(&[1, 1, k], Dtype::BF16, pattern(k as u32, 8)).unwrap();
+    let ids = u32s((0..rows).map(|r| r * e / rows).collect());
+    let map = u32s(vec![0; rows as usize]);
+    let read = |a: &Array| {
+        ops::evaluate(&streams, std::slice::from_ref(a), gpu, false).unwrap();
+        let mut v = vec![0; a.layout().nbytes()];
+        a.buffer().read(a.offset(), &mut v).unwrap();
+        v
+    };
+    for (failure, group) in [(1, 32), (2, 64), (3, 128)] {
+        for separate in [false, true] {
+            let mut q = mlp::fixture_linear(e * n, k as u32, 4, group, 7);
+            for a in [&mut q.weight, &mut q.scales, &mut q.biases] {
+                let width = a.layout().dim(-1).unwrap();
+                *a = ops::reshape(a, &[e as i32, n as i32, width], gpu).unwrap();
+            }
+            let up = separate.then_some(&q);
+            let reference = read(&stock(&x, &map, &q, up, &ids, true, &streams, gpu).unwrap());
+            let before = stats();
+            FAILURE.store(failure, std::sync::atomic::Ordering::Relaxed);
+            let actual = read(&swiglu_rows(&x, &map, &q, up, &ids, true, &streams, gpu).unwrap());
+            FAILURE.store(0, std::sync::atomic::Ordering::Relaxed);
+            assert_eq!(actual, reference);
+            let rejected = stats();
+            assert_eq!(
+                rejected.failures - before.failures,
+                if separate { 1 } else { 2 }
+            );
+            assert_eq!(
+                read(&swiglu_rows(&x, &map, &q, up, &ids, true, &streams, gpu).unwrap()),
+                reference
+            );
+            let after = stats();
+            assert_eq!(after.canaries, rejected.canaries);
+            assert_eq!(after.calls, before.calls);
+            assert_eq!(after.failures, rejected.failures);
+            println!(
+                "{}",
+                serde_json::json!({"case":"sticky_fallback","failure":failure,"separate":separate,"canaries":after.canaries,"failures":after.failures,"fallbacks":after.fallbacks,"exact":true})
+            );
+        }
+    }
+}
+
+#[test]
+fn row_map_applicability() {
+    let a = |shape: &[i32], dtype| Array::unallocated(shape, dtype).unwrap();
+    let q = QuantizedLinear {
+        weight: a(&[128, 256, 32], Dtype::U32),
+        scales: a(&[128, 256, 4], Dtype::BF16),
+        biases: a(&[128, 256, 4], Dtype::BF16),
+        bits: 4,
+        group: 64,
+    };
+    let x = a(&[512, 1, 256], Dtype::BF16);
+    for rows in [1, 40, 63, 64, 4095, 4096, 32769] {
+        let map = a(&[rows], Dtype::U32);
+        let ids = a(&[rows], Dtype::U32);
+        assert_eq!(applies(&x, &map, &q, None, &ids, true, true), rows >= 4096);
+        assert!(!applies(&x, &map, &q, None, &ids, false, true));
+        assert!(!applies(&x, &map, &q, None, &ids, true, false));
+    }
+    let ids = a(&[4096], Dtype::U32);
+    for map in [
+        a(&[4096], Dtype::I32),
+        a(&[1, 4096], Dtype::U32),
+        a(&[4095], Dtype::U32),
+    ] {
+        assert!(!applies(&x, &map, &q, None, &ids, true, true));
+    }
+    for x in [
+        a(&[512, 1, 256], Dtype::F32),
+        a(&[512, 256], Dtype::BF16),
+        a(&[512, 2, 256], Dtype::BF16),
+        a(&[512, 1, 128], Dtype::BF16),
+    ] {
+        assert!(!applies(&x, &ids, &q, None, &ids, true, true));
+    }
+    let mut bad = q.clone();
+    bad.group = 16;
+    assert!(!applies(&x, &ids, &bad, None, &ids, true, true));
+    bad = q.clone();
+    bad.bits = 2;
+    assert!(!applies(&x, &ids, &bad, None, &ids, true, true));
+    bad = q.clone();
+    bad.scales = a(&[128, 256, 4], Dtype::F16);
+    assert!(!applies(&x, &ids, &bad, None, &ids, true, true));
+    bad = q.clone();
+    bad.weight = a(&[128, 255, 32], Dtype::U32);
+    assert!(!applies(&x, &ids, &bad, None, &ids, true, true));
+    assert!(!applies(&x, &ids, &q, Some(&bad), &ids, true, true));
+}
+
+#[test]
+#[ignore = "installed projection weights; external supervisor required"]
+fn issue112_installed_gate_up() {
+    use super::super::{
+        gpu::Context,
+        qwen_mtplx_tests::{hex_digest, pattern},
+    };
+    super::configure_sources().unwrap();
+    let _context = Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(super::stream::Device::Gpu).unwrap();
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    let weights = super::weights::load_shards(&root, &streams).unwrap();
+    let parts = |name: &str| {
+        ["weight", "scales", "biases"].map(|p| {
+            weights[&format!("language_model.model.layers.0.mlp.switch_mlp.{name}.{p}")].clone()
+        })
+    };
+    let gate = parts("gate_proj");
+    let up = parts("up_proj");
+    let gu = std::array::from_fn(|i| {
+        ops::concatenate(&[gate[i].clone(), up[i].clone()], 1, gpu).unwrap()
+    });
+    let pack = |[weight, scales, biases]: [Array; 3]| QuantizedLinear {
+        weight,
+        scales,
+        biases,
+        bits: 4,
+        group: 64,
+    };
+    let fused = pack(gu);
+    let gate = pack(gate);
+    let up = pack(up);
+    let u32s = |v: Vec<u32>| {
+        Array::new(
+            &[v.len() as i32],
+            Dtype::U32,
+            super::scalar_buffer(&v.into_iter().flat_map(u32::to_le_bytes).collect::<Vec<_>>())
+                .unwrap(),
+        )
+        .unwrap()
+    };
+    let read = |a: &Array| {
+        ops::evaluate(&streams, std::slice::from_ref(a), gpu, false).unwrap();
+        let mut v = vec![0; a.layout().nbytes()];
+        a.buffer().read(a.offset(), &mut v).unwrap();
+        v
+    };
+    for tokens in [410, 2048, 3277, 4095, 4096, 8192] {
+        let rows = tokens * 10;
+        let x = Array::new(
+            &[tokens as i32, 1, 2560],
+            Dtype::BF16,
+            pattern(tokens * 2560, 7),
+        )
+        .unwrap();
+        let ids = u32s((0..rows).map(|r| r * 512 / rows).collect());
+        let map = u32s((0..rows).map(|r| r % tokens).collect());
+        let mut packed = None;
+        for separate in [false, true] {
+            let (g, u) = if separate {
+                (&gate, Some(&up))
+            } else {
+                (&fused, None)
+            };
+            let before = stats();
+            let y = swiglu_rows(&x, &map, g, u, &ids, true, &streams, gpu).unwrap();
+            let got = read(&y);
+            assert_eq!(
+                got,
+                read(&stock(&x, &map, g, u, &ids, true, &streams, gpu).unwrap())
+            );
+            let hash = hex_digest(&got);
+            if let Some(packed) = &packed {
+                assert_eq!(&hash, packed);
+            } else {
+                packed = Some(hash.clone());
+            }
+            let after = stats();
+            assert_eq!(after.failures, before.failures);
+            assert_eq!(after.calls, before.calls + 1);
+            println!(
+                "{}",
+                serde_json::json!({"case":"installed_gate_up","tokens":tokens,"rows":rows,"separate":separate,"sha256":hash,"canaries":after.canaries,"calls":after.calls,"exact":true})
+            );
+        }
+    }
+}
+
+#[test]
+#[ignore = "full installed model; validates all cold/warm state receipts against the pre-port oracle"]
+fn issue112_full_model_state() {
+    let before = stats();
+    super::issue110::issue110_large_cold_and_warm_prefill();
+    let after = stats();
+    assert!(after.calls > before.calls);
+    assert_eq!(after.failures, before.failures);
+    let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE110_DIR").unwrap());
+    for chunk in [2048, 4096] {
+        for phase in ["cold", "warm"] {
+            let actual: serde_json::Value = serde_json::from_slice(
+                &std::fs::read(directory.join(format!("{chunk}-{phase}.json"))).unwrap(),
+            )
+            .unwrap();
+            let expected: serde_json::Value = serde_json::from_slice(
+                &std::fs::read(
+                    std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(format!(
+                        "docs/issue110-sorted-gather/{chunk}-{phase}-state.json"
+                    )),
+                )
+                .unwrap(),
+            )
+            .unwrap();
+            assert_eq!(actual["state"], expected["state"]);
+            assert_eq!(actual["cached"], expected["cached"]);
+        }
+    }
+    println!(
+        "{}",
+        serde_json::json!({"case":"model_state","calls":after.calls,"canaries":after.canaries,"failures":after.failures,"cold_warm_exact":true})
+    );
+}

src/engine/metal/qwen_mtplx/ops.rs

@@ -6,6 +6,8 @@ use super::stream::{Device, Stream, Streams};
 use super::{eval, views};
 
 pub(super) enum Operation {
+    MoeSortedGather(super::moe_sorted::Kernel),
+    MoeSortedCounts,
     Vision(super::vision_kernels::Kernel),
     Compiled(super::fused::Expression),
     Load(super::load::Load),
@@ -60,7 +62,7 @@ pub(super) enum Operation {
         bits: u32,
         group: u32,
     },
-    GatherQmmBf16 {
+    GatherQmm {
         bits: u32,
         group: u32,
         right_sorted: bool,
@@ -134,6 +136,9 @@ pub(super) enum Operation {
 impl Operation {
     fn evaluate(&self, input: &Array, output: &Array, stream: Stream) -> Result<(), String> {
         match self {
+            Self::MoeSortedGather(_) | Self::MoeSortedCounts => {
+                Err("sorted MoE requires captured inputs".into())
+            }
             Self::Compiled(_) => Err("compiled expression requires all inputs/outputs".into()),
             Self::QsaCompiledScalar(_) => {
                 Err("compiled scalar requires its captured inputs".into())
@@ -215,7 +220,7 @@ impl Operation {
             | Self::ArangeI64 => Err("invalid indexing arity".into()),
             Self::Vision(_) | Self::RmsNorm(_) => Err("RMSNorm requires two inputs".into()),
             Self::QuantizedLinearBf16 { .. } => Err("quantized linear requires four inputs".into()),
-            Self::GatherQmmBf16 { .. } => Err("GatherQMM requires six inputs".into()),
+            Self::GatherQmm { .. } => Err("GatherQMM requires six inputs".into()),
             Self::AffineQuantize { .. } => Err("quantize requires three sibling outputs".into()),
             Self::AffineDequantize { .. } => Err("dequantize requires three inputs".into()),
             Self::HyperRead(_) | Self::HyperV3R1 | Self::HyperV3R2 => {
@@ -262,6 +267,12 @@ pub(super) fn evaluate(
                 .operation()
                 .ok_or("array has no executable primitive")?;
             match &*operation {
+                Operation::MoeSortedGather(kernel) => {
+                    return kernel.evaluate(inputs, outputs);
+                }
+                Operation::MoeSortedCounts => {
+                    return super::moe_sorted::counts_evaluate(inputs, &outputs[0]);
+                }
                 Operation::Vision(k) => {
                     return super::vision_kernels::evaluate(*k, inputs, &outputs[0]);
                 }
@@ -526,7 +537,7 @@ pub(super) fn evaluate(
                 }
                 return super::linear::evaluate(inputs, &outputs[0], *bits, *group);
             }
-            if let Operation::GatherQmmBf16 {
+            if let Operation::GatherQmm {
                 bits,
                 group,
                 right_sorted,

src/engine/metal/qwen_mtplx/policy.rs

@@ -25,6 +25,11 @@ fn falsey(value: Option<&str>) -> bool {
     })
 }
 
+pub(super) fn sorted_gather_enabled(get: impl Fn(&'static str) -> Option<String>) -> bool {
+    !falsey(get("MTPLX_MOE_SORTED_GATHER_KERNEL").as_deref())
+        && !truthy(get("MTPLX_FORCE_GPU_FAMILY_FALLBACK").as_deref())
+}
+
 fn strict_bool(name: &str, raw: Option<String>) -> Result<bool, String> {
     match raw
         .as_deref()
@@ -323,3 +328,19 @@ fn qwen_server_policy_matches_original_functions() {
     }
     assert_eq!(count, 268);
 }
+
+#[test]
+fn sorted_moe_switch_and_family_rehearsal() {
+    assert!(sorted_gather_enabled(|_| None));
+    for off in ["0", " false ", "NO", "off"] {
+        assert!(!sorted_gather_enabled(|key| (key
+            == "MTPLX_MOE_SORTED_GATHER_KERNEL")
+            .then(|| off.to_owned())));
+    }
+    for force in ["1", " true ", "YES", "on"] {
+        assert!(!sorted_gather_enabled(|key| (key
+            == "MTPLX_FORCE_GPU_FAMILY_FALLBACK")
+            .then(|| force.to_owned())));
+    }
+    assert!(sorted_gather_enabled(|_| Some("unexpected".into())));
+}