[Test] Consolidate test cleanup and CI taxonomy (net -11.4K lines) (#37436)
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Mick Qian
Mick Qian
parent
6a1ff90f2d
commit
4d23a4fa6d
@@ -1,228 +0,0 @@
|
||||
"""Guards that keep the ``diffusion`` package's import surface from eroding.
|
||||
|
||||
The reorganization only stays useful if two invariants hold:
|
||||
|
||||
1. runtime code imports from ``sglang.kernels.ops.diffusion`` and not from a
|
||||
submodule, so the internal layout can move without touching call sites;
|
||||
2. the facade's ``_EXPORTS`` table and the registry's ``_SPECS`` table both
|
||||
point at symbols that actually exist.
|
||||
|
||||
Neither is checkable by the type system, and both fail silently -- a stale
|
||||
``_EXPORTS`` entry only raises when some model happens to call that kernel, on
|
||||
a GPU, at serving time. These are pure-CPU tests: they read the tables and
|
||||
resolve them with ``importlib``/``ast`` without importing torch backends.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import functools
|
||||
import importlib
|
||||
import pathlib
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.kernels.ops.diffusion import _EXPORTS, _SPECS
|
||||
from sglang.kernels.registry import registry
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=16, suite="base-a-test-cpu")
|
||||
|
||||
PACKAGE = "sglang.kernels.ops.diffusion"
|
||||
_PACKAGE_DIR = pathlib.Path(importlib.import_module(PACKAGE).__file__ or "").parent
|
||||
_REPO_ROOT = _PACKAGE_DIR.parents[4] # <repo>/python/sglang/kernels/ops/diffusion
|
||||
|
||||
# Backend-specific test files may name a leaf module on purpose; everything
|
||||
# else -- all runtime code -- must go through the facade.
|
||||
_DEEP_IMPORT_ALLOWLIST = {
|
||||
"python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py",
|
||||
"test/registered/kernels/ops/diffusion/test_model_fast_paths.py",
|
||||
"test/registered/kernels/ops/diffusion/test_sites.py",
|
||||
# This test exercises the pure-Torch fallback implementation directly.
|
||||
"test/registered/unit/utils/test_diffusion_torch_fallback.py",
|
||||
}
|
||||
|
||||
|
||||
def _module_defines(module_path: str) -> set[str]:
|
||||
"""Top-level names bound by a submodule, without importing it.
|
||||
|
||||
Importing would pull in Triton / CuTe-DSL / FlyDSL, none of which are
|
||||
installed on the CPU CI lane -- so this reads the source instead.
|
||||
"""
|
||||
if module_path.startswith("sglang."):
|
||||
spec = importlib.util.find_spec(module_path)
|
||||
assert spec is not None and spec.origin is not None, module_path
|
||||
path = pathlib.Path(spec.origin)
|
||||
else:
|
||||
path = _PACKAGE_DIR / (module_path.replace(".", "/") + ".py")
|
||||
if not path.exists():
|
||||
path = _PACKAGE_DIR / module_path.replace(".", "/") / "__init__.py"
|
||||
assert path.exists(), f"{PACKAGE}.{module_path} does not exist"
|
||||
|
||||
names: set[str] = set()
|
||||
for node in ast.parse(path.read_text(encoding="utf-8")).body:
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
names.add(node.name)
|
||||
elif isinstance(node, ast.Assign):
|
||||
names.update(t.id for t in node.targets if isinstance(t, ast.Name))
|
||||
elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
|
||||
names.add(node.target.id)
|
||||
elif isinstance(node, (ast.Import, ast.ImportFrom)):
|
||||
names.update((a.asname or a.name).split(".")[0] for a in node.names)
|
||||
elif isinstance(node, (ast.If, ast.Try)):
|
||||
# Platform-conditional rebinds (``x = select_impl(...)``) and
|
||||
# guarded defs still bind a public name.
|
||||
for inner in ast.walk(node):
|
||||
if isinstance(inner, (ast.FunctionDef, ast.ClassDef)):
|
||||
names.add(inner.name)
|
||||
elif isinstance(inner, ast.Assign):
|
||||
names.update(t.id for t in inner.targets if isinstance(t, ast.Name))
|
||||
return names
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=None)
|
||||
def _scan_root(root: str) -> tuple[frozenset[str], tuple[str, ...]]:
|
||||
unexported: set[str] = set()
|
||||
offenders: list[str] = []
|
||||
root_dir = _REPO_ROOT / root
|
||||
if not root_dir.exists():
|
||||
return frozenset(), ()
|
||||
|
||||
for path in root_dir.rglob("*.py"):
|
||||
rel = path.relative_to(_REPO_ROOT).as_posix()
|
||||
if rel.startswith(
|
||||
(
|
||||
"python/sglang/kernels/ops/diffusion/",
|
||||
"python/sglang/kernels/kda_kernels/",
|
||||
)
|
||||
):
|
||||
continue
|
||||
try:
|
||||
source = path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
if PACKAGE not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
continue
|
||||
allowlisted = rel in _DEEP_IMPORT_ALLOWLIST
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.ImportFrom):
|
||||
if node.module == PACKAGE:
|
||||
unexported.update(
|
||||
a.name
|
||||
for a in node.names
|
||||
if a.name not in _EXPORTS and not a.name.startswith("_")
|
||||
)
|
||||
elif (
|
||||
not allowlisted
|
||||
and node.module
|
||||
and node.module.startswith(f"{PACKAGE}.")
|
||||
):
|
||||
offenders.append(f"{rel}:{node.lineno} imports {node.module}")
|
||||
elif isinstance(node, ast.Import) and not allowlisted:
|
||||
offenders.extend(
|
||||
f"{rel}:{node.lineno} imports {a.name}"
|
||||
for a in node.names
|
||||
if a.name.startswith(f"{PACKAGE}.")
|
||||
)
|
||||
return frozenset(unexported), tuple(offenders)
|
||||
|
||||
|
||||
def test_every_export_resolves_to_a_real_symbol():
|
||||
missing = [
|
||||
f"{symbol} -> {module}"
|
||||
for symbol, module in sorted(_EXPORTS.items())
|
||||
if symbol not in _module_defines(module)
|
||||
]
|
||||
assert not missing, f"stale _EXPORTS entries: {missing}"
|
||||
|
||||
|
||||
def test_every_symbol_imported_from_the_facade_is_exported():
|
||||
"""The reverse of the check above, and the one that actually bites.
|
||||
|
||||
A missing ``_EXPORTS`` entry raises ``ImportError`` at module import, so a
|
||||
module-level ``from ...diffusion import x`` fails loudly. A *function-local*
|
||||
one -- the pattern used for optional backends -- fails only when that test
|
||||
or code path runs, on the platform that has the backend. Enumerating the
|
||||
call sites catches it here instead.
|
||||
"""
|
||||
unexported: set[str] = set()
|
||||
for root in ("python/sglang", "test", "benchmark"):
|
||||
unexported.update(_scan_root(root)[0])
|
||||
assert not unexported, f"imported but not in _EXPORTS: {sorted(unexported)}"
|
||||
|
||||
|
||||
def test_every_registered_spec_target_resolves():
|
||||
missing = []
|
||||
for _op, _backend, target, _caps, _description in _SPECS:
|
||||
module, _, attr = target.partition(":")
|
||||
if attr not in _module_defines(module):
|
||||
missing.append(target)
|
||||
assert not missing, f"stale _SPECS targets: {missing}"
|
||||
|
||||
|
||||
def test_registry_holds_the_diffusion_ops():
|
||||
# Registration happens at package import, is metadata-only, and is what
|
||||
# ``select_kernel`` / the tracing tools read.
|
||||
registered = {op for op in registry.ops() if op.startswith("diffusion.")}
|
||||
assert {op for op, *_ in _SPECS} <= registered
|
||||
|
||||
|
||||
def test_facade_rejects_unknown_attributes():
|
||||
module = sys.modules[PACKAGE]
|
||||
with pytest.raises(AttributeError):
|
||||
module.definitely_not_a_kernel
|
||||
assert set(module.__all__) == set(_EXPORTS)
|
||||
assert set(_EXPORTS) <= set(dir(module))
|
||||
|
||||
|
||||
def test_importing_the_package_does_not_import_any_leaf_module():
|
||||
"""The reason ``__getattr__`` is lazy rather than a block of re-exports.
|
||||
|
||||
The backends have disjoint, heavy, mutually-exclusive dependencies --
|
||||
Triton (CUDA/ROCm), CUTLASS/CuTe-DSL, and FlyDSL (gfx950). If
|
||||
``_EXPORTS`` ever degrades into eager ``from .norm.x import y`` lines, all
|
||||
of them become import-time requirements on every platform, which is how a
|
||||
CPU-only or Apple install starts failing at ``import sglang``.
|
||||
|
||||
Asserted on this package's own leaf modules rather than on ``triton`` in
|
||||
``sys.modules``: sibling operator groups import Triton for their own
|
||||
reasons, so a global check would not isolate this package's behavior.
|
||||
Run in a fresh interpreter because this process has already resolved
|
||||
exports through the facade.
|
||||
"""
|
||||
code = (
|
||||
"import importlib, sys\n"
|
||||
f"importlib.import_module('{PACKAGE}')\n"
|
||||
f"prefix = '{PACKAGE}.'\n"
|
||||
"leaves = [m for m in sys.modules if m.startswith(prefix)"
|
||||
" and not m.endswith('__init__')]\n"
|
||||
"print(','.join(sorted(m for m in leaves if '.' in m[len(prefix):]"
|
||||
" or sys.modules[m].__file__ and not sys.modules[m].__file__"
|
||||
".endswith('__init__.py'))))\n"
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code], capture_output=True, text=True, timeout=600
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
leaked = [m for m in result.stdout.strip().split(",") if m]
|
||||
assert not leaked, f"importing {PACKAGE} eagerly imported: {leaked}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("root", ["python/sglang", "test", "benchmark"])
|
||||
def test_runtime_code_imports_only_through_the_facade(root):
|
||||
if not (_REPO_ROOT / root).exists(): # source checkouts only
|
||||
pytest.skip(f"{root} not present in this install")
|
||||
|
||||
offenders = _scan_root(root)[1]
|
||||
assert not offenders, (
|
||||
"import from sglang.kernels.ops.diffusion instead of a submodule:\n "
|
||||
+ "\n ".join(offenders)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -1,7 +1,6 @@
|
||||
"""GPU-free import / registry / selector tests for ``sglang.kernels`` (RFC #29630)."""
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
@@ -9,127 +8,19 @@ import pytest
|
||||
|
||||
import sglang.kernels as K
|
||||
import sglang.kernels.fused_op as fo
|
||||
import sglang.kernels.ops # noqa: F401 -- populate the registry
|
||||
import sglang.kernels.selector as sel
|
||||
from sglang.kernels import DeviceType, KernelBackend, PlatformInfo
|
||||
from sglang.kernels import KernelBackend, PlatformInfo
|
||||
from sglang.kernels.spec import CapabilityRequirement as Cap
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=24, suite="base-a-test-cpu")
|
||||
|
||||
GROUPS = K.ops.__all__
|
||||
|
||||
# Representative ops checked as a subset (the registry holds many more).
|
||||
EXPECTED = {
|
||||
"activation.silu_and_mul": {"aot", "jit", "aiter", "torch", "torch_compile"},
|
||||
"activation.relu2": {"jit", "torch", "torch_compile"},
|
||||
"layernorm.rmsnorm": {"aot", "jit", "aiter", "torch_npu", "torch", "torch_compile"},
|
||||
"layernorm.gemma_rmsnorm": {"aot", "jit", "torch_npu", "torch", "torch_compile"},
|
||||
"gemm.fp8_scaled_mm": {"aot", "torch", "torch_compile"},
|
||||
"moe.moe_align_block_size": {"aot", "jit"},
|
||||
"quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
|
||||
"kvcache.reshape_and_cache_flash": {"triton"},
|
||||
"diffusion.apply_group_norm_silu": {"triton"},
|
||||
"diffusion.norm_scale_shift": {"KDA", "cute_dsl", "flydsl"},
|
||||
"diffusion.scale_residual_norm_scale_shift": {
|
||||
"KDA",
|
||||
"triton",
|
||||
"cute_dsl",
|
||||
"flydsl",
|
||||
},
|
||||
"diffusion.residual_gate_add": {"KDA"},
|
||||
"diffusion.ltx2_qknorm_split_rope": {"KDA"},
|
||||
"diffusion.causal_conv3d_cat_pad": {"KDA", "triton"},
|
||||
"diffusion.flux2_layernorm_modulate_fp8_quant": {"KDA"},
|
||||
"diffusion.flux2_qkv_epilogue": {"KDA"},
|
||||
"diffusion.flux2_token_cat_fp8": {"KDA"},
|
||||
"gemm.qwen3x_nvfp4": {"KDA"},
|
||||
"gemm.sm120_fp8_linear": {"KDA"},
|
||||
}
|
||||
|
||||
_CPU = PlatformInfo(device_type="cpu")
|
||||
_SM90 = PlatformInfo(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
|
||||
_SM100 = PlatformInfo(device_type="cuda", cuda_arch_major=10, cuda_arch_minor=0)
|
||||
_HIP = PlatformInfo(device_type="hip")
|
||||
|
||||
|
||||
def test_top_level_exports():
|
||||
for name in (
|
||||
"KernelSpec",
|
||||
"KernelBackend",
|
||||
"FormatSignature",
|
||||
"CapabilityRequirement",
|
||||
"PlatformInfo",
|
||||
"registry",
|
||||
"get_kernel",
|
||||
"select_kernel",
|
||||
):
|
||||
assert hasattr(K, name), name
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group", GROUPS)
|
||||
def test_group_importable(group):
|
||||
assert importlib.import_module(f"sglang.kernels.ops.{group}") is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op, backends", list(EXPECTED.items()))
|
||||
def test_registry_backends(op, backends):
|
||||
assert {s.backend.value for s in K.registry.get(op)} == backends
|
||||
|
||||
|
||||
def test_specs_well_formed():
|
||||
for spec in K.registry.all_specs():
|
||||
assert spec.op == f"{spec.group}.{spec.name}"
|
||||
mod, sep, attr = spec.target.partition(":")
|
||||
assert sep == ":" and mod and attr, spec.target
|
||||
|
||||
|
||||
def test_internal_registry_target_modules_exist():
|
||||
for spec in K.registry.all_specs():
|
||||
module, _, _ = spec.target.partition(":")
|
||||
if module.startswith("sglang.kernels."):
|
||||
assert importlib.util.find_spec(module) is not None, spec.target
|
||||
|
||||
|
||||
def test_sparse_linear_attention_registry_targets_forward_kernel():
|
||||
spec = K.registry.get_backend(
|
||||
"diffusion.sparse_linear_attn_fwd", KernelBackend.TRITON
|
||||
)
|
||||
assert spec.target.endswith(":_attn_fwd")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"op, target_suffix",
|
||||
(
|
||||
("diffusion.norm_scale_shift", ":kda_norm_scale_shift"),
|
||||
(
|
||||
"diffusion.scale_residual_norm_scale_shift",
|
||||
":kda_scale_residual_norm_scale_shift",
|
||||
),
|
||||
("diffusion.residual_gate_add", ":residual_gate_add"),
|
||||
(
|
||||
"diffusion.ltx2_qknorm_split_rope",
|
||||
":ltx2_qknorm_split_rope_cuda",
|
||||
),
|
||||
(
|
||||
"diffusion.causal_conv3d_cat_pad",
|
||||
":fused_causal_conv3d_cat_pad_cuda",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_merged_diffusion_kda_provenance_backend(op, target_suffix):
|
||||
spec = K.registry.get_backend(op, KernelBackend.KDA)
|
||||
assert spec.target.endswith(target_suffix)
|
||||
|
||||
|
||||
def test_kda_backend_implementations_live_in_kda_home():
|
||||
specs = [
|
||||
spec for spec in K.registry.all_specs() if spec.backend is KernelBackend.KDA
|
||||
]
|
||||
assert specs
|
||||
assert all(spec.target.startswith("sglang.kernels.kda_kernels.") for spec in specs)
|
||||
|
||||
|
||||
def test_single_backend_resolves_without_backend():
|
||||
assert (
|
||||
K.select_kernel("kvcache.reshape_and_cache_flash").backend
|
||||
@@ -193,15 +84,6 @@ def test_layernorm_default_backend(monkeypatch, op_attr, device, expect):
|
||||
assert getattr(ln, op_attr).auto_selected_backend().value == expect
|
||||
|
||||
|
||||
def test_per_op_backend_subset():
|
||||
# silu_and_mul ships an aiter (HIP) kernel; the gelu siblings deliberately
|
||||
# do not -- ROCm coverage is a per-(op, backend) subset.
|
||||
from sglang.kernels.ops.activation import _GELU_AND_MUL, _SILU_AND_MUL
|
||||
|
||||
assert KernelBackend.AITER in _SILU_AND_MUL.available_backends()
|
||||
assert KernelBackend.AITER not in _GELU_AND_MUL.available_backends()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"req, plat, ok",
|
||||
[
|
||||
@@ -227,20 +109,6 @@ def test_capabilities_or_semantics():
|
||||
assert K.capabilities_satisfied(Cap.CUDA, _SM90) # single tolerated
|
||||
|
||||
|
||||
def test_capability_shortcuts():
|
||||
assert Cap.CUDA == Cap(device=DeviceType.CUDA)
|
||||
assert Cap.HIP == Cap(device=DeviceType.HIP)
|
||||
assert Cap.NPU == Cap(device=DeviceType.NPU)
|
||||
assert {Cap.CUDA, Cap.HIP} == {Cap.HIP, Cap.CUDA}
|
||||
assert Cap.cuda(min_sm=(10, 0)) == Cap(
|
||||
device=DeviceType.CUDA, min_cuda_arch=(10, 0)
|
||||
)
|
||||
|
||||
|
||||
def test_platform_detect_does_not_raise():
|
||||
assert PlatformInfo.detect().device_type in ("cpu", "cuda", "hip", "npu")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"relative_path",
|
||||
(
|
||||
|
||||
@@ -1,239 +0,0 @@
|
||||
"""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=13, 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}"
|
||||
|
||||
|
||||
# `load_jit` takes in-tree names and absolute paths on the same keyword, so this
|
||||
# check can only reach the declarations spelled out in the source. A module that
|
||||
# assembles its file list at runtime from a package outside `jit/csrc` has no
|
||||
# in-tree name to verify and belongs here; there is none at the moment.
|
||||
_RUNTIME_JIT_SOURCE_MODULES: set[str] = set()
|
||||
|
||||
|
||||
def test_jit_source_declarations_exist():
|
||||
missing = []
|
||||
unsupported = []
|
||||
for python_file in OPS_ROOT.rglob("*.py"):
|
||||
if python_file.relative_to(OPS_ROOT).as_posix() in _RUNTIME_JIT_SOURCE_MODULES:
|
||||
continue
|
||||
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