File size: 2,905 Bytes
ce8679e 4adadc1 ce8679e 4adadc1 ce8679e 4adadc1 ce8679e 4adadc1 ce8679e 4adadc1 ce8679e 4adadc1 ce8679e 4adadc1 ce8679e 4adadc1 ce8679e | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | 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)
@pytest.mark.kernels_ci
@pytest.mark.parametrize("M", [1, 4, 16, 128])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
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}"
@pytest.mark.kernels_ci
@pytest.mark.parametrize("M", [1, 4, 15])
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}"
@pytest.mark.kernels_ci
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)
@pytest.mark.kernels_ci
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()
|