358 lines
12 KiB
Python
358 lines
12 KiB
Python
"""
|
|
Benchmark & Correctness: Triton KDA vs CuTeDSL KDA (prefill, SM100 Blackwell).
|
|
|
|
Compares:
|
|
- Triton: sglang's chunk_kda (FLA chunkwise gated delta rule, per-channel gate)
|
|
- CuteDSL: kda_blackwell pipeline (fused Triton prologue -> kkt_inv_uw -> h -> o)
|
|
|
|
KDA differs from GDN by a PER-CHANNEL decay gate (g is [T, H, K], not scalar).
|
|
The cutedsl pipeline externalizes the per-channel decay into five pre-scaled
|
|
key/query tensors computed by a fused Triton prologue; the chunk metadata is
|
|
computed once and shared across layers in a real forward, so the benchmarked
|
|
cutedsl path precomputes it outside the timed region (the realistic ceiling).
|
|
|
|
Correctness is checked against the token-by-token fused_recurrent_kda ground
|
|
truth. Reports performance (ms, approx TFLOPS, TB/s, speedup).
|
|
|
|
Usage:
|
|
python bench_kda_prefill_cutedsl.py # default sweep
|
|
python bench_kda_prefill_cutedsl.py --mode bench # benchmark only
|
|
python bench_kda_prefill_cutedsl.py --mode correctness # correctness only
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python"))
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from sglang.kernels.ops.attention.fla.kda import chunk_kda, fused_recurrent_kda
|
|
from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata
|
|
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import (
|
|
kda_h_cutedsl,
|
|
)
|
|
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_kkt_inv_uw import (
|
|
kkt_inv_uw_cutedsl,
|
|
)
|
|
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_o import (
|
|
kda_o_cutedsl,
|
|
)
|
|
from sglang.kernels.ops.attention.linear.kda_blackwell.prologue import (
|
|
kda_prologue,
|
|
)
|
|
|
|
BT = 64 # chunk size
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _l2norm(x: torch.Tensor) -> torch.Tensor:
|
|
return F.normalize(x.float(), p=2, dim=-1)
|
|
|
|
|
|
def kda_flops(total_seq_len, num_heads, head_k, head_v):
|
|
"""Per-token-per-head: k@v outer (2*K*V) + q@state (2*K*V), plus the intra-chunk
|
|
KKT (2*K*K averaged over the chunk). Approximate (ignores the inverse)."""
|
|
return total_seq_len * num_heads * (4 * head_k * head_v + 2 * head_k * head_k)
|
|
|
|
|
|
def kda_bytes(total_seq_len, num_heads, head_k, head_v, num_seqs, dtype):
|
|
elem = dtype.itemsize
|
|
q_b = total_seq_len * num_heads * head_k * elem
|
|
k_b = total_seq_len * num_heads * head_k * elem
|
|
v_b = total_seq_len * num_heads * head_v * elem
|
|
o_b = total_seq_len * num_heads * head_v * elem
|
|
g_b = total_seq_len * num_heads * head_k * 4 # per-channel gate, fp32
|
|
beta_b = total_seq_len * num_heads * 4
|
|
state_b = 2 * num_seqs * num_heads * head_k * head_v * 4 # fp32 r/w
|
|
return q_b + k_b + v_b + o_b + g_b + beta_b + state_b
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Input factory (single sequence per benchmark point, B=1)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def make_inputs(T, H, K, V, device, dtype, seed=42):
|
|
torch.manual_seed(seed)
|
|
q = _l2norm(torch.randn(1, T, H, K, device=device)).to(dtype)
|
|
k = _l2norm(torch.randn(1, T, H, K, device=device)).to(dtype)
|
|
v = torch.randn(1, T, H, V, device=device).to(dtype)
|
|
# Mild per-channel gate (real Kimi-Linear regime; keeps exp() in fp32 range).
|
|
A_log = torch.randn(H, device=device) * 0.5 - 1.5
|
|
dt_bias = torch.randn(H, K, device=device) * 0.1
|
|
g_raw = torch.randn(1, T, H, K, device=device)
|
|
g_act = (
|
|
-A_log.exp().view(1, 1, H, 1) * F.softplus(g_raw + dt_bias.view(1, 1, H, K))
|
|
).float()
|
|
beta = torch.sigmoid(torch.randn(1, T, H, device=device)).float()
|
|
return dict(q=q, k=k, v=v, g_act=g_act, beta=beta, T=T, H=H, K=K, V=V)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runners
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_recurrent(inp, scale):
|
|
"""Token-by-token ground truth. Returns (o [T,H,V], state [1,H,V,K])."""
|
|
cu = torch.tensor([0, inp["T"]], dtype=torch.int64, device=inp["q"].device)
|
|
h0 = torch.zeros(1, inp["H"], inp["V"], inp["K"], device=inp["q"].device)
|
|
o, state = fused_recurrent_kda(
|
|
q=inp["q"],
|
|
k=inp["k"],
|
|
v=inp["v"],
|
|
g=inp["g_act"],
|
|
beta=inp["beta"],
|
|
scale=scale,
|
|
initial_state=h0,
|
|
inplace_final_state=False,
|
|
use_qk_l2norm_in_kernel=False,
|
|
cu_seqlens=cu,
|
|
)
|
|
return o[0], state
|
|
|
|
|
|
def cutedsl_buffers(inp, num_sms, device):
|
|
"""Precompute metadata + preallocate (shared across layers in a real forward)."""
|
|
T, H, K, V = inp["T"], inp["H"], inp["K"], inp["V"]
|
|
cu = torch.tensor([0, T], dtype=torch.int32, device=device)
|
|
ci, co, tc, total = prepare_metadata(cu)
|
|
pad_t = total * BT
|
|
return dict(
|
|
cu=cu,
|
|
ci=ci,
|
|
co=co,
|
|
tc=tc,
|
|
total=total,
|
|
num_sms=num_sms,
|
|
h0=torch.zeros(1, H, V, K, device=device, dtype=torch.float32),
|
|
U=torch.empty(pad_t, H, V, device=device, dtype=torch.bfloat16),
|
|
W=torch.empty(pad_t, H, K, device=device, dtype=torch.bfloat16),
|
|
V_new=torch.empty(pad_t, H, V, device=device, dtype=torch.bfloat16),
|
|
h_chunks=torch.empty(total, H, V, K, device=device, dtype=torch.bfloat16),
|
|
ht=torch.empty(1, H, V, K, device=device, dtype=torch.float32),
|
|
o=torch.empty(T, H, V, device=device, dtype=torch.bfloat16),
|
|
)
|
|
|
|
|
|
def run_cutedsl_pipeline(inp, buf, scale):
|
|
"""The fused-prologue + 3 cutedsl kernels (metadata precomputed in buf)."""
|
|
q3, k3, v3 = inp["q"][0], inp["k"][0], inp["v"][0]
|
|
g3, beta3 = inp["g_act"][0], inp["beta"][0].contiguous()
|
|
KL, KR, KG, qg, qg2, g_cu = kda_prologue(
|
|
q3, k3, g3, scale, buf["cu"], buf["ci"], buf["total"]
|
|
)
|
|
kkt_inv_uw_cutedsl(
|
|
KL,
|
|
KR,
|
|
KG,
|
|
v3,
|
|
buf["U"],
|
|
buf["W"],
|
|
beta3,
|
|
buf["cu"],
|
|
buf["ci"],
|
|
buf["tc"],
|
|
num_sms=buf["num_sms"],
|
|
)
|
|
kda_h_cutedsl(
|
|
KR,
|
|
buf["U"],
|
|
buf["W"],
|
|
buf["V_new"],
|
|
g_cu,
|
|
buf["h_chunks"],
|
|
buf["h0"],
|
|
buf["ht"],
|
|
buf["cu"],
|
|
buf["co"],
|
|
)
|
|
kda_o_cutedsl(
|
|
qg,
|
|
qg2,
|
|
KR,
|
|
buf["V_new"],
|
|
buf["h_chunks"],
|
|
buf["o"],
|
|
buf["cu"],
|
|
buf["ci"],
|
|
buf["tc"],
|
|
num_sms=buf["num_sms"],
|
|
)
|
|
return buf["o"], buf["ht"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Correctness
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def check_shape(T, H, K, V, device, dtype, num_sms):
|
|
tag = f"T={T:>5} H={H:>2} K={K:>3} V={V:>3}"
|
|
if K != 128 or V != 128:
|
|
print(f" [SKIP] {tag} (cutedsl requires K=V=128)")
|
|
return True
|
|
|
|
scale = K**-0.5
|
|
inp = make_inputs(T, H, K, V, device, dtype)
|
|
o_ref, state_ref = run_recurrent(inp, scale)
|
|
|
|
try:
|
|
buf = cutedsl_buffers(inp, num_sms, device)
|
|
o, ht = run_cutedsl_pipeline(inp, buf, scale)
|
|
torch.cuda.synchronize()
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" [SKIP] {tag} (cutedsl error: {e})")
|
|
return True
|
|
|
|
finite = bool(torch.isfinite(o).all() and torch.isfinite(ht).all())
|
|
o_err = (o.float() - o_ref.float()).abs().max().item()
|
|
s_err = (ht.float() - state_ref.float()).abs().max().item()
|
|
ok = finite and o_err < 1e-2 and s_err < 5e-2
|
|
status = "PASS" if ok else "FAIL"
|
|
print(
|
|
f" [{status}] {tag} | o_err {o_err:.2e} state_err {s_err:.2e} finite={finite}"
|
|
)
|
|
return ok
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Benchmark
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def bench_shape(T, H, K, V, device, dtype, num_sms):
|
|
import triton.testing
|
|
|
|
if K != 128 or V != 128:
|
|
print(f" [SKIP] T={T} H={H} K={K} V={V} (cutedsl K=V=128 only)")
|
|
return
|
|
|
|
scale = K**-0.5
|
|
inp = make_inputs(T, H, K, V, device, dtype)
|
|
q, k, v = inp["q"], inp["k"], inp["v"]
|
|
g_act, beta = inp["g_act"], inp["beta"]
|
|
|
|
h0f = torch.zeros(1, H, K, V, device=device, dtype=torch.float32)
|
|
idx = torch.zeros(1, dtype=torch.int32, device=device)
|
|
|
|
def fn_triton():
|
|
chunk_kda(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
g=g_act,
|
|
beta=beta,
|
|
scale=scale,
|
|
initial_state=h0f,
|
|
initial_state_indices=idx,
|
|
use_qk_l2norm_in_kernel=False,
|
|
cu_seqlens=None,
|
|
A_log=None,
|
|
dt_bias=None,
|
|
lower_bound=None,
|
|
)
|
|
|
|
buf = cutedsl_buffers(inp, num_sms, device)
|
|
|
|
def fn_cutedsl():
|
|
run_cutedsl_pipeline(inp, buf, scale)
|
|
|
|
quantiles = [0.5, 0.2, 0.8]
|
|
fn_triton()
|
|
fn_cutedsl()
|
|
torch.cuda.synchronize()
|
|
|
|
ms_triton, _, _ = triton.testing.do_bench_cudagraph(fn_triton, quantiles=quantiles)
|
|
ms_cutedsl, _, _ = triton.testing.do_bench_cudagraph(
|
|
fn_cutedsl, quantiles=quantiles
|
|
)
|
|
|
|
flops = kda_flops(T, H, K, V)
|
|
mem_bytes = kda_bytes(T, H, K, V, 1, dtype)
|
|
speedup = ms_triton / ms_cutedsl if ms_cutedsl > 0 else float("inf")
|
|
print(
|
|
f" {H:>3} {T:>7} | "
|
|
f"{ms_triton:>8.3f} {flops / ms_triton / 1e9:>7.2f} {mem_bytes / ms_triton / 1e9:>7.2f} | "
|
|
f"{ms_cutedsl:>8.3f} {flops / ms_cutedsl / 1e9:>7.2f} {mem_bytes / ms_cutedsl / 1e9:>7.2f} | "
|
|
f"{speedup:>7.2f}x"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_correctness(device, dtype, H, num_sms):
|
|
print("=" * 72)
|
|
print("Correctness: cutedsl pipeline vs fused_recurrent_kda (ground truth)")
|
|
print("=" * 72)
|
|
all_pass = True
|
|
for T in (128, 192, 256, 512, 1024):
|
|
if not check_shape(T, H, 128, 128, device, dtype, num_sms):
|
|
all_pass = False
|
|
print("\nALL PASSED." if all_pass else "\nSOME FAILED.")
|
|
return all_pass
|
|
|
|
|
|
def run_benchmark(device, dtype, args, num_sms):
|
|
print()
|
|
print("=" * 92)
|
|
print("Benchmark: Triton chunk_kda vs CuTeDSL pipeline (do_bench_cudagraph)")
|
|
print("=" * 92)
|
|
print(f" Device SMs={num_sms}, K=V=128, dtype={dtype}, metadata precomputed")
|
|
print(
|
|
f" {'H':>3} {'T':>7} | "
|
|
f"{'tri(ms)':>8} {'TFLOP':>7} {'TB/s':>7} | "
|
|
f"{'cute(ms)':>8} {'TFLOP':>7} {'TB/s':>7} | {'speedup':>8}"
|
|
)
|
|
print(" " + "-" * 84)
|
|
for H in args.num_heads:
|
|
for T in args.seq_lens:
|
|
bench_shape(T, H, 128, 128, device, dtype, num_sms)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Benchmark & Correctness: Triton KDA vs CuTeDSL KDA (SM100)"
|
|
)
|
|
parser.add_argument(
|
|
"--mode", choices=["all", "correctness", "bench"], default="all"
|
|
)
|
|
parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16")
|
|
parser.add_argument("--num-heads", type=int, nargs="+", default=[32])
|
|
parser.add_argument(
|
|
"--seq-lens", type=int, nargs="+", default=[512, 1024, 2048, 4096, 8192]
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
device = "cuda"
|
|
dtype = getattr(torch, args.dtype)
|
|
cap = torch.cuda.get_device_capability()
|
|
print(f"Device: {torch.cuda.get_device_name()} (SM {cap[0]}{cap[1]})")
|
|
if cap[0] < 10:
|
|
print("ERROR: CuTeDSL KDA prefill requires SM100+ (Blackwell). Exiting.")
|
|
return 1
|
|
num_sms = torch.cuda.get_device_properties(0).multi_processor_count
|
|
|
|
if args.mode in ("all", "correctness"):
|
|
all_pass = run_correctness(device, dtype, args.num_heads[0], num_sms)
|
|
if not all_pass and args.mode == "all":
|
|
print("\nSkipping benchmark due to correctness failures.")
|
|
return 1
|
|
|
|
if args.mode in ("all", "bench"):
|
|
run_benchmark(device, dtype, args, num_sms)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|