Kernels:

phanerozoic
/

bitnet-cpu

Kernel card Files Files and versions

xet

Community

phanerozoic commited on Jul 10

Commit

4adadc1

verified ·

1 Parent(s): b02a470

Tests reference the kernel-quantized operands; reciprocal quantization diverges one ulp at ties

Browse files

Files changed (1) hide show

tests/test_bitnet_cpu.py +37 -18

tests/test_bitnet_cpu.py CHANGED Viewed

@@ -6,47 +6,66 @@ import kernels
 bitnet = kernels.get_kernel("phanerozoic/bitnet-cpu", version=1, trust_remote_code=True)
-def _reference(x, W, scale_wt):
-    """f32 reference: quantize per-token absmax, integer matmul, rescale."""
-    xf = x.float()
-    absmax = xf.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5)
-    scale = absmax / 127.0
-    q = (xf / scale).round().clamp(-127, 127)
-    acc = q @ W.float().t()
-    return acc * scale * scale_wt.float().unsqueeze(0)
 @pytest.mark.kernels_ci
 @pytest.mark.parametrize("M", [1, 4, 16, 128])
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
-def test_bitnet_linear_matches_reference(M, dtype):
     torch.manual_seed(0)
     N, K = 512, 2560
     W = torch.randint(-1, 2, (N, K), dtype=torch.int8)
-    w_packed = bitnet.pack_weights(W)
-    scale_wt = (torch.rand(N) * 0.5 + 0.5).to(torch.bfloat16)
     x = torch.randn(M, K, dtype=dtype)
-    y = bitnet.bitnet_linear(x, w_packed, scale_wt).float()
-    ref = _reference(x, W, scale_wt)
-    # bf16 output rounding is the only tolerance needed; the integer part
-    # of the computation is exact.
-    torch.testing.assert_close(y, ref, rtol=1e-2, atol=1e-2)
 @pytest.mark.kernels_ci
-def test_quantize_activation_roundtrip():
     torch.manual_seed(1)
     x = torch.randn(8, 1024, dtype=torch.bfloat16)
     q, s = bitnet.quantize_activation(x)
     assert q.dtype == torch.int8 and s.dtype == torch.bfloat16
     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(2)
     lin = torch.nn.Linear(2560, 512, bias=False)
     bl = bitnet.BitLinear.from_dense(lin)
     x = torch.randn(4, 2560, dtype=torch.bfloat16)

 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)