|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""ExecuTorch-internal dp4a-planar INT5 tensor for the CUDA W5A8 dp4a decode kernel. |
| 8 | +
|
| 9 | +``CudaDp4aPlanarInt5Tensor`` is an ExecuTorch-internal tensor subclass that stores |
| 10 | +a genuine 5-bit packed weight (0.625 B/elem) in a dp4a-friendly *planar* layout: |
| 11 | +the 5 bits are split into a low nibble plane (``ql``) and a 1-bit high plane |
| 12 | +(``qh``) so the decode kernel can run the W5A8 dp4a matvec directly. Used for |
| 13 | +GGUF Q5_K weights. Unlike the symmetric INT6 path (``CudaDp4aPlanarInt6Tensor``), |
| 14 | +Q5_K is *asymmetric* (it has both ``d`` and ``dmin``), so this format carries a |
| 15 | +zero tensor exactly like the INT4 path. |
| 16 | +
|
| 17 | +Build one with :meth:`from_exportable_gguf` (from a native Q5_K |
| 18 | +``ExportableGGUFTensor`` — it reuses the shared Q5_K block decode in |
| 19 | +``extension/llm/export/gguf.py`` and then feeds the internal ql/qh packer |
| 20 | +:meth:`_from_intx_int8`). This class owns only the 5-bit pack and the metadata |
| 21 | +re-encoding, never the Q5_K block decode. |
| 22 | +
|
| 23 | +The stored value is the raw unsigned ``u = q`` in ``[0, 31]`` (no offset — the |
| 24 | +per-group zero point is subtracted in the decode kernel, mirroring INT4). The 5 |
| 25 | +bits are split into two planes that mirror the INT4 nibble layout so the kernel |
| 26 | +can reuse the INT4 even/odd extraction verbatim: |
| 27 | +
|
| 28 | + ql : (N, K/2) uint8 — low-nibble plane, nibble-packed even/odd |
| 29 | + (``ql[:, j] = lo[:, 2j] | (lo[:, 2j+1] << 4)``, ``lo = u & 0xF``). |
| 30 | + qh : (N, K/8) uint8 — high-1-bit plane, 8 values/byte, arranged per |
| 31 | + 32-weight chunk as 4 bytes (one per dp4a word); each byte holds |
| 32 | + the four 1-bit highs of that word's even weights in the low |
| 33 | + nibble and its odd weights in the high nibble |
| 34 | + (``hi_even_nibble | (hi_odd_nibble << 4)``, ``hi = (u >> 4) & 1``). |
| 35 | + scale : (N, n_groups) uint8 — per-group scale *codes*, coalesced. |
| 36 | + scale_step : (N, K/256) fp16 — per-256-super-block scale step; the real |
| 37 | + per-group scale is ``scale_code * scale_step[:, g // 8]``. |
| 38 | + zero_point : (N, n_groups) uint8 — per-group zero *codes*, coalesced. |
| 39 | + zero_point_step : (N, K/256) fp16 — per-256-super-block zero step; the real |
| 40 | + per-group zero is ``zero_code * zero_point_step[:, g // 8]``. Both |
| 41 | + per-256 fp16 steps mirror GGUF Q5_K's per-super-block fp16 |
| 42 | + ``d`` / ``dmin`` and are packed into ONE 32-bit warp-shuffle |
| 43 | + word by the decode kernel (z_pack), exactly like the INT4 path. |
| 44 | + The finer per-256 step (vs the previous per-row step) improves |
| 45 | + whole-weight dequant SNR at ~5.625 bpw. |
| 46 | +
|
| 47 | +The pack/unpack helpers (:func:`pack_int5`, :func:`unpack_int5`) must stay in |
| 48 | +lockstep with ``int5_plain_mm.cuh`` (the decode kernel) — the per-32-weight |
| 49 | +even/odd high-bit byte order is the single most error-prone detail and is |
| 50 | +covered by the pack round-trip and the C++ gtest. |
| 51 | +""" |
| 52 | + |
| 53 | +from typing import List, Optional, Tuple |
| 54 | + |
| 55 | +import torch |
| 56 | +from torchao.utils import TorchAOBaseTensor |
| 57 | + |
| 58 | +__all__ = [ |
| 59 | + "CudaDp4aPlanarInt5Tensor", |
| 60 | + "pack_int5", |
| 61 | + "unpack_int5", |
| 62 | +] |
| 63 | + |
| 64 | +_CODE_MAX = 255 # uint8 code range [0, 255] (both scale and zero) |
| 65 | +_SUPER_BLOCK = 256 # weights per super-block (GGUF Q5_K QK_K); steps are per this |
| 66 | + |
| 67 | + |
| 68 | +def pack_int5(q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| 69 | + """Pack unsigned Q5_K values ``u`` in ``[0, 31]`` into the (ql, qh) planes. |
| 70 | +
|
| 71 | + Args: |
| 72 | + q: (N, K) integer tensor with values in ``[0, 31]``. |
| 73 | +
|
| 74 | + Returns: |
| 75 | + ``(ql, qh)`` where ``ql`` is ``(N, K/2)`` uint8 and ``qh`` is |
| 76 | + ``(N, K/8)`` uint8 (see the module docstring for the layout). |
| 77 | + """ |
| 78 | + if q.dim() != 2: |
| 79 | + raise ValueError(f"pack_int5 expects a 2-D tensor, got shape {tuple(q.shape)}") |
| 80 | + N, K = int(q.shape[0]), int(q.shape[1]) |
| 81 | + if K % 32 != 0: |
| 82 | + raise ValueError(f"K={K} must be a multiple of 32 for INT5 packing") |
| 83 | + |
| 84 | + # All intermediates are uint8 (values fit in a byte) to keep peak memory low. |
| 85 | + u = q.to(torch.uint8) # [0, 31] |
| 86 | + lo = u & 0xF # low nibble (uint8) |
| 87 | + hi = (u >> 4) & 0x1 # high 1 bit (uint8) |
| 88 | + |
| 89 | + # ql: nibble-pack the low plane even/odd, exactly like the INT4 path. |
| 90 | + ql = lo[:, 0::2] | (lo[:, 1::2] << 4) # (N, K/2) uint8 |
| 91 | + |
| 92 | + # qh: per 32-weight chunk -> 4 bytes (one per dp4a word); each byte packs the |
| 93 | + # four 1-bit highs of that word's even weights (low nibble) and odd weights |
| 94 | + # (high nibble), field j at bit j. |
| 95 | + chunks = K // 32 |
| 96 | + hw = hi.reshape(N, chunks, 4, 8) # (N, chunk, word, pos-in-word) |
| 97 | + even = hw[..., 0::2] # (N, chunk, 4, 4) positions 0,2,4,6 |
| 98 | + odd = hw[..., 1::2] # (N, chunk, 4, 4) positions 1,3,5,7 |
| 99 | + # Explicit OR (not sum) keeps the result uint8 (torch.sum would promote). |
| 100 | + hi_even_nib = ( |
| 101 | + even[..., 0] | (even[..., 1] << 1) | (even[..., 2] << 2) | (even[..., 3] << 3) |
| 102 | + ) # (N, chunk, 4) uint8, one nibble per word |
| 103 | + hi_odd_nib = ( |
| 104 | + odd[..., 0] | (odd[..., 1] << 1) | (odd[..., 2] << 2) | (odd[..., 3] << 3) |
| 105 | + ) |
| 106 | + qh = hi_even_nib | (hi_odd_nib << 4) # (N, chunk, 4) uint8 |
| 107 | + qh = qh.reshape(N, K // 8) |
| 108 | + return ql.contiguous(), qh.contiguous() |
| 109 | + |
| 110 | + |
| 111 | +def unpack_int5(ql: torch.Tensor, qh: torch.Tensor, N: int, K: int) -> torch.Tensor: |
| 112 | + """Inverse of :func:`pack_int5`. Returns ``(N, K)`` uint8 ``u`` in ``[0, 31]``. |
| 113 | +
|
| 114 | + All intermediates stay uint8 — the value is unsigned (the zero point is |
| 115 | + applied in dequant, not here). |
| 116 | + """ |
| 117 | + qlu = ql.to(torch.uint8) |
| 118 | + lo_even = qlu & 0xF # low nibble -> even weights |
| 119 | + lo_odd = (qlu >> 4) & 0xF # high nibble -> odd weights |
| 120 | + lo = torch.stack([lo_even, lo_odd], dim=-1).reshape(N, K) # uint8 |
| 121 | + |
| 122 | + chunks = K // 32 |
| 123 | + qhu = qh.to(torch.uint8).reshape(N, chunks, 4) # one byte per word |
| 124 | + hi_even_nib = qhu & 0xF # (N, chunk, 4) |
| 125 | + hi_odd_nib = (qhu >> 4) & 0xF |
| 126 | + hi_even = torch.stack( |
| 127 | + [(hi_even_nib >> s) & 0x1 for s in range(4)], dim=-1 |
| 128 | + ) # (N, chunk, 4, 4) uint8 |
| 129 | + hi_odd = torch.stack([(hi_odd_nib >> s) & 0x1 for s in range(4)], dim=-1) |
| 130 | + hi = torch.empty(N, chunks, 4, 8, dtype=torch.uint8, device=ql.device) |
| 131 | + hi[..., 0::2] = hi_even |
| 132 | + hi[..., 1::2] = hi_odd |
| 133 | + hi = hi.reshape(N, K) |
| 134 | + |
| 135 | + return lo | (hi << 4) # [0, 31] uint8 |
| 136 | + |
| 137 | + |
| 138 | +def _encode_uint8_per_super( |
| 139 | + x: torch.Tensor, |
| 140 | + group_size: int, |
| 141 | +) -> Tuple[torch.Tensor, torch.Tensor]: |
| 142 | + """Encode a (N, n_groups) non-negative per-group tensor to per-256 uint8 codes. |
| 143 | +
|
| 144 | + Used for both the scale and the zero. A super-block is ``_SUPER_BLOCK`` |
| 145 | + (256) weights = ``groups_per_super = 256 // group_size`` groups (8 for |
| 146 | + group_size=32). Q5_K per-group scales/zeros are non-negative (``d``/``dmin`` |
| 147 | + >= 0), so an unsigned code is exact at the endpoints. Returns |
| 148 | + ``(codes, step)`` where ``codes`` is ``(N, n_groups)`` uint8 and ``step`` is |
| 149 | + ``(N, n_super)`` fp16 with ``n_super = n_groups // groups_per_super = K // |
| 150 | + 256``, such that ``code * step[:, g // groups_per_super] ≈ x``. The step is |
| 151 | + the per-256-super-block max / 255 rounded to fp16 (what the kernel reads), so |
| 152 | + encode and decode agree. Mirrors the INT4 ``_encode_uint8_per_super`` (the |
| 153 | + int5 scale/zero are already row-major (N, n_groups), so no transpose). |
| 154 | + """ |
| 155 | + xf = x.contiguous().float() # (N, n_groups), non-negative |
| 156 | + N, n_groups = int(xf.shape[0]), int(xf.shape[1]) |
| 157 | + groups_per_super = _SUPER_BLOCK // int(group_size) |
| 158 | + if groups_per_super < 1: |
| 159 | + raise ValueError( |
| 160 | + f"group_size={group_size} must be <= {_SUPER_BLOCK} for the per-256 step" |
| 161 | + ) |
| 162 | + if n_groups % groups_per_super != 0: |
| 163 | + raise ValueError( |
| 164 | + f"n_groups={n_groups} must be a multiple of {groups_per_super} " |
| 165 | + f"(K must be a multiple of {_SUPER_BLOCK}) for group_size={group_size}" |
| 166 | + ) |
| 167 | + n_super = n_groups // groups_per_super |
| 168 | + xb = xf.reshape(N, n_super, groups_per_super) # (N, n_super, gps) |
| 169 | + block_max = xb.amax(dim=2, keepdim=True).clamp_min(1e-30) # (N, n_super, 1) |
| 170 | + step = (block_max / _CODE_MAX).to(torch.float16) # (N, n_super, 1) fp16 |
| 171 | + step_f = step.float().clamp_min(1e-30) |
| 172 | + codes = torch.round(xb / step_f).clamp_(0, _CODE_MAX).to(torch.uint8) |
| 173 | + codes = codes.reshape(N, n_groups).contiguous() |
| 174 | + return codes, step.squeeze(2).contiguous() |
| 175 | + |
| 176 | + |
| 177 | +class CudaDp4aPlanarInt5Tensor(TorchAOBaseTensor): |
| 178 | + """Dp4a-planar 5-bit weight (ql/qh split bit-planes + per-group scale/zero), asymmetric. |
| 179 | +
|
| 180 | + ExecuTorch-internal; see the module docstring. The CUDA decode/prefill |
| 181 | + dispatch (``int5_dispatch.py``) is selected by *type* — it is registered on |
| 182 | + this class only. |
| 183 | + """ |
| 184 | + |
| 185 | + tensor_data_names = [ |
| 186 | + "ql", |
| 187 | + "qh", |
| 188 | + "scale", |
| 189 | + "scale_step", |
| 190 | + "zero_point", |
| 191 | + "zero_point_step", |
| 192 | + ] |
| 193 | + tensor_attribute_names = ["block_size", "shape"] |
| 194 | + |
| 195 | + def __new__( |
| 196 | + cls, |
| 197 | + ql: torch.Tensor, |
| 198 | + qh: torch.Tensor, |
| 199 | + scale: torch.Tensor, |
| 200 | + scale_step: torch.Tensor, |
| 201 | + zero_point: torch.Tensor, |
| 202 | + zero_point_step: torch.Tensor, |
| 203 | + block_size: List[int], |
| 204 | + shape: torch.Size, |
| 205 | + ): |
| 206 | + kwargs = {} |
| 207 | + kwargs["device"] = ql.device |
| 208 | + kwargs["dtype"] = torch.bfloat16 |
| 209 | + kwargs["requires_grad"] = False |
| 210 | + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] |
| 211 | + |
| 212 | + def __init__( |
| 213 | + self, |
| 214 | + ql: torch.Tensor, |
| 215 | + qh: torch.Tensor, |
| 216 | + scale: torch.Tensor, |
| 217 | + scale_step: torch.Tensor, |
| 218 | + zero_point: torch.Tensor, |
| 219 | + zero_point_step: torch.Tensor, |
| 220 | + block_size: List[int], |
| 221 | + shape: torch.Size, |
| 222 | + ): |
| 223 | + super().__init__() |
| 224 | + self.ql = ql |
| 225 | + self.qh = qh |
| 226 | + self.scale = scale |
| 227 | + self.scale_step = scale_step |
| 228 | + self.zero_point = zero_point |
| 229 | + self.zero_point_step = zero_point_step |
| 230 | + self.block_size = block_size |
| 231 | + |
| 232 | + def _quantization_type(self): |
| 233 | + return f"shape={self.shape}, block_size={self.block_size}, device={self.device}" |
| 234 | + |
| 235 | + @classmethod |
| 236 | + def _from_intx_int8(cls, t: torch.Tensor) -> "CudaDp4aPlanarInt5Tensor": |
| 237 | + """Build from an asymmetric int5 ``IntxUnpackedToInt8Tensor`` decoded from Q5_K. |
| 238 | +
|
| 239 | + The source stores centered quants ``qdata`` in ``[-16, 15]`` (int8) with a |
| 240 | + (float) bf16 ``zero_point`` and per-group ``scale``, both ``(N, n_groups)`` |
| 241 | + (dequant = ``scale * (qdata - zero_point)``). We shift back to the raw |
| 242 | + unsigned ``u = qdata + 16`` in ``[0, 31]`` and the unsigned zero |
| 243 | + ``zero_u = zero_point + 16`` (both non-negative), bit-pack the ql/qh |
| 244 | + planes, and re-encode scale/zero to per-group uint8 codes each with a |
| 245 | + per-256-super-block fp16 step (mirroring the INT4 z_pack). All of this is |
| 246 | + baked into the serialized weight constant here, once at pack time. |
| 247 | + """ |
| 248 | + q_centered = t.qdata |
| 249 | + q_min, q_max = int(q_centered.min()), int(q_centered.max()) |
| 250 | + if q_min < -16 or q_max > 15: |
| 251 | + raise ValueError( |
| 252 | + f"Q5_K centered values must be in [-16, 15], got [{q_min}, {q_max}]" |
| 253 | + ) |
| 254 | + u = (q_centered.to(torch.int16) + 16).to(torch.uint8) # [0, 31] |
| 255 | + # Unsigned zero in q-space (>= 0 for Q5_K's non-negative affine min). |
| 256 | + zero_u = (t.zero_point.float() + 16.0).clamp_min(0.0) # (N, n_groups) |
| 257 | + scale = t.scale.float() # (N, n_groups), non-negative |
| 258 | + |
| 259 | + gs = int(t.block_size[-1]) |
| 260 | + ql, qh = pack_int5(u) |
| 261 | + scale_codes, scale_step = _encode_uint8_per_super(scale, gs) |
| 262 | + zero_codes, zero_point_step = _encode_uint8_per_super(zero_u, gs) |
| 263 | + return cls( |
| 264 | + ql, |
| 265 | + qh, |
| 266 | + scale_codes, |
| 267 | + scale_step, |
| 268 | + zero_codes, |
| 269 | + zero_point_step, |
| 270 | + list(t.block_size), |
| 271 | + t.shape, |
| 272 | + ) |
| 273 | + |
| 274 | + @classmethod |
| 275 | + def from_exportable_gguf(cls, gt) -> "CudaDp4aPlanarInt5Tensor": |
| 276 | + """Build from a native Q5_K ``ExportableGGUFTensor``. |
| 277 | +
|
| 278 | + Reuses the shared Q5_K block decode in |
| 279 | + ``extension/llm/export/gguf.py`` (``to_intx_unpacked_to_int8_tensor`` -> |
| 280 | + an asymmetric int5 tensor centered in ``[-16, 15]``), then bit-packs into |
| 281 | + the ql/qh planes via :meth:`_from_intx_int8`. The Q5_K decode lives in one |
| 282 | + place; this class only owns the 5-bit pack + metadata re-encode. |
| 283 | + """ |
| 284 | + if gt.ggml_type != "q5_k": |
| 285 | + raise ValueError( |
| 286 | + "CudaDp4aPlanarInt5Tensor.from_exportable_gguf requires a q5_k " |
| 287 | + f"ExportableGGUFTensor, got {gt.ggml_type!r}" |
| 288 | + ) |
| 289 | + return cls._from_intx_int8(gt.to_intx_unpacked_to_int8_tensor()) |
| 290 | + |
| 291 | + def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor: |
| 292 | + """Dequantize to a dense tensor (asymmetric: ``w = scale * (u - zero)``). |
| 293 | +
|
| 294 | + Reconstructs the per-group scale/zero from the uint8 codes and their |
| 295 | + per-256-super-block fp16 steps (broadcast over the 8 groups per |
| 296 | + super-block). Used for the tied lm_head / token embedding (which can't |
| 297 | + gather a packed tensor) and as the numerical reference. |
| 298 | + """ |
| 299 | + dtype = output_dtype if output_dtype is not None else torch.bfloat16 |
| 300 | + N, K = int(self.shape[0]), int(self.shape[1]) |
| 301 | + gs = self.block_size[-1] |
| 302 | + n_groups = K // gs |
| 303 | + n_super = int(self.scale_step.shape[1]) |
| 304 | + groups_per_super = n_groups // n_super |
| 305 | + |
| 306 | + u = unpack_int5(self.ql, self.qh, N, K).to(torch.float32) |
| 307 | + scale_code = self.scale.to(torch.float32) # (N, n_groups) |
| 308 | + scale_step = self.scale_step.float().repeat_interleave(groups_per_super, dim=1) |
| 309 | + scale = (scale_code * scale_step).repeat_interleave(gs, dim=1) # (N, K) |
| 310 | + |
| 311 | + zero_code = self.zero_point.to(torch.float32) # (N, n_groups) |
| 312 | + zero_point_step = self.zero_point_step.float().repeat_interleave( |
| 313 | + groups_per_super, dim=1 |
| 314 | + ) |
| 315 | + zero = (zero_code * zero_point_step).repeat_interleave(gs, dim=1) # (N, K) |
| 316 | + return (scale * (u - zero)).to(dtype) |
| 317 | + |
| 318 | + |
| 319 | +# Allow a model with CudaDp4aPlanarInt5Tensor weights to be loaded with |
| 320 | +# `weights_only=True` (mirrors torchao quantized tensors). |
| 321 | +torch.serialization.add_safe_globals([CudaDp4aPlanarInt5Tensor]) |
0 commit comments