v2: Initial clean branch with unit tests + CI/CD pipeline
- Started from ac75faa (initial E4B-MarkBase integration)
- Kept Sources/ (all engine code) + Package.swift + .gitignore
- Removed all ad-hoc tests, documentation, scripts, Python files
- Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest)
- Added .gitea/workflows/ci.yaml (build + unit tests + lint)
- Added Scripts/check_resources.sh (memory-aware test runner)
- Added Tests/Manifest.json (resource requirements for all tests)
- Focus: 4-bit quantized models only
This commit is contained in:
MarkBase Admin
@@ -0,0 +1,177 @@
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// Batch Metal Kernels - Process multiple tokens simultaneously
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
|
||||
// Batch quantized matmul - process N tokens with shared weights
|
||||
// Expected improvement: 8-15x for batch inference
|
||||
kernel void quantized_matmul_batch(
|
||||
device float* inputs [[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* outputs [[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]])
|
||||
{
|
||||
// Each thread processes one output dimension for one batch element
|
||||
uint batchIdx = gid.x; // [0, batchSize)
|
||||
uint outIdx = gid.y; // [0, outDim)
|
||||
|
||||
if (batchIdx >= batchSize || outIdx >= outDim) return;
|
||||
|
||||
// Get input for this batch element
|
||||
device float* input = inputs + batchIdx * inDim;
|
||||
|
||||
// Compute dot product for this output dimension
|
||||
float sum = biases[outIdx];
|
||||
uint groupIdx = outIdx * (inDim / groupSize);
|
||||
|
||||
for (uint i = 0; i < inDim; i += 4) {
|
||||
// Load 4 input values
|
||||
float4 inVals = float4(input[i], input[i+1], input[i+2], input[i+3]);
|
||||
|
||||
// Load 4 packed weights (uint8 packed as uint32)
|
||||
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;
|
||||
|
||||
// Get scale for this group
|
||||
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];
|
||||
|
||||
// Dequantize and multiply
|
||||
sum += inVals.x * (w0 - 128) * scale0;
|
||||
sum += inVals.y * (w1 - 128) * scale1;
|
||||
sum += inVals.z * (w2 - 128) * scale2;
|
||||
sum += inVals.w * (w3 - 128) * scale3;
|
||||
}
|
||||
|
||||
outputs[batchIdx * outDim + outIdx] = sum;
|
||||
}
|
||||
|
||||
// Batch RMS norm - process N tokens simultaneously
|
||||
kernel void rms_norm_batch(
|
||||
device float* inputs [[buffer(0)]], // [batchSize, N]
|
||||
device float* weights [[buffer(1)]], // [N]
|
||||
device float* outputs [[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;
|
||||
|
||||
// Compute sum of squares for this batch element
|
||||
threadgroup float sharedSqSum[256];
|
||||
uint threadIdx = elemIdx % 256;
|
||||
|
||||
device float* input = inputs + batchIdx * N;
|
||||
|
||||
float sqSum = 0.0;
|
||||
for (uint i = 0; i < N; i++) {
|
||||
sqSum += input[i] * input[i];
|
||||
}
|
||||
|
||||
// RMS
|
||||
float rms = sqrt(sqSum / float(N) + eps);
|
||||
|
||||
// Normalize
|
||||
outputs[batchIdx * N + elemIdx] = input[elemIdx] / rms * weights[elemIdx];
|
||||
}
|
||||
|
||||
// Batch attention - process N tokens with shared KV cache
|
||||
// This is the most complex batch operation
|
||||
kernel void sliding_attention_batch(
|
||||
device float* queries [[buffer(0)], // [batchSize, nHeads, headDim]
|
||||
device float* kvCache [[buffer(1)], // [maxSeqLen, 2, nKvHeads, headDim]
|
||||
device float* outputs [[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 = queries + batchIdx * nHeads * headDim + headIdx * headDim;
|
||||
|
||||
// Sliding window attention
|
||||
uint start = max(0u, pos - windowSize);
|
||||
uint end = min(pos, maxSeqLen);
|
||||
|
||||
float sum = 0.0;
|
||||
float maxScore = -1e10;
|
||||
|
||||
// Compute attention scores
|
||||
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);
|
||||
}
|
||||
|
||||
// Softmax
|
||||
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);
|
||||
}
|
||||
|
||||
// Compute weighted sum of values
|
||||
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];
|
||||
}
|
||||
|
||||
outputs[batchIdx * nHeads * headDim + headIdx * headDim + dimIdx] = output;
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// Batch Layer Processing Kernels
|
||||
// Process entire layer for multiple tokens simultaneously
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
|
||||
// Batch RMS Norm for layer input
|
||||
// Process [batchSize, hiddenSize] with shared weights
|
||||
kernel void batch_layer_rms_norm(
|
||||
device float* batchInput [[buffer(0)]], // [batchSize, hiddenSize]
|
||||
device float* weights [[buffer(1)]], // [hiddenSize]
|
||||
device float* batchOutput [[buffer(2)]], // [batchSize, hiddenSize]
|
||||
constant uint32_t& hiddenSize [[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 >= hiddenSize) return;
|
||||
|
||||
device float* input = batchInput + batchIdx * hiddenSize;
|
||||
device float* output = batchOutput + batchIdx * hiddenSize;
|
||||
|
||||
// Compute sum of squares for this batch element
|
||||
float ss = 0.0;
|
||||
for (uint i = 0; i < hiddenSize; i++) {
|
||||
ss += input[i] * input[i];
|
||||
}
|
||||
|
||||
float rms = sqrt(ss / float(hiddenSize) + eps);
|
||||
output[elemIdx] = input[elemIdx] / rms * weights[elemIdx];
|
||||
}
|
||||
|
||||
// Batch Quantized Matmul for layer projections
|
||||
// Process [batchSize, outDim] with shared quantized weights
|
||||
kernel void batch_layer_quantized_matmul(
|
||||
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;
|
||||
device float* output = batchOutput + batchIdx * outDim;
|
||||
|
||||
float sum = biases[outIdx];
|
||||
uint groupIdx = outIdx * (inDim / groupSize);
|
||||
|
||||
// Process in groups for quantization
|
||||
for (uint i = 0; i < inDim; i++) {
|
||||
// Load weight (8-bit quantized)
|
||||
uint8_t w = weights[outIdx * inDim + i];
|
||||
|
||||
// Get scale for this group
|
||||
uint g = i / groupSize;
|
||||
float scale = scales[groupIdx + g];
|
||||
|
||||
// Dequantize and accumulate
|
||||
sum += input[i] * (w - 128) * scale;
|
||||
}
|
||||
|
||||
output[outIdx] = sum;
|
||||
}
|
||||
|
||||
// Batch Elementwise Add for residual connections
|
||||
// Process [batchSize, size]
|
||||
kernel void batch_eltwise_add(
|
||||
device float* batchA [[buffer(0)]], // [batchSize, size]
|
||||
device float* batchB [[buffer(1)]], // [batchSize, size]
|
||||
device float* batchOutput [[buffer(2)]], // [batchSize, size]
|
||||
constant uint32_t& size [[buffer(3)]],
|
||||
constant uint32_t& batchSize [[buffer(4)]],
|
||||
uint2 gid [[thread_position_in_grid]])
|
||||
{
|
||||
uint batchIdx = gid.x;
|
||||
uint elemIdx = gid.y;
|
||||
|
||||
if (batchIdx >= batchSize || elemIdx >= size) return;
|
||||
|
||||
uint offset = batchIdx * size + elemIdx;
|
||||
batchOutput[offset] = batchA[offset] + batchB[offset];
|
||||
}
|
||||
|
||||
// Batch Gated FFN (fused gate + up projection)
|
||||
// Process [batchSize, intermediateSize]
|
||||
kernel void batch_fused_gate_up(
|
||||
device float* batchInput [[buffer(0)]], // [batchSize, hiddenSize]
|
||||
device uint8_t* gateWeights [[buffer(1)]], // [intermediateSize, hiddenSize]
|
||||
device float* gateScales [[buffer(2)]],
|
||||
device float* gateBiases [[buffer(3)]],
|
||||
device uint8_t* upWeights [[buffer(4)]], // [intermediateSize, hiddenSize]
|
||||
device float* upScales [[buffer(5)]],
|
||||
device float* upBiases [[buffer(6)]],
|
||||
device float* batchOutput [[buffer(7)]], // [batchSize, intermediateSize]
|
||||
constant uint32_t& hiddenSize [[buffer(8)]],
|
||||
constant uint32_t& intermediateSize [[buffer(9)]],
|
||||
constant uint32_t& groupSize [[buffer(10)]],
|
||||
constant uint32_t& batchSize [[buffer(11)]],
|
||||
uint3 gid [[thread_position_in_grid]])
|
||||
{
|
||||
uint batchIdx = gid.x;
|
||||
uint interIdx = gid.y;
|
||||
|
||||
if (batchIdx >= batchSize || interIdx >= intermediateSize) return;
|
||||
|
||||
device float* input = batchInput + batchIdx * hiddenSize;
|
||||
device float* output = batchOutput + batchIdx * intermediateSize;
|
||||
|
||||
// Compute gate
|
||||
float gate = gateBiases[interIdx];
|
||||
uint gateGroupIdx = interIdx * (hiddenSize / groupSize);
|
||||
for (uint i = 0; i < hiddenSize; i++) {
|
||||
uint8_t w = gateWeights[interIdx * hiddenSize + i];
|
||||
uint g = i / groupSize;
|
||||
float scale = gateScales[gateGroupIdx + g];
|
||||
gate += input[i] * (w - 128) * scale;
|
||||
}
|
||||
|
||||
// Compute up
|
||||
float up = upBiases[interIdx];
|
||||
uint upGroupIdx = interIdx * (hiddenSize / groupSize);
|
||||
for (uint i = 0; i < hiddenSize; i++) {
|
||||
uint8_t w = upWeights[interIdx * hiddenSize + i];
|
||||
uint g = i / groupSize;
|
||||
float scale = upScales[upGroupIdx + g];
|
||||
up += input[i] * (w - 128) * scale;
|
||||
}
|
||||
|
||||
// Fused activation: gate * sigmoid(gate) * up
|
||||
float sigmoidGate = 1.0 / (1.0 + exp(-gate));
|
||||
output[interIdx] = gate * sigmoidGate * up;
|
||||
}
|
||||
|
||||
// Batch Down Projection (FFN output)
|
||||
// Process [batchSize, hiddenSize]
|
||||
kernel void batch_down_projection(
|
||||
device float* batchInter [[buffer(0)]], // [batchSize, intermediateSize]
|
||||
device uint8_t* downWeights [[buffer(1)]], // [hiddenSize, intermediateSize]
|
||||
device float* downScales [[buffer(2)]],
|
||||
device float* downBiases [[buffer(3)]],
|
||||
device float* batchOutput [[buffer(4)]], // [batchSize, hiddenSize]
|
||||
constant uint32_t& hiddenSize [[buffer(5)]],
|
||||
constant uint32_t& intermediateSize [[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 >= hiddenSize) return;
|
||||
|
||||
device float* inter = batchInter + batchIdx * intermediateSize;
|
||||
device float* output = batchOutput + batchIdx * hiddenSize;
|
||||
|
||||
float sum = downBiases[outIdx];
|
||||
uint groupIdx = outIdx * (intermediateSize / groupSize);
|
||||
|
||||
for (uint i = 0; i < intermediateSize; i++) {
|
||||
uint8_t w = downWeights[outIdx * intermediateSize + i];
|
||||
uint g = i / groupSize;
|
||||
float scale = downScales[groupIdx + g];
|
||||
sum += inter[i] * (w - 128) * scale;
|
||||
}
|
||||
|
||||
output[outIdx] = sum;
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// ════════════════════════════════════════════════════════
|
||||
// Float16 Metal Kernels
|
||||
// ════════════════════════════════════════════════════════
|
||||
|
||||
// ── Float16 Quantized Matmul ──────────────────────────
|
||||
// Uses half precision for input/weights
|
||||
kernel void quantized_matmul_f16(
|
||||
device const half *x [[buffer(0)]], // Input [inDim]
|
||||
device const uint *w [[buffer(1)]], // Packed weights [outDim, inDim/8]
|
||||
device const half *s [[buffer(2)]], // Scales [outDim, inDim/64]
|
||||
device const half *b [[buffer(3)]], // Biases [outDim, inDim/64]
|
||||
device float *out [[buffer(4)]], // Output [outDim] - Float32 for accuracy
|
||||
constant uint &inDim [[buffer(5)]],
|
||||
constant uint &outDim [[buffer(6)]],
|
||||
constant uint &groupSize [[buffer(7)]],
|
||||
threadgroup half *shared_x [[threadgroup(0)]], // Input cache in half
|
||||
uint gid [[thread_position_in_grid]],
|
||||
uint tid [[thread_position_in_threadgroup]],
|
||||
uint tgSize [[threads_per_threadgroup]]
|
||||
) {
|
||||
uint outRow = gid;
|
||||
if (outRow >= outDim) return;
|
||||
|
||||
// Cooperative loading of input vector
|
||||
for (uint i = tid; i < inDim; i += tgSize) {
|
||||
shared_x[i] = x[i];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Compute dot product
|
||||
uint numGroups = inDim / groupSize;
|
||||
float sum = 0.0;
|
||||
|
||||
for (uint g = 0; g < numGroups; g++) {
|
||||
half scale = s[outRow * numGroups + g];
|
||||
half bias = b[outRow * numGroups + g];
|
||||
|
||||
uint packedBase = outRow * (inDim / 8) + g * (groupSize / 8);
|
||||
|
||||
// Process 8 packed uint32 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;
|
||||
|
||||
// Load 16 half values
|
||||
half4 xVec0 = half4(shared_x[xBase+0], shared_x[xBase+1], shared_x[xBase+2], shared_x[xBase+3]);
|
||||
half4 xVec1 = half4(shared_x[xBase+4], shared_x[xBase+5], shared_x[xBase+6], shared_x[xBase+7]);
|
||||
half4 xVec2 = half4(shared_x[xBase+8], shared_x[xBase+9], shared_x[xBase+10], shared_x[xBase+11]);
|
||||
half4 xVec3 = half4(shared_x[xBase+12], shared_x[xBase+13], shared_x[xBase+14], shared_x[xBase+15]);
|
||||
|
||||
// Dequantize
|
||||
half4 qVec0 = half4(
|
||||
half((packed0 >> 0) & 0xF) * scale + bias,
|
||||
half((packed0 >> 4) & 0xF) * scale + bias,
|
||||
half((packed0 >> 8) & 0xF) * scale + bias,
|
||||
half((packed0 >> 12) & 0xF) * scale + bias
|
||||
);
|
||||
half4 qVec1 = half4(
|
||||
half((packed0 >> 16) & 0xF) * scale + bias,
|
||||
half((packed0 >> 20) & 0xF) * scale + bias,
|
||||
half((packed0 >> 24) & 0xF) * scale + bias,
|
||||
half((packed0 >> 28) & 0xF) * scale + bias
|
||||
);
|
||||
half4 qVec2 = half4(
|
||||
half((packed1 >> 0) & 0xF) * scale + bias,
|
||||
half((packed1 >> 4) & 0xF) * scale + bias,
|
||||
half((packed1 >> 8) & 0xF) * scale + bias,
|
||||
half((packed1 >> 12) & 0xF) * scale + bias
|
||||
);
|
||||
half4 qVec3 = half4(
|
||||
half((packed1 >> 16) & 0xF) * scale + bias,
|
||||
half((packed1 >> 20) & 0xF) * scale + bias,
|
||||
half((packed1 >> 24) & 0xF) * scale + bias,
|
||||
half((packed1 >> 28) & 0xF) * scale + bias
|
||||
);
|
||||
|
||||
// Accumulate in Float32 for accuracy
|
||||
sum += float(dot(qVec0, xVec0)) + float(dot(qVec1, xVec1)) +
|
||||
float(dot(qVec2, xVec2)) + float(dot(qVec3, xVec3));
|
||||
}
|
||||
}
|
||||
|
||||
out[outRow] = sum;
|
||||
}
|
||||
|
||||
// ── Float16 RMS Norm ──────────────────────────────────
|
||||
kernel void rms_norm_f16(
|
||||
device const half *x [[buffer(0)]], // Input [N]
|
||||
device const half *w [[buffer(1)]], // Weight [N]
|
||||
device half *y [[buffer(2)]], // Output [N]
|
||||
constant uint &N [[buffer(3)]],
|
||||
constant half &eps [[buffer(4)]],
|
||||
threadgroup half *partial_sums [[threadgroup(0)]],
|
||||
uint tid [[thread_position_in_threadgroup]],
|
||||
uint tgSize [[threads_per_threadgroup]]
|
||||
) {
|
||||
// Phase 1: Each thread computes partial sum of squares
|
||||
half localSum = 0.0;
|
||||
for (uint i = tid; i < N; i += tgSize) {
|
||||
localSum += x[i] * x[i];
|
||||
}
|
||||
partial_sums[tid] = localSum;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Phase 2: Parallel reduction
|
||||
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);
|
||||
}
|
||||
|
||||
// Phase 3: Compute RMS and normalize
|
||||
half ss = partial_sums[0];
|
||||
half rms = rsqrt(ss / half(N) + eps);
|
||||
|
||||
// Each thread outputs its portion
|
||||
for (uint i = tid; i < N; i += tgSize) {
|
||||
y[i] = x[i] * rms * (w ? w[i] : half(1.0));
|
||||
}
|
||||
}
|
||||
|
||||
// ── Float16 Elementwise Operations ────────────────────
|
||||
|
||||
kernel void eltwise_mul_f16(
|
||||
device const half *a,
|
||||
device const half *b,
|
||||
device half *out,
|
||||
constant uint &count,
|
||||
uint id [[thread_position_in_grid]]
|
||||
) {
|
||||
uint idx = id * 4;
|
||||
if (idx >= count) return;
|
||||
|
||||
half4 aVec = half4(a[idx], a[idx+1], a[idx+2], a[idx+3]);
|
||||
half4 bVec = half4(b[idx], b[idx+1], b[idx+2], b[idx+3]);
|
||||
half4 outVec = aVec * bVec;
|
||||
|
||||
if (idx < count) out[idx] = outVec.x;
|
||||
if (idx+1 < count) out[idx+1] = outVec.y;
|
||||
if (idx+2 < count) out[idx+2] = outVec.z;
|
||||
if (idx+3 < count) out[idx+3] = outVec.w;
|
||||
}
|
||||
|
||||
kernel void eltwise_add_f16(
|
||||
device const half *a,
|
||||
device const half *b,
|
||||
device half *out,
|
||||
constant uint &count,
|
||||
uint id [[thread_position_in_grid]]
|
||||
) {
|
||||
uint idx = id * 4;
|
||||
if (idx >= count) return;
|
||||
|
||||
half4 aVec = half4(a[idx], a[idx+1], a[idx+2], a[idx+3]);
|
||||
half4 bVec = half4(b[idx], b[idx+1], b[idx+2], b[idx+3]);
|
||||
half4 outVec = aVec + bVec;
|
||||
|
||||
if (idx < count) out[idx] = outVec.x;
|
||||
if (idx+1 < count) out[idx+1] = outVec.y;
|
||||
if (idx+2 < count) out[idx+2] = outVec.z;
|
||||
if (idx+3 < count) out[idx+3] = outVec.w;
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// Kernel Fusion: Combine multiple operations into single kernels
|
||||
// Goal: Reduce kernel dispatches for common patterns
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
// ── Fused Embedding Dequantize + Scale ──
|
||||
// Combines: dequantize_row + eltwise_scale
|
||||
// Eliminates one kernel dispatch
|
||||
kernel void fused_dequantize_scale(
|
||||
device const uint32_t* weight [[buffer(0)]],
|
||||
device const float* scales [[buffer(1)]],
|
||||
device const float* biases [[buffer(2)]],
|
||||
device float* output [[buffer(3)]],
|
||||
constant uint& nCols [[buffer(4)]],
|
||||
constant int& row [[buffer(5)]],
|
||||
constant uint& groupSize [[buffer(6)]],
|
||||
constant float& scale [[buffer(7)]], // Extra scale to apply
|
||||
uint id [[thread_position_in_grid]]
|
||||
) {
|
||||
if (id >= nCols) return;
|
||||
|
||||
uint numGroups = nCols / groupSize;
|
||||
uint groupIdx = id / groupSize;
|
||||
uint inGroupIdx = id % groupSize;
|
||||
|
||||
uint weightRowOffset = row * (nCols / 8);
|
||||
uint packedIdx = weightRowOffset + id / 8;
|
||||
uint subIdx = id % 8;
|
||||
|
||||
uint32_t packed = weight[packedIdx];
|
||||
uint32_t qval = (packed >> (subIdx * 4)) & 0xF;
|
||||
|
||||
float scale_val = scales[groupIdx];
|
||||
float bias_val = biases[groupIdx];
|
||||
|
||||
float val = float(qval) * scale_val + bias_val;
|
||||
|
||||
// Apply extra scale (embedding scale or per-layer scale)
|
||||
val *= scale;
|
||||
|
||||
output[id] = val;
|
||||
}
|
||||
|
||||
// ── Fused RMS Norm + Residual Add ──
|
||||
// Combines: rmsNorm + eltwiseAdd
|
||||
// Eliminates one kernel dispatch
|
||||
kernel void fused_rms_norm_residual(
|
||||
device const float* input [[buffer(0)]],
|
||||
device const float* residual [[buffer(1)]],
|
||||
device const float* weight [[buffer(2)]],
|
||||
device float* output [[buffer(3)]],
|
||||
constant uint& N [[buffer(4)]],
|
||||
constant float& eps [[buffer(5)]],
|
||||
uint tid [[thread_position_in_grid]],
|
||||
uint threadgroupId [[threadgroup_position_in_grid]],
|
||||
uint threadgroupSize [[threads_per_threadgroup]]
|
||||
) {
|
||||
// Parallel RMS computation
|
||||
threadgroup float sharedSum[256];
|
||||
|
||||
uint laneId = tid % threadgroupSize;
|
||||
uint groupId = tid / threadgroupSize;
|
||||
|
||||
float sumSq = 0.0;
|
||||
uint start = groupId * (N / 256);
|
||||
uint end = min((groupId + 1) * (N / 256), N);
|
||||
|
||||
for (uint i = start; i < end; i++) {
|
||||
float val = input[i];
|
||||
sumSq += val * val;
|
||||
}
|
||||
|
||||
sharedSum[laneId] = sumSq;
|
||||
|
||||
// Simplified RMS (proper implementation would use SIMD reduction)
|
||||
float rms = sqrt(sharedSum[0] / N + eps);
|
||||
|
||||
if (tid < N) {
|
||||
float normed = input[tid] / rms * weight[tid];
|
||||
output[tid] = residual[tid] + normed; // Residual add
|
||||
}
|
||||
}
|
||||
|
||||
// ── Fused Matmul + GELU + Residual ──
|
||||
// Combines: quantized_matmul + gelu + eltwiseAdd
|
||||
kernel void fused_matmul_gelu_residual(
|
||||
device const float* input [[buffer(0)]],
|
||||
device const uint32_t* weight [[buffer(1)]],
|
||||
device const float* scales [[buffer(2)]],
|
||||
device const float* biases [[buffer(3)]],
|
||||
device const float* residual [[buffer(4)]],
|
||||
device float* output [[buffer(5)]],
|
||||
constant uint& inDim [[buffer(6)]],
|
||||
constant uint& outDim [[buffer(7)]],
|
||||
constant uint& groupSize [[buffer(8)]],
|
||||
uint id [[thread_position_in_grid]]
|
||||
) {
|
||||
if (id >= outDim) return;
|
||||
|
||||
uint numGroups = inDim / groupSize;
|
||||
float sum = 0.0;
|
||||
|
||||
for (uint g = 0; g < numGroups; g++) {
|
||||
float scale = scales[id * numGroups + g];
|
||||
float bias = biases[id * numGroups + g];
|
||||
|
||||
for (uint j = 0; j < groupSize / 8; j++) {
|
||||
uint weightIdx = id * (inDim / 8) + g * (groupSize / 8) + j;
|
||||
uint32_t packed = weight[weightIdx];
|
||||
|
||||
for (uint k = 0; k < 8; k++) {
|
||||
uint inputIdx = g * groupSize + j * 8 + k;
|
||||
uint32_t qval = (packed >> (k * 4)) & 0xF;
|
||||
float wval = float(qval) * scale + bias;
|
||||
sum += input[inputIdx] * wval;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply GELU approximation
|
||||
float gelu = sum * 0.5 * (1.0 + tanh(sum * 0.7978845608 * (1.0 + 0.044715 * sum * sum)));
|
||||
|
||||
// Residual add
|
||||
output[id] = residual[id] + gelu;
|
||||
}
|
||||
|
||||
// ── Batch RMS Norm for Multiple Layers ──
|
||||
// Process 42 layers' norm operations in one dispatch
|
||||
kernel void batch_rms_norm_layers(
|
||||
device const float* inputs [[buffer(0)]], // [numLayers * hiddenSize] flattened
|
||||
device const float* weights [[buffer(1)]], // [numLayers * hiddenSize] flattened
|
||||
device float* outputs [[buffer(2)]], // [numLayers * hiddenSize] flattened
|
||||
constant uint& hiddenSize [[buffer(3)]],
|
||||
constant uint& numLayers [[buffer(4)]],
|
||||
constant float& eps [[buffer(5)]],
|
||||
uint2 id [[thread_position_in_grid]]
|
||||
) {
|
||||
uint layerIdx = id.y;
|
||||
uint dimIdx = id.x;
|
||||
|
||||
if (layerIdx >= numLayers || dimIdx >= hiddenSize) return;
|
||||
|
||||
uint offset = layerIdx * hiddenSize;
|
||||
|
||||
// Simplified RMS computation (proper would need threadgroup reduction)
|
||||
float sumSq = 0.0;
|
||||
for (uint i = 0; i < hiddenSize; i++) {
|
||||
float val = inputs[offset + i];
|
||||
sumSq += val * val;
|
||||
}
|
||||
float rms = sqrt(sumSq / hiddenSize + eps);
|
||||
|
||||
outputs[offset + dimIdx] = inputs[offset + dimIdx] / rms * weights[offset + dimIdx];
|
||||
}
|
||||
|
||||
// ── Fused Quantized Matmul + Bias Add ──
|
||||
kernel void fused_quantized_matmul_bias(
|
||||
device const float* input [[buffer(0)]],
|
||||
device const uint32_t* weight [[buffer(1)]],
|
||||
device const float* scales [[buffer(2)]],
|
||||
device const float* biases_quant [[buffer(3)]],
|
||||
device const float* bias_unquant [[buffer(4)]], // Optional unquantized bias
|
||||
device float* output [[buffer(5)]],
|
||||
constant uint& inDim [[buffer(6)]],
|
||||
constant uint& outDim [[buffer(7)]],
|
||||
constant uint& groupSize [[buffer(8)]],
|
||||
constant bool& hasBias [[buffer(9)]],
|
||||
uint id [[thread_position_in_grid]]
|
||||
) {
|
||||
if (id >= outDim) return;
|
||||
|
||||
uint numGroups = inDim / groupSize;
|
||||
float sum = 0.0;
|
||||
|
||||
for (uint g = 0; g < numGroups; g++) {
|
||||
float scale = scales[id * numGroups + g];
|
||||
float bias = biases_quant[id * numGroups + g];
|
||||
|
||||
for (uint j = 0; j < groupSize / 8; j++) {
|
||||
uint weightIdx = id * (inDim / 8) + g * (groupSize / 8) + j;
|
||||
uint32_t packed = weight[weightIdx];
|
||||
|
||||
for (uint k = 0; k < 8; k++) {
|
||||
uint inputIdx = g * groupSize + j * 8 + k;
|
||||
uint32_t qval = (packed >> (k * 4)) & 0xF;
|
||||
float wval = float(qval) * scale + bias;
|
||||
sum += input[inputIdx] * wval;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Add unquantized bias if present
|
||||
if (hasBias) {
|
||||
sum += bias_unquant[id];
|
||||
}
|
||||
|
||||
output[id] = sum;
|
||||
}
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,236 @@
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// ═══════════════════════════════════════════════
|
||||
// Numerically Stable RMSNorm Kernel
|
||||
// ═══════════════════════════════════════════════
|
||||
|
||||
// Optimized RMSNorm with numerical stability
|
||||
// Uses threadgroup parallel reduction to avoid overflow
|
||||
kernel void rms_norm_stable(
|
||||
device const float *x [[buffer(0)]], // [N]
|
||||
device const float *w [[buffer(1)]], // [N] weight (can be null)
|
||||
device float *y [[buffer(2)]], // [N]
|
||||
constant uint &N [[buffer(3)]],
|
||||
constant float &eps [[buffer(4)]],
|
||||
uint tid [[thread_position_in_threadgroup]],
|
||||
uint gid [[thread_position_in_grid]],
|
||||
uint tgsize [[threads_per_threadgroup]]
|
||||
) {
|
||||
// Early exit for out-of-range threads
|
||||
if (gid >= N) return;
|
||||
|
||||
// Threadgroup shared memory for partial sums
|
||||
threadgroup float partialSums[256];
|
||||
|
||||
// Step 1: Each thread computes partial sum with numerical stability
|
||||
float localSum = 0.0;
|
||||
uint chunkSize = N / tgsize;
|
||||
uint start = tid * chunkSize;
|
||||
uint end = min(start + chunkSize, N);
|
||||
|
||||
// Optimized SIMD batch clamp for performance
|
||||
// Process 4 values at once using SIMD
|
||||
for (uint i = start; i < end; i += 4) {
|
||||
// Load 4 values
|
||||
float4 xiVec = float4(
|
||||
i < end ? x[i] : 0.0f,
|
||||
i+1 < end ? x[i+1] : 0.0f,
|
||||
i+2 < end ? x[i+2] : 0.0f,
|
||||
i+3 < end ? x[i+3] : 0.0f
|
||||
);
|
||||
|
||||
// Single clamp operation (SIMD)
|
||||
xiVec = clamp(xiVec, -20.0f, 20.0f);
|
||||
|
||||
// Compute sum of squares
|
||||
float4 sqVec = xiVec * xiVec;
|
||||
localSum += sqVec[0] + sqVec[1] + sqVec[2] + sqVec[3];
|
||||
}
|
||||
|
||||
// Store partial sum
|
||||
if (tid < 256) {
|
||||
partialSums[tid] = localSum;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Step 2: Parallel reduction in threadgroup
|
||||
// Reduce to single sum
|
||||
for (uint stride = tgsize / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
partialSums[tid] += partialSums[tid + stride];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
// Step 3: Compute RMS from total sum
|
||||
float totalSum = partialSums[0];
|
||||
float meanSq = totalSum / float(N);
|
||||
|
||||
// Numerical stability: ensure meanSq is positive and reasonable
|
||||
meanSq = max(meanSq, eps);
|
||||
meanSq = min(meanSq, 10000.0f); // Prevent extreme RMS values
|
||||
|
||||
float rms = rsqrt(meanSq + eps);
|
||||
|
||||
// Numerical stability: clamp RMS to reasonable range
|
||||
rms = clamp(rms, 0.01f, 100.0f);
|
||||
|
||||
// Step 4: Apply normalization
|
||||
float xi = x[gid];
|
||||
float yi = xi * rms;
|
||||
|
||||
// Apply weight if provided
|
||||
if (w) {
|
||||
yi *= w[gid];
|
||||
}
|
||||
|
||||
// Final numerical stability: aggressive clamp output
|
||||
// Progressive output clamp
|
||||
float yiFinal = yi;
|
||||
if (yiFinal > 50.0f) yiFinal = 50.0f;
|
||||
else if (yiFinal < -50.0f) yiFinal = -50.0f;
|
||||
else if (yiFinal > 20.0f) yiFinal = 20.0f + (yiFinal - 20.0f) * 0.2f;
|
||||
else if (yiFinal < -20.0f) yiFinal = -20.0f + (yiFinal + 20.0f) * 0.2f;
|
||||
y[gid] = yiFinal;
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════
|
||||
// Numerically Stable Softmax Kernel
|
||||
// ═══════════════════════════════════════════════
|
||||
|
||||
// Stable softmax with numerical overflow protection
|
||||
kernel void softmax_stable(
|
||||
device const float *logits [[buffer(0)]], // [N]
|
||||
device float *probs [[buffer(1)]], // [N]
|
||||
constant uint &N [[buffer(2)]],
|
||||
uint tid [[thread_position_in_threadgroup]],
|
||||
uint gid [[thread_position_in_grid]],
|
||||
uint tgsize [[threads_per_threadgroup]]
|
||||
) {
|
||||
if (gid >= N) return;
|
||||
|
||||
threadgroup float sharedMax[256];
|
||||
threadgroup float sharedSumExp[256];
|
||||
|
||||
// Step 1: Find max using threadgroup parallel reduction
|
||||
float localMax = -INFINITY;
|
||||
uint chunkSize = N / tgsize;
|
||||
uint start = tid * chunkSize;
|
||||
uint end = min(start + chunkSize, N);
|
||||
|
||||
for (uint i = start; i < end; i++) {
|
||||
// More aggressive logits clamp
|
||||
float li = logits[i];
|
||||
if (li > 30.0f) li = 30.0f;
|
||||
else if (li < -30.0f) li = -30.0f;
|
||||
else if (li > 10.0f) li = 10.0f + (li - 10.0f) * 0.3f;
|
||||
else if (li < -10.0f) li = -10.0f + (li + 10.0f) * 0.3f;
|
||||
localMax = max(localMax, li);
|
||||
}
|
||||
|
||||
if (tid < 256) {
|
||||
sharedMax[tid] = localMax;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Parallel reduction to find global max
|
||||
for (uint stride = tgsize / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
sharedMax[tid] = max(sharedMax[tid], sharedMax[tid + stride]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
float globalMax = sharedMax[0];
|
||||
|
||||
// Optimized SIMD batch softmax
|
||||
float localSumExp = 0.0;
|
||||
for (uint i = start; i < end; i += 4) {
|
||||
float4 liVec = float4(
|
||||
i < end ? logits[i] : 0.0f,
|
||||
i+1 < end ? logits[i+1] : 0.0f,
|
||||
i+2 < end ? logits[i+2] : 0.0f,
|
||||
i+3 < end ? logits[i+3] : 0.0f
|
||||
);
|
||||
|
||||
// SIMD clamp
|
||||
liVec = clamp(liVec, -30.0f, 30.0f);
|
||||
|
||||
// SIMD compute diff
|
||||
float4 diffVec = liVec - globalMax;
|
||||
diffVec = clamp(diffVec, -10.0f, 10.0f);
|
||||
|
||||
// SIMD exp
|
||||
float4 expVec = exp(diffVec);
|
||||
localSumExp += expVec[0] + expVec[1] + expVec[2] + expVec[3];
|
||||
}
|
||||
|
||||
if (tid < 256) {
|
||||
sharedSumExp[tid] = localSumExp;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Parallel reduction to compute total sumExp
|
||||
for (uint stride = tgsize / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
sharedSumExp[tid] += sharedSumExp[tid + stride];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
float totalSumExp = sharedSumExp[0];
|
||||
totalSumExp = max(totalSumExp, 1e-6f); // Prevent division by zero
|
||||
|
||||
// Step 3: Compute output
|
||||
float li = logits[gid];
|
||||
if (li > 30.0f) li = 30.0f;
|
||||
else if (li < -30.0f) li = -30.0f;
|
||||
else if (li > 10.0f) li = 10.0f + (li - 10.0f) * 0.3f;
|
||||
else if (li < -10.0f) li = -10.0f + (li + 10.0f) * 0.3f;
|
||||
|
||||
float diff = li - globalMax;
|
||||
if (diff > 10.0f) diff = 10.0f;
|
||||
else if (diff < -10.0f) diff = -10.0f;
|
||||
probs[gid] = exp(diff) / totalSumExp;
|
||||
}
|
||||
|
||||
// Alternative: Block-wise RMSNorm for very large N
|
||||
kernel void rms_norm_blockwise(
|
||||
device const float *x [[buffer(0)]],
|
||||
device const float *w [[buffer(1)]],
|
||||
device float *y [[buffer(2)]],
|
||||
constant uint &N [[buffer(3)]],
|
||||
constant float &eps [[buffer(4)]],
|
||||
constant uint &blockSize [[buffer(5)]],
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
if (gid >= N) return;
|
||||
|
||||
// Compute block index
|
||||
uint blockIdx = gid / blockSize;
|
||||
uint blockStart = blockIdx * blockSize;
|
||||
uint blockEnd = min(blockStart + blockSize, N);
|
||||
|
||||
// Compute sum of squares for this block only
|
||||
float blockSum = 0.0;
|
||||
for (uint i = blockStart; i < blockEnd; i++) {
|
||||
float xi = clamp(x[i], -100.0f, 100.0f);
|
||||
blockSum += xi * xi;
|
||||
}
|
||||
|
||||
// Normalize by block size
|
||||
float meanSq = blockSum / float(blockEnd - blockStart);
|
||||
meanSq = max(meanSq, eps);
|
||||
|
||||
float rms = rsqrt(meanSq + eps);
|
||||
rms = clamp(rms, 0.01f, 100.0f);
|
||||
|
||||
// Apply normalization
|
||||
float xi = x[gid];
|
||||
float yi = xi * rms;
|
||||
|
||||
if (w) yi *= w[gid];
|
||||
|
||||
y[gid] = clamp(yi, -1000.0f, 1000.0f);
|
||||
}
|
||||
Reference in New Issue
Block a user