mlx-native 0.13.0

#include <metal_stdlib>
using namespace metal;

constant uint QK5_0 = 32;
constant uint VALUES_PER_THREAD = 16;

typedef struct {
    half d;
    uchar qh[4];
    uchar qs[QK5_0 / 2];
} block_q5_0;
static_assert(sizeof(block_q5_0) == 22, "wrong q5_0 block size");

struct EmbeddingQ5_0Params {
    uint vocab_size;
    uint embed_dim;
    uint blocks_per_row;
    uint n_tokens;
};

kernel void embedding_gather_q5_0_f32(
    device const block_q5_0 *weights [[buffer(0)]],
    device const uint *token_ids [[buffer(1)]],
    device float *output [[buffer(2)]],
    constant EmbeddingQ5_0Params &p [[buffer(3)]],
    uint2 gid [[thread_position_in_grid]]) {
    const uint chunk = gid.x;
    const uint token_index = gid.y;
    const uint chunks_per_row = p.embed_dim / VALUES_PER_THREAD;
    if (chunk >= chunks_per_row || token_index >= p.n_tokens) {
        return;
    }
    const uint token = token_ids[token_index];
    if (token >= p.vocab_size) {
        const uint output_base = token_index * p.embed_dim + chunk * VALUES_PER_THREAD;
        for (uint lane = 0; lane < VALUES_PER_THREAD; ++lane) {
            output[output_base + lane] = 0.0f;
        }
        return;
    }

    const uint column_base = chunk * VALUES_PER_THREAD;
    const uint block_index = column_base / QK5_0;
    const bool high_nibble = (column_base % QK5_0) >= VALUES_PER_THREAD;
    device const block_q5_0 &block = weights[token * p.blocks_per_row + block_index];
    const uint qh = *((device const uint *)block.qh);
    const float scale = float(block.d);
    const uint output_base = token_index * p.embed_dim + column_base;
    for (uint lane = 0; lane < VALUES_PER_THREAD; ++lane) {
        const uchar packed_quant = block.qs[lane];
        const uint value_index = lane + (high_nibble ? VALUES_PER_THREAD : 0);
        const uint low = high_nibble ? packed_quant >> 4 : packed_quant & 0x0f;
        const uint quant = low | (((qh >> value_index) & 1) << 4);
        output[output_base + lane] = scale * (float(quant) - 16.0f);
    }
}