| 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) |
|
|
