|
| 1 | +# |
| 2 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | +# All rights reserved. |
| 4 | +# |
| 5 | +# This source code is licensed under the BSD-style license found in the |
| 6 | +# LICENSE file in the root directory of this source tree. |
| 7 | +# |
| 8 | + |
| 9 | +"""Shared GGUF **Q5_K** primitives for the MLX backend. |
| 10 | +
|
| 11 | +This module holds the pieces common to every Q5_K kernel (linear matmul/matvec |
| 12 | +and the embedding gather), so format-specific op modules import from here: |
| 13 | +
|
| 14 | +* ``QK_K`` / ``K_SCALE_SIZE`` / ``Q5K_BLOCK_BYTES`` and the per-super-block byte |
| 15 | + layout constants. |
| 16 | +* ``_Q5K_HEADER`` -- the Metal header (the ``block_q5_K`` struct plus the |
| 17 | + per-element and vectorized dequant helpers) shared by all Q5_K Metal kernels. |
| 18 | +
|
| 19 | +Q5_K layout (per 256-element super-block, 176 bytes, see llama.cpp |
| 20 | +``block_q5_K`` in ``ggml-common.h``):: |
| 21 | +
|
| 22 | + half d # super-block scale for the quantized scales |
| 23 | + half dmin # super-block scale for the quantized mins |
| 24 | + uint8 scales[12] # 6-bit packed sub-block scales + mins (as Q4_K) |
| 25 | + uint8 qh[QK_K/8 = 32] # quants, high bit (the 5th bit) |
| 26 | + uint8 qs[QK_K/2 = 128] # quants, low 4 bits |
| 27 | +
|
| 28 | +Q5_K combines Q4_K's affine super-block (``d``/``dmin`` with 6-bit packed |
| 29 | +scales/mins unpacked via ``get_scale_min_k4``) with a Q6_K-style high-bit array: |
| 30 | +each weight is a 5-bit code ``q`` (0..31) whose low 4 bits come from ``qs`` and |
| 31 | +whose 5th bit comes from ``qh``. The dequantized value for code ``q`` in |
| 32 | +sub-block ``s`` is ``d * scale[s] * q - dmin * min[s]`` (affine, same shape as |
| 33 | +Q4_K but with ``q`` in 0..31). |
| 34 | +
|
| 35 | +Attribution |
| 36 | +----------- |
| 37 | +The Q5_K block layout and the Metal dequant helpers in ``_Q5K_HEADER`` follow |
| 38 | +llama.cpp (``ggml-common.h`` / ``ggml-metal.metal``: ``block_q5_K``, |
| 39 | +``dequantize_q5_K``, ``dequantize_row_q5_K``, ``get_scale_min_k4``), which is |
| 40 | +MIT-licensed (Copyright (c) 2023-2024 The ggml authors). |
| 41 | +""" |
| 42 | + |
| 43 | +from __future__ import annotations |
| 44 | + |
| 45 | +# --------------------------------------------------------------------------- |
| 46 | +# Q5_K constants |
| 47 | +# --------------------------------------------------------------------------- |
| 48 | + |
| 49 | +QK_K = 256 |
| 50 | +K_SCALE_SIZE = 12 |
| 51 | +_Q5K_D_BYTES = 2 |
| 52 | +_Q5K_DMIN_BYTES = 2 |
| 53 | +_Q5K_SCALES_BYTES = K_SCALE_SIZE |
| 54 | +_Q5K_QH_BYTES = QK_K // 8 # 32 |
| 55 | +_Q5K_QS_BYTES = QK_K // 2 # 128 |
| 56 | +Q5K_BLOCK_BYTES = ( |
| 57 | + _Q5K_D_BYTES + _Q5K_DMIN_BYTES + _Q5K_SCALES_BYTES + _Q5K_QH_BYTES + _Q5K_QS_BYTES |
| 58 | +) # 176 |
| 59 | + |
| 60 | + |
| 61 | +# --------------------------------------------------------------------------- |
| 62 | +# Shared Metal header |
| 63 | +# --------------------------------------------------------------------------- |
| 64 | + |
| 65 | +# Ported from llama.cpp ggml-common.h (block_q5_K, get_scale_min_k4) and |
| 66 | +# ggml-metal.metal (dequantize_q5_K). Struct field order matches GGUF bytes: |
| 67 | +# d(0:2), dmin(2:4), scales(4:16), qh(16:48), qs(48:176). |
| 68 | +_Q5K_HEADER = """ |
| 69 | +#include <metal_simdgroup> |
| 70 | +#include <metal_simdgroup_matrix> |
| 71 | +using namespace metal; |
| 72 | +
|
| 73 | +#define QK_K 256 |
| 74 | +#define K_SCALE_SIZE 12 |
| 75 | +
|
| 76 | +typedef struct { |
| 77 | + half d; // super-block scale for quantized scales |
| 78 | + half dmin; // super-block scale for quantized mins |
| 79 | + uint8_t scales[K_SCALE_SIZE]; // 6-bit packed scales + mins (same as Q4_K) |
| 80 | + uint8_t qh[QK_K/8]; // quants, high bit (the 5th bit) |
| 81 | + uint8_t qs[QK_K/2]; // quants, low 4 bits |
| 82 | +} block_q5_K; // 176 bytes |
| 83 | +
|
| 84 | +// Unpack 6-bit scale and min for sub-block index j (0..7). |
| 85 | +// Ported from llama.cpp get_scale_min_k4 (ggml-quants.c). Identical to Q4_K. |
| 86 | +inline void get_scale_min_k4(int j, device const uint8_t * q, |
| 87 | + thread uint8_t & sc, thread uint8_t & m) { |
| 88 | + if (j < 4) { |
| 89 | + sc = q[j] & 63; |
| 90 | + m = q[j + 4] & 63; |
| 91 | + } else { |
| 92 | + sc = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); |
| 93 | + m = (q[j + 4] >> 4) | ((q[j] >> 6) << 4); |
| 94 | + } |
| 95 | +} |
| 96 | +
|
| 97 | +// Metal variant used by dequantize_q5_K_16 (matches ggml-metal.metal). Same as Q4_K. |
| 98 | +inline uchar2 get_scale_min_k4_just2(int j, int k, device const uint8_t * q) { |
| 99 | + return j < 4 |
| 100 | + ? uchar2{uint8_t(q[j + 0 + k] & 63), uint8_t(q[j + 4 + k] & 63)} |
| 101 | + : uchar2{ |
| 102 | + uint8_t((q[j + 4 + k] & 0xF) | ((q[j - 4 + k] >> 6) << 4)), |
| 103 | + uint8_t((q[j + 4 + k] >> 4) | ((q[j + 0 + k] >> 6) << 4))}; |
| 104 | +} |
| 105 | +
|
| 106 | +// Dequantize a single element at within-block position p (0..255). |
| 107 | +// Mirrors dequantize_row_q5_K (ggml-quants.c): 64-element chunks, 32 lows then |
| 108 | +// 32 highs per chunk; the 5th bit comes from qh at bit (2*chunk + half). |
| 109 | +inline float dequant_q5k_elem(device const block_q5_K * blk, int p) { |
| 110 | + const int c = p >> 6; // 0..3 (64-element chunk) |
| 111 | + const int within = p & 63; // 0..63 within chunk |
| 112 | + const int l = within & 31; // 0..31 (also the qh byte index) |
| 113 | + const int half = within >> 5; // 0 low nibble, 1 high nibble |
| 114 | + const int sub = 2 * c + half; // sub-block 0..7 (also the qh bit index) |
| 115 | +
|
| 116 | + const uint8_t byte = blk->qs[c * 32 + l]; |
| 117 | + const int nib = (half == 0) ? (byte & 0xF) : (byte >> 4); |
| 118 | + const int hibit = (blk->qh[l] >> sub) & 1; |
| 119 | + const int q = nib + (hibit << 4); // 0..31 |
| 120 | +
|
| 121 | + uint8_t sc, mn; |
| 122 | + get_scale_min_k4(sub, blk->scales, sc, mn); |
| 123 | +
|
| 124 | + const float d = (float) blk->d; |
| 125 | + const float dm = (float) blk->dmin; |
| 126 | + return d * (float) sc * (float) q - dm * (float) mn; |
| 127 | +} |
| 128 | +
|
| 129 | +// Vectorized Q5_K dequantize: decodes 16 values into half4x4. |
| 130 | +// Ported from llama.cpp dequantize_q5_K (ggml-metal.metal): Q4_K's vectorized |
| 131 | +// decode plus the Q6_K-style 5th bit. mat-mat uses NL=16 (QK_K/16); for sub-block |
| 132 | +// `il` the 16 outputs map to block positions [il*16 : il*16+16]. |
| 133 | +inline void dequantize_q5_K_16(device const block_q5_K * xb, short il, |
| 134 | + thread half4x4 & reg) { |
| 135 | + device const uint8_t * q = xb->qs; |
| 136 | + device const uint8_t * qh = xb->qh; |
| 137 | +
|
| 138 | + short is = (il / 4) * 2; |
| 139 | + q = q + (il / 4) * 32 + 16 * (il & 1); |
| 140 | + qh = qh + 16 * (il & 1); |
| 141 | + const uint8_t ul = 1 << (il / 2); // 5th-bit mask (uses original il) |
| 142 | + il = il & 3; |
| 143 | +
|
| 144 | + const uchar2 sc = get_scale_min_k4_just2(is, il / 2, xb->scales); |
| 145 | + const float d = il < 2 ? (float) xb->d : (float) xb->d / 16.f; |
| 146 | + const float dm = (float) xb->dmin; |
| 147 | + const float dl = d * (float) sc[0]; |
| 148 | + const float ml = dm * (float) sc[1]; |
| 149 | +
|
| 150 | + const ushort mask = il < 2 ? 0x0F : 0xF0; |
| 151 | + // Low nibble (mask 0x0F): the 5th bit adds 16. High nibble (mask 0xF0) is the |
| 152 | + // value pre-shifted by 4 with d/16, so the 5th bit adds 16*16 = 256. |
| 153 | + const float qh_val = il < 2 ? 16.f : 256.f; |
| 154 | + for (int i = 0; i < 16; ++i) { |
| 155 | + const float v = (float)(q[i] & mask) + ((qh[i] & ul) ? qh_val : 0.f); |
| 156 | + reg[i / 4][i % 4] = (half)(dl * v - ml); |
| 157 | + } |
| 158 | +} |
| 159 | +""" |
0 commit comments