File size: 1,284 Bytes
f0c91ec | 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 | from __future__ import annotations
import torch
from orbitquant_packed_matmul import matmul_packed_weight
def _pack(values: torch.Tensor, bits: int) -> torch.Tensor:
flat = values.detach().to(device="cpu", dtype=torch.uint8).flatten()
packed = torch.zeros((flat.numel() * bits + 7) // 8, dtype=torch.uint8)
for value_index, value in enumerate(flat.tolist()):
bit_start = value_index * bits
byte_index = bit_start // 8
shift = bit_start % 8
packed[byte_index] |= (value << shift) & 0xFF
if shift + bits > 8:
packed[byte_index + 1] |= value >> (8 - shift)
return packed
device = "cuda" if torch.cuda.is_available() else "mps"
bits = 4
rows = 8
in_features = 16
out_features = 6
x = torch.randn(rows, in_features, device=device, dtype=torch.float16)
indices = torch.arange(out_features * in_features, dtype=torch.uint8).reshape(
out_features, in_features
) % (2**bits)
packed = _pack(indices, bits).to(device)
row_norms = torch.linspace(0.5, 1.5, out_features, device=device)
centroids = torch.linspace(-1.0, 1.0, 2**bits, device=device)
out = matmul_packed_weight(
x,
packed,
row_norms,
centroids,
bits=bits,
out_features=out_features,
in_features=in_features,
)
print(out.shape)
|