| import torch |
|
|
| from orbitquant_wan_a2.nibbles import pack_uint4 |
| from orbitquant_wan_a2.reference import w4a4_linear_reference |
| from orbitquant_wan_a2.triton_w4a4 import PackedActivation, w4a4_linear_triton |
|
|
|
|
| def test_gemm_native_transposed_weight_layout_cpu_reference(): |
| torch.manual_seed(17) |
| m, n, k = 5, 13, 64 |
| cb = torch.linspace(-0.2, 0.2, 16, dtype=torch.float32) |
| ac = torch.randint(0, 16, (m, k), dtype=torch.uint8) |
| wc = torch.randint(0, 16, (n, k), dtype=torch.uint8) |
| ap = pack_uint4(ac) |
| wp_row = pack_uint4(wc) |
| wp_runtime = wp_row.transpose(0, 1).contiguous() |
| a_scale = torch.rand(m) + 0.1 |
| w_scale = torch.rand(n) + 0.1 |
| bias = torch.randn(n, dtype=torch.bfloat16) |
| packed_a = PackedActivation(ap, a_scale, k) |
| got = w4a4_linear_triton(packed_a, wp_runtime, w_scale, cb, bias) |
| ref = w4a4_linear_reference(ap, a_scale, wp_row, w_scale, cb, k, bias) |
| assert torch.equal(got, ref) |
|
|
