K-Quantization¶

pub const QK_K: usize = 256;
pub const K_SUB_BLOCKS: usize = 8;
pub const K_SUB_BLOCK_SIZE: usize = QK_K / K_SUB_BLOCKS; // 32

pub const BlockQ4K = extern struct {
    d: f16,              // super-block scale
    dmin: f16,           // super-block minimum
    scales: [12]u8,      // packed sub-block scales and minimums
    qs: [128]u8,         // packed 4-bit quantized values

    pub fn dequantize(self: BlockQ4K, output: *[QK_K]f32) void {
        const d: f32 = @floatCast(self.d);
        const dmin: f32 = @floatCast(self.dmin);
        const sm = extractScales(self.scales);

        for (0..K_SUB_BLOCKS) |j| {
            const sc: f32 = @floatFromInt(sm.sub_scales[j]);
            const mn: f32 = @floatFromInt(sm.sub_mins[j]);
            const base = j * K_SUB_BLOCK_SIZE;

            for (0..K_SUB_BLOCK_SIZE / 2) |k| {
                const byte = self.qs[base / 2 + k];
                const q_lo: f32 = @floatFromInt(@as(u4, @truncate(byte)));
                const q_hi: f32 = @floatFromInt(@as(u4, @truncate(byte >> 4)));

                output[base + 2 * k] = d * sc * q_lo - dmin * mn;
                output[base + 2 * k + 1] = d * sc * q_hi - dmin * mn;
            }
        }
    }
};

pub const BlockQ5K = extern struct {
    d: f16,
    dmin: f16,
    scales: [12]u8,
    qh: [32]u8,     // high bit for each of 256 values
    qs: [128]u8,     // low 4 bits, packed

    pub fn dequantize(self: BlockQ5K, output: *[QK_K]f32) void {
        const d_val: f32 = @floatCast(self.d);
        const dmin_val: f32 = @floatCast(self.dmin);
        const sm = extractScales(self.scales);

        for (0..K_SUB_BLOCKS) |j| {
            const sc: f32 = @floatFromInt(sm.sub_scales[j]);
            const mn: f32 = @floatFromInt(sm.sub_mins[j]);
            const base = j * K_SUB_BLOCK_SIZE;

            for (0..K_SUB_BLOCK_SIZE / 2) |k| {
                const idx = base + 2 * k;
                const byte = self.qs[base / 2 + k];
                const lo: u8 = byte & 0x0F;
                const hi: u8 = byte >> 4;

                // Extract high bit from qh
                const h_lo: u8 = (self.qh[idx / 8] >> @truncate(idx % 8)) & 1;
                const h_hi: u8 = (self.qh[(idx + 1) / 8] >> @truncate((idx + 1) % 8)) & 1;

                const q5_lo: f32 = @floatFromInt(@as(u8, lo + 16 * h_lo));
                const q5_hi: f32 = @floatFromInt(@as(u8, hi + 16 * h_hi));

                output[idx] = d_val * sc * q5_lo - dmin_val * mn;
                output[idx + 1] = d_val * sc * q5_hi - dmin_val * mn;
            }
        }
    }
};

pub const BlockQ6K = extern struct {
    ql: [128]u8,       // low 4 bits (packed, 2 per byte)
    qh: [64]u8,        // high 2 bits (packed, 4 per byte)
    scales: [16]i8,     // sub-block scales (16 sub-blocks of 16 values)
    d: f16,             // super-block scale

    pub fn dequantize(self: BlockQ6K, output: *[QK_K]f32) void {
        const d_val: f32 = @floatCast(self.d);

        for (0..QK_K) |i| {
            // Extract low 4 bits
            const ql_byte = self.ql[i / 2];
            const lo: u8 = if (i % 2 == 0) ql_byte & 0x0F else ql_byte >> 4;

            // Extract high 2 bits
            const qh_byte = self.qh[i / 4];
            const shift: u3 = @truncate((i % 4) * 2);
            const hi: u8 = (qh_byte >> shift) & 0x03;

            // Reconstruct 6-bit value
            const q6: i32 = @as(i32, lo + 16 * hi) - 32;

            // Sub-block scale (16 sub-blocks of 16 values each)
            const sub_idx = i / 16;
            const sc: f32 = @floatFromInt(self.scales[sub_idx]);

            output[i] = d_val * sc * @as(f32, @floatFromInt(q6));
        }
    }
};

pub const KQuantizer = struct {
    allocator: std.mem.Allocator,

    pub fn init(allocator: std.mem.Allocator) KQuantizer {
        return .{ .allocator = allocator };
    }

    /// Quantize a f32 slice to Q4_K blocks.
    pub fn quantizeQ4K(
        self: KQuantizer,
        data: []const f32,
    ) ![]BlockQ4K {
        const n_blocks = (data.len + QK_K - 1) / QK_K;
        const blocks = try self.allocator.alloc(BlockQ4K, n_blocks);

        for (0..n_blocks) |bi| {
            const start = bi * QK_K;
            const end = @min(start + QK_K, data.len);
            blocks[bi] = quantizeBlockQ4K(data[start..end]);
        }

        return blocks;
    }

    /// Quantize a f32 slice to Q5_K blocks.
    pub fn quantizeQ5K(
        self: KQuantizer,
        data: []const f32,
    ) ![]BlockQ5K {
        const n_blocks = (data.len + QK_K - 1) / QK_K;
        const blocks = try self.allocator.alloc(BlockQ5K, n_blocks);

        for (0..n_blocks) |bi| {
            const start = bi * QK_K;
            const end = @min(start + QK_K, data.len);
            blocks[bi] = quantizeBlockQ5K(data[start..end]);
        }

        return blocks;
    }

    /// Quantize a f32 slice to Q6_K blocks.
    pub fn quantizeQ6K(
        self: KQuantizer,
        data: []const f32,
    ) ![]BlockQ6K {
        const n_blocks = (data.len + QK_K - 1) / QK_K;
        const blocks = try self.allocator.alloc(BlockQ6K, n_blocks);

        for (0..n_blocks) |bi| {
            const start = bi * QK_K;
            const end = @min(start + QK_K, data.len);
            blocks[bi] = quantizeBlockQ6K(data[start..end]);
        }

        return blocks;
    }

    /// Dequantize an entire Q4_K block array to f32.
    pub fn dequantizeQ4K(
        blocks: []const BlockQ4K,
        output: []f32,
    ) void {
        for (blocks, 0..) |block, bi| {
            block.dequantize(output[bi * QK_K ..][0..QK_K]);
        }
    }

    pub fn deinit(self: *KQuantizer, blocks: anytype) void {
        self.allocator.free(blocks);
    }
};