Metal MoE: memoryBarrier between phases (MLX pattern)

Applied MLX's Metal pattern:
- Single command buffer for all 3 phases (was 3 separate)
- memoryBarrierWithScope:MTLBarrierScopeBuffers between phases
- GPU-side sync instead of CPU waitUntilCompleted per phase
- Removed all NSLog debug overhead

35B no longer hangs with single cmdBuf. But per-layer dispatch
still slower than CPU fused IQ2 dot for MoE experts.
0.8B: 7.5 tok/s (Metal build). 35B: ~1 tok/s (loading overhead).

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

Author: unamedkr

src/backend/metal/tq_matmul.metal

@@ -516,3 +516,67 @@ kernel void matmul_q4_k(
         output[row] = total;
     }
 }
+
+
+/* ============================================================
+ * TurboQuant self Q4 matmul: block_size=32, 16 packed bytes + 1 float scale
+ * dequant: (nibble - 8) * scale
+ * Optimized: 4-byte unroll, SIMD reduce
+ * ============================================================ */
+kernel void matmul_tq_q4(
+    device const float*   input       [[buffer(0)]],
+    device float*         output      [[buffer(1)]],
+    device const uint8_t* weight_qs   [[buffer(2)]],
+    device const float*   weight_sc   [[buffer(3)]],
+    constant uint&        in_dim_u    [[buffer(4)]],
+    constant uint&        out_dim_u   [[buffer(5)]],
+    uint                  row         [[threadgroup_position_in_grid]],
+    uint                  tid         [[thread_index_in_threadgroup]],
+    uint                  tg_size     [[threads_per_threadgroup]])
+{
+    if (row >= out_dim_u) return;
+
+    const uint in_dim = in_dim_u;
+    const uint n_blocks = in_dim / 32;
+    const uint blocks_per_thread = (n_blocks + tg_size - 1) / tg_size;
+    const uint block_start = tid * blocks_per_thread;
+    const uint block_end = min(block_start + blocks_per_thread, n_blocks);
+
+    const uint qs_row = row * n_blocks * 16;
+    const uint sc_row = row * n_blocks;
+    float sum = 0.0f;
+
+    for (uint b = block_start; b < block_end; b++) {
+        const float sc = weight_sc[sc_row + b];
+        device const uint8_t* qs = weight_qs + qs_row + b * 16;
+        const uint base = b * 32;
+        for (uint k = 0; k < 16; k += 4) {
+            uint8_t p0 = qs[k], p1 = qs[k+1], p2 = qs[k+2], p3 = qs[k+3];
+            sum += (float(int(p0 & 0xF) - 8) * input[base + k]
+                 +  float(int(p0 >> 4)  - 8) * input[base + k + 16]
+                 +  float(int(p1 & 0xF) - 8) * input[base + k + 1]
+                 +  float(int(p1 >> 4)  - 8) * input[base + k + 17]
+                 +  float(int(p2 & 0xF) - 8) * input[base + k + 2]
+                 +  float(int(p2 >> 4)  - 8) * input[base + k + 18]
+                 +  float(int(p3 & 0xF) - 8) * input[base + k + 3]
+                 +  float(int(p3 >> 4)  - 8) * input[base + k + 19]) * sc;
+        }
+    }
+
+    sum += simd_shuffle_down(sum, 16);
+    sum += simd_shuffle_down(sum, 8);
+    sum += simd_shuffle_down(sum, 4);
+    sum += simd_shuffle_down(sum, 2);
+    sum += simd_shuffle_down(sum, 1);
+
+    threadgroup float simd_sums[8];
+    if (tid % 32 == 0) simd_sums[tid / 32] = sum;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tid == 0) {
+        uint n_simd = (tg_size + 31) / 32;
+        float total = 0.0f;
+        for (uint s = 0; s < n_simd; s++) total += simd_sums[s];
+        output[row] = total;
+    }
+}

src/backend/metal/tq_metal_dispatch.m

@@ -903,7 +903,20 @@ int tq_metal_moe_forward(
                                                               options:MTLResourceStorageModeShared];
         if (!params_buf) return -1;
 
