mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
from examples.dequantize_gemm.quantize.quantization import _tir_u32_to_int_to_float
|
|
|
|
|
|
def _pack_signed_values(values: list[int], nbit: int) -> torch.Tensor:
|
|
lanes_per_word = 32 // nbit
|
|
words = []
|
|
mask = (1 << nbit) - 1
|
|
for word_start in range(0, len(values), lanes_per_word):
|
|
packed = 0
|
|
for pos, value in enumerate(values[word_start : word_start + lanes_per_word]):
|
|
packed |= (value & mask) << (pos * nbit)
|
|
words.append(packed)
|
|
return torch.tensor(words, dtype=torch.uint32, device="cuda")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"nbit,values",
|
|
[
|
|
(2, [-2, -1, 0, 1]),
|
|
(4, list(range(-8, 8))),
|
|
(8, [-128, -17, -1, 0, 1, 42, 127]),
|
|
],
|
|
)
|
|
@tilelang.testing.requires_cuda
|
|
def test_u32_to_int_to_float_sign_extends_subword_values(nbit, values):
|
|
"""Signed sub-word decode should preserve negative values from uint32 storage."""
|
|
|
|
lanes_per_word = 32 // nbit
|
|
num_words = math.ceil(len(values) / lanes_per_word)
|
|
num_values = len(values)
|
|
|
|
@T.prim_func
|
|
def main(packed_values: T.Tensor((num_words,), "uint32"), decoded_values: T.Tensor((num_values,), "float32")):
|
|
with T.Kernel(1, threads=1) as _:
|
|
for i in T.serial(num_values):
|
|
decoded_values[i] = _tir_u32_to_int_to_float(nbit, packed_values[i // lanes_per_word], i % lanes_per_word, "float32")
|
|
|
|
kernel = tilelang.compile(main, target="cuda")
|
|
packed = _pack_signed_values(values, nbit)
|
|
out = torch.empty(num_values, dtype=torch.float32, device="cuda")
|
|
kernel(packed, out)
|
|
torch.testing.assert_close(out.cpu(), torch.tensor(values, dtype=torch.float32), rtol=0, atol=0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|