#include <metal_stdlib>
using namespace metal;

// ── 8-bit Quantized Matmul ───────────────
// Standard quantized matmul for 8-bit weights (4 values per uint32, mask 0xFF)
kernel void quantized_matmul_8bit(
    device const float *x      [[buffer(0)]],  // [inDim]
    device const uint *w       [[buffer(1)]],  // [outDim, inDim/4]
    device const float *s      [[buffer(2)]],  // [outDim, numGroups]
    device const float *b      [[buffer(3)]],  // [outDim, numGroups]
    device float *out          [[buffer(4)]],  // [outDim]
    constant uint &inDim       [[buffer(5)]],
    constant uint &outDim      [[buffer(6)]],
    constant uint &groupSize   [[buffer(7)]],
    uint id [[thread_position_in_grid]]
) {
    if (id >= outDim) return;
    
    uint numGroups = inDim / groupSize;
    uint packedPerOut = inDim / 4;  // 8-bit: 4 vals per uint32
    
    float sum = 0.0;
    for (uint g = 0; g < numGroups; g++) {
        float scale = s[id * numGroups + g];
        float bias  = b[id * numGroups + g];
        for (uint j = 0; j < groupSize; j++) {
            // 8-bit: groupSize/4 packed values per group
            uint packedIdx = g * (groupSize / 4) + j / 4;
            uint shift = (j % 4) * 8;  // 8-bit shift
            uint qval = (w[id * packedPerOut + packedIdx] >> shift) & 0xFF;  // 8-bit mask
            sum += (float(qval) * scale + bias) * x[g * groupSize + j];
        }
    }
    
    out[id] = sum;
}