Metal MoE: hybrid GPU/CPU implemented, disabled (per-layer sync too slow)

Phase 1+2 (gate+up+SwiGLU) on GPU works correctly.
Phase 3 (IQ2_S down) hangs — needs shader fix.
Hybrid approach: GPU Phase 1+2 + CPU down — functional but slow
due to per-layer Metal commit+wait synchronization overhead.

Metal infrastructure fully in place for future IQ2_S fix.
CPU fused IQ2 dot remains fastest: 3.7 tok/s (100 tokens).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: unamedkr

include/turboquant/tq_gguf.h

@@ -289,23 +289,31 @@ void tq_moe_advise(const tq_moe_layer_t* layer,
 /* ============================================================
  * Fused MoE Metal GPU dispatch
  *
- * Processes ALL active experts for one MoE layer in 3 GPU
- * dispatches instead of per-expert per-matmul dispatch:
- *   Phase 1: gate + up projections (all experts, parallel)
- *   Phase 2: SwiGLU activation (all experts, parallel)
- *   Phase 3: down projection + weighted accumulate
+ * Hybrid GPU/CPU MoE dispatch:
+ *   Phase 1: gate + up projections on GPU (all experts, parallel)
+ *   Phase 2: SwiGLU activation on GPU (all experts, parallel)
+ *   Phase 3: down projection + weighted accumulate on CPU
+ *            (IQ2_S shader hangs on Metal; CPU fallback is reliable)
  *
- * Requires Apple Silicon with Metal. Returns -1 if unavailable.
+ * GPU handles the larger gate+up matmuls (~2/3 of compute).
+ * Returns 1 = partial (hb_output filled, caller does down+accum on CPU).
+ * Returns -1 if unavailable.
  * ============================================================ */
 
 /* Check if fused MoE Metal dispatch is available */
 int tq_metal_moe_available(void);
 
-/* Fused MoE forward: 3 Metal dispatches for all active experts.
- * Returns 0 on success, -1 on failure (caller falls back to CPU). */
+/* Fused MoE forward: GPU handles gate+up+SwiGLU (Phases 1+2).
+ * Returns:
+ *   0  = full success (all phases on GPU, output[] filled)
+ *   1  = partial success: gate+up+SwiGLU done on GPU,
+ *         hb_output[] filled with [num_active * expert_dim] SwiGLU'd activations,
+ *         caller must do down projection + accumulate on CPU.
+ *  -1  = failure (caller falls back to full CPU path). */
 int tq_metal_moe_forward(
     const float*    input,
     float*          output,
+    float*          hb_output,      /* [num_active * expert_dim] — GPU writes SwiGLU results here */
     const void*     weight_base,
     size_t          weight_size,
     const uint64_t* gate_offsets,

src/backend/metal/tq_metal_dispatch.m

@@ -756,10 +756,12 @@ int tq_metal_moe_available(void) {
 }
 
 /**
- * Fused MoE forward: 3 dispatches for all active experts.
+ * Hybrid GPU/CPU MoE forward: GPU does gate+up+SwiGLU (Phases 1+2),
+ * CPU handles down projection (Phase 3 IQ2_S shader hangs on Metal).
  *
  * @param input        Input hidden state [hidden_dim] (CPU pointer, will be uploaded)
- * @param output       Output hidden state [hidden_dim] (CPU pointer, will be downloaded)
+ * @param output       Output hidden state [hidden_dim] (unused in hybrid mode)
+ * @param hb_output    SwiGLU'd activations [num_active * expert_dim] (filled by GPU)
  * @param weight_base  Base pointer for all expert weights (mmap'd GGUF data)
  * @param weight_size  Total size of weight region (for zero-copy buffer)
  * @param gate_offsets Byte offsets from weight_base to each active expert's gate weights [num_active]
@@ -775,11 +777,13 @@ int tq_metal_moe_available(void) {
  * @param gate_types_in  Per-expert gate quant types [num_active] (NULL = use weight_type for all)
  * @param up_types_in    Per-expert up quant types [num_active] (NULL = use weight_type for all)
  * @param down_types_in  Per-expert down quant types [num_active] (NULL = use weight_type for all)
- * @return 0 on success, -1 on failure
+ * @return 1 on partial success (hb_output filled, caller does down+accum on CPU),
+ *        -1 on failure
  */
 int tq_metal_moe_forward(
     const float*    input,
     float*          output,
+    float*          hb_output,
     const void*     weight_base,
     size_t          weight_size,
     const uint64_t* gate_offsets,
@@ -812,11 +816,21 @@ int tq_metal_moe_forward(
         params.num_experts_total = num_experts_total;
         params.gate_type = (uint32_t)weight_type;
 
-        /* Default block geometry uses IQ2_XXS (66 bytes per 256 elements) */
+        /* Block geometry: blocks_per_row is type-independent (always 256 elements/block).
+         * row_bytes_gate/down use IQ2_XXS (66) as default; shader overrides per-expert
+         * for IQ2_S (82) via the gate_types/down_types arrays. */
         params.blocks_per_row_gate = hidden_dim / 256;
-        params.row_bytes_gate = params.blocks_per_row_gate * 66;
         params.blocks_per_row_down = expert_dim / 256;
-        params.row_bytes_down = params.blocks_per_row_down * 66;
+
+        /* Determine default row bytes from first expert's type */
+        int first_gate_type = gate_types_in ? gate_types_in[0] : weight_type;
+        int first_down_type = down_types_in ? down_types_in[0] : weight_type;
+        int gate_blk_bytes = (first_gate_type == 22) ? 82 : 66; /* IQ2_S=82, IQ2_XXS=66 */
+        int down_blk_bytes = (first_down_type == 22) ? 82 : 66;
+        params.row_bytes_gate = params.blocks_per_row_gate * gate_blk_bytes;
+        params.row_bytes_down = params.blocks_per_row_down * down_blk_bytes;
+        params.gate_row_bytes = params.row_bytes_gate;
+        params.down_row_bytes = params.row_bytes_down;
 
         for (int k = 0; k < num_active; k++) {
             params.active_experts[k] = active_expert_ids[k];
@@ -850,11 +864,8 @@ int tq_metal_moe_forward(
                                                             options:MTLResourceStorageModeShared];
         if (!gate_buf || !up_buf) return -1;
 
-        /* --- Create output buffer --- */
-        size_t output_bytes = (size_t)hidden_dim * sizeof(float);
-        id<MTLBuffer> output_buf = [tq_mtl_device newBufferWithLength:output_bytes
-                                                               options:MTLResourceStorageModeShared];
-        if (!output_buf) return -1;
+        /* output_buf removed: hybrid mode returns SwiGLU results via hb_output,
+         * down projection + accumulate happens on CPU in tq_moe.c */
 
         /* --- Create params buffer --- */
         id<MTLBuffer> params_buf = [tq_mtl_device newBufferWithBytes:&params
@@ -866,16 +877,10 @@ int tq_metal_moe_forward(
         id<MTLCommandBuffer> cmdBuf = [tq_mtl_queue commandBuffer];
         if (!cmdBuf) return -1;
 
-        /* Shared memory sizes:
-         * Phase 1 (gate_up): hidden_dim floats for input + 8 floats for SIMD sums
-         * Phase 3 (down):    expert_dim floats for hb + 8 floats for SIMD sums */
+        /* Shared memory for Phase 1 (gate_up): hidden_dim floats for input + 8 for SIMD sums */
         NSUInteger shared_phase1 = ((NSUInteger)hidden_dim + 8) * sizeof(float);
-        NSUInteger shared_phase3 = ((NSUInteger)expert_dim + 8) * sizeof(float);
         NSUInteger max_shared = [tq_mtl_device maxThreadgroupMemoryLength];
-
-        /* Cap shared memory to device limit */
         if (shared_phase1 > max_shared) shared_phase1 = max_shared;
-        if (shared_phase3 > max_shared) shared_phase3 = max_shared;
 
         /* ======== Phase 1: gate + up (fused) ======== */
         {
@@ -908,6 +913,38 @@ int tq_metal_moe_forward(
         }
         NSLog(@"TurboQuant MoE: Phase 1 (gate+up) completed OK");
 
+#ifdef TQ_MOE_DEBUG_VALIDATE
+        /* === Debug: compare GPU gate output for expert 0 vs CPU tq_matmul_gguf === */
+        {
+            /* tq_matmul_gguf declared in tq_gguf.h (already included) */
+            float* gpu_gate = (float*)[gate_buf contents];
+            float* cpu_gate = (float*)malloc((size_t)expert_dim * sizeof(float));
+            if (cpu_gate) {
+                /* CPU matmul for expert 0's gate weights */
+                const uint8_t* gate_w = (const uint8_t*)weight_base + gate_offsets[0];
+                tq_ggml_dtype gt0 = gate_types_in ? (tq_ggml_dtype)gate_types_in[0]
+                                                  : (tq_ggml_dtype)weight_type;
+                tq_matmul_gguf(cpu_gate, input, gate_w, gt0, expert_dim, hidden_dim);
+
+                /* Compare first 8 and last 8 values */
+                NSLog(@"TurboQuant MoE DEBUG: gate expert 0 comparison (first 8):");
+                float max_err = 0.0f;
+                for (int i = 0; i < expert_dim; i++) {
+                    float err = fabsf(gpu_gate[i] - cpu_gate[i]);
+                    if (err > max_err) max_err = err;
+                    if (i < 8 || i >= expert_dim - 4) {
+                        NSLog(@"  [%d] GPU=%.6f CPU=%.6f err=%.6f", i, gpu_gate[i], cpu_gate[i], err);
+                    }
+                }
+                NSLog(@"TurboQuant MoE DEBUG: gate max_err=%.6f across %d elements", max_err, expert_dim);
+                if (max_err > 0.01f) {
+                    NSLog(@"TurboQuant MoE DEBUG: *** MISMATCH DETECTED *** — weight offset or decoding bug");
+                }
+                free(cpu_gate);
+            }
+        }
+#endif /* TQ_MOE_DEBUG_VALIDATE */
+
         /* --- New command buffer for Phase 2 --- */
         cmdBuf = [tq_mtl_queue commandBuffer];
         if (!cmdBuf) return -1;
@@ -941,42 +978,13 @@ int tq_metal_moe_forward(
         }
         NSLog(@"TurboQuant MoE: Phase 2 (SwiGLU) completed OK");
 
-        /* --- New command buffer for Phase 3 --- */
-        cmdBuf = [tq_mtl_queue commandBuffer];
-        if (!cmdBuf) return -1;
-
-        /* ======== Phase 3: down projection + weighted accumulation ======== */
-        {
-            id<MTLComputeCommandEncoder> enc = [cmdBuf computeCommandEncoder];
-            if (!enc) return -1;
-
-            [enc setComputePipelineState:tq_pipe_moe_down_accum];
-            [enc setBuffer:weight_buf  offset:0 atIndex:0];
-            [enc setBuffer:gate_buf    offset:0 atIndex:1];  /* hb_all (post-SwiGLU) */
-            [enc setBuffer:output_buf  offset:0 atIndex:2];
-            [enc setBuffer:params_buf  offset:0 atIndex:3];
-            [enc setThreadgroupMemoryLength:shared_phase3 atIndex:0];
-
-            /* One threadgroup per output row */
-            NSUInteger n_tgs = (NSUInteger)hidden_dim;
-            MTLSize gridSize = MTLSizeMake(n_tgs, 1, 1);
-            MTLSize tgSize   = MTLSizeMake(TQ_MATMUL_TG_SIZE, 1, 1);
-            [enc dispatchThreadgroups:gridSize threadsPerThreadgroup:tgSize];
-            [enc endEncoding];
-        }
-
-        /* --- Phase 3: commit and wait --- */
-        [cmdBuf commit];
-        [cmdBuf waitUntilCompleted];
-
-        if (cmdBuf.status == MTLCommandBufferStatusError) {
-            NSLog(@"TurboQuant MoE: Phase 3 (down+accum) FAILED: %@", cmdBuf.error);
-            return -1;
-        }
-        NSLog(@"TurboQuant MoE: Phase 3 (down+accum) completed OK");
-
-        memcpy(output, [output_buf contents], output_bytes);
-        return 0;
+        /* ======== Phase 3: SKIP GPU down projection (IQ2_S shader hangs) ========
+         * Instead, copy the SwiGLU'd activations from gate_buf back to CPU.
+         * The caller (tq_moe.c) will do down projection + accumulate on CPU. */
+        memcpy(hb_output, [gate_buf contents], inter_bytes);
+        NSLog(@"TurboQuant MoE: Hybrid — Phase 1+2 on GPU done, returning hb for CPU down proj");
+        (void)output; /* unused in hybrid mode */
+        return 1; /* partial success: hb_output filled, caller does down+accum */
     }
 }
 

src/engine/tq_moe.c

@@ -662,16 +662,17 @@ void tq_moe_forward(const tq_moe_layer_t* layer,
 #ifdef TQ_HAS_METAL
     extern int tq_metal_moe_available(void);
     extern int tq_metal_moe_forward(
-        const float* input, float* output,
+        const float* input, float* output, float* hb_output,
         const void* weight_base, size_t weight_size,
         const uint64_t* gate_offsets, const uint64_t* up_offsets, const uint64_t* down_offsets,
         const int* active_expert_ids, const float* expert_routing_weights,
         int num_active, int expert_dim, int hidden_dim, int num_experts_total, int weight_type,
         const int* gate_types, const int* up_types, const int* down_types);
 
-    /* Fused Metal MoE — re-enabled for per-phase hang isolation debugging.
-     * Each phase runs with waitUntilCompleted + NSLog to identify which phase hangs. */
-    if (tq_metal_moe_available() && num_active > 0) {
+    /* Metal MoE disabled: IQ2_S Phase 3 hangs, hybrid too slow.
+     * CPU fused IQ2 dot at 3.7 tok/s is the current fastest path.
+     * GPU acceleration needs IQ2_S shader fix or Q4_K model format. */
+    if (0 && tq_metal_moe_available() && num_active > 0) {
         /* Check that all active experts use IQ2_XXS and have valid weights */
         int can_fuse = 1;
         const void* base_ptr = NULL;
@@ -743,8 +744,12 @@ void tq_moe_forward(const tq_moe_layer_t* layer,
                 per_down_types[k] = (int)exp->down_type;
             }
 
+            /* Allocate buffer for GPU SwiGLU results [num_active * expert_dim] */
+            float* hb_gpu = (float*)malloc((size_t)num_active * (size_t)expert_dim * sizeof(float));
+            if (!hb_gpu) goto moe_cpu_fallback;
+
             int rc = tq_metal_moe_forward(
-                input, output,
+                input, output, hb_gpu,
                 base_ptr, base_size,
                 gate_offs, up_offs, down_offs,
                 expert_ids, routing_w,
@@ -753,15 +758,39 @@ void tq_moe_forward(const tq_moe_layer_t* layer,
                 0, /* weight_type=0 means use per-expert types */
                 per_gate_types, per_up_types, per_down_types);
 
+            if (rc == 1) {
+                /* Hybrid: GPU did gate+up+SwiGLU, we do down+accum on CPU */
+                for (int k = 0; k < num_active; k++) {
+                    int eid = state->top_experts[k];
+                    float w = state->expert_weights[k];
+                    const tq_expert_weights_t* exp = &layer->experts[eid];
+                    float* hb_k = hb_gpu + k * expert_dim;
+
+                    /* down = hb_k @ w_down^T -> [hidden_dim] */
+                    tq_matmul_gguf(state->expert_out, hb_k,
+                                   exp->w_down, exp->down_type,
+                                   hidden_dim, expert_dim);
+
+                    /* Weighted accumulation: output += weight * down_proj */
+                    for (int i = 0; i < hidden_dim; i++)
+                        output[i] += w * state->expert_out[i];
+                }
+                free(hb_gpu);
+                goto moe_shared_expert;
+            }
+            free(hb_gpu);
             if (rc == 0) {
-                /* Fused MoE succeeded — skip to shared expert */
+                /* Full GPU success (unlikely in hybrid mode, but handle it) */
                 goto moe_shared_expert;
             }
-            /* else: fall through to per-expert CPU path */
+            /* else: rc == -1, fall through to per-expert CPU path */
         }
     }
 #endif /* TQ_HAS_METAL */
 
+#ifdef TQ_HAS_METAL
+moe_cpu_fallback: ;
+#endif
     /* Step 3: For each selected expert, compute SwiGLU FFN and accumulate */
     for (int k = 0; k < num_active; k++) {
         int eid = state->top_experts[k];