refine mla v4 co

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(

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}")

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}")

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")