File size: 2,928 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | from __future__ import annotations
import math
import torch
import torch.nn.functional as F
from .nibbles import pack_uint4, unpack_uint4
EPS = 1e-10
def fwht_last_dim(x: torch.Tensor, block_size: int) -> torch.Tensor:
if block_size <= 0 or block_size & (block_size - 1):
raise ValueError("block_size must be a power of two")
if x.shape[-1] % block_size:
raise ValueError("last dimension must divide block_size")
y = x.reshape(*x.shape[:-1], x.shape[-1] // block_size, block_size)
step = 1
while step < block_size:
shape = y.shape
z = y.reshape(*shape[:-1], -1, 2, step)
a, b = z[..., 0, :], z[..., 1, :]
y = torch.stack((a + b, a - b), dim=-2).reshape(*shape)
step *= 2
return (y / math.sqrt(block_size)).reshape_as(x)
def rpbh(x: torch.Tensor, perm: torch.Tensor, signs: torch.Tensor, block_size: int) -> torch.Tensor:
p = perm.to(device=x.device, dtype=torch.long)
s = signs.to(device=x.device, dtype=torch.float32)
y = x.float().index_select(-1, p) * s
return fwht_last_dim(y, block_size)
def nearest_codes(unit: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor:
cb = codebook.to(device=unit.device, dtype=torch.float32)
thresholds = 0.5 * (cb[:-1] + cb[1:])
return torch.bucketize(unit.contiguous(), thresholds).to(torch.uint8)
def a4_pack_reference(
x: torch.Tensor,
perm: torch.Tensor,
signs: torch.Tensor,
block_size: int,
codebook: torch.Tensor,
eps: float = EPS,
) -> tuple[torch.Tensor, torch.Tensor]:
rot = rpbh(x, perm, signs, block_size)
scale = torch.linalg.vector_norm(rot, ord=2, dim=-1).float()
unit = rot / (scale.unsqueeze(-1) + eps)
codes = nearest_codes(unit, codebook)
return pack_uint4(codes), scale
def dequant_packed(
packed: torch.Tensor,
scale: torch.Tensor,
codebook: torch.Tensor,
logical_k: int,
dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
codes = unpack_uint4(packed, logical_k).long()
cb = codebook.to(device=packed.device, dtype=torch.float32)
return (cb[codes] * scale.float().unsqueeze(-1)).to(dtype)
def w4a4_linear_reference(
a_packed: torch.Tensor,
a_scale: torch.Tensor,
w_packed: torch.Tensor,
w_scale: torch.Tensor,
codebook: torch.Tensor,
logical_k: int,
bias: torch.Tensor | None = None,
out_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
a = dequant_packed(a_packed, a_scale, codebook, logical_k, torch.bfloat16)
w = dequant_packed(w_packed, w_scale, codebook, logical_k, torch.bfloat16)
b = bias.to(torch.bfloat16) if bias is not None else None
# Match the Project-A fake-quant runtime: the quantized operands enter
# F.linear as BF16. CUDA tensor cores accumulate internally; the public
# numerical contract is BF16 operands and BF16 output.
out = F.linear(a, w, b)
return out.to(out_dtype)
|