Save inference parity implementation and evaluation harness

src/engine/metal/qwen_mtplx/fused.rs

@@ -0,0 +1,831 @@
+//! Runtime-generated Compiled Metal expressions from the pinned
+//! metal/compiled.cpp emitter. The shared functor/header source is unchanged.
+use super::array::{Array, Dtype};
+use super::{compiled, ops::Operation};
+use sha2::{Digest, Sha256};
+use std::fmt::Write;
+
+#[derive(Debug)]
+struct Input {
+    dtype: Dtype,
+    scalar: bool,
+    constant: Option<String>,
+}
+
+#[derive(Debug)]
+pub(super) struct Statement {
+    pub(super) dtype: Dtype,
+    /// Reference primitive name; Broadcast/AsType emit a static cast.
+    pub(super) operator: &'static str,
+    pub(super) inputs: Vec<usize>,
+}
+
+pub(super) struct Expression {
+    inputs: Vec<Input>,
+    statements: Vec<Statement>,
+    outputs: Vec<usize>,
+    prefix: String,
+}
+
+pub(super) fn operator(op: &Operation) -> Option<&'static str> {
+    use super::{binary::Binary as B, unary::Unary as U};
+    Some(match op {
+        Operation::Broadcast => "Broadcast",
+        Operation::AsType => "AsType",
+        Operation::Select => "Select",
+        Operation::Unary(op) => match op {
+            U::LogicalNot => "LogicalNot",
+            U::Abs => "Abs",
+            U::Sign => "Sign",
+            U::Negative => "Negative",
+            U::Square => "Square",
+            U::Exp => "Exp",
+            U::Log => "Log",
+            U::Sigmoid => "Sigmoid",
+            U::Tanh => "Tanh",
+            U::Sqrt => "Sqrt",
+            U::Rsqrt => "Rsqrt",
+            U::Sin => "Sin",
+            U::Cos => "Cos",
+        },
+        Operation::Binary(op) => match op {
+            B::Equal => "Equal",
+            B::GreaterEqual => "GreaterEqual",
+            B::Greater => "Greater",
+            B::Less => "Less",
+            B::LessEqual => "LessEqual",
+            B::LogicalAnd => "LogicalAnd",
+            B::LogicalOr => "LogicalOr",
+            B::BitwiseAnd => "BitwiseAnd",
+            B::BitwiseOr => "BitwiseOr",
+            B::BitwiseXor => "BitwiseXor",
+            B::Remainder => "Remainder",
+            B::Add => "Add",
+            B::Subtract => "Subtract",
+            B::Multiply => "Multiply",
+            B::Divide => "Divide",
+            B::Maximum => "Maximum",
+            B::Minimum => "Minimum",
+            B::LogAddExp => "LogAddExp",
+            B::Power => "Power",
+        },
+        _ => return None,
+    })
+}
+
+fn dtype(d: Dtype) -> &'static str {
+    match d {
+        Dtype::Bool => "bool",
+        Dtype::U8 => "uint8_t",
+        Dtype::U16 => "uint16_t",
+        Dtype::U32 => "uint32_t",
+        Dtype::U64 => "uint64_t",
+        Dtype::I8 => "int8_t",
+        Dtype::I16 => "int16_t",
+        Dtype::I32 => "int32_t",
+        Dtype::I64 => "int64_t",
+        Dtype::F16 => "float16_t",
+        Dtype::BF16 => "bfloat16_t",
+        Dtype::F32 => "float",
+        Dtype::F64 => "double",
+        Dtype::C64 => "complex64_t",
+    }
+}
+
+// ostream defaultfloat with digits10+1 (7 for F32/F16/BF16, 16 for F64).
+fn float_constant(value: f64, precision: usize) -> String {
+    if value.is_nan() {
+        return "NAN".into();
+    }
+    if value.is_infinite() {
+        return if value.is_sign_negative() {
+            "-INFINITY"
+        } else {
+            "INFINITY"
+        }
+        .into();
+    }
+    let scientific = format!("{:.*e}", precision - 1, value);
+    let (mantissa, exponent) = scientific.split_once('e').unwrap();
+    let exponent: i32 = exponent.parse().unwrap();
+    let trim = |s: String| {
+        if s.contains('.') {
+            s.trim_end_matches('0').trim_end_matches('.').to_owned()
+        } else {
+            s
+        }
+    };
+    if exponent < -4 || exponent >= precision as i32 {
+        format!("{}e{:+03}", trim(mantissa.to_owned()), exponent)
+    } else {
+        trim(format!(
+            "{:.*}",
+            (precision as i32 - 1 - exponent).max(0) as usize,
+            value
+        ))
+    }
+}
+
+fn constant(a: &Array) -> Result<String, String> {
+    a.wait()?;
+    if !a.has_data() {
+        return Err("compiled constant is unavailable".into());
+    }
+    let mut bits = [0u8; 8];
+    a.buffer().read(a.offset(), &mut bits[..a.itemsize()])?;
+    let u16v = u16::from_le_bytes(bits[..2].try_into().unwrap());
+    let u32v = u32::from_le_bytes(bits[..4].try_into().unwrap());
+    Ok(match a.layout().dtype() {
+        Dtype::Bool => (bits[0] != 0).to_string(),
+        Dtype::U8 => bits[0].to_string(),
+        Dtype::I8 => (bits[0] as i8).to_string(),
+        Dtype::U16 => u16v.to_string(),
+        Dtype::I16 => (u16v as i16).to_string(),
+        Dtype::U32 => u32v.to_string(),
+        Dtype::I32 => (u32v as i32).to_string(),
+        Dtype::U64 => u64::from_le_bytes(bits).to_string(),
+        Dtype::I64 => i64::from_le_bytes(bits).to_string(),
+        Dtype::BF16 => float_constant(f32::from_bits(u32::from(u16v) << 16) as f64, 7),
+        Dtype::F16 => float_constant(super::super::half_to_f32(u16v) as f64, 7),
+        Dtype::F32 => float_constant(f32::from_bits(u32v) as f64, 7),
+        Dtype::F64 => float_constant(f64::from_le_bytes(bits), 16),
+        Dtype::C64 => format!(
+            "complex64_t({}, {})",
+            float_constant(f32::from_bits(u32v) as f64, 7),
+            float_constant(f32::from_le_bytes(bits[4..].try_into().unwrap()) as f64, 7)
+        ),
+    })
+}
+
+struct Names {
+    slots: Vec<Option<String>>,
+    count: usize,
+}
+impl Names {
+    fn get(&mut self, slot: usize) -> String {
+        self.slots[slot]
+            .get_or_insert_with(|| {
+                self.count += 1;
+                let mut n = self.count;
+                let mut letters = Vec::new();
+                while n > 0 {
+                    letters.push(b'A' + ((n - 1) % 26) as u8);
+                    n = (n - 1) / 26;
+                }
+                letters.reverse();
+                String::from_utf8(letters).unwrap()
+            })
+            .clone()
+    }
+}
+
+impl Expression {
+    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
+    pub(super) fn name(&self) -> String {
+        std::iter::once("Compiled")
+            .chain(self.statements.iter().map(|s| s.operator))
+            .collect()
+    }
+
+    pub(super) fn new(
+        inputs: &[Array],
+        statements: Vec<Statement>,
+        outputs: Vec<usize>,
+        constants: &[bool],
+    ) -> Result<Self, String> {
+        if inputs.len() != constants.len() || statements.is_empty() || outputs.is_empty() {
+            return Err("invalid compiled expression arity".into());
+        }
+        for (i, s) in statements.iter().enumerate() {
+            if s.inputs.is_empty() || s.inputs.iter().any(|&j| j >= inputs.len() + i) {
+                return Err("compiled expression is not topological".into());
+            }
+        }
+        if outputs
+            .iter()
+            .any(|&o| o < inputs.len() || o >= inputs.len() + statements.len())
+        {
+            return Err("invalid compiled expression output".into());
+        }
+        let inputs = inputs
+            .iter()
+            .zip(constants)
+            .map(|(a, &embed)| {
+                Ok(Input {
+                    dtype: a.layout().dtype(),
+                    scalar: a.layout().shape().is_empty(),
+                    constant: if embed { Some(constant(a)?) } else { None },
+                })
+            })
+            .collect::<Result<Vec<_>, String>>()?;
+        // Symbol renaming only; expression spelling and argument ordering
+        // follow the original emitter. Hash includes embedded constant text.
+        let hash = Sha256::digest(format!("{inputs:?}{statements:?}{outputs:?}"));
+        let hash = hash.iter().map(|b| format!("{b:02x}")).collect::<String>();
+        let prefix = format!("kernel_qwen_mtplx_compiled_{hash}");
+        Ok(Self {
+            inputs,
+            statements,
+            outputs,
+            prefix,
+        })
+    }
+
+    pub(super) fn evaluate(&self, inputs: &[Array], outputs: &[Array]) -> Result<(), String> {
+        super::compile_custom(&format!("{}_contiguous", self.prefix), || self.source())?;
+        compiled::evaluate_many(inputs, outputs, &self.prefix, |i| {
+            self.inputs[i].constant.is_some()
+        })
+    }
+
+    fn value_dtype(&self, slot: usize) -> Dtype {
+        if slot < self.inputs.len() {
+            self.inputs[slot].dtype
+        } else {
+            self.statements[slot - self.inputs.len()].dtype
+        }
+    }
+
+    pub(super) fn source(&self) -> Result<String, String> {
+        let mut source =
+            include_str!("../../../../metal/mtplx-qsa-compiled-header.metal").to_owned();
+        let work = (8 / self.value_dtype(self.outputs[0]).itemsize()).max(1);
+        source.push_str(&self.kernel(true, 0, false, false, 1)?);
+        if work > 1 {
+            source.push_str(&self.kernel(true, 0, false, false, work)?);
+        }
+        source.push_str(&self.kernel(true, 0, false, true, work)?);
+        for rank in 1..8 {
+            source.push_str(&self.kernel(
+                false,
+                rank,
+                false,
+                false,
+                if rank > 3 { 2 } else { 1 },
+            )?);
+            source.push_str(&self.kernel(
+                false,
+                rank,
+                false,
+                true,
+                if rank > 3 { 4 } else { 1 },
+            )?);
+        }
+        source.push_str(&self.kernel(false, 0, true, false, 2)?);
+        source.push_str(&self.kernel(false, 0, true, true, 4)?);
+        Ok(source)
+    }
+
+    fn kernel(
+        &self,
+        contiguous: bool,
+        rank: usize,
+        dynamic: bool,
+        large: bool,
+        work: usize,
+    ) -> Result<String, String> {
+        let suffix = if contiguous {
+            format!("contiguous{}", if work > 1 && !large { "_n" } else { "" })
+        } else {
+            format!(
+                "strided_{}",
+                if dynamic {
+                    "dynamic".to_owned()
+                } else {
+                    rank.to_string()
+                }
+            )
+        };
+        let name = format!(
+            "{}_{}{}",
+            self.prefix,
+            suffix,
+            if large { "_large" } else { "" }
+        );
+        let mut os = format!("[[host_name(\"{name}\")]]\n[[kernel]] void {name}(\n");
+        let mut names = Names {
+            slots: vec![None; self.inputs.len() + self.statements.len()],
+            count: 0,
+        };
+        let mut count = 0;
+        let mut indices = false;
+        for (i, input) in self.inputs.iter().enumerate() {
+            if input.constant.is_some() {
+                continue;
+            }
+            if !input.scalar && !contiguous {
+                indices = true;
+            }
+            writeln!(
+                os,
+                "    device const {}* {} [[buffer({count})]],",
+                dtype(input.dtype),
+                names.get(i)
+            )
+            .unwrap();
+            count += 1;
+        }
+        let index_type = if large { "int64_t" } else { "uint" };
+        if indices {
+            writeln!(
+                os,
+                "    constant const int64_t* in_strides [[buffer({count})]],"
+            )
+            .unwrap();
+            count += 1;
+        }
+        for &o in &self.outputs {
+            writeln!(
+                os,
+                "    device {}* {} [[buffer({count})]],",
+                dtype(self.value_dtype(o)),
+                names.get(o)
+            )
+            .unwrap();
+            count += 1;
+        }
+        if contiguous {
+            writeln!(
+                os,
+                "    constant const {index_type}& size [[buffer({count})]],"
+            )
+            .unwrap();
+        } else {
+            writeln!(
+                os,
+                "    constant const int* output_shape [[buffer({count})]],"
+            )
+            .unwrap();
+        }
+        count += 1;
+        if dynamic {
+            writeln!(os, "    constant const int& ndim [[buffer({count})]],").unwrap();
+            count += 1;
+        }
+        if count > 31 {
+            return Err("compiled expression exceeds Metal's 31 argument buffers".into());
+        }
+        os.push_str(
+            "    uint3 pos [[thread_position_in_grid]],\n    uint3 grid [[threads_per_grid]]) {\n",
+        );
+        writeln!(os, "  constexpr int N_ = {work};").unwrap();
+        if contiguous && large {
+            os.push_str("  int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));\n");
+        } else if contiguous {
+            os.push_str("  uint index = N_ * pos.x;\n");
+        } else if work > 1 {
+            writeln!(
+                os,
+                "  int xshape = output_shape[{}];",
+                if dynamic {
+                    "ndim - 1".to_owned()
+                } else {
+                    (rank - 1).to_string()
+                }
+            )
+            .unwrap();
+            writeln!(os, "  {index_type} index = N_ * pos.x + xshape * (pos.y + {index_type}(grid.y) * pos.z);").unwrap();
+        } else {
+            writeln!(
+                os,
+                "  {index_type} index = pos.x + grid.x * (pos.y + {index_type}(grid.y) * pos.z);"
+            )
+            .unwrap();
+        }
+        if work > 1 && contiguous {
+            os.push_str("  for (int i = 0; i < N_ && index < size; ++i) {\n");
+        }
+        let mut nc = Vec::new();
+        for (i, input) in self.inputs.iter().enumerate() {
+            let n = names.get(i);
+            if let Some(c) = &input.constant {
+                writeln!(
+                    os,
+                    "  auto tmp_{n} = static_cast<{}>({c});",
+                    dtype(input.dtype)
+                )
+                .unwrap();
+            } else if input.scalar {
+                writeln!(os, "  {} tmp_{n} = {n}[0];", dtype(input.dtype)).unwrap();
+            } else if contiguous {
+                writeln!(os, "  {} tmp_{n} = {n}[index];", dtype(input.dtype)).unwrap();
+            } else {
+                nc.push(i);
+            }
+        }
+        for (i, &slot) in nc.iter().enumerate() {
+            let n = names.get(slot);
+            write!(os, "  {index_type} index_{n} = ").unwrap();
+            match rank {
+                1 => writeln!(os, "elem_to_loc_1<{index_type}>(pos.x, in_strides[{}]);", i * rank).unwrap(),
+                2 => writeln!(os, "elem_to_loc_2<{index_type}>({{pos.x, pos.y}}, in_strides + {});", i * rank).unwrap(),
+                3 => writeln!(os, "elem_to_loc_3<{index_type}>(pos, in_strides + {});", i * rank).unwrap(),
+                _ if !dynamic => writeln!(os, "N_ * pos.x * {index_type}(in_strides[{}]) + pos.y * {index_type}(in_strides[{}]);", (i+1)*rank-1, (i+1)*rank-2).unwrap(),
+                _ => writeln!(os, "N_ * pos.x * {index_type}(in_strides[ndim * {i} + ndim - 1]) + pos.y * {index_type}(in_strides[ndim * {i} + ndim - 2]);").unwrap(),
+            }
+        }
+        if !nc.is_empty() && (rank > 3 || dynamic) {
+            os.push_str("  uint zpos = pos.z;\n");
+            writeln!(
+                os,
+                "  for (int d = {}; d >= 0; --d) {{",
+                if dynamic {
+                    "ndim - 3".to_owned()
+                } else {
+                    (rank - 3).to_string()
+                }
+            )
+            .unwrap();
+            os.push_str("    uint l = zpos % output_shape[d];\n");
+            for (i, &slot) in nc.iter().enumerate() {
+                let n = names.get(slot);
+                if dynamic {
+                    writeln!(
+                        os,
+                        "    index_{n} += l * {index_type}(in_strides[{i} * ndim + d]);"
+                    )
+                    .unwrap();
+                } else {
+                    writeln!(
+                        os,
+                        "    index_{n} += l * {index_type}(in_strides[{} + d]);",
+                        i * rank
+                    )
+                    .unwrap();
+                }
+            }
+            os.push_str("    zpos /= output_shape[d];\n  }\n");
+        }
+        if work > 1 && !contiguous {
+            os.push_str("  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {\n");
+        }
+        for &slot in &nc {
+            let n = names.get(slot);
+            writeln!(
+                os,
+                "  {} tmp_{n} = {n}[index_{n}];",
+                dtype(self.inputs[slot].dtype)
+            )
+            .unwrap();
+        }
+        for (i, s) in self.statements.iter().enumerate() {
+            let n = names.get(self.inputs.len() + i);
+            write!(os, "  {} tmp_{n} = ", dtype(s.dtype)).unwrap();
+            if matches!(s.operator, "Broadcast" | "AsType") {
+                writeln!(
+                    os,
+                    "cast_to<{}>(tmp_{});",
+                    dtype(s.dtype),
+                    names.get(s.inputs[0])
+                )
+                .unwrap();
+            } else {
+                let args = s
+                    .inputs
+                    .iter()
+                    .map(|&s| format!("tmp_{}", names.get(s)))
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                writeln!(os, "{}()({args});", s.operator).unwrap();
+            }
+        }
+        for &slot in &self.outputs {
+            let n = names.get(slot);
+            writeln!(os, "  {n}[index] = tmp_{n};").unwrap();
+        }
+        if work > 1 {
+            for (i, &slot) in nc.iter().enumerate() {
+                let n = names.get(slot);
+                if dynamic {
+                    writeln!(os, "  index_{n} += in_strides[{i} * ndim + ndim - 1];").unwrap();
+                } else {
+                    writeln!(os, "  index_{n} += in_strides[{}];", i * rank + rank - 1).unwrap();
+                }
+            }
+            os.push_str("  index++;\n  }\n");
+        }
+        os.push_str("}\n");
+        Ok(os)
+    }
+}
+
+#[test]
+fn mtplx_compiled_constant_spelling() {
+    for (value, text) in [
+        (0., "0"),
+        (-0., "-0"),
+        (0.1, "0.1"),
+        (1e-4, "0.0001"),
+        (1e-5, "1e-05"),
+        (1e7, "1e+07"),
+        (1234567., "1234567"),
+        (12345678., "1.234568e+07"),
+        (f64::NAN, "NAN"),
+        (f64::INFINITY, "INFINITY"),
+        (f64::NEG_INFINITY, "-INFINITY"),
+    ] {
+        assert_eq!(float_constant(value, 7), text);
+    }
+    assert_eq!(float_constant(1.2345678901234567, 16), "1.234567890123457");
+}
+
+#[test]
+fn mtplx_compiled_emitter_matches_all_pinned_silu_bodies() {
+    // The pinned shapeless SiLU has two explicit Broadcast primitives;
+    // compare the emitter, independently of shape-fixed trace selection.
+    let input = Array::unallocated(&[8], Dtype::BF16).unwrap();
+    let s = |operator, inputs| Statement {
+        dtype: Dtype::BF16,
+        operator,
+        inputs,
+    };
+    let expression = Expression::new(
+        &[input],
+        vec![
+            s("Sigmoid", vec![0]),
+            s("Broadcast", vec![0, 1]),
+            s("Broadcast", vec![1, 0]),
+            s("Multiply", vec![2, 3]),
+        ],
+        vec![4],
+        &[false],
+    )
+    .unwrap();
+    let source = expression.source().unwrap();
+    let reference = include_str!("../../../../metal/mtplx_qwen.metal");
+    let mut compared = 0;
+    for section in reference
+        .split("[[host_name(\"kernel_qwen_mtplx_silu_bf16_")
+        .skip(1)
+    {
+        let suffix = section.split_once('"').unwrap().0;
+        let original = format!("[[host_name(\"kernel_qwen_mtplx_silu_bf16_{section}");
+        let original = &original[..original.find("\n}\n").unwrap() + 3];
+        let symbol = original
+            .split_once("void ")
+            .unwrap()
+            .1
+            .split_once('(')
+            .unwrap()
+            .0;
+        let expected = original
+            .replace(symbol, &format!("{}_{suffix}", expression.prefix))
+            .replace(
+                &format!("kernel_qwen_mtplx_silu_bf16_{suffix}"),
+                &format!("{}_{suffix}", expression.prefix),
+            );
+        let marker = format!("[[host_name(\"{}_{suffix}\")]]", expression.prefix);
+        let start = source.find(&marker).unwrap();
+        let generated = &source[start..];
+        let generated = &generated[..generated.find("\n}\n").unwrap() + 3];
+        assert_eq!(generated, expected, "{suffix}");
+        compared += 1;
+    }
+    assert_eq!(compared, 19);
+}
+
+#[test]
+fn mtplx_compile_fusion_preserves_parent_and_output_boundaries() {
+    use super::{
+        binary::{Binary, binary},
+        compiled::Replay,
+        ops,
+        stream::{Device, Stream},
+        unary::{Unary, unary},
+    };
+    let gpu = Stream::new(0, Device::Gpu);
+    let other = Stream::new(1, Device::Gpu);
+    let x = Array::unallocated(&[4], Dtype::F32).unwrap();
+    let scalar = Array::unallocated(&[], Dtype::F32).unwrap();
+    let inputs = [x, scalar];
+    let results = Replay::trace_compiled(&inputs, |a| {
+        let square = unary(&a[0], Unary::Square, gpu)?;
+        let sum = binary(&square, &a[1], Binary::Add, gpu)?;
+        Ok(vec![square, sum.clone(), sum])
+    })
+    .unwrap()
+    .run(&inputs)
+    .unwrap();
+    assert!(matches!(
+        &*results[0].operation().unwrap(),
+        Operation::Compiled(_)
+    ));
+    assert_eq!(results[0].outputs().len(), 2);
+    assert_eq!(results[1].id(), results[2].id());
+    assert!(std::rc::Rc::ptr_eq(
+        &results[0].operation().unwrap(),
+        &results[1].operation().unwrap()
+    ));
+    // Mixed-parent Broadcast is split; the outside consumer retains its view.
+    let results = Replay::trace_compiled(&inputs, |a| {
+        let broadcast = ops::broadcast_to(&a[1], &[4], gpu)?;
+        let sum = binary(&a[0], &broadcast, Binary::Add, gpu)?;
+        Ok(vec![ops::reshape(&broadcast, &[2, 2], gpu)?, sum])
+    })
+    .unwrap()
+    .run(&inputs)
+    .unwrap();
+    assert!(matches!(
+        &*results[0].inputs()[0].operation().unwrap(),
+        Operation::Broadcast
+    ));
+    assert!(results[1].inputs().iter().any(|a| a.id() == inputs[1].id()));
+    // A non-Broadcast intermediate with outside parents cannot be split.
+    let results = Replay::trace_compiled(&inputs, |a| {
+        let square = unary(&a[0], Unary::Square, gpu)?;
+        Ok(vec![
+            ops::reshape(&square, &[2, 2], gpu)?,
+            binary(&square, &a[1], Binary::Add, gpu)?,
+        ])
+    })
+    .unwrap()
+    .run(&inputs)
+    .unwrap();
+    assert_eq!(results[0].inputs()[0].id(), results[1].inputs()[0].id());
+    // Stream mismatch is a fusion boundary even though CSE ignores streams.
+    let results = Replay::trace_compiled(&inputs, |a| {
+        let square = unary(&a[0], Unary::Square, gpu)?;
+        let negative = unary(&square, Unary::Negative, other)?;
+        Ok(vec![binary(&negative, &a[1], Binary::Add, other)?])
+    })
+    .unwrap()
+    .run(&inputs)
+    .unwrap();
+    assert_eq!(results[0].inputs()[0].stream(), Some(gpu));
+    // A 25-node chain splits at the original depth=11 boundary: 11+11+3.
+    let results = Replay::trace_compiled(&inputs, |a| {
+        let mut x = a[0].clone();
+        for _ in 0..25 {
+            x = unary(&x, Unary::Negative, gpu)?;
+        }
+        Ok(vec![x])
+    })
+    .unwrap()
+    .run(&inputs)
+    .unwrap();
+    let mut x = results[0].clone();
+    let mut lengths = Vec::new();
+    while let Some(op) = x.operation() {
+        let Operation::Compiled(expression) = &*op else {
+            panic!("unfused chain");
+        };
+        lengths.push(expression.statements.len());
+        let next = x.inputs()[0].clone();
+        x = next;
+    }
+    assert_eq!(lengths, [11, 11, 3]);
+    // Broadcast-only outputs remain views rather than materialized kernels.
+    let result = Replay::trace_compiled(&inputs, |a| {
+        Ok(vec![ops::broadcast_to(&a[1], &[4096], gpu)?])
+    })
+    .unwrap()
+    .run(&inputs)
+    .unwrap();
+    assert!(matches!(
+        &*result[0].operation().unwrap(),
+        Operation::Broadcast
+    ));
+}
+
+#[test]
+#[ignore = "requires Metal and external supervisor; generated fusion vs pinned SiLU plus multi-output/scalar dispatch"]
+fn mtplx_compiled_fusion_gpu_matches_reference() {
+    use super::{
+        super::{
+            gpu::Context,
+            qwen_mtplx_tests::{capture_dispatches, pattern},
+        },
+        allocator,
+        compiled::Replay,
+        configure_sources, gdn, mlp, ops,
+        stream::{Device, Streams},
+    };
+    use super::{
+        binary::{Binary, binary},
+        unary::{Unary, unary},
+    };
+    configure_sources().unwrap();
+    let _context = Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    for rank in [1, 4, 8] {
+        for reverse in [false, true] {
+            eprintln!("compiled SiLU rank={rank} reverse={reverse}");
+            let shape = vec![2; rank];
+            let input = Array::new(&shape, Dtype::BF16, pattern(1 << rank, 37)).unwrap();
+            let axes = (0..rank as i32)
+                .step_by(2)
+                .chain((1..rank as i32).step_by(2))
+                .collect::<Vec<_>>();
+            let input = ops::transpose(&input, &axes, gpu).unwrap();
+            let input = if reverse {
+                ops::slice(
+                    &input,
+                    &vec![1; rank],
+                    &vec![-3; rank],
+                    &vec![-1; rank],
+                    gpu,
+                )
+                .unwrap()
+            } else {
+                input
+            };
+            ops::evaluate(&streams, std::slice::from_ref(&input), gpu, false).unwrap();
+            let expected = mlp::silu(&input, gpu).unwrap();
+            let graph = Replay::trace_compiled(std::slice::from_ref(&input), |a| {
+                Ok(vec![mlp::silu(&a[0], gpu)?])
+            })
+            .unwrap();
+            let actual = graph.run(std::slice::from_ref(&input)).unwrap().remove(0);
+            let (_, calls) = capture_dispatches(|| {
+                ops::evaluate(&streams, std::slice::from_ref(&actual), gpu, false).unwrap()
+            });
+            assert_eq!(calls.len(), 1);
+            assert!(calls[0].0.starts_with("kernel_qwen_mtplx_compiled_"));
+            let actual = ops::contiguous(&actual, false, gpu).unwrap();
+            let expected = ops::contiguous(&expected, false, gpu).unwrap();
+            ops::evaluate(&streams, &[actual.clone(), expected.clone()], gpu, false).unwrap();
+            let mut a = vec![0; actual.layout().nbytes()];
+            let mut b = vec![0; expected.layout().nbytes()];
+            actual.buffer().read(actual.offset(), &mut a).unwrap();
+            expected.buffer().read(expected.offset(), &mut b).unwrap();
+            assert_eq!(a, b);
+        }
+    }
+    for rows in [1, 4, 17] {
+        eprintln!("compiled nested compute_g rows={rows}");
+        let args = [
+            Array::new(&[48], Dtype::BF16, pattern(48, 62)).unwrap(),
+            Array::new(&[1, rows, 48], Dtype::BF16, pattern(rows as u32 * 48, 60)).unwrap(),
+            Array::new(&[48], Dtype::BF16, pattern(48, 63)).unwrap(),
+        ];
+        let expected = gdn::compute_g(&args[0], &args[1], &args[2], gpu).unwrap();
+        let graph = Replay::trace_compiled(&args, |a| {
+            Ok(vec![gdn::compute_g(&a[0], &a[1], &a[2], gpu)?])
+        })
+        .unwrap();
+        let actual = graph.run(&args).unwrap().remove(0);
+        let (_, calls) = capture_dispatches(|| {
+            ops::evaluate(&streams, std::slice::from_ref(&actual), gpu, false).unwrap()
+        });
+        assert_eq!(calls.len(), 1, "{calls:?}");
+        assert!(calls[0].0.starts_with("kernel_qwen_mtplx_compiled_"));
+        ops::evaluate(&streams, std::slice::from_ref(&expected), gpu, false).unwrap();
+        let mut a = vec![0; actual.layout().nbytes()];
+        let mut b = vec![0; expected.layout().nbytes()];
+        actual.buffer().read(actual.offset(), &mut a).unwrap();
+        expected.buffer().read(expected.offset(), &mut b).unwrap();
+        assert_eq!(a, b, "nested compute_g rows={rows}");
+    }
+    for n in [8, 65536] {
+        for embedded in [false, true] {
+            eprintln!("compiled multiple outputs n={n} embedded={embedded}");
+            let data = (0..n)
+                .flat_map(|i| ((i % 8) as f32).to_le_bytes())
+                .collect::<Vec<_>>();
+            let input = Array::new(&[n], Dtype::F32, super::scalar_buffer(&data).unwrap()).unwrap();
+            let scalar = Array::new(
+                &[],
+                Dtype::F32,
+                super::scalar_buffer(&2f32.to_le_bytes()).unwrap(),
+            )
+            .unwrap();
+            let args = if embedded {
+                vec![input]
+            } else {
+                vec![input, scalar.clone()]
+            };
+            let graph = Replay::trace_compiled(&args, |a| {
+                let square = unary(&a[0], Unary::Square, gpu)?;
+                let sum = binary(
+                    &square,
+                    if embedded { &scalar } else { &a[1] },
+                    Binary::Add,
+                    gpu,
+                )?;
+                Ok(vec![square, sum])
+            })
+            .unwrap();
+            let roots = graph.run(&args).unwrap();
+            // The pinned evaluator selects its pipeline before allocating
+            // outputs (metal/compiled.cpp:384,419). Fresh ArrayDesc data_size
+            // is zero, even when the logical shape exceeds 65536 elements.
+            assert!(roots.iter().all(|a| a.data_size() == 0));
+            let (_, calls) =
+                capture_dispatches(|| ops::evaluate(&streams, &roots, gpu, false).unwrap());
+            assert_eq!(calls.len(), 1, "{calls:?}");
+            assert!(calls[0].0.ends_with("_contiguous"), "{calls:?}");
+            assert_eq!(calls[0].1, [n as u32, 1, 1]);
+            for (j, root) in roots.iter().enumerate() {
+                let mut bytes = vec![0; root.layout().nbytes()];
+                root.buffer().read(root.offset(), &mut bytes).unwrap();
+                for (i, b) in bytes.chunks_exact(4).enumerate() {
+                    assert_eq!(
+                        f32::from_le_bytes(b.try_into().unwrap()),
+                        ((i % 8) * (i % 8)) as f32 + (j * 2) as f32
+                    );
+                }
+            }
+        }
+    }
+}