File size: 8,682 Bytes
8996239 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """Qualification harness: load causal-conv1d kernel from Nix store bundle, assert provenance, run boundary tests."""
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 cc_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}"
cc = kernels.get_local_kernel(RESULT_DIR.resolve(), backend="rocm")
return cc
@pytest.fixture(scope="module")
def causal_conv1d_fn(cc_bundle):
"""Extract the causal_conv1d_fn from the bundle."""
return cc_bundle.causal_conv1d_fn
@pytest.fixture(scope="module")
def causal_conv1d_update(cc_bundle):
"""Extract the causal_conv1d_update from the bundle."""
return cc_bundle.causal_conv1d_update
class TestNixStoreProvenance:
"""Fail-hard: kernel must come from the immutable Nix store, not editable source."""
def test_module_file_is_in_nix_store(self, cc_bundle):
module_path = Path(cc_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, cc_bundle):
module_path = Path(cc_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, cc_bundle):
name = cc_bundle.__name__
assert "rocm" in name.lower(), f"Module name missing rocm suffix: {name}"
def test_not_from_torch_ext_editable(self, cc_bundle):
module_path = Path(cc_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 TestCausalConv1dAPIBoundary:
"""Verify the causal_conv1d_fn and causal_conv1d_update API signatures."""
def test_causal_conv1d_fn_signature(self, causal_conv1d_fn):
import inspect
params = tuple(inspect.signature(causal_conv1d_fn).parameters)
expected = (
"x", "weight", "bias", "seq_idx", "initial_states",
"return_final_states", "final_states_out", "activation",
)
assert params == expected, f"Signature mismatch: {params}"
def test_causal_conv1d_update_signature(self, causal_conv1d_update):
import inspect
params = tuple(inspect.signature(causal_conv1d_update).parameters)
expected = (
"x", "conv_state", "weight", "bias", "activation",
"cache_seqlens", "conv_state_indices",
)
assert params == expected, f"Signature mismatch: {params}"
def test_causal_conv1d_fn_rejects_cpu(self, causal_conv1d_fn):
"""Triton kernel must reject CPU tensors — no silent fallback in the bundle."""
torch = pytest.importorskip("torch")
x = torch.zeros((1, 4, 8), dtype=torch.float32)
weight = torch.zeros((4, 4), dtype=torch.float32)
bias = torch.zeros((4,), dtype=torch.float32)
# The Nix bundle's functional.py routes through Triton directly.
# CPU tensors must fail — no silent CPU fallback in the qualified bundle.
with pytest.raises((ValueError, RuntimeError, AssertionError)):
causal_conv1d_fn(x, weight, bias, activation="silu")
class TestGfx1030BoundaryMatrices:
"""Run boundary shapes against the Nix bundle on GPU."""
@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_fp32(self, causal_conv1d_fn):
torch = pytest.importorskip("torch")
b, d, seq_len, width = 1, 4, 8, 4
x = torch.randn(b, d, seq_len, device="cuda", dtype=torch.float32)
weight = torch.randn(d, width, device="cuda", dtype=torch.float32)
bias = torch.randn(d, device="cuda", dtype=torch.float32)
out = causal_conv1d_fn(x, weight, bias, activation="silu")
assert out.shape == x.shape
assert not torch.isnan(out).any(), "NaN in output"
assert not torch.isinf(out).any(), "Inf in output"
def test_small_fp16(self, causal_conv1d_fn):
torch = pytest.importorskip("torch")
b, d, seq_len, width = 1, 16, 32, 4
x = torch.randn(b, d, seq_len, device="cuda", dtype=torch.float16)
weight = torch.randn(d, width, device="cuda", dtype=torch.float16)
bias = torch.randn(d, device="cuda", dtype=torch.float16)
out = causal_conv1d_fn(x, weight, bias, activation="silu")
assert out.shape == x.shape
assert out.dtype == torch.float16
assert not torch.isnan(out).any(), "NaN in output"
def test_initial_states(self, causal_conv1d_fn):
torch = pytest.importorskip("torch")
b, d, seq_len, width = 2, 8, 16, 4
x = torch.randn(b, d, seq_len, device="cuda")
weight = torch.randn(d, width, device="cuda")
bias = torch.randn(d, device="cuda")
init = torch.randn(b, d, width - 1, device="cuda")
out, final = causal_conv1d_fn(
x, weight, bias, initial_states=init, return_final_states=True
)
# The Triton kernel prepends initial_states to x before convolution,
# so the output length is seq_len + (width - 1). The final_states
# are the last (width - 1) elements of the original x (not the
# prepended version).
expected_len = seq_len + (width - 1)
assert out.shape == (b, d, expected_len), f"Expected {(b, d, expected_len)}, got {out.shape}"
assert not torch.isnan(out).any()
assert final.shape == (b, d, width - 1)
def test_update_matches_prefill(self, causal_conv1d_fn, causal_conv1d_update):
"""Decode-step update must equal the prefill output for the last token."""
torch = pytest.importorskip("torch")
b, d, seq_len, width = 1, 8, 6, 4
x = torch.randn(b, d, seq_len, device="cuda")
weight = torch.randn(d, width, device="cuda")
bias = torch.randn(d, device="cuda")
prefill = causal_conv1d_fn(x, weight, bias, activation="silu")
last_token = prefill[:, :, -1]
conv_state = x[:, :, : width - 1]
for t in range(width - 1, seq_len):
token = x[:, :, t]
out = causal_conv1d_update(
token, conv_state, weight, bias, activation="silu"
)
torch.testing.assert_close(out, prefill[:, :, t], atol=1e-4, rtol=1e-4)
conv_state = torch.cat([conv_state[:, :, 1:], token.unsqueeze(-1)], dim=-1)
torch.testing.assert_close(out, last_token, atol=1e-4, rtol=1e-4)
def test_backward_finite_gradients(self, causal_conv1d_fn):
"""Autograd backward produces finite gradients."""
torch = pytest.importorskip("torch")
b, d, seq_len, width = 1, 8, 16, 4
x = torch.randn(b, d, seq_len, device="cuda", dtype=torch.float32, requires_grad=True)
weight = torch.randn(d, width, device="cuda", dtype=torch.float32, requires_grad=True)
bias = torch.randn(d, device="cuda", dtype=torch.float32, requires_grad=True)
out = causal_conv1d_fn(x, weight, bias, activation="silu")
loss = out.float().square().mean()
loss.backward()
torch.cuda.synchronize()
assert torch.isfinite(x.grad).all(), "dx contains non-finite values"
assert torch.isfinite(weight.grad).all(), "dweight contains non-finite values"
assert torch.isfinite(bias.grad).all(), "dbias contains non-finite values"
def _record_qualification_artifact(cc_bundle, result_dir: Path) -> Path:
"""Write qualification artifact JSON with provenance metadata."""
module_path = Path(cc_bundle.__file__).resolve()
artifact = {
"kernel_source": "nix-local-bundle",
"kernel_module": cc_bundle.__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
|