[AMD] Add Work-Centric (Lean) Attention: a persistent-CTA decode kernel for long-context serving (#33576)
Signed-off-by: Valerie Chen <valechen@amd.com> Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
valechen
Executable
+163
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Kernel-level Lean vs Standard (SplitK) decode-attention sweep, both GQA head
|
||||
configs, full batch x context grid, at the shipped 1xCU persistent grid.
|
||||
|
||||
Generalizes benchmark/lean_kernel_qwen_gqa.py (which is Qwen batch=1 only) to
|
||||
sweep Qwen2.5-7B (28Q/4KV) and Llama-3.1-8B (32Q/8KV) over
|
||||
batch in {1,2,4,8,16,32} x context in {8K,16K,32K,64K,128K}, reporting per-call
|
||||
kernel latency, speedup (std / lean), cosine parity vs SplitK, and whether the
|
||||
eager auto-gate would enable Lean. Writes a CSV for the PR tables.
|
||||
|
||||
Grid is whatever SGLANG_FORCE_LEAN_GRID_CU_MULT resolves to (default 1.0 = one
|
||||
CTA per CU); set it to A/B other grids without a rebuild.
|
||||
|
||||
Usage: python3 benchmark/lean_kernel_sweep.py [out.csv]
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention.decode_attention import (
|
||||
_LEAN_BLOCK_M,
|
||||
_lean_decode_launch_params,
|
||||
decode_attention_fwd,
|
||||
decode_attention_fwd_grouped,
|
||||
lean_decode_seqlen_gate,
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
("qwen2.5-7b", 28, 4),
|
||||
("llama3.1-8b", 32, 8),
|
||||
]
|
||||
D = D_V = 128
|
||||
MAX_KV_SPLITS = 8
|
||||
BATCHES = [1, 2, 4, 8, 16, 32]
|
||||
CONTEXTS = [8192, 16384, 32768, 65536, 131072]
|
||||
|
||||
|
||||
def bench(fn, warmup=15, iters=100):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(iters):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
return (time.perf_counter() - t0) / iters * 1000.0 # ms
|
||||
|
||||
|
||||
def run(H_Q, H_KV, B, S):
|
||||
dev, dt = "cuda", torch.float16
|
||||
kvg = H_Q // H_KV
|
||||
sm = 1.0 / (D**0.5)
|
||||
tot = B * S
|
||||
|
||||
total_programs, _, _ = _lean_decode_launch_params(H_KV, kvg)
|
||||
lean_Mp = torch.empty(
|
||||
(total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev
|
||||
)
|
||||
lean_Lp = torch.empty(
|
||||
(total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev
|
||||
)
|
||||
lean_Op = torch.empty(
|
||||
(total_programs, _LEAN_BLOCK_M, D_V), dtype=torch.float32, device=dev
|
||||
)
|
||||
lean_locks = torch.zeros((total_programs,), dtype=torch.int32, device=dev)
|
||||
|
||||
kvi = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32)
|
||||
kvx = torch.arange(0, tot, device=dev, dtype=torch.int32)
|
||||
q = torch.randn(B, H_Q, D, dtype=dt, device=dev)
|
||||
k = torch.randn(tot, H_KV, D, dtype=dt, device=dev)
|
||||
v = torch.randn(tot, H_KV, D_V, dtype=dt, device=dev)
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, MAX_KV_SPLITS, D_V), dtype=torch.float32, device=dev
|
||||
)
|
||||
attn_lse = torch.empty((B, H_Q, MAX_KV_SPLITS), dtype=torch.float32, device=dev)
|
||||
nks = torch.full((B,), MAX_KV_SPLITS, dtype=torch.int32, device=dev)
|
||||
|
||||
o_std = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
|
||||
std_ms = bench(
|
||||
lambda: decode_attention_fwd_grouped(
|
||||
q, k, v, o_std, kvi, kvx, attn_logits, attn_lse, nks, MAX_KV_SPLITS, sm, 1.0
|
||||
)
|
||||
)
|
||||
|
||||
attn_logits2 = torch.empty_like(attn_logits)
|
||||
attn_lse2 = torch.empty_like(attn_lse)
|
||||
o_lean = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
|
||||
lean_ms = bench(
|
||||
lambda: decode_attention_fwd(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o_lean,
|
||||
kvi,
|
||||
kvx,
|
||||
attn_logits2,
|
||||
attn_lse2,
|
||||
nks,
|
||||
MAX_KV_SPLITS,
|
||||
sm,
|
||||
1.0,
|
||||
1.0,
|
||||
enable_lean=True,
|
||||
lean_Mp=lean_Mp,
|
||||
lean_Lp=lean_Lp,
|
||||
lean_Op=lean_Op,
|
||||
lean_locks=lean_locks,
|
||||
)
|
||||
)
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
o_lean.flatten().float(), o_std.flatten().float(), dim=0
|
||||
).item()
|
||||
gate = lean_decode_seqlen_gate(H_Q, kvg, B, B * S, is_mla=False)
|
||||
return std_ms, lean_ms, cos, gate
|
||||
|
||||
|
||||
def main():
|
||||
out = sys.argv[1] if len(sys.argv) > 1 else "grid_out/kernel_sweep_1xcu.csv"
|
||||
mult = float(os.environ.get("SGLANG_FORCE_LEAN_GRID_CU_MULT", "1.0"))
|
||||
print(f"\nKernel sweep GPU={torch.cuda.get_device_name(0)} grid_mult={mult}")
|
||||
print("=" * 78)
|
||||
rows = ["model,H_Q,H_KV,batch,context,std_ms,lean_ms,speedup,cos,gate"]
|
||||
for name, H_Q, H_KV in MODELS:
|
||||
print(f"\n{name} ({H_Q}Q/{H_KV}KV)")
|
||||
print(
|
||||
f"{'batch':>5} {'ctx':>6} {'std_ms':>9} {'lean_ms':>9} {'speedup':>8} {'gate':>5} {'cos':>7}"
|
||||
)
|
||||
for B in BATCHES:
|
||||
for S in CONTEXTS:
|
||||
# b32 x 128K on the 8-KV-head config exceeds the microbench's single
|
||||
# contiguous KV tensor (faults the GPU); real serving uses a paged pool.
|
||||
if H_KV == 8 and B == 32 and S == 131072:
|
||||
print(f"{B:>5} {S//1024:>5}K {'skipped (contiguous-KV limit)':>30}")
|
||||
rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,")
|
||||
continue
|
||||
try:
|
||||
std, lean, cos, gate = run(H_Q, H_KV, B, S)
|
||||
except torch.cuda.OutOfMemoryError:
|
||||
torch.cuda.empty_cache()
|
||||
print(f"{B:>5} {S//1024:>5}K {'OOM':>9}")
|
||||
rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,OOM,")
|
||||
continue
|
||||
sp = std / lean
|
||||
print(
|
||||
f"{B:>5} {S//1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}"
|
||||
)
|
||||
rows.append(
|
||||
f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}"
|
||||
)
|
||||
os.makedirs(os.path.dirname(out), exist_ok=True)
|
||||
open(out, "w").write("\n".join(rows) + "\n")
|
||||
print(f"\nwrote {out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user