Tests reference the kernel-quantized operands; reciprocal quantization diverges one ulp at ties
4adadc1 verified | import pytest | |
| import torch | |
| import kernels | |
| bitnet = kernels.get_kernel("phanerozoic/bitnet-cpu", version=1, trust_remote_code=True) | |
| def unpack_ternary(wp): | |
| cols = [((wp >> (2 * j)) & 3).to(torch.int16) - 2 for j in range(4)] | |
| return torch.stack(cols, dim=-1).reshape(wp.shape[0], wp.shape[1] * 4) | |
| def test_gemm_matches_exact_reference(M, dtype): | |
| """The integer path is exact: feeding the kernel's own quantized operands | |
| into an f32 reference bounds the residual by bf16 output rounding alone | |
| (one ulp, 2^-8 relative). An independent Python re-quantization is not a | |
| valid reference: dividing by amax/127 and multiplying by 127/amax differ | |
| by one ulp at rounding boundaries and flip occasional codes.""" | |
| torch.manual_seed(0) | |
| N, K = 512, 2560 | |
| W = torch.randint(-1, 2, (N, K), dtype=torch.int8) | |
| wp = bitnet.pack_weights(W) | |
| sw = (torch.rand(N) * 0.5 + 0.5).to(torch.bfloat16) | |
| x = torch.randn(M, K, dtype=dtype) | |
| q, s = bitnet.quantize_activation(x) | |
| y = bitnet.bitnet_gemm(q, wp, s, sw).float() | |
| ref = (q.float() @ unpack_ternary(wp).float().t()) * s.float().unsqueeze(-1) * sw.float().unsqueeze(0) | |
| rel = ((y - ref).abs() / ref.abs().clamp(min=1.0)).max().item() | |
| assert rel < 8e-3, f"max rel {rel}" | |
| def test_fused_path_matches_gemm_path(M): | |
| """The fused (M<16) path quantizes internally with the same code as | |
| quantize_activation; outputs agree to bf16 rounding of the scale.""" | |
| torch.manual_seed(1) | |
| N, K = 1024, 4096 | |
| W = torch.randint(-1, 2, (N, K), dtype=torch.int8) | |
| wp = bitnet.pack_weights(W) | |
| sw = torch.ones(N, dtype=torch.bfloat16) | |
| x = torch.randn(M, K, dtype=torch.bfloat16) | |
| y_fused = bitnet.bitnet_gemv_fused(x, wp, sw).float() | |
| q, s = bitnet.quantize_activation(x) | |
| y_split = bitnet.bitnet_gemm(q, wp, s, sw).float() | |
| rel = ((y_fused - y_split).abs() / y_split.abs().clamp(min=1.0)).max().item() | |
| assert rel < 8e-3, f"max rel {rel}" | |
| def test_quantize_activation_roundtrip(): | |
| torch.manual_seed(2) | |
| x = torch.randn(8, 1024, dtype=torch.bfloat16) | |
| q, s = bitnet.quantize_activation(x) | |
| assert q.dtype == torch.int8 and s.dtype == torch.bfloat16 | |
| assert (q.abs() <= 127).all() | |
| recon = q.float() * s.float().unsqueeze(-1) | |
| torch.testing.assert_close(recon, x.float(), rtol=2e-2, atol=2e-2) | |
| def test_bitlinear_module(): | |
| torch.manual_seed(3) | |
| lin = torch.nn.Linear(2560, 512, bias=False) | |
| bl = bitnet.BitLinear.from_dense(lin) | |
| x = torch.randn(4, 2560, dtype=torch.bfloat16) | |
| y = bl(x) | |
| assert y.shape == (4, 512) and y.dtype == torch.bfloat16 | |
| assert torch.isfinite(y.float()).all() | |