fix: broadcast ref/old model state_dict for FSDP

- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks
- Fix create_ref_model returning None on non-rank-0 under FSDP
- Fix sync_old_model only updating old_model on rank-0 under FSDP
- Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest
- Add distributed tests for broadcast_state_dict and create_ref_model

tests/extension/conftest.py

@@ -4,13 +4,8 @@ import pytest
 import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.extension import is_available
 from astrai.model.transformer import AutoRegressiveLM
-
-CUDA_AVAILABLE = torch.cuda.is_available() and is_available("attn_paged_decode")
-skip_no_cuda = pytest.mark.skipif(
-    not CUDA_AVAILABLE, reason="CUDA not available or kernels not built"
-)
+from tests.conftest import skip_no_kernel
 
 D = 64
 CFG = dict(

tests/extension/test_backend_equivalence.py

@@ -8,10 +8,10 @@ import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
 from astrai.inference.core.cache import PagePool
-from tests.extension.conftest import D, skip_no_cuda
+from tests.extension.conftest import D, skip_no_kernel
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
     """Training forward (kv_cache=None) should produce identical logits."""
     model, _ = cuda_model
@@ -27,7 +27,7 @@ def test_training_forward_matches_torch(cuda_model):
     assert diff == 0.0, f"Training forward diff {diff} should be 0"
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_prefill_with_kv_cache_matches_torch(cuda_model):
     """Inference prefill with KV cache should match torch backend."""
     model, _ = cuda_model
@@ -93,7 +93,7 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
         assert d == 0.0, f"Prefill diff for sample {i}: {d}"
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     """Decode with mixed seq_lens in batch — padding mask must produce correct output."""
     model, _ = cuda_model
@@ -152,7 +152,7 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_run_batch_cuda_matches_torch_greedy(cuda_model):
     """Greedy decode (temperature=0) should produce identical tokens."""
     from astrai.inference.core.scheduler import InferenceScheduler

tests/extension/test_kernel_mask.py

@@ -2,10 +2,10 @@
 
 import torch
 
-from tests.extension.conftest import D, skip_no_cuda
+from tests.extension.conftest import D, skip_no_kernel
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_kernel_accepts_2d_mask():
     """Kernel should accept 2D mask [batch, kv_len]."""
     from astrai.extension.attention_ops import attn_prefill
@@ -22,7 +22,7 @@ def test_kernel_accepts_2d_mask():
     assert out.shape == (batch, q_len, n_heads, D)
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_kernel_accepts_3d_mask():
     """Kernel should accept 3D mask [batch, q_len, kv_len]."""
     from astrai.extension.attention_ops import attn_prefill
@@ -38,7 +38,7 @@ def test_kernel_accepts_3d_mask():
     assert out.shape == (batch, q_len, n_heads, D)
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_kernel_accepts_4d_mask():
     """Kernel should accept 4D mask [batch, n_heads, q_len, kv_len]."""
     from astrai.extension.attention_ops import attn_prefill
@@ -55,7 +55,7 @@ def test_kernel_accepts_4d_mask():
     assert out.shape == (batch, q_len, n_heads, D)
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_4d_mask_matches_no_mask_when_all_true():
     """A 4D all-True mask should produce the same output as no mask."""
     from astrai.extension.attention_ops import attn_prefill