diff --git a/aiter/mla.py b/aiter/mla.py index 01430c6090d..80046655aee 100644 --- a/aiter/mla.py +++ b/aiter/mla.py @@ -317,11 +317,9 @@ def mla_decode_fwd( # Per-batch valid KV split count writeback buffer. Always allocated (and # passed to stage1) so the asm kernel has a valid destination; whether # stage2 actually uses it is gated by use_valid_split_count_reduce. - # Initialized to num_kv_splits so a min() against it is a no-op until the - # kernel overwrites it with the real (smaller) valid count. - valid_split_count = torch.full( - (bs,), num_kv_splits, dtype=dtypes.i32, device=device - ) + # Left uninitialized: the asm kernel writes the real (valid) count before + # stage2 reads it, so pre-filling would be dead work. + valid_split_count = torch.empty((bs,), dtype=dtypes.i32, device=device) use_valid_split_count_reduce = int(num_kv_splits > 1) aiter.mla_decode_stage1_asm_fwd( diff --git a/hsa/gfx1250/mla_v4/mla_a8w8_qh16_1tg_16mx4_64nx1_sparse.co b/hsa/gfx1250/mla_v4/mla_a8w8_qh16_1tg_16mx4_64nx1_sparse.co index 6e20774d7a1..f7fdd6d6b86 100755 Binary files a/hsa/gfx1250/mla_v4/mla_a8w8_qh16_1tg_16mx4_64nx1_sparse.co and b/hsa/gfx1250/mla_v4/mla_a8w8_qh16_1tg_16mx4_64nx1_sparse.co differ diff --git a/hsa/gfx1250/mla_v4/mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co b/hsa/gfx1250/mla_v4/mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co index faa5674d275..08aa808a9b4 100755 Binary files a/hsa/gfx1250/mla_v4/mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co and b/hsa/gfx1250/mla_v4/mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co differ diff --git a/op_tests/_failanalyze.py b/op_tests/_failanalyze.py new file mode 100644 index 00000000000..8643e87b8f5 --- /dev/null +++ b/op_tests/_failanalyze.py @@ -0,0 +1,73 @@ +import sys, torch +sys.path.insert(0, "op_tests") +import test_mla_v4_kargpreld as T +import aiter, aiter.mla +from aiter import dtypes + +BATCH = int(sys.argv[1]) if len(sys.argv) > 1 else 64 +KV = int(sys.argv[2]) if len(sys.argv) > 2 else 271 +SPLITS = int(sys.argv[3]) if len(sys.argv) > 3 else 8 +GQA, Q, SINK = 64, 1, True +ATOL = RTOL = 3e-2 +T._SEED = 0 + + +def build_and_run(): + inp = T._build_bf16_inputs(batch=BATCH, kv_seq_lens=KV, q_seq_logical=Q, seed=0, + gqa_ratio=GQA, attn_sink=SINK) + sm = 1.0 / (T._QUANT_D ** 0.5) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + ref, _ = T._torch_attn_decode_fp8_dequant_ref( + qp, qr, kvp, kvr, inp["qo_indptr"], inp["kv_indptr"], + inp["kv_page_indices"], inp["kv_last_page_lens"], sm, attn_sink=inp["sink"]) + total_q = inp["q_bf16"].size(0); ns = inp["qo_indptr"].size(0) - 1 + nh = T.NUM_KV_HEADS * GQA; dev = "cuda" + ob = torch.empty((total_q, GQA, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + sidx = torch.tensor([i * SPLITS for i in range(ns + 1)], dtype=torch.int32, device=dev) + lb = torch.empty((total_q, SPLITS, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev) + eb = torch.empty((total_q, SPLITS, nh, 1), dtype=dtypes.fp32, device=dev) + aiter.mla.mla_decode_fwd_v4_nm( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob, + qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=sm, out_16_nosplit=0, num_kv_splits=SPLITS, logits=lb, attn_lse=eb) + torch.cuda.synchronize() + return ref.float(), ob.float() + + +tries = int(sys.argv[4]) if len(sys.argv) > 4 else 20 +for t in range(tries): + ref, asm = build_and_run() + close = torch.isclose(ref, asm, rtol=RTOL, atol=ATOL) + pct = (~close).float().mean().item() + if pct > 0.02: + print(f"[got FAIL on try {t}] mismatch={pct:.2%}") + break +else: + print(f"no >2% fail in {tries} tries; using last run (mismatch={pct:.2%})") + +m = ~close +rv = ref[m]; av = asm[m] +delta = (av - rv).abs() +print(f"\n===== ctx={KV} batch={BATCH} split={SPLITS} : {m.sum().item()} mismatched elems =====") +print(f"ref (golden) at mismatches: min={rv.min():.3f} max={rv.max():.3f} mean|.|={rv.abs().mean():.3f}") +print(f"asm (kernel) at mismatches: min={av.min():.3f} max={av.max():.3f} mean|.|={av.abs().mean():.3f}") +print(f"|delta| at mismatches : min={delta.min():.4f} max={delta.max():.4f} mean={delta.mean():.4f}") +print(f"whole-tensor magnitude : ref |.|max={ref.abs().max():.3f} asm |.|max={asm.abs().max():.3f}") +print(f"NaN/Inf in asm? {torch.isnan(asm).any().item()} / {torch.isinf(asm).any().item()}") +# neighbor vs wild: how big is delta relative to typical value scale? +scale = ref.abs().mean().item() +print(f"\ntypical |ref| scale (all)= {scale:.3f}") +print(f"delta/scale ratio: median={ (delta/ (rv.abs()+1e-6)).median():.2f} max={(delta/(rv.abs()+1e-6)).max():.2f}") +# histogram of delta +import numpy as np +d = delta.cpu().numpy() +for thr in [0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 1e3, 1e6]: + print(f" |delta| <= {thr:>8}: {100.0*(d<=thr).mean():5.1f}%") +# sample pairs +idx = torch.argsort(delta, descending=True)[:12] +print("\ntop-12 worst (ref -> asm, delta):") +for i in idx.tolist(): + print(f" {rv[i].item():+.4f} -> {av[i].item():+.4f} d={delta[i].item():.4f}") diff --git a/op_tests/_pat.py b/op_tests/_pat.py new file mode 100644 index 00000000000..76030f54b58 --- /dev/null +++ b/op_tests/_pat.py @@ -0,0 +1,46 @@ +import sys, torch +sys.path.insert(0, "op_tests") +import test_mla_v4_kargpreld as T +import aiter, aiter.mla +from aiter import dtypes + +BATCH = int(sys.argv[1]) if len(sys.argv) > 1 else 60 +KV = int(sys.argv[2]) if len(sys.argv) > 2 else 271 +SPLITS = int(sys.argv[3]) if len(sys.argv) > 3 else 8 +Q, GQA, SINK = 1, 64, True + + +def run_once(): + inp = T._build_bf16_inputs(batch=BATCH, kv_seq_lens=KV, q_seq_logical=Q, seed=0, + gqa_ratio=GQA, attn_sink=SINK) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0); nh = T.NUM_KV_HEADS * GQA; dev = "cuda" + ns = inp["qo_indptr"].numel() - 1 + sidx = torch.tensor([i * SPLITS for i in range(ns + 1)], dtype=torch.int32, device=dev) + ob = torch.empty((total_q, GQA, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + lb = torch.zeros((total_q, SPLITS, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev) + eb = torch.full((total_q, SPLITS, nh, 1), float("-inf"), dtype=dtypes.fp32, device=dev) + aiter.mla.mla_decode_fwd_v4_nm( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob, + qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=1.0 / (T._QUANT_D ** 0.5), out_16_nosplit=0, num_kv_splits=SPLITS, + logits=lb, attn_lse=eb) + torch.cuda.synchronize() + return lb.float().clone(), eb.float().clone() + + +print(f"===== batch={BATCH} : per-split stage1 nondeterminism (3 runs, same input) =====") +L1, E1 = run_once() +L2, E2 = run_once() +L3, E3 = run_once() +print("split : logits max|r2-r1| logits max|r3-r1| lse max|r2-r1| status") +for s in range(SPLITS): + d2 = (L2[:, s] - L1[:, s]).abs().max().item() + d3 = (L3[:, s] - L1[:, s]).abs().max().item() + fe = torch.isfinite(E2[:, s]) & torch.isfinite(E1[:, s]) + de = (E2[:, s] - E1[:, s])[fe].abs().max().item() if fe.any() else 0.0 + status = "RACE" if max(d2, d3) > 1e-6 else "stable" + print(f" {s} : {d2:>13.6g} {d3:>13.6g} {de:>13.6g} {status}") diff --git a/op_tests/_stage1_bench.py b/op_tests/_stage1_bench.py new file mode 100644 index 00000000000..9fdf4fbd928 --- /dev/null +++ b/op_tests/_stage1_bench.py @@ -0,0 +1,101 @@ +"""Standalone stage1-only perf comparison: asm sparse decode kernel vs ATOM's +triton split kernel (bf16). + +- asm stage1 = the `..._sparse` .co kernel (fp8, its native dtype), timed alone. +- triton stage1 = ATOM `_paged_decode_split_kernel` (bf16), forced via kv_splits>1. + +Only the stage1 GPU kernel device-time is compared (profiler filtered by name). +NOT bit-exact / not memory-fair: triton is bf16 while asm is fp8, and D=512 here +(MLA QK is really 576-wide, so triton QK is ~11% under-counted). Perf only. + +Usage: + ENABLE_CK=0 python op_tests/_stage1_bench.py # default sweep +""" +import sys, os +sys.path.insert(0, "op_tests") +sys.path.insert(0, "/home/amd/feifei/ATOM") + +import torch +import test_mla_v4_kargpreld as T +import aiter, aiter.mla +from aiter import dtypes +from aiter.test_common import run_perftest +from atom.model_ops.v4_kernels.paged_decode import ( + _sparse_attn_v4_paged_decode_triton as triton_decode, +) + +Q = 1 +SINK = True + +# Mirror test_mla_v4_kargpreld.py sweep grids + active kernel variants. +_GQA_LIST = [64, 128] +_CTX_LENS = [256, 512, 1024] +_BATCH_SIZES = [64] +_SPLITS = [1] + +# Perf iteration counts for run_perftest (mirrors test_mla_v4_kargpreld.py's +# _PERF usage). Kept as module constants so both bench fns share them. +_PERF = {"num_iters": 50, "num_warmup": 3} + + +def bench_asm_stage1(batch, ctx, split, gqa): + inp = T._build_bf16_inputs(batch=batch, kv_seq_lens=ctx, q_seq_logical=Q, + seed=0, gqa_ratio=gqa, attn_sink=SINK) + sm = 1.0 / (T._QUANT_D ** 0.5) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0) + ns = inp["qo_indptr"].size(0) - 1 + nh = T.NUM_KV_HEADS * gqa + dev = "cuda" + ob = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + sidx = torch.tensor([i * split for i in range(ns + 1)], dtype=torch.int32, device=dev) + lb = torch.empty((total_q, split, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev) + eb = torch.empty((total_q, split, nh, 1), dtype=dtypes.fp32, device=dev) + kw = dict( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob, + qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=sm, out_16_nosplit=0, num_kv_splits=split, logits=lb, attn_lse=eb, + ) + _, us = run_perftest( + aiter.mla.mla_decode_fwd_v4_nm, + **kw, + num_iters=_PERF["num_iters"], + num_warmup=_PERF["num_warmup"], + ) + return us + + +def bench_triton_stage1(batch, ctx, split, gqa): + dev = "cuda" + D = T.V_HEAD_DIM # 512 + Tn = batch # one decode token per sequence + H = gqa + q = torch.randn((Tn, H, D), dtype=torch.bfloat16, device=dev) + total_pages = batch * ctx + unified_kv = torch.randn((total_pages, D), dtype=torch.bfloat16, device=dev) * 0.1 + kv_indices = torch.arange(total_pages, dtype=torch.int32, device=dev) + kv_indptr = torch.arange(0, (Tn + 1) * ctx, ctx, dtype=torch.int32, device=dev) + attn_sink = torch.randn((H,), dtype=torch.float32, device=dev) + sm = 1.0 / (D ** 0.5) + _, us = run_perftest( + triton_decode, + q, unified_kv, kv_indices, kv_indptr, attn_sink, sm, kv_splits=split, + num_iters=_PERF["num_iters"], + num_warmup=_PERF["num_warmup"], + ) + return us + + +import itertools + +print(f"{'gqa':>4} {'batch':>6} {'ctx':>6} {'split':>6} | {'asm_stage1_us':>14} {'triton_stage1_us':>17} {'triton/asm':>11}") +print("-" * 80) +# gqa outermost so rows mirror the kargpreld summary-table grouping. +for gqa, batch, ctx, split in itertools.product(_GQA_LIST, _BATCH_SIZES, _CTX_LENS, _SPLITS): + asm_us = bench_asm_stage1(batch, ctx, split, gqa) + tri_us = bench_triton_stage1(batch, ctx, split, gqa) + ratio = tri_us / asm_us if asm_us > 0 else float("nan") + print(f"{gqa:>4} {batch:>6} {ctx:>6} {split:>6} | {asm_us:>14.2f} {tri_us:>17.2f} {ratio:>10.2f}x") diff --git a/op_tests/_stage1_bench_pa.py b/op_tests/_stage1_bench_pa.py new file mode 100644 index 00000000000..79dafdcbbfe --- /dev/null +++ b/op_tests/_stage1_bench_pa.py @@ -0,0 +1,183 @@ +"""Standalone perf comparison: asm sparse decode kernel vs ATOM's gfx1250 gluon +`pa_decode_sparse`, with pa invoked EXACTLY the way ATOM invokes it. + +- asm = the `..._sparse` .co kernel (fp8, native), num_kv_splits=1 so it + produces its final output directly; we report its stage1 device time + (there is no real stage2 at split=1). +- pa (full) = `pa_decode_sparse(...)` called the ATOM way (see ATOM + model_ops/v4_kernels/paged_decode.py gfx1250 path): NO kv_splits (so + it is auto-inferred) and NO skip_reduce (so stage1 split-K + stage2 + reduce BOTH run when the inferred split > 1). This is what ATOM + actually pays. The inferred split is printed as `pa_split`. +- pa (s1) = the same call forced with skip_reduce=True at the same inferred + split -> ONLY the stage1 split kernel (kept for context, so the + stage2 reduce cost = pa_full_us - pa_s1_us is visible). + +NOT bit-exact / not memory-fair: pa is bf16 while asm is fp8, and D=512 here +(MLA QK is really 576-wide, so pa QK is ~11% under-counted). Perf only. +asm stays at split=1 (its natural fast path; the poc's asm stage2 merge is +pure-torch and unrepresentative), so the pa_full/asm ratio is "ATOM's real +full-pipeline cost vs asm's single-pass cost". + +Usage: + ENABLE_CK=0 python op_tests/_stage1_bench_pa.py # default sweep + ENABLE_CK=0 python op_tests/_stage1_bench_pa.py 64 512 # one combo: batch ctx + ENABLE_CK=0 python op_tests/_stage1_bench_pa.py 64 512 128 # batch ctx gqa + ENABLE_CK=0 python op_tests/_stage1_bench_pa.py 64 512 128 4 # batch ctx gqa asm_split + +The asm_split only sets the asm pass's num_kv_splits (sweep the default list via +_ASM_SPLITS). pa's split is always auto-inferred the ATOM way and is unaffected. +""" +import sys +sys.path.insert(0, "op_tests") +sys.path.insert(0, "/home/amd/feifei/ATOM") + +import torch +import test_mla_v4_kargpreld as T +import aiter, aiter.mla +from aiter import dtypes +from aiter.test_common import run_perftest +from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse + +Q = 1 +SINK = True + +# gqa + ctx mirror the kargpreld sweep; batch spans small->large so the +# auto-inferred pa split covers ATOM's large-split (stage2) and split=1 ends. +_GQA_LIST = [64, 128] +#_CTX_LENS = [256, 512, 1024] +#_BATCH_SIZES = [64] +#_ASM_SPLITS = [1] + +_CTX_LENS = [1024] +_BATCH_SIZES = [64] +_ASM_SPLITS = [1, 2, 4] + +# gfx1250/gluon constants used by pa_decode_sparse's kv_splits auto-infer. +_PA_BLOCK_K = 16 +_PA_MAX_NUM_WG = 1024 + +# Perf iteration counts for run_perftest (mirrors test_mla_v4_kargpreld.py's +# _PERF usage). Kept as module constants so both bench fns share them. +_PERF = {"num_iters": 50, "num_warmup": 2} + + +def _next_pow2(x): + return 1 << (int(x) - 1).bit_length() if x > 1 else 1 + + +def _infer_kv_splits(n_tokens, n_heads, n_indices): + """Replicate pa_decode_sparse's gfx1250 kv_splits auto-inference (so we can + force the same split for the stage1-only timing and print it).""" + block_h = max(_next_pow2(min(n_heads, 16)), 16) + n_head_blocks = -(-n_heads // block_h) # ceil + max_kv_splits = max(1, -(-n_indices // _PA_BLOCK_K)) + ks = max(1, _PA_MAX_NUM_WG // max(1, n_tokens * n_head_blocks)) + ks = min(max_kv_splits, ks) + return _next_pow2(ks) + + +def bench_asm(batch, ctx, gqa, split=1): + # split=1 is asm's natural fast path (single pass produces the final output); + # split>1 exercises the split-K stage1 + torch stage2 merge. + inp = T._build_bf16_inputs(batch=batch, kv_seq_lens=ctx, q_seq_logical=Q, + seed=0, gqa_ratio=gqa, attn_sink=SINK) + sm = 1.0 / (T._QUANT_D ** 0.5) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0) + ns = inp["qo_indptr"].size(0) - 1 + nh = T.NUM_KV_HEADS * gqa + dev = "cuda" + ob = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + sidx = torch.tensor([i * split for i in range(ns + 1)], dtype=torch.int32, device=dev) + lb = torch.empty((total_q, split, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev) + eb = torch.empty((total_q, split, nh, 1), dtype=dtypes.fp32, device=dev) + kw = dict( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), output=ob, + qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=sm, out_16_nosplit=0, num_kv_splits=split, logits=lb, attn_lse=eb, + ) + # asm mla_decode_fwd_v4_nm launches stage1 (sparse) + stage2 (merge). Profile + # both device times separately so we can report stage1 / stage2 / total and + # compare each against pa. (stage1 still lines up with the kargpreld table.) + s1_us, s2_us = T._profile_stage_times( + lambda: aiter.mla.mla_decode_fwd_v4_nm(**kw), + iters=_PERF["num_iters"], + warmup=_PERF["num_warmup"], + ) + return s1_us, s2_us + + +def bench_pa(batch, ctx, gqa): + """Time pa the ATOM way (auto split, full stage1+stage2) and, for context, + the same call forced to stage1-only. Returns (pa_s1_us, pa_full_us, split).""" + dev = "cuda" + D = T.V_HEAD_DIM # 512 + Tn = batch + H = gqa + q = torch.randn((Tn, H, D), dtype=torch.bfloat16, device=dev) + total_pages = batch * ctx + unified_kv = torch.randn((total_pages, D), dtype=torch.bfloat16, device=dev) * 0.1 + kv_indices = torch.arange(total_pages, dtype=torch.int32, device=dev) + kv_indptr = torch.arange(0, (Tn + 1) * ctx, ctx, dtype=torch.int32, device=dev) + attn_sink = torch.randn((H,), dtype=torch.float32, device=dev) + sm = 1.0 / (D ** 0.5) + + split = _infer_kv_splits(Tn, H, kv_indices.numel()) + + # pa full: EXACTLY ATOM's gfx1250 call — no kv_splits (auto-inferred to + # `split`), no skip_reduce => stage1 split + stage2 reduce both run. + _, full_us = run_perftest( + pa_decode_sparse, + q, unified_kv, kv_indices, kv_indptr, attn_sink, sm, + has_invalid=False, + num_iters=_PERF["num_iters"], + num_warmup=_PERF["num_warmup"], + ) + # pa stage1-only: same split but skip_reduce=True => only the split kernel + # (skip_reduce is a no-op when the inferred split==1, so s1==full there). + _, s1_us = run_perftest( + pa_decode_sparse, + q, unified_kv, kv_indices, kv_indptr, attn_sink, sm, + kv_splits=split, has_invalid=False, skip_reduce=True, + num_iters=_PERF["num_iters"], + num_warmup=_PERF["num_warmup"], + ) + return s1_us, full_us, split + + +import itertools + +def _ratio(pa, asm): + return pa / asm if asm > 0 else float("nan") + + +# Three side-by-side comparisons: stage1, stage2, and total (asm s1+s2 vs pa full). +print(f"{'gqa':>4} {'batch':>6} {'ctx':>6} {'asm_split':>9} {'pa_split':>8} | " + f"{'asm_s1':>8} {'pa_s1':>8} {'s1 pa/asm':>10} | " + f"{'asm_s2':>8} {'pa_s2':>8} {'s2 pa/asm':>10} | " + f"{'asm_tot':>8} {'pa_tot':>8} {'tot pa/asm':>11}") +print("-" * 128) + +if len(sys.argv) > 1: + batch = int(sys.argv[1]) + ctx = int(sys.argv[2]) if len(sys.argv) > 2 else 512 + gqa = int(sys.argv[3]) if len(sys.argv) > 3 else 64 + asm_split = int(sys.argv[4]) if len(sys.argv) > 4 else 1 + combos = [(gqa, batch, ctx, asm_split)] +else: + combos = list(itertools.product(_GQA_LIST, _BATCH_SIZES, _CTX_LENS, _ASM_SPLITS)) + +# gqa outermost so rows mirror the kargpreld summary-table grouping. +for gqa, batch, ctx, asm_split in combos: + asm_s1, asm_s2 = bench_asm(batch, ctx, gqa, asm_split) + pa_s1_us, pa_full_us, split = bench_pa(batch, ctx, gqa) + asm_tot = asm_s1 + asm_s2 + pa_s2_us = pa_full_us - pa_s1_us # pa stage2 reduce = full - stage1 + print(f"{gqa:>4} {batch:>6} {ctx:>6} {asm_split:>9} {split:>8} | " + f"{asm_s1:>8.2f} {pa_s1_us:>8.2f} {_ratio(pa_s1_us, asm_s1):>9.2f}x | " + f"{asm_s2:>8.2f} {pa_s2_us:>8.2f} {_ratio(pa_s2_us, asm_s2):>9.2f}x | " + f"{asm_tot:>8.2f} {pa_full_us:>8.2f} {_ratio(pa_full_us, asm_tot):>10.2f}x") diff --git a/op_tests/_stage1_probe.py b/op_tests/_stage1_probe.py new file mode 100644 index 00000000000..91d9207d770 --- /dev/null +++ b/op_tests/_stage1_probe.py @@ -0,0 +1,338 @@ +"""Stage1 probe for mla_decode_fwd_v4_nm. + +Pre-fills the stage1 partial buffers (logits / attn_lse) with a sentinel, +runs the v4 nm kernel at the reproducing config, then classifies every cell of +the STAGE1 output as: sane / NaN / +-Inf / illegal-huge (|x|>1e30 finite) / +unwritten (still == sentinel). Reported per split index so we can tell whether +the garbage lives in valid split slots (=> stage1 asm bug) or only in +skipped/unwritten slots (=> stage2 merge reading uninitialized memory). +""" + +import sys + +import torch + +sys.path.insert(0, "op_tests") +import test_mla_v4_kargpreld as T # noqa: E402 + +import aiter # noqa: E402 +import aiter.mla # noqa: E402 +from aiter import dtypes # noqa: E402 + +SENT = 123456.0 + + +def classify(x, name): + x = x.float() + tot = x.numel() + nan = torch.isnan(x) + inf = torch.isinf(x) + finite = torch.isfinite(x) + huge = (x.abs() > 1e30) & finite + unwritten = x == SENT + print( + f" [{name}] total={tot} nan={nan.sum().item()} inf={inf.sum().item()} " + f"illegal_huge_finite={huge.sum().item()} unwritten(sentinel)={unwritten.sum().item()}" + ) + if finite.any(): + f = x[finite & ~unwritten] + if f.numel(): + print( + f" finite(written) min={f.min().item():.4g} max={f.max().item():.4g} " + f"absmax={f.abs().max().item():.4g}" + ) + return nan, inf, huge, unwritten + + +def probe(batch, kv, q, gqa, splits, sink, seed): + inp = T._build_bf16_inputs( + batch=batch, + kv_seq_lens=kv, + q_seq_logical=q, + seed=seed, + gqa_ratio=gqa, + attn_sink=sink, + ) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + + total_q = inp["q_bf16"].size(0) + num_heads = T.NUM_KV_HEADS * gqa + dev = "cuda" + num_seqs = inp["qo_indptr"].numel() - 1 + + logits = torch.full( + (total_q, splits, num_heads, T.V_HEAD_DIM), SENT, dtype=dtypes.fp32, device=dev + ) + lse = torch.full( + (total_q, splits, num_heads, 1), SENT, dtype=dtypes.fp32, device=dev + ) + out = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + split_indptr = torch.tensor( + [i * splits for i in range(num_seqs + 1)], dtype=torch.int32, device=dev + ) + + ret_logits, ret_lse = aiter.mla.mla_decode_fwd_v4_nm( + q=qp, + qrope=qr.contiguous(), + kv_buffer=kvp, + kvrope=kvr.contiguous(), + output=out, + qo_indptr=inp["qo_indptr"], + kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], + kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=split_indptr, + max_seqlen_q=inp["max_seqlen_q"], + sink=inp["sink"], + sm_scale=1.0 / (T._QUANT_D**0.5), + out_16_nosplit=0, + num_kv_splits=splits, + logits=logits, + attn_lse=lse, + ) + torch.cuda.synchronize() + + print( + f"=== config batch={batch} kv={kv} q={q} gqa={gqa} splits={splits} " + f"sink={sink} seed={seed} ===" + ) + print(" STAGE1 logits (partials, kernel-native [total_q, splits, H, dv]):") + classify(ret_logits, "logits ALL") + print(" STAGE1 attn_lse:") + classify(ret_lse, "lse ALL") + + # Per-split breakdown of the logits partials. + L = ret_logits.float() + for s in range(splits): + Ls = L[:, s] + nan = torch.isnan(Ls).sum().item() + inf = torch.isinf(Ls).sum().item() + fin = torch.isfinite(Ls) + huge = ((Ls.abs() > 1e30) & fin).sum().item() + unw = (Ls == SENT).sum().item() + print( + f" split {s}: nan={nan} inf={inf} illegal_huge={huge} unwritten={unw}" + f" (of {Ls.numel()})" + ) + + # Merged output classification. + print(" MERGED output (stage2 -> BF16):") + classify(out, "output") + print() + + +def leak_test(batch=64, kv=271, q=1, gqa=64, splits=8, sink=True, iters=20): + """Reproduce the test's buffer-reuse: allocate logits/lse ONCE, poison every + cell with an extreme value (mimicking uninitialized torch.empty garbage that + can be ~1e35/NaN), then call the kernel `iters` times reusing the SAME + buffers WITHOUT re-poisoning. If the merged output ever shows huge/NaN/Inf, + the stage2 merge is leaking the unwritten (invalid) split slots.""" + inp = T._build_bf16_inputs( + batch=batch, kv_seq_lens=kv, q_seq_logical=q, seed=0, gqa_ratio=gqa, + attn_sink=sink, + ) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0) + num_heads = T.NUM_KV_HEADS * gqa + dev = "cuda" + num_seqs = inp["qo_indptr"].numel() - 1 + split_indptr = torch.tensor( + [i * splits for i in range(num_seqs + 1)], dtype=torch.int32, device=dev + ) + logits = torch.empty( + (total_q, splits, num_heads, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev + ) + lse = torch.empty((total_q, splits, num_heads, 1), dtype=dtypes.fp32, device=dev) + out = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + # Poison: half +1e35, and inject NaN/Inf too. + logits.fill_(1e35) + lse.fill_(1e35) + logits.view(-1)[::2] = float("nan") + lse.view(-1)[::3] = float("inf") + + print(f"=== LEAK TEST (poison=1e35/NaN/Inf, reuse buffers) x{iters} ===") + bad = 0 + for it in range(iters): + aiter.mla.mla_decode_fwd_v4_nm( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), + output=out, qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], + kv_last_page_lens=inp["kv_last_page_lens"], split_indptr=split_indptr, + max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=1.0 / (T._QUANT_D**0.5), out_16_nosplit=0, num_kv_splits=splits, + logits=logits, attn_lse=lse, + ) + torch.cuda.synchronize() + of = out.float() + nan = torch.isnan(of).sum().item() + inf = torch.isinf(of).sum().item() + huge = ((of.abs() > 1e30) & torch.isfinite(of)).sum().item() + absmax = of[torch.isfinite(of)].abs().max().item() if torch.isfinite(of).any() else float("nan") + flag = " <== LEAK" if (nan or inf or huge) else "" + if nan or inf or huge: + bad += 1 + print( + f" iter {it:2d}: out nan={nan} inf={inf} huge={huge} absmax={absmax:.4g}{flag}" + ) + print(f" => {bad}/{iters} iterations leaked garbage into merged output\n") + + +def sync_test(batch=64, kv=271, q=1, gqa=64, splits=8, sink=True, iters=10): + """Isolate the cross-iteration race: call the wrapper `iters` times reusing + the SAME buffers, once WITH torch.cuda.synchronize() between calls and once + WITHOUT. If only the no-sync run corrupts stage1, the bug is a missing + barrier between iteration N's stage2 read and iteration N+1's stage1 write + on the shared logits buffer.""" + inp = T._build_bf16_inputs( + batch=batch, kv_seq_lens=kv, q_seq_logical=q, seed=0, gqa_ratio=gqa, + attn_sink=sink, + ) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0) + num_heads = T.NUM_KV_HEADS * gqa + dev = "cuda" + num_seqs = inp["qo_indptr"].numel() - 1 + split_indptr = torch.tensor( + [i * splits for i in range(num_seqs + 1)], dtype=torch.int32, device=dev + ) + out = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + logits = torch.empty( + (total_q, splits, num_heads, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev + ) + lse = torch.empty((total_q, splits, num_heads, 1), dtype=dtypes.fp32, device=dev) + common = dict( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), + output=out, qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], + kv_last_page_lens=inp["kv_last_page_lens"], split_indptr=split_indptr, + max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=1.0 / (T._QUANT_D**0.5), out_16_nosplit=0, num_kv_splits=splits, + logits=logits, attn_lse=lse, + ) + + for label, per_iter_sync in [ + ("per-iter sync", True), + ("NO per-iter sync (back-to-back)", False), + ]: + # Fresh buffers for each phase so state can't leak between phases. + logits.fill_(0) + lse.fill_(0) + out.fill_(0) + torch.cuda.synchronize() + for _ in range(iters): + aiter.mla.mla_decode_fwd_v4_nm(**common) + if per_iter_sync: + torch.cuda.synchronize() + torch.cuda.synchronize() + Lf = logits.float() + per = [] + for s in range(splits): + xs = Lf[:, s] + fin = torch.isfinite(xs) + per.append(round(xs[fin].abs().max().item(), 1) if fin.any() else float("nan")) + of = out.float() + oh = ((of.abs() > 1e30) & torch.isfinite(of)).sum().item() + oam = of[torch.isfinite(of)].abs().max().item() + print(f"[{label}] stage1 per-split absmax={per}") + print(f" merged out huge={oh} absmax={oam:.4g}" + f"{' <== GARBAGE' if oh else ' (clean)'}") + print() + + +def perf_repro(batch=64, kv=271, q=1, gqa=64, splits=8, sink=True, reps=6): + """Reproduce the EXACT test path (run_perftest under a profiler, buffer + reuse) and classify BOTH the returned stage1 partials and the merged output + after the perf loop, to localize where the garbage enters.""" + from aiter.test_common import run_perftest + + for r in range(reps): + inp = T._build_bf16_inputs( + batch=batch, kv_seq_lens=kv, q_seq_logical=q, seed=0, gqa_ratio=gqa, + attn_sink=sink, + ) + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0) + num_heads = T.NUM_KV_HEADS * gqa + dev = "cuda" + num_seqs = inp["qo_indptr"].numel() - 1 + split_indptr = torch.tensor( + [i * splits for i in range(num_seqs + 1)], dtype=torch.int32, device=dev + ) + output_buf = torch.empty( + (total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev + ) + logits_buf = torch.empty( + (total_q, splits, num_heads, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev + ) + lse_buf = torch.empty( + (total_q, splits, num_heads, 1), dtype=dtypes.fp32, device=dev + ) + common = dict( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), + output=output_buf, qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], + kv_last_page_lens=inp["kv_last_page_lens"], split_indptr=split_indptr, + max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=1.0 / (T._QUANT_D**0.5), num_kv_splits=splits, + logits=logits_buf, attn_lse=lse_buf, + ) + # fp8-dequant reference (same as the real test), for the exact compare. + out_fp8_ref, _ = T._torch_attn_decode_fp8_dequant_ref( + qp, qr, kvp, kvr, inp["qo_indptr"], inp["kv_indptr"], + inp["kv_page_indices"], inp["kv_last_page_lens"], + 1.0 / (T._QUANT_D**0.5), attn_sink=inp["sink"], + ) + + (logits, _lse), us = run_perftest( + aiter.mla.mla_decode_fwd_v4_nm, out_16_nosplit=0, **common, + num_iters=10, num_warmup=2, num_rotate_args=1, + ) + torch.cuda.synchronize() + + Lf = logits.float() + of = output_buf.float() + # per-split absmax of stage1 partials + per_am = [] + for s in range(splits): + xs = Lf[:, s] + fin = torch.isfinite(xs) + per_am.append(round(xs[fin].abs().max().item(), 1) if fin.any() else float("nan")) + # merged output garbage + on = torch.isnan(of).sum().item() + oi = torch.isinf(of).sum().item() + ofin = torch.isfinite(of) + oh = ((of.abs() > 1e30) & ofin).sum().item() + oam = of[ofin].abs().max().item() if ofin.any() else float("nan") + # exact test-style delta + delta = (of - out_fp8_ref.float()).abs() + dmax = delta.max().item() + nbad = (delta > 3e-2).sum().item() + tot = delta.numel() + print( + f"rep {r}: stage1 per-split absmax={per_am}" + ) + print( + f" MERGED out: nan={on} inf={oi} huge={oh} absmax={oam:.4g} " + f"| vs ref: max_delta={dmax:.4g} bad={100*nbad/tot:.2f}% " + f"{'<== GARBAGE' if (on or oi or oh or dmax > 1e3) else ''}" + ) + print() + + +if __name__ == "__main__": + mode = sys.argv[1] if len(sys.argv) > 1 else "probe" + if mode == "leak": + leak_test(iters=int(sys.argv[2]) if len(sys.argv) > 2 else 20) + elif mode == "perf": + perf_repro(reps=int(sys.argv[2]) if len(sys.argv) > 2 else 6) + elif mode == "sync": + for _ in range(int(sys.argv[2]) if len(sys.argv) > 2 else 3): + sync_test() + else: + reps = int(sys.argv[1]) if len(sys.argv) > 1 else 5 + for r in range(reps): + probe(batch=64, kv=271, q=1, gqa=64, splits=8, sink=True, seed=r) diff --git a/op_tests/_stage1_verify_pa.py b/op_tests/_stage1_verify_pa.py new file mode 100644 index 00000000000..e7c107ad5d4 --- /dev/null +++ b/op_tests/_stage1_verify_pa.py @@ -0,0 +1,173 @@ +"""Standalone *correctness* comparison: asm sparse decode kernel vs ATOM's +gfx1250 gluon `pa_decode_sparse`, invoked EXACTLY the way ATOM invokes it. + +Companion to `_stage1_bench_pa.py` (which only compares stage1 *timing*). Here we +feed BOTH kernels the SAME bf16 q/kv and check their final outputs against a +shared pure-torch bf16 golden. + +`pa_decode_sparse` is called the ATOM way (see ATOM +model_ops/v4_kernels/paged_decode.py::sparse_attn_v4_paged_decode gfx1250 path): +NO `kv_splits` (so it is auto-inferred from T/H/len(kv_indices)) and NO +`skip_reduce` (so BOTH stage1 split-K AND stage2 reduce run whenever the inferred +split > 1). The auto-inferred split is printed as the `pa_split` column. The asm +side stays at num_kv_splits=1 (its logits[:,0] is already the final result); +split count is just a parallelization strategy, so the merged pa output is still +mathematically the same attention as asm and the golden. + +Why this is a fair *numeric* comparison (unlike the perf bench): + * In this poc MLA both QK score and V output run over the full 512-wide latent + (`_torch_attn_decode_bf16_golden` dots q.kv over all 512 dims for scores AND + for the value). `pa_decode_sparse` with D=512 does exactly the same math, so + D=512 is not an under-count here. + * The only intrinsic difference is dtype: asm quantizes q/kv to fp8+e8m0 + (native), while pa runs in bf16. So: + golden(bf16) vs pa(bf16) -> should be TIGHT (both bf16, same math) + golden(bf16) vs asm(fp8) -> fp8 quant noise floor + pa(bf16) vs asm(fp8) -> the two kernels' mutual difference + The last row is the "are the results consistent?" number the user asked for; + it is expected to sit near the fp8 quant floor, NOT at bf16 machine eps. + +The batch grid includes small batches (1, 16) so the auto-inferred split lands +well above 1 there, exercising ATOM's real small-batch / large-split + stage2 +path (batch=64 falls back toward split=1). + +Usage: + ENABLE_CK=0 python op_tests/_stage1_verify_pa.py # default sweep + ENABLE_CK=0 python op_tests/_stage1_verify_pa.py 64 512 # one combo: batch ctx + ENABLE_CK=0 python op_tests/_stage1_verify_pa.py 64 512 128 # batch ctx gqa +""" +import sys +sys.path.insert(0, "op_tests") +sys.path.insert(0, "/home/amd/feifei/ATOM") + +import itertools + +import torch +import test_mla_v4_kargpreld as T +import aiter, aiter.mla +from aiter import dtypes +from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse + +Q = 1 # q_seq_logical (decode) +SINK = True # exercise the per-head softmax-denom sink on both paths + +# gqa + ctx mirror _stage1_bench_pa.py; batch spans small->large so the +# auto-inferred pa split covers ATOM's large-split (stage2) and split=1 ends. +_GQA_LIST = [64, 128] +_CTX_LENS = [256, 512, 1024] +_BATCH_SIZES = [1, 16, 64] + +# gfx1250/gluon constants used by pa_decode_sparse's kv_splits auto-infer. +_PA_BLOCK_K = 16 +_PA_MAX_NUM_WG = 1024 + + +def _next_pow2(x): + return 1 << (int(x) - 1).bit_length() if x > 1 else 1 + + +def _infer_kv_splits(n_tokens, n_heads, n_indices): + """Replicate pa_decode_sparse's gfx1250 kv_splits auto-inference (for display + only; the real value is chosen inside the kernel launcher identically).""" + block_h = max(_next_pow2(min(n_heads, 16)), 16) + n_head_blocks = -(-n_heads // block_h) # ceil + max_kv_splits = max(1, -(-n_indices // _PA_BLOCK_K)) + ks = max(1, _PA_MAX_NUM_WG // max(1, n_tokens * n_head_blocks)) + ks = min(max_kv_splits, ks) + return _next_pow2(ks) + + +def _delta_stats(a, b): + """max |a-b| and mismatch ratio (rtol=atol=3e-2) as a compact (float,float).""" + a = a.float() + b = b.float() + maxd = float((a - b).abs().max().item()) + close = torch.isclose(a, b, rtol=3e-2, atol=3e-2) + ratio = float((~close).sum().item()) / max(a.numel(), 1) + return maxd, ratio + + +def run_golden(inp, sm): + out, _ = T._torch_attn_decode_bf16_golden( + inp["q_bf16"], inp["kv_bf16"], inp["qo_indptr"], inp["kv_indptr"], + inp["kv_page_indices"], inp["kv_last_page_lens"], sm, attn_sink=inp["sink"], + ) + return out # [total_q, gqa, 512] bf16 + + +def run_asm(inp, gqa, sm): + """asm stage1 with num_kv_splits=1 -> logits[:,0] is the final fp32 partial.""" + dev = "cuda" + qp, qr = T._native_to_2buff_for_asm(inp["q_bf16"]) + kvp, kvr = T._native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = inp["q_bf16"].size(0) + ns = inp["qo_indptr"].size(0) - 1 + nh = T.NUM_KV_HEADS * gqa + ob = torch.empty((total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=dev) + sidx = torch.arange(0, ns + 1, dtype=torch.int32, device=dev) # split=1 + lb = torch.empty((total_q, 1, nh, T.V_HEAD_DIM), dtype=dtypes.fp32, device=dev) + eb = torch.empty((total_q, 1, nh, 1), dtype=dtypes.fp32, device=dev) + logits, _lse = aiter.mla.mla_decode_fwd_v4_nm( + q=qp, qrope=qr.contiguous(), kv_buffer=kvp, kvrope=kvr.contiguous(), + output=ob, qo_indptr=inp["qo_indptr"], kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=sidx, max_seqlen_q=inp["max_seqlen_q"], sink=inp["sink"], + sm_scale=sm, out_16_nosplit=0, num_kv_splits=1, logits=lb, attn_lse=eb, + ) + return logits[:, 0].to(dtypes.bf16) # [total_q, gqa, 512] + + +def run_pa(inp, gqa, sm): + """pa called EXACTLY as ATOM does on gfx1250: no kv_splits (auto-inferred) + and no skip_reduce, so stage1 split-K + stage2 reduce both run when the + inferred split > 1. Returns (final bf16 output, inferred_kv_splits). + + Same bf16 q/kv as the golden/asm: kv is the 512-latent pool, kv_indices/indptr + are the asm page tables (page_size=1, last_page_len=1 => pages == kv tokens). + """ + q = inp["q_bf16"].contiguous() # [total_q, gqa, 512] + unified_kv = inp["kv_bf16"].reshape(-1, T.V_HEAD_DIM).contiguous() # [pages, 512] + kv_indices = inp["kv_page_indices"].to(torch.int32).contiguous() + kv_indptr = inp["kv_indptr"].to(torch.int32).contiguous() + out = pa_decode_sparse( + q, unified_kv, kv_indices, kv_indptr, inp["sink"].contiguous(), sm, + has_invalid=False, + ) + ks = _infer_kv_splits(q.size(0), gqa, kv_indices.numel()) + return out, ks # [total_q, gqa, 512] bf16, int + + +def verify(batch, ctx, gqa): + inp = T._build_bf16_inputs(batch=batch, kv_seq_lens=ctx, q_seq_logical=Q, + seed=0, gqa_ratio=gqa, attn_sink=SINK) + sm = 1.0 / (T._QUANT_D ** 0.5) + golden = run_golden(inp, sm) + asm = run_asm(inp, gqa, sm) + pa, pa_split = run_pa(inp, gqa, sm) + g_pa = _delta_stats(golden, pa) + g_asm = _delta_stats(golden, asm) + pa_asm = _delta_stats(pa, asm) + return pa_split, g_pa, g_asm, pa_asm + + +hdr = (f"{'gqa':>4} {'batch':>6} {'ctx':>6} {'pa_split':>8} | " + f"{'golden~pa (max|Δ|/mism%)':>26} | " + f"{'golden~asm (max|Δ|/mism%)':>27} | " + f"{'pa~asm (max|Δ|/mism%)':>24}") +print(hdr) +print("-" * len(hdr)) + +if len(sys.argv) > 1: + batch = int(sys.argv[1]) + ctx = int(sys.argv[2]) if len(sys.argv) > 2 else 512 + gqa = int(sys.argv[3]) if len(sys.argv) > 3 else 64 + combos = [(gqa, batch, ctx)] +else: + combos = list(itertools.product(_GQA_LIST, _BATCH_SIZES, _CTX_LENS)) + +for gqa, batch, ctx in combos: + pa_split, g_pa, g_asm, pa_asm = verify(batch, ctx, gqa) + print(f"{gqa:>4} {batch:>6} {ctx:>6} {pa_split:>8} | " + f"{g_pa[0]:>13.4g}/{g_pa[1]*100:>10.3f}% | " + f"{g_asm[0]:>13.4g}/{g_asm[1]*100:>11.3f}% | " + f"{pa_asm[0]:>10.4g}/{pa_asm[1]*100:>10.3f}%") diff --git a/op_tests/_v4_nm_stream_graph_check.py b/op_tests/_v4_nm_stream_graph_check.py new file mode 100644 index 00000000000..402578badfd --- /dev/null +++ b/op_tests/_v4_nm_stream_graph_check.py @@ -0,0 +1,147 @@ +"""Ad-hoc probe: verify the gfx1250 v4-nm Python .co launcher is correct on a +non-default stream AND under CUDA-graph capture/replay. + +Run: ENABLE_CK=0 python op_tests/_v4_nm_stream_graph_check.py +""" +import os +import sys + +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import aiter # noqa: E402 +import aiter.mla # noqa: E402 +from aiter import dtypes # noqa: E402 +from test_mla_v4_kargpreld import ( # noqa: E402 + _build_bf16_inputs, + _native_to_2buff_for_asm, + V_HEAD_DIM, + NUM_KV_HEADS, +) + +torch.set_default_device("cuda") + + +def _run(inp, q_packed, q_rope, kv_packed, kv_rope, num_kv_splits, logits, lse, out): + return aiter.mla.mla_decode_fwd_v4_nm( + q=q_packed, + qrope=q_rope.contiguous(), + kv_buffer=kv_packed, + kvrope=kv_rope.contiguous(), + output=out, + qo_indptr=inp["qo_indptr"], + kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], + kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=torch.tensor( + [i * num_kv_splits for i in range(inp["qo_indptr"].numel())], + dtype=torch.int32, + ), + max_seqlen_q=inp["max_seqlen_q"], + sink=inp["sink"], + sm_scale=1.0 / (512**0.5), + out_16_nosplit=0, + num_kv_splits=num_kv_splits, + logits=logits, + attn_lse=lse, + ) + + +def _alloc(total_q, num_heads, num_kv_splits): + logits = torch.empty( + (total_q, num_kv_splits, num_heads, V_HEAD_DIM), dtype=dtypes.fp32 + ) + lse = torch.empty((total_q, num_kv_splits, num_heads, 1), dtype=dtypes.fp32) + out = torch.empty((total_q, num_heads, V_HEAD_DIM), dtype=dtypes.bf16) + return logits, lse, out + + +def main(): + gqa, qlen, batch, kv = 64, 1, 8, 271 + inp = _build_bf16_inputs( + batch=batch, kv_seq_lens=kv, q_seq_logical=qlen, gqa_ratio=gqa, attn_sink=True + ) + q_packed, q_rope = _native_to_2buff_for_asm(inp["q_bf16"]) + kv_packed, kv_rope = _native_to_2buff_for_asm(inp["kv_bf16"]) + total_q = batch * qlen + num_heads = NUM_KV_HEADS * gqa + + def result_for(splits, on_stream=None): + logits, lse, out = _alloc(total_q, num_heads, splits) + if on_stream is not None: + with torch.cuda.stream(on_stream): + lg, _ = _run(inp, q_packed, q_rope, kv_packed, kv_rope, splits, logits, lse, out) + on_stream.synchronize() + else: + lg, _ = _run(inp, q_packed, q_rope, kv_packed, kv_rope, splits, logits, lse, out) + torch.cuda.synchronize() + return (out if splits > 1 else lg[:, 0].to(dtypes.bf16)).clone() + + print("== test 1: default vs non-default stream ==") + side = torch.cuda.Stream() + for splits in (1, 8): + ref = result_for(splits) + alt = result_for(splits, on_stream=side) + # NaN-safe compare (fp8 kv pad can be nan on unused split tails). + eq = torch.equal(torch.nan_to_num(ref), torch.nan_to_num(alt)) + md = (torch.nan_to_num(ref) - torch.nan_to_num(alt)).abs().max().item() + print(f" splits={splits}: side-stream bit_exact={eq} max_abs_diff={md:.3e}") + + print("== test 2: CUDA graph capture + replay ==") + splits = 1 + logits, lse, out = _alloc(total_q, num_heads, splits) + split_indptr = torch.tensor( + [i * splits for i in range(batch + 1)], dtype=torch.int32 + ) + valid = torch.full((batch,), splits, dtype=dtypes.i32) + + def call(): + aiter.mla.mla_decode_fwd_v4_nm( + q=q_packed, + qrope=q_rope.contiguous(), + kv_buffer=kv_packed, + kvrope=kv_rope.contiguous(), + output=out, + qo_indptr=inp["qo_indptr"], + kv_indptr=inp["kv_indptr"], + kv_page_indices=inp["kv_page_indices"], + kv_last_page_lens=inp["kv_last_page_lens"], + split_indptr=split_indptr, + max_seqlen_q=inp["max_seqlen_q"], + sink=inp["sink"], + sm_scale=1.0 / (512**0.5), + out_16_nosplit=0, + num_kv_splits=splits, + logits=logits, + attn_lse=lse, + ) + + # Warmup (loads the .co module OUTSIDE capture) on a side stream, per the + # torch cuda-graph recipe. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + call() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + eager = logits[:, 0].to(dtypes.bf16).clone() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + call() + # Perturb the output buffer, replay, and confirm the graph re-ran the kernel. + logits.zero_() + g.replay() + torch.cuda.synchronize() + graphed = logits[:, 0].to(dtypes.bf16).clone() + + eq = torch.equal(torch.nan_to_num(eager), torch.nan_to_num(graphed)) + md = (torch.nan_to_num(eager) - torch.nan_to_num(graphed)).abs().max().item() + print(f" graph replay bit_exact_vs_eager={eq} max_abs_diff={md:.3e}") + print("\n>>> stream+graph probe done.") + + +if __name__ == "__main__": + main() diff --git a/op_tests/mla_v4_nm_stage1_split_bug.md b/op_tests/mla_v4_nm_stage1_split_bug.md new file mode 100644 index 00000000000..4a202a82f64 --- /dev/null +++ b/op_tests/mla_v4_nm_stage1_split_bug.md @@ -0,0 +1,204 @@ +# MLA v4 nm — stage1 多 split partial 污染问题 排查报告 + +- 内核：`aiter.mla.mla_decode_fwd_v4_nm`（gfx1250 v4 nm，asm kernel `mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co`） +- 变体：`qh64-q1-16mx4-64nx1-np`（`gqa_ratio=64`, `q_seq_logical=1`, `num_kv_heads=1`） +- 测试脚本：`aiter/op_tests/test_mla_v4_kargpreld.py` +- 探针脚本：`aiter/op_tests/_stage1_probe.py`（临时诊断工具） +- 运行环境：容器 `ff_mla`，`ENABLE_CK=0` + +--- + +## 1. 复现配置 + +`test_mla_v4_kargpreld.py` 中的扫描参数： + +```python +_gfx1250_CTX_LENS = [271] # kv_seq_lens +_gfx1250_BATCH_SIZES = [60] # 触发；[30] 不触发 +_gfx1250_SPLIT_PER_BATCH = [8] # num_kv_splits +``` + +运行命令： + +```bash +ENABLE_CK=0 python op_tests/test_mla_v4_kargpreld.py +``` + +精度判据：`checkAllclose(atol=0.03, rtol=0.03, tol_err_ratio=0.02)`，即不匹配元素比例 > 2% 判 FAILED。 + +--- + +## 2. 现象总览 + +| 现象 | 结论 | +|---|---| +| 稳定复现 | 是，跨进程稳定复现，但**幅度非确定性** | +| 参考侧 | `golden_bf16 vs fp8_ref` 恒定干净（约 14 元素、delta 0.055）—— torch 参考无辜 | +| 污染位置 | stage1 输出的 **有效 split 1/2/3**（中间 split） | +| 污染值 | 恒为 **≈ `2.94e+35`** 的非法巨值（非 NaN/Inf，而是溢出/未初始化量级） | +| 头尾 split | split 0、split 4 在**所有 batch 上恒正确** | +| 无效 split | split 5/6/7（`valid_split_count=5`）—— 开关开时残留旧数据，开关关时被清 0 | +| batch 依赖 | **batch=30 不触发；batch=60 触发** | +| 分布特征 | 稀疏、随机（部分 batch + 部分中间 split + 少量 cell） | +| 根因推断 | interior split partial 写回缺同步的**竞争条件（race）** | + +--- + +## 3. 现象一：稳定复现、幅度非确定性（10 次进程级测试） + +- `golden_bf16 vs fp8_ref`（参考自检）：每次都稳定干净 → 参考侧没问题。 +- `v4_nm [fp8_dequant_ref vs asm]`（ASM 输出对比）：每次都不匹配，且幅度随机。 + +典型 10 次结果（batch=60）： + +| 指标 | 观测范围 | +|---|---| +| 不匹配比例 | 0.2% ~ 8%，间歇性越过 2% 阈值 → 有时 warning，有时 **FAILED** | +| max abs delta | 常见 **1e30 ~ 1e35 量级巨值**（正常输出应 ~1） | + +即：错误**稳定复现**，但每次污染幅度/比例都不同，是典型的非确定性数值污染。 + +--- + +## 4. 现象二：孤立单次调用干净，负载下才污染 + +stage1 探针（预填哨兵值 + 逐 split 分类）： + +- **孤立单次调用**：stage1 完全干净——split 0~4 写入正常值（~±3），无 nan/inf/巨值；split 5~7 未写入。 +- **`run_perftest` / 重复调用路径下**：stage1 出现污染。 + +说明：污染只在多次背靠背 launch / 真实显存压力下暴露，单次隔离调用不出现 → 状态相关，指向竞争 / 未初始化读。 + +--- + +## 5. 现象三：污染精确定位 —— "头尾对、中间错" + +逐 split 的 `absmax`（run_perftest 路径，典型值）： + +``` +split : 0 1 2 3 4 | 5 6 7 +absmax: 2.2 2.94e+35 2.94e+35 2.94e+35 3.0 | 0 / 0 / 0 + ─┬─ ─────────── 中间三个坏 ─────────── ─┬─ + 首端 尾端 +``` + +- **两端**（第一个有效 split 0、最后一个有效 split 4）：数值正确、稳定。 +- **中间**（split 1/2/3）：被 `≈2.94e+35` 的非法巨值污染。 +- 污染值 `2.94e+35` 在多次运行、多个 split 间高度一致（非随机内存噪声）。 + +诊断价值：很多 kernel 对首段/尾段做特殊边界处理、中间段走统一循环。"头尾对、中间错"把嫌疑缩小到 **stage1 处理 interior split 的写回/累加代码**。 + +--- + +## 6. 现象四：无效 split 5~7 与 `valid_split_count` + +### KV-split 机制 + +271 个 KV token 被切成最多 8 个 split，每个 split 由不同 workgroup 算一份 partial（stage1，写入 `logits[total_q, 8, num_heads, v_head_dim]`），再由 stage2 合并。 + +### valid_split_count = 5 + +受 kernel 粒度限制，271 token 实际只填满 **5 个 split**，于是 `valid_split_count=5`，stage2 只合并前 5 个（0~4），跳过 5/6/7。 + +### split 5~7 的两种状态 + +`logits` 由 `torch.empty` 分配（不清零）。split 5~7 是否被写取决于 `use_valid_split_count_reduce` 开关： + +| 开关状态（`aiter/mla.py` v4_nm 路径 line ~1315） | split 5~7 状态 | +|---|---| +| ON：`use_valid_split_count_reduce = int(num_kv_splits>1)` | **未写入**，保留 `torch.empty` 残留（值为 0 或 `7.16e-43`） | +| OFF：`use_valid_split_count_reduce = 0` | **被 kernel 主动清 0**（每个 split 全部 cell = 0） | + +### `7.16e-43` 的来历 + +不是随机噪声，而是**小整数被按 float32 比特重新解释**：整数 `512`（0x00000200）当 float32 读 = `512 × 2⁻¹⁴⁹ ≈ 7.17e-43`（denormal）。说明那块显存里之前躺着一个 ≈512 的整数（元数据残留）。坐实"未初始化显存旧数据"。 + +--- + +## 7. 现象五：关闭 `valid_split_count` 开关的影响 + +在 `aiter/mla.py` v4_nm 路径强制 `use_valid_split_count_reduce = 0` 后（注意：V3 路径的 line 326 测试用不到，真正生效的是 v4_nm 路径 line ~1315）： + +- **split 5~7 被完整清 0**（预填哨兵值全被覆盖为 0，`zeros = 每split全部cell`）。 +- **但 interior split 1/2/3 的污染依旧存在**（run_perftest 路径下 huge > 0，absmax 仍 `2.94e+35`）。 + +结论：关开关只改变"无效 split 是否清 0"，**不修复** interior split 污染 —— 根因与该开关无关。 + +--- + +## 8. 现象六：batch 规模依赖 + +同为 kv=271、split=8、开关关闭： + +| batch | 完整测试结果 | stage1 split 1/2/3 | +|---|---|---| +| **30** | 5/5 **passed** | 干净（huge=0，含 run_perftest 路径） | +| **60** | 间歇 **FAILED** | 被 `2.94e+35` 污染 | + +结论：污染**与 batch 规模相关**，batch=30 不触发、batch=60 触发。指向 workgroup 数量 / CU 占用 / launch geometry（`gdx = ceil(gqa*max_seqlen_q/64) × batch × split` 的 workgroup 总数）越过某临界规模后才暴露的竞争或越界。 + +--- + +## 9. 现象七：逐 batch × split 的污染分布（batch=60） + +对 60 个 batch 行（`total_q = batch × q_seq_logical = 60`）逐一统计每个 split 的 huge cell 数： + +### 头尾恒对 +- **split 0 列：60/60 全 0** +- **split 4 列：60/60 全 0** +- split 5/6/7：全 0（无效，已清 0） + +### 中间随机、部分污染 +| split | 被污染 batch 数 | 干净 batch 举例 | +|---|---|---| +| split1 | **51 / 60** | row 2,3,4,9,10,15,23,45,59 | +| split2 | **47 / 60** | row 2,3,4,5,8,9 … | +| split3 | **55 / 60** | row 11,15,21,23,53 | + +进一步： +- 几乎每个 batch（60/60 行）**至少有一个**中间 split 被污染，但**不是三个中间 split 全坏**（如 row 2 只有 split3 坏；row 11 只有 split1 坏；row 23 只有 split2 坏）。 +- 污染**稀疏**：单个 (batch, split) 内被污染 cell 数从 **2 到 ~700** 不等，最多仅占该 split 的 ~2%（每 split 每 batch 共 64×512=32768 cell）。 + +--- + +## 10. 根因推断 + +综合以上： + +1. 污染只打中 **interior split（1/2/3）**，头尾（0/4）恒对 → 边界代码正确，问题在中间段统一写回逻辑。 +2. 污染**稀疏、随机**（每次坏的 batch/split/cell 都不同、只有一小撮）→ **不是确定性索引/写覆盖 bug**（那会每次固定坏同一批位置且整段坏）。 +3. **状态相关**：单次隔离干净，重复 launch / 大 batch 才触发。 +4. 与 `valid_split_count` 开关无关。 + +→ **强烈指向竞争条件（race）**：多个 workgroup 对 interior split partial 的写回 / 累加缺少同步（缺 barrier，或对同一 LDS / 寄存器 / 全局区域存在跨 wave 的读写竞争），在 batch 规模大（workgroup 多）时被触发，随机污染少量 cell，残留出溢出量级的 `2.94e+35`。 + +--- + +## 11. 复现与探针工具 + +### 完整测试（10 次看稳定性） +```bash +docker exec ff_mla bash -lc 'cd /home/amd/feifei/aiter && \ + for i in $(seq 1 10); do \ + ENABLE_CK=0 python -u op_tests/test_mla_v4_kargpreld.py 2>&1 \ + | grep -E "fp8_dequant_ref vs asm|max abs delta.*elements"; \ + done' +``` + +### stage1 探针（`aiter/op_tests/_stage1_probe.py`） +- `probe`：预填哨兵 + 逐 split 分类 nan/inf/巨值/未写入 +- `leak`：给无效 split 灌毒（1e35/NaN/Inf），验证 stage2 是否泄漏（结论：不泄漏） +- `perf`：复刻 `run_perftest` 路径并对齐 fp8 参考 +- `sync`：对比 per-iter sync / 背靠背 + +```bash +ENABLE_CK=0 python -u op_tests/_stage1_probe.py perf 8 +``` + +--- + +## 12. 建议后续 + +1. 扫临界 batch（如 31 / 40 / 48 / 56），卡出触发阈值，佐证 launch geometry / 资源规模相关。 +2. 进入 `csrc/py_itfs_cu/asm_mla_v4.cu` 与 `.s`/`.co`，检查 interior split partial 写回是否缺 barrier、是否存在跨 wave 共享写 / LDS 未初始化读。 +3. 对照 `num_kv_splits=1`（单 split 直出）是否恒定正确，进一步锁死到多 split 写回路径。 diff --git a/op_tests/run_stage1_bench.sh b/op_tests/run_stage1_bench.sh new file mode 100755 index 00000000000..362875cad04 --- /dev/null +++ b/op_tests/run_stage1_bench.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Run the stage1-only perf comparison (_stage1_bench.py: asm sparse decode +# stage1 vs ATOM triton split stage1) INSIDE the ff_mla container, launched +# from the HOST (bare metal). +# +# ./run_stage1_bench.sh # default sweep from the .py file +# CONTAINER=ff_mla ./run_stage1_bench.sh +# +# Every argument after the script name is forwarded verbatim to the python +# script. +set -euo pipefail + +CONTAINER="${CONTAINER:-ff_mla}" +AITER_DIR="${AITER_DIR:-/home/amd/feifei/aiter}" + +# ensure container is running +if ! docker ps --format '{{.Names}}' | grep -qx "$CONTAINER"; then + echo "[host] starting container $CONTAINER ..." + docker start "$CONTAINER" >/dev/null + sleep 2 +fi + +echo "[host] container=$CONTAINER aiter_dir=$AITER_DIR args=[$*]" + +# Forward all extra args to the python script. "$@" is expanded here (on the +# host) into the remote bash -lc command string. +docker exec "$CONTAINER" bash -lc ' + cd "'"$AITER_DIR"'" + # clear JIT cache before each run + sudo rm -rf ../aiter/jit/built + sudo rm -f ../aiter/jit/*.so + ENABLE_CK=0 python3 op_tests/_stage1_bench.py '"$*"' +' diff --git a/op_tests/run_stage1_bench_pa.sh b/op_tests/run_stage1_bench_pa.sh new file mode 100755 index 00000000000..bfc1c3e7e6f --- /dev/null +++ b/op_tests/run_stage1_bench_pa.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Run the stage1-only perf comparison (_stage1_bench_pa.py: asm sparse decode +# stage1 vs ATOM gluon pa_decode_sparse stage1) INSIDE the ff_mla container, +# launched from the HOST (bare metal). +# +# ./run_stage1_bench_pa.sh # default sweep from the .py file +# ./run_stage1_bench_pa.sh 128 2048 4 # one combo: batch ctx split +# CONTAINER=ff_mla ./run_stage1_bench_pa.sh +# +# Every argument after the script name is forwarded verbatim to the python +# script, so its positional CLI args (batch ctx split) work unchanged. +set -euo pipefail + +CONTAINER="${CONTAINER:-ff_mla}" +AITER_DIR="${AITER_DIR:-/home/carhuang/feifei/aiter}" + +# ensure container is running +if ! docker ps --format '{{.Names}}' | grep -qx "$CONTAINER"; then + echo "[host] starting container $CONTAINER ..." + docker start "$CONTAINER" >/dev/null + sleep 2 +fi + +echo "[host] container=$CONTAINER aiter_dir=$AITER_DIR args=[$*]" + +# Forward all extra args to the python script. "$@" is expanded here (on the +# host) into the remote bash -lc command string. +docker exec "$CONTAINER" bash -lc ' + cd "'"$AITER_DIR"'" + # clear JIT cache before each run + sudo rm -rf ../aiter/jit/built + sudo rm -f ../aiter/jit/*.so + ENABLE_CK=0 python3 op_tests/_stage1_bench_pa.py '"$*"' +' diff --git a/op_tests/run_stage1_verify_pa.sh b/op_tests/run_stage1_verify_pa.sh new file mode 100755 index 00000000000..933d84168f4 --- /dev/null +++ b/op_tests/run_stage1_verify_pa.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Run the stage1-only CORRECTNESS comparison (_stage1_verify_pa.py: asm sparse +# decode stage1 vs ATOM gluon pa_decode_sparse stage1, both vs a torch golden) +# INSIDE the ff_mla container, launched from the HOST (bare metal). +# +# ./run_stage1_verify_pa.sh # default sweep from the .py file +# ./run_stage1_verify_pa.sh 64 512 # one combo: batch ctx +# ./run_stage1_verify_pa.sh 64 512 128 # one combo: batch ctx gqa +# CONTAINER=ff_mla ./run_stage1_verify_pa.sh +# +# Every argument after the script name is forwarded verbatim to the python +# script, so its positional CLI args (batch ctx gqa) work unchanged. +set -euo pipefail + +CONTAINER="${CONTAINER:-ff_mla}" +AITER_DIR="${AITER_DIR:-/home/amd/feifei/aiter}" + +# ensure container is running +if ! docker ps --format '{{.Names}}' | grep -qx "$CONTAINER"; then + echo "[host] starting container $CONTAINER ..." + docker start "$CONTAINER" >/dev/null + sleep 2 +fi + +echo "[host] container=$CONTAINER aiter_dir=$AITER_DIR args=[$*]" + +# Forward all extra args to the python script. "$@" is expanded here (on the +# host) into the remote bash -lc command string. +docker exec "$CONTAINER" bash -lc ' + cd "'"$AITER_DIR"'" + # clear JIT cache before each run + sudo rm -rf ../aiter/jit/built + sudo rm -f ../aiter/jit/*.so + ENABLE_CK=0 python3 op_tests/_stage1_verify_pa.py '"$*"' +' diff --git a/op_tests/run_test_mla_v4_kargpreld.sh b/op_tests/run_test_mla_v4_kargpreld.sh new file mode 100755 index 00000000000..b61eda2b9d6 --- /dev/null +++ b/op_tests/run_test_mla_v4_kargpreld.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Run the v4 nm MLA decode accuracy + perf sweep (test_mla_v4_kargpreld.py) +# INSIDE the ff_mla container, launched from the HOST (bare metal). +# +# ./run_test_mla_v4_kargpreld.sh # default sweep +# ./run_test_mla_v4_kargpreld.sh -b 64 -c 1024 2048 # override shapes +# ./run_test_mla_v4_kargpreld.sh --split-kv 1 2 4 # sweep splits +# CONTAINER=ff_mla ./run_test_mla_v4_kargpreld.sh +# +# Every argument after the script name is forwarded verbatim to the python +# script, so its CLI args (-b/--batch, -c/--kv-seq-lens, --variant, --split-kv, +# --attn-sink, --seed, --iters, --warmup) work unchanged. +set -euo pipefail + +CONTAINER="${CONTAINER:-ff_mla}" +AITER_DIR="${AITER_DIR:-/home/amd/feifei/aiter}" + +# ensure container is running +if ! docker ps --format '{{.Names}}' | grep -qx "$CONTAINER"; then + echo "[host] starting container $CONTAINER ..." + docker start "$CONTAINER" >/dev/null + sleep 2 +fi + +echo "[host] container=$CONTAINER aiter_dir=$AITER_DIR args=[$*]" + +# Forward all extra args to the python script. "$@" is expanded here (on the +# host) into the remote bash -lc command string. +docker exec "$CONTAINER" bash -lc ' + cd "'"$AITER_DIR"'" + # clear JIT cache before each run + sudo rm -rf ../aiter/jit/built + sudo rm -f ../aiter/jit/*.so + ENABLE_CK=0 python3 op_tests/test_mla_v4_kargpreld.py '"$*"' +' diff --git a/op_tests/test_mla_v4_kargpreld.py b/op_tests/test_mla_v4_kargpreld.py index 9ef76387311..47641e1bc73 100644 --- a/op_tests/test_mla_v4_kargpreld.py +++ b/op_tests/test_mla_v4_kargpreld.py @@ -82,7 +82,7 @@ class Mlagfx1250KernelVariant: _gfx1250_KERNEL_VARIANTS = [ - Mlagfx1250KernelVariant(name="qh16-q4-16mx4-64nx1-np", nhead=16, decode_qlen=4), + #Mlagfx1250KernelVariant(name="qh16-q4-16mx4-64nx1-np", nhead=16, decode_qlen=4), Mlagfx1250KernelVariant(name="qh64-q1-16mx4-64nx1-np", nhead=64, decode_qlen=1), Mlagfx1250KernelVariant(name="qh128-q1-16mx4-64nx1-np", nhead=128, decode_qlen=1), ] @@ -98,9 +98,9 @@ class Mlagfx1250KernelVariant: _SHIPPED_TILE_VARIANTS = {(16, 4), (64, 1), (128, 1)} # Default sweep grids (mirrors test_mla_gfx1250_triton.py). -_gfx1250_CTX_LENS = [13, 61, 128 + 3, 256 + 67, 1024, 4096, 16384] -_gfx1250_BATCH_SIZES = [3, 17, 32, 64] -_gfx1250_SPLIT_PER_BATCH = [1, 2, 4, 8] +_gfx1250_CTX_LENS = [1024, 1024+512, 2048] +_gfx1250_BATCH_SIZES = [64] +_gfx1250_SPLIT_PER_BATCH = [2, 4] # --------------------------------------------------------------------------- @@ -365,6 +365,43 @@ def _build_bf16_inputs( ) +# --------------------------------------------------------------------------- +# Per-stage timing: v4 nm runs two distinct GPU kernels — stage1 (the asm +# sparse decode, symbol `..._sparse`) and stage2 (the triton cross-split merge +# `_fwd_kernel_stage2_asm`). We profile a few iters and split total device time +# by kernel name so the summary table can report stage1_us / stage2_us. (single +# -pass num_kv_splits==1 has no stage2 -> stage2_us stays 0.) +# --------------------------------------------------------------------------- +_STAGE2_KERNEL_HINT = "stage2" # matches `_fwd_kernel_stage2_asm` +_STAGE1_KERNEL_HINT = "sparse" # matches the asm `..._sparse` decode symbol + + +def _profile_stage_times(fn, iters=5, warmup=2): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + s1 = s2 = 0.0 + for ev in prof.events(): + dev_us = getattr(ev, "self_device_time_total", 0) or 0 + if dev_us <= 0: + continue + nm = ev.name.lower() + if _STAGE2_KERNEL_HINT in nm: + s2 += dev_us + elif _STAGE1_KERNEL_HINT in nm: + s1 += dev_us + return s1 / iters, s2 / iters + + # --------------------------------------------------------------------------- # @benchmark perf + accuracy fn — the summary-table producer. # Its call args become the table's left-hand columns (SKILL rule 2). @@ -526,6 +563,11 @@ def test_mla_v4_nm( num_iters=_PERF["num_iters"], num_warmup=_PERF["num_warmup"], ) + stage1_us, stage2_us = _profile_stage_times( + lambda: aiter.mla.mla_decode_fwd_v4_nm( + out_16_nosplit=int(o16), **common_kwargs + ) + ) # Resolve the buffer the wrapper actually populated: # o16=1 -> wrapper unpacked packed-BF16 into output_buf. # single-pass -> kernel wrote one FP32 partial to logits[:, 0]. @@ -546,6 +588,8 @@ def test_mla_v4_nm( msg=f"{name}: mla_v4_nm [fp8_dequant_ref vs asm]", ) ret[f"{name} us"] = us + ret[f"{name} stage1_us"] = stage1_us + ret[f"{name} stage2_us"] = stage2_us ret[f"{name} TFLOPS"] = flops / us / 1e6 ret[f"{name} TB/s"] = nbytes / us / 1e6 ret[f"{name} err"] = err diff --git a/op_tests/test_mla_v4_kargpreld_oob.py b/op_tests/test_mla_v4_kargpreld_oob.py new file mode 100644 index 00000000000..fbca5e611b8 --- /dev/null +++ b/op_tests/test_mla_v4_kargpreld_oob.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Guard-page out-of-bounds (OOB) memory detector for the gfx1250 v4 'nm' MLA +decode asm kernel (mla_decode_fwd_v4_nm), generalizing the single-buffer / +single-shape probe in test_mla_v4_nm.py: + + test_mla_v4_nm.py -> only gqa=128 (tg_idx=1), only the `qrope` buffer. + THIS file -> ALL tested kernel variants {(16,4),(64,1),(128,1)} + x ALL device buffers the kernel reads AND writes. + +Principle (same as test_mla_v4_nm.py's guard page, made general): + ROCm 7.x has no compute-sanitizer, so we synthesize a guard page. Every + buffer the kernel touches is re-allocated so its DATA ends exactly on a + 2 MiB boundary (the tail of its own tight, non-caching allocation). Any + read OR write even one byte past a buffer's logical end crosses into the + next (unmapped) page and raises a GPU memory-access fault. This catches the + class of "OOB but numerically silent / non-faulting on small tensors" bugs + that accuracy / perf / plain-crash checks all miss. + + Because a GPU fault is unrecoverable (it would abort the whole process / + pytest session), each kernel launch runs in a SUBPROCESS. A clean exit 0 + + "COMPLETED no fault" means no OOB; a nonzero exit / "Memory access fault" / + "HSA_STATUS_ERROR" means the kernel over-ran a buffer. + +Buffers guarded (every device pointer handed to mla_decode_v4_asm): + reads : q, qrope, kv_buffer, kvrope, qo_indptr, kv_indptr, + kv_page_indices, kv_last_page_lens, split_indptr, sink + writes: logits, attn_lse, output + +Usage (run INSIDE the ff_mla container, gfx1250): + # default edge-focused sweep, all buffers guarded at once: + ENABLE_CK=0 python op_tests/test_mla_v4_kargpreld_oob.py + + # on any fault, re-run per-buffer to name the culprit buffer: + ENABLE_CK=0 python op_tests/test_mla_v4_kargpreld_oob.py --localize + + # widen / narrow the sweep: + ENABLE_CK=0 python op_tests/test_mla_v4_kargpreld_oob.py \ + --variant qh128-q1-16mx4-64nx1-np -b 256 -c 61 131 323 --split-kv 1 2 + + # run the FULL test_mla_v4_kargpreld.py grid (slow: one subprocess/combo): + ENABLE_CK=0 python op_tests/test_mla_v4_kargpreld_oob.py --full +""" + +import argparse +import itertools +import os +import subprocess +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import torch + +import aiter +import aiter.mla # noqa: F401 (explicit; main no longer auto-imports submodules) +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx + +import test_mla_v4_kargpreld as T + +torch.set_default_device("cuda") + +SUPPORTED_GFX = T.SUPPORTED_GFX # ["gfx1250"] + +# 2 MiB: the ROCm coarse-grain allocation granularity. Padding every buffer up +# to a whole number of 2 MiB pages and placing its data at the tail puts the +# logical end on a page boundary, so the first OOB byte lands in the (separate, +# almost-always-unmapped) next allocation and faults. +_PAGE = 2 * 1024 * 1024 + +# Every device pointer mla_decode_v4_asm receives (see aiter/mla.py). "all" +# guards every one at once (any-buffer OOB -> fault); --localize re-runs each +# in isolation to name the culprit. +_BUFFERS = [ + "q", + "qrope", + "kv_buffer", + "kvrope", + "qo_indptr", + "kv_indptr", + "kv_page_indices", + "kv_last_page_lens", + "split_indptr", + "sink", + "logits", + "attn_lse", + "output", +] + + +# --------------------------------------------------------------------------- +# Guard-page allocator. +# --------------------------------------------------------------------------- +def _tail_guarded(t): + """Return a copy of `t` (same shape/dtype/contents) whose storage ends + exactly on a 2 MiB boundary, so any read/write past its logical end hits + the unmapped next page and faults the GPU. + + Relies on PYTORCH_NO_HIP_MEMORY_CACHING=1 so the padded buffer is its own + tight hipMalloc (page-aligned base AND size); the tail slice then abuts the + allocation end. + """ + t = t.contiguous() + n = t.numel() + if n == 0: + return t + s = t.element_size() + nbytes = n * s + pad_bytes = ((nbytes + _PAGE - 1) // _PAGE) * _PAGE + pad_elems = pad_bytes // s + full = torch.empty(pad_elems, dtype=t.dtype, device=t.device) + flat = full[pad_elems - n :] # [n], contiguous, ends at the page boundary + flat.copy_(t.reshape(-1)) + return flat.view(t.shape) # shares `full`'s storage -> kept alive by view + + +def _maybe_guard(t, name, guard_set): + return _tail_guarded(t) if name in guard_set else t.contiguous() + + +# --------------------------------------------------------------------------- +# Subprocess worker body: build one config, guard the requested buffers, run. +# --------------------------------------------------------------------------- +def _oob_worker_main(gqa, q_seq_logical, batch, kv_seq_lens, num_kv_splits, + out_16_nosplit, guard_set, inject="none"): + device = "cuda" + sm_scale = 1.0 / (T._QUANT_D**0.5) # kernel ignores it; ABI needs a value + + inp = T._build_bf16_inputs( + batch=batch, + kv_seq_lens=kv_seq_lens, + q_seq_logical=q_seq_logical, + seed=T._SEED, + gqa_ratio=gqa, + attn_sink=True, + ) + q_packed, q_rope = T._native_to_2buff_for_asm(inp["q_bf16"]) + kv_packed, kv_rope = T._native_to_2buff_for_asm(inp["kv_bf16"]) + + # Self-test hook: deliberately drop the LAST kv page so the kernel's + # legitimate read of the highest page index lands exactly on the guarded + # tail's page boundary -> must fault. Proves the guard page is live (not a + # silent always-pass). kv_page_indices still references the dropped page. + if inject == "kv_short": + kv_packed = kv_packed[:-1].contiguous() + kv_rope = kv_rope[:-1].contiguous() + + total_q = inp["q_bf16"].size(0) + num_seqs = inp["qo_indptr"].size(0) - 1 + num_heads = T.NUM_KV_HEADS * gqa + + output_buf = torch.empty( + (total_q, gqa, T.V_HEAD_DIM), dtype=dtypes.bf16, device=device + ) + split_indptr = torch.tensor( + [i * num_kv_splits for i in range(num_seqs + 1)], + dtype=torch.int32, + device=device, + ) + logits_buf = torch.empty( + (total_q, num_kv_splits, num_heads, T.V_HEAD_DIM), + dtype=dtypes.fp32, + device=device, + ) + lse_buf = torch.empty( + (total_q, num_kv_splits, num_heads, 1), dtype=dtypes.fp32, device=device + ) + + kwargs = dict( + q=_maybe_guard(q_packed, "q", guard_set), + qrope=_maybe_guard(q_rope, "qrope", guard_set), + kv_buffer=_maybe_guard(kv_packed, "kv_buffer", guard_set), + kvrope=_maybe_guard(kv_rope, "kvrope", guard_set), + output=_maybe_guard(output_buf, "output", guard_set), + qo_indptr=_maybe_guard(inp["qo_indptr"], "qo_indptr", guard_set), + kv_indptr=_maybe_guard(inp["kv_indptr"], "kv_indptr", guard_set), + kv_page_indices=_maybe_guard( + inp["kv_page_indices"], "kv_page_indices", guard_set + ), + kv_last_page_lens=_maybe_guard( + inp["kv_last_page_lens"], "kv_last_page_lens", guard_set + ), + split_indptr=_maybe_guard(split_indptr, "split_indptr", guard_set), + max_seqlen_q=inp["max_seqlen_q"], + sink=_maybe_guard(inp["sink"], "sink", guard_set), + sm_scale=sm_scale, + out_16_nosplit=int(out_16_nosplit), + num_kv_splits=num_kv_splits, + logits=_maybe_guard(logits_buf, "logits", guard_set), + attn_lse=_maybe_guard(lse_buf, "attn_lse", guard_set), + ) + + aiter.mla.mla_decode_fwd_v4_nm(**kwargs) + torch.cuda.synchronize() + + +# --------------------------------------------------------------------------- +# Config enumeration. +# --------------------------------------------------------------------------- +def _valid_split(kv_seq_lens, num_kv_splits): + """Mirror test_mla_v4_kargpreld's guard: the .co inner KV loop processes + pass_size=16 tokens/iter, so the smallest split must be >= 16 (page_size=1), + else that split tail-drops (a correctness, not OOB, concern -> skip).""" + if num_kv_splits <= 1: + return True + return (kv_seq_lens // num_kv_splits) >= 16 + + +def _enumerate_configs(variants, batches, ctx_lens, splits): + """Yield (gqa, qlen, batch, kv, split, o16) for every shipped, valid combo. + o16=1 is added only where split==1 (bf16-direct is single-pass only).""" + for name in variants: + v = T._gfx1250_VARIANT_BY_KEY_NAME[name] + gqa, qlen = v.nhead, v.decode_qlen + if (gqa, qlen) not in T._SHIPPED_TILE_VARIANTS: + continue + for batch, kv, split in itertools.product(batches, ctx_lens, splits): + if not _valid_split(kv, split): + continue + yield (gqa, qlen, batch, kv, split, 0) + if split == 1: + yield (gqa, qlen, batch, kv, split, 1) + + +# --------------------------------------------------------------------------- +# Parent: launch one guard-page worker subprocess per (config, guard_set). +# --------------------------------------------------------------------------- +def _run_worker(cfg, guard_spec, timeout=300, inject="none"): + """Run a single guard-page probe in a subprocess. Returns (ok, combined_out). + `guard_spec` is "all", "none", or a single buffer name.""" + gqa, qlen, batch, kv, split, o16 = cfg + env = dict(os.environ) + env["PYTORCH_NO_HIP_MEMORY_CACHING"] = "1" # each buffer -> own tight hipMalloc + env["HSA_XNACK"] = "0" # page fault instead of demand-paging the OOB read + env["AMD_LOG_LEVEL"] = "0" # suppress noisy fault dumps + _op_dir = os.path.dirname(os.path.abspath(__file__)) + _repo_root = os.path.dirname(_op_dir) + env["PYTHONPATH"] = os.pathsep.join( + [_repo_root, _op_dir, env.get("PYTHONPATH", "")] + ).rstrip(os.pathsep) + + def _no_core_dump(): + import resource + + resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) + + argv = [ + sys.executable, __file__, "--oob-worker", + str(gqa), str(qlen), str(batch), str(kv), str(split), str(o16), + guard_spec, inject, + ] + try: + proc = subprocess.run( + argv, + env=env, + capture_output=True, + text=True, + timeout=timeout, + preexec_fn=_no_core_dump, + ) + except subprocess.TimeoutExpired as e: + # A GPU memory-access fault can wedge the HSA queue so the worker hangs + # at synchronize() instead of aborting. subprocess.run has already + # killed the child; treat the hang as a (probable) OOB fault. + out = (e.stdout or "") + (e.stderr or "") + if isinstance(out, bytes): + out = out.decode(errors="replace") + return False, out + f"\n[timeout/hang after {timeout}s -> treated as FAULT]" + + combined = proc.stdout + proc.stderr + faulted = ( + "Memory access fault" in combined + or "HSA_STATUS_ERROR" in combined + or "page fault" in combined.lower() + or proc.returncode != 0 + ) + ok = (not faulted) and ("COMPLETED no fault" in combined) + return ok, combined + + +def _cfg_label(cfg): + gqa, qlen, batch, kv, split, o16 = cfg + return ( + f"gqa={gqa:<3} qlen={qlen} batch={batch:<4} kv={kv:<6} " + f"split={split} o16={o16}" + ) + + +def run_sweep(variants, batches, ctx_lens, splits, localize=False, timeout=300): + configs = list(_enumerate_configs(variants, batches, ctx_lens, splits)) + print( + f"[oob] gfx={get_gfx()} configs={len(configs)} " + f"buffers_guarded=all({len(_BUFFERS)}) localize={localize}" + ) + print("-" * 88) + + faults = [] + for cfg in configs: + ok, out = _run_worker(cfg, "all", timeout=timeout) + status = "OK " if ok else "FAULT" + print(f"[{status}] {_cfg_label(cfg)}") + if not ok: + tail = "\n ".join(out.strip().splitlines()[-6:]) + print(f" --- worker tail ---\n {tail}") + culprit = None + if localize: + culprit = _localize(cfg, timeout=timeout) + faults.append((cfg, culprit)) + + print("-" * 88) + if not faults: + print(f"[oob] PASS: no OOB across {len(configs)} configs (all buffers).") + return True + + print(f"[oob] FAIL: {len(faults)} of {len(configs)} configs faulted:") + for cfg, culprit in faults: + extra = f" -> buffer(s): {culprit}" if culprit else "" + print(f" {_cfg_label(cfg)}{extra}") + return False + + +def _localize(cfg, timeout=300): + """A config faulted with all buffers guarded; re-run guarding ONE buffer at + a time to name which buffer(s) the kernel over-runs.""" + culprits = [] + for name in _BUFFERS: + ok, _ = _run_worker(cfg, name, timeout=timeout) + if not ok: + culprits.append(name) + return ",".join(culprits) if culprits else "unknown(none-in-isolation)" + + +def _run_selftest(timeout=300): + """Prove the guard page actually faults on a real OOB (not a silent + always-pass). Uses gqa=64 as a representative shipped variant.""" + cfg = (64, 1, 64, 131, 1, 0) # gqa, qlen, batch, kv, split, o16 + print(f"[oob][selftest] cfg: {_cfg_label(cfg)}") + + ok_short, out_short = _run_worker(cfg, "all", timeout=timeout, inject="kv_short") + print(f"[oob][selftest] kv_buffer short-by-one-page -> " + f"{'FAULT (good)' if not ok_short else 'NO FAULT (BAD: guard is dead)'}") + + ok_full, _ = _run_worker(cfg, "all", timeout=timeout, inject="none") + print(f"[oob][selftest] same config, unmodified -> " + f"{'OK (good)' if ok_full else 'FAULT (BAD: false positive)'}") + + passed = (not ok_short) and ok_full + if ok_short: # the short run should have faulted but didn't -> show why + tail = "\n ".join(out_short.strip().splitlines()[-6:]) + print(f" --- short-run worker tail ---\n {tail}") + print(f"[oob][selftest] {'PASS: guard page is live.' if passed else 'FAIL.'}") + return 0 if passed else 1 + + +# --------------------------------------------------------------------------- +# main. +# --------------------------------------------------------------------------- +def main(): + if get_gfx() not in SUPPORTED_GFX: + print(f"[oob] skip: v4 nm shipped only for {SUPPORTED_GFX}; " + f"current device is {get_gfx()}. Exiting 0.") + return 0 + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="Guard-page OOB detector for the v4 nm asm kernel: all " + "variants x all read/write buffers.", + ) + parser.add_argument( + "--variant", + nargs="*", + choices=[v.name for v in T._gfx1250_KERNEL_VARIANTS], + default=[v.name for v in T._gfx1250_KERNEL_VARIANTS], + help="Kernel variant name(s) to probe (default: all shipped).", + ) + parser.add_argument( + "-b", "--batch", type=int, nargs="*", default=[256], + help="Batch size(s). Large batch => big buffers => OOB reaches farther.", + ) + parser.add_argument( + "-c", "--kv-seq-lens", type=int, nargs="*", + default=[61, 131, 323, 1024], + help="KV context length(s). Odd/boundary values stress tail handling.", + ) + parser.add_argument( + "--split-kv", type=int, nargs="*", default=[1, 2, 4], + help="num_kv_splits value(s) (invalid small splits auto-skipped).", + ) + parser.add_argument( + "--full", action="store_true", + help="Use the full test_mla_v4_kargpreld.py grid (slow: 1 subprocess/combo).", + ) + parser.add_argument( + "--localize", action="store_true", + help="On fault, re-run per-buffer to name the over-run buffer(s).", + ) + parser.add_argument( + "--selftest", action="store_true", + help="Validate the guard page is live: a kv_buffer short by one page " + "MUST fault, the same config unmodified MUST pass. Then exit.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--timeout", type=int, default=90) + args = parser.parse_args() + + T._SEED = args.seed + + if args.selftest: + return _run_selftest(timeout=args.timeout) + + if args.full: + batches = T._gfx1250_BATCH_SIZES + ctx_lens = T._gfx1250_CTX_LENS + splits = T._gfx1250_SPLIT_PER_BATCH + else: + batches, ctx_lens, splits = args.batch, args.kv_seq_lens, args.split_kv + + ok = run_sweep( + args.variant, batches, ctx_lens, splits, + localize=args.localize, timeout=args.timeout, + ) + return 0 if ok else 1 + + +if __name__ == "__main__": + # Subprocess worker entry: `--oob-worker gqa qlen batch kv split o16 guard`. + if len(sys.argv) >= 2 and sys.argv[1] == "--oob-worker": + _gqa, _qlen, _batch, _kv, _split, _o16 = (int(x) for x in sys.argv[2:8]) + _guard_spec = sys.argv[8] if len(sys.argv) >= 9 else "all" + _inject = sys.argv[9] if len(sys.argv) >= 10 else "none" + if _guard_spec == "all": + _guard_set = set(_BUFFERS) + elif _guard_spec == "none": + _guard_set = set() + else: + _guard_set = set(_guard_spec.split(",")) + _oob_worker_main( + gqa=_gqa, q_seq_logical=_qlen, batch=_batch, kv_seq_lens=_kv, + num_kv_splits=_split, out_16_nosplit=_o16, guard_set=_guard_set, + inject=_inject, + ) + print("COMPLETED no fault") + sys.exit(0) + + sys.exit(main())