-        /* --- Create command buffer and encoder --- */
+        /* --- Create output buffer for Phase 3 (allocated once with other buffers) --- */
+        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) {
+            /* Fallback to hybrid if buffer creation fails */
+            memcpy(hb_output, [gate_buf contents], inter_bytes);
+            return 1;
+        }
+
+        /* --- Single command buffer for all 3 phases (MLX pattern) ---
+         * Metal guarantees sequential execution of compute encoders within
+         * one command buffer. memoryBarrierWithScope ensures buffer writes
+         * from one encoder are visible to the next. */
         id<MTLCommandBuffer> cmdBuf = [tq_mtl_queue commandBuffer];
         if (!cmdBuf) return -1;
 
@@ -933,56 +946,13 @@ int tq_metal_moe_forward(
             MTLSize gridSize = MTLSizeMake(n_tgs, 1, 1);
             MTLSize tgSize   = MTLSizeMake(TQ_MATMUL_TG_SIZE, 1, 1);
             [enc dispatchThreadgroups:gridSize threadsPerThreadgroup:tgSize];
-            [enc endEncoding];
-        }
-
-        /* --- Phase 1: commit and wait to isolate hang --- */
-        [cmdBuf commit];
-        [cmdBuf waitUntilCompleted];
-
-        if (cmdBuf.status == MTLCommandBufferStatusError) {
-            NSLog(@"TurboQuant MoE: Phase 1 (gate+up) FAILED: %@", cmdBuf.error);
-            return -1;
-        }
-        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);
-            }
+            /* Memory barrier: ensure gate_buf/up_buf writes visible to Phase 2 */
+            [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+            [enc endEncoding];
         }
-#endif /* TQ_MOE_DEBUG_VALIDATE */
 
-        /* --- New command buffer for Phase 2 --- */
-        cmdBuf = [tq_mtl_queue commandBuffer];
-        if (!cmdBuf) return -1;
-
-        /* ======== Phase 2: SwiGLU ======== */
+        /* ======== Phase 2: SwiGLU (reads gate_buf/up_buf from Phase 1) ======== */
         {
             id<MTLComputeCommandEncoder> enc = [cmdBuf computeCommandEncoder];
             if (!enc) return -1;
@@ -998,40 +968,15 @@ int tq_metal_moe_forward(
             MTLSize gridSize = MTLSizeMake(n_tgs, 1, 1);
             MTLSize tgSize   = MTLSizeMake(tg, 1, 1);
             [enc dispatchThreadgroups:gridSize threadsPerThreadgroup:tgSize];
-            [enc endEncoding];
-        }
-
-        /* --- Phase 2: commit and wait to isolate hang --- */
-        [cmdBuf commit];
-        [cmdBuf waitUntilCompleted];
 
-        if (cmdBuf.status == MTLCommandBufferStatusError) {
-            NSLog(@"TurboQuant MoE: Phase 2 (SwiGLU) FAILED: %@", cmdBuf.error);
-            return -1;
+            /* Memory barrier: ensure gate_buf writes visible to Phase 3 */
+            [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+            [enc endEncoding];
         }
-        NSLog(@"TurboQuant MoE: Phase 2 (SwiGLU) completed OK");
 
         /* ======== Phase 3: down projection + weighted accumulate (GPU) ========
-         * Previously skipped due to IQ2_S shader hanging with constant array.
-         * Now fixed: IQ2_S codebook passed as device buffer (buffer 4). */
+         * IQ2_S codebook passed as device buffer (buffer 4). */
         {
-            /* Create output buffer for hidden_dim results */
-            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) {
-                /* Fallback to hybrid if buffer creation fails */
-                memcpy(hb_output, [gate_buf contents], inter_bytes);
-                return 1;
-            }
-
-            /* New command buffer for Phase 3 */
-            cmdBuf = [tq_mtl_queue commandBuffer];
-            if (!cmdBuf) {
-                memcpy(hb_output, [gate_buf contents], inter_bytes);
-                return 1;
-            }
-
             id<MTLComputeCommandEncoder> enc = [cmdBuf computeCommandEncoder];
             if (!enc) {
                 memcpy(hb_output, [gate_buf contents], inter_bytes);
@@ -1057,26 +1002,26 @@ int tq_metal_moe_forward(
             MTLSize tgSize3   = MTLSizeMake(TQ_MATMUL_TG_SIZE, 1, 1);
             [enc dispatchThreadgroups:gridSize3 threadsPerThreadgroup:tgSize3];
             [enc endEncoding];
+        }
 
-            [cmdBuf commit];
-            [cmdBuf waitUntilCompleted];
+        /* ONE commit + wait for all 3 phases */
+        [cmdBuf commit];
+        [cmdBuf waitUntilCompleted];
 
-            if (cmdBuf.status == MTLCommandBufferStatusError) {
-                NSLog(@"TurboQuant MoE: Phase 3 (down+accum) FAILED: %@", cmdBuf.error);
-                /* Fallback to hybrid on failure */
-                memcpy(hb_output, [gate_buf contents], inter_bytes);
-                return 1;
-            }
-            NSLog(@"TurboQuant MoE: Phase 3 (down+accum) completed OK");
+        if (cmdBuf.status == MTLCommandBufferStatusError) {
+            NSLog(@"TurboQuant MoE: GPU dispatch FAILED: %@", cmdBuf.error);
+            /* Fallback to hybrid on failure */
+            memcpy(hb_output, [gate_buf contents], inter_bytes);
+            return 1;
+        }
 
-            /* Copy result to output */
-            memcpy(output, [output_buf contents], output_bytes);
+        /* Copy result to output */
+        memcpy(output, [output_buf contents], output_bytes);
 
-            /* Also copy hb for potential caller use */
-            memcpy(hb_output, [gate_buf contents], inter_bytes);
+        /* Also copy hb for potential caller use */
+        memcpy(hb_output, [gate_buf contents], inter_bytes);
 
-            return 0; /* Full GPU success */
-        }
+        return 0; /* Full GPU success */
     }
 }
 

