# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. from __future__ import annotations import pytest import torch from tokenspeed_kernel.numerics.comparison import compare_outputs, format_comparison from tokenspeed_kernel.numerics.inputs import get_input_generator from tokenspeed_kernel.numerics.tolerance import Tolerance from tokenspeed_kernel.numerics.verify import ( _verification_signature_and_reference, verify_kernel, ) from tokenspeed_kernel.platform import Platform from tokenspeed_kernel.registry import KernelRegistry, KernelSpec, load_builtin_kernels from tokenspeed_kernel.signature import ScaleFormat, format_signatures _fp8_dtype = Platform.get().fp8e4m3fn.dtype class TestCompareOutputs: def test_treats_nan_as_mismatch(self) -> None: actual = torch.tensor([1.0, float("nan")], dtype=torch.float32) expected = torch.tensor([1.0, 1.0], dtype=torch.float32) result = compare_outputs( actual, expected, tolerance=Tolerance(atol=1e6, rtol=1e6), ) assert not result.passed assert result.num_mismatches == 1 def test_treats_inf_as_mismatch(self) -> None: actual = torch.tensor([float("inf"), 1.0], dtype=torch.float32) expected = torch.tensor([float("inf"), 1.0], dtype=torch.float32) result = compare_outputs( actual, expected, tolerance=Tolerance(atol=1e6, rtol=1e6), ) assert not result.passed assert result.num_mismatches == 1 def test_gemm_input_generator_uses_signature_scale_metadata() -> None: scale = ScaleFormat( storage_dtype=torch.float32, granularity="block", block_shape=(128, 128), ) signature = next( iter(format_signatures(("a", "b"), "mxfp8", {_fp8_dtype}, scale=scale)) ) generator = get_input_generator( "gemm", "mm", dtype=_fp8_dtype, traits={}, format_signature=signature, device="cpu", ) inputs = generator.generate(M=4, N=256, K=128) assert inputs["A"].dtype == _fp8_dtype assert inputs["B"].dtype == _fp8_dtype assert inputs["A_scales"].shape == (4, 1) assert inputs["B_scales"].shape == (2, 1) assert inputs["A_scales"].dtype == torch.float32 assert inputs["B_scales"].dtype == torch.float32 assert inputs["block_size"] == [128, 128] def test_gemm_input_generator_requires_mxfp8_block_shape() -> None: scale = ScaleFormat( storage_dtype=torch.float32, granularity="block", dynamic_block_shape=True, ) signature = next( iter(format_signatures(("a", "b"), "mxfp8", {_fp8_dtype}, scale=scale)) ) generator = get_input_generator( "gemm", "mm", dtype=_fp8_dtype, traits={}, format_signature=signature, device="cpu", ) with pytest.raises(ValueError, match="requires concrete block_shape"): generator.generate(M=4, N=256, K=128) def test_verification_uses_signature_with_compatible_reference(fresh_registry) -> None: tensor_scale = ScaleFormat(storage_dtype=torch.float32, granularity="tensor") channel_scale = ScaleFormat(storage_dtype=torch.float32, granularity="channel") tensor_signature = next( iter( format_signatures( ("a", "b"), "scaled-fp8", {_fp8_dtype}, scale=tensor_scale ) ) ) channel_signature = next( iter( format_signatures( ("a", "b"), "scaled-fp8", {_fp8_dtype}, scale=channel_scale ) ) ) ref_spec = KernelSpec( name="test_tensor_scale_reference", family="gemm", mode="mm", solution="reference", format_signatures=frozenset({tensor_signature}), traits={"b_layout": frozenset({"KN"})}, ) test_spec = KernelSpec( name="test_fp8_scaled", family="gemm", mode="mm", solution="triton", format_signatures=frozenset({channel_signature, tensor_signature}), traits={"b_layout": frozenset({"KN"})}, ) registry = KernelRegistry.get() registry.register(ref_spec, lambda **_kwargs: None) registry.register(test_spec, lambda **_kwargs: None) signature, reference = _verification_signature_and_reference( registry, test_spec, _fp8_dtype, "a" ) assert signature == tensor_signature assert reference is ref_spec class TestNumericsVerification: def _get_verifiable_specs( dtype: torch.dtype, dtype_role: str, family: str | None = None ) -> list[KernelSpec]: load_builtin_kernels() registry = KernelRegistry.get() platform = Platform.get() specs: list[KernelSpec] = [] for family_name, mode in registry.list_operators(): if family and family_name != family: continue # Only run kernels that have a paired reference for this dtype; # otherwise verify_kernel raises ValueError and the test errors. op_specs = registry.get_for_operator(family_name, mode) dtype_specs = [ s for s in op_specs if s.format_signatures_for_storage_dtype(dtype, dtype_role) ] has_reference = any(s.solution == "reference" for s in dtype_specs) if not has_reference: continue for spec in dtype_specs: if spec.solution == "reference": continue if spec.solution == "deep_gemm": continue if not spec.capability.satisfied_by(platform): continue specs.append(spec) specs.sort(key=lambda s: (s.family, s.mode, s.name)) return specs def _verify(self, spec: KernelSpec, dtype: torch.dtype, dtype_role: str) -> None: if not torch.cuda.is_available(): pytest.skip("CUDA is required for numerics verification") try: results = verify_kernel( spec.name, dtype=dtype, dtype_role=dtype_role, verbose=False ) except Exception as exc: pytest.fail( f"Kernel {spec.name} raised an exception during verification: {exc}" ) for i, result in enumerate(results): if not result.passed: pytest.fail( f"Kernel {spec.name} failed numerics verification for shape set {i}:\n" f"{format_comparison(result, kernel_name=spec.name)}" ) @pytest.mark.parametrize( "spec", _get_verifiable_specs(_fp8_dtype, "a", family="gemm"), ids=lambda s: f"{s.family}.{s.mode}:{s.name}", ) def test_gemm_fp8(self, spec: KernelSpec): self._verify(spec, _fp8_dtype, "a") @pytest.mark.parametrize( "spec", _get_verifiable_specs(torch.bfloat16, "a", family="gemm"), ids=lambda s: f"{s.family}.{s.mode}:{s.name}", ) def test_gemm_bf16(self, spec: KernelSpec): self._verify(spec, torch.bfloat16, "a") @pytest.mark.parametrize( "spec", _get_verifiable_specs(torch.bfloat16, "x", family="quantize"), ids=lambda s: f"{s.family}.{s.mode}:{s.name}", ) def test_quantize_bf16(self, spec: KernelSpec): self._verify(spec, torch.bfloat16, "x") @pytest.mark.parametrize( "spec", _get_verifiable_specs(torch.int32, "indices", family="moe"), ids=lambda s: f"{s.family}.{s.mode}:{s.name}", ) def test_moe_int32(self, spec: KernelSpec): self._verify(spec, torch.int32, "indices")