File size: 3,812 Bytes
f0c91ec fe97b1d f0c91ec fe97b1d f0c91ec fe97b1d f0c91ec fe97b1d f0c91ec fe97b1d f0c91ec fe97b1d f0c91ec fe97b1d | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | from __future__ import annotations
import pytest
import torch
from orbitquant_packed_matmul import matmul_packed_w4a4_int8, supports_device
@pytest.mark.kernels_ci
@pytest.mark.parametrize(
"rows,n,k",
[
(1, 129, 64),
(8, 3, 64),
(1, 129, 512),
(1, 2048, 1024),
(1, 32768, 2048),
(8, 129, 1024),
(2, 1024, 2048),
(2, 2047, 1024),
(2, 2048, 1024),
(3, 2049, 2048),
(4, 2048, 960),
(4, 2048, 1024),
(8, 2049, 16384),
(9, 2049, 1024),
(8, 2048, 6144),
(1, 257, 16384),
(9, 129, 128),
],
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize(
"offset,bias_enabled,k_major",
[(0, False, False), (0, True, False), (1, True, False), (3, False, False), (0, True, True)],
)
def test_decode_against_integer_oracle(rows, n, k, dtype, offset, bias_enabled, k_major):
if not torch.cuda.is_available() or not supports_device("cuda"):
pytest.skip("CUDA kernel required")
torch.manual_seed(42)
x = torch.randint(0, 256, (rows * k // 2 + offset,), dtype=torch.uint8)[offset:]
w = torch.randint(0, 256, (n * k // 2 + offset,), dtype=torch.uint8)[offset:]
# Exercise every code, including signed INT8 extremes and mixed nibble signs.
ac = torch.tensor(
[-128, 127, -1, 0, 1, -64, 64, -32, 32, -16, 16, -8, 8, -4, 4, 2], dtype=torch.int8
)
wc = ac.flip(0)
def unpack(values, count):
values = values.reshape(count, k // 2)
return torch.stack((values & 15, values >> 4), -1).reshape(count, k).long()
sums = ac.long()[unpack(x, rows)] @ wc.long()[unpack(w, n)].T
# Transfer whole storage to retain the deliberately unaligned device view.
def device_slice(values):
storage = torch.zeros(values.numel() + offset, dtype=torch.uint8)
storage[offset:] = values
return storage.cuda()[offset:]
xd = device_slice(x).reshape(rows, k // 2)
wd = device_slice(w)
if k_major:
wd = wd.reshape(n, k // 2).T.contiguous()
xn = torch.rand(rows, device="cuda")
wn = torch.rand(n, device="cuda", dtype=torch.bfloat16)
bias = torch.rand(n, device="cuda", dtype=dtype) if bias_enabled else None
# Binary-exact scales isolate the integer accumulation and norm epilogue.
expected = sums.cuda().float() * (xn[:, None] * wn.float()[None, :] * 0.001953125)
alternate = expected
if bias is not None:
# Existing kernels permit both contracted and separate FP32 multiply/
# add epilogues. Check exactly those two results, not a loose tolerance.
alternate = expected + bias.float()
scale = xn[:, None] * wn.float()[None, :] * 0.001953125
expected = (sums.cuda().float().double() * scale.double() + bias.double()).float()
expected = expected.to(dtype)
alternate = alternate.to(dtype)
args = (xd, wd, xn, wn, ac.cuda(), wc.cuda())
kwargs = dict(
activation_scale=0.03125,
weight_scale=0.0625,
out_features=n,
in_features=k,
bias=bias,
output_dtype=dtype,
weight_k_major=k_major,
)
actual = matmul_packed_w4a4_int8(*args, **kwargs)
assert torch.all((actual == expected) | (actual == alternate))
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
for _ in range(3):
matmul_packed_w4a4_int8(*args, **kwargs)
torch.cuda.current_stream().wait_stream(stream)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured = matmul_packed_w4a4_int8(*args, **kwargs)
graph.replay()
torch.cuda.synchronize()
torch.testing.assert_close(captured, actual, rtol=0, atol=0)
|