File size: 784 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | import torch
from orbitquant_wan_a2.nibbles import pack_uint4
from orbitquant_wan_a2.reference import w4a4_linear_reference
def test_reference_matches_explicit_bf16_dequant():
torch.manual_seed(3)
m, n, k = 7, 11, 64
cb = torch.linspace(-0.2, 0.2, 16)
ac = torch.randint(0,16,(m,k),dtype=torch.uint8)
wc = torch.randint(0,16,(n,k),dtype=torch.uint8)
a_s = torch.rand(m) + 0.1
w_s = torch.rand(n) + 0.1
bias = torch.randn(n, dtype=torch.bfloat16)
out = w4a4_linear_reference(pack_uint4(ac), a_s, pack_uint4(wc), w_s, cb, k, bias)
a = (cb[ac.long()] * a_s[:,None]).to(torch.bfloat16)
w = (cb[wc.long()] * w_s[:,None]).to(torch.bfloat16)
ref = (a.float() @ w.float().T + bias.float()).to(torch.bfloat16)
assert torch.equal(out, ref)
|