From 75304d084df776f0f246887732a69413fa5dc9ba Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 22 Aug 2026 21:08:06 +0800 Subject: [PATCH] refactor: unify CUDA skip guards in tests - add skip_no_fp8 (CUDA + fp8_mm kernel + cc 8.9+) to tests/conftest.py - use skip_no_cuda / skip_no_kernel / skip_no_fp8 directly in test modules - drop _GPU alias and tests.extension.conftest re-exports - remove unused imports (Union in hf_adapter, make_grpo_config in data conftest) --- astrai/serialization/hf_adapter.py | 2 +- tests/conftest.py | 9 ++++++++ tests/data/conftest.py | 1 - tests/extension/conftest.py | 1 - tests/extension/test_backend_equivalence.py | 3 ++- tests/extension/test_fp8_mma.py | 24 ++++++++------------- tests/extension/test_kernel_mask.py | 3 ++- tests/optim/test_optimizer_distributed.py | 6 ++---- 8 files changed, 25 insertions(+), 24 deletions(-) diff --git a/astrai/serialization/hf_adapter.py b/astrai/serialization/hf_adapter.py index aa82cf9..3fc2481 100644 --- a/astrai/serialization/hf_adapter.py +++ b/astrai/serialization/hf_adapter.py @@ -21,7 +21,7 @@ Not supported: import logging import re -from typing import Any, Dict, Mapping, Union +from typing import Any, Dict, Mapping import torch diff --git a/tests/conftest.py b/tests/conftest.py index 12b0c21..e6c6348 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,8 +15,17 @@ from tests.helpers import ( CUDA_AVAIL = torch.cuda.is_available() KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES) +FP8_AVAIL = ( + CUDA_AVAIL + and is_available("fp8_mm") + and torch.cuda.get_device_capability() >= (8, 9) +) skip_no_cuda = pytest.mark.skipif(not CUDA_AVAIL, reason="CUDA not available") skip_no_kernel = pytest.mark.skipif(not KERNEL_AVAIL, reason="CUDA kernels not built") +skip_no_fp8 = pytest.mark.skipif( + not FP8_AVAIL, + reason="fused FP8 MMA requires a built kernel and compute capability 8.9+", +) def pytest_configure(config): diff --git a/tests/data/conftest.py b/tests/data/conftest.py index a8ab651..60c7463 100644 --- a/tests/data/conftest.py +++ b/tests/data/conftest.py @@ -8,7 +8,6 @@ from astrai.preprocessing.builder import ( SectionedMaskBuilder, SingleOutputMaskBuilder, ) -from tests.data.factories import make_grpo_config from tests.helpers import build_test_tokenizer _SPECIAL_TOKENS_CONFIG = { diff --git a/tests/extension/conftest.py b/tests/extension/conftest.py index e0c18ac..8e3e373 100644 --- a/tests/extension/conftest.py +++ b/tests/extension/conftest.py @@ -5,7 +5,6 @@ import torch from astrai.config.model_config import AutoRegressiveLMConfig from astrai.model.transformer import AutoRegressiveLM -from tests.conftest import skip_no_kernel # noqa: F401 re-export for test modules D = 64 CFG = dict( diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index baeede1..73b383c 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/tests/extension/test_backend_equivalence.py @@ -12,7 +12,8 @@ from astrai.inference.cache import PagePool, TaskCacheManager from astrai.inference.runtime.graph import CudaGraphContext from astrai.inference.scheduler import InferenceScheduler from astrai.inference.workspace import InferenceWorkspace -from tests.extension.conftest import D, skip_no_kernel +from tests.conftest import skip_no_kernel +from tests.extension.conftest import D from tests.helpers import FakeTokenizer diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index 511700b..e172280 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -16,20 +16,14 @@ from astrai.extension.fp8 import ( fp8_autocast, fp8_state, ) -from astrai.extension.loader import get_module, is_available +from astrai.extension.loader import get_module from astrai.extension.ops.fp8 import ( linear_backward_fp8, linear_forward_fp8, mm_fp8, quantize_bf16, ) - -_GPU = pytest.mark.skipif( - not torch.cuda.is_available() - or torch.cuda.get_device_capability() < (8, 9) - or not is_available("fp8_mm"), - reason="fused FP8 MMA requires a built kernel and compute capability 8.9+", -) +from tests.conftest import skip_no_fp8 def _scale(tensor): @@ -45,7 +39,7 @@ def _quantize(tensor, scale): # -------------------------------------------------------------------------- -@_GPU +@skip_no_fp8 @pytest.mark.parametrize( ("m", "n", "k"), [(16, 8, 32), (17, 9, 33), (31, 15, 64), (32, 48, 96)], @@ -67,7 +61,7 @@ def test_fused_fp8_mma_matches_explicit_quantization(m, n, k): torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01) -@_GPU +@skip_no_fp8 def test_quantize_bf16_returns_amax(): """quantize_bf16 returns (x8, amax); amax tracks the *raw* values and the caller never clears it (zero-initialized inside the kernel entry).""" @@ -83,7 +77,7 @@ def test_quantize_bf16_returns_amax(): assert torch.equal(x8, ref) -@_GPU +@skip_no_fp8 def test_quantize_bf16_e5m2_format(): x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) x8, amax = quantize_bf16(x, torch.tensor([0.1], device="cuda"), "e5m2") @@ -91,7 +85,7 @@ def test_quantize_bf16_e5m2_format(): torch.testing.assert_close(amax, x.abs().amax().float().reshape(1)) -@_GPU +@skip_no_fp8 def test_fused_fp8_linear_forward_and_backward(): torch.manual_seed(7) m, n, k = 19, 13, 37 @@ -122,7 +116,7 @@ def test_fused_fp8_linear_forward_and_backward(): torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1)) -@_GPU +@skip_no_fp8 def test_linear_backward_e5m2_gradients(): """Hybrid backward: gradient GEMMs run in E5M2 (larger dynamic range).""" torch.manual_seed(5) @@ -151,7 +145,7 @@ def test_linear_backward_e5m2_gradients(): torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1)) -@_GPU +@skip_no_fp8 def test_mm_fp8_matches_scaled_mm(): torch.manual_seed(11) m, n, k = 512, 4096, 4096 @@ -180,7 +174,7 @@ def test_mm_fp8_matches_scaled_mm(): ) -@_GPU +@skip_no_fp8 def test_mm_fp8_fp8_output(): """mm_fp8 with out_dtype='e4m3' produces an FP8 output (layer-to-layer).""" torch.manual_seed(12) diff --git a/tests/extension/test_kernel_mask.py b/tests/extension/test_kernel_mask.py index ab1ed76..1b29d31 100644 --- a/tests/extension/test_kernel_mask.py +++ b/tests/extension/test_kernel_mask.py @@ -3,7 +3,8 @@ import torch from astrai.extension.ops.attention import attn_prefill -from tests.extension.conftest import D, skip_no_kernel +from tests.conftest import skip_no_kernel +from tests.extension.conftest import D @skip_no_kernel diff --git a/tests/optim/test_optimizer_distributed.py b/tests/optim/test_optimizer_distributed.py index d7ae652..9fd1b84 100644 --- a/tests/optim/test_optimizer_distributed.py +++ b/tests/optim/test_optimizer_distributed.py @@ -1,4 +1,3 @@ -import pytest import torch import torch.distributed as dist from torch.distributed.fsdp import fully_shard @@ -8,11 +7,10 @@ from torch.nn.parallel import DistributedDataParallel as DDP from astrai.model import AutoRegressiveLM from astrai.optim import NoraNAdamW from astrai.parallel.setup import find_free_port +from tests.conftest import skip_no_cuda from tests.helpers import make_tiny_config -pytestmark = pytest.mark.skipif( - torch.cuda.device_count() < 1, reason="CUDA device required" -) +pytestmark = skip_no_cuda def _assign_grads_and_step(model):