229 lines
8.5 KiB
Python
229 lines
8.5 KiB
Python
"""CPU-only structural checks for the unified kernel tree."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import sglang.kernels as kernels
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
KERNELS_ROOT = REPO_ROOT / "python" / "sglang" / "kernels"
|
|
OPS_ROOT = KERNELS_ROOT / "ops"
|
|
JIT_CSRC_ROOT = KERNELS_ROOT / "jit" / "csrc"
|
|
AOT_ROOT = KERNELS_ROOT / "aot"
|
|
|
|
|
|
def _directory_names(root: Path) -> set[str]:
|
|
return {
|
|
path.name
|
|
for path in root.iterdir()
|
|
if path.is_dir()
|
|
and not path.name.startswith((".", "__"))
|
|
and any(path.rglob("*.py"))
|
|
}
|
|
|
|
|
|
def _target_names(target: ast.expr) -> set[str]:
|
|
if isinstance(target, ast.Name):
|
|
return {target.id}
|
|
if isinstance(target, (ast.List, ast.Tuple)):
|
|
return {name for element in target.elts for name in _target_names(element)}
|
|
return set()
|
|
|
|
|
|
def _bound_names(statements: list[ast.stmt]) -> set[str]:
|
|
"""Collect names a module can bind without importing it."""
|
|
names: set[str] = set()
|
|
for statement in statements:
|
|
if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
|
names.add(statement.name)
|
|
elif isinstance(statement, ast.Assign):
|
|
for target in statement.targets:
|
|
names.update(_target_names(target))
|
|
elif isinstance(statement, (ast.AnnAssign, ast.AugAssign)):
|
|
names.update(_target_names(statement.target))
|
|
elif isinstance(statement, (ast.Import, ast.ImportFrom)):
|
|
for alias in statement.names:
|
|
names.add(alias.asname or alias.name.split(".", 1)[0])
|
|
elif isinstance(statement, (ast.For, ast.AsyncFor)):
|
|
names.update(_target_names(statement.target))
|
|
names.update(_bound_names(statement.body))
|
|
names.update(_bound_names(statement.orelse))
|
|
elif isinstance(statement, ast.If):
|
|
names.update(_bound_names(statement.body))
|
|
names.update(_bound_names(statement.orelse))
|
|
elif isinstance(statement, (ast.With, ast.AsyncWith)):
|
|
names.update(_bound_names(statement.body))
|
|
elif isinstance(statement, ast.Try):
|
|
names.update(_bound_names(statement.body))
|
|
names.update(_bound_names(statement.orelse))
|
|
names.update(_bound_names(statement.finalbody))
|
|
for handler in statement.handlers:
|
|
names.update(_bound_names(handler.body))
|
|
elif isinstance(statement, ast.Match):
|
|
for case in statement.cases:
|
|
names.update(_bound_names(case.body))
|
|
return names
|
|
|
|
|
|
def _module_string_constants(tree: ast.Module) -> dict[str, str]:
|
|
constants: dict[str, str] = {}
|
|
for statement in tree.body:
|
|
if not isinstance(statement, (ast.Assign, ast.AnnAssign)):
|
|
continue
|
|
value = statement.value
|
|
if not isinstance(value, ast.Constant) or not isinstance(value.value, str):
|
|
continue
|
|
targets = (
|
|
statement.targets
|
|
if isinstance(statement, ast.Assign)
|
|
else [statement.target]
|
|
)
|
|
for target in targets:
|
|
for name in _target_names(target):
|
|
constants[name] = value.value
|
|
return constants
|
|
|
|
|
|
def _source_patterns(expression: ast.expr, constants: dict[str, str]) -> list[str]:
|
|
if isinstance(expression, (ast.List, ast.Tuple)):
|
|
return [
|
|
pattern
|
|
for element in expression.elts
|
|
for pattern in _source_patterns(element, constants)
|
|
]
|
|
if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
|
|
return [expression.value]
|
|
if isinstance(expression, ast.Name) and expression.id in constants:
|
|
return [constants[expression.id]]
|
|
if isinstance(expression, ast.JoinedStr):
|
|
parts = []
|
|
for value in expression.values:
|
|
if isinstance(value, ast.Constant):
|
|
parts.append(str(value.value))
|
|
elif isinstance(value, ast.FormattedValue):
|
|
parts.append("*")
|
|
else:
|
|
raise AssertionError(f"Unsupported f-string segment: {ast.dump(value)}")
|
|
return ["".join(parts)]
|
|
raise AssertionError(
|
|
f"Unsupported JIT source declaration: {ast.unparse(expression)}"
|
|
)
|
|
|
|
|
|
def test_declared_operator_groups_match_packages():
|
|
assert set(kernels.ops.__all__) == _directory_names(OPS_ROOT)
|
|
|
|
|
|
def test_registered_kernel_test_groups_are_known():
|
|
declared_groups = set(kernels.ops.__all__)
|
|
registered_root = REPO_ROOT / "test" / "registered" / "kernels"
|
|
for kind in ("ops", "benchmark"):
|
|
unknown = _directory_names(registered_root / kind) - declared_groups
|
|
assert (
|
|
not unknown
|
|
), f"Unknown {kind} kernel group directories: {sorted(unknown)}"
|
|
|
|
|
|
def test_internal_registry_target_attributes_are_declared():
|
|
missing = []
|
|
for spec in kernels.registry.all_specs():
|
|
module_name, _, attribute_path = spec.target.partition(":")
|
|
if not module_name.startswith("sglang.kernels."):
|
|
continue
|
|
module_spec = importlib.util.find_spec(module_name)
|
|
if (
|
|
module_spec is None
|
|
or module_spec.origin is None
|
|
or not module_spec.origin.endswith(".py")
|
|
):
|
|
continue
|
|
tree = ast.parse(Path(module_spec.origin).read_text())
|
|
root_attribute = attribute_path.split(".", 1)[0]
|
|
if root_attribute not in _bound_names(tree.body):
|
|
missing.append(spec.target)
|
|
assert not missing, f"KernelSpec targets missing attributes: {missing}"
|
|
|
|
|
|
def test_jit_source_declarations_exist():
|
|
missing = []
|
|
unsupported = []
|
|
for python_file in OPS_ROOT.rglob("*.py"):
|
|
tree = ast.parse(python_file.read_text())
|
|
constants = _module_string_constants(tree)
|
|
for call in (node for node in ast.walk(tree) if isinstance(node, ast.Call)):
|
|
function_name = (
|
|
call.func.id
|
|
if isinstance(call.func, ast.Name)
|
|
else call.func.attr if isinstance(call.func, ast.Attribute) else None
|
|
)
|
|
if function_name != "load_jit":
|
|
continue
|
|
for keyword in call.keywords:
|
|
if keyword.arg not in {"cpp_files", "cuda_files"}:
|
|
continue
|
|
try:
|
|
patterns = _source_patterns(keyword.value, constants)
|
|
except AssertionError as exc:
|
|
unsupported.append(f"{python_file.relative_to(REPO_ROOT)}: {exc}")
|
|
continue
|
|
for pattern in patterns:
|
|
matches = list(JIT_CSRC_ROOT.glob(pattern))
|
|
if not matches:
|
|
missing.append(
|
|
f"{python_file.relative_to(REPO_ROOT)} -> {pattern}"
|
|
)
|
|
assert not unsupported, "Unsupported JIT source declarations:\n" + "\n".join(
|
|
unsupported
|
|
)
|
|
assert not missing, "Missing JIT sources:\n" + "\n".join(missing)
|
|
|
|
|
|
def test_aot_compilation_units_are_accounted_for():
|
|
manifests = [
|
|
AOT_ROOT / "CMakeLists.txt",
|
|
AOT_ROOT / "setup_metal.py",
|
|
AOT_ROOT / "setup_musa.py",
|
|
AOT_ROOT / "setup_rocm.py",
|
|
AOT_ROOT / "csrc" / "cpu" / "CMakeLists.txt",
|
|
*sorted((AOT_ROOT / "cmake").rglob("*.cmake")),
|
|
]
|
|
manifest_text = "\n".join(path.read_text() for path in manifests)
|
|
source_text = {
|
|
path: path.read_text(errors="ignore")
|
|
for path in (AOT_ROOT / "csrc").rglob("*")
|
|
if path.is_file()
|
|
}
|
|
compilation_suffixes = {".cc", ".cpp", ".cu", ".hip", ".metal", ".mu"}
|
|
missing = []
|
|
for source in source_text:
|
|
if source.suffix not in compilation_suffixes:
|
|
continue
|
|
if AOT_ROOT / "csrc" / "cpu" in source.parents:
|
|
# The CPU build intentionally uses file(GLOB_RECURSE ... *.cpp).
|
|
continue
|
|
relative_path = source.relative_to(AOT_ROOT).as_posix()
|
|
if relative_path in manifest_text:
|
|
continue
|
|
if any(
|
|
source.name in text
|
|
for other_source, text in source_text.items()
|
|
if other_source != source
|
|
):
|
|
# Some CUDA translation units are included by another source.
|
|
continue
|
|
missing.append(relative_path)
|
|
assert not missing, f"AOT compilation units missing from build manifests: {missing}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__, "-v"]))
|