Save inference parity implementation and evaluation harness

src/engine/metal/qwen_mtplx/moe.rs

@@ -0,0 +1,491 @@
+//! Qwen SparseMoeBlock as ordinary Array nodes, with the original two
+//! CustomKernel nodes only where MTPLX selects fused decode/verify.
+use super::array::{Array, Dtype, Layout};
+use super::binary::{Binary, binary};
+use super::mlp::{self, QuantizedLinear};
+use super::stream::{Device, Stream, Streams};
+use super::{indexing, linear, normalization, ops, reduce, sort, unary};
+
+pub(super) enum GateUp {
+    Fused(QuantizedLinear),
+    Separate {
+        gate: QuantizedLinear,
+        up: QuantizedLinear,
+    },
+}
+
+pub(super) struct FeedForward {
+    pub(super) input: GateUp,
+    pub(super) down: QuantizedLinear,
+}
+
+fn gather(
+    input: &Array,
+    layer: &QuantizedLinear,
+    ids: &Array,
+    sorted: bool,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    linear::gather_qmm_bf16(
+        input,
+        [&layer.weight, &layer.scales, &layer.biases],
+        [None, Some(ids)],
+        [layer.bits, layer.group],
+        sorted,
+        streams,
+        stream,
+    )
+}
+
+impl FeedForward {
+    fn shared(&self, input: &Array, stream: Stream) -> Result<Array, String> {
+        match &self.input {
+            GateUp::Fused(gu) => mlp::fused_gate_up_mlp(input, gu, &self.down, stream),
+            GateUp::Separate { gate, up } => {
+                // Qwen3NextMLP evaluates gate first, unlike SwitchGLU below.
+                let gate = gate.apply(input, stream)?;
+                let up = up.apply(input, stream)?;
+                self.down.apply(&mlp::swiglu(&gate, &up, stream)?, stream)
+            }
+        }
+    }
+
+    fn switched(
+        &self,
+        input: &Array,
+        indices: &Array,
+        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 ids = indices.clone();
+        let mut inverse = None;
+        if do_sort {
+            let size =
+                i32::try_from(indices.layout().size()).map_err(|_| "MoE index size overflow")?;
+            ids = ops::reshape(&ids, &[size], stream)?;
+            let order = sort::sort(&ids, Some(-1), true, None, stream)?;
+            inverse = Some(sort::sort(&order, Some(-1), true, None, stream)?);
+            let topk = *indices.layout().shape().last().ok_or("MoE indices rank")? as u32;
+            let buffer = super::super::gpu::Buffer::mtplx_bytes(4)?;
+            buffer.write(0, &topk.to_le_bytes())?;
+            let divisor = Array::new(&[], Dtype::U32, buffer)?;
+            let rows = super::binary::floor_divide_integer(&order, &divisor, stream)?;
+            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,
+            )?;
+            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 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)?;
+        if let Some(inverse) = inverse {
+            output = indexing::take(&output, &inverse, 0, stream)?;
+            let mut shape = indices.layout().shape().to_vec();
+            shape.extend_from_slice(&output.layout().shape()[1..]);
+            output = ops::reshape(&output, &shape, stream)?;
+        }
+        ops::squeeze(&output, Some(&[-2]), stream)
+    }
+}
+
+pub(super) struct SparseMoe {
+    pub(super) experts: FeedForward,
+    pub(super) shared: FeedForward,
+    pub(super) router: QuantizedLinear,
+    pub(super) shared_gate: QuantizedLinear,
+    pub(super) topk: i32,
+    pub(super) normalize: bool,
+    pub(super) fused_decode: bool,
+    pub(super) fused_verify: bool,
+}
+
+impl SparseMoe {
+    fn fused_rows(&self, input: &Array) -> Option<u32> {
+        let GateUp::Fused(gu) = &self.experts.input else {
+            return None;
+        };
+        if gu.bits != 4
+            || !matches!(gu.group, 32 | 64)
+            || self.experts.down.bits != 4
+            || !matches!(self.experts.down.group, 32 | 64)
+        {
+            return None;
+        }
+        let l = input.layout();
+        let shape = l.shape();
+        let rows = shape[shape.len() - 2];
+        let k = shape[shape.len() - 1];
+        if l.size() != rows as usize * k as usize {
+            return None;
+        }
+        ((rows == 1 && self.fused_decode) || ((2..=4).contains(&rows) && self.fused_verify))
+            .then_some(rows as u32)
+    }
+
+    pub(super) fn apply(
+        &self,
+        input: &Array,
+        streams: &Streams,
+        stream: Stream,
+    ) -> Result<Array, String> {
+        if stream.device() != Device::Gpu
+            || input.layout().dtype() != Dtype::BF16
+            || input.layout().shape().len() < 2
+            || input.layout().size() == 0
+        {
+            return Err("SparseMoe requires nonempty BF16 GPU rows".into());
+        }
+        let fused = self.fused_rows(input);
+        let gates = normalization::softmax_last(&self.router.apply(input, stream)?, true, stream)?;
+        let shape = gates.layout().shape().to_vec();
+        let rank = shape.len();
+        let experts = shape[rank - 1];
+        if self.topk < 1 || self.topk > experts {
+            return Err("invalid MoE top-k".into());
+        }
+        let partition = sort::sort(&gates, Some(-1), true, Some(-self.topk), stream)?;
+        let mut starts = vec![0; rank];
+        starts[rank - 1] = experts - self.topk;
+        let indices = ops::slice(&partition, &starts, &shape, &vec![1; rank], stream)?;
+        let mut scores = indexing::take_along_axis(&gates, &indices, -1, stream)?;
+        if self.normalize {
+            scores = binary(
+                &scores,
+                &reduce::sum(&scores, &[-1], true, stream)?,
+                Binary::Divide,
+                stream,
+            )?;
+        }
+        let y = if let Some(rows) = fused {
+            let GateUp::Fused(gu) = &self.experts.input else {
+                unreachable!()
+            };
+            fused_experts(
+                input,
+                gu,
+                &self.experts.down,
+                &indices,
+                &scores,
+                rows,
+                self.topk,
+                stream,
+            )?
+        } else {
+            let switched = self.experts.switched(input, &indices, streams, stream)?;
+            let scores = ops::expand_dims(&scores, &[-1], stream)?;
+            reduce::sum(
+                &binary(&switched, &scores, Binary::Multiply, stream)?,
+                &[-2],
+                false,
+                stream,
+            )?
+        };
+        let gate = || {
+            unary::unary(
+                &self.shared_gate.apply(input, stream)?,
+                unary::Unary::Sigmoid,
+                stream,
+            )
+        };
+        let (gate, shared) = if fused.is_some() {
+            (gate()?, self.shared.shared(input, stream)?)
+        } else {
+            let shared = self.shared.shared(input, stream)?;
+            (gate()?, shared)
+        };
+        let y = binary(
+            &y,
+            &binary(&gate, &shared, Binary::Multiply, stream)?,
+            Binary::Add,
+            stream,
+        )?;
+        if fused.is_some() {
+            ops::astype(&y, input.layout().dtype(), false, stream)
+        } else {
+            Ok(y)
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)] // Original two CustomKernel call inputs and template dimensions.
+fn fused_experts(
+    input: &Array,
+    gu: &QuantizedLinear,
+    down: &QuantizedLinear,
+    indices: &Array,
+    scores: &Array,
+    rows: u32,
+    topk: i32,
+    stream: Stream,
+) -> Result<Array, String> {
+    if topk != 10
+        || gu.weight.layout().shape().get(1) != Some(&1280)
+        || down.weight.layout().shape().get(1) != Some(&2560)
+        || input.layout().shape().last() != Some(&2560)
+    {
+        return Err("fused MoE exports require the installed top10/2560/640 shape".into());
+    }
+    let flatten = |a: &Array| ops::reshape(a, &[a.layout().size() as i32], stream);
+    let x = if rows == 1 {
+        flatten(input)?
+    } else {
+        ops::reshape(input, &[rows as i32, 2560], stream)?
+    };
+    let ids = ops::astype(&flatten(indices)?, Dtype::U32, false, stream)?;
+    let scores = ops::astype(&flatten(scores)?, Dtype::F32, false, stream)?;
+    let hshape = if rows == 1 {
+        vec![10, 640]
+    } else {
+        vec![rows as i32, 10, 640]
+    };
+    let h = Array::make_operation_with_inputs(
+        stream,
+        &[
+            x,
+            gu.weight.clone(),
+            gu.scales.clone(),
+            gu.biases.clone(),
+            ids.clone(),
+        ],
+        Layout::new(&hshape, Dtype::BF16)?,
+        ops::Operation::MoeGluH {
+            rows,
+            group: gu.group,
+        },
+    )?;
+    let yshape = if rows == 1 {
+        vec![2560]
+    } else {
+        vec![rows as i32, 2560]
+    };
+    let y = Array::make_operation_with_inputs(
+        stream,
+        &[
+            h,
+            down.weight.clone(),
+            down.scales.clone(),
+            down.biases.clone(),
+            ids,
+            scores,
+        ],
+        Layout::new(&yshape, Dtype::BF16)?,
+        ops::Operation::MoeDownY {
+            rows,
+            group: down.group,
+        },
+    )?;
+    ops::reshape(&y, input.layout().shape(), stream)
+}
+
+pub(super) fn fused_evaluate(
+    inputs: &[Array],
+    outputs: &[Array],
+    rows: u32,
+    group: u32,
+    down: bool,
+) -> Result<(), String> {
+    if !(1..=4).contains(&rows) || !matches!(group, 32 | 64) {
+        return Err("invalid fused MoE specialization".into());
+    }
+    let suffix = if rows == 1 {
+        format!("g{group}_bf16")
+    } else {
+        format!("m_m{rows}_g{group}_bf16")
+    };
+    let name = format!(
+        "kernel_qwen_mtplx_moe_{}_{suffix}",
+        if down { "down_y" } else { "glu_h" }
+    );
+    let count = rows * if down { 2560 } else { 10 * 640 };
+    super::gdn::custom_evaluate(inputs, outputs, |bindings| {
+        super::dispatch_geometry(
+            &name,
+            bindings,
+            &[],
+            [count.div_ceil(32) * 1024, 1, 1],
+            [1024, 1, 1],
+            true,
+        )
+    })
+}
+
+#[test]
+#[ignore = "requires Apple M5 Max; full MoE Array graph, no model weights; use test-supervisor"]
+fn mtplx_graph_sparse_moe_matches_reference() {
+    use super::super::{
+        gpu::Context,
+        qwen_mtplx_tests::{capture_dispatches, pattern, verify},
+    };
+    use super::{allocator, configure_sources};
+    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();
+    let fixtures = include_str!("../../../../tests/fixtures/mtplx-custom-kernels.jsonl")
+        .lines()
+        .map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
+        .filter(|row| row["kernel"].as_str().unwrap().starts_with("moe_"))
+        .collect::<Vec<_>>();
+    assert_eq!(fixtures.len(), 220);
+    for prepared in [false, true] {
+        for row in &fixtures {
+            let [tokens, experts, k, intermediate, gu_group, down_group] = [
+                "tokens",
+                "experts",
+                "k",
+                "intermediate",
+                "gu_group",
+                "down_group",
+            ]
+            .map(|key| row[key].as_u64().unwrap() as u32);
+            eprintln!("MoE Array graph prepared={prepared} {}", row["kernel"]);
+            let pack = |e: u32, n: u32, k: u32, group: u32, salt: u32, expert: bool| {
+                let mut q = mlp::fixture_linear(e * n, k, 4, group, salt);
+                if expert {
+                    for a in [&mut q.weight, &mut q.scales, &mut q.biases] {
+                        let width = *a.layout().shape().last().unwrap();
+                        *a = ops::reshape(a, &[e as i32, n as i32, width], gpu).unwrap();
+                    }
+                }
+                q
+            };
+            let ff = |expert: bool, separate: bool| {
+                let e = if expert { experts } else { 1 };
+                let salt = if expert { 39 } else { 46 };
+                FeedForward {
+                    input: if separate {
+                        GateUp::Separate {
+                            gate: pack(e, intermediate, k, gu_group, salt, expert),
+                            up: pack(
+                                e,
+                                intermediate,
+                                k,
+                                down_group,
+                                if expert { 45 } else { 54 },
+                                expert,
+                            ),
+                        }
+                    } else {
+                        GateUp::Fused(pack(e, 2 * intermediate, k, gu_group, salt, expert))
+                    },
+                    down: pack(
+                        e,
+                        k,
+                        intermediate,
+                        down_group,
+                        if expert { 41 } else { 48 },
+                        expert,
+                    ),
+                }
+            };
+            let model = SparseMoe {
+                experts: ff(true, row["separate_experts"].as_bool().unwrap_or(false)),
+                shared: ff(false, row["separate_shared"].as_bool().unwrap_or(false)),
+                router: mlp::fixture_linear(experts, k, 8, 64, 50),
+                shared_gate: mlp::fixture_linear(1, k, 4, 64, 52),
+                topk: 10,
+                normalize: row["normalize"].as_bool().unwrap(),
+                fused_decode: row["fused"].as_bool().unwrap_or(false),
+                fused_verify: row["fused"].as_bool().unwrap_or(false),
+            };
+            if prepared {
+                let mut parameters = Vec::new();
+                let mut add = |q: &QuantizedLinear| {
+                    parameters.extend([q.weight.clone(), q.scales.clone(), q.biases.clone()]);
+                };
+                for ff in [&model.experts, &model.shared] {
+                    match &ff.input {
+                        GateUp::Fused(gu) => add(gu),
+                        GateUp::Separate { gate, up } => {
+                            add(gate);
+                            add(up);
+                        }
+                    }
+                    add(&ff.down);
+                }
+                add(&model.router);
+                add(&model.shared_gate);
+                ops::evaluate(&streams, &parameters, gpu, false).unwrap();
+            }
+            let input = Array::new(
+                &[1, tokens as i32, k as i32],
+                Dtype::BF16,
+                pattern(tokens * k, 44),
+            )
+            .unwrap();
+            let output = model.apply(&input, &streams, gpu).unwrap();
+            let fused = row["fused"].as_bool().unwrap_or(false)
+                && !row["separate_experts"].as_bool().unwrap_or(false)
+                && tokens <= 4;
+            assert_eq!(model.fused_rows(&input).is_some(), fused);
+            // Graph edges, not an external model/parameter table, keep inputs alive.
+            drop((input, model));
+            let (_, dispatches) = capture_dispatches(|| {
+                ops::evaluate(&streams, std::slice::from_ref(&output), gpu, false).unwrap();
+            });
+            let actual = dispatches
+                .into_iter()
+                .filter(|call| call.0.starts_with("kernel_qwen_mtplx_moe_"))
+                .collect::<Vec<_>>();
+            let expected = if fused {
+                [
+                    ("glu_h", gu_group, tokens * 10 * 640),
+                    ("down_y", down_group, tokens * 2560),
+                ]
+                .map(|(name, group, rows)| {
+                    let suffix = if tokens == 1 {
+                        format!("g{group}_bf16")
+                    } else {
+                        format!("m_m{tokens}_g{group}_bf16")
+                    };
+                    (
+                        format!("kernel_qwen_mtplx_moe_{name}_{suffix}"),
+                        [rows.div_ceil(32) * 1024, 1, 1],
+                        [1024, 1, 1],
+                        true,
+                    )
+                })
+                .to_vec()
+            } else {
+                Vec::new()
+            };
+            assert_eq!(
+                actual, expected,
+                "original CustomKernel dispatch geometry/order"
+            );
+            verify(
+                row["kernel"].as_str().unwrap(),
+                &[(&output.buffer(), tokens * k)],
+            );
+        }
+    }
+    streams.clear_streams().unwrap();
+}