src/engine/tq_moe.c

@@ -669,10 +669,9 @@ void tq_moe_forward(const tq_moe_layer_t* layer,
         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);
 
-    /* Metal MoE: IQ2_S hang fixed! But per-phase waitUntilCompleted
-     * makes it slow. Need single command buffer (was 9.5 tok/s).
-     * Re-enable after merging single-cmdBuf dispatch. */
-    if (0 && tq_metal_moe_available() && num_active > 0) {
+    /* Metal MoE: single command buffer with memoryBarrier between phases.
+     * Eliminates per-phase waitUntilCompleted overhead. */
+    if (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;

src/engine/tq_transformer.c

@@ -1576,12 +1576,12 @@ float* tq_forward(tq_model_t* model, tq_state_t* s, int token, int pos) {
                 tq_matmul_q2_preq(s->hb2, layer->w_up_q2, layer->w_up_q2s,
                                    s->xb_q8, s->xb_q8s, c->intermediate_dim, dim);
             } else if (layer->w_gate_q4) {
-                tq_quantize_row_q8(s->xb, s->xb_q8, s->xb_q8s, dim);
-
-                tq_matmul_q4_preq(s->hb, layer->w_gate_q4, layer->w_gate_q4s,
-                                   s->xb_q8, s->xb_q8s, c->intermediate_dim, dim);
-                tq_matmul_q4_preq(s->hb2, layer->w_up_q4, layer->w_up_q4s,
-                                   s->xb_q8, s->xb_q8s, c->intermediate_dim, dim);
+                /* FFN gate+up: batch 2 matmuls on GPU if Metal available,
+                 * otherwise use Q4×Q8 preq fast path on CPU */
+                tq_metal_batch_begin_if_available();
+                tq_matmul_q4(s->hb, s->xb, layer->w_gate_q4, layer->w_gate_q4s, c->intermediate_dim, dim);
+                tq_matmul_q4(s->hb2, s->xb, layer->w_up_q4, layer->w_up_q4s, c->intermediate_dim, dim);
+                tq_metal_batch_flush_if_available();
             } else if (layer->gguf_w_gate) {
                 /* Batch gate+up into one GPU command buffer (2 matmuls, 1 dispatch) */
                 tq_metal_batch_begin_if_available();