| """Lossless storage for LoopQ signed group-32 W4 QDQ weights. |
| |
| This is a storage codec, not an INT4 GEMM kernel. Dequantization materializes |
| one dense weight; resident-memory and runtime-speed claims require integration. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import math |
| import torch |
| from .quantization import quantize_weight |
|
|
| _DTYPES = {str(dtype): dtype for dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64)} |
|
|
|
|
| @dataclass(frozen=True) |
| class PackedLoopQWeight: |
| codes: torch.Tensor |
| scales: torch.Tensor |
| shape: tuple[int, int] |
| output_dtype: str |
| format_version: int = 2 |
|
|
| def validate(self): |
| if len(self.shape) != 2 or any(type(x) is not int or x <= 0 for x in self.shape): |
| raise ValueError('packed weight shape must contain two positive integers') |
| rows, width = self.shape |
| if self.codes.dtype != torch.uint8 or self.codes.shape != (rows, math.ceil(width/2)): |
| raise ValueError('packed code shape/dtype mismatch') |
| if self.scales.dtype != torch.float32 or self.scales.shape != (rows, math.ceil(width/32)): |
| raise ValueError('group-32 FP32 scale shape/dtype mismatch') |
| if not torch.isfinite(self.scales).all() or not (self.scales > 0).all(): |
| raise ValueError('scales must be finite and positive') |
| if self.output_dtype not in _DTYPES: |
| raise ValueError('unsupported output dtype') |
| if width % 2 and (self.codes[:, -1] >> 4).any(): |
| raise ValueError('nonzero unused high nibble') |
| if self.format_version not in (1, 2): |
| raise ValueError('unsupported packed-weight version') |
| if self.format_version == 1 and ( |
| ((self.codes & 15) == 8).any() or ((self.codes >> 4) == 8).any()): |
| raise ValueError('reserved -8 code is outside LoopQ narrow signed range') |
|
|
| @property |
| def payload_bytes(self): |
| """Tensor payload only; excludes serialization/container overhead.""" |
| return self.codes.numel() + self.scales.numel() * 4 |
|
|
| def dequantize(self, *, device=None): |
| self.validate() |
| return self._dequantize_validated(device=device) |
|
|
| def _dequantize_validated(self, *, device=None): |
| """Internal immutable-dispatch path; the owner validates before use.""" |
| codes = self.codes.to(device=device) if device is not None else self.codes |
| pairs = torch.stack((codes & 15, codes >> 4), dim=-1).flatten(-2) |
| values = pairs[:, :self.shape[1]].to(torch.int16) |
| values = torch.where(values >= 8, values - 16, values).float() |
| scales = self.scales.to(device=values.device).repeat_interleave(32, dim=-1)[:, :self.shape[1]] |
| return (values * scales).to(_DTYPES[self.output_dtype]) |
|
|
| def state_dict(self): |
| self.validate() |
| return dict(format='loopq_packed_w4_group32', format_version=self.format_version, |
| nibble_order='even_column_low', shape=list(self.shape), |
| output_dtype=self.output_dtype, codes=self.codes.detach().cpu(), |
| scales=self.scales.detach().cpu()) |
|
|
| @classmethod |
| def from_state_dict(cls, state): |
| if state.get('format') != 'loopq_packed_w4_group32' or state.get('format_version') not in (1, 2) \ |
| or state.get('nibble_order') != 'even_column_low': |
| raise ValueError('unsupported packed-weight contract') |
| result = cls(state['codes'], state['scales'], tuple(state['shape']), |
| state['output_dtype'], state['format_version']) |
| result.validate() |
| return result |
|
|
|
|
| def pack_weight(weight: torch.Tensor) -> PackedLoopQWeight: |
| if weight.ndim != 2 or min(weight.shape) <= 0: |
| raise ValueError('weight must be a nonempty matrix') |
| if not torch.isfinite(weight).all(): |
| raise ValueError('weight must be finite') |
| quantized = quantize_weight(weight.detach()) |
| integers = quantized.integers |
| unsigned = (integers.to(torch.int16) & 15).to(torch.uint8) |
| if weight.shape[1] % 2: |
| unsigned = torch.nn.functional.pad(unsigned, (0, 1)) |
| codes = unsigned[:, 0::2] | (unsigned[:, 1::2] << 4) |
| result = PackedLoopQWeight(codes, quantized.scales.detach().float(), tuple(weight.shape), str(weight.dtype)) |
| result.validate() |
| return result |
|
|
