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