"""Qualification harness: load kernel from Nix store bundle, assert provenance, run boundary matrices.""" from __future__ import annotations import json import sys from datetime import UTC, datetime from pathlib import Path import pytest RESULT_DIR = Path(__file__).resolve().parents[1] / "result" @pytest.fixture(scope="module") def ssd_bundle(): """Load the kernel from the Nix store result/ symlink.""" kernels = pytest.importorskip("kernels") assert RESULT_DIR.is_dir(), f"result/ not found at {RESULT_DIR}" ssd = kernels.get_local_kernel(RESULT_DIR.resolve(), backend="rocm") return ssd @pytest.fixture(scope="module") def mamba_combined(ssd_bundle): """Extract the combined SSD function from the bundle.""" return ssd_bundle.mamba_chunk_scan_combined class TestNixStoreProvenance: """Fail-hard: kernel must come from the immutable Nix store, not editable source.""" def test_module_file_is_in_nix_store(self, ssd_bundle): module_path = Path(ssd_bundle.__file__).resolve() assert "/nix/store/" in str(module_path), ( f"Kernel loaded from editable source, not Nix store: {module_path}" ) def test_module_parent_is_torch_rocm(self, ssd_bundle): module_path = Path(ssd_bundle.__file__).resolve() assert module_path.parent.name == "torch-rocm", ( f"Expected torch-rocm directory, got: {module_path.parent.name}" ) def test_module_name_has_rocm_suffix(self, ssd_bundle): name = ssd_bundle.__name__ assert "rocm" in name.lower(), f"Module name missing rocm suffix: {name}" def test_not_from_torch_ext_editable(self, ssd_bundle): module_path = Path(ssd_bundle.__file__).resolve() path_str = str(module_path) # Exclude the Nix store derivation name ("torch-ext-bundle") — # only flag if torch-ext appears as a source tree path component. assert "/nix/store/" in path_str or "torch-ext/" not in path_str, ( f"Kernel loaded from editable torch-ext, not Nix bundle: {module_path}" ) class TestCombinedScanAPIBoundary: """Verify the combined SSD API signature and basic forward pass.""" def test_combined_scan_signature(self, mamba_combined): import inspect params = tuple(inspect.signature(mamba_combined).parameters) expected = ( "x", "dt", "A", "B", "C", "chunk_size", "D", "z", "dt_bias", "initial_states", "seq_idx", "cu_seqlens", "dt_softplus", "dt_limit", "return_final_states", "return_varlen_states", "state_dtype", ) assert params == expected, f"Signature mismatch: {params}" def test_combined_scan_forward_smoke(self, mamba_combined): """Minimal forward pass on CPU-sized inputs (should raise on CPU).""" torch = pytest.importorskip("torch") x = torch.zeros((1, 8, 2, 4), dtype=torch.float32) dt = torch.zeros((1, 8, 2), dtype=torch.float32) A = -torch.ones((2,), dtype=torch.float32) B = torch.zeros((1, 8, 1, 4), dtype=torch.float32) C = torch.zeros_like(B) with pytest.raises((AssertionError, RuntimeError, ValueError), match="(?i)cuda|hip|gpu"): mamba_combined(x, dt, A, B, C, 4) class TestGfx1030BoundaryMatrices: """Run the same boundary shapes used in editable-source tests against the Nix bundle.""" @pytest.fixture(autouse=True) def _require_gpu(self): torch = pytest.importorskip("torch") if not torch.cuda.is_available(): pytest.skip("No GPU available for gfx1030 boundary test") def test_small_tile_64x32x32(self, mamba_combined): torch = pytest.importorskip("torch") B, L, H, D = 1, 8, 2, 4 chunk_size = 4 x = torch.randn(B, L, H, D, device="cuda", dtype=torch.float32) dt = torch.rand(B, L, H, device="cuda", dtype=torch.float32).abs() + 0.01 A = -torch.rand(H, device="cuda", dtype=torch.float32).abs() - 0.01 B_mat = torch.randn(B, L, 1, D, device="cuda", dtype=torch.float32) C_mat = torch.randn(B, L, 1, D, device="cuda", dtype=torch.float32) out = mamba_combined(x, dt, A, B_mat, C_mat, chunk_size) assert out.shape == x.shape, f"Output shape mismatch: {out.shape} != {x.shape}" assert not torch.isnan(out).any(), "NaN in output" assert not torch.isinf(out).any(), "Inf in output" def test_varlen_single_sequence(self, mamba_combined): torch = pytest.importorskip("torch") B, L, H, D = 1, 16, 2, 4 chunk_size = 4 x = torch.randn(B, L, H, D, device="cuda", dtype=torch.float32) dt = torch.rand(B, L, H, device="cuda", dtype=torch.float32).abs() + 0.01 A = -torch.rand(H, device="cuda", dtype=torch.float32).abs() - 0.01 B_mat = torch.randn(B, L, 1, D, device="cuda", dtype=torch.float32) C_mat = torch.randn(B, L, 1, D, device="cuda", dtype=torch.float32) cu_seqlens = torch.tensor([0, L], device="cuda", dtype=torch.int32) out, varlen_states = mamba_combined( x, dt, A, B_mat, C_mat, chunk_size, cu_seqlens=cu_seqlens, return_varlen_states=True, ) assert out.shape == x.shape, f"Output shape mismatch: {out.shape} != {x.shape}" assert not torch.isnan(out).any(), "NaN in output" assert varlen_states is not None, "varlen_states should not be None" def test_large_state_dim(self, mamba_combined): torch = pytest.importorskip("torch") B, L, H, D = 1, 32, 2, 16 chunk_size = 8 x = torch.randn(B, L, H, D, device="cuda", dtype=torch.float32) dt = torch.rand(B, L, H, device="cuda", dtype=torch.float32).abs() + 0.01 A = -torch.rand(H, device="cuda", dtype=torch.float32).abs() - 0.01 B_mat = torch.randn(B, L, 1, D, device="cuda", dtype=torch.float32) C_mat = torch.randn(B, L, 1, D, device="cuda", dtype=torch.float32) out = mamba_combined(x, dt, A, B_mat, C_mat, chunk_size) assert out.shape == x.shape assert not torch.isnan(out).any() def _record_qualification_artifact(ssd_bundle, result_dir: Path) -> Path: """Write qualification artifact JSON with provenance metadata.""" module_path = Path(ssd_bundle.__file__).resolve() module_name = type(ssd_bundle).__module__ artifact = { "kernel_source": "nix-local-bundle", "kernel_module": module_name, "kernel_path": str(module_path), "backend": "rocm", "triton_target": "hip:gfx1030", "result_dir": str(result_dir.resolve()), "timestamp": datetime.now(UTC).isoformat(), "python": sys.version, } out_path = result_dir / "qualification_artifact.json" out_path.write_text(json.dumps(artifact, indent=2) + "\n") return out_path