[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",
|
||||
(
|
||||
|
||||
Reference in New Issue
Block a user