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,154 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// ═══════════════════════════════════════════════════════════════
+// Batch Metal Kernels - Process multiple tokens simultaneously
+// ═══════════════════════════════════════════════════════════════
+
+// Batch quantized matmul - process N tokens with shared weights
+kernel void quantized_matmul_batch(
+    device float* batchInput [[buffer(0)]],     // [batchSize, inDim]
+    device uint8_t* weights [[buffer(1)]],      // [outDim, inDim] packed
+    device float* scales [[buffer(2)]],         // [outDim, groups]
+    device float* biases [[buffer(3)]],         // [outDim]
+    device float* batchOutput [[buffer(4)]],    // [batchSize, outDim]
+    constant uint32_t& inDim [[buffer(5)]],
+    constant uint32_t& outDim [[buffer(6)]],
+    constant uint32_t& groupSize [[buffer(7)]],
+    constant uint32_t& batchSize [[buffer(8)]],
+    uint3 gid [[thread_position_in_grid]])
+{
+    uint batchIdx = gid.x;
+    uint outIdx = gid.y;
+    
+    if (batchIdx >= batchSize || outIdx >= outDim) return;
+    
+    device float* input = batchInput + batchIdx * inDim;
+    float sum = biases[outIdx];
+    uint groupIdx = outIdx * (inDim / groupSize);
+    
+    for (uint i = 0; i < inDim; i += 4) {
+        float4 inVals = float4(input[i], input[i+1], input[i+2], input[i+3]);
+        
+        uint packedWeight = weights[outIdx * inDim + i];
+        uint8_t w0 = (packedWeight >> 0) & 0xFF;
+        uint8_t w1 = (packedWeight >> 8) & 0xFF;
+        uint8_t w2 = (packedWeight >> 16) & 0xFF;
+        uint8_t w3 = (packedWeight >> 24) & 0xFF;
+        
+        uint g0 = (i + 0) / groupSize;
+        uint g1 = (i + 1) / groupSize;
+        uint g2 = (i + 2) / groupSize;
+        uint g3 = (i + 3) / groupSize;
+        
+        float scale0 = scales[groupIdx + g0];
+        float scale1 = scales[groupIdx + g1];
+        float scale2 = scales[groupIdx + g2];
+        float scale3 = scales[groupIdx + g3];
+        
+        sum += inVals.x * (w0 - 128) * scale0;
+        sum += inVals.y * (w1 - 128) * scale1;
+        sum += inVals.z * (w2 - 128) * scale2;
+        sum += inVals.w * (w3 - 128) * scale3;
+    }
+    
+    batchOutput[batchIdx * outDim + outIdx] = sum;
+}
+
+// Batch RMS norm - process N tokens simultaneously
+kernel void rms_norm_batch(
+    device float* batchInput [[buffer(0)]],     // [batchSize, N]
+    device float* weights [[buffer(1)]],        // [N]
+    device float* batchOutput [[buffer(2)]],    // [batchSize, N]
+    constant uint32_t& N [[buffer(3)]],
+    constant float& eps [[buffer(4)]],
+    constant uint32_t& batchSize [[buffer(5)]],
+    uint3 gid [[thread_position_in_grid]])
+{
+    uint batchIdx = gid.x;
+    uint elemIdx = gid.y;
+    
+    if (batchIdx >= batchSize || elemIdx >= N) return;
+    
+    device float* input = batchInput + batchIdx * N;
+    float sqSum = 0.0;
+    for (uint i = 0; i < N; i++) {
+        sqSum += input[i] * input[i];
+    }
+    
+    float rms = sqrt(sqSum / float(N) + eps);
+    batchOutput[batchIdx * N + elemIdx] = input[elemIdx] / rms * weights[elemIdx];
+}
+
+// Batch attention (simplified - for demonstration)
+// Full implementation would require complex KV cache management
+kernel void sliding_attention_batch(
+    device float* batchQuery [[buffer(0)]],     // [batchSize, nHeads, headDim]
+    device float* kvCache [[buffer(1)]],        // [maxSeqLen, 2, nKvHeads, headDim]
+    device float* batchOutput [[buffer(2)]],    // [batchSize, nHeads, headDim]
+    constant uint32_t* positions [[buffer(3)]], // [batchSize]
+    constant uint32_t& nHeads [[buffer(4)]],
+    constant uint32_t& nKvHeads [[buffer(5)]],
+    constant uint32_t& headDim [[buffer(6)]],
+    constant uint32_t& batchSize [[buffer(7)]],
+    constant uint32_t& windowSize [[buffer(8)]],
+    uint3 gid [[thread_position_in_grid]])
+{
+    uint batchIdx = gid.x;
+    uint headIdx = gid.y;
+    uint dimIdx = gid.z;
+    
+    if (batchIdx >= batchSize || headIdx >= nHeads || dimIdx >= headDim) return;
+    
+    uint pos = positions[batchIdx];
+    uint kvHeadIdx = headIdx / (nHeads / nKvHeads);
+    
+    device float* query = batchQuery + batchIdx * nHeads * headDim + headIdx * headDim;
+    
+    uint start = max(0u, pos - windowSize);
+    uint end = pos;
+    
+    float maxScore = -1e10;
+    for (uint t = start; t < end; t++) {
+        device float* key = kvCache + t * 2 * nKvHeads * headDim + kvHeadIdx * headDim;
+        
+        float score = 0.0;
+        for (uint d = 0; d < headDim; d++) {
+            score += query[d] * key[d];
+        }
+        
+        score /= sqrt(float(headDim));
+        maxScore = max(maxScore, score);
+    }
+    
+    float expSum = 0.0;
+    for (uint t = start; t < end; t++) {
+        device float* key = kvCache + t * 2 * nKvHeads * headDim + kvHeadIdx * headDim;
+        
+        float score = 0.0;
+        for (uint d = 0; d < headDim; d++) {
+            score += query[d] * key[d];
+        }
+        
+        score /= sqrt(float(headDim));
+        expSum += exp(score - maxScore);
+    }
+    
+    float output = 0.0;
+    for (uint t = start; t < end; t++) {
+        device float* key = kvCache + t * 2 * nKvHeads * headDim + kvHeadIdx * headDim;
+        device float* value = kvCache + t * 2 * nKvHeads * headDim + nKvHeads * headDim + kvHeadIdx * headDim;
+        
+        float score = 0.0;
+        for (uint d = 0; d < headDim; d++) {
+            score += query[d] * key[d];
+        }
+        
+        score /= sqrt(float(headDim));
+        float weight = exp(score - maxScore) / expSum;
+        
+        output += weight * value[dimIdx];
+    }
+    
+    batchOutput[batchIdx * nHeads * headDim + headIdx * headDim + dimIdx] = output;
+}