Spaces:
Runtime error
Runtime error
File size: 5,107 Bytes
e317e25 | 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 | """Tests for kernel stubs.
Verifies that:
1. Every kernel stub file exists on disk.
2. Python stub files contain a module-level docstring.
3. Python stub files do NOT define a callable with that name
(they are stubs — Phase 2 will implement them).
Run:
uv run pytest tests/test_kernels.py -v
"""
import os
import pytest
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
KERNEL_DIR = os.path.join(_REPO, "kernels")
# ---------------------------------------------------------------------------
# Existence checks — one per stub file
# ---------------------------------------------------------------------------
_ALL_STUBS = [
("triton", "ssd_exp_trap.py"),
("triton", "sinkhorn_fused.py"),
("triton", "bcnorm_fused.py"),
("triton", "oja_update.py"),
("tilelang", "ssd_mimo_prefill.py"),
("tilelang", "mhc_kernels.py"),
("cuda", "hash_kernel.cu"),
("cuda", "decode_kernels.cu"),
]
_PYTHON_STUBS = [
("triton", "ssd_exp_trap.py"),
("triton", "sinkhorn_fused.py"),
("triton", "bcnorm_fused.py"),
("triton", "oja_update.py"),
("tilelang", "ssd_mimo_prefill.py"),
("tilelang", "mhc_kernels.py"),
]
_CUDA_STUBS = [
("cuda", "hash_kernel.cu"),
("cuda", "decode_kernels.cu"),
]
@pytest.mark.parametrize("subdir,filename", _ALL_STUBS)
def test_kernel_stub_exists(subdir: str, filename: str) -> None:
"""Each kernel stub file must exist on disk."""
path = os.path.join(KERNEL_DIR, subdir, filename)
assert os.path.exists(path), (
f"Missing kernel stub: kernels/{subdir}/{filename}\n"
f"(Full path: {path})"
)
@pytest.mark.parametrize("subdir,filename", _PYTHON_STUBS)
def test_python_stub_has_docstring(subdir: str, filename: str) -> None:
"""Python kernel stubs must have a module-level docstring."""
path = os.path.join(KERNEL_DIR, subdir, filename)
with open(path) as fh:
content = fh.read()
assert '"""' in content or "'''" in content, (
f"No docstring found in kernels/{subdir}/{filename}"
)
@pytest.mark.parametrize("subdir,filename", _PYTHON_STUBS)
def test_python_stub_is_non_empty(subdir: str, filename: str) -> None:
"""Python stub files must contain at least some text (not empty)."""
path = os.path.join(KERNEL_DIR, subdir, filename)
assert os.path.getsize(path) > 0, (
f"kernels/{subdir}/{filename} is empty"
)
@pytest.mark.parametrize("subdir,filename", _CUDA_STUBS)
def test_cuda_stub_has_comment(subdir: str, filename: str) -> None:
"""CUDA stub files must contain a comment describing their purpose."""
path = os.path.join(KERNEL_DIR, subdir, filename)
with open(path) as fh:
content = fh.read()
assert "/*" in content or "//" in content, (
f"No comment found in kernels/{subdir}/{filename}"
)
def test_kernel_dir_structure() -> None:
"""kernels/ directory contains triton/, tilelang/, and cuda/ subdirectories."""
for subdir in ("triton", "tilelang", "cuda"):
path = os.path.join(KERNEL_DIR, subdir)
assert os.path.isdir(path), f"Missing kernels/{subdir}/ directory"
def test_triton_stub_count() -> None:
"""kernels/triton/ contains exactly the expected number of stubs."""
triton_dir = os.path.join(KERNEL_DIR, "triton")
py_files = [f for f in os.listdir(triton_dir) if f.endswith(".py")]
expected = {name for _, name in _PYTHON_STUBS if _ == "triton"}
assert expected.issubset(set(py_files)), (
f"Missing triton stubs: {expected - set(py_files)}"
)
def test_tilelang_stub_count() -> None:
"""kernels/tilelang/ contains exactly the expected number of stubs."""
tilelang_dir = os.path.join(KERNEL_DIR, "tilelang")
py_files = [f for f in os.listdir(tilelang_dir) if f.endswith(".py")]
expected = {name for _, name in _PYTHON_STUBS if _ == "tilelang"}
assert expected.issubset(set(py_files)), (
f"Missing tilelang stubs: {expected - set(py_files)}"
)
def test_cuda_stub_count() -> None:
"""kernels/cuda/ contains exactly the expected number of stubs."""
cuda_dir = os.path.join(KERNEL_DIR, "cuda")
cu_files = [f for f in os.listdir(cuda_dir) if f.endswith(".cu")]
expected = {name for _, name in _CUDA_STUBS}
assert expected.issubset(set(cu_files)), (
f"Missing CUDA stubs: {expected - set(cu_files)}"
)
# ---------------------------------------------------------------------------
# Content-quality checks for Python stubs
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("subdir,filename", _PYTHON_STUBS)
def test_stub_mentions_phase(subdir: str, filename: str) -> None:
"""Python stubs should document which Phase will implement them."""
path = os.path.join(KERNEL_DIR, subdir, filename)
with open(path) as fh:
content = fh.read()
assert "Phase" in content, (
f"kernels/{subdir}/{filename} should mention 'Phase 1' or 'Phase 2' in its docs"
)
|