| 10 | |
| 11 | |
| 12 | class TestLLMInt8Ops: |
| 13 | @pytest.mark.parametrize("device", get_available_devices()) |
| 14 | def test_int8_linear_matmul(self, device): |
| 15 | A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device) |
| 16 | B = torch.randint(-128, 127, (30, 20), dtype=torch.int8, device=device) |
| 17 | out = torch.ops.bitsandbytes.int8_linear_matmul.default(A, B) |
| 18 | |
| 19 | assert out.shape == (10, 30) |
| 20 | assert out.dtype == torch.int32 |
| 21 | assert out.device == A.device |
| 22 | |
| 23 | opcheck(torch.ops.bitsandbytes.int8_linear_matmul.default, (A, B)) |
| 24 | |
| 25 | @pytest.mark.parametrize("device", get_available_devices()) |
| 26 | def test_int8_linear_matmul_out(self, device): |
| 27 | A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device) |
| 28 | B = torch.randint(-128, 127, (30, 20), dtype=torch.int8, device=device) |
| 29 | |
| 30 | out = torch.empty((10, 30), dtype=torch.int32, device=device) |
| 31 | torch.ops.bitsandbytes.int8_linear_matmul.out(A, B, out) |
| 32 | |
| 33 | assert out.shape == (10, 30) |
| 34 | assert out.dtype == torch.int32 |
| 35 | assert out.device == A.device |
| 36 | |
| 37 | opcheck(torch.ops.bitsandbytes.int8_linear_matmul.out, (A, B, out)) |
| 38 | |
| 39 | @pytest.mark.parametrize("threshold", [0.0, 6.0]) |
| 40 | @pytest.mark.parametrize("device", get_available_devices()) |
| 41 | def test_int8_vectorwise_quant(self, threshold, device): |
| 42 | A = torch.randn(10, 20, dtype=torch.float16, device=device) |
| 43 | A[1][0] = 1000.0 |
| 44 | |
| 45 | out_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant(A, threshold=threshold) |
| 46 | |
| 47 | assert out_row.shape == (10, 20) |
| 48 | assert out_row.dtype == torch.int8 |
| 49 | assert out_row.device == A.device |
| 50 | assert row_stats.shape == (10,) |
| 51 | assert row_stats.dtype == torch.float32 |
| 52 | assert row_stats.device == A.device |
| 53 | |
| 54 | if threshold > 0.0: |
| 55 | assert outlier_cols is not None |
| 56 | assert outlier_cols.dim() == 1 |
| 57 | assert outlier_cols.shape[0] <= A.shape[1] |
| 58 | assert outlier_cols.device == A.device |
| 59 | else: |
| 60 | assert outlier_cols is None |
| 61 | |
| 62 | opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,)) |
| 63 | opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold)) |
| 64 | |
| 65 | @pytest.mark.parametrize("device", get_available_devices()) |
| 66 | def test_int8_mm_dequant(self, device): |
| 67 | A = torch.randint(-128, 127, (256, 256), dtype=torch.int32, device=device) |
| 68 | row_stats = torch.randn(256, dtype=torch.float32, device=device) |
| 69 | col_stats = torch.randn(256, dtype=torch.float32, device=device) |
nothing calls this directly
no outgoing calls
no test coverage detected