File size: 1,489 Bytes
cc96dd8 | 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 | import kernels
import pytest
import torch
k = kernels.get_kernel("GoedelMachines/dg-w4-kernels")
@pytest.mark.kernels_ci
def test_fused_entropy_matches_categorical():
x = torch.randn(64, 8192, device="cuda", dtype=torch.float32) * 3
ref = torch.distributions.Categorical(logits=x).entropy()
got = k.fused_entropy(x)
torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4)
@pytest.mark.kernels_ci
def test_gumbel_argmax_returns_exact_argmax():
x = torch.randn(32, 4096, device="cuda", dtype=torch.float32)
_, amax = k.gumbel_argmax_sample(x, seed=1234)
assert (amax == x.argmax(-1)).all()
@pytest.mark.kernels_ci
def test_w4a16_matches_dequant_reference():
W = torch.randn(512, 1024, device="cuda") * 0.02
qw, s, z = k.quantize_w4(W, 128)
Wdq = k.dequant_w4(qw, s, z, 128)
x = torch.randn(64, 1024, device="cuda", dtype=torch.float16)
ref = torch.nn.functional.linear(x, Wdq)
got = k.w4a16_linear(x, qw, s, z, BK=128)
rel = (got.float() - ref.float()).abs().max() / ref.float().abs().max()
assert rel < 2e-2, rel
@pytest.mark.kernels_ci
def test_fused_rmsnorm_matches_reference():
x = torch.randn(128, 2816, device="cuda", dtype=torch.bfloat16)
w = torch.randn(2816, device="cuda", dtype=torch.bfloat16)
ref = (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6) * w.float())
got = k.fused_rmsnorm(x, w, 1e-6).float()
assert ((got - ref).abs().mean() / ref.abs().mean()) < 5e-3
|