Skip to content

Commit b834050

Browse files
feifei14119feifei14119
committed
shabi
1 parent eaef170 commit b834050

15 files changed

Lines changed: 1616 additions & 9 deletions

‎aiter/mla.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -317,11 +317,9 @@ def mla_decode_fwd(
317317
# Per-batch valid KV split count writeback buffer. Always allocated (and
318318
# passed to stage1) so the asm kernel has a valid destination; whether
319319
# stage2 actually uses it is gated by use_valid_split_count_reduce.
320-
# Initialized to num_kv_splits so a min() against it is a no-op until the
321-
# kernel overwrites it with the real (smaller) valid count.
322-
valid_split_count = torch.full(
323-
(bs,), num_kv_splits, dtype=dtypes.i32, device=device
324-
)
320+
# Left uninitialized: the asm kernel writes the real (valid) count before
321+
# stage2 reads it, so pre-filling would be dead work.
322+
valid_split_count = torch.empty((bs,), dtype=dtypes.i32, device=device)
325323
use_valid_split_count_reduce = int(num_kv_splits > 1)
326324

327325
aiter.mla_decode_stage1_asm_fwd(

‎hsa/gfx1250/mla_v4/mla_a8w8_qh16_1tg_16mx4_64nx1_sparse.co‎

512 Bytes
Binary file not shown.

‎hsa/gfx1250/mla_v4/mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co‎

512 Bytes
Binary file not shown.

‎op_tests/_failanalyze.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
import sys, torch
2+
sys.path.insert(0, "op_tests")
3+
import test_mla_v4_kargpreld as T
4+
import aiter, aiter.mla
5+
from aiter import dtypes
6+
7+
BATCH = int(sys.argv[1]) if len(sys.argv) > 1 else 64
8+
KV = int(sys.argv[2]) if len(sys.argv) > 2 else 271
9+
SPLITS = int(sys.argv[3]) if len(sys.argv) > 3 else 8
10+
GQA, Q, SINK = 64, 1, True
11+
ATOL = RTOL = 3e-2
12+
T._SEED = 0
13+
14+
15+
def build_and_run():
16+
inp = T._build_bf16_inputs(batch=BATCH, kv_seq_lens=KV, q_seq_logical=Q, seed=0,
17+
gqa_ratio=GQA, attn_sink=SINK)
18+
sm = 1.0 / (T._QUANT_D ** 0.5)
19+
qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"])
20+
kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"])
21+
ref, _ = T._torch_attn_decode_fp8_dequant_ref(
22+
qp, qr, kvp, kvr, inp["qo_indptr"], inp["kv_indptr"],
23+
inp["kv_page_indices"], inp["kv_last_page_lens"], sm, attn_sink=inp["sink"])
24+
total_q = inp["q_bf16"].size(0); ns = inp["qo_indptr"].size(0) - 1
25+
nh = T.NUM_KV_HEADS * GQA; dev = "cuda"
26+
ob = torch.empty((total_q, GQA, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev)
27+
sidx = torch.tensor([i * SPLITS for i in range(ns + 1)], dtype=torch.int32, device=dev)
28+
lb = torch.empty((total_q, SPLITS, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev)
29+
eb = torch.empty((total_q, SPLITS, nh, 1), dtype=dtypes.fp32, device=dev)
30+
aiter.mla.mla_decode_fwd_v4_nm(
31+
q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob,
32+
qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"],
33+
kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"],
34+
split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"],
35+
sm_scale=sm, out_16_nosplit=0, num_kv_splits=SPLITS, logits=lb, attn_lse=eb)
36+
torch.cuda.synchronize()
37+
return ref.float(), ob.float()
38+
39+
40+
tries = int(sys.argv[4]) if len(sys.argv) > 4 else 20
41+
for t in range(tries):
42+
ref, asm = build_and_run()
43+
close = torch.isclose(ref, asm, rtol=RTOL, atol=ATOL)
44+
pct = (~close).float().mean().item()
45+
if pct > 0.02:
46+
print(f"[got FAIL on try {t}] mismatch={pct:.2%}")
47+
break
48+
else:
49+
print(f"no >2% fail in {tries} tries; using last run (mismatch={pct:.2%})")
50+
51+
m = ~close
52+
rv = ref[m]; av = asm[m]
53+
delta = (av - rv).abs()
54+
print(f"\n===== ctx={KV} batch={BATCH} split={SPLITS} : {m.sum().item()} mismatched elems =====")
55+
print(f"ref (golden) at mismatches: min={rv.min():.3f} max={rv.max():.3f} mean|.|={rv.abs().mean():.3f}")
56+
print(f"asm (kernel) at mismatches: min={av.min():.3f} max={av.max():.3f} mean|.|={av.abs().mean():.3f}")
57+
print(f"|delta| at mismatches : min={delta.min():.4f} max={delta.max():.4f} mean={delta.mean():.4f}")
58+
print(f"whole-tensor magnitude : ref |.|max={ref.abs().max():.3f} asm |.|max={asm.abs().max():.3f}")
59+
print(f"NaN/Inf in asm? {torch.isnan(asm).any().item()} / {torch.isinf(asm).any().item()}")
60+
# neighbor vs wild: how big is delta relative to typical value scale?
61+
scale = ref.abs().mean().item()
62+
print(f"\ntypical |ref| scale (all)= {scale:.3f}")
63+
print(f"delta/scale ratio: median={ (delta/ (rv.abs()+1e-6)).median():.2f} max={(delta/(rv.abs()+1e-6)).max():.2f}")
64+
# histogram of delta
65+
import numpy as np
66+
d = delta.cpu().numpy()
67+
for thr in [0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 1e3, 1e6]:
68+
print(f" |delta| <= {thr:>8}: {100.0*(d<=thr).mean():5.1f}%")
69+
# sample pairs
70+
idx = torch.argsort(delta, descending=True)[:12]
71+
print("\ntop-12 worst (ref -> asm, delta):")
72+
for i in idx.tolist():
73+
print(f" {rv[i].item():+.4f} -> {av[i].item():+.4f} d={delta[i].item():.4f}")

‎op_tests/_pat.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import sys, torch
2+
sys.path.insert(0, "op_tests")
3+
import test_mla_v4_kargpreld as T
4+
import aiter, aiter.mla
5+
from aiter import dtypes
6+
7+
BATCH = int(sys.argv[1]) if len(sys.argv) > 1 else 60
8+
KV = int(sys.argv[2]) if len(sys.argv) > 2 else 271
9+
SPLITS = int(sys.argv[3]) if len(sys.argv) > 3 else 8
10+
Q, GQA, SINK = 1, 64, True
11+
12+
13+
def run_once():
14+
inp = T._build_bf16_inputs(batch=BATCH, kv_seq_lens=KV, q_seq_logical=Q, seed=0,
15+
gqa_ratio=GQA, attn_sink=SINK)
16+
qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"])
17+
kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"])
18+
total_q = inp["q_bf16"].size(0); nh = T.NUM_KV_HEADS * GQA; dev = "cuda"
19+
ns = inp["qo_indptr"].numel() - 1
20+
sidx = torch.tensor([i * SPLITS for i in range(ns + 1)], dtype=torch.int32, device=dev)
21+
ob = torch.empty((total_q, GQA, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev)
22+
lb = torch.zeros((total_q, SPLITS, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev)
23+
eb = torch.full((total_q, SPLITS, nh, 1), float("-inf"), dtype=dtypes.fp32, device=dev)
24+
aiter.mla.mla_decode_fwd_v4_nm(
25+
q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob,
26+
qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"],
27+
kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"],
28+
split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"],
29+
sm_scale=1.0 / (T._QUANT_D ** 0.5), out_16_nosplit=0, num_kv_splits=SPLITS,
30+
logits=lb, attn_lse=eb)
31+
torch.cuda.synchronize()
32+
return lb.float().clone(), eb.float().clone()
33+
34+
35+
print(f"===== batch={BATCH} : per-split stage1 nondeterminism (3 runs, same input) =====")
36+
L1, E1 = run_once()
37+
L2, E2 = run_once()
38+
L3, E3 = run_once()
39+
print("split : logits max|r2-r1| logits max|r3-r1| lse max|r2-r1| status")
40+
for s in range(SPLITS):
41+
d2 = (L2[:, s] - L1[:, s]).abs().max().item()
42+
d3 = (L3[:, s] - L1[:, s]).abs().max().item()
43+
fe = torch.isfinite(E2[:, s]) & torch.isfinite(E1[:, s])
44+
de = (E2[:, s] - E1[:, s])[fe].abs().max().item() if fe.any() else 0.0
45+
status = "RACE" if max(d2, d3) > 1e-6 else "stable"
46+
print(f" {s} : {d2:>13.6g} {d3:>13.6g} {de:>13.6g} {status}")

‎op_tests/_stage1_bench.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
"""Standalone stage1-only perf comparison: asm sparse decode kernel vs ATOM's
2+
triton split kernel (bf16).
3+
4+
- asm stage1 = the `..._sparse` .co kernel (fp8, its native dtype), timed alone.
5+
- triton stage1 = ATOM `_paged_decode_split_kernel` (bf16), forced via kv_splits>1.
6+
7+
Only the stage1 GPU kernel device-time is compared (profiler filtered by name).
8+
NOT bit-exact / not memory-fair: triton is bf16 while asm is fp8, and D=512 here
9+
(MLA QK is really 576-wide, so triton QK is ~11% under-counted). Perf only.
10+
11+
Usage:
12+
ENABLE_CK=0 python op_tests/_stage1_bench.py # default sweep
13+
"""
14+
import sys, os
15+
sys.path.insert(0, "op_tests")
16+
sys.path.insert(0, "/home/amd/feifei/ATOM")
17+
18+
import torch
19+
import test_mla_v4_kargpreld as T
20+
import aiter, aiter.mla
21+
from aiter import dtypes
22+
from aiter.test_common import run_perftest
23+
from atom.model_ops.v4_kernels.paged_decode import (
24+
_sparse_attn_v4_paged_decode_triton as triton_decode,
25+
)
26+
27+
Q = 1
28+
SINK = True
29+
30+
# Mirror test_mla_v4_kargpreld.py sweep grids + active kernel variants.
31+
_GQA_LIST = [64, 128]
32+
_CTX_LENS = [256, 512, 1024]
33+
_BATCH_SIZES = [64]
34+
_SPLITS = [1]
35+
36+
# Perf iteration counts for run_perftest (mirrors test_mla_v4_kargpreld.py's
37+
# _PERF usage). Kept as module constants so both bench fns share them.
38+
_PERF = {"num_iters": 50, "num_warmup": 3}
39+
40+
41+
def bench_asm_stage1(batch, ctx, split, gqa):
42+
inp = T._build_bf16_inputs(batch=batch, kv_seq_lens=ctx, q_seq_logical=Q,
43+
seed=0, gqa_ratio=gqa, attn_sink=SINK)
44+
sm = 1.0 / (T._QUANT_D ** 0.5)
45+
qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"])
46+
kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"])
47+
total_q = inp["q_bf16"].size(0)
48+
ns = inp["qo_indptr"].size(0) - 1
49+
nh = T.NUM_KV_HEADS * gqa
50+
dev = "cuda"
51+
ob = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev)
52+
sidx = torch.tensor([i * split for i in range(ns + 1)], dtype=torch.int32, device=dev)
53+
lb = torch.empty((total_q, split, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev)
54+
eb = torch.empty((total_q, split, nh, 1), dtype=dtypes.fp32, device=dev)
55+
kw = dict(
56+
q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob,
57+
qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"],
58+
kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"],
59+
split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"],
60+
sm_scale=sm, out_16_nosplit=0, num_kv_splits=split, logits=lb, attn_lse=eb,
61+
)
62+
_, us = run_perftest(
63+
aiter.mla.mla_decode_fwd_v4_nm,
64+
**kw,
65+
num_iters=_PERF["num_iters"],
66+
num_warmup=_PERF["num_warmup"],
67+
)
68+
return us
69+
70+
71+
def bench_triton_stage1(batch, ctx, split, gqa):
72+
dev = "cuda"
73+
D = T.V_HEAD_DIM # 512
74+
Tn = batch # one decode token per sequence
75+
H = gqa
76+
q = torch.randn((Tn, H, D), dtype=torch.bfloat16, device=dev)
77+
total_pages = batch * ctx
78+
unified_kv = torch.randn((total_pages, D), dtype=torch.bfloat16, device=dev) * 0.1
79+
kv_indices = torch.arange(total_pages, dtype=torch.int32, device=dev)
80+
kv_indptr = torch.arange(0, (Tn + 1) * ctx, ctx, dtype=torch.int32, device=dev)
81+
attn_sink = torch.randn((H,), dtype=torch.float32, device=dev)
82+
sm = 1.0 / (D ** 0.5)
83+
_, us = run_perftest(
84+
triton_decode,
85+
q, unified_kv, kv_indices, kv_indptr, attn_sink, sm, kv_splits=split,
86+
num_iters=_PERF["num_iters"],
87+
num_warmup=_PERF["num_warmup"],
88+
)
89+
return us
90+
91+
92+
import itertools
93+
94+
print(f"{'gqa':>4} {'batch':>6} {'ctx':>6} {'split':>6} | {'asm_stage1_us':>14} {'triton_stage1_us':>17} {'triton/asm':>11}")
95+
print("-" * 80)
96+
# gqa outermost so rows mirror the kargpreld summary-table grouping.
97+
for gqa, batch, ctx, split in itertools.product(_GQA_LIST, _BATCH_SIZES, _CTX_LENS, _SPLITS):
98+
asm_us = bench_asm_stage1(batch, ctx, split, gqa)
99+
tri_us = bench_triton_stage1(batch, ctx, split, gqa)
100+
ratio = tri_us / asm_us if asm_us > 0 else float("nan")
101+
print(f"{gqa:>4} {batch:>6} {ctx:>6} {split:>6} | {asm_us:>14.2f} {tri_us:>17.2f} {ratio:>10.2f}x")

‎op_tests/_stage1_bench_pa.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
"""Standalone stage1-only perf comparison: asm sparse decode kernel vs ATOM's
2+
gfx1250 gluon `pa_decode_sparse` stage1 (bf16).
3+
4+
- asm stage1 = the `..._sparse` .co kernel (fp8, its native dtype), timed alone.
5+
- pa stage1 = aiter `pa_decode_sparse(..., skip_reduce=True)` with kv_splits>1,
6+
which launches ONLY the split (stage1) gluon kernel (no reduce).
7+
8+
Only stage1 GPU device-time is compared. NOT bit-exact / not memory-fair: pa is
9+
bf16 while asm is fp8, and D=512 here (MLA QK is really 576-wide, so pa QK is
10+
~11% under-counted). Perf only.
11+
12+
Usage:
13+
ENABLE_CK=0 python op_tests/_stage1_bench_pa.py # default sweep
14+
ENABLE_CK=0 python op_tests/_stage1_bench_pa.py 128 2048 4 # one combo
15+
"""
16+
import sys
17+
sys.path.insert(0, "op_tests")
18+
sys.path.insert(0, "/home/amd/feifei/ATOM")
19+
20+
import torch
21+
import test_mla_v4_kargpreld as T
22+
import aiter, aiter.mla
23+
from aiter import dtypes
24+
from aiter.test_common import run_perftest
25+
from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse
26+
27+
Q = 1
28+
SINK = True
29+
30+
# Mirror test_mla_v4_kargpreld.py sweep grids + active kernel variants.
31+
_GQA_LIST = [64, 128]
32+
_CTX_LENS = [256, 512, 1024]
33+
_BATCH_SIZES = [64]
34+
_SPLITS = [1]
35+
36+
# Perf iteration counts for run_perftest (mirrors test_mla_v4_kargpreld.py's
37+
# _PERF usage). Kept as module constants so both bench fns share them.
38+
_PERF = {"num_iters": 50, "num_warmup": 2}
39+
40+
41+
def bench_asm_stage1(batch, ctx, split, gqa):
42+
inp = T._build_bf16_inputs(batch=batch, kv_seq_lens=ctx, q_seq_logical=Q,
43+
seed=0, gqa_ratio=gqa, attn_sink=SINK)
44+
sm = 1.0 / (T._QUANT_D ** 0.5)
45+
qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"])
46+
kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"])
47+
total_q = inp["q_bf16"].size(0)
48+
ns = inp["qo_indptr"].size(0) - 1
49+
nh = T.NUM_KV_HEADS * gqa
50+
dev = "cuda"
51+
ob = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev)
52+
sidx = torch.tensor([i * split for i in range(ns + 1)], dtype=torch.int32, device=dev)
53+
lb = torch.empty((total_q, split, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev)
54+
eb = torch.empty((total_q, split, nh, 1), dtype=dtypes.fp32, device=dev)
55+
kw = dict(
56+
q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob,
57+
qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"],
58+
kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"],
59+
split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"],
60+
sm_scale=sm, out_16_nosplit=0, num_kv_splits=split, logits=lb, attn_lse=eb,
61+
)
62+
# asm mla_decode_fwd_v4_nm launches stage1 (sparse) + stage2 (merge). Match
63+
# test_mla_v4_kargpreld.py's stage1_us: profile and keep ONLY the stage1
64+
# "sparse" kernel's device time (drop stage2), so this column lines up with
65+
# the kargpreld summary table.
66+
s1_us, _s2_us = T._profile_stage_times(
67+
lambda: aiter.mla.mla_decode_fwd_v4_nm(**kw),
68+
iters=_PERF["num_iters"],
69+
warmup=_PERF["num_warmup"],
70+
)
71+
return s1_us
72+
73+
74+
def bench_pa_stage1(batch, ctx, split, gqa):
75+
dev = "cuda"
76+
D = T.V_HEAD_DIM # 512
77+
Tn = batch
78+
H = gqa
79+
q = torch.randn((Tn, H, D), dtype=torch.bfloat16, device=dev)
80+
total_pages = batch * ctx
81+
unified_kv = torch.randn((total_pages, D), dtype=torch.bfloat16, device=dev) * 0.1
82+
kv_indices = torch.arange(total_pages, dtype=torch.int32, device=dev)
83+
kv_indptr = torch.arange(0, (Tn + 1) * ctx, ctx, dtype=torch.int32, device=dev)
84+
attn_sink = torch.randn((H,), dtype=torch.float32, device=dev)
85+
sm = 1.0 / (D ** 0.5)
86+
87+
# skip_reduce=True + kv_splits>1 => only the stage1 (split) kernel runs.
88+
_, us = run_perftest(
89+
pa_decode_sparse,
90+
q, unified_kv, kv_indices, kv_indptr, attn_sink, sm,
91+
kv_splits=split, has_invalid=False, skip_reduce=True,
92+
num_iters=_PERF["num_iters"],
93+
num_warmup=_PERF["num_warmup"],
94+
)
95+
return us
96+
97+
98+
import itertools
99+
100+
print(f"{'gqa':>4} {'batch':>6} {'ctx':>6} {'split':>6} | {'asm_stage1_us':>14} {'pa_stage1_us':>13} {'pa/asm':>8}")
101+
print("-" * 74)
102+
# gqa outermost so rows mirror the kargpreld summary-table grouping.
103+
for gqa, batch, ctx, split in itertools.product(_GQA_LIST, _BATCH_SIZES, _CTX_LENS, _SPLITS):
104+
asm_us = bench_asm_stage1(batch, ctx, split, gqa)
105+
pa_us = bench_pa_stage1(batch, ctx, split, gqa)
106+
ratio = pa_us / asm_us if asm_us > 0 else float("nan")
107+
print(f"{gqa:>4} {batch:>6} {ctx:>6} {split:>6} | {asm_us:>14.2f} {pa_us:>13.2f} {ratio:>7.2f}x")

0 commit comments

Comments
 (0)