"""IQ-quant formats (llama.cpp i-quants): packed slice-dequant + quantizers. All i-quants are codebook formats: weight values are indices into a lattice grid; the grid entry (8 or 4 values) is scaled by a block delta and optionally sign-flipped. Grids come from llama.cpp ggml-common.h (see _iq_tables.py). Formats (block QK_K = 256, IQ4_NL uses QK4_NL = 32): IQ2_XXS 2.0625 bpw d(f16) + qs[64B] = 66 B IQ2_XS 2.3125 bpw d(f16) + qs[64B] + scales[8] = 74 B IQ2_S 2.5625 bpw d(f16) + qs[64] + qh[8] + scales[8] = 82 B IQ3_XXS 3.0625 bpw d(f16) + qs[96] = 98 B IQ3_S 3.4375 bpw d(f16) + qs[64] + qh[8] + signs[32] + scales[4] = 110 B IQ1_S 1.5625 bpw d(f16) + qs[32] + qh[16]u16 = 66 B IQ1_M 1.75 bpw qs[32] + qh[16] + scales[8]u16 = 56 B IQ4_NL 4.5 bpw d(f16) + qs[16] (block 32) = 18 B IQ4_XS 4.25 bpw d(f16) + scales_h[2] + scales_l[4] + qs[128] = 136 B Dequant follows the CPU reference (ggml-quants.c, dequantize_row_iq*_*) so results match llama.cpp bit-for-bit on the same packed bytes. All functions slice output rows [start:end] and are GPU-friendly (torch ops only). Quantization (quantize_row_iq*_ref): for each 32-value group the sign pattern is chosen so the flipped magnitude vector lies on the grid; the best scale is found by weighted least squares over candidate scales; grid rows are looked up via the precomputed map + neighbour lists (same as iq2xs_init_impl). Author: AGIWS NeuralQuant team License: Apache 2.0 """ from __future__ import annotations import torch from agiws_neural_quant.kquant._iq_tables import iq_table from agiws_neural_quant.kquant.gguf_pack import _cat_blocks QK_K = 256 QK4_NL = 32 IQ1S_DELTA = 0.125 IQ1M_DELTA = 0.125 # GGML dtype ids (see converters/gguf_reader.py). GGML_TYPE_IQ2_XXS = 16 GGML_TYPE_IQ2_XS = 17 GGML_TYPE_IQ3_XXS = 18 GGML_TYPE_IQ1_S = 19 GGML_TYPE_IQ4_NL = 20 GGML_TYPE_IQ3_S = 21 GGML_TYPE_IQ2_S = 22 GGML_TYPE_IQ4_XS = 23 GGML_TYPE_IQ1_M = 29 _IQ_BYTES_PER_BLOCK = { GGML_TYPE_IQ2_XXS: 66, GGML_TYPE_IQ2_XS: 74, GGML_TYPE_IQ2_S: 82, GGML_TYPE_IQ3_XXS: 98, GGML_TYPE_IQ3_S: 110, GGML_TYPE_IQ1_S: 66, GGML_TYPE_IQ1_M: 56, GGML_TYPE_IQ4_NL: 18, GGML_TYPE_IQ4_XS: 136, } def bytes_per_row_iq(ggml_dtype: int, cols: int) -> int: block = QK4_NL if ggml_dtype == GGML_TYPE_IQ4_NL else QK_K nbl = (cols + block - 1) // block return nbl * _IQ_BYTES_PER_BLOCK[ggml_dtype] def _f16(u8: torch.Tensor) -> torch.Tensor: """Interpret raw [..., 2] bytes as fp16 -> fp32 (LE). Returns [...]. """ return u8.view(torch.float16)[..., 0].to(torch.float32) def _le_u16(u8: torch.Tensor) -> torch.Tensor: """[..., 2] bytes -> int64 (LE).""" return (u8[..., 0].to(torch.int64) | (u8[..., 1].to(torch.int64) << 8)) def _signs_apply(vals: torch.Tensor, signs: torch.Tensor) -> torch.Tensor: """Flip sign of last-dim 8-tuples where the sign byte bit is set. vals [..., 8], signs uint8 [..., 1] (last dim = 8-tuple index). """ mask = iq_table("kmask_iq2xs").to(signs.device) # [8] bit = (signs.unsqueeze(-1) & mask).to(vals.dtype) # [..., 1, 8] -> [..., 8] return torch.where(bit != 0, -vals, vals) def _ksigns(idx: torch.Tensor) -> torch.Tensor: """128-entry ksigns table lookup.""" t = iq_table("ksigns_iq2xs").to(idx.device) return t[idx] # --------------------------------------------------------------------------- # IQ2_XXS -- block 66 B: d(2) + qs(32 x u16 = 64 B) # per 32-group ib (8 groups): # q2 = qs[4*ib..4*ib+3] (4 u16) # aux32 = q2[2] | q2[3]<<16 (upper half carries scale+signs) # db = d * (0.5 + (aux32 >> 28)) * 0.25 # signs = ksigns[(aux32 >> 7*l) & 127] # grid = iq2xxs_grid[ (u8)q2[l] ] (l=0..3) -> 8 values # --------------------------------------------------------------------------- def dequant_iq2_xxs_packed_rows( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) # [cr, nbl*66] cr = chunk.shape[0] nblk = chunk.shape[1] // 66 chunk = chunk.view(cr, nblk, 66) d = _f16(chunk[..., 0:2]) # [cr, nblk] qs = chunk[..., 2:66].view(cr, nblk, 32, 2) # 32 u16 LE q2 = _le_u16(qs).view(cr, nblk, 8, 4) # [cr, nblk, 8 grp, 4 u16] aux32 = q2[..., 2] | (q2[..., 3] << 16) # [cr, nblk, 8] db = d.unsqueeze(-1) * (0.5 + ((aux32 >> 28) & 0xF).to(torch.float32)) * 0.25 # [cr,nblk,8] l4 = torch.arange(4, device=raw.device) signs7 = (aux32.unsqueeze(-1) >> (7 * l4).view(1, 1, 1, 4)) & 127 # [cr,nblk,8,4] signs = _ksigns(signs7) # [cr,nblk,8,4] g = iq_table("iq2xxs_grid").to(raw.device) # [256, 8] int8 q0 = q2[..., 0] q1 = q2[..., 1] idx = torch.stack([q0 & 0xFF, q0 >> 8, q1 & 0xFF, q1 >> 8], dim=-1) # [cr,nblk,8,4] vals = g[idx.to(torch.int64)] # [cr,nblk,8,4,8] vals = _signs_apply(vals, signs) out = db.unsqueeze(-1).unsqueeze(-1) * vals # [cr,nblk,8,4,8] return out.reshape(cr, nblk * QK_K)[:, :cols] # --------------------------------------------------------------------------- # IQ2_XS -- block 74 B: d(2) + qs(64 B) + scales(8) # qs: 32 u16; per group ib (8): q2 = qs16[4*ib+il] (il 0..3) # grid = iq2xs_grid[q2 & 511], signs = ksigns[q2 >> 9] # db = d * (0.5 + ((scales[ib] >> 4*(il/2)) & 0xF)) * 0.25 # --------------------------------------------------------------------------- def dequant_iq2_xs_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 74 chunk = chunk.view(cr, nblk, 74) d = _f16(chunk[..., 0:2]) qs = chunk[..., 2:66].view(cr, nblk, 32, 2) q2 = _le_u16(qs) # [cr, nblk, 32] q2 = q2.view(cr, nblk, 8, 4) # [cr, nblk, 8, 4] scales = chunk[..., 66:74] # [cr, nblk, 8] gr = iq_table("iq2xs_grid").to(raw.device) # [512, 8] vals = gr[(q2 & 0x1FF).to(torch.int64)] # [cr, nblk, 8, 4, 8] signs = _ksigns((q2 >> 9).to(torch.int64)) vals = _signs_apply(vals, signs) il = torch.arange(4, device=raw.device) code = (scales.to(torch.int64).unsqueeze(-1) >> (4 * (il // 2)).view(1, 1, 1, 4)) & 0xF # [cr, nblk, 8, 4] db = d.unsqueeze(-1).unsqueeze(-1) * (0.5 + code.to(torch.float32)) * 0.25 # [cr,nblk,8,4] out = (db.unsqueeze(-1) * vals).reshape(cr, nblk * QK_K)[:, :cols] return out # --------------------------------------------------------------------------- # IQ2_S -- block 82 B: d(2) + qs(64) + qh(8) + scales(8) # per group ib (8), il 0..3: # grid = iq2s_grid[ qs[4*ib+il] | ((qh[ib] << (8-2*il)) & 0x300) ] # signs = qs[32 + 4*ib + il] # db = d * (0.5 + ((scales[8] >> 4*(il/2)) & 0xF)) * 0.25 # --------------------------------------------------------------------------- def dequant_iq2_s_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 82 chunk = chunk.view(cr, nblk, 82) d = _f16(chunk[..., 0:2]) qs = chunk[..., 2:66].to(torch.int64) # [cr, nblk, 64] qh = chunk[..., 66:74].to(torch.int64) # [cr, nblk, 8] scales = chunk[..., 74:82].to(torch.int64) # [cr, nblk, 8] il = torch.arange(4, device=raw.device) # Layout: indexes = qs[0..31] (4 per super-group), signs = qs[32..63]. idx_b = qs[..., :32].view(cr, nblk, 8, 4) sign_b = qs[..., 32:64].view(cr, nblk, 8, 4) gi = idx_b + ((qh.view(cr, nblk, 8, 1) << (8 - 2 * il.view(1, 1, 1, 4))) & 0x300) gr = iq_table("iq2s_grid").to(raw.device) # [1024, 8] vals = gr[gi] # [cr, nblk, 8, 4, 8] vals = _signs_apply(vals, sign_b) code = (scales.view(cr, nblk, 8, 1) >> (4 * (il.view(1, 1, 1, 4) // 2))) & 0xF db = d.unsqueeze(-1).unsqueeze(-1) * (0.5 + code.to(torch.float32)) * 0.25 out = (db.unsqueeze(-1) * vals).reshape(cr, nblk * QK_K)[:, :cols] return out # --------------------------------------------------------------------------- # IQ3_XXS -- block 98 B: d(2) + qs(96) # qs[0:64] = 64 grid index bytes (8 groups x 8) # qs[64:96] = scales_and_signs (8 groups x 4 = 32 B) # per group ib: aux32 = sas[4*ib..4*ib+3] (LE u32) # db = d * (0.5 + (aux32 >> 28)) * 0.5 # signs = ksigns[(aux32 >> 7*il) & 127] (il 0..3) # grid1 = iq3xxs_grid[qs[8*ib + 2*il]], grid2 = [qs[8*ib+2*il+1]] # y[j] = db*g1[j]*sgn, y[j+4] = db*g2[j]*sgn (j 0..3) # --------------------------------------------------------------------------- def dequant_iq3_xxs_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 98 chunk = chunk.view(cr, nblk, 98) d = _f16(chunk[..., 0:2]) qs = chunk[..., 2:66].to(torch.int64) # [cr, nblk, 64] gas = chunk[..., 66:98].view(cr, nblk, 16, 2) # scales+signs 16 u16 aux = _le_u16(gas).view(cr, nblk, 8, 2) # [cr, nblk, 8 grp, 2 u16] aux32 = (aux[..., 0] | (aux[..., 1] << 16)).to(torch.int64) # [cr, nblk, 8] db = d.unsqueeze(-1) * (0.5 + ((aux32 >> 28) & 0xF).to(torch.float32)) * 0.5 l4 = torch.arange(4, device=raw.device) signs7 = (aux32.unsqueeze(-1) >> (7 * l4).view(1, 1, 1, 4)) & 127 # [cr, nblk, 8, 4] signs = _ksigns(signs7) gr = iq_table("iq3xxs_grid").to(raw.device) # [256, 4] idx = qs.view(cr, nblk, 8, 8) # [cr, nblk, 8 grp, 8 idx] g1 = gr[idx[..., 0::2].to(torch.int64)] # [cr, nblk, 8, 4, 4] g2 = gr[idx[..., 1::2].to(torch.int64)] vals = torch.cat([g1, g2], dim=-1) # [cr, nblk, 8, 4, 8] vals = _signs_apply(vals, signs) out = (db.unsqueeze(-1).unsqueeze(-1) * vals).reshape(cr, nblk * QK_K)[:, :cols] return out # --------------------------------------------------------------------------- # IQ3_S -- block 110 B: d(2) + qs(64) + qh(8) + signs(32) + scales(4) # per group ib (0..7), il (0..3): # grid1 = iq3s_grid[ qs[8*ib + 2*il] | ((qh[ib] << (8-2*il)) & 256) ] # grid2 = iq3s_grid[ qs[8*ib + 2*il+1] | ((qh[ib] << (7-2*il)) & 256) ] # db = d * (1 + 2*((scales[ib/2] >> 4*(ib%2)) & 0xF)) # signs = signs[4*ib + il] # y[j] = db*grid1[j]*sgn, y[j+4] = db*grid2[j]*sgn (j 0..3) # --------------------------------------------------------------------------- def dequant_iq3_s_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 110 chunk = chunk.view(cr, nblk, 110) d = _f16(chunk[..., 0:2]) # [cr, nblk] qs = chunk[..., 2:66].to(torch.int64) # [cr, nblk, 64] qh = chunk[..., 66:74].to(torch.int64) # [cr, nblk, 8] signs = chunk[..., 74:106] # [cr, nblk, 32] scales = chunk[..., 106:110].to(torch.int64) # [cr, nblk, 4] ib = torch.arange(8, device=raw.device) il = torch.arange(4, device=raw.device) qs8 = qs.view(cr, nblk, 8, 8) # [cr, nblk, grp, 8 idx bytes] qh_sel = qh.view(cr, nblk, 8, 1) gi1 = qs8[..., 0::2] + ((qh_sel << (8 - 2 * il.view(1, 1, 1, 4))) & 256) # [cr,nblk,8,4] gi2 = qs8[..., 1::2] + ((qh_sel << (7 - 2 * il.view(1, 1, 1, 4))) & 256) gr = iq_table("iq3s_grid").to(raw.device) # [512, 4] v1 = gr[gi1] # [cr, nblk, 8, 4, 4] v2 = gr[gi2] vals = torch.cat([v1, v2], dim=-1) # [cr, nblk, 8, 4, 8] # db per group (8 groups: scales[ib/2], nibble by ib%2) code = (scales.view(cr, nblk, 4).repeat_interleave(2, dim=2) >> (4 * (ib % 2).view(1, 1, 8))) & 0xF # [cr, nblk, 8] db = d.unsqueeze(-1) * (1 + 2 * code.to(torch.float32)) # [cr, nblk, 8] # signs per group/il sgn = signs.view(cr, nblk, 8, 4) # [cr, nblk, 8, 4] vals = _signs_apply(vals, sgn) out = db.unsqueeze(-1).unsqueeze(-1) * vals # [cr, nblk, 8, 4, 8] return out.reshape(cr, nblk * QK_K)[:, :cols] # --------------------------------------------------------------------------- # IQ1_S -- block 66 B: d(2) + qs(32) + qh(16 u16) # per 32-group ib (8): # dl = d * (2*((qh[ib] >> 12) & 7) + 1) # delta = qh[ib] & 0x8000 ? -IQ1S_DELTA : +IQ1S_DELTA # grid idx = qs[4*ib+il] | (((qh[ib] >> 3*il) & 7) << 8) (il 0..3) # grid = iq1s_grid[int8] (signed! values in {-1, 0, 1}-ish) # y[j] = dl * (grid[j] + delta) # --------------------------------------------------------------------------- def dequant_iq1_s_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 66 chunk = chunk.view(cr, nblk, 66) d = _f16(chunk[..., 0:2]) # [cr, nblk] qs = chunk[..., 2:34].to(torch.int64) # [cr, nblk, 32] qh = chunk[..., 34:66].view(cr, nblk, 16, 2) qh = _le_u16(qh) # [cr, nblk, 16] ib = torch.arange(8, device=raw.device) il = torch.arange(4, device=raw.device) qs8 = qs.view(cr, nblk, 8, 4) qh8 = qh.view(cr, nblk, 8, 2) # 2 u16 per group # per group: uses qh8[..., 0] for scale/delta, and qh8[..., 0] for index too dl = d.unsqueeze(-1) * (2 * ((qh8[..., 0] >> 12) & 7) + 1).to(torch.float32) # [cr,nblk,8] delta = torch.where((qh8[..., 0] & 0x8000) != 0, -IQ1S_DELTA, IQ1S_DELTA) # [cr,nblk,8] idx = qs8 + (((qh8[..., 0].unsqueeze(-1) >> (3 * il.view(1, 1, 1, 4))) & 7) << 8) gr = iq_table("iq1s_grid").to(raw.device) # [2048, 8] int8 signed vals = gr[idx].to(torch.float32) # [cr,nblk,8,4,8] out = dl.unsqueeze(-1).unsqueeze(-1) * (vals + delta.unsqueeze(-1).unsqueeze(-1)) return out.reshape(cr, nblk * QK_K)[:, :cols] # --------------------------------------------------------------------------- # IQ1_M -- block 56 B: qs(32) + qh(16) + scales(8 u16) # No fp16 d! scale is f16 packed into 4 bytes of scales. # scale.f16 = (sc[0]>>12) | ((sc[1]>>8)&0xF0) | ((sc[2]>>4)&0xF00) | (sc[3]&0xF000) # per group ib: dl1 = d*(2*((sc[ib/2] >> 6*(ib%2)+0) & 0x7) + 1) # dl2 = d*(2*((sc[ib/2] >> 6*(ib%2)+3) & 0x7) + 1) # idx[0] = qs[0] | ((qh[0]<<8)&0x700), idx[1] = qs[1] | ((qh[0]<<4)&0x700) # idx[2] = qs[2] | ((qh[1]<<8)&0x700), idx[3] = qs[3] | ((qh[1]<<4)&0x700) # delta[l] = qh[l/2] & (0x08 << 4*(l%2)) ? -IQ1M_DELTA : +IQ1M_DELTA # y = dl1 * (grid[idx[l]] + delta[l]) (l 0,1) dl2 for l 2,3 # --------------------------------------------------------------------------- def dequant_iq1_m_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 56 chunk = chunk.view(cr, nblk, 56) qs = chunk[..., 0:32].to(torch.int64) # [cr,nblk,32] qh = chunk[..., 32:48].to(torch.int64) # [cr,nblk,16] sc = chunk[..., 48:56].view(cr, nblk, 4, 2) sc = _le_u16(sc) # [cr,nblk,4] # reassemble f16: nibbles from sc[0..3] u = (sc[..., 0] >> 12) | ((sc[..., 1] >> 8) & 0x00F0) \ | ((sc[..., 2] >> 4) & 0x0F00) | (sc[..., 3] & 0xF000) # [cr,nblk] d = u.to(torch.int16).view(torch.float16).to(torch.float32) ib = torch.arange(8, device=raw.device) il = torch.arange(4, device=raw.device) sc8 = sc.repeat_interleave(2, dim=2) # [cr,nblk,8] # dl1 = d*(2*((sc8[ib] >> 6*(ib%2)) & 7)+1) dl1 = d.unsqueeze(-1) * (2 * ((sc8 >> (6 * (ib % 2).view(1, 1, 8))) & 0x7) + 1).to(torch.float32) dl2 = d.unsqueeze(-1) * (2 * ((sc8 >> (6 * (ib % 2).view(1, 1, 8) + 3)) & 0x7) + 1).to(torch.float32) qh2 = qh.view(cr, nblk, 8, 2) # [cr,nblk,8,2] qs8 = qs.view(cr, nblk, 8, 4) idx0 = qs8[..., 0] | ((qh2[..., 0] << 8) & 0x700) idx1 = qs8[..., 1] | ((qh2[..., 0] << 4) & 0x700) idx2 = qs8[..., 2] | ((qh2[..., 1] << 8) & 0x700) idx3 = qs8[..., 3] | ((qh2[..., 1] << 4) & 0x700) gr = iq_table("iq1s_grid").to(torch.int64) # [2048, 8] g0 = gr[idx0].to(torch.float32) g1 = gr[idx1].to(torch.float32) g2 = gr[idx2].to(torch.float32) g3 = gr[idx3].to(torch.float32) delta0 = torch.where((qh2[..., 0] & 0x08) != 0, -IQ1M_DELTA, IQ1M_DELTA) delta1 = torch.where((qh2[..., 0] & 0x80) != 0, -IQ1M_DELTA, IQ1M_DELTA) delta2 = torch.where((qh2[..., 1] & 0x08) != 0, -IQ1M_DELTA, IQ1M_DELTA) delta3 = torch.where((qh2[..., 1] & 0x80) != 0, -IQ1M_DELTA, IQ1M_DELTA) out = torch.cat([ (dl1.unsqueeze(-1) * (g0 + delta0.unsqueeze(-1))).unsqueeze(-2), (dl1.unsqueeze(-1) * (g1 + delta1.unsqueeze(-1))).unsqueeze(-2), (dl2.unsqueeze(-1) * (g2 + delta2.unsqueeze(-1))).unsqueeze(-2), (dl2.unsqueeze(-1) * (g3 + delta3.unsqueeze(-1))).unsqueeze(-2), ], dim=-2) # [cr,nblk,8,4,8] return out.reshape(cr, nblk * QK_K)[:, :cols] # --------------------------------------------------------------------------- # IQ4_NL -- block 18 B: d(2) + qs(16) # y[j] = d * kvalues_iq4nl[qs[j] & 0xF], y[j+16] = d * kvalues_iq4nl[qs[j] >> 4] # --------------------------------------------------------------------------- def dequant_iq4_nl_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 18 chunk = chunk.view(cr, nblk, 18) d = _f16(chunk[..., 0:2]) # [cr,nblk] qs = chunk[..., 2:18] # [cr,nblk,16] kv = iq_table("kvalues_iq4nl").to(raw.device) # [16] lo = kv[qs.to(torch.int64) & 0xF] # [cr,nblk,16] hi = kv[(qs >> 4).to(torch.int64)] vals = torch.stack([lo, hi], dim=-1).reshape(cr, nblk, 32) return (d.unsqueeze(-1) * vals).reshape(cr, nblk * QK4_NL)[:, :cols] # --------------------------------------------------------------------------- # IQ4_XS -- block 136 B: d(2) + scales_h(2) + scales_l(4) + qs(128) # per 32-group ib (8): ls = ((scales_l[ib/2] >> 4*(ib%2)) & 0xF) # | (((scales_h >> 2*ib) & 3) << 4) # dl = d * (ls - 32) # y[j] = dl * kvalues_iq4nl[qs[j] & 0xF], y[j+16] = dl * kvalues_iq4nl[qs[j] >> 4] # --------------------------------------------------------------------------- def dequant_iq4_xs_packed( raw: torch.Tensor, start: int, end: int, cols: int ) -> torch.Tensor: chunk = raw[start:end].to(torch.uint8) cr = chunk.shape[0] nblk = chunk.shape[1] // 136 chunk = chunk.view(cr, nblk, 136) d = _f16(chunk[..., 0:2]) scales_h = _le_u16(chunk[..., 2:4]) # [cr, nblk] uint16 scales_l = chunk[..., 4:8].to(torch.int64) # [cr, nblk, 4] qs = chunk[..., 8:136] # [cr, nblk, 128] ib = torch.arange(8, device=raw.device) ls = (scales_l.gather(2, (ib // 2).view(1, 1, 8).expand(cr, nblk, 8)) >> (4 * (ib % 2)).view(1, 1, 8)) & 0xF ls = ls | (((scales_h.unsqueeze(-1) >> (2 * ib).view(1, 1, 8)) & 3) << 4) # [cr,nblk,8] dl = d.unsqueeze(-1) * (ls.to(torch.float32) - 32) kv = iq_table("kvalues_iq4nl").to(raw.device) qs32 = qs.view(cr, nblk, 8, 16) lo = kv[qs32.to(torch.int64) & 0xF] # [cr,nblk,8,16] hi = kv[(qs32 >> 4).to(torch.int64)] vals = torch.stack([lo, hi], dim=-1).reshape(cr, nblk, 8, 32) return (dl.unsqueeze(-1) * vals).reshape(cr, nblk * QK_K)[:, :cols] # =========================================================================== # Quantizers (self-consistent with the dequantizers above) # =========================================================================== _SCALE_CANDIDATES = 17 def _find_nearest( x: torch.Tensor, # [M, W] fp32 (magnitudes if not signed) grid: torch.Tensor, # [G, W] int8 grid values (W = 4 or 8) signed: bool = False, # grid rows carry signs (IQ1-style: {-1,0,1}) ) -> tuple[torch.Tensor, torch.Tensor]: """Best (grid row, scale) for each W-vector by weighted nearest search. Works for any grid row width W (4 for IQ3-style grids, 8 for IQ1/IQ2). Weighted (by |x|) squared distance to every grid row over a log-spaced candidate-scale grid; chunked over M. Returns (idx [M] int64, scale [M] fp32). scale is the per-element multiplier: dequant gives y = scale * grid values (scaled by the block d later). """ M = x.shape[0] G = grid.shape[0] gf = grid.to(torch.float32) # [G, W] xabs = x.abs() gmax = gf.abs().amax(dim=-1).max().clamp(min=1e-8) amax = xabs.amax(dim=-1).clamp(min=1e-8) # [M] cand = (amax / gmax).unsqueeze(-1) * torch.logspace( -2.0, 1.0, _SCALE_CANDIDATES, device=x.device) # [M, C] ~0.01..10 x amax w = (xabs + 1e-6) # importance weight wnorm = w / w.sum(-1, keepdim=True).clamp(min=1e-8) # [M, W] idx = torch.empty(M, dtype=torch.int64, device=x.device) scale = torch.empty(M, dtype=torch.float32, device=x.device) chunk = max(1, min(M, (1 << 21) // max(1, G))) for s0 in range(0, M, chunk): e = min(s0 + chunk, M) xs = x[s0:e] # [K, W] wn = wnorm[s0:e] # [K, W] cs = cand[s0:e] # [K, C] # d2[K, C, G] = sum_i wn_i * (x_i - s*g_i)^2 xg = (wn * xs) @ gf.T # [K, G] g2w = (wn.unsqueeze(1) * gf.unsqueeze(0) ** 2).sum(-1) # [K, G] x2w = (wn * xs * xs).sum(-1) # [K] d2 = x2w.unsqueeze(-1).unsqueeze(-1) \ - 2 * cs.unsqueeze(-1) * xg.unsqueeze(1) \ + (cs.unsqueeze(-1) ** 2) * g2w.unsqueeze(1) # [K, C, G] d2 = d2.clamp(min=0) flat = d2.view(e - s0, -1) best = flat.argmin(dim=-1) # [K] scale_idx = best // G grid_idx = best % G scale[s0:e] = cs.view(e - s0, -1).gather(1, scale_idx.unsqueeze(-1)).squeeze(-1) idx[s0:e] = grid_idx # Refine scale by weighted least squares on the chosen row. gv = grid[idx].to(torch.float32) num = (w * x * gv).sum(-1) den = (w * gv * gv).sum(-1).clamp(min=1e-8) scale = torch.where(den > 0, num / den, scale) return idx, scale def _pick_signs8(x8: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Sign bits with even parity (mirror llama.cpp ksigns trick). Returns (sign_bits [.., 8] int64 0/1, xabs [.., 8] fp32). For odd-popcount vectors the smallest-magnitude element is flipped. """ sgn = (x8 < 0).to(torch.int64) odd = (sgn.sum(-1) % 2) == 1 if odd.any(): mn = x8.abs().argmin(-1) # [..] o = odd.unsqueeze(-1) m = mn.unsqueeze(-1) pos = torch.arange(x8.shape[-1], device=x8.device).view(1, -1) == m sgn = torch.where(o & pos, 1 - sgn, sgn) return sgn, x8.abs() def _ksigns_encode(sgn: torch.Tensor) -> torch.Tensor: """sgn [.., 8] int64 bits -> ksigns table index [..] int64.""" sb = (sgn * (1 << torch.arange(8, device=sgn.device))).sum(-1) # [..] 0..255 ks = iq_table("ksigns_iq2xs").to(sgn.device) # [128] uint8 lut = torch.full((256,), -1, dtype=torch.int64, device=sgn.device) lut[ks.to(torch.int64)] = torch.arange(128, device=sgn.device) return lut[sb.to(torch.int64)] def _pack_u16(v: torch.Tensor) -> torch.Tensor: """int64 [.., N] -> [.., N, 2] LE bytes.""" return torch.stack([v & 0xFF, (v >> 8) & 0xFF], dim=-1).to(torch.uint8) def _pack_f16(d: torch.Tensor) -> torch.Tensor: """fp32 [N] -> [N, 2] LE bytes.""" return _pack_u16(d.to(torch.float16).view(torch.int16).to(torch.int64) & 0xFFFF) def quantize_iq2_xxs(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ2_XXS bytes [out, 66*nblk] (d + 32 u16). Block 256 = 8 super-groups x 32; each super-group = 4 8-groups. Per super-group: 4 u16 = [idx0|idx1, idx2|idx3, aux_lo, aux_hi]; aux32 = (ls << 28) | (sgn_idx[0] << 0) | (sgn_idx[1] << 7) | (sgn_idx[2] << 14) | (sgn_idx[3] << 21). """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) sgn, xa = _pick_signs8(x32.view(-1, 8)) sgn = sgn.view(out_f, nblk, 8, 4, 8) # [out, nblk, sg, 4x8, 8] xa = xa.view(out_f, nblk, 8, 4, 8) grid = iq_table("iq2xxs_grid").to(W.device) # [256, 8] idx, scale = _find_nearest(xa.reshape(-1, 8), grid) # [out*nblk*32] idx = idx.view(out_f, nblk, 8, 4) scale = scale.view(out_f, nblk, 8, 4) sc_g = scale.amax(dim=-1).clamp(min=1e-8) # [out, nblk, 8] per super-group # llama.cpp: scale is "unit" scale (for q=1); dequant gives y = (scale/8)*grid. # Our per-8-group scale is x/grid -> unit scale = 8*scale. d = 8.0 * sc_g.amax(dim=-1) / 31.0 # [out, nblk] per block ls = torch.clamp(torch.round(0.5 * (8.0 * scale / d.unsqueeze(-1).unsqueeze(-1) - 1)), 0, 15).to(torch.int64) sgn_idx = _ksigns_encode(sgn) # [out, nblk, 8, 4] 0..127 aux32 = (ls << 28) | (sgn_idx << (7 * torch.arange(4, device=W.device)).view(1, 1, 1, 4)) aux32 = aux32[..., 0] | aux32[..., 1] | aux32[..., 2] | aux32[..., 3] # [out, nblk, 8] qb = idx.to(torch.int64) & 0xFF q2 = (qb[..., 0] | (qb[..., 1] << 8)) & 0xFFFF # [out, nblk, 8] q2h = (qb[..., 2] | (qb[..., 3] << 8)) & 0xFFFF aux_lo = aux32 & 0xFFFF aux_hi = (aux32 >> 16) & 0xFFFF u16 = torch.stack([q2, q2h, aux_lo, aux_hi], dim=-1) # [out, nblk, 8, 4] return _cat_blocks([_pack_f16(d), _pack_u16(u16).reshape(out_f, nblk, 64)], out_f, nblk, 66) def quantize_iq2_xs(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ2_XS bytes [out, 74*nblk] (d + 32 u16 + scales[8]). Each 32-group: 4 u16, each = 9-bit grid index | (7-bit sign index << 9). scales[ib] = 2 4-bit codes (lo = sub-group<2, hi = sub-group>=2) of the unit scale / 8. """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) sgn, xa = _pick_signs8(x32.view(-1, 8)) sgn = sgn.view(out_f, nblk, 8, 4, 8) xa = xa.view(out_f, nblk, 8, 4, 8) grid = iq_table("iq2xs_grid").to(W.device) # [512, 8] idx, scale = _find_nearest(xa.reshape(-1, 8), grid) idx = idx.view(out_f, nblk, 8, 4) scale = scale.view(out_f, nblk, 8, 4) unit = 8.0 * scale # [out, nblk, 8, 4] pair = unit.view(out_f, nblk, 8, 2, 2).amax(-1) # [out, nblk, 8, 2] d = pair.amax(dim=(-1, -2)) / 31.0 # [out, nblk] per block d4 = d.unsqueeze(-1).unsqueeze(-1) # [out, nblk, 1, 1] code = torch.clamp(torch.round(0.5 * (pair / d4 - 1)), 0, 15).to(torch.int64) scales = code[..., 0] | (code[..., 1] << 4) # [out, nblk, 8] sgn_idx = _ksigns_encode(sgn) # [out, nblk, 8, 4] 0..127 q2 = (idx.to(torch.int64) & 0x1FF) | (sgn_idx << 9) # [out, nblk, 8, 4] return _cat_blocks([_pack_f16(d), _pack_u16(q2.reshape(out_f, nblk, 32)).reshape(out_f, nblk, 64), scales.to(torch.uint8)], out_f, nblk, 74) def quantize_iq2_s(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ2_S bytes [out, 82*nblk]. qs[4*ib+il] = low byte of grid index, qh[ib] = 2 high bits per sub-group, signs byte at qs[32+4*ib+il], scales[8] = 4-bit pair. """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) sgn, xa = _pick_signs8(x32.view(-1, 8)) sgn = sgn.view(out_f, nblk, 8, 4, 8) xa = xa.view(out_f, nblk, 8, 4, 8) grid = iq_table("iq2s_grid").to(W.device) # [1024, 8] idx, scale = _find_nearest(xa.reshape(-1, 8), grid) idx = idx.view(out_f, nblk, 8, 4) scale = scale.view(out_f, nblk, 8, 4) unit = 8.0 * scale unit_p = unit.view(out_f, nblk, 8, 2, 2).amax(-1) # [out, nblk, 8, 2] d = unit_p.amax(dim=(-1, -2)) / 31.0 d4 = d.unsqueeze(-1).unsqueeze(-1) code = torch.clamp(torch.round(0.5 * (unit_p / d4 - 1)), 0, 15).to(torch.int64) scales = code[..., 0] | (code[..., 1] << 4) # signs: plain 8-bit byte (NOT ksigns!) at qs[32+4*ib+il] signs_b = (sgn * (1 << torch.arange(8, device=W.device))).sum(-1).to(torch.uint8) qs_low = (idx.to(torch.int64) & 0xFF).to(torch.uint8) qs_hi = signs_b qh = (((idx.to(torch.int64) >> 8) << (2 * torch.arange(4, device=W.device)).view(1, 1, 1, 4)).sum(-1) & 0xFF).to(torch.uint8) return _cat_blocks([_pack_f16(d), qs_low.reshape(out_f, nblk, 32), qs_hi.reshape(out_f, nblk, 32), qh, scales.to(torch.uint8)], out_f, nblk, 82) def quantize_iq3_xxs(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ3_XXS bytes [out, 98*nblk]. qs[0:64] = 64 grid-index bytes (8 groups x 8), qs[64:96] = scales+signs (8 groups x 4 u16 = aux32 per group). grid = iq3xxs_grid [256, 4]. """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) sgn, xa = _pick_signs8(x32.view(-1, 8)) sgn = sgn.view(out_f, nblk, 8, 4, 8) xa = xa.view(out_f, nblk, 8, 4, 8) grid = iq_table("iq3xxs_grid").to(W.device) # [256, 4] # Each 8-group = 2 interleaved 4-halves (lo at idx[2*il], hi at idx[2*il+1]). xa_4 = xa.view(out_f, nblk, 8, 4, 2, 4) # [.,.,8,4,2 halves,4] xa_lo = xa_4[..., 0, :].reshape(-1, 4) xa_hi = xa_4[..., 1, :].reshape(-1, 4) idx_lo, s_lo = _find_nearest(xa_lo, grid) idx_hi, s_hi = _find_nearest(xa_hi, grid) scale = torch.maximum(s_lo, s_hi).view(out_f, nblk, 8, 4) idx = (idx_lo.view(out_f, nblk, 8, 4).to(torch.int64) | (idx_hi.view(out_f, nblk, 8, 4).to(torch.int64) << 8)) # 2 x 8-bit per sub-group # per-32-group scale (4 sub-groups): max over the 4 sg = scale.amax(dim=-1).clamp(min=1e-8) # [out, nblk, 8] d = sg.amax(dim=-1) / 15.0 # [out, nblk] d4 = d.unsqueeze(-1).unsqueeze(-1) # db = d*(0.5+(aux>>28))*0.5 -> aux = 4*scale/d - 1... solve: aux = round(scale/d/0.5 - 0.5) ls = torch.clamp(torch.round(2.0 * scale / d4 - 0.5), 0, 15).to(torch.int64) sgn_idx = _ksigns_encode(sgn) # [out, nblk, 8, 4] 0..127 aux32 = (ls << 28) | (sgn_idx << (7 * torch.arange(4, device=W.device)).view(1, 1, 1, 4)) aux32 = aux32[..., 0] | aux32[..., 1] | aux32[..., 2] | aux32[..., 3] # [out, nblk, 8] # Per sub-group: 2 grid codes of 4 -> 2 bytes, interleaved (lo, hi). # Layout of qs[0:64] per 8-group: [lo_il0, hi_il0, lo_il1, hi_il1, ...]. lo_b = (idx & 0xFF).view(out_f, nblk, 8, 4) # [.,.,8,4] hi_b = ((idx >> 8) & 0xFF).view(out_f, nblk, 8, 4) qs_idx = torch.stack([lo_b, hi_b], dim=-1).reshape(out_f, nblk, 64).to(torch.uint8) aux_pairs = torch.stack([aux32 & 0xFFFF, aux32 >> 16], dim=-1) # [out, nblk, 8, 2] return _cat_blocks([_pack_f16(d), qs_idx, _pack_u16(aux_pairs.reshape(out_f, nblk, 16)).reshape(out_f, nblk, 32)], out_f, nblk, 98) def quantize_iq3_s(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ3_S bytes [out, 110*nblk]. qs[0:64] grid-index bytes (lo/hi interleaved), qh[8] high bits, signs[32], scales[4] (nibble per group of a pair). """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) sgn, xa = _pick_signs8(x32.view(-1, 8)) sgn = sgn.view(out_f, nblk, 8, 4, 8) xa = xa.view(out_f, nblk, 8, 4, 8) grid = iq_table("iq3s_grid").to(W.device) # [512, 4] xa_4 = xa.view(out_f, nblk, 8, 4, 2, 4) idx_lo, s_lo = _find_nearest(xa_4[..., 0, :].reshape(-1, 4), grid) idx_hi, s_hi = _find_nearest(xa_4[..., 1, :].reshape(-1, 4), grid) idx_lo = idx_lo.view(out_f, nblk, 8, 4).to(torch.int64) idx_hi = idx_hi.view(out_f, nblk, 8, 4).to(torch.int64) scale = torch.stack([s_lo, s_hi], -1).amax(-1).view(out_f, nblk, 8, 4) # db per group (8 groups; scales[ib/2] nibble by ib%2, code = (unit/d-1)/2) scale_g = scale.amax(dim=-1) # [out, nblk, 8] sg_pairs = scale_g.view(out_f, nblk, 4, 2) # [out, nblk, 4, 2] (even, odd) d = scale_g.amax(dim=-1) / 15.0 # [out, nblk] d4 = d.unsqueeze(-1).unsqueeze(-1) # [out, nblk, 1, 1] code_pair = torch.clamp(torch.round((sg_pairs / d4 - 1.0) * 0.5), 0, 15).to(torch.int64) scales = (code_pair[..., 0] | (code_pair[..., 1] << 4)).to(torch.uint8) # [out, nblk, 4] # qs bytes: interleave lo/hi codes per sub-group; qh carries high bits. lo_code = (idx_lo & 0xFF).to(torch.uint8) hi_code = (idx_hi & 0xFF).to(torch.uint8) qs_idx = torch.stack([lo_code, hi_code], dim=-1).reshape(out_f, nblk, 64) il2 = torch.arange(4, device=W.device) qh_lo = (idx_lo >> 8) & 1 qh_hi = (idx_hi >> 8) & 1 qh = ((qh_lo << (2 * il2).view(1, 1, 1, 4)) | (qh_hi << (2 * il2 + 1).view(1, 1, 1, 4))).sum(-1) & 0xFF # [out, nblk, 8] qh = qh.to(torch.uint8) signs_b = (sgn * (1 << torch.arange(8, device=W.device))).sum(-1).to(torch.uint8) # [out,nblk,8,4] return _cat_blocks([_pack_f16(d), qs_idx, qh, signs_b.reshape(out_f, nblk, 32), scales], out_f, nblk, 110) def quantize_iq1_s(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ1_S bytes [out, 66*nblk]. d(f16) + qs[32] + qh[16]u16. grid = iq1s_grid [2048, 8] (signed, {-1,0,1}). Per 32-group ib: dl = d*(2*((qh[ib]>>12)&7)+1); delta = +-IQ1S_DELTA (sign from qh[ib]&0x8000); idx = qs[4*ib+il] | (((qh[ib]>>3*il)&7)<<8); y[j] = dl * (grid[idx][j] + delta). """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) grid = iq_table("iq1s_grid").to(W.device) # [2048, 8] int8 signed x8 = x32.view(out_f * nblk * 32, 8) # Search signed grid: dequant y = dl*(g+delta); for a candidate delta # the effective target is x/dl - delta. Do a 2-pass: first find the best # (idx, scale=dl) on the signed grid, then decide delta by the residual # sign and re-pick if needed. Approx is fine at 1.5 bpw. idx, scale = _find_nearest(x8, grid, signed=True) idx = idx.view(out_f, nblk, 8, 4) scale = scale.view(out_f, nblk, 8, 4) # dl per 8-group # group multiplier: dl = d*(2*m+1), m 0..7 g_max = scale.amax(dim=-1).clamp(min=1e-8) # [out, nblk, 8] m = torch.clamp(torch.round((g_max / g_max.amax(dim=-1).unsqueeze(-1) * 7 - 1) * 0.5), 0, 7).to(torch.int64) d = (g_max.amax(dim=-1) / (2 * 7 + 1)).clamp(min=1e-8) # [out, nblk] dl = d.unsqueeze(-1) * (2 * m + 1) # [out, nblk, 8] per super-group id_dl = 1.0 / dl.unsqueeze(-1).unsqueeze(-1) gv = grid[idx].to(torch.float32) # [out, nblk, 8, 4, 8] rec = (dl.unsqueeze(-1).unsqueeze(-1) * gv) # [out, nblk, 8, 4, 8] res = x32.view(out_f, nblk, 8, 4, 8) - rec # delta: one sign per 32-group (aggregate residual over all 32 values) delta = torch.where(res.mean(dim=(-1, -2)) < 0, -IQ1S_DELTA, IQ1S_DELTA) # [out, nblk, 8] # adjust idx for the delta (approximate: re-search on x/dl - delta) tgt = x32.view(out_f, nblk, 8, 4, 8) * id_dl - delta.unsqueeze(-1).unsqueeze(-1) idx2, _ = _find_nearest(tgt.reshape(-1, 8), grid, signed=True) idx = idx2.view(out_f, nblk, 8, 4) # encode: qh[ib] = (sign delta << 15) | (m << 12) | (3 high bits of idx per il) qh = (delta < 0).to(torch.int64) << 15 # [out, nblk, 8] qh = qh | (m << 12) # [out, nblk, 8] qs = (idx.to(torch.int64) & 0xFF) # low 8 bits per (ib, il) hi3 = (idx.to(torch.int64) >> 8) & 7 qh = qh | (hi3 << (3 * torch.arange(4, device=W.device)).view(1, 1, 1, 4)).sum(-1) # qh storage: 16 u16 per block (2 per group); decoder reads first of each pair. qh16 = torch.zeros(out_f, nblk, 16, dtype=torch.int64, device=W.device) qh16[..., 0::2] = qh return _cat_blocks([_pack_f16(d), qs.to(torch.uint8).reshape(out_f, nblk, 32), _pack_u16(qh16).reshape(out_f, nblk, 32)], out_f, nblk, 66) def quantize_iq4_xs(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ4_XS bytes [out, 136*nblk]. d(f16) + scales_h[2] + scales_l[4] + qs[128]. """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W bl = Wp.view(out_f, nblk, 8, 32) amax = bl.abs().amax(dim=-1).clamp(min=1e-8) # [out, nblk, 8] d = amax.amax(dim=-1) / 31.0 # [out, nblk] kv = iq_table("kvalues_iq4nl").to(W.device) # [16] id_d = 1.0 / d.unsqueeze(-1).unsqueeze(-1) al = bl * id_d # [out, nblk, 8, 32] codes = (al.unsqueeze(-1) - kv.view(1, 1, 1, 1, 16)).abs().argmin(-1) # [out,nblk,8,32] # scales: 6-bit per group = (scales_l nibble) | (scales_h 2-bit << 4), -32 offset ls = torch.clamp(torch.round(amax / d.unsqueeze(-1)), 0, 63).to(torch.int64) # [out,nblk,8] ls = ls - 32 scales_l = (ls & 0xF).view(out_f, nblk, 4, 2) scales_l = (scales_l[..., 0] | (scales_l[..., 1] << 4)).view(out_f, nblk, 4) hi2 = (ls >> 4) & 3 # [out,nblk,8] 2-bit each h0 = hi2[..., 0] | (hi2[..., 1] << 2) | (hi2[..., 2] << 4) | (hi2[..., 3] << 6) h1 = hi2[..., 4] | (hi2[..., 5] << 2) | (hi2[..., 6] << 4) | (hi2[..., 7] << 6) scales_h = torch.stack([h0, h1], dim=-1).view(out_f, nblk, 2).to(torch.int64) packed = (codes[..., 0::2].to(torch.int64) | (codes[..., 1::2].to(torch.int64) << 4)).to(torch.uint8) # [out,nblk,8,16] return _cat_blocks([_pack_f16(d), scales_h.to(torch.uint8), scales_l.to(torch.uint8), packed.reshape(out_f, nblk, 128)], out_f, nblk, 136) def quantize_iq4_nl(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ4_NL bytes [out, 18*nblk32]. Block = 32 values: d(f16) + qs[16]. y[j] = d*kvalues[qs&0xF], y[j+16] = d*kvalues[qs>>4]. kvalues = [-127,-104,-83,-65,-49,-35,-22, -10,1,13,25,38,53,69,89,113]. """ out_f, in_f = W.shape nblk = (in_f + QK4_NL - 1) // QK4_NL Wp = torch.zeros(out_f, nblk * QK4_NL, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W bl = Wp.view(out_f, nblk, QK4_NL) # [out, nblk, 32] kv = iq_table("kvalues_iq4nl").to(W.device) # [16] fp32 amax = bl.abs().amax(dim=-1).clamp(min=1e-8) # [out, nblk] d = amax / kv.abs().max().clamp(min=1e-8) # [out, nblk] id_d = 1.0 / d.unsqueeze(-1) al = bl * id_d # [out, nblk, 32] codes = (al.unsqueeze(-1) - kv.view(1, 1, 1, 16)).abs().argmin(-1) # [out,nblk,32] packed = (codes[..., 0::2].to(torch.int64) | (codes[..., 1::2].to(torch.int64) << 4)).to(torch.uint8) # [out,nblk,16] return _cat_blocks([_pack_f16(d), packed], out_f, nblk, 18) def quantize_iq1_m(W: torch.Tensor) -> torch.Tensor: """[out, in] fp32 -> IQ1_M bytes [out, 56*nblk]. qs[32] + qh[16] + scales[8 u16]. No fp16 d: scale is reassembled from the top nibbles of the 8 scale u16 (see dequant_iq1_m_packed). """ out_f, in_f = W.shape nblk = (in_f + QK_K - 1) // QK_K Wp = torch.zeros(out_f, nblk * QK_K, dtype=torch.float32, device=W.device) Wp[:, :in_f] = W x32 = Wp.view(out_f, nblk, QK_K) grid = iq_table("iq1s_grid").to(W.device) # [2048, 8] signed {-1,0,1} x8 = x32.view(out_f, nblk, 8, 4, 8) # per 8-group: find best (idx, scale) on the signed grid idx, scale = _find_nearest(x8.reshape(-1, 8), grid, signed=True) idx = idx.view(out_f, nblk, 8, 4) scale = scale.view(out_f, nblk, 8, 4) # group dl (per 4 sub-groups of a 32-group... actually per 32-group pair): # dl1 = d*(2*sc3+1), dl2 = d*(2*sc3+3) -> sc3 = 3-bit scale per 16-group. # Approx: per 32-group, dl = max scale over its 4 sub-groups. g_max = scale.amax(dim=-1).clamp(min=1e-8) # [out, nblk, 8] per 32-group # per 32-group 2 codes: dl1 = max(scale of sub-groups 0,1), dl2 = max(2,3) pair_scale = scale.view(out_f, nblk, 8, 2, 2).amax(dim=-1) # [out, nblk, 8, 2] dmax = pair_scale.amax(dim=-1).amax(dim=-1).clamp(min=1e-8) # [out, nblk] code2 = torch.clamp(torch.round((pair_scale / dmax.unsqueeze(-1).unsqueeze(-1) * 7 - 1) * 0.5), 0, 7).to(torch.int64) # [out,nblk,8,2] dl = dmax.unsqueeze(-1).unsqueeze(-1) * (2 * code2 + 1) # [out, nblk, 8, 2] # sc[p] u16 (p=0..3): 2 codes per 32-group, groups 2p (even, bits 0-5) and # 2p+1 (odd, bits 6-11): [dl1, dl2] of each group. c = code2.view(out_f, nblk, 4, 2, 2) # [out,nblk,4 u16,2 groups,2 dl] gA1, gA2 = c[..., 0, 0], c[..., 0, 1] gB1, gB2 = c[..., 1, 0], c[..., 1, 1] sc16 = gA1 | (gA2 << 3) | (gB1 << 6) | (gB2 << 9) # [out, nblk, 4] # fp16 d in the top nibbles across the 4 u16: # d16 = (sc0>>12) | ((sc1>>8)&0xF0) | ((sc2>>4)&0xF00) | (sc3&0xF000) d16 = dmax.to(torch.float16).view(torch.int16).to(torch.int64) & 0xFFFF # [out,nblk] d16e = d16.unsqueeze(-1) # [out, nblk, 1] sc16 = sc16 | (((d16e >> 12) & 0xF) << 12) sc16 = sc16 | ((((d16e >> 8) & 0xF) << 8) & 0x0F00) sc16 = sc16 | ((((d16e >> 4) & 0xF) << 4) & 0x00F0) sc16 = sc16 | (d16e & 0x000F) # delta per group (from qh 0x08 / 0x80 bits) + idx high bits in qh delta = torch.where(g_max / dmax.unsqueeze(-1) < 0.5, -IQ1M_DELTA, IQ1M_DELTA) # placeholder delta = torch.full_like(g_max, IQ1M_DELTA) # sign from residual below # delta: reconstruct with dl per (group, pair): dl1 for il 0..1, dl2 for 2..3 dl_pair = dl.view(out_f, nblk, 8, 2, 1) # [out,nblk,8,2 dl,1] dl_full = dl_pair.expand(-1, -1, -1, -1, 2).reshape(out_f, nblk, 8, 4) # per sub-group gv = grid[idx].to(torch.float32) # [out,nblk,8,4,8] res = x8 - dl_full.unsqueeze(-1) * gv delta = torch.where(res.mean(dim=(-1, -2)) < 0, -IQ1M_DELTA, IQ1M_DELTA) # [out,nblk,8] # qh: 16 bytes per block (2 per 32-group). Per 32-group: # qh[0] bits: idx0 hi3<<8, idx1 hi3<<4, delta0<<3, delta1<<7 (via 0x08/0x80) # qh[1] bits: idx2 hi3<<8, idx3 hi3<<4, delta2<<3, delta3<<7 hi3 = (idx.to(torch.int64) >> 8) & 7 # [out,nblk,8,4] qh2 = torch.zeros(out_f, nblk, 8, 2, dtype=torch.int64, device=W.device) qh2[..., 0] |= (hi3[..., 0] << 8) | (hi3[..., 1] << 4) qh2[..., 1] |= (hi3[..., 2] << 8) | (hi3[..., 3] << 4) dlt = (delta < 0).to(torch.int64) # [out,nblk,8] qh2[..., 0] |= (dlt << 3) & 0x08 qh2[..., 0] |= (dlt << 7) & 0x80 qh2[..., 1] |= (dlt << 3) & 0x08 qh2[..., 1] |= (dlt << 7) & 0x80 qs = (idx.to(torch.int64) & 0xFF) return _cat_blocks([qs.to(torch.uint8).reshape(out_f, nblk, 32), qh2.to(torch.uint8).reshape(out_f, nblk, 16), _pack_u16(sc16).reshape(out_f, nblk, 8)], out_f, nblk, 56)