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

@@ -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
 

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):

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 = {

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(

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
 
 

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,
 )
-
+from tests.conftest import skip_no_fp8
-_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+",
-)
 
 
 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)

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

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(
+pytestmark = skip_no_cuda
-    torch.cuda.device_count() < 1, reason="CUDA device required"
-)
 
 
 def _assign_grads_and_step(model):