Initial commit: E4B-MarkBase model integration with passing tests

- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
2026-06-23 18:12:35 +08:00
commit ac75faa0cc
301 changed files with 63426 additions and 0 deletions
@@ -0,0 +1,133 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// ════════════════════════════════════════════════════════
+// Kernel Fusion Optimizations - Reduce dispatch overhead
+// ════════════════════════════════════════════════════════
+
+// Use SIMD_WIDTH from OptimizedKernels.metal (already defined as uint = 4)
+
+// ── Fused RMS Norm + Quantized Matmul ────────────────
+// Combines norm and projection in single kernel
+// Saves 1 dispatch per layer (42 layers = 42 fewer dispatches)
+kernel void rms_norm_matmul_fused(
+    device const float *x     [[buffer(0)]],       // Input [inDim]
+    device const float *normW [[buffer(1)]],       // Norm weight [inDim]
+    device const uint *w      [[buffer(2)]],       // Packed weights [outDim, inDim/8]
+    device const float *s     [[buffer(3)]],       // Scales [outDim, inDim/64]
+    device const float *b     [[buffer(4)]],       // Biases [outDim, inDim/64]
+    device float *out         [[buffer(5)]],       // Output [outDim]
+    constant uint &inDim      [[buffer(6)]],
+    constant uint &outDim     [[buffer(7)]],
+    constant float &eps       [[buffer(8)]],
+    constant uint &groupSize  [[buffer(9)]],
+    threadgroup float *shared_norm_x [[threadgroup(0)]],  // Normed input cache
+    uint gid [[thread_position_in_grid]],
+    uint tid [[thread_position_in_threadgroup]],
+    uint tgSize [[threads_per_threadgroup]]
+) {
+    uint outRow = gid;
+    if (outRow >= outDim) return;
+    
+    // ── Phase 1: RMS Norm (cooperative) ───────────────────────
+    // Compute sum of squares in threadgroup
+    float localSum = 0.0;
+    for (uint i = tid; i < inDim; i += tgSize) {
+        float val = x[i];
+        localSum += val * val;
+    }
+    
+    // Parallel reduction (simplified - single threadgroup)
+    threadgroup float partial_sums[256];
+    partial_sums[tid] = localSum;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    
+    // Reduce to single sum
+    for (uint stride = tgSize/2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            partial_sums[tid] += partial_sums[tid + stride];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+    
+    // Compute RMS and normalize
+    float rms = rsqrt(partial_sums[0] / float(inDim) + eps);
+    
+    // Store normed values in threadgroup cache
+    for (uint i = tid; i < inDim; i += tgSize) {
+        shared_norm_x[i] = x[i] * rms * (normW ? normW[i] : 1.0);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    
+    // ── Phase 2: Quantized Matmul ─────────────────────────────
+    // Each thread processes one output row
+    uint numGroups = inDim / groupSize;
+    float sum = 0.0;
+    
+    for (uint g = 0; g < numGroups; g++) {
+        float scale = s[outRow * numGroups + g];
+        float bias = b[outRow * numGroups + g];
+        
+        uint packedBase = outRow * (inDim / 8) + g * (groupSize / 8);
+        
+        // SIMD processing (batch 2 packed values)
+        for (uint p = 0; p < 8; p += 2) {
+            uint packed0 = w[packedBase + p];
+            uint packed1 = w[packedBase + p + 1];
+            
+            uint xBase = g * groupSize + p * 8;
+            
+            float4 xVec0 = float4(
+                shared_norm_x[xBase + 0], shared_norm_x[xBase + 1],
+                shared_norm_x[xBase + 2], shared_norm_x[xBase + 3]
+            );
+            float4 xVec1 = float4(
+                shared_norm_x[xBase + 4], shared_norm_x[xBase + 5],
+                shared_norm_x[xBase + 6], shared_norm_x[xBase + 7]
+            );
+            float4 xVec2 = float4(
+                shared_norm_x[xBase + 8], shared_norm_x[xBase + 9],
+                shared_norm_x[xBase + 10], shared_norm_x[xBase + 11]
+            );
+            float4 xVec3 = float4(
+                shared_norm_x[xBase + 12], shared_norm_x[xBase + 13],
+                shared_norm_x[xBase + 14], shared_norm_x[xBase + 15]
+            );
+            
+            float4 qVec0 = float4(
+                float((packed0 >> 0) & 0xF) * scale + bias,
+                float((packed0 >> 4) & 0xF) * scale + bias,
+                float((packed0 >> 8) & 0xF) * scale + bias,
+                float((packed0 >> 12) & 0xF) * scale + bias
+            );
+            float4 qVec1 = float4(
+                float((packed0 >> 16) & 0xF) * scale + bias,
+                float((packed0 >> 20) & 0xF) * scale + bias,
+                float((packed0 >> 24) & 0xF) * scale + bias,
+                float((packed0 >> 28) & 0xF) * scale + bias
+            );
+            float4 qVec2 = float4(
+                float((packed1 >> 0) & 0xF) * scale + bias,
+                float((packed1 >> 4) & 0xF) * scale + bias,
+                float((packed1 >> 8) & 0xF) * scale + bias,
+                float((packed1 >> 12) & 0xF) * scale + bias
+            );
+            float4 qVec3 = float4(
+                float((packed1 >> 16) & 0xF) * scale + bias,
+                float((packed1 >> 20) & 0xF) * scale + bias,
+                float((packed1 >> 24) & 0xF) * scale + bias,
+                float((packed1 >> 28) & 0xF) * scale + bias
+            );
+            
+            sum += dot(qVec0, xVec0);
+            sum += dot(qVec1, xVec1);
+            sum += dot(qVec2, xVec2);
+            sum += dot(qVec3, xVec3);
+        }
+    }
+    
+    out[outRow] = sum;
+}
+
+// Note: batch_matmul_8 not possible in Metal - pointer arrays not supported as parameters
+// Alternative: Use Argument Buffer (Metal 2.0+) or separate dispatches