123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960 |
- import pytest
- import torch
- from aphrodite._custom_ops import scaled_int8_quant
- from tests.kernels.quant_utils import ref_dynamic_per_token_quant
- DTYPES = [torch.half, torch.bfloat16, torch.float]
- HIDDEN_SIZES = [16, 67, 768, 2048, 5120, 5137, 8192,
- 8193] # Arbitrary values for testing
- NUM_TOKENS = [1, 7, 83, 4096] # Arbitrary values for testing
- SEEDS = [0]
- SCALE = [0.1, 0.5, 0.8, 1.2, 2.1]
- @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
- @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
- @pytest.mark.parametrize("dtype", DTYPES)
- @pytest.mark.parametrize("seed", SEEDS)
- @torch.inference_mode()
- def test_dynamic_scaled_int8_quant(num_tokens: int, hidden_size: int,
- dtype: torch.dtype, seed: int) -> None:
- torch.random.manual_seed(seed)
- torch.cuda.manual_seed(seed)
- x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000
- # reference
- ref_out, ref_scales = ref_dynamic_per_token_quant(x, torch.int8)
- # kernel
- ops_out, ops_scales = scaled_int8_quant(x)
- torch.testing.assert_close(ops_scales, ref_scales)
- torch.testing.assert_close(
- ops_out, ref_out, atol=1,
- rtol=0.0) # big atol to account for rounding errors
- @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
- @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
- @pytest.mark.parametrize("dtype", DTYPES)
- @pytest.mark.parametrize("seed", SEEDS)
- @pytest.mark.parametrize("scale", SCALE)
- @torch.inference_mode()
- def test_static_scaled_int8_quant(num_tokens: int, hidden_size: int,
- dtype: torch.dtype, seed: int,
- scale: float) -> None:
- torch.random.manual_seed(seed)
- torch.cuda.manual_seed(seed)
- int8_traits = torch.iinfo(torch.int8)
- x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000
- scale = torch.tensor([scale], dtype=torch.float32, device="cuda")
- out1 = (x / scale).round().clamp(int8_traits.min,
- int8_traits.max).to(torch.int8)
- out2, _ = scaled_int8_quant(x, scale)
- torch.testing.assert_close(
- out1, out2, atol=1,
- rtol=0.0) # big atol to account for rounding errors
|