Quantization Fundamentals¶

pub const BlockQ8_0 = extern struct {
    d: f16,             // scale
    qs: [32]i8,         // quantized values

    pub fn dequantize(self: BlockQ8_0, output: *[32]f32) void {
        const scale: f32 = @floatCast(self.d);
        for (0..32) |i| {
            output[i] = @as(f32, @floatFromInt(self.qs[i])) * scale;
        }
    }
};

const low: u4 = @truncate(byte);       // value at even index
const high: u4 = @truncate(byte >> 4); // value at odd index

pub const BlockQ4_0 = extern struct {
    d: f16,             // scale
    qs: [16]u8,         // packed 4-bit values (2 per byte)

    pub fn dequantize(self: BlockQ4_0, output: *[32]f32) void {
        const scale: f32 = @floatCast(self.d);
        for (0..16) |j| {
            const byte = self.qs[j];
            const low: i32 = @as(i32, @as(u4, @truncate(byte))) - 8;
            const high: i32 = @as(i32, @as(u4, @truncate(byte >> 4))) - 8;
            output[2 * j] = @as(f32, @floatFromInt(low)) * scale;
            output[2 * j + 1] = @as(f32, @floatFromInt(high)) * scale;
        }
    }
};

pub const QuantizedINT8 = struct {
    scale: f32,
    zero_point: i32,
    data: []const u8,

    pub fn dequantize(self: QuantizedINT8, output: []f32) void {
        for (self.data, 0..) |q, i| {
            output[i] = self.scale *
                @as(f32, @floatFromInt(@as(i32, q) - self.zero_point));
        }
    }
};

pub fn QuantizedTensor(comptime quant_type: QuantType) type {
    const BlockType = quant_type.BlockType();

    return struct {
        const Self = @This();

        blocks: []const BlockType,
        n_elements: usize,

        /// Dequantize the entire tensor into a pre-allocated f32 buffer.
        pub fn dequantize(self: Self, output: []f32) void {
            std.debug.assert(output.len >= self.n_elements);
            const block_size = quant_type.blockSize();
            for (self.blocks, 0..) |block, bi| {
                const start = bi * block_size;
                const end = @min(start + block_size, self.n_elements);
                block.dequantize(output[start..end]);
            }
        }

        /// Dequantize a single block at the given index.
        pub fn dequantizeBlock(self: Self, block_idx: usize, output: []f32) void {
            self.blocks[block_idx].dequantize(output);
        }

        /// Return the quantized dot product with a f32 vector.
        /// Dequantizes on the fly to avoid materialising the full tensor.
        pub fn dot(self: Self, vec: []const f32) f32 {
            var sum: f32 = 0.0;
            const block_size = quant_type.blockSize();
            var buf: [block_size]f32 = undefined;

            for (self.blocks, 0..) |block, bi| {
                const start = bi * block_size;
                block.dequantize(&buf);
                const len = @min(block_size, self.n_elements - start);
                for (0..len) |i| {
                    sum += buf[i] * vec[start + i];
                }
            }
            return sum;
        }

        /// Number of quantized blocks.
        pub fn numBlocks(self: Self) usize {
            return self.blocks.len;
        }

        /// Storage size in bytes.
        pub fn sizeBytes(self: Self) usize {
            return self.blocks.len * @sizeOf(BlockType);
        }
    };
}

pub const QuantType = enum {
    q8_0,
    q4_0,
    int8,
    q4_k,
    q5_k,
    q6_k,
    iq1_s,
    iq2_xs,
    iq3_s,
    iq4_xs,
    iq4_nl,
    // ... additional formats

    pub fn BlockType(comptime self: QuantType) type {
        return switch (self) {
            .q8_0 => BlockQ8_0,
            .q4_0 => BlockQ4_0,
            .q4_k => BlockQ4K,
            // ... etc.
        };
    }

    pub fn blockSize(comptime self: QuantType) usize {
        return switch (self) {
            .q8_0, .q4_0 => 32,
            .q4_k, .q5_k, .q6_k => 256,
            // ... etc.
        };
    }

    pub fn bitsPerWeight(self: QuantType) f32 {
        return switch (self) {
            .q8_0 => 8.5,
            .q4_0 => 4.5,
            .int8 => 8.0,
            .q4_k => 4.5,
            .q5_k => 5.5,
            .q6_k => 6.5,
            // ... etc.
        };
    }
};