From 1d9c2925473c2462c4a78b3033f8a34fc0c56e17 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 30 Jul 2026 09:03:24 +0800 Subject: [PATCH] [Kernel] Add inventory guards and clean benchmark layout (#32788) --- .../attention/bench_flash_attention_fp8.py | 4 +- .../attention}/bench_gdn_replayssm_decode.py | 2 +- .../kernels/attention/fa4_benchmark_utils.py | 0 .../kernels/attention}/sm90_config_search.py | 21 +- .../attention/flash_attn/cute/bench_utils.py | 261 ------------------ python/sglang/kernels/registry.py | 15 +- .../bench_add_constant.py | 0 .../bench_hadamard.py | 0 .../test_inkling_rel_proj.py | 0 .../test_add_constant.py | 0 .../kernels/ops/layernorm/test_fused_op.py | 14 +- .../ops/layernorm/test_kernels_namespace.py | 21 +- .../test_hadamard_jit.py | 0 .../kernels/test_kernel_inventory.py | 228 +++++++++++++++ 14 files changed, 269 insertions(+), 297 deletions(-) rename python/sglang/kernels/ops/attention/flash_attn/cute/benchmark_flash_attention_fp8.py => benchmark/kernels/attention/bench_flash_attention_fp8.py (99%) rename {python/sglang/kernels/ops/attention/fla => benchmark/kernels/attention}/bench_gdn_replayssm_decode.py (98%) rename python/sglang/kernels/ops/attention/flash_attn/cute/benchmark.py => benchmark/kernels/attention/fa4_benchmark_utils.py (100%) rename {python/sglang/kernels/ops/attention/flash_attn/cute => benchmark/kernels/attention}/sm90_config_search.py (95%) delete mode 100644 python/sglang/kernels/ops/attention/flash_attn/cute/bench_utils.py rename test/registered/kernels/benchmark/{attention => elementwise}/bench_add_constant.py (100%) rename test/registered/kernels/benchmark/{attention => quantization}/bench_hadamard.py (100%) rename test/registered/kernels/ops/{model => attention}/test_inkling_rel_proj.py (100%) rename test/registered/kernels/ops/{attention => elementwise}/test_add_constant.py (100%) rename test/registered/kernels/ops/{attention => quantization}/test_hadamard_jit.py (100%) create mode 100644 test/registered/kernels/test_kernel_inventory.py diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/benchmark_flash_attention_fp8.py b/benchmark/kernels/attention/bench_flash_attention_fp8.py similarity index 99% rename from python/sglang/kernels/ops/attention/flash_attn/cute/benchmark_flash_attention_fp8.py rename to benchmark/kernels/attention/bench_flash_attention_fp8.py index 844b2ab82..5ff291b1d 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/benchmark_flash_attention_fp8.py +++ b/benchmark/kernels/attention/bench_flash_attention_fp8.py @@ -1,7 +1,7 @@ # Benchmark FP8 attention for FA4 (CuTe-DSL) on SM100. # # Run (recommended): -# python -m flash_attn.cute.benchmark_flash_attention_fp8 +# python benchmark/kernels/attention/bench_flash_attention_fp8.py # # Notes: # - This is intended to be used while bringing up FP8 support for SM100. @@ -20,8 +20,8 @@ from typing import Iterable import torch from einops import rearrange +from fa4_benchmark_utils import benchmark_forward -from sglang.kernels.ops.attention.flash_attn.cute.benchmark import benchmark_forward from sglang.kernels.ops.attention.flash_attn.cute.interface import ( _flash_attn_fwd as flash_attn_cute_fwd, ) diff --git a/python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py b/benchmark/kernels/attention/bench_gdn_replayssm_decode.py similarity index 98% rename from python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py rename to benchmark/kernels/attention/bench_gdn_replayssm_decode.py index 58ecd41f3..c27ef70c4 100644 --- a/python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py +++ b/benchmark/kernels/attention/bench_gdn_replayssm_decode.py @@ -15,7 +15,7 @@ traffic ratio is reported per L. Run:: - python -m sglang.kernels.ops.attention.fla.bench_gdn_replayssm_decode + python benchmark/kernels/attention/bench_gdn_replayssm_decode.py Requires a GPU (Triton). """ diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/benchmark.py b/benchmark/kernels/attention/fa4_benchmark_utils.py similarity index 100% rename from python/sglang/kernels/ops/attention/flash_attn/cute/benchmark.py rename to benchmark/kernels/attention/fa4_benchmark_utils.py diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/sm90_config_search.py b/benchmark/kernels/attention/sm90_config_search.py similarity index 95% rename from python/sglang/kernels/ops/attention/flash_attn/cute/sm90_config_search.py rename to benchmark/kernels/attention/sm90_config_search.py index 1e1a38997..7686a8024 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/sm90_config_search.py +++ b/benchmark/kernels/attention/sm90_config_search.py @@ -4,9 +4,11 @@ Enumerates tile sizes, swap modes, atom layouts, and staging options. Checks GMMA divisibility, register budget, and shared memory budget. Usage: - python flash_attn/cute/sm90_config_search.py --headdim 128 - python flash_attn/cute/sm90_config_search.py --mode fwd --headdim 192-128 - python flash_attn/cute/sm90_config_search.py --mode bwd --headdim 192 --tile-n 64,96 + python benchmark/kernels/attention/sm90_config_search.py --headdim 128 + python benchmark/kernels/attention/sm90_config_search.py \ + --mode fwd --headdim 192-128 + python benchmark/kernels/attention/sm90_config_search.py \ + --mode bwd --headdim 192 --tile-n 64,96 """ import math @@ -17,6 +19,10 @@ REG_LIMITS = {2: 216, 3: 128} # per-WG budget: 2WG=240-24, 3WG=160-32 THREADS_PER_WG = 128 +def _bool_flag(value): + return "T" if value else "F" + + def _divisors(n): return [d for d in range(1, n + 1) if n % d == 0] @@ -242,11 +248,12 @@ def print_bwd_configs(configs, max_results=20): ) print(hdr) print("-" * len(hdr)) - B = lambda b: "T" if b else "F" for c in configs[:max_results]: print( f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3} " - f"{B(c['SdP_swapAB']):>3} {B(c['dKV_swapAB']):>3} {B(c['dQ_swapAB']):>3} " + f"{_bool_flag(c['SdP_swapAB']):>3} " + f"{_bool_flag(c['dKV_swapAB']):>3} " + f"{_bool_flag(c['dQ_swapAB']):>3} " f"{c['AtomLayoutMSdP']:>4} {c['AtomLayoutNdKV']:>4} {c['AtomLayoutMdQ']:>4} " f"{c['Q_stage']:>2} {c['dO_stage']:>3} " f"{c['regs_SdP']:>3} {c['regs_dK']:>3} {c['regs_dV']:>3} {c['regs_dQ']:>3} " @@ -354,11 +361,11 @@ def print_fwd_configs(configs, max_results=20): ) print(hdr) print("-" * len(hdr)) - B = lambda b: "T" if b else "F" for c in configs[:max_results]: print( f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3} " - f"{B(c['pv_is_rs']):>2} {B(c['overlap_wg']):>4} " + f"{_bool_flag(c['pv_is_rs']):>2} " + f"{_bool_flag(c['overlap_wg']):>4} " f"{c['regs_S']:>3} {c['regs_P']:>3} {c['regs_O']:>3} " f"{c['total_regs']:>4}/{c['reg_limit']:<3} " f"{c['smem_kb']:>4.0f}K " diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/bench_utils.py b/python/sglang/kernels/ops/attention/flash_attn/cute/bench_utils.py deleted file mode 100644 index 537cc2e85..000000000 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/bench_utils.py +++ /dev/null @@ -1,261 +0,0 @@ -"""Shared benchmark utilities: attention_ref, cuDNN helpers, flops calculation.""" - -import math - -import torch - -try: - import cudnn -except ImportError: - cudnn = None - - -# ── FLOPS calculation ──────────────────────────────────────────────────────── - - -def flops( - batch, - nheads, - seqlen_q, - seqlen_k, - headdim, - headdim_v, - causal=False, - window_size=(None, None), - has_qv=False, -): - if causal: - avg_seqlen = (max(0, seqlen_k - seqlen_q) + seqlen_k) / 2 - else: - if window_size == (None, None): - avg_seqlen = seqlen_k - else: - row_idx = torch.arange(seqlen_q, device="cuda") - col_left = ( - torch.maximum( - row_idx + seqlen_k - seqlen_q - window_size[0], torch.tensor(0) - ) - if window_size[0] is not None - else torch.zeros_like(row_idx) - ) - col_right = ( - torch.minimum( - row_idx + seqlen_k - seqlen_q + window_size[1], - torch.tensor(seqlen_k - 1), - ) - if window_size[1] is not None - else torch.full_like(row_idx, seqlen_k - 1) - ) - avg_seqlen = (col_right - col_left + 1).float().mean().item() - eff_headdim = headdim + headdim_v if has_qv else headdim - return batch * nheads * 2 * seqlen_q * avg_seqlen * (eff_headdim + headdim_v) - - -# ── Bandwidth calculation ──────────────────────────────────────────────────── - - -def bandwidth_fwd_bytes( - batch, - nheads, - nheads_kv, - seqlen_q, - seqlen_k, - headdim, - headdim_v, - dtype_bytes=2, - has_qv=False, -): - """HBM traffic for one attention pass: read Q,K,V + write O.""" - q = batch * nheads * seqlen_q * headdim - qv = batch * nheads * seqlen_q * headdim_v if has_qv else 0 - k = batch * nheads_kv * seqlen_k * headdim - v = batch * nheads_kv * seqlen_k * headdim_v - o = batch * nheads * seqlen_q * headdim_v - return (q + qv + k + v + o) * dtype_bytes - - -def bandwidth_bwd_bytes( - batch, nheads, nheads_kv, seqlen_q, seqlen_k, headdim, headdim_v, dtype_bytes=2 -): - """HBM traffic for one attention pass: read Q,K,V,dO + write dQ,dK,dV.""" - q = batch * nheads * seqlen_q * headdim - k = batch * nheads_kv * seqlen_k * headdim - v = batch * nheads_kv * seqlen_k * headdim_v - do = batch * nheads * seqlen_q * headdim_v - dq = q - dk = k - dv = v - return (q + k + v + do + dq + dk + dv) * dtype_bytes - - -# ── Reference attention ───────────────────────────────────────────────────── - -_attention_ref_mask_cache = {} - - -def attention_ref(q, k, v, causal=False): - """Standard attention reference implementation. - - Args: - q, k, v: (batch, seqlen, nheads, headdim) tensors. - causal: whether to apply causal mask. - """ - softmax_scale = 1.0 / math.sqrt(q.shape[-1]) - scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k) - if causal: - if scores.shape[-2] not in _attention_ref_mask_cache: - mask = torch.tril( - torch.ones(scores.shape[-2:], device=scores.device, dtype=torch.bool), - diagonal=0, - ) - _attention_ref_mask_cache[scores.shape[-2]] = mask - else: - mask = _attention_ref_mask_cache[scores.shape[-2]] - scores = scores.masked_fill(mask, float("-inf")) - attn = torch.softmax(scores, dim=-1) - return torch.einsum("bhts,bshd->bthd", attn, v) - - -# ── cuDNN graph helpers ───────────────────────────────────────────────────── - -_TORCH_TO_CUDNN_DTYPE = { - torch.float16: "HALF", - torch.bfloat16: "BFLOAT16", - torch.float32: "FLOAT", - torch.int32: "INT32", - torch.int64: "INT64", -} - - -def _build_cudnn_graph(io_dtype, tensors, build_fn): - """Build a cuDNN graph. Returns (graph, variant_pack, workspace).""" - assert cudnn is not None, "cuDNN is not available" - cudnn_dtype = getattr(cudnn.data_type, _TORCH_TO_CUDNN_DTYPE[io_dtype]) - graph = cudnn.pygraph( - io_data_type=cudnn_dtype, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) - graph_tensors = {name: graph.tensor_like(t.detach()) for name, t in tensors.items()} - variant_pack = build_fn(graph, graph_tensors) - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - graph.build_plans() - workspace = torch.empty( - graph.get_workspace_size(), device="cuda", dtype=torch.uint8 - ) - return graph, variant_pack, workspace - - -def cudnn_fwd_setup(q, k, v, causal=False, window_size_left=None): - """Build a cuDNN forward SDPA graph. - - Args: - q, k, v: (batch, nheads, seqlen, headdim) tensors (cuDNN layout). - causal: whether to apply causal mask. - window_size_left: sliding window size (None for no window). - - Returns: - (fwd_fn, o_gpu, stats_gpu) where fwd_fn is a zero-arg callable. - """ - b, nheads, seqlen_q, headdim = q.shape - headdim_v = v.shape[-1] - o_gpu = torch.empty(b, nheads, seqlen_q, headdim_v, dtype=q.dtype, device=q.device) - stats_gpu = torch.empty( - b, nheads, seqlen_q, 1, dtype=torch.float32, device=q.device - ) - - def build(graph, gt): - o, stats = graph.sdpa( - name="sdpa", - q=gt["q"], - k=gt["k"], - v=gt["v"], - is_inference=False, - attn_scale=1.0 / math.sqrt(headdim), - use_causal_mask=causal or window_size_left is not None, - sliding_window_length=( - window_size_left - if window_size_left is not None and not causal - else None - ), - ) - o.set_output(True).set_dim(o_gpu.shape).set_stride(o_gpu.stride()) - stats.set_output(True).set_data_type(cudnn.data_type.FLOAT) - return {gt["q"]: q, gt["k"]: k, gt["v"]: v, o: o_gpu, stats: stats_gpu} - - graph, variant_pack, workspace = _build_cudnn_graph( - q.dtype, {"q": q, "k": k, "v": v}, build - ) - - def fwd_fn(): - graph.execute(variant_pack, workspace) - return o_gpu - - return fwd_fn, o_gpu, stats_gpu - - -def cudnn_bwd_setup(q, k, v, o, g, lse, causal=False, window_size_left=None): - """Build a cuDNN backward SDPA graph. - - Args: - q, k, v, o, g, lse: (batch, nheads, seqlen, dim) tensors (cuDNN layout). - causal: whether to apply causal mask. - window_size_left: sliding window size (None for no window). - - Returns: - bwd_fn: zero-arg callable that returns (dq, dk, dv). - """ - headdim = q.shape[-1] - dq_gpu, dk_gpu, dv_gpu = ( - torch.empty_like(q), - torch.empty_like(k), - torch.empty_like(v), - ) - - def build(graph, gt): - dq, dk, dv = graph.sdpa_backward( - name="sdpa_backward", - q=gt["q"], - k=gt["k"], - v=gt["v"], - o=gt["o"], - dO=gt["g"], - stats=gt["lse"], - attn_scale=1.0 / math.sqrt(headdim), - use_causal_mask=causal or window_size_left is not None, - sliding_window_length=( - window_size_left - if window_size_left is not None and not causal - else None - ), - use_deterministic_algorithm=False, - ) - dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride()) - dk.set_output(True).set_dim(dk_gpu.shape).set_stride(dk_gpu.stride()) - dv.set_output(True).set_dim(dv_gpu.shape).set_stride(dv_gpu.stride()) - return { - gt["q"]: q, - gt["k"]: k, - gt["v"]: v, - gt["o"]: o, - gt["g"]: g, - gt["lse"]: lse, - dq: dq_gpu, - dk: dk_gpu, - dv: dv_gpu, - } - - graph, variant_pack, workspace = _build_cudnn_graph( - q.dtype, - {"q": q, "k": k, "v": v, "o": o, "g": g, "lse": lse}, - build, - ) - - def bwd_fn(): - graph.execute(variant_pack, workspace) - return dq_gpu, dk_gpu, dv_gpu - - return bwd_fn diff --git a/python/sglang/kernels/registry.py b/python/sglang/kernels/registry.py index bbf8b3bf1..dc93b319d 100644 --- a/python/sglang/kernels/registry.py +++ b/python/sglang/kernels/registry.py @@ -23,13 +23,20 @@ class KernelRegistry: def register(self, spec: KernelSpec) -> KernelSpec: """Register ``spec``. - Re-registering the same ``(op, backend)`` pair replaces the previous - entry so that module reloads during tests stay idempotent. + Re-registering an identical spec is idempotent so that module reloads + during tests remain safe. A different spec for the same ``(op, + backend)`` pair is rejected because silently replacing it makes the + selected implementation depend on import order. """ existing = self._by_op[spec.op] - for i, other in enumerate(existing): + for other in existing: if other.backend == spec.backend: - existing[i] = spec + if other != spec: + raise ValueError( + f"Conflicting kernel registration for op {spec.op!r}, " + f"backend {spec.backend.value!r}: " + f"{other.target!r} != {spec.target!r}" + ) return spec existing.append(spec) return spec diff --git a/test/registered/kernels/benchmark/attention/bench_add_constant.py b/test/registered/kernels/benchmark/elementwise/bench_add_constant.py similarity index 100% rename from test/registered/kernels/benchmark/attention/bench_add_constant.py rename to test/registered/kernels/benchmark/elementwise/bench_add_constant.py diff --git a/test/registered/kernels/benchmark/attention/bench_hadamard.py b/test/registered/kernels/benchmark/quantization/bench_hadamard.py similarity index 100% rename from test/registered/kernels/benchmark/attention/bench_hadamard.py rename to test/registered/kernels/benchmark/quantization/bench_hadamard.py diff --git a/test/registered/kernels/ops/model/test_inkling_rel_proj.py b/test/registered/kernels/ops/attention/test_inkling_rel_proj.py similarity index 100% rename from test/registered/kernels/ops/model/test_inkling_rel_proj.py rename to test/registered/kernels/ops/attention/test_inkling_rel_proj.py diff --git a/test/registered/kernels/ops/attention/test_add_constant.py b/test/registered/kernels/ops/elementwise/test_add_constant.py similarity index 100% rename from test/registered/kernels/ops/attention/test_add_constant.py rename to test/registered/kernels/ops/elementwise/test_add_constant.py diff --git a/test/registered/kernels/ops/layernorm/test_fused_op.py b/test/registered/kernels/ops/layernorm/test_fused_op.py index 83b9b37d6..c98469068 100644 --- a/test/registered/kernels/ops/layernorm/test_fused_op.py +++ b/test/registered/kernels/ops/layernorm/test_fused_op.py @@ -143,11 +143,19 @@ def test_registry_register_and_get(): assert reg.get("no.such") == [] -def test_registry_reregister_replaces(): +def test_registry_reregister_is_idempotent(): + reg = KernelRegistry() + spec = _spec(target="math:sqrt") + reg.register(spec) + reg.register(spec) + assert reg.get("g.n") == [spec] + + +def test_registry_rejects_conflicting_backend_registration(): reg = KernelRegistry() reg.register(_spec(target="math:sqrt")) - reg.register(_spec(target="math:floor")) - assert [s.target for s in reg.get("g.n")] == ["math:floor"] + with pytest.raises(ValueError, match="Conflicting kernel registration"): + reg.register(_spec(target="math:floor")) def test_registry_get_backend_missing_raises(): diff --git a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py index 8380e5ce2..a61f79124 100644 --- a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py +++ b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py @@ -17,24 +17,7 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") -GROUPS = [ - "activation", - "attention", - "communication", - "diffusion", - "elementwise", - "embeddings", - "gemm", - "grammar", - "kvcache", - "layernorm", - "mamba", - "memory", - "moe", - "quantization", - "sampling", - "speculative", -] +GROUPS = K.ops.__all__ # Representative ops checked as a subset (the registry holds many more). EXPECTED = { @@ -71,7 +54,7 @@ def test_top_level_exports(): @pytest.mark.parametrize("group", GROUPS) def test_group_importable(group): - assert hasattr(importlib.import_module(f"sglang.kernels.ops.{group}"), "__all__") + assert importlib.import_module(f"sglang.kernels.ops.{group}") is not None @pytest.mark.parametrize("op, backends", list(EXPECTED.items())) diff --git a/test/registered/kernels/ops/attention/test_hadamard_jit.py b/test/registered/kernels/ops/quantization/test_hadamard_jit.py similarity index 100% rename from test/registered/kernels/ops/attention/test_hadamard_jit.py rename to test/registered/kernels/ops/quantization/test_hadamard_jit.py diff --git a/test/registered/kernels/test_kernel_inventory.py b/test/registered/kernels/test_kernel_inventory.py new file mode 100644 index 000000000..e2f3dd6b2 --- /dev/null +++ b/test/registered/kernels/test_kernel_inventory.py @@ -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"]))