Port Qwen row-map SwiGLU expert gathers

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),
             };