| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0 |
| 8 | |
| 9 | |
| 10 | def test_wo_a_block_scale_to_e8m0_from_float(): |
| 11 | scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32) |
| 12 | |
| 13 | encoded = _wo_a_block_scale_to_e8m0(scale) |
| 14 | |
| 15 | assert encoded is not None |
| 16 | torch.testing.assert_close( |
| 17 | encoded, |
| 18 | torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8), |
| 19 | ) |
| 20 | assert encoded.is_contiguous() |
| 21 | |
| 22 | |
| 23 | def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales(): |
| 24 | raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8) |
| 25 | encoded = raw.view(torch.float8_e8m0fnu) |
| 26 | |
| 27 | converted = _wo_a_block_scale_to_e8m0(encoded) |
| 28 | |
| 29 | assert converted is not None |
| 30 | torch.testing.assert_close(converted, raw) |
| 31 | |
| 32 | |
| 33 | @pytest.mark.parametrize( |
| 34 | "scale", |
| 35 | [ |
| 36 | torch.tensor([[0.0, 1.0]]), |
| 37 | torch.tensor([[-1.0, 1.0]]), |
| 38 | torch.tensor([[0.75, 1.0]]), |
| 39 | torch.tensor([[float("inf"), 1.0]]), |
| 40 | torch.ones(1, dtype=torch.int32), |
| 41 | ], |
| 42 | ) |
| 43 | def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor): |
| 44 | assert _wo_a_block_scale_to_e8m0(scale) is None |