File size: 1,610 Bytes
d72b5a3 | 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 | import torch
@torch.no_grad()
def unpack_sc_lut_indices(packed: torch.Tensor, n_tiles: int, bits: int) -> torch.Tensor:
R = packed.shape[0]
if bits == 10:
p = packed.view(R, n_tiles // 4, 5).to(torch.int32)
i0 = p[:, :, 0] | ((p[:, :, 1] & 0x03) << 8)
i1 = (p[:, :, 1] >> 2) | ((p[:, :, 2] & 0x0F) << 6)
i2 = (p[:, :, 2] >> 4) | ((p[:, :, 3] & 0x3F) << 4)
i3 = (p[:, :, 3] >> 6) | (p[:, :, 4] << 2)
return torch.stack([i0, i1, i2, i3], dim=-1).reshape(R, n_tiles)
if bits == 9:
p = packed.view(R, n_tiles // 8, 9).to(torch.int64)
lo = torch.zeros((R, n_tiles // 8), dtype=torch.int64, device=packed.device)
for k in range(8):
lo |= (p[:, :, k] & 0xFF) << (8 * k)
out = [((lo >> (9 * t)) & 0x1FF).to(torch.int32) for t in range(7)]
i7 = (((lo >> 63) & 0x01) | ((p[:, :, 8] & 0xFF) << 1)).to(torch.int32)
out.append(i7)
return torch.stack(out, dim=-1).reshape(R, n_tiles)
raise ValueError(f"Unsupported bit width: {bits}")
@torch.no_grad()
def decode_sc_lut_matrix(idx_or_packed, scales, u_svd, v_svd, lut, q_basis, bits=0, out_dtype=torch.bfloat16):
R, C = u_svd.shape[0], v_svd.shape[1]
tile = lut.shape[1]
n_tiles = C // tile
idx = unpack_sc_lut_indices(idx_or_packed, n_tiles, bits) if bits > 0 else idx_or_packed
w_rot = lut.float()[idx.long()].reshape(R, scales.shape[1], -1) * scales.float().unsqueeze(-1)
w_main = (w_rot.reshape(R, C // 64, 64) @ q_basis.float()).reshape(R, C)
return (w_main + (u_svd.float() @ v_svd.float())).to(out_dtype)
|