Save inference parity implementation and evaluation harness

src/engine/metal/qwen_mtplx/model.rs

@@ -0,0 +1,806 @@
+//! Qwen4ExpTextModel/TextModel graph orchestration from the pinned MTPLX.
+//! GDN runs use the reference simplify/fuse passes. Synthetic run graphs are
+//! checked; real-weight and serving integration remain outside this test port.
+use super::array::{Array, Dtype};
+use super::binary::{Binary, binary};
+use super::decoder::{Attention, AttentionCache, Cache, DecoderLayer, Options};
+use super::mlp::Linear;
+use super::stream::{Stream, Streams};
+use super::{cache::QsaCache, compiled::Replay, gdn, hyper, indexing, linear, ops, ple};
+use std::collections::HashMap;
+use std::ops::Range;
+
+#[derive(Clone, Copy, Default)]
+pub(super) struct CompilePolicy {
+    explicit_off: bool,
+    environment: bool,
+    lane: bool,
+}
+
+impl CompilePolicy {
+    pub(super) fn from_values(gdn: Option<&str>, family: Option<&str>) -> Self {
+        let truthy = |s: &str| {
+            matches!(
+                s.trim().to_ascii_lowercase().as_str(),
+                "1" | "true" | "yes" | "on"
+            )
+        };
+        let explicit_off = [gdn, family].into_iter().flatten().any(|s| !truthy(s));
+        Self {
+            explicit_off,
+            environment: !explicit_off && [gdn, family].into_iter().flatten().any(truthy),
+            lane: false,
+        }
+    }
+
+    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
+    pub(super) fn from_env() -> Self {
+        Self::from_values(
+            std::env::var("MTPLX_COMPILED_GDN").ok().as_deref(),
+            std::env::var("MTPLX_QWEN4EXP_COMPILE").ok().as_deref(),
+        )
+    }
+
+    fn enabled(self) -> bool {
+        !self.explicit_off && (self.environment || self.lane)
+    }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum Run {
+    Eager(usize),
+    Gdn(Range<usize>),
+}
+
+fn runs(layers: impl IntoIterator<Item = (bool, bool)>) -> Vec<Run> {
+    let mut result = Vec::new();
+    let mut start = None;
+    let mut end = 0;
+    for (i, (linear, ple)) in layers.into_iter().enumerate() {
+        end = i + 1;
+        if linear && !ple {
+            start.get_or_insert(i);
+        } else {
+            if let Some(start) = start.take() {
+                result.push(Run::Gdn(start..i));
+            }
+            result.push(Run::Eager(i));
+        }
+    }
+    if let Some(start) = start {
+        result.push(Run::Gdn(start..end));
+    }
+    result
+}
+
+#[derive(Hash, PartialEq, Eq)]
+struct RunKey {
+    range: Range<usize>,
+    capture: bool,
+    stream: Stream,
+    signature: Vec<(Vec<i32>, Dtype)>,
+}
+
+pub(super) struct TextModel<'a> {
+    /// Original Module.parameters tree order, populated by checkpoint binding.
+    pub(super) parameter_order: Vec<String>,
+    pub(super) embedding: Linear<'a>,
+    pub(super) layers: Vec<DecoderLayer<'a>>,
+    pub(super) mixer: hyper::GatedResidual<'a>,
+    pub(super) head: Option<Linear<'a>>, // None means tied embedding weights.
+    pub(super) mtp: Option<super::mtp::Mtp<'a>>,
+    pub(super) last_widened: Option<Array>,
+    policy: CompilePolicy,
+    decode_runs: Option<Vec<Run>>,
+    replays: HashMap<RunKey, Replay>,
+    commit_refusals: usize,
+}
+
+/// ArraysCache.state before verification; QSA entries need no snapshot here.
+pub(super) type ArraySnapshot = Vec<Option<Array>>;
+
+#[derive(Clone, Copy)]
+pub(super) struct OutputOptions {
+    pub(super) emit_logits: bool,
+    pub(super) logits_keep: i32,
+    pub(super) return_hidden: bool,
+}
+
+pub(super) struct Output {
+    pub(super) logits: Option<Array>,
+    pub(super) hidden: Option<Array>,
+}
+
+fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result<Array, String> {
+    match layer {
+        Linear::Dense(weight) => indexing::take_rows(weight, ids, stream),
+        Linear::Quantized(layer) => linear::dequantize(
+            &indexing::take_rows(&layer.weight, ids, stream)?,
+            &indexing::take_rows(&layer.scales, ids, stream)?,
+            &indexing::take_rows(&layer.biases, ids, stream)?,
+            layer.group,
+            layer.bits,
+            stream,
+        ),
+    }
+}
+
+fn mask(cache: Option<&Cache>, rows: i32, stream: Stream) -> Result<Option<Array>, String> {
+    let Some(Cache {
+        attention: AttentionCache::Linear(cache),
+        ..
+    }) = cache
+    else {
+        return Ok(None);
+    };
+    let (metadata, op) = if let Some(padding) = &cache.left_padding {
+        (padding, Binary::GreaterEqual)
+    } else if let Some(lengths) = &cache.lengths {
+        (lengths, Binary::Less)
+    } else {
+        return Ok(None);
+    };
+    binary(
+        &indexing::arange_i32(rows, stream)?,
+        &ops::expand_dims(metadata, &[-1], stream)?,
+        op,
+        stream,
+    )
+    .map(Some)
+}
+
+impl<'a> TextModel<'a> {
+    #[cfg(test)]
+    pub(super) fn replay_count(&self) -> usize {
+        self.replays.len()
+    }
+    pub(super) fn new(
+        embedding: Linear<'a>,
+        layers: Vec<DecoderLayer<'a>>,
+        mixer: hyper::GatedResidual<'a>,
+        head: Option<Linear<'a>>,
+        policy: CompilePolicy,
+    ) -> Result<Self, String> {
+        if layers.is_empty() || mixer.inject.is_some() {
+            return Err("Qwen requires decoder layers and a non-combining final mixer".into());
+        }
+        Ok(Self {
+            embedding,
+            layers,
+            mixer,
+            head,
+            parameter_order: Vec::new(),
+            mtp: None,
+            last_widened: None,
+            policy,
+            decode_runs: None,
+            replays: HashMap::new(),
+            commit_refusals: 0,
+        })
+    }
+
+    pub(super) fn make_cache(&self, ratio: i32) -> Vec<Option<Cache>> {
+        self.layers
+            .iter()
+            .map(|layer| {
+                Some(Cache {
+                    attention: match layer.attention {
+                        Attention::Linear(_) => AttentionCache::Linear(gdn::Cache::default()),
+                        Attention::Qsa(_) => {
+                            AttentionCache::Qsa(QsaCache::new(if ratio == 0 { 4 } else { ratio }))
+                        }
+                    },
+                    ple: ple::Cache::default(),
+                    capture_ple: None,
+                })
+            })
+            .collect()
+    }
+
+    pub(super) fn make_mtp_cache(&self, ratio: i32) -> [Option<Cache>; 1] {
+        super::mtp::make_cache(ratio)
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn mtp_forward(
+        &mut self,
+        hidden: &Array,
+        ids: &Array,
+        cache: Option<&mut [Option<Cache>]>,
+        return_hidden: bool,
+        options: &Options,
+        vision: Option<(Option<&Array>, i32)>,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<Output, String> {
+        let emb = embedding(self.embedding, ids, stream)?;
+        let mtp = self
+            .mtp
+            .as_mut()
+            .ok_or("Qwen MTP weights are not attached")?;
+        let h = mtp.fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)?;
+        let mixed = mtp
+            .mixer
+            .forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
+            .mixed;
+        let logits = self.head.unwrap_or(self.embedding).apply(&mixed, stream)?;
+        Ok(Output {
+            logits: Some(logits),
+            hidden: return_hidden.then_some(h),
+        })
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn mtp_update_cache(
+        &mut self,
+        hidden: &Array,
+        ids: &Array,
+        cache: Option<&mut [Option<Cache>]>,
+        input_embeddings: Option<&Array>,
+        options: &Options,
+        vision: Option<(Option<&Array>, i32)>,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<Array, String> {
+        let emb = match input_embeddings {
+            Some(emb) => emb.clone(),
+            None => embedding(self.embedding, ids, stream)?,
+        };
+        self.mtp
+            .as_mut()
+            .ok_or("Qwen MTP weights are not attached")?
+            .fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)
+    }
+
+    /// Model.post_weight_load: attach each PLE table after trunk weight loading.
+    pub(super) fn post_weight_load(
+        &mut self,
+        root: &std::path::Path,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<(), String> {
+        let path = root.join("ngram-table.safetensors");
+        if !path.exists() {
+            return Ok(());
+        }
+        let resident = super::ngram::resident_policy();
+        let mut sidecar = None;
+        for layer in &mut self.layers {
+            let Some(ple) = &mut layer.ple else { continue };
+            let table = &mut ple.embedding.table;
+            table.attach_sidecar(&path)?;
+            if resident {
+                // Original attach_resident reports failure and retains the SSD fallback.
+                table.attach_resident(&path, streams, stream);
+            }
+            sidecar = table.sidecar.as_ref();
+        }
+        if !resident && let Some(sidecar) = sidecar {
+            println!(
+                "[qwen4_exp] ngram table: streamed from SSD (file-backed pages, hot cache {}M, prefetch {}) — resident only at >=160G RAM or MTPLX_NGRAM_RESIDENT=1",
+                sidecar.hot_mebibytes(),
+                if sidecar.prefetch_enabled() {
+                    "on"
+                } else {
+                    "off"
+                }
+            );
+            let _ = std::io::Write::flush(&mut std::io::stdout());
+        }
+        Ok(())
+    }
+
+    pub(super) fn set_ar_pipeline_mode(&mut self, enabled: bool) -> bool {
+        let mut ready = true;
+        for layer in &mut self.layers {
+            let Some(ple) = &mut layer.ple else { continue };
+            if enabled && ple.embedding.table.resident.is_none() {
+                ready = false;
+                continue;
+            }
+            ple.embedding.stage_disabled = enabled;
+            ple.embedding.table.prefer_lazy = enabled;
+        }
+        if ready {
+            self.policy.lane = enabled && !self.policy.explicit_off;
+        }
+        ready
+    }
+
+    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
+    pub(super) fn clear_verify_capture(cache: &mut [Option<Cache>]) {
+        for cache in cache.iter_mut().flatten() {
+            cache.capture_ple = None;
+            if let AttentionCache::Linear(c) = &mut cache.attention {
+                c.capture = None;
+            }
+        }
+    }
+
+    fn refuse_commit(&mut self, layer: usize, reason: &str) -> bool {
+        if self.commit_refusals < 3 {
+            use std::io::Write;
+            println!(
+                "[qwen4_exp] capture-commit refused (layer {layer}: {reason}) — falling back to re-forward"
+            );
+            let _ = std::io::stdout().flush();
+        }
+        self.commit_refusals += 1;
+        false
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn commit_verified_window(
+        &mut self,
+        cache: &mut [Option<Cache>],
+        snapshots: Option<&[Option<ArraySnapshot>]>,
+        keep: i32,
+        verified: i32,
+        options: &Options,
+        vision: Option<(Option<&Array>, i32)>,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<bool, String> {
+        if keep < 1 || keep > verified || cache.len() != self.layers.len() {
+            return Ok(false);
+        }
+        // First validate EVERY layer, without cache mutations or GPU work.
+        // A refusal must leave the full verify window available for fallback.
+        for (i, (layer, entry)) in self.layers.iter().zip(cache.iter()).enumerate() {
+            let Some(entry) = entry else {
+                return Ok(self.refuse_commit(i, "entry_missing"));
+            };
+            let AttentionCache::Linear(c) = &entry.attention else {
+                continue;
+            };
+            let snapshot = snapshots.and_then(|s| s.get(i)).and_then(Option::as_ref);
+            let validation = if layer.ple.is_some() {
+                match (snapshot, &entry.capture_ple) {
+                    (None, _) => Err("snapshot_missing".to_owned()),
+                    (_, None) => Err("ple_rows_missing".to_owned()),
+                    (Some(_), Some((h, _))) if h.layout().shape()[1] != verified => Err(format!(
+                        "ple_rows_width_{}_vs_{verified}",
+                        h.layout().shape()[1]
+                    )),
+                    (Some(pre), _) if pre.len() < 4 => Err("ple_snapshot_short".to_owned()),
+                    _ => Ok(()),
+                }
+            } else {
+                c.validate_verified_window(snapshot.map(Vec::as_slice), verified)
+            };
+            if let Err(reason) = validation {
+                return Ok(self.refuse_commit(i, &reason));
+            }
+        }
+        for (i, (layer, entry)) in self.layers.iter_mut().zip(cache).enumerate() {
+            let entry = entry.as_mut().unwrap();
+            if let AttentionCache::Qsa(c) = &mut entry.attention {
+                if keep != verified {
+                    c.trim(verified - keep);
+                }
+                continue;
+            }
+            let snapshot = snapshots.unwrap()[i].as_ref().unwrap();
+            if layer.ple.is_some() {
+                let (hidden, ids) = entry.capture_ple.as_ref().unwrap();
+                let prefix = |a: &Array| {
+                    let mut end = a.layout().shape().to_vec();
+                    end[1] = keep;
+                    ops::slice(a, &vec![0; end.len()], &end, &vec![1; end.len()], stream)
+                };
+                let hidden = prefix(hidden)?;
+                let ids = prefix(ids)?;
+                let AttentionCache::Linear(c) = &mut entry.attention else {
+                    unreachable!()
+                };
+                c.conv = snapshot[0].clone();
+                c.delta = snapshot[1].clone();
+                entry.ple.conv = snapshot[2].clone();
+                entry.ple.history = snapshot[3].clone();
+                layer.ple.as_mut().unwrap().embedding.stage(
+                    &ids,
+                    entry.ple.history.as_ref(),
+                    &options.stage,
+                    streams,
+                    stream,
+                );
+                layer.call(
+                    &hidden,
+                    Some(&ids),
+                    None,
+                    Some(entry),
+                    options,
+                    vision,
+                    streams,
+                    stream,
+                )?;
+                entry.capture_ple = None;
+            } else {
+                let AttentionCache::Linear(c) = &mut entry.attention else {
+                    unreachable!()
+                };
+                let Attention::Linear(weights) = &layer.attention else {
+                    return Err("GDN capture-commit layer/cache mismatch".into());
+                };
+                c.replay_verified_prefix(snapshot, keep, weights.a_log, weights.dt_bias, stream)?;
+            }
+        }
+        Ok(true)
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn call(
+        &mut self,
+        ids: &Array,
+        cache: Option<&mut [Option<Cache>]>,
+        input_embeddings: Option<&Array>,
+        output: OutputOptions,
+        options: &Options,
+        vision: Option<(Option<&Array>, i32)>,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<Output, String> {
+        let h = match input_embeddings {
+            Some(h) => h.clone(),
+            None => embedding(self.embedding, ids, stream)?,
+        };
+        let shape = h.layout().shape().to_vec();
+        if shape.len() != 3
+            || shape[0] < 1
+            || shape[1] < 1
+            || shape[2] != 2560
+            || ids.layout().shape() != &shape[..2]
+        {
+            return Err(
+                "installed Qwen requires [B,S,2560] embeddings and matching token IDs".into(),
+            );
+        }
+        let mut empty = Vec::new();
+        let cache = match cache {
+            Some(cache) => cache,
+            None => {
+                empty.resize_with(self.layers.len(), || None);
+                &mut empty
+            }
+        };
+        if cache.len() != self.layers.len() {
+            return Err("Qwen cache/layer count mismatch".into());
+        }
+        let ssm = self
+            .layers
+            .iter()
+            .position(|l| matches!(l.attention, Attention::Linear(_)))
+            .unwrap_or(0);
+        let ssm_mask = mask(cache[ssm].as_ref(), shape[1], stream)?;
+        if let Some(i) = self.layers.iter().position(|l| l.ple.is_some()) {
+            self.layers[i].ple.as_mut().unwrap().embedding.stage(
+                ids,
+                cache[i].as_ref().and_then(|c| c.ple.history.as_ref()),
+                &options.stage,
+                streams,
+                stream,
+            );
+        }
+        // tile(h, (1,1,4)): original reshape -> broadcast -> reshape.
+        let h = ops::reshape(&h, &[shape[0], shape[1], 1, 2560], stream)?;
+        let h = ops::broadcast_to(&h, &[shape[0], shape[1], 4, 2560], stream)?;
+        let mut h = ops::reshape(&h, &[shape[0], shape[1], 10240], stream)?;
+        if (1..=4).contains(&shape[1])
+            && ssm_mask.is_none()
+            && self.policy.enabled()
+            && cache[ssm].is_some()
+        {
+            h = self.decode(h, ids, cache, options, vision, streams, stream)?;
+        } else {
+            for (layer, cache) in self.layers.iter_mut().zip(cache) {
+                h = eager(
+                    layer,
+                    &h,
+                    ids,
+                    ssm_mask.as_ref(),
+                    cache.as_mut(),
+                    options,
+                    vision,
+                    streams,
+                    stream,
+                )?;
+            }
+        }
+        self.last_widened = Some(h.clone());
+        let mixed = self
+            .mixer
+            .forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
+            .mixed;
+        let logits = if output.emit_logits {
+            let mixed = if output.logits_keep != 0 {
+                let keep = output.logits_keep.max(1).min(shape[1]);
+                ops::slice(
+                    &mixed,
+                    &[0, shape[1] - keep, 0],
+                    &[shape[0], shape[1], 2560],
+                    &[1; 3],
+                    stream,
+                )?
+            } else {
+                mixed
+            };
+            Some(self.head.unwrap_or(self.embedding).apply(&mixed, stream)?)
+        } else {
+            None
+        };
+        Ok(Output {
+            logits,
+            hidden: output.return_hidden.then_some(h),
+        })
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn decode(
+        &mut self,
+        mut h: Array,
+        ids: &Array,
+        cache: &mut [Option<Cache>],
+        options: &Options,
+        vision: Option<(Option<&Array>, i32)>,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<Array, String> {
+        let runs = self.decode_runs.get_or_insert_with(|| {
+            runs(
+                self.layers
+                    .iter()
+                    .map(|l| (matches!(l.attention, Attention::Linear(_)), l.ple.is_some())),
+            )
+        });
+        for run in runs {
+            let range = match run {
+                Run::Eager(i) => {
+                    h = eager(
+                        &mut self.layers[*i],
+                        &h,
+                        ids,
+                        None,
+                        cache[*i].as_mut(),
+                        options,
+                        vision,
+                        streams,
+                        stream,
+                    )?;
+                    continue;
+                }
+                Run::Gdn(range) => range.clone(),
+            };
+            let mut flat = vec![h.clone()];
+            for c in &cache[range.clone()] {
+                let Some(Cache {
+                    attention: AttentionCache::Linear(c),
+                    ..
+                }) = c
+                else {
+                    break;
+                };
+                let (Some(conv), Some(delta)) = (&c.conv, &c.delta) else {
+                    break;
+                };
+                flat.extend([conv.clone(), delta.clone()]);
+            }
+            if flat.len() != 1 + 2 * range.len() {
+                for i in range {
+                    h = eager(
+                        &mut self.layers[i],
+                        &h,
+                        ids,
+                        None,
+                        cache[i].as_mut(),
+                        options,
+                        vision,
+                        streams,
+                        stream,
+                    )?;
+                }
+                continue;
+            }
+            let key = RunKey {
+                range: range.clone(),
+                capture: options.gdn.capture,
+                stream,
+                signature: flat
+                    .iter()
+                    .map(|a| (a.layout().shape().to_vec(), a.layout().dtype()))
+                    .collect(),
+            };
+            let replay = match self.replays.entry(key) {
+                std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
+                std::collections::hash_map::Entry::Vacant(entry) => {
+                    let replay = Replay::trace_compiled(&flat, |args| {
+                        let mut h = args[0].clone();
+                        let mut states = Vec::new();
+                        let mut captured = Vec::new();
+                        for (n, i) in range.clone().enumerate() {
+                            let mut c = Cache {
+                                attention: AttentionCache::Linear(gdn::Cache {
+                                    conv: Some(args[1 + 2 * n].clone()),
+                                    delta: Some(args[2 + 2 * n].clone()),
+                                    ..Default::default()
+                                }),
+                                ple: ple::Cache::default(),
+                                capture_ple: None,
+                            };
+                            h = self.layers[i].call(
+                                &h,
+                                None,
+                                None,
+                                Some(&mut c),
+                                options,
+                                vision,
+                                streams,
+                                stream,
+                            )?;
+                            let AttentionCache::Linear(c) = c.attention else {
+                                unreachable!()
+                            };
+                            states.extend([
+                                c.conv.ok_or("compiled GDN conv missing")?,
+                                c.delta.ok_or("compiled GDN state missing")?,
+                            ]);
+                            if options.gdn.capture {
+                                captured.extend(c.capture.ok_or("compiled GDN capture missing")?);
+                            }
+                        }
+                        Ok(std::iter::once(h).chain(states).chain(captured).collect())
+                    })?;
+                    entry.insert(replay)
+                }
+            };
+            let out = replay.run(&flat)?;
+            h = out[0].clone();
+            for (n, i) in range.clone().enumerate() {
+                let AttentionCache::Linear(c) = &mut cache[i].as_mut().unwrap().attention else {
+                    unreachable!()
+                };
+                c.conv = Some(out[1 + 2 * n].clone());
+                c.delta = Some(out[2 + 2 * n].clone());
+                if options.gdn.capture {
+                    let start = 1 + 2 * range.len() + 6 * n;
+                    c.capture = Some(std::array::from_fn(|j| out[start + j].clone()));
+                }
+            }
+        }
+        Ok(h)
+    }
+}
+
+#[cfg(test)]
+impl TextModel<'_> {
+    pub(super) fn run_contract(&self, rows: i32, capture: bool) -> serde_json::Value {
+        use super::ops::Operation as O;
+        let replay = &self
+            .replays
+            .iter()
+            .find(|(key, _)| key.capture == capture && key.signature[0].0 == [1, rows, 10240])
+            .expect("compiled fixture run missing")
+            .1;
+        let mut counts = std::collections::BTreeMap::<String, usize>::new();
+        let contract = replay.contract(
+            &["hidden", "conv", "delta"],
+            |op| {
+                let name = match op {
+                    O::Compiled(expression) => expression.name(),
+                    _ => match op {
+                        O::Reshape => "Reshape",
+                        O::Flatten { .. } => "Flatten",
+                        O::Unflatten { .. } => "Unflatten",
+                        O::Transpose(_) => "Transpose",
+                        O::ExpandDims(_) => "ExpandDims",
+                        O::Squeeze(_) => "Squeeze",
+                        O::Slice { .. } => "Slice",
+                        O::Contiguous(_) => "Contiguous",
+                        O::RmsNorm(_) => "RMSNorm",
+                        O::Softmax(_) => "Softmax",
+                        O::QuantizedLinearBf16 { .. } => "QuantizedMatmul",
+                        O::GatherQmmBf16 { .. } => "GatherQMM",
+                        O::DenseMatmul => "Matmul",
+                        O::Sum(_) => "Sum",
+                        O::Sort {
+                            indices: true,
+                            kth: Some(_),
+                            ..
+                        } => "ArgPartition",
+                        O::Split { .. } => "Split",
+                        O::HyperRead(_)
+                        | O::GdnStepBf16
+                        | O::GdnConvNorm(_)
+                        | O::GatedDeltaBf16(_) => "CustomKernel",
+                        O::Concatenate(_) => "Concatenate",
+                        O::DepthwiseConv1d => "Convolution",
+                        O::GatherAxis(_) => "GatherAxis",
+                        O::ArangeI32(_) | O::ArangeU32 => "Arange",
+                        _ => super::fused::operator(op).expect("uncaptured GDN primitive"),
+                    }
+                    .to_owned(),
+                };
+                *counts.entry(name.clone()).or_default() += 1;
+                name
+            },
+            if capture { 9 } else { 3 },
+        );
+        serde_json::json!({"contract":contract,"ops":counts})
+    }
+}
+
+#[test]
+fn mtplx_model_compile_policy_and_run_boundaries() {
+    assert_eq!(
+        runs([
+            (true, false),
+            (true, false),
+            (true, true),
+            (true, false),
+            (false, false),
+            (true, false)
+        ]),
+        vec![
+            Run::Gdn(0..2),
+            Run::Eager(2),
+            Run::Gdn(3..4),
+            Run::Eager(4),
+            Run::Gdn(5..6)
+        ]
+    );
+    for on in ["1", " true ", "YES", "On"] {
+        assert!(CompilePolicy::from_values(Some(on), None).enabled());
+        assert!(CompilePolicy::from_values(None, Some(on)).enabled());
+        for off in ["0", "", "false", "off", "anything"] {
+            for mut policy in [
+                CompilePolicy::from_values(Some(on), Some(off)),
+                CompilePolicy::from_values(Some(off), Some(on)),
+            ] {
+                policy.lane = true;
+                assert!(!policy.enabled());
+            }
+        }
+    }
+    let mut automatic = CompilePolicy::from_values(None, None);
+    assert!(!automatic.enabled());
+    automatic.lane = true;
+    assert!(automatic.enabled());
+    let environment = CompilePolicy::from_env();
+    assert_eq!(
+        environment.enabled(),
+        CompilePolicy::from_values(
+            std::env::var("MTPLX_COMPILED_GDN").ok().as_deref(),
+            std::env::var("MTPLX_QWEN4EXP_COMPILE").ok().as_deref(),
+        )
+        .enabled()
+    );
+}
+
+#[allow(clippy::too_many_arguments)]
+fn eager(
+    layer: &mut DecoderLayer<'_>,
+    hidden: &Array,
+    ids: &Array,
+    mask: Option<&Array>,
+    mut cache: Option<&mut Cache>,
+    options: &Options,
+    vision: Option<(Option<&Array>, i32)>,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    if options.gdn.capture
+        && layer.ple.is_some()
+        && let Some(c) = cache.as_deref_mut()
+    {
+        c.capture_ple = Some((hidden.clone(), ids.clone()));
+    }
+    layer.call(
+        hidden,
+        Some(ids),
+        mask,
+        cache,
+        options,
+        vision,
+        streams,
+        stream,
+    )
+}