forked from ccf-ai-infra/TileOPs-Metax
53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from tests.test_base import FixtureBase, TestBase
|
|
from tileops.ops import FP8QuantOp
|
|
from workloads.fp8_quant import FP8QuantTest as _FP8QuantTestWorkload
|
|
|
|
|
|
class FP8QuantTest(_FP8QuantTestWorkload, TestBase):
|
|
def ref_program(self, input_tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
# input_tensor: (batch, seq_len_kv, kv_group, index_dim)
|
|
amax_value = torch.abs(input_tensor).amax(dim=-1, keepdim=True).clamp(min=1e-4)
|
|
scale_tensor = amax_value / 448.0
|
|
output_tensor = torch.clamp(input_tensor / scale_tensor, min=-448.0, max=448.0)
|
|
output_tensor = output_tensor.to(torch.float8_e4m3fn)
|
|
return scale_tensor.squeeze(dim=-1), output_tensor
|
|
|
|
|
|
class FP8QuantFixture(FixtureBase):
|
|
PARAMS = [
|
|
("batch, seq_len_kv, kv_group, index_dim, in_dtype, tune", [
|
|
pytest.param(1, 8192, 1, 64, torch.float16, False, marks=pytest.mark.smoke),
|
|
pytest.param(1, 8192, 1, 64, torch.bfloat16, False, marks=pytest.mark.full),
|
|
pytest.param(1, 4096, 1, 128, torch.float32, False, marks=pytest.mark.full),
|
|
pytest.param(1, 16384, 1, 32, torch.float32, False, marks=pytest.mark.full),
|
|
pytest.param(1, 1024, 4, 64, torch.float16, False, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
def _cosine_compare(output: torch.Tensor, output_ref: torch.Tensor) -> None:
|
|
"""Compare using cosine similarity (fp8 quantization needs looser checks)."""
|
|
output = output.to(torch.float32)
|
|
output_ref = output_ref.to(torch.float32)
|
|
cos_sim = F.cosine_similarity(output.flatten(), output_ref.flatten(), dim=0)
|
|
assert cos_sim >= 0.99, \
|
|
f"Cosine similarity too low: {cos_sim.item()}"
|
|
|
|
|
|
@FP8QuantFixture
|
|
def test_fp8_quant_op(batch: int, seq_len_kv: int, kv_group: int, index_dim: int,
|
|
in_dtype: torch.dtype, tune: bool) -> None:
|
|
test = FP8QuantTest(batch, seq_len_kv, kv_group, index_dim, in_dtype)
|
|
op = FP8QuantOp(tune=tune)
|
|
test.check(op, *test.gen_inputs(), compare=_cosine_compare)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-vvs"])
|