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