mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
J.W (Jun)
25908d48a6
_tir_packed_to_unsigned_convert_with_zeros performed the zero-point subtraction in the unsigned packed storage dtype before casting to the requested output dtype. Values below the zero-point therefore wrapped modulo the storage width. Cast the unpacked field and zero-point to the output dtype before subtraction so asymmetric dequantization preserves negative values. Add CUDA regression coverage for uint32-packed 4-bit values decoded with a non-zero zero-point. Fixes #2947
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from examples.dequantize_gemm.quantize.quantization import _tir_packed_to_unsigned_convert_with_zeros
|
|
from tvm import tirx
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_packed_unsigned_convert_with_zeros_uses_signed_output_domain():
|
|
"""Zero-point subtraction must not underflow in uint32 packed storage."""
|
|
|
|
n = 16
|
|
zero = 8
|
|
convert = _tir_packed_to_unsigned_convert_with_zeros("uint", 32)
|
|
|
|
@T.prim_func
|
|
def dequantize(packed: T.Tensor((n,), "uint32"), decoded: T.Tensor((n,), "float16")):
|
|
with T.Kernel(1, threads=n):
|
|
thread_id = T.get_thread_binding()
|
|
decoded[thread_id] = convert(
|
|
4,
|
|
packed[thread_id],
|
|
tirx.const(0, "int32"),
|
|
tirx.const(zero, "uint32"),
|
|
"float16",
|
|
)
|
|
|
|
kernel = tilelang.compile(dequantize, target="cuda", execution_backend="nvrtc")
|
|
packed = torch.tensor(np.arange(n, dtype=np.uint32), device="cuda")
|
|
decoded = torch.empty(n, dtype=torch.float16, device="cuda")
|
|
kernel(packed, decoded)
|
|
|
|
expected = torch.arange(n, dtype=torch.float16) - zero
|
|
torch.testing.assert_close(decoded.cpu(), expected, rtol=0, atol=0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|