ferrum-kernels 0.8.8

#include <metal_stdlib>
using namespace metal;

constant uint THREADS_PER_GROUP = 256;
constant uint QK_K = 256;
constant uint QK8_0 = 32;

struct EmbeddingParams {
    uint token_count;
    uint hidden_size;
    uint vocabulary_size;
};

struct RmsNormParams {
    uint rows;
    uint hidden_size;
    float epsilon;
};

struct ResidualAddParams {
    uint elements;
};

struct LastTokenMaskedArgmaxParams {
    uint vocabulary_size;
    uint repetition_capacity;
};

struct block_q4_K {
    half d;
    half dmin;
    uchar scales[12];
    uchar qs[QK_K / 2];
};

struct block_q6_K {
    uchar ql[QK_K / 2];
    uchar qh[QK_K / 4];
    char scales[QK_K / 16];
    half d;
};

struct block_q8_0 {
    half d;
    char qs[QK8_0];
};

static inline float q4_k_value(device const block_q4_K & block, uint index) {
    const uint subblock = index / 32;
    uchar scale;
    uchar minimum;
    if (subblock < 4) {
        scale = block.scales[subblock] & 63;
        minimum = block.scales[subblock + 4] & 63;
    } else {
        scale = (block.scales[subblock + 4] & 0x0f)
            | ((block.scales[subblock - 4] >> 6) << 4);
        minimum = (block.scales[subblock + 4] >> 4)
            | ((block.scales[subblock] >> 6) << 4);
    }
    const uint packed_index = (subblock / 2) * 32 + index % 32;
    const uchar packed = block.qs[packed_index];
    const uint quantized = (subblock & 1) == 0 ? packed & 0x0f : packed >> 4;
    return float(block.d) * float(scale) * float(quantized)
        - float(block.dmin) * float(minimum);
}

static inline float q6_k_value(device const block_q6_K & block, uint index) {
    const uint half_block = index / 128;
    const uint position = index % 128;
    const uint lane = position % 32;
    const uint ql_base = half_block * 64;
    const uint qh_base = half_block * 32;
    const uint scale_base = half_block * 8;
    const uchar high = block.qh[qh_base + lane];

    uint low;
    uint upper;
    uint scale;
    if (position < 32) {
        low = block.ql[ql_base + lane] & 0x0f;
        upper = (high >> 0) & 0x03;
        scale = scale_base + lane / 16;
    } else if (position < 64) {
        low = block.ql[ql_base + lane + 32] & 0x0f;
        upper = (high >> 2) & 0x03;
        scale = scale_base + lane / 16 + 2;
    } else if (position < 96) {
        low = block.ql[ql_base + lane] >> 4;
        upper = (high >> 4) & 0x03;
        scale = scale_base + lane / 16 + 4;
    } else {
        low = block.ql[ql_base + lane + 32] >> 4;
        upper = (high >> 6) & 0x03;
        scale = scale_base + lane / 16 + 6;
    }
    const int quantized = int(low | (upper << 4)) - 32;
    return float(block.d) * float(block.scales[scale]) * float(quantized);
}

kernel void vnext_embedding_dense_f16(
    device const half * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    output[output_index] = token_id < params.vocabulary_size
        ? table[ulong(token_id) * params.hidden_size + column]
        : half(0.0h);
}

kernel void vnext_embedding_q4_k_f16(
    device const block_q4_K * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    if (token_id >= params.vocabulary_size) {
        output[output_index] = half(0.0h);
        return;
    }
    const uint blocks_per_row = params.hidden_size / QK_K;
    const ulong block_index = ulong(token_id) * blocks_per_row + column / QK_K;
    output[output_index] = half(q4_k_value(table[block_index], column % QK_K));
}

kernel void vnext_embedding_q6_k_f16(
    device const block_q6_K * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    if (token_id >= params.vocabulary_size) {
        output[output_index] = half(0.0h);
        return;
    }
    const uint blocks_per_row = params.hidden_size / QK_K;
    const ulong block_index = ulong(token_id) * blocks_per_row + column / QK_K;
    output[output_index] = half(q6_k_value(table[block_index], column % QK_K));
}

kernel void vnext_embedding_q8_0_f16(
    device const block_q8_0 * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    if (token_id >= params.vocabulary_size) {
        output[output_index] = half(0.0h);
        return;
    }
    const uint blocks_per_row = params.hidden_size / QK8_0;
    const ulong block_index = ulong(token_id) * blocks_per_row + column / QK8_0;
    const block_q8_0 block = table[block_index];
    output[output_index] = half(float(block.d) * float(block.qs[column % QK8_0]));
}

