Gemma MOE support

.devin/config.local.json

@@ -0,0 +1,12 @@
+{
+  "permissions": {
+    "allow": [
+      "Exec(git status)",
+      "Exec(gh pr)",
+      "Exec(gh run)",
+      "Exec(cargo fmt)",
+      "Exec(git diff)",
+      "Exec(cargo clippy)"
+    ]
+  }
+}

fusor-ml/core/src/quantized/matmul/mod.rs

@@ -911,8 +911,9 @@ fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32,
     if format.is_q6k_family() && n <= 4096 && k > 4096 {
         return 4; // was Q6KLargeNarrow4
     }
-    // Q8_0 wide.
-    if format.is_q8_0_family() && k <= 1024 && n >= 8192 {
+    // Q8_0 wide. Keep this aligned with qgemv_tile_with_epilogue, which
+    // switches Q8_0/Q8_0Native to 4x8 output columns for wide outputs.
+    if format.is_q8_0_family() && n >= 8192 {
         return 32; // was Q8WideAccelerated32
     }
     // FormatAccelerated: Q5_0 mid (K,N both 2048..=4096), Q4K/Q6K general,
@@ -929,7 +930,10 @@ fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32,
     if format.is_q5_0_family() && k <= 1024 && n <= 4096 {
         return 8; // was Q5Small8
     }
-    4 // was Default4
+    // The shader-side qgemv shape table is the source of truth for generic
+    // formats. Returning a smaller value here over-dispatches masked columns,
+    // and every extra workgroup still pays the full K-loop cost.
+    tile_ir_kernels::qgemv_cols_per_workgroup_for_shape(format, k, n)
 }
 
 fn qmatmul_m_pad_target_for_caps(m: usize, n: usize, caps: KernelDeviceCaps) -> Option<usize> {

fusor-ml/core/src/quantized/mod.rs

@@ -255,21 +255,26 @@ fn qmatrix_storage_layout_for_parts_with_env(
     // the raw-word ggml qgemv dot (`qgemv_q6k_ggml`) cannot address. The
     // f32-scale layout is 212 bytes (word-aligned) and feeds that amortized
     // decode; the +2 bytes/block (~0.9%) is paid back many times over by the
-    // faster kernel. Q4K/Q5K keep their native f16-scale layout (their native
-    // blocks are word-aligned, and native f16 scales read less memory).
+    // faster kernel.
     if ty == GgmlType::Q6K {
         return QMatrixStorageLayout::GpuF32Scales;
     }
-    if matches!(
-        ty,
-        GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K
-    ) && native_half_scale_storage_enabled(shader_f16_supported, env_override)
+    // The small 32-element block formats save only two bytes/block with native
+    // f16 scales, but native Q4_0/Q5_0/Q8_0 blocks are byte-addressed. Decode is
+    // much faster with word-aligned f32-scale storage on Apple GPUs.
+    if matches!(ty, GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0) {
+        if env_override.is_some()
+            && native_half_scale_storage_enabled(shader_f16_supported, env_override)
+        {
+            QMatrixStorageLayout::Native
+        } else {
+            QMatrixStorageLayout::GpuF32Scales
+        }
+    } else if matches!(ty, GgmlType::Q4K | GgmlType::Q5K)
+        && native_half_scale_storage_enabled(shader_f16_supported, env_override)
     {
         QMatrixStorageLayout::Native
-    } else if matches!(
-        ty,
-        GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K
-    ) {
+    } else if matches!(ty, GgmlType::Q4K | GgmlType::Q5K) {
         QMatrixStorageLayout::GpuF32Scales
     } else {
         QMatrixStorageLayout::Native
@@ -536,14 +541,22 @@ mod tests {
     use super::*;
 
     #[test]
-    fn native_capable_quant_storage_defaults_to_native_when_shader_f16_is_supported() {
-        for ty in [
-            GgmlType::Q4_0,
-            GgmlType::Q5_0,
-            GgmlType::Q8_0,
-            GgmlType::Q4K,
-            GgmlType::Q5K,
-        ] {
+    fn small_block_quant_storage_defaults_to_f32_scales() {
+        for ty in [GgmlType::Q4_0, GgmlType::Q5_0, GgmlType::Q8_0] {
+            assert_eq!(
+                qmatrix_storage_layout_for_parts_with_env(ty, true, None),
+                QMatrixStorageLayout::GpuF32Scales
+            );
+            assert_eq!(
+                qmatrix_storage_layout_for_parts_with_env(ty, false, None),
+                QMatrixStorageLayout::GpuF32Scales
+            );
+        }
+    }
+
+    #[test]
+    fn k_quant_storage_defaults_to_native_when_shader_f16_is_supported() {
+        for ty in [GgmlType::Q4K, GgmlType::Q5K] {
             assert_eq!(
                 qmatrix_storage_layout_for_parts_with_env(ty, true, None),
                 QMatrixStorageLayout::Native
@@ -564,7 +577,15 @@ mod tests {
     }
 
     #[test]
-    fn q4k_storage_env_can_force_native_or_f32_expanded() {
+    fn quant_storage_env_can_force_native_or_f32_expanded() {
+        assert_eq!(
+            qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4_0, true, Some("1")),
+            QMatrixStorageLayout::Native
+        );
+        assert_eq!(
+            qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4_0, true, Some("0")),
+            QMatrixStorageLayout::GpuF32Scales
+        );
         assert_eq!(
             qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4K, false, Some("1")),
             QMatrixStorageLayout::Native

fusor-ml/core/src/row_program.rs

@@ -46,11 +46,6 @@ use crate::{
 
 const BLOCK: u32 = 256;
 
-/// Below this many rows a long axis is split across workgroups (one tile
-/// each) with a combine kernel folding the spans — decode has too few rows
-/// to fill the device with one workgroup per row.
-const SPLIT_ROWS_TARGET: u32 = 256;
-
 /// Workgroup buckets for dynamic-axis row programs, smallest first. The
 /// kernel monomorphizes per bucket; the active axis length rides in the
 /// params input.
@@ -606,10 +601,12 @@ fn build_row_program_kernel(
     let block = workgroup_shape.x();
     let lanes_own_axis = operation.dynamic_axis.is_some();
 
-    // Long axes with few rows fan out across workgroups: each split runs
-    // the online body over one tile, writing its unnormalized accumulator
-    // and softmax statistics to scratch; a combine kernel folds the spans
-    // with the online monoid.
+    // Long reducing outputs fan out across workgroups: each split owns one
+    // axis tile, writes its unnormalized accumulator and softmax statistics
+    // to scratch, then a combine kernel folds the spans with the online
+    // monoid. This avoids the single-workgroup multi-tile loop for attention
+    // prefill, which is both slower to fill the GPU and less robust at Gemma's
+    // unscaled vision-attention logits.
     let output_kind = operation.output_step().clone();
     let phase_steps = operation.phase_steps().to_vec();
     let free_dim_out = match &output_kind {
@@ -619,11 +616,7 @@ fn build_row_program_kernel(
     let tiles = k.div_ceil(block);
     let splits: u32 = match free_dim_out {
         Some(free)
-            if lanes_own_axis
-                && rows < SPLIT_ROWS_TARGET
-                && tiles > 1
-                && tiles <= block
-                && (free as u32 + 2) <= block =>
+            if lanes_own_axis && tiles > 1 && tiles <= block && (free as u32 + 2) <= block =>
         {
             tiles
         }