style: fix ruff lint warnings

- Remove unused local variable b in attention_backend.py
- Remove unused variable rank0_sd in test_broadcast_state_dict.py
- Remove unused imports across test files

tests/data/test_dataset.py

@@ -6,7 +6,6 @@ import numpy as np
 import pytest
 import torch
 
-from astrai.config.preprocess_config import PipelineConfig
 from astrai.dataset.dataset import (
     DatasetFactory,
     GRPODataset,

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
 
 D = 64
 CFG = dict(

tests/extension/test_backend.py

@@ -10,7 +10,6 @@ from astrai.extension import (
     ATTN_BACKEND,
     AttentionBackendFactory,
     CudaBackend,
-    TorchNativeBackend,
     attn_backend,
     get_backend,
 )

tests/extension/test_backend_equivalence.py

@@ -31,7 +31,6 @@ def test_training_forward_matches_torch(cuda_model):
     or non-bf16 inputs it falls back to torch SDPA.  Verify the fallback
     path matches the torch-native forward exactly.
     """
-    import pytest
 
     model, _ = cuda_model
     input_ids = torch.randint(0, 1000, (2, 16), device="cuda")

tests/parallel/test_broadcast_state_dict.py

@@ -5,14 +5,12 @@ multi-rank environment without requiring multiple GPUs.
 """
 
 import torch
-import torch.distributed as dist
-import torch.nn as nn
 
 from astrai.model.transformer import AutoRegressiveLM
 from astrai.parallel import get_rank, spawn_parallel_fn
 from astrai.parallel.executor import broadcast_state_dict, create_ref_model
 from astrai.trainer.strategy import GRPOStrategy
-from tests.helpers import FakeExecutor, make_rollout_config
+from tests.helpers import make_rollout_config
 
 
 def _broadcast_worker():
@@ -77,8 +75,6 @@ def _create_ref_model_worker():
     )
 
     assert ref is not None, f"rank {rank}: ref model is None"
-    # Every rank should have rank-0's weights, not its own
-    rank0_sd = model.state_dict() if rank == 0 else None
     # Broadcast rank-0's original weights for comparison
     if rank == 0:
         expected_sd = {k: v.clone() for k, v in model.state_dict().items()}