[Kernel] Add inventory guards and clean benchmark layout (#32788)
This commit is contained in:
Xiaoyu Zhang
+2
-2
@@ -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,
|
||||
)
|
||||
+1
-1
@@ -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).
|
||||
"""
|
||||
+14
-7
@@ -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 "
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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()))
|
||||
|
||||
@@ -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