kernel void vnext_rms_norm_f16(
    device const half * input [[buffer(0)]],
    device const half * weight [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant RmsNormParams & params [[buffer(3)]],
    threadgroup float * partial [[threadgroup(0)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_group [[simdgroup_index_in_threadgroup]]) {
    const uint row = group.x;
    if (row >= params.rows) {
        return;
    }
    const ulong base = ulong(row) * params.hidden_size;
    float sum = 0.0f;
    for (uint column = lane; column < params.hidden_size; column += THREADS_PER_GROUP) {
        const float value = float(input[base + column]);
        sum += value * value;
    }
    sum = simd_sum(sum);
    if (simd_lane == 0) {
        partial[simd_group] = sum;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (simd_group == 0) {
        const uint simdgroups = THREADS_PER_GROUP / 32;
        float total = simd_lane < simdgroups ? partial[simd_lane] : 0.0f;
        total = simd_sum(total);
        if (simd_lane == 0) {
            partial[0] = rsqrt(total / float(params.hidden_size) + params.epsilon);
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const float inverse_rms = partial[0];
    for (uint column = lane; column < params.hidden_size; column += THREADS_PER_GROUP) {
        output[base + column] = half(
            float(input[base + column]) * inverse_rms * float(weight[column])
        );
    }
}

kernel void vnext_residual_add_f16(
    device const half * left [[buffer(0)]],
    device const half * right [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant ResidualAddParams & params [[buffer(3)]],
    uint index [[thread_position_in_grid]]) {
    if (index < params.elements) {
        output[index] = half(float(left[index]) + float(right[index]));
    }
}

kernel void vnext_last_token_masked_argmax_f16(
    device const half * logits [[buffer(0)]],
    device half * scratch [[buffer(1)]],
    device const uchar * valid_mask [[buffer(2)]],
    device const uint * repetition_token_ids [[buffer(3)]],
    device const uint * repetition_offsets [[buffer(4)]],
    device const float * repetition_penalty [[buffer(5)]],
    device uint * output [[buffer(6)]],
    constant LastTokenMaskedArgmaxParams & params [[buffer(7)]],
    uint lane [[thread_index_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_group [[simdgroup_index_in_threadgroup]]) {
    const float penalty = repetition_penalty[0];
    const uint repetition_start = min(repetition_offsets[0], params.repetition_capacity);
    const uint repetition_end = min(repetition_offsets[1], params.repetition_capacity);
    const bool apply_penalty = penalty != 1.0f && repetition_start < repetition_end;
    if (apply_penalty) {
        for (uint token = lane; token < params.vocabulary_size; token += THREADS_PER_GROUP) {
            scratch[token] = logits[token];
        }
        threadgroup_barrier(mem_flags::mem_device);
        for (uint entry = repetition_start + lane;
             entry < repetition_end;
             entry += THREADS_PER_GROUP) {
            const uint token = repetition_token_ids[entry];
            if (token >= params.vocabulary_size) {
                continue;
            }
            const float value = float(scratch[token]);
            if (!isfinite(value)) {
                continue;
            }
            scratch[token] = half(value > 0.0f ? value / penalty : value * penalty);
        }
        threadgroup_barrier(mem_flags::mem_device);
    }

    float local_maximum = -INFINITY;
    int local_index = -1;
    for (uint token = lane; token < params.vocabulary_size; token += THREADS_PER_GROUP) {
        if (valid_mask[token] == 0) {
            continue;
        }
        const float value = float(apply_penalty ? scratch[token] : logits[token]);
        if (!isfinite(value)) {
            continue;
        }
        if (local_index < 0 || value > local_maximum ||
            (value == local_maximum && int(token) < local_index)) {
            local_maximum = value;
            local_index = int(token);
        }
    }

    const float simd_maximum = simd_max(local_maximum);
    const int simd_index = simd_min(
        local_index >= 0 && local_maximum == simd_maximum ? local_index : 0x7fffffff
    );
    threadgroup float partial_maximum[THREADS_PER_GROUP / 32];
    threadgroup int partial_index[THREADS_PER_GROUP / 32];
    if (simd_lane == 0) {
        partial_maximum[simd_group] = simd_maximum;
        partial_index[simd_group] = simd_index;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (lane == 0) {
        float maximum = -INFINITY;
        int index = 0x7fffffff;
        for (uint group = 0; group < THREADS_PER_GROUP / 32; ++group) {
            const float candidate_maximum = partial_maximum[group];
            const int candidate_index = partial_index[group];
            if (candidate_index != 0x7fffffff &&
                (index == 0x7fffffff || candidate_maximum > maximum ||
                 (candidate_maximum == maximum && candidate_index < index))) {
                maximum = candidate_maximum;
                index = candidate_index;
            }
        }
        output[0] = index == 0x7fffffff ? 0xffffffffu : uint(index);
    }
}

kernel void vnext_embedding_dense_f32(
    device const half * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device float * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    output[output_index] = token_id < params.vocabulary_size
        ? float(table[ulong(token_id) * params.hidden_size + column])
        : 0.0f;
}

kernel void vnext_embedding_q4_k_f32(
    device const block_q4_K * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device float * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    if (token_id >= params.vocabulary_size) {
        output[output_index] = 0.0f;
        return;
    }
    const uint blocks_per_row = params.hidden_size / QK_K;
    const ulong block_index = ulong(token_id) * blocks_per_row + column / QK_K;
    output[output_index] = q4_k_value(table[block_index], column % QK_K);
}

kernel void vnext_embedding_q6_k_f32(
    device const block_q6_K * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device float * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    if (token_id >= params.vocabulary_size) {
        output[output_index] = 0.0f;
        return;
    }
    const uint blocks_per_row = params.hidden_size / QK_K;
    const ulong block_index = ulong(token_id) * blocks_per_row + column / QK_K;
    output[output_index] = q6_k_value(table[block_index], column % QK_K);
}

kernel void vnext_embedding_q8_0_f32(
    device const block_q8_0 * table [[buffer(0)]],
    device const uint * token_ids [[buffer(1)]],
    device float * output [[buffer(2)]],
    constant EmbeddingParams & params [[buffer(3)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
    const uint token = group.y;
    const uint column = group.x * THREADS_PER_GROUP + lane;
    if (token >= params.token_count || column >= params.hidden_size) {
        return;
    }
    const uint token_id = token_ids[token];
    const ulong output_index = ulong(token) * params.hidden_size + column;
    if (token_id >= params.vocabulary_size) {
        output[output_index] = 0.0f;
        return;
    }
    const uint blocks_per_row = params.hidden_size / QK8_0;
    const ulong block_index = ulong(token_id) * blocks_per_row + column / QK8_0;
    const block_q8_0 block = table[block_index];
    output[output_index] = float(block.d) * float(block.qs[column % QK8_0]);
}

kernel void vnext_rms_norm_f32_to_f16(
    device const float * input [[buffer(0)]],
    device const half * weight [[buffer(1)]],
    device half * output [[buffer(2)]],
    constant RmsNormParams & params [[buffer(3)]],
    threadgroup float * partial [[threadgroup(0)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_group [[simdgroup_index_in_threadgroup]]) {
    const uint row = group.x;
    if (row >= params.rows) {
        return;
    }
    const ulong base = ulong(row) * params.hidden_size;
    float sum = 0.0f;
    for (uint column = lane; column < params.hidden_size; column += THREADS_PER_GROUP) {
        const float value = input[base + column];
        sum += value * value;
    }
    sum = simd_sum(sum);
    if (simd_lane == 0) {
        partial[simd_group] = sum;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (simd_group == 0) {
        const uint simdgroups = THREADS_PER_GROUP / 32;
        float total = simd_lane < simdgroups ? partial[simd_lane] : 0.0f;
        total = simd_sum(total);
        if (simd_lane == 0) {
            partial[0] = rsqrt(total / float(params.hidden_size) + params.epsilon);
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const float inverse_rms = partial[0];
    for (uint column = lane; column < params.hidden_size; column += THREADS_PER_GROUP) {
        output[base + column] = half(input[base + column] * inverse_rms * float(weight[column]));
    }
}

kernel void vnext_rms_norm_f32(
    device const float * input [[buffer(0)]],
    device const half * weight [[buffer(1)]],
    device float * output [[buffer(2)]],
    constant RmsNormParams & params [[buffer(3)]],
    threadgroup float * partial [[threadgroup(0)]],
    uint3 group [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_group [[simdgroup_index_in_threadgroup]]) {
    const uint row = group.x;
    if (row >= params.rows) {
        return;
    }
    const ulong base = ulong(row) * params.hidden_size;
    float sum = 0.0f;
    for (uint column = lane; column < params.hidden_size; column += THREADS_PER_GROUP) {
        const float value = input[base + column];
        sum += value * value;
    }
    sum = simd_sum(sum);
    if (simd_lane == 0) {
        partial[simd_group] = sum;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (simd_group == 0) {
        const uint simdgroups = THREADS_PER_GROUP / 32;
        float total = simd_lane < simdgroups ? partial[simd_lane] : 0.0f;
        total = simd_sum(total);
        if (simd_lane == 0) {
            partial[0] = rsqrt(total / float(params.hidden_size) + params.epsilon);
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const float inverse_rms = partial[0];
    for (uint column = lane; column < params.hidden_size; column += THREADS_PER_GROUP) {
        output[base + column] = input[base + column] * inverse_rms * float(weight[column]);
    }
}

kernel void vnext_residual_add_f32_f16(
    device const float * left [[buffer(0)]],
    device const half * right [[buffer(1)]],
    device float * output [[buffer(2)]],
    constant ResidualAddParams & params [[buffer(3)]],
    uint index [[thread_position_in_grid]]) {
    if (index < params.elements) {
        output[index] = left[index] + float(right[index]);
    }
}

kernel void vnext_last_token_masked_argmax_f32(
    device const float * logits [[buffer(0)]],
    device float * scratch [[buffer(1)]],
    device const uchar * valid_mask [[buffer(2)]],
    device const uint * repetition_token_ids [[buffer(3)]],
    device const uint * repetition_offsets [[buffer(4)]],
    device const float * repetition_penalty [[buffer(5)]],
    device uint * output [[buffer(6)]],
    constant LastTokenMaskedArgmaxParams & params [[buffer(7)]],
    uint lane [[thread_index_in_threadgroup]],
    uint simd_lane [[thread_index_in_simdgroup]],
    uint simd_group [[simdgroup_index_in_threadgroup]]) {
    const float penalty = repetition_penalty[0];
    const uint repetition_start = min(repetition_offsets[0], params.repetition_capacity);
    const uint repetition_end = min(repetition_offsets[1], params.repetition_capacity);
    const bool apply_penalty = penalty != 1.0f && repetition_start < repetition_end;
    if (apply_penalty) {
        for (uint token = lane; token < params.vocabulary_size; token += THREADS_PER_GROUP) {
            scratch[token] = logits[token];
        }
        threadgroup_barrier(mem_flags::mem_device);
        for (uint entry = repetition_start + lane;
             entry < repetition_end;
             entry += THREADS_PER_GROUP) {
            const uint token = repetition_token_ids[entry];
            if (token >= params.vocabulary_size) {
                continue;
            }
            const float value = scratch[token];
            if (!isfinite(value)) {
                continue;
            }
            scratch[token] = value > 0.0f ? value / penalty : value * penalty;
        }
        threadgroup_barrier(mem_flags::mem_device);
    }

    float local_maximum = -INFINITY;
    int local_index = -1;
    for (uint token = lane; token < params.vocabulary_size; token += THREADS_PER_GROUP) {
        if (valid_mask[token] == 0) {
            continue;
        }
        const float value = apply_penalty ? scratch[token] : logits[token];
        if (!isfinite(value)) {
            continue;
        }
        if (local_index < 0 || value > local_maximum ||
            (value == local_maximum && int(token) < local_index)) {
            local_maximum = value;
            local_index = int(token);
        }
    }

    const float simd_maximum = simd_max(local_maximum);
    const int simd_index = simd_min(
        local_index >= 0 && local_maximum == simd_maximum ? local_index : 0x7fffffff
    );
    threadgroup float partial_maximum[THREADS_PER_GROUP / 32];
    threadgroup int partial_index[THREADS_PER_GROUP / 32];
    if (simd_lane == 0) {
        partial_maximum[simd_group] = simd_maximum;
        partial_index[simd_group] = simd_index;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (lane == 0) {
        float maximum = -INFINITY;
        int index = 0x7fffffff;
        for (uint group = 0; group < THREADS_PER_GROUP / 32; ++group) {
            const float candidate_maximum = partial_maximum[group];
            const int candidate_index = partial_index[group];
            if (candidate_index != 0x7fffffff &&
                (index == 0x7fffffff || candidate_maximum > maximum ||
                 (candidate_maximum == maximum && candidate_index < index))) {
                maximum = candidate_maximum;
                index = candidate_index;
            }
        }
        output[0] = index == 0x7fffffff ? 0xffffffffu : uint(index);
    }
}