| """Shared bit-twiddling helpers for GGUF k-quant packed dequantization.""" | |
| from __future__ import annotations | |
| import torch | |
| def f16_view_to_f32(u8: torch.Tensor) -> torch.Tensor: | |
| """View uint8 bytes as float16 then cast to float32. u8: [..., 2] uint8.""" | |
| flat = u8.reshape(-1, 2).to(torch.int16) | |
| val = flat[:, 0] | (flat[:, 1].to(torch.int16) << 8) | |
| f16 = val.view(torch.float16) | |
| return f16.to(torch.float32).reshape(u8.shape[:-1]) | |
| def unpack_scales_k4(scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Unpack 8 (scale, min) 6-bit pairs from 12 bytes. scales: [..., 12] uint8. | |
| Returns (sc, m) each [..., 8] (6-bit values as uint8). | |
| """ | |
| q = scales.to(torch.int32) | |
| sh = q.shape[:-1] | |
| sc = torch.zeros(*sh, 8, dtype=torch.int32, device=q.device) | |
| m = torch.zeros(*sh, 8, dtype=torch.int32, device=q.device) | |
| sc[..., 0:4] = q[..., 0:4] & 0x3F | |
| m[..., 0:4] = q[..., 4:8] & 0x3F | |
| sc[..., 4:8] = (q[..., 8:12] & 0x0F) | ((q[..., 0:4] >> 6) << 4) | |
| m[..., 4:8] = (q[..., 8:12] >> 4) | ((q[..., 4:8] >> 6) << 4) | |
| return sc.to(torch.uint8), m.to(torch.uint8) |