minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Kernel] Add inventory guards and clean benchmark layout (#32788)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-07-30 09:03:24 +08:00
committed by GitHub
parent 5efbb18a6f
commit 1d9c292547
14 changed files with 269 additions and 297 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py
-166
View File
@@ -1,166 +0,0 @@
"""Microbenchmark: buffered output-only GDN decode (ReplaySSM Part A) vs. the
existing packed GDN decode kernel.
Compares per-step decode latency of
``fused_recurrent_gated_delta_rule_packed_decode`` (writes the full recurrent
state S every step) against ``fused_recurrent_gdn_replayssm_decode`` at
L in {1, 8, 16} (writes the full state only every L steps) across batch sizes
{1, 16, 64, 256} for a realistic GDN config (HV=32, K=V=128).
The win is per-step HBM *state* traffic: the packed kernel reads + writes S
(~2 * num_slots * HV * V * K * 4 bytes / step for an fp32 state), while the
ReplaySSM kernel reads S every step but writes it only 1-in-L steps, plus a
small ring append (d:[HV,V], k:[H,K], g:[HV] per step). The amortized state
traffic ratio is reported per L.
Run::
python -m sglang.kernels.ops.attention.fla.bench_gdn_replayssm_decode
Requires a GPU (Triton).
"""
from __future__ import annotations
import argparse
import torch
import triton
from sglang.kernels.ops.attention.fla.fused_recurrent import (
fused_recurrent_gated_delta_rule_packed_decode,
)
from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import (
fused_recurrent_gdn_replayssm_decode,
)
def _make_static(B, H, HV, K, V, dtype, device):
qk_dim = 2 * H * K
v_dim = HV * V
mixed_qkv = torch.randn(B, qk_dim + v_dim, device=device, dtype=dtype)
a = torch.randn(B, HV, device=device, dtype=dtype) * 0.5
b = torch.randn(B, HV, device=device, dtype=dtype)
A_log = (torch.randn(HV, device=device, dtype=torch.float32) * 0.3).contiguous()
dt_bias = (torch.randn(HV, device=device, dtype=torch.float32) * 0.1).contiguous()
return mixed_qkv, a, b, A_log, dt_bias
def _state_bytes_per_step(B, HV, K, V, L, dtype):
"""Amortized per-step HBM *state* traffic (bytes), state in fp32.
packed: read S + write S every step.
replay: read S every step; write S once per L steps; append ring records
(d:[HV,V] in `dtype`, k:[H,K] in `dtype` shared across HV//H, g:[HV]
fp32) every step. We report the dominant fp32-state terms; ring
appends are tiny by comparison and shown separately.
"""
fp32 = 4
state_elems = B * HV * V * K # one record per active request slot
packed = (state_elems * fp32) * 2 # read + write
replay = (state_elems * fp32) * (1 + 1.0 / L) # read every step + write 1/L
return packed, replay
def _bench_cfg(B, H, HV, K, V, Ls, dtype, device, num_slots=None, warmup=25, rep=100):
num_slots = num_slots or B
mixed_qkv, a, b, A_log, dt_bias = _make_static(B, H, HV, K, V, dtype, device)
scale = K**-0.5
cache_indices = torch.arange(B, device=device, dtype=torch.int32)
# packed decode
state = torch.randn(num_slots, HV, V, K, device=device, dtype=torch.float32)
out = mixed_qkv.new_empty(B, 1, HV, V)
def run_packed():
fused_recurrent_gated_delta_rule_packed_decode(
mixed_qkv=mixed_qkv,
a=a,
b=b,
A_log=A_log,
dt_bias=dt_bias,
scale=scale,
initial_state=state,
out=out,
ssm_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True,
)
t_packed = triton.testing.do_bench(run_packed, warmup=warmup, rep=rep)
rows = []
for L in Ls:
rstate = torch.randn(num_slots, HV, V, K, device=device, dtype=torch.float32)
d_cache = torch.zeros(num_slots, HV, L, V, device=device, dtype=dtype)
k_cache = torch.zeros(num_slots, H, L, K, device=device, dtype=dtype)
g_cache = torch.zeros(num_slots, HV, L, device=device, dtype=torch.float32)
write_pos = torch.zeros(B, device=device, dtype=torch.int32)
rout = mixed_qkv.new_empty(B, 1, HV, V)
nk = 1 if L == 1 else 2
def run_replay():
fused_recurrent_gdn_replayssm_decode(
mixed_qkv=mixed_qkv,
a=a,
b=b,
A_log=A_log,
dt_bias=dt_bias,
scale=scale,
initial_state=rstate,
d_cache=d_cache,
k_cache=k_cache,
g_cache=g_cache,
out=rout,
ssm_state_indices=cache_indices,
write_pos=write_pos,
use_qk_l2norm_in_kernel=True,
nk=nk,
)
t_replay = triton.testing.do_bench(run_replay, warmup=warmup, rep=rep)
packed_bytes, replay_bytes = _state_bytes_per_step(B, HV, K, V, L, dtype)
rows.append((L, t_replay, t_packed / t_replay, replay_bytes / packed_bytes))
return t_packed, rows
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--hv", type=int, default=32, help="num value heads")
parser.add_argument("--h", type=int, default=16, help="num key/query heads")
parser.add_argument("--k", type=int, default=128)
parser.add_argument("--v", type=int, default=128)
parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 16, 64, 256])
parser.add_argument("--ls", type=int, nargs="+", default=[1, 8, 16])
parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16")
args = parser.parse_args()
if not torch.cuda.is_available():
raise SystemExit("CUDA / Triton required for this microbenchmark.")
device = "cuda"
dtype = {
"bf16": torch.bfloat16,
"fp16": torch.float16,
"fp32": torch.float32,
}[args.dtype]
print(
f"GDN ReplaySSM decode microbench HV={args.hv} H={args.h} "
f"K={args.k} V={args.v} dtype={args.dtype}\n"
"per-step latency (ms); speedup = packed/replay; "
"state-traffic = replay/packed (lower is better)"
)
for B in args.batch_sizes:
t_packed, rows = _bench_cfg(
B, args.h, args.hv, args.k, args.v, args.ls, dtype, device
)
print(f"\nB={B:<4d} packed={t_packed:.4f} ms")
for L, t_replay, speedup, traffic_ratio in rows:
print(
f" L={L:<3d} replay={t_replay:.4f} ms "
f"speedup={speedup:5.2f}x "
f"state-traffic={traffic_ratio:5.2f}x"
)
if __name__ == "__main__":
main()
python/sglang/kernels/ops/attention/flash_attn/cute/bench_utils.py
-261
View File
@@ -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
python/sglang/kernels/ops/attention/flash_attn/cute/benchmark.py
-281
View File
@@ -1,281 +0,0 @@
# Copyright (c) 2023, Tri Dao.
"""Useful functions for writing test code."""
import torch
import torch.utils.benchmark as benchmark
def benchmark_forward(
fn,
*inputs,
repeats=10,
desc="",
verbose=True,
amp=False,
amp_dtype=torch.float16,
**kwinputs,
):
"""Use Pytorch Benchmark on the forward pass of an arbitrary function."""
if verbose:
print(desc, "- Forward pass")
def amp_wrapper(*inputs, **kwinputs):
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
fn(*inputs, **kwinputs)
t = benchmark.Timer(
stmt="fn_amp(*inputs, **kwinputs)",
globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
num_threads=torch.get_num_threads(),
)
m = t.timeit(repeats)
if verbose:
print(m)
return t, m
def benchmark_backward(
fn,
*inputs,
grad=None,
repeats=10,
desc="",
verbose=True,
amp=False,
amp_dtype=torch.float16,
**kwinputs,
):
"""Use Pytorch Benchmark on the backward pass of an arbitrary function."""
if verbose:
print(desc, "- Backward pass")
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
y = fn(*inputs, **kwinputs)
if type(y) is tuple:
y = y[0]
if grad is None:
grad = torch.randn_like(y)
else:
if grad.shape != y.shape:
raise RuntimeError("Grad shape does not match output shape")
def f(*inputs, y, grad):
# Set .grad to None to avoid extra operation of gradient accumulation
for x in inputs:
if isinstance(x, torch.Tensor):
x.grad = None
y.backward(grad, retain_graph=True)
t = benchmark.Timer(
stmt="f(*inputs, y=y, grad=grad)",
globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
num_threads=torch.get_num_threads(),
)
m = t.timeit(repeats)
if verbose:
print(m)
return t, m
def benchmark_combined(
fn,
*inputs,
grad=None,
repeats=10,
desc="",
verbose=True,
amp=False,
amp_dtype=torch.float16,
**kwinputs,
):
"""Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
if verbose:
print(desc, "- Forward + Backward pass")
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
y = fn(*inputs, **kwinputs)
if type(y) is tuple:
y = y[0]
if grad is None:
grad = torch.randn_like(y)
else:
if grad.shape != y.shape:
raise RuntimeError("Grad shape does not match output shape")
def f(grad, *inputs, **kwinputs):
for x in inputs:
if isinstance(x, torch.Tensor):
x.grad = None
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
y = fn(*inputs, **kwinputs)
if type(y) is tuple:
y = y[0]
y.backward(grad, retain_graph=True)
t = benchmark.Timer(
stmt="f(grad, *inputs, **kwinputs)",
globals={
"f": f,
"fn": fn,
"inputs": inputs,
"grad": grad,
"kwinputs": kwinputs,
},
num_threads=torch.get_num_threads(),
)
m = t.timeit(repeats)
if verbose:
print(m)
return t, m
def benchmark_fwd_bwd(
fn,
*inputs,
grad=None,
repeats=10,
desc="",
verbose=True,
amp=False,
amp_dtype=torch.float16,
**kwinputs,
):
"""Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
return (
benchmark_forward(
fn,
*inputs,
repeats=repeats,
desc=desc,
verbose=verbose,
amp=amp,
amp_dtype=amp_dtype,
**kwinputs,
),
benchmark_backward(
fn,
*inputs,
grad=grad,
repeats=repeats,
desc=desc,
verbose=verbose,
amp=amp,
amp_dtype=amp_dtype,
**kwinputs,
),
)
def benchmark_all(
fn,
*inputs,
grad=None,
repeats=10,
desc="",
verbose=True,
amp=False,
amp_dtype=torch.float16,
**kwinputs,
):
"""Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
return (
benchmark_forward(
fn,
*inputs,
repeats=repeats,
desc=desc,
verbose=verbose,
amp=amp,
amp_dtype=amp_dtype,
**kwinputs,
),
benchmark_backward(
fn,
*inputs,
grad=grad,
repeats=repeats,
desc=desc,
verbose=verbose,
amp=amp,
amp_dtype=amp_dtype,
**kwinputs,
),
benchmark_combined(
fn,
*inputs,
grad=grad,
repeats=repeats,
desc=desc,
verbose=verbose,
amp=amp,
amp_dtype=amp_dtype,
**kwinputs,
),
)
def pytorch_profiler(
fn,
*inputs,
trace_filename=None,
backward=False,
amp=False,
amp_dtype=torch.float16,
cpu=False,
verbose=True,
**kwinputs,
):
"""Wrap benchmark functions in Pytorch profiler to see CUDA information."""
if backward:
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
out = fn(*inputs, **kwinputs)
if type(out) is tuple:
out = out[0]
g = torch.randn_like(out)
for _ in range(30): # Warm up
if backward:
for x in inputs:
if isinstance(x, torch.Tensor):
x.grad = None
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
out = fn(*inputs, **kwinputs)
if type(out) is tuple:
out = out[0]
# Backward should be done outside autocast
if backward:
out.backward(g, retain_graph=True)
activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
torch.profiler.ProfilerActivity.CUDA
]
with torch.profiler.profile(
activities=activities,
record_shapes=True,
# profile_memory=True,
with_stack=True,
) as prof:
if backward:
for x in inputs:
if isinstance(x, torch.Tensor):
x.grad = None
with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
out = fn(*inputs, **kwinputs)
if type(out) is tuple:
out = out[0]
if backward:
out.backward(g, retain_graph=True)
if verbose:
# print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
print(prof.key_averages().table(row_limit=50))
if trace_filename is not None:
prof.export_chrome_trace(trace_filename)
def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
fn(*inputs, **kwinputs)
torch.cuda.synchronize()
mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
if verbose:
print(f"{desc} max memory: {mem}GB")
torch.cuda.empty_cache()
return mem
python/sglang/kernels/ops/attention/flash_attn/cute/benchmark_flash_attention_fp8.py
-487
View File
@@ -1,487 +0,0 @@
# Benchmark FP8 attention for FA4 (CuTe-DSL) on SM100.
#
# Run (recommended):
# python -m flash_attn.cute.benchmark_flash_attention_fp8
#
# Notes:
# - This is intended to be used while bringing up FP8 support for SM100.
# - FP8 correctness depends on descales + max-offset scaling being implemented in the SM100 kernel.
# This script optionally checks output vs a BF16 PyTorch baseline on dequantized FP8 inputs.
#
# Adapted from: `hopper/benchmark_flash_attention_fp8.py`
from __future__ import annotations
import argparse
import inspect
import math
import time
from typing import Iterable
import torch
from einops import rearrange
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,
)
try:
import cudnn
except ImportError:
cudnn = None
def _torch_float8_dtype(name: str) -> torch.dtype:
if name in ("fp8", "fp8_e4m3", "fp8_e4m3fn"):
return torch.float8_e4m3fn
if name in ("fp8_e5m2", "fp8_e5m2fn"):
return torch.float8_e5m2
raise ValueError(f"Unsupported fp8 dtype name: {name}")
def _parse_int_list(csv: str) -> list[int]:
out: list[int] = []
for part in csv.split(","):
part = part.strip()
if not part:
continue
out.append(int(part))
return out
def attention_pytorch(qkv: torch.Tensor, causal: bool) -> torch.Tensor:
"""
qkv: (batch, seqlen, 3, nheads, headdim)
out: (batch, seqlen, nheads, headdim)
"""
batch_size, seqlen, _, nheads, d = qkv.shape
q, k, v = qkv.unbind(dim=2)
q = rearrange(q, "b t h d -> (b h) t d")
k = rearrange(k, "b s h d -> (b h) d s")
softmax_scale = 1.0 / math.sqrt(d)
scores = torch.empty(
batch_size * nheads, seqlen, seqlen, dtype=qkv.dtype, device=qkv.device
)
scores = rearrange(
torch.baddbmm(scores, q, k, beta=0, alpha=softmax_scale),
"(b h) t s -> b h t s",
h=nheads,
)
if causal:
causal_mask = torch.triu(
torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1
)
scores = scores + causal_mask.to(dtype=scores.dtype)
attention = torch.softmax(scores, dim=-1)
output = torch.einsum("bhts,bshd->bthd", attention, v)
return output.to(dtype=qkv.dtype)
def flops(batch: int, seqlen: int, headdim: int, nheads: int, causal: bool) -> int:
# Matches the hopper benchmark’s convention.
return 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
def efficiency(flop: int, seconds: float) -> float:
return (flop / seconds / 1e12) if not math.isnan(seconds) else 0.0
def time_fwd(fn, *args, repeats: int, **kwargs) -> float:
time.sleep(1) # reduce residual throttling effects between benchmarks
_, m = benchmark_forward(fn, *args, repeats=repeats, verbose=False, **kwargs)
return float(m.mean)
def convert_to_cudnn_type(torch_type):
if torch_type == torch.float16:
return cudnn.data_type.HALF
if torch_type == torch.bfloat16:
return cudnn.data_type.BFLOAT16
if torch_type == torch.float32:
return cudnn.data_type.FLOAT
if torch_type == torch.int32:
return cudnn.data_type.INT32
if torch_type == torch.int64:
return cudnn.data_type.INT64
if torch_type == torch.float8_e4m3fn:
return cudnn.data_type.FP8_E4M3
if torch_type == torch.float8_e5m2:
return cudnn.data_type.FP8_E5M2
raise ValueError("Unsupported tensor data type.")
def cudnn_sdpa_fp8_setup(qkv: torch.Tensor, seqlen_q: int, seqlen_k: int, causal: bool):
"""Minimal cudnn.fp8 sdpa runner (optional)."""
assert cudnn is not None, "cudnn python bindings not available"
b, _, _, nheads, headdim = qkv.shape
o_gpu = torch.zeros(
b, seqlen_q, nheads, headdim, dtype=qkv.dtype, device=qkv.device
)
o_gpu_transposed = torch.as_strided(
o_gpu,
[b, nheads, seqlen_q, headdim],
[nheads * seqlen_q * headdim, headdim, nheads * headdim, 1],
)
amax_s_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
amax_o_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
graph = cudnn.pygraph(
io_data_type=convert_to_cudnn_type(qkv.dtype),
intermediate_data_type=cudnn.data_type.FLOAT,
compute_data_type=cudnn.data_type.FLOAT,
)
new_q = torch.as_strided(
qkv,
[b, nheads, seqlen_q, headdim],
[seqlen_q * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
storage_offset=0,
)
q = graph.tensor(
name="Q",
dim=list(new_q.shape),
stride=list(new_q.stride()),
data_type=convert_to_cudnn_type(qkv.dtype),
)
new_k = torch.as_strided(
qkv,
[b, nheads, seqlen_k, headdim],
[seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
storage_offset=nheads * headdim,
)
k = graph.tensor(
name="K",
dim=list(new_k.shape),
stride=list(new_k.stride()),
data_type=convert_to_cudnn_type(qkv.dtype),
)
new_v = torch.as_strided(
qkv,
[b, nheads, seqlen_k, headdim],
[seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
storage_offset=nheads * headdim * 2,
)
v = graph.tensor(
name="V",
dim=list(new_v.shape),
stride=list(new_v.stride()),
data_type=convert_to_cudnn_type(qkv.dtype),
)
def _scale_tensor():
return graph.tensor(
dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT
)
default_scale_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float32, device="cuda")
descale_q = _scale_tensor()
descale_k = _scale_tensor()
descale_v = _scale_tensor()
descale_s = _scale_tensor()
scale_s = _scale_tensor()
scale_o = _scale_tensor()
o, _, amax_s, amax_o = graph.sdpa_fp8(
q=q,
k=k,
v=v,
descale_q=descale_q,
descale_k=descale_k,
descale_v=descale_v,
descale_s=descale_s,
scale_s=scale_s,
scale_o=scale_o,
is_inference=True,
attn_scale=1.0 / math.sqrt(headdim),
use_causal_mask=causal,
name="sdpa",
)
o.set_output(True).set_dim(o_gpu_transposed.shape).set_stride(
o_gpu_transposed.stride()
)
amax_s.set_output(False).set_dim(amax_s_gpu.shape).set_stride(amax_s_gpu.stride())
amax_o.set_output(False).set_dim(amax_o_gpu.shape).set_stride(amax_o_gpu.stride())
graph.validate()
graph.build_operation_graph()
graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
graph.check_support()
graph.build_plans()
variant_pack = {
q: new_q,
k: new_k,
v: new_v,
descale_q: default_scale_gpu,
descale_k: default_scale_gpu,
descale_v: default_scale_gpu,
descale_s: default_scale_gpu,
scale_s: default_scale_gpu,
scale_o: default_scale_gpu,
o: o_gpu_transposed,
amax_s: amax_s_gpu,
amax_o: amax_o_gpu,
}
workspace = torch.empty(
graph.get_workspace_size(), device="cuda", dtype=torch.uint8
)
def run():
graph.execute(variant_pack, workspace)
return o_gpu
return run
def _maybe_pass_descales(callable_, **kwargs):
sig = inspect.signature(callable_)
return {k: v for k, v in kwargs.items() if k in sig.parameters}
def main(argv: Iterable[str] | None = None) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--repeats", type=int, default=30)
parser.add_argument("--dim", type=int, default=2048)
parser.add_argument("--headdims", default="64,128")
parser.add_argument("--dtype", default="fp8_e4m3fn")
parser.add_argument("--seed", type=int, default=0)
parser.add_argument(
"--check",
action=argparse.BooleanOptionalAction,
default=True,
help="Enable correctness checks vs BF16 PyTorch baseline.",
)
parser.add_argument(
"--check-quantization-only",
action="store_true",
help="Check FP8 kernel vs dequantized-FP8 baseline (quantization error only).",
)
parser.add_argument("--atol-bf16", type=float, default=0.10)
parser.add_argument("--rtol-bf16", type=float, default=0.10)
parser.add_argument("--atol-fp8", type=float, default=0.50)
parser.add_argument("--rtol-fp8", type=float, default=0.50)
parser.add_argument("--run-cudnn", action="store_true")
args = parser.parse_args(list(argv) if argv is not None else None)
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required")
major, minor = torch.cuda.get_device_capability()
if major != 10:
raise RuntimeError(
f"This benchmark is for SM100 (compute capability 10.x). Got {major}.{minor}."
)
torch.manual_seed(args.seed)
device = "cuda"
fp8_dtype = _torch_float8_dtype(args.dtype)
headdim_vals = _parse_int_list(args.headdims)
bs_seqlen_vals = [
(32, 512),
(16, 1024),
(8, 2048),
(4, 4096),
(2, 8192),
(1, 16384),
]
methods = ["Pytorch", "FA4-CuTe-BF16", "FA4-CuTe-FP8"] + (
["cuDNN-FP8"] if args.run_cudnn and cudnn is not None else []
)
fp8_failures = []
for headdim in headdim_vals:
for causal in (False, True):
for batch, seqlen in bs_seqlen_vals:
torch.cuda.empty_cache()
nheads = args.dim // headdim
if args.dim % headdim != 0:
raise ValueError(
f"--dim must be divisible by headdim ({args.dim=} {headdim=})"
)
q_bf16 = torch.randn(
batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
)
k_bf16 = torch.randn(
batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
)
v_bf16 = torch.randn(
batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
)
qkv_bf16 = torch.stack([q_bf16, k_bf16, v_bf16], dim=2)
times = {}
speeds = {}
out_ref_bf16 = None
try:
out_ref_bf16 = attention_pytorch(
qkv_bf16, causal=causal
) # warmup / reference
t = time_fwd(
attention_pytorch, qkv_bf16, causal=causal, repeats=args.repeats
)
times["Pytorch"] = t
except RuntimeError as e:
if "out of memory" in str(e).lower():
times["Pytorch"] = float("nan")
out_ref_bf16 = None
else:
raise
# FA4 / CuTe BF16 baseline
try:
softmax_scale = headdim**-0.5
out_fa4_bf16, _ = flash_attn_cute_fwd(
q_bf16,
k_bf16,
v_bf16,
softmax_scale=softmax_scale,
causal=causal,
) # warmup / compile
t = time_fwd(
flash_attn_cute_fwd,
q_bf16,
k_bf16,
v_bf16,
softmax_scale=softmax_scale,
causal=causal,
repeats=args.repeats,
)
times["FA4-CuTe-BF16"] = t
if args.check and out_ref_bf16 is not None:
torch.testing.assert_close(
out_fa4_bf16,
out_ref_bf16,
atol=args.atol_bf16,
rtol=args.rtol_bf16,
)
except Exception as e:
# Treat as fatal: BF16 kernel should be usable for basic sanity checking.
raise RuntimeError("FA4-CuTe BF16 baseline failed") from e
# FA4 / CuTe FP8
q_fp8 = q_bf16.to(fp8_dtype)
k_fp8 = k_bf16.to(fp8_dtype)
v_fp8 = v_bf16.to(fp8_dtype)
# Placeholder descales (FA3-style: per-(batch, kv_head)).
q_descale = torch.ones(
batch, nheads, device=device, dtype=torch.float32
)
k_descale = torch.ones(
batch, nheads, device=device, dtype=torch.float32
)
v_descale = torch.ones(
batch, nheads, device=device, dtype=torch.float32
)
# Optional: FP8 reference baseline (dequantized FP8 -> PyTorch) for quantization-error-only checks
out_ref_fp8 = None
if args.check and args.check_quantization_only:
try:
# Dequantize FP8 inputs back to BF16 (applying descales)
q_ref_fp8 = (
q_fp8.to(torch.bfloat16) * q_descale[:, None, :, None]
).to(torch.bfloat16)
k_ref_fp8 = (
k_fp8.to(torch.bfloat16) * k_descale[:, None, :, None]
).to(torch.bfloat16)
v_ref_fp8 = (
v_fp8.to(torch.bfloat16) * v_descale[:, None, :, None]
).to(torch.bfloat16)
qkv_ref_fp8 = torch.stack(
[q_ref_fp8, k_ref_fp8, v_ref_fp8], dim=2
)
out_ref_fp8 = attention_pytorch(qkv_ref_fp8, causal=causal)
except RuntimeError as e:
if "out of memory" in str(e).lower():
out_ref_fp8 = None
else:
raise
fa4_kwargs = dict(softmax_scale=softmax_scale, causal=causal)
fa4_kwargs.update(
_maybe_pass_descales(
flash_attn_cute_fwd,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
)
)
try:
# Warmup/compile (will raise until FP8 is implemented)
out_fa4_fp8, _ = flash_attn_cute_fwd(
q_fp8, k_fp8, v_fp8, **fa4_kwargs
)
t = time_fwd(
flash_attn_cute_fwd,
q_fp8,
k_fp8,
v_fp8,
repeats=args.repeats,
**fa4_kwargs,
)
times["FA4-CuTe-FP8"] = t
if args.check:
# Choose baseline: quantization-only (dequantized FP8) or full (BF16)
if args.check_quantization_only:
ref_baseline = out_ref_fp8
else:
ref_baseline = out_ref_bf16
if ref_baseline is not None:
torch.testing.assert_close(
out_fa4_fp8,
ref_baseline,
atol=args.atol_fp8,
rtol=args.rtol_fp8,
)
except Exception as e:
fp8_failures.append((causal, headdim, batch, seqlen, repr(e)))
times["FA4-CuTe-FP8"] = float("nan")
if args.run_cudnn and cudnn is not None:
qkv_fp8 = qkv_bf16.to(fp8_dtype)
runner = cudnn_sdpa_fp8_setup(
qkv_fp8, seqlen, seqlen, causal=causal
)
_ = runner() # warmup
t = time_fwd(lambda: runner(), repeats=args.repeats)
times["cuDNN-FP8"] = t
print(
f"### causal={causal}, headdim={headdim}, batch={batch}, seqlen={seqlen} ###"
)
for method in methods:
t = times.get(method, float("nan"))
speeds[method] = efficiency(
flops(batch, seqlen, headdim, nheads, causal), t
)
if math.isnan(t):
print(f"{method} fwd: (skipped)")
else:
print(
f"{method} fwd: {speeds[method]:.2f} TFLOPs/s, {t * 1e3:.3f} ms"
)
if math.isnan(times.get("FA4-CuTe-FP8", float("nan"))):
print("FA4-CuTe-FP8 status: FAILED")
if fp8_failures:
print(f"\nFP8 failures: {len(fp8_failures)} (showing first 5)")
for causal, headdim, batch, seqlen, err in fp8_failures[:5]:
print(
f"- causal={causal} headdim={headdim} batch={batch} seqlen={seqlen}: {err}"
)
return 1
return 0
if __name__ == "__main__":
raise SystemExit(main())
python/sglang/kernels/ops/attention/flash_attn/cute/sm90_config_search.py
-416
View File
@@ -1,416 +0,0 @@
"""Search feasible SM90 fwd/bwd attention configs for given (head_dim, head_dim_v).
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
"""
import math
# H100 hardware limits
SMEM_LIMIT = 224 * 1024 # 228 KB minus ~3 KB for LSE, dPsum, mbarriers
REG_LIMITS = {2: 216, 3: 128} # per-WG budget: 2WG=240-24, 3WG=160-32
THREADS_PER_WG = 128
def _divisors(n):
return [d for d in range(1, n + 1) if n % d == 0]
def _acc_regs(M, N, num_wg):
"""Accumulator registers per thread per WG."""
return M * N // (num_wg * THREADS_PER_WG)
def _check_mma(M, N, num_wg, atom_layout_m, swap_AB):
"""Check MMA feasibility. Returns regs per WG, or None if infeasible.
GMMA atom M=64. Swap exchanges (M, N) and atom layout.
Requires: M divisible by (atom_layout_m * 64), N by (atom_layout_n * 8).
"""
if swap_AB:
M, N = N, M
atom_layout_m = num_wg // atom_layout_m
atom_layout_n = num_wg // atom_layout_m
if M % (atom_layout_m * 64) != 0 or N % (atom_layout_n * 8) != 0:
return None
return _acc_regs(M, N, num_wg)
def _mma_traffic(M_eff, N_eff, K_red, num_wg, wg_n, is_rs=False):
"""Total SMEM read traffic for one MMA (all WGs combined).
num_instr = (M_eff / 64) * wg_n instructions total.
Each reads A(64, K_red) and B(N_eff/wg_n, K_red) from smem (bf16).
"""
num_instr = (M_eff // 64) * wg_n
A_per = 64 * K_red * 2 if not is_rs else 0
B_per = (N_eff // wg_n) * K_red * 2
return num_instr * (A_per + B_per)
# ============================================================================
# Backward
# ============================================================================
def _check_bwd_config(
hdim,
hdimv,
tile_m,
tile_n,
num_wg,
SdP_swapAB,
dKV_swapAB,
dQ_swapAB,
AtomLayoutMSdP,
AtomLayoutNdKV,
AtomLayoutMdQ,
):
reg_limit = REG_LIMITS[num_wg]
# MMA feasibility
regs_SdP = _check_mma(tile_m, tile_n, num_wg, AtomLayoutMSdP, SdP_swapAB)
regs_dK = _check_mma(tile_n, hdim, num_wg, AtomLayoutNdKV, dKV_swapAB)
regs_dV = _check_mma(tile_n, hdimv, num_wg, AtomLayoutNdKV, dKV_swapAB)
regs_dQ = _check_mma(tile_m, hdim, num_wg, AtomLayoutMdQ, dQ_swapAB)
if any(r is None for r in (regs_SdP, regs_dK, regs_dV, regs_dQ)):
return None
# Peak regs: max(S+dP, dQ) + dK + dV
total_regs = max(2 * regs_SdP, regs_dQ) + regs_dK + regs_dV
if total_regs > reg_limit:
return None
# SMEM
mma_dkv_is_rs = (
AtomLayoutMSdP == 1
and AtomLayoutNdKV == num_wg
and SdP_swapAB
and not dKV_swapAB
)
Q_stage, PdS_stage = 2, 1
for dO_stage in (2, 1):
sQ = tile_m * hdim * 2 * Q_stage
sK = tile_n * hdim * 2
sV = tile_n * hdimv * 2
sdO = tile_m * hdimv * 2 * dO_stage
sPdS = tile_m * tile_n * 2 * PdS_stage
sP = sPdS if not mma_dkv_is_rs else 0
sdQaccum = tile_m * hdim * 4
smem = sQ + sK + sV + sdO + sP + sPdS + sdQaccum
if smem <= SMEM_LIMIT:
break
else:
return None
# SMEM traffic
def _swap(a, b, s):
return (b, a) if s else (a, b)
def _wg_n(al_m, s):
return al_m if s else num_wg // al_m
M_s, N_s = _swap(tile_m, tile_n, SdP_swapAB)
wn_SdP = _wg_n(AtomLayoutMSdP, SdP_swapAB)
traffic_S = _mma_traffic(M_s, N_s, hdim, num_wg, wn_SdP)
traffic_dP = _mma_traffic(M_s, N_s, hdimv, num_wg, wn_SdP)
wn_dKV = _wg_n(AtomLayoutNdKV, dKV_swapAB)
M_dv, N_dv = _swap(tile_n, hdimv, dKV_swapAB)
traffic_dV = _mma_traffic(M_dv, N_dv, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs)
M_dk, N_dk = _swap(tile_n, hdim, dKV_swapAB)
traffic_dK = _mma_traffic(M_dk, N_dk, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs)
M_dq, N_dq = _swap(tile_m, hdim, dQ_swapAB)
wn_dQ = _wg_n(AtomLayoutMdQ, dQ_swapAB)
traffic_dQ = _mma_traffic(M_dq, N_dq, tile_n, num_wg, wn_dQ)
traffic_P_store = tile_m * tile_n * 2 if not mma_dkv_is_rs else 0
traffic_dS_store = tile_m * tile_n * 2
traffic_dQ_smem = tile_m * hdim * 4 * 2 # store + TMA load
smem_traffic = (
traffic_S
+ traffic_dP
+ traffic_dV
+ traffic_dK
+ traffic_dQ
+ traffic_P_store
+ traffic_dS_store
+ traffic_dQ_smem
)
return dict(
tile_m=tile_m,
tile_n=tile_n,
num_wg=num_wg,
Q_stage=Q_stage,
dO_stage=dO_stage,
PdS_stage=PdS_stage,
SdP_swapAB=SdP_swapAB,
dKV_swapAB=dKV_swapAB,
dQ_swapAB=dQ_swapAB,
AtomLayoutMSdP=AtomLayoutMSdP,
AtomLayoutNdKV=AtomLayoutNdKV,
AtomLayoutMdQ=AtomLayoutMdQ,
mma_dkv_is_rs=mma_dkv_is_rs,
regs_SdP=regs_SdP,
regs_dK=regs_dK,
regs_dV=regs_dV,
regs_dQ=regs_dQ,
total_regs=total_regs,
reg_limit=reg_limit,
smem_bytes=smem,
smem_kb=smem / 1024,
smem_traffic=smem_traffic,
smem_traffic_kb=smem_traffic / 1024,
smem_traffic_per_block=smem_traffic / (tile_m * tile_n),
)
def find_feasible_bwd_configs(
head_dim,
head_dim_v=None,
tile_m_choices=(64, 80, 96, 112, 128),
tile_n_choices=(64, 80, 96, 112, 128),
):
if head_dim_v is None:
head_dim_v = head_dim
hdim = int(math.ceil(head_dim / 32) * 32)
hdimv = int(math.ceil(head_dim_v / 32) * 32)
results = []
for num_wg in (2, 3):
divs = _divisors(num_wg)
for tile_m in tile_m_choices:
for tile_n in tile_n_choices:
for SdP_swap in (False, True):
if (tile_n if SdP_swap else tile_m) % 64 != 0:
continue
for dKV_swap in (False, True):
if not dKV_swap and tile_n % 64 != 0:
continue
if dKV_swap and (hdim % 64 != 0 or hdimv % 64 != 0):
continue
for dQ_swap in (False, True):
if (hdim if dQ_swap else tile_m) % 64 != 0:
continue
for a1 in divs:
for a2 in divs:
for a3 in divs:
cfg = _check_bwd_config(
hdim,
hdimv,
tile_m,
tile_n,
num_wg,
SdP_swap,
dKV_swap,
dQ_swap,
a1,
a2,
a3,
)
if cfg is not None:
results.append(cfg)
results.sort(
key=lambda c: (-c["tile_n"], -c["tile_m"], c["smem_traffic_per_block"])
)
return results
def print_bwd_configs(configs, max_results=20):
if not configs:
print("No feasible configs found!")
return
n = min(len(configs), max_results)
print(f"Found {len(configs)} feasible configs (showing top {n}):\n")
hdr = (
f"{'wg':>2} {'tm':>3} {'tn':>3} "
f"{'SdP':>3} {'dKV':>3} {'dQ':>3} "
f"{'aSdP':>4} {'adKV':>4} {'adQ':>4} "
f"{'Qs':>2} {'dOs':>3} "
f"{'rS':>3} {'rdK':>3} {'rdV':>3} {'rdQ':>3} {'tot':>4}/{'':<3} "
f"{'smem':>5} {'traffic':>7} {'tr/blk':>6}"
)
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"{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} "
f"{c['total_regs']:>4}/{c['reg_limit']:<3} "
f"{c['smem_kb']:>4.0f}K "
f"{c['smem_traffic_kb']:>6.0f}K "
f"{c['smem_traffic_per_block']:>6.1f}"
)
# ============================================================================
# Forward
# ============================================================================
def _check_fwd_config(hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg):
reg_limit = REG_LIMITS[num_wg]
tile_m = num_wg * 64
if tile_n % 8 != 0:
return None
regs_S = _acc_regs(tile_m, tile_n, num_wg)
regs_O = _acc_regs(tile_m, hdimv, num_wg)
regs_P = regs_S // 2 # bf16 = half of f32
if overlap_wg:
total_regs = regs_S + regs_P + regs_O
else:
total_regs = regs_S + regs_O
if total_regs > reg_limit:
return None
# SMEM: 1 stage Q, 2 stages K/V, O overlaps Q, sP if not RS
sQ = tile_m * hdim * 2
sK = tile_n * hdim * 2 * 2
sV = tile_n * hdimv * 2 * 2
sO = tile_m * hdimv * 2
sP = tile_m * tile_n * 2 if not pv_is_rs else 0
smem = max(sQ, sO) + sK + sV + sP
if smem > SMEM_LIMIT:
return None
# SMEM traffic: num_instr = num_wg (all WGs in M, wg_n=1)
traffic_S = num_wg * (64 * hdim * 2 + tile_n * hdim * 2)
A_pv = 64 * tile_n * 2 if not pv_is_rs else 0
traffic_O = num_wg * (A_pv + hdimv * tile_n * 2)
traffic_P_store = tile_m * tile_n * 2 if not pv_is_rs else 0
smem_traffic = traffic_S + traffic_O + traffic_P_store
return dict(
tile_m=tile_m,
tile_n=tile_n,
num_wg=num_wg,
pv_is_rs=pv_is_rs,
overlap_wg=overlap_wg,
regs_S=regs_S,
regs_O=regs_O,
regs_P=regs_P,
total_regs=total_regs,
reg_limit=reg_limit,
smem_bytes=smem,
smem_kb=smem / 1024,
smem_traffic=smem_traffic,
smem_traffic_kb=smem_traffic / 1024,
smem_traffic_per_block=smem_traffic / (tile_m * tile_n),
)
def find_feasible_fwd_configs(
head_dim, head_dim_v=None, tile_n_choices=(64, 80, 96, 112, 128, 144, 160, 176, 192)
):
if head_dim_v is None:
head_dim_v = head_dim
hdim = int(math.ceil(head_dim / 32) * 32)
hdimv = int(math.ceil(head_dim_v / 32) * 32)
results = []
for num_wg in (2, 3):
for tile_n in tile_n_choices:
for pv_is_rs in (True, False):
for overlap_wg in (True, False):
cfg = _check_fwd_config(
hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg
)
if cfg is not None:
results.append(cfg)
results.sort(key=lambda c: (-c["tile_n"], c["smem_traffic_per_block"]))
return results
def print_fwd_configs(configs, max_results=20):
if not configs:
print("No feasible configs found!")
return
n = min(len(configs), max_results)
print(f"Found {len(configs)} feasible configs (showing top {n}):\n")
hdr = (
f"{'wg':>2} {'tm':>3} {'tn':>3} "
f"{'RS':>2} {'olap':>4} "
f"{'rS':>3} {'rP':>3} {'rO':>3} {'tot':>4}/{'':<3} "
f"{'smem':>5} {'traffic':>7} {'tr/blk':>6}"
)
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"{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 "
f"{c['smem_traffic_kb']:>6.0f}K "
f"{c['smem_traffic_per_block']:>6.1f}"
)
# ============================================================================
# CLI
# ============================================================================
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Search feasible SM90 MMA configs")
parser.add_argument("--mode", choices=["fwd", "bwd", "both"], default="both")
parser.add_argument(
"--headdim",
type=str,
default="128",
help="Head dim, or hdim-hdimv (e.g. 192-128)",
)
parser.add_argument(
"--tile-m", type=str, default="64,80,96,112,128", help="Bwd tile_m choices"
)
parser.add_argument(
"--tile-n",
type=str,
default=None,
help="tile_n choices (default: fwd up to 192, bwd up to 128)",
)
parser.add_argument("-n", "--num-results", type=int, default=30)
args = parser.parse_args()
parts = args.headdim.split("-")
hdim = int(parts[0])
hdimv = int(parts[1]) if len(parts) > 1 else hdim
TN_FWD = "64,80,96,112,128,144,160,176,192"
TN_BWD = "64,80,96,112,128"
if args.mode in ("fwd", "both"):
tn = tuple(int(x) for x in (args.tile_n or TN_FWD).split(","))
print(f"=== FWD configs: hdim={hdim}, hdimv={hdimv} ===\n")
print_fwd_configs(find_feasible_fwd_configs(hdim, hdimv, tn), args.num_results)
print()
if args.mode in ("bwd", "both"):
tm = tuple(int(x) for x in args.tile_m.split(","))
tn = tuple(int(x) for x in (args.tile_n or TN_BWD).split(","))
print(f"=== BWD configs: hdim={hdim}, hdimv={hdimv} ===\n")
print_bwd_configs(
find_feasible_bwd_configs(hdim, hdimv, tm, tn), args.num_results
)
python/sglang/kernels/registry.py
+11 -4
View File
@@ -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