[Kernel] Add inventory guards and clean benchmark layout (#32788)
This commit is contained in:
Xiaoyu Zhang
@@ -0,0 +1,228 @@
|
||||
"""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"]))
|
||||
Reference in New Issue
Block a user