diff --git a/aiter/test_common.py b/aiter/test_common.py index 75b0b3be21..4f523e23e2 100644 --- a/aiter/test_common.py +++ b/aiter/test_common.py @@ -564,6 +564,330 @@ def checkAllclose( return percent +# --------------------------------------------------------------------------- # +# DATA / SCALE init. +# gen = make_generator(seed) +# x = fill(shape, dist, gen, dtype=...) # bf16 / fp32 / float8 +# s = fill_scale(shape, dist, gen) # float32 block scale +# xq = fill_fp4(shape, dist, gen) # MXFP4 packed e2m1 +# x8 = fill_fp8(shape, dist, gen) # MXFP8 e4m3 +# s8 = fill_scale_e8m0(shape, dist, gen) # MX E8M0 on-wire +# s4 = fill_scale_e4m3(shape, dist, gen) # NVFP4 E4M3 on-wire +# +# OCP (Open Compute Project) published the MX microscaling formats: e2m1/e4m3 +# data plus a tiny E8M0 (or E4M3) scale per block. fill_fp* emit those on-wire +# buffers. Large 2-D tensors are filled in row chunks (~1 GiB f32 staging). +# --------------------------------------------------------------------------- # +DATA_DISTS = ("zero", "constant", "uniform", "norm") +SCALE_DISTS = DATA_DISTS +SCALE_UNIFORM = (0.5, 2.0) +SCALE_NORM_MEAN, SCALE_NORM_STD = 1.0, 0.25 +FP8_E4M3 = torch.float8_e4m3fn +FP4_UNIFORM = (-3.0, 3.0) # e2m1 max is 6.0; keep headroom +FP8_UNIFORM = (-6.0, 6.0) +E8M0_BIAS = 127 +E8M0_NEUTRAL = 0x7F # 2^0 = 1.0 +E4M3_NEUTRAL = 0x38 # e4m3 exp bias -> 1.0 +E4M3_SCALE_MEAN, E4M3_SCALE_STD = 0.34375, 0.08 +POW2_BINOMIAL_N = 10 +E8M0_SCALE_DISTS = ("zero", "constant", "uniform", "norm", "auto", "pow2_binomial") +E4M3_SCALE_DISTS = ("zero", "constant", "uniform", "norm", "auto") +_STAGE_ELEMS = 1 << 28 # 256M f32 = 1 GiB per chunk + + +def make_generator(seed, device="cuda"): + """Seeded ``torch.Generator`` -- same seed => bit-identical buffers.""" + return torch.Generator(device=device).manual_seed(int(seed)) + + +def add_data_init_args( + parser, *, default_dist="uniform", default_scale="constant", default_seed=0 +): + """Attach ``--data-init``, ``--scale-init`` and ``--seed``.""" + parser.add_argument( + "--data-init", + dest="data_init", + nargs="*", + choices=list(DATA_DISTS), + default=[default_dist], + help="DATA init: zero | constant | uniform | norm (N(0,1)). " + "e.g.: --data-init uniform norm", + ) + parser.add_argument( + "--scale-init", + dest="scale_init", + nargs="*", + choices=list(SCALE_DISTS), + default=[default_scale], + help="SCALE init (non-negative float): zero | constant(=1) | " + "uniform U(0.5,2) | norm N(1,0.25). Independent of --data-init.", + ) + parser.add_argument( + "--seed", + type=int, + default=default_seed, + help="RNG seed; same seed -> bit-identical uniform/norm buffers", + ) + return parser + + +def _row_chunks(rows, cols): + """Row slices whose f32 staging stays around _STAGE_ELEMS elements.""" + step = max(_STAGE_ELEMS // max(cols, 1), 1) + for start in range(0, rows, step): + yield start, min(start + step, rows) + + +def _canon_dist(dist, allowed): + if dist == "gaussian": + dist = "norm" + if dist not in allowed: + raise ValueError(f"dist {dist!r}; choose from {allowed}") + return dist + + +def _sample_data_f32(shape, dist, gen, *, lo, hi, device): + if dist == "uniform": + return torch.empty(shape, dtype=torch.float32, device=device).uniform_( + lo, hi, generator=gen + ) + if dist == "norm": + return torch.empty(shape, dtype=torch.float32, device=device).normal_( + 0.0, 1.0, generator=gen + ) + raise ValueError(f"data dist {dist!r} is not continuous; use fill dispatch") + + +def _sample_scale_f32(shape, dist, gen, *, lo, hi, device): + if dist == "uniform": + v = torch.empty(shape, dtype=torch.float32, device=device).uniform_( + lo, hi, generator=gen + ) + elif dist == "norm": + v = torch.empty(shape, dtype=torch.float32, device=device).normal_( + SCALE_NORM_MEAN, SCALE_NORM_STD, generator=gen + ) + else: + raise ValueError(f"scale dist {dist!r} is not continuous; use fill_scale") + v.clamp_(min=0.0) + return v + + +def _fill_sampled(shape, dist, gen, *, dtype, device, uniform, constant, sample_fn): + if dist == "zero": + return torch.zeros(shape, dtype=dtype, device=device) + if dist == "constant": + return torch.full(shape, constant, dtype=dtype, device=device) + lo, hi = uniform + if len(shape) != 2: + return sample_fn(shape, dist, gen, lo=lo, hi=hi, device=device).to(dtype) + rows, cols = shape + out = torch.empty(shape, dtype=dtype, device=device) + for r0, r1 in _row_chunks(rows, cols): + v = sample_fn((r1 - r0, cols), dist, gen, lo=lo, hi=hi, device=device) + out[r0:r1] = v.to(dtype) + del v + return out + + +def fill( + shape, + dist, + gen, + *, + dtype=torch.float32, + device="cuda", + uniform=(-1.0, 1.0), + constant=1.0, +): + """Return a ``dtype`` DATA tensor of ``shape``. + + ``dist`` in {zero, constant, uniform, norm}. ``uniform`` is U(lo, hi); + ``norm`` / ``gaussian`` is N(0, 1). ``zero`` / ``constant`` ignore ``gen``. + """ + dist = _canon_dist(dist, DATA_DISTS) + return _fill_sampled( + shape, + dist, + gen, + dtype=dtype, + device=device, + uniform=uniform, + constant=constant, + sample_fn=_sample_data_f32, + ) + + +def fill_scale( + shape, + dist, + gen, + *, + dtype=torch.float32, + device="cuda", + uniform=SCALE_UNIFORM, + constant=1.0, +): + """Return a non-negative float SCALE tensor of ``shape``. + + Same dist names as ``fill``, sampled independently. ``constant`` defaults + to 1.0 (neutral). ``norm`` is N(1, 0.25) clamped >= 0 -- not DATA's N(0,1). + For MX on-wire scales use ``fill_scale_e8m0`` / ``fill_scale_e4m3``. + """ + dist = _canon_dist(dist, SCALE_DISTS) + return _fill_sampled( + shape, + dist, + gen, + dtype=dtype, + device=device, + uniform=uniform, + constant=constant, + sample_fn=_sample_scale_f32, + ) + + +def _f32_to_e8m0(v: torch.Tensor) -> torch.Tensor: + """Round positive floats to the nearest E8M0 on-wire byte (bias 127).""" + e = torch.zeros_like(v, dtype=torch.int32) + pos = v > 0 + e[pos] = v[pos].log2().round().to(torch.int32) + E8M0_BIAS + return e.clamp_(0, 255).to(torch.uint8) + + +def _popcount64(x: torch.Tensor) -> torch.Tensor: + """Population count for a non-negative int64 tensor (SWAR bit-hack).""" + x = x - ((x >> 1) & 0x5555555555555555) + x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333) + x = (x + (x >> 4)) & 0x0F0F0F0F0F0F0F0F + return (x * 0x0101010101010101) >> 56 + + +def fill_fp4(shape, dist, gen, *, uniform=FP4_UNIFORM, device="cuda", constant=0): + """MXFP4 on-wire: packed e2m1 ``uint8`` of shape ``(rows, cols // 2)``. + + Samples with the same DATA dists as ``fill``, then round-to-nearest e2m1. + ``shape`` is the logical ``(rows, cols)``; ``cols`` must be even. + """ + dist = _canon_dist(dist, DATA_DISTS) + rows, cols = shape + assert cols % 2 == 0, f"FP4 needs even columns, got {cols}" + packed = (rows, cols // 2) + if dist == "zero": + return torch.zeros(packed, dtype=torch.uint8, device=device) + if dist == "constant": + return torch.full(packed, int(constant), dtype=torch.uint8, device=device) + + from aiter.utility import fp4_utils # local: fp4_utils pulls in triton + + out = torch.empty(packed, dtype=torch.uint8, device=device) + for r0, r1 in _row_chunks(rows, cols): + v = _sample_data_f32( + (r1 - r0, cols), + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + ) + out[r0:r1] = fp4_utils.f32_to_mxfp4(v).view(torch.uint8) + del v + return out + + +def fill_fp8(shape, dist, gen, *, uniform=FP8_UNIFORM, device="cuda", constant=0.5): + """MXFP8 on-wire: e4m3 tensor of ``shape``.""" + dist = _canon_dist(dist, DATA_DISTS) + if dist == "zero": + return torch.zeros(shape, dtype=FP8_E4M3, device=device) + if dist == "constant": + return torch.full( + shape, float(constant), dtype=torch.float32, device=device + ).to(FP8_E4M3) + if len(shape) != 2: + v = _sample_data_f32( + shape, dist, gen, lo=uniform[0], hi=uniform[1], device=device + ) + return v.to(FP8_E4M3) + rows, cols = shape + out = torch.empty(shape, dtype=FP8_E4M3, device=device) + for r0, r1 in _row_chunks(rows, cols): + v = _sample_data_f32( + (r1 - r0, cols), + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + ) + out[r0:r1] = v.to(FP8_E4M3) + del v + return out + + +def fill_scale_e8m0( + shape, + dist="auto", + gen=None, + *, + device="cuda", + n=POW2_BINOMIAL_N, + constant=E8M0_NEUTRAL, +): + """MX E8M0 on-wire ``uint8`` (biased exponent, bias 127). + + ``zero`` / ``constant`` / ``uniform`` / ``norm`` map from our SCALE dists + (float then round to nearest power-of-two byte). ``auto`` / + ``pow2_binomial`` match the MX GEMM default: 2^(Binomial(21,0.5)-11). + """ + if dist == "gaussian": + dist = "norm" + if dist not in E8M0_SCALE_DISTS: + raise ValueError(f"E8M0 scale dist {dist!r}; choose from {E8M0_SCALE_DISTS}") + if dist == "zero": + return torch.zeros(shape, dtype=torch.uint8, device=device) + if dist == "constant": + return torch.full(shape, int(constant), dtype=torch.uint8, device=device) + if dist in ("uniform", "norm"): + v = fill_scale(shape, dist, gen, device=device) + return _f32_to_e8m0(v) + # auto / pow2_binomial: Binomial(k, 0.5) == popcount of a uniform k-bit int + trials = 2 * n + 1 + assert trials <= 24, "pow2_binomial popcount path assumes <= 24 trials" + bits = torch.randint( + 0, 1 << trials, shape, dtype=torch.int64, device=device, generator=gen + ) + e = _popcount64(bits).to(torch.int32) - (n + 1) + return (e + E8M0_BIAS).clamp_(0, 255).to(torch.uint8) + + +def fill_scale_e4m3( + shape, dist="auto", gen=None, *, device="cuda", constant=E4M3_NEUTRAL +): + """NVFP4 / E4M3 on-wire ``uint8``. + + ``auto`` -> N(0.34375, 0.08) clamped >= 0, then cast e4m3 (MX GEMM default). + ``uniform`` / ``norm`` use ``fill_scale`` then cast. ``constant`` is 0x38 + (1.0). + """ + if dist == "gaussian": + dist = "auto" + if dist not in E4M3_SCALE_DISTS: + raise ValueError(f"E4M3 scale dist {dist!r}; choose from {E4M3_SCALE_DISTS}") + if dist == "zero": + return torch.zeros(shape, dtype=torch.uint8, device=device) + if dist == "constant": + return torch.full(shape, int(constant), dtype=torch.uint8, device=device) + if dist in ("uniform", "norm"): + v = fill_scale(shape, dist, gen, device=device) + return v.to(FP8_E4M3).view(torch.uint8) + v = torch.empty(shape, dtype=torch.float32, device=device).normal_( + E4M3_SCALE_MEAN, E4M3_SCALE_STD, generator=gen + ) + v.clamp_(min=0.0) + return v.to(FP8_E4M3).view(torch.uint8) + + def tensor_dump(x: torch.Tensor, name: str, dir="./"): x_cpu = x.cpu().view(torch.uint8) filename = f"{dir}/{name}.bin" diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py new file mode 100644 index 0000000000..32b511714f --- /dev/null +++ b/op_tests/bench_gfx1250_combo.py @@ -0,0 +1,1847 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Combined gfx1250 asm-kernel perf bench. + +Imports the top-level @benchmark sweep fns from the aiter op_tests (which the +aiter-op-test skill keeps importable for exactly this kind of combination +testing) and runs each over its own shape axes. + +Output discipline: this script prints ONLY the per-op summary tables. All the +underlying noise (per-config "calling ..." logs, JIT build output, aiter import +banners, pandas/torch/ROCTracer warnings, including C-level fd writes) is +silenced via os-level fd redirection while the kernels run; the markdown tables +are then printed to real stdout. + +Run from the aiter repo root so `op_tests/` siblings import cleanly: + + cd /app/aiter + # Hardware-oriented single-op performance (no model-specific shape contract). + python op_tests/bench_gfx1250_combo.py --perf # all perf ops + python op_tests/bench_gfx1250_combo.py --perf --ops mha # MHA + python op_tests/bench_gfx1250_combo.py --perf --ops moe # grouped MoE FC1/FC2 (a4w4 + a8w4) + python op_tests/bench_gfx1250_combo.py --perf --ops gemm # F4GEMM + python op_tests/bench_gfx1250_combo.py --perf --ops f8gemm # F8GEMM + python op_tests/bench_gfx1250_combo.py --perf --ops mla_v4_decode # MLA v4 decode + + # DeepSeek-V4 operators at the model shapes used by the DSv4 workload. + python op_tests/bench_gfx1250_combo.py --dsv4 # all DSv4 ops + python op_tests/bench_gfx1250_combo.py --dsv4 --ops moe # grouped MoE FC1/FC2 (DSv4 a8w4) + python op_tests/bench_gfx1250_combo.py --dsv4 --ops a8w8_blockscale # DSv4 FP8 linears + python op_tests/bench_gfx1250_combo.py --dsv4 --ops a16w16 # DSv4 BF16 linears + python op_tests/bench_gfx1250_combo.py --dsv4 --ops mla_v4_decode # sparse MLA v4 decode + python op_tests/bench_gfx1250_combo.py --dsv4 --ops inverse_rope # inverse RoPE + group quant + python op_tests/bench_gfx1250_combo.py --dsv4 --ops mla_v4_prefill # MLA v4 prefill + python op_tests/bench_gfx1250_combo.py --dsv4 --ops mhc # mHC fused RMSNorm + python op_tests/bench_gfx1250_combo.py --dsv4 --ops qk_norm # QK norm + RoPE + python op_tests/bench_gfx1250_combo.py --dsv4 --ops score_qk # FP8 paged MQA logits + python op_tests/bench_gfx1250_combo.py --dsv4 --ops mori_ep # MORI EPv2 dispatch/combine + python op_tests/bench_gfx1250_combo.py --dsv4 --ops mega_moe # Mega on/off, 4 GPUs + +Environment +----------- + +One variable, applied to every op that sweeps a token count: + + AITER_BENCH_TOKENS=1,128,512 python op_tests/bench_gfx1250_combo.py --dsv4 + +Leave it unset and every op runs the default in this file, which is the tested +configuration -- the shapes below are what the suites are expected to pass on. +Set it and it wins everywhere, with no second-guessing: an explicit request is +the caller's to make, including for shapes an op is known to fail. + +The defaults are not one list, because a single token count does not mean the +same thing to every op: + + mla_v4_decode 1..1024. Decode carries one token per sequence, so the + axis is really the batch, and 65536 is not a shape the + model runs. + inverse_rope 1..16384. The axis is -s at a fixed -b 128,16, and 65536 + faults -- in the triton reference the UT compares against, + not in the kernel under test. + mega_moe 1..2048. 65536 cannot allocate its symmetric arena; see + _MEGA_MOE_TOKENS. + a8w8_blockscale 1024..65536. Below 1024 it walks into a UT bug; see the + note in DSV4_OPS. + mla_v4_prefill 1024..16384, the DSv4 prefill chunk. 65536 faults; see + _MLA_PREFILL_TOKENS. + +With the variable unset, the child-UT ops (score_qk, a8w8_blockscale) pass no +shape flag at all, so each UT sweeps the range its owner maintains. The +in-process ops (moe, a16w16, mha, mla_v4_prefill) iterate shapes here and take +their default from the module. + +Other variables: + + ENABLE_CK=0 set before importing aiter; the module + already setdefault()s it. + GPU_ARCHS / CU_NUM detected once here and exported to every + child, so no child runs rocminfo. Four + ranks starting at once contend for + rocminfo's rocm_smi mutex and a rank can + lose it outright -- see _pin_arch. Set + either yourself and yours wins. + +mega_moe at tokens/rank=65536 fails in setup(), asking 7.5 GB for cco's VMM +arena against a 4 GiB default. MORI_SHMEM_HEAP_SIZE does not reach that arena +(see run_mega_moe), so exporting it changes nothing -- and exporting it +sweep-wide takes the machine down, because that heap is preallocated per rank +for every case. The tier is out of the sweep; fixing it means passing +per_rank_vmm at Communicator.init(). + +Failures do not stop the sweep: a case that aborts is recorded and the run +moves to the next one, with a "N failed, M ops selected" list at the end and a +non-zero exit code. A GPU fault inside this process is the exception -- it +takes the interpreter down and no handler runs, which is why the child-UT ops +are the ones that survive their own crashes. + +``--ops`` accepts any op, including one held out of a suite's defaults because +it is broken on the current arch, so it can be re-checked on a newer image. + +The ``mori_ep`` op runs the EPv2 benchmark from ``${MORI:-/app/mori}`` as the +image provides it -- this script never updates or installs mori. Environment +variables select backend, token tiers, eager/graph modes, EP size, dispatch +dtype, and correctness checking: + + TOKENS=512 MODES=graph \ + python op_tests/bench_gfx1250_combo.py --dsv4 --ops mori_ep + +It sweeps two dispatch wires by default, bf16 and fp4, one child process each +(bench_ep.py reads $DISP at import). fp4 is the wire DSv4 serves on; bf16 is +the reference. $DISP overrides, comma-separated -- DISP=fp8 or DISP=bf16,fp8,fp4. +Note that mori disables its own correctness check on fp4, so those rows are +unchecked rather than verified; the table label says so. + +The ``mhc`` op runs: + + python3 op_tests/test_mhc.py -n 7168 -m 512 --fuse_rmsnorm + +The ``qk_norm`` op runs both DSv4 phases: + + python3 op_tests/test_flydsl_qk_norm_rope_quant.py \ + -T 16384 --H 128 --D 512 --RD 64 --no-quant --qweight + + python3 op_tests/test_flydsl_qk_norm_rope_quant.py \ + -T 1 2 4 8 16 32 64 --H 128 --D 512 --RD 64 --no-quant --qweight + +The ``score_qk`` op runs two decode KV lengths at batch 512: + + python3 op_tests/op_benchmarks/triton/bench_deepgemm_attention.py \ + --batch 512 --heads 64 --index_dim 128 -kv_length 384 -mtp 0 \ + --kv_preshuffle --blocksize 64 + + python3 op_tests/op_benchmarks/triton/bench_deepgemm_attention.py \ + --batch 512 --heads 64 --index_dim 128 -kv_length 10240 -mtp 0 \ + --kv_preshuffle --blocksize 64 + +score-QK is a decode op, so its KV length is the average context one decode +step scans -- input + output/2 -- after CSA's 4x KV compression: + + 1K in / 1K out -> (1024 + 512) / 4 = 384 + 16K in / 4K out -> (16384 + 2048) / 4 = 4608 + 32K in / 16K out -> (32768 + 8192) / 4 = 10240 + +The DSv4 ``mla_v4_decode`` op runs sparse decode with GQA/H=128, batch=512 and +q_seq=1 (M=512), sweeping KV lengths 256/512/1024 and split counts 1/2/4. + +The DSv4 ``mla_v4_prefill`` op runs eight performance cases at H=128 and +D=512: compressed prefix-pool rows 4096/16384, crossed with dense/sparse CSR +modes, crossed with fp8/bf16. The current chunk remains uncompressed: + + python3 op_tests/test_pa_sparse_prefill.py \ + -n --h_q 128 -d 512 \ + --total_pages 4096 16384 --total_tokens \ + --prec fp8 bf16 --mode dense sparse --no-verify + +The UT compares the backends it has for each precision -- opus and asm on fp8, +opus and triton on bf16 -- so one run covers both precisions and all three +backends. There is no nnz axis to sweep: the CSR is generated from --mode +(sparse draws a random nnz per row, dense fills every row) under --seed, so +nnz is an outcome, not an input. + +The ``inverse_rope`` op runs the tp1 attention-output shape (-b is +(n_local_heads, n_local_groups); 128,16 is V4-Pro at dp/tp1): + + python3 op_tests/test_inverse_rope_group_quant.py \ + -b 128,16 -s -l n32k4 --group-size 32 + +The ``a8w8_blockscale`` op runs: + + python3 op_tests/test_gemm_a8w8_blockscale.py \ + -m 512 \ + -nk 2048,7168 7168,16384 6144,7168 \ + 7168,3072 65536,1536 8192,1536 \ + --ck_preshuffle True --flydsl + +The ``a16w16`` op uses ``test_opus_a16w16_gemm.py`` with batch=1, M=512, +K=7168 and N=64,384,1024,2048,32320,129280. + +The ``mega_moe`` op runs both sides of the comparison: + + MORI_V2_KERNEL_BACKEND=hip MEGA_DISPATCH=mori \ + torchrun --standalone --nproc_per_node=4 \ + op_tests/multigpu_tests/test_mega_moe_gfx1250.py \ + -e 384 -k 6 -hd 7168 -id 3072 \ + --layers 61 -tpr 512 --combine scatter_fused \ + --acc_verify 0 --profile_table 1 + + MORI_V2_KERNEL_BACKEND=hip MEGA_DISPATCH=mori \ + torchrun --standalone --nproc_per_node=4 \ + op_tests/multigpu_tests/test_mega_moe_gfx1250.py \ + -e 384 -k 6 -hd 7168 -id 3072 \ + --layers 61 -tpr 512 --combine gather \ + --acc_verify 0 --profile_table 1 + +Token sweeps come from one variable, AITER_BENCH_TOKENS (see Environment +above). Unset, each op runs its own default -- the ops do not share a supported +range, so those defaults differ and each says why at its constant. Set, it +applies to every op that sweeps tokens, and the file does not argue with it. + +``a16w16`` is not one of them: its range is a function of what opus has tuned, +not a fixed limit. Shapes with no tuned winner fall back to a split-K kid whose +launcher is 32-bit gmem-descriptor bound, which is both slow and, at M=65536, +wrong. Re-tuning through csrc/gemm_a16w16/gemm_a16w16_tune.py --libtype opus is +what widens the range, so the bench predicts nothing and reports what it gets. + +gfx1250's bundled CK does not compile, so the asm JIT modules must be built with +ENABLE_CK=0. The script sets it (before importing aiter) so a plain run just +works; an explicit env override still wins. +""" + +import os + +# Must be set BEFORE `import aiter` so the JIT build picks it up. setdefault => +# an explicitly-exported ENABLE_CK from the caller is respected. +os.environ.setdefault("ENABLE_CK", "0") + +# FlyDSL MoE env vars — must be set before importing aiter / moe test module. +os.environ.setdefault("AITER_USE_GROUPED_GEMM", "1") +os.environ.setdefault("AITER_GROUPED_DEBUG", "0") +os.environ.setdefault("FLYDSL_DUMP_IR", "1") +os.environ.setdefault("AITER_LOG_MORE", "1") +os.environ.setdefault("AITER_MOE_EXPERT_BALANCE", "true") +os.environ.setdefault("AITER_FLYDSL_MOE_EXPERT_SCHEDULING_MODE", "1") +os.environ.setdefault("AITER_FORCE_GFX1250", "1") + +import argparse +import contextlib +import itertools +import subprocess +import sys +import tempfile +import warnings + +warnings.filterwarnings("ignore") + + +@contextlib.contextmanager +def _silence(): + """Discard everything written to stdout/stderr — including native (C/C++) + fd writes (ROCTracer, hipcc, aiter logger) — for the duration of the block. + Redirects at the OS fd level so it catches more than sys.stdout swapping.""" + devnull = os.open(os.devnull, os.O_WRONLY) + # Flush any buffered Python-level output to the REAL fds BEFORE redirecting. + # stdout is block-buffered when piped/redirected, so an earlier _print_table() + # can still be sitting in the buffer; without this flush it would drain to + # devnull once fd 1 is redirected here and the printed table would be lost. + sys.stdout.flush() + sys.stderr.flush() + old1, old2 = os.dup(1), os.dup(2) + try: + os.dup2(devnull, 1) + os.dup2(devnull, 2) + yield + finally: + # Flush again BEFORE restoring so anything printed inside the block goes + # to devnull (not the real stdout after we restore it). + sys.stdout.flush() + sys.stderr.flush() + os.dup2(old1, 1) + os.dup2(old2, 2) + os.close(devnull) + os.close(old1) + os.close(old2) + + +# Import aiter + the op-test modules quietly (import-time banners suppressed). +with _silence(): + import pandas as pd + import test_f4gemm as gemm_mod + import test_flydsl_grouped_gemm_gfx1250 as moe_mod + import test_fmha_fwd_with_sink_asm as mha_mod # has __main__ guard + import test_mla_v4_kargpreld as mla_v4_kargpreld_mod + import test_mxfp8fp4gemm as f8gemm_mod + import test_opus_a16w16_gemm as a16w16_mod + import torch + from triton_tests.attention import test_mla_v4_triton as mla_v4_triton_mod + + import aiter + from aiter import dtypes + from aiter.jit.utils.chip_info import get_cu_num, get_gfx + from aiter.test_common import run_perftest + +SUPPORTED_GFX = ["gfx1250"] +# a16w16 N shapes at K=7168: attention/router projections, then lm_head twice +# (129280 is the DeepSeek vocab, 32320 is that sharded over TP4). +_A16W16_NS = (64, 384, 1024, 2048, 32320, 129280) +# lm_head cap. A shape judgement, not a kernel limit: lm_head runs one row per +# sequence, so M past this is not something the model produces, and M*N alone +# is 16 GB of bf16 output at (65536, 129280). +_A16W16_WIDE_N = 2048 +_A16W16_WIDE_N_MAX_M = 2048 +# a16w16 returns its own error ratio, and a wrong answer here is silent: the UT +# neither raises nor prints a warning. Measured on gfx1250 / 20260827, every +# shape that computed correctly came back 0 or ~1e-5, while M=65536 came back +# 0.96-0.99 on all four of its N -- an unrelated result, not a tolerance miss. +# Anything above this is reported as a failed op rather than printed as data. +_A16W16_MAX_ERR = 1e-2 + + +def _tokens(default=None): + """Token sweep from AITER_BENCH_TOKENS, else the op's own default. + + One variable for the whole bench. Ops whose axis is not a token count, or + whose usable range is fixed, ignore it and pin their sweep in the source + instead -- see _MLA_DECODE_TOKENS and _INVERSE_ROPE_TOKENS. + + Returns None when the variable is unset and the op has no default of its + own: the op then passes no shape flag at all and the UT sweeps its own + default, which is the range its owner keeps working. + """ + raw = os.environ.get("AITER_BENCH_TOKENS") + if raw: + return tuple(int(t) for t in raw.replace(",", " ").split()) + return tuple(default) if default is not None else None + + +# The in-process ops call their UT per shape, so they need a sweep to iterate; +# keep one here. The ops that shell out pass no shape flag unless asked, letting +# each UT sweep the range its owner maintains. +_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024, 2048, 65536)) +# a16w16's M is the token count, and the global sweep jumps 2048 -> 65536, so +# the prefill chunk sizes never got measured on the BF16 linears. Add them to +# this op's default; AITER_BENCH_TOKENS still overrides the whole thing. +_A16W16_MS = _tokens(tuple(sorted({*_TOKENS, 4096, 8192, 16384}))) +# This axis is -s (sequence length) at a fixed -b 128,16, not the token count +# the other ops sweep, so the default is its own rather than _TOKENS. 65536 is +# left off it: that value faults, and in the triton reference the UT compares +# against rather than in the kernel under test. AITER_BENCH_TOKENS still wins if +# set -- what an explicit request sweeps is the caller's business. +_INVERSE_ROPE_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024, 2048, 16384)) +_SCORE_QK_TOKENS = _tokens() +# score_qk is decode, so its KV length is the average context a decode step +# scans: input + output/2, then CSA's 4x compression. +# 1K in / 1K out -> (1024 + 512) / 4 = 384 +# 16K in / 4K out -> (16384 + 2048) / 4 = 4608 +# 32K in / 16K out-> (32768 + 8192) / 4 = 10240 +_SCORE_QK_KV_LENGTHS = ( + ("1K/1K average", "384"), + ("16K/4K average", "4608"), + ("32K/16K average", "10240"), +) +# Was unset, which let the UT sweep its own 27-value default down to M=1. Two +# reasons to set it. First, M here is the token count of one step, so the small +# end of that default is decode batch and the large end is prefill chunk; this +# list is the prefill side, up to the 65536 the other DSv4 ops sweep and past +# the UT default's own ceiling of 10240. Second, the small M are what walk into +# the UT bug described at "a8w8_blockscale" below: get_CKGEMM_config retries the +# lookup as M -> get_padded_m(gl=0) -> nextPow2, so anything in [1, 16] or +# [33, 64] can land on one of #4773's M=16/M=64 gluon rows (gemm_common.cu:13). +# Starting at 1024 clears both ranges by a wide margin. +# +# Two things stop being covered, both worth remembering: decode-side M, and the +# 11 tuned rows that are the only shapes dispatching to gluon. This is a way +# around the UT bug, not a fix for it. +_A8W8_BLOCKSCALE_TOKENS = _tokens((1024, 2048, 4096, 8192, 16384, 65536)) +# Decode carries one token per sequence, so this axis is the batch, not a token +# count; past 1024 it stops being a shape the model runs, hence its own default +# rather than _TOKENS. AITER_BENCH_TOKENS overrides it like everywhere else. +_MLA_DECODE_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024)) +# Up to 16384, the DSv4 prefill chunk. Re-measured on the #5084 UT (20260901, +# b45-2), one process per tier, all with the --no-verify below: +# 1024 .. 16384 clean, no coredump (4096/8192/16384 faulted on the old UT) +# 65536 Memory access fault, and an 89 GB coredump with it +# 65536 stays out: it is past the chunk size the model prefills, and one fault +# costs a third of the host's free disk. Its fault also looks unrelated to the +# others -- address 0x7f2ddbec0000, a mapped high address, where the old UT's +# faults were low wild pointers like 0xc00000. +# +# "Clean" here means the kernel did not fault, NOT that it computed correctly. +# Correctness cannot be checked on gfx1250 at all right now: drop --no-verify +# and even n=1024 dies at the first case (fp8/dense, fault at 0x43000), so the +# reference or the comparison is what breaks, not the kernel under test. Until +# that is fixed these are timings from an unverified kernel -- the same footing +# as a16w16's M=65536 rows before _A16W16_MAX_ERR caught them. +_MLA_PREFILL_TOKENS = _tokens((1024, 2048, 4096, 8192, 16384)) +# Default stops at 2048: tokens/rank=65536 dies in pipe.setup() building the +# symmetric arena -- cco sizes it from Communicator.DEFAULT_PER_RANK_VMM (4 GiB) +# and asks for 7.5 GB. That is a per_rank_vmm the UT never passes, not something +# MORI_SHMEM_HEAP_SIZE reaches, so the tier cannot run from here. Ask for it via +# AITER_BENCH_TOKENS anyway and you get it, along with that failure. +_MEGA_MOE_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024, 2048)) +# What dispatch puts on the wire; combine is always bf16, so anything but bf16 +# is an asymmetric pair. fp4 is the wire DSv4 actually serves on -- the receiver +# hands the payload straight to the expert GEMM as its A operand, and that GEMM +# is a4w4 (ATOM's serve script pins MEGA_WIRE=fp4, AITER_FORCE_A8W4=0) -- so +# measuring only bf16 measures a leg the model does not run. $DISP overrides, +# comma-separated, and is passed through unvalidated: mori owns the value set. +_MORI_EP_DISP = tuple( + d.strip() for d in os.environ.get("DISP", "bf16,fp4").split(",") if d.strip() +) +# bench_ep.py forces its own correctness check off on fp4 ("fp4 combine is too +# lossy to compare"), so a passing fp4 row is unchecked, not verified. Labelled +# in the table rather than left for the reader to know that from mori's source. +_MORI_EP_UNCHECKED = ("fp4",) + + +def _int_quad(s): + """Parse 'a,b,c,d' -> (int, int, int, int) — MLA v4 kargpreld shape tuples.""" + a, b, c, d = s.split(",") + return int(a), int(b), int(c), int(d) + + +def _tflops(flop, us): + """TFLOPS from a FLOP count and microseconds (None-safe).""" + return round(flop / us / 1e6, 2) if us else None + + +def _bw(nbytes, us): + """Bandwidth (TB/s) from a byte count and microseconds (None-safe). + bytes / (us*1e-6) / 1e12 == bytes / us / 1e6.""" + return round(nbytes / us / 1e6, 3) if us else None + + +# bytes-per-VALUE for the MoE quant formats (dims below are logical value counts, +# so fp4 must be 0.5 B/value, not the 1 B/element of the packed fp4x2 dtype). +# a4w4 : fp4 act (0.5) x fp4 weight (0.5) +# a8w4 : fp8 act (1.0) x fp4 weight (0.5) (mxfp8 x mxfp4) +# The bf16 stage output is 2 B/value. (act_bpe, weight_bpe) per data_format. +_MOE_BPE = {"a4w4": (0.5, 0.5), "a8w4": (1.0, 0.5)} +_OUT_BPE = 2 # bf16 stage outputs + + +def _moe_stage_flops(token, topk, model_dim, inter_dim, use_g1u1=True): + """Per-stage FLOP counts for the fused 2-stage MoE (matches gemm_moe_tune.py): + stage1 GEMM: [token, model_dim] x [E, n, model_dim] -> token*n*model_dim*topk*2 + n = inter_dim*2 (g1u1 gate+up) or inter_dim + stage2 GEMM: [token, topk, inter_dim] x [E, model_dim, inter_dim] + -> topk*token*model_dim*inter_dim*2 + Returns (flop1, flop2).""" + n = inter_dim * 2 if use_g1u1 else inter_dim + flop1 = token * n * model_dim * topk * 2 + flop2 = topk * token * model_dim * inter_dim * 2 + return flop1, flop2 + + +# per_1x32 microscale: every 32 quantized values share one e8m0 (1B) scale, so +# each quantized value carries an extra 1/32 B of scale traffic, on top of its +# own bpe. Applies to BOTH activations and weights (fp4 => bpe 0.5 => 17/16; +# fp8 => bpe 1.0 => 33/32). Output stays bf16 and is not microscaled. +# (gemm_moe_tune.py's stage1/stage2 omit scale entirely; we include it.) +_SCALE_PER_VALUE = 1 / 32 + + +def _moe_stage_bytes( + token, topk, model_dim, inter_dim, experts, aq_bpe, wq_bpe, use_g1u1=True +): + """Per-stage MoE traffic (bytes), including per_1x32 e8m0 scale on every + quantized operand (act + weight). The stage1 output / stage2 input is the + expanded [token*topk, n] / [token*topk, inter] intermediate, so both carry + topk; the stage1 input act is read once per token (reused across its topk + experts): + stage1: act[token,model_dim]@aq + out[token,topk,n]@bf16 + w1[E,n,model_dim]@wq + stage2: act[token,topk,inter_dim]@aq + out[token,model_dim]@bf16 + + w2[E,model_dim,inter_dim]@wq + n = inter_dim*2 (g1u1) or inter_dim. + Returns (bytes1, bytes2).""" + n = inter_dim * 2 if use_g1u1 else inter_dim + bo = _OUT_BPE + aq = aq_bpe + _SCALE_PER_VALUE # quantized act: data + e8m0 scale per value + wq = wq_bpe + _SCALE_PER_VALUE # quantized weight: data + e8m0 scale per value + bytes1 = ( + token * model_dim * aq + token * topk * n * bo + experts * n * model_dim * wq + ) + bytes2 = ( + token * topk * inter_dim * aq + + token * model_dim * bo + + experts * model_dim * inter_dim * wq + ) + return bytes1, bytes2 + + +# Per-op column whitelists: keep shape identifiers + perf, drop the constant +# config/correctness columns @benchmark echoes (gfx/dtype/err/cos_diff/...). +_MHA_KEEP = [ + "dtype", + "head_dim", + "hq", + "hk", + "sq", + "sk", + "batch", + "is_causal", + "init", + "asm us", + "asm TFLOPS", + "asm TB/s", +] +# Curated (head_dim, seqlen, is_causal) grid — hq=64, hk=8(d64)/4(d128), batch=1. +_MHA_SHAPES = [ + (head_dim, tokens, causal) + for head_dim in (64, 128) + for tokens in _TOKENS + for causal in (True, False) +] +_MOE_KEEP = [ + "data_format", + "act", + "token", + "model_dim", + "inter_dim", + "E", + "topk", + "pass", + "gemm1_us", + "gemm1 TFLOPS", + "gemm1 TB/s", + "gemm2_us", + "gemm2 TFLOPS", + "gemm2 TB/s", + "total us", + "total TFLOPS", + "total TB/s", + "kernel", +] +# Fixed kernel-bench config (mirrors test_flydsl_grouped_gemm_gfx1250.py --scenario kernel). +_MOE_DATA_FORMATS = ["a4w4", "a8w4"] +_MOE_CONFIG = { + "experts": 96, + "tokens": _TOKENS, + "topk": 6, + "model_dim": 7168, + "inter_dim": 3072, + "activation": "silu", # ActivationType.Silu + "use_bias": False, +} +_GEMM_KEEP = [ + "workload", + "intype", + "M", + "N", + "K", + "apre", + "outtype", + "data_init", + "scale_init", + "knl_name", + "asm us", + "asm TFLOPS", + "asm TB/s", + "asm err", + "asm result", +] +# gemm_a4w4 throughput square. +_GEMM_A4W4_SHAPES = [(tokens, 16384, 16384) for tokens in _TOKENS] + +_F8GEMM_PERF_SHAPES = { + "a8w8": [(tokens, 16384, 8192) for tokens in _TOKENS] + + [(tokens, 1048576, 16384) for tokens in _TOKENS], + "a8w4": [(tokens, 16384, 16384) for tokens in _TOKENS] + + [(tokens, 1048576, 16384) for tokens in _TOKENS], +} +# Curated (gqa_ratio, batch, kv_seq_lens, num_kv_splits) grid for MLA v4 nm +# kernarg-preload perf (mirrors op_tests/test_mla_v4_kargpreld.py sweep subset). +_MLA_V4_KARGPRELD_SHAPES = [ + (64, 64, 256, 1), + (64, 64, 256, 2), + (64, 64, 256, 4), + (64, 64, 512, 1), + (64, 64, 512, 2), + (64, 64, 512, 4), + (64, 64, 1024, 1), + (64, 64, 1024, 2), + (64, 64, 1024, 4), + (128, 64, 256, 1), + (128, 64, 256, 2), + (128, 64, 256, 4), + (128, 64, 512, 1), + (128, 64, 512, 2), + (128, 64, 512, 4), + (128, 64, 1024, 1), + (128, 64, 1024, 2), + (128, 64, 1024, 4), +] + [ + (gqa, tokens, kv_seq_lens, num_kv_splits) + for gqa in (64, 128) + for tokens in _MLA_DECODE_TOKENS + if tokens != 64 + for kv_seq_lens in (256, 512, 1024) + for num_kv_splits in (1, 2, 4) +] +_MLA_V4_DSV4_SHAPES = [ + (128, 512, kv_seq_lens, num_kv_splits) + for kv_seq_lens in (256, 512, 1024) + for num_kv_splits in (1, 2, 4) +] + [ + (128, tokens, kv_seq_lens, num_kv_splits) + for tokens in _MLA_DECODE_TOKENS + if tokens != 512 + for kv_seq_lens in (256, 512, 1024) + for num_kv_splits in (1, 2, 4) +] +_MLA_V4_COMPARE_KEEP = [ + "dtype", + "gqa_ratio", + "batch", + "kv_seq_lens", + "num_kv_splits", + "asm_s1", + "triton_s1", + "s1 triton/asm", + "asm_s2", + "triton_s2", + "s2 triton/asm", + "asm_tot", + "triton_tot", + "tot triton/asm", +] + + +@contextlib.contextmanager +def _capture(): + """Like _silence, but hand the block's fd-level output back to the caller. + + Yields a one-element list that holds the captured text once the block ends. + Backed by a temp file rather than a pipe: an op that emits more than the + pipe buffer (64K) would otherwise deadlock with nobody draining it. + """ + sys.stdout.flush() + sys.stderr.flush() + old1, old2 = os.dup(1), os.dup(2) + box = [] + with tempfile.TemporaryFile(mode="w+") as tmp: + try: + os.dup2(tmp.fileno(), 1) + os.dup2(tmp.fileno(), 2) + yield box + finally: + sys.stdout.flush() + sys.stderr.flush() + os.dup2(old1, 1) + os.dup2(old2, 2) + os.close(old1) + os.close(old2) + tmp.seek(0) + box.append(tmp.read()) + + +def _print_table(name, rows, keep=None): + df = pd.DataFrame([r for r in rows if r is not None]) + if not df.empty: + # Drop columns that are entirely empty, then whitelist/order via `keep`. + # The @benchmark decorator dumps every call arg as a column, which makes + # the tables wide; `keep` trims to shape ids + perf. ALWAYS surface any + # err_msg / *err column so failures never get silently hidden. + df = df.replace("", pd.NA).dropna(axis=1, how="all") + if keep is not None: + cols = [c for c in keep if c in df.columns] + cols += [c for c in df.columns if "err_msg" in c and c not in cols] + df = df[cols] + print(f"\n===== {name} =====") + print(df.to_markdown(index=False)) + + +# Compiler / logger / IR-dump chatter the child UTs interleave with results. +_NOISE = ( + "[flydsl.compile]", + "[aiter INFO]", + "[aiter WARNING]", + "import [module_", + "In file included from", + "torch/distributed/run.py", + "Building extension", + "Emitting ninja", + "hipcc", + "warning:", + "UserWarning", + "_warn_once", +) + + +def _md_row(line): + """Markdown table row emitted by a child UT.""" + return line.startswith("|") + + +def _quiet(line): + """Any non-empty line that is not compiler/logger noise.""" + return bool(line.strip()) and not any(n in line for n in _NOISE) + + +def _lines(pred): + """Adapt a per-line predicate into a block extractor.""" + return lambda lines: [ln for ln in lines if pred(ln)] + + +def _md_tables(*labels): + """Keep the markdown tables, labelling each by the columns it carries. + + A child UT often emits several tables in a row with different columns and + nothing saying which is which. `labels` is ((column, ...), title) pairs; the + first entry whose columns all appear in a header row names that table. + """ + + def extract(lines): + md = [ln for ln in lines if _md_row(ln)] + out = [] + for i, line in enumerate(md): + is_header = i + 1 < len(md) and set(md[i + 1]) <= set("|-: ") + if is_header: + title = next( + (t for cols, t in labels if all(c in line for c in cols)), + None, + ) + if title: + out.append(f"\n----- {title} -----") + elif out: + out.append("") + out.append(line) + return out + + return extract + + +def _isnum(field): + """Does this field parse as a number (thousands separators allowed)?""" + try: + float(field.replace(",", "")) + except ValueError: + return False + return True + + +def _md_kernel_table(lines): + """Render a rank-major kernel table as markdown, keeping the summary lines. + + mega_moe prints '[cfg] ...' / '# MEGA-MOE ...' lines around a space-aligned + 'Name rank0 rank1 rank2 rank3 avg calls' table. Kernel names contain spaces + ("void at::native::reduce_kernel<512, 1, ...>"), so split from the right: + the column count is fixed even when the name is not. + """ + out, rows, cols = [], [], None + seen = set() + + def flush(): + if cols and rows: + out.append(pd.DataFrame(rows, columns=cols).to_markdown(index=False)) + rows.clear() + + for line in lines: + if not _quiet(line): + continue + if line.startswith("Name") and "rank0" in line: + flush() + cols = line.rsplit(maxsplit=6) + continue + fields = line.rsplit(maxsplit=6) + if cols and len(fields) == 7 and all(_isnum(f) for f in fields[1:]): + rows.append(fields) + continue + flush() + # Non-table lines are kept for the "[cfg] ..." summary, but the child + # also emits "no grouped CSV config matched (...)" once per layer per + # rank -- hundreds of byte-identical lines around one table. Keep the + # first of each; a repeat carries nothing the first did not. + if line in seen: + continue + seen.add(line) + out.append(line) + flush() + return "\n".join(out).splitlines() + + +def _md_from_pandas(marker, columns): + """Re-emit a pandas-printed block as markdown. + + Some UTs print their result with DataFrame.__str__ (space aligned, leading + index column) right after a marker line, which reads nothing like the + markdown every other op produces. + """ + + def extract(lines): + for i, line in enumerate(lines): + if line.strip() != marker or i + 2 >= len(lines): + continue + values = lines[i + 2].split()[1 : len(columns) + 1] + if len(values) != len(columns): + continue + df = pd.DataFrame([values], columns=list(columns)) + return df.to_markdown(index=False).splitlines() + return [] + + return extract + + +def _space_table(header_col): + """Keep a UT's own aligned summary table, verbatim. + + Anchors on the header row carrying `header_col` and takes the data rows that + follow, so the table survives the trace fragments and compiler warnings + interleaved before it. + + Emitted as the UT formatted it rather than rebuilt as markdown: pandas + writes multi-word column names ("opus us", "asm TFLOPS"), so the header + splits into 33 words against 21 data fields and cannot be mapped back to + columns. A column name also appears in a UT's argument echo + ("total_tokens = 1024,"), so require the next line to look like data. + + "Looks like data" counts numeric fields rather than testing the first one: + the sparse-prefill table leads with prec/mode (bf16, dense), so a + first-field test drops the whole table. An argument echo carries one + number, a data row carries most of a row of them. + """ + + def is_data(fields): + return sum(1 for f in fields if _isnum(f)) >= 4 + + def extract(lines): + for i, line in enumerate(lines): + if header_col not in line.split() or i + 1 >= len(lines): + continue + first = lines[i + 1].split() + if not first or not is_data(first): + continue + width = len(first) + out = [line] + for follower in lines[i + 1 :]: + fields = follower.split() + if len(fields) != width or not is_data(fields): + break + out.append(follower) + return out + return [] + + return extract + + +def _table_row(*headers): + """Whitespace-aligned table: the header line plus its numeric rows.""" + + def keep(line): + if not _quiet(line): + return False + if any(h in line for h in headers): + return True + # A data row starts with a bare number; "100% |####|" (pip) does not. + head = line.split(maxsplit=1)[0] + return head.strip("-").replace(",", "").replace(".", "").isdigit() + + return keep + + +def _gpu_trace_rows(lines): + """(kernel, calls, device_us) for every GPU row in a profiler fragment. + + Rows look like ' CUDA '; + the kernel name carries spaces, the trailing column count does not. host_us + is 0 on GPU rows, so the time to report is device_us. + """ + for line in lines: + fields = line.rsplit(maxsplit=6) + if len(fields) != 7 or fields[-2] != "CUDA": + continue + _, cnt, _host, device_us, _avg, _, _ = fields + if not (_isnum(cnt) and _isnum(device_us)): + continue + head = fields[0].split(None, 1) + name = head[1].strip() if len(head) == 2 and head[0].isdigit() else fields[0] + if name: + yield name, float(cnt.replace(",", "")), float(device_us.replace(",", "")) + + +def _kernel_names(lines): + """Distinct GPU kernel names in the order they first appear.""" + names = [] + for name, _, _ in _gpu_trace_rows(lines): + if name not in names: + names.append(name) + return names + + +def _kernel_digest(lines): + """Which GPU kernels actually ran, from the trace fragments in the output. + + A table of microseconds does not say which code path produced them, so a + silent fallback (or a shape that quietly picked another kernel) reads as a + normal result. The profiler fragments name every kernel that reached the + GPU -- roll them up so each op states what it actually ran. + """ + total, calls = {}, {} + for name, n, us in _gpu_trace_rows(lines): + total[name] = total.get(name, 0.0) + us + calls[name] = calls.get(name, 0.0) + n + if not total: + return [] + ranked = sorted(total, key=total.get, reverse=True) + table = pd.DataFrame( + [ + {"kernel": k, "calls": round(calls[k]), "device us": round(total[k], 1)} + for k in ranked + ] + ) + return ["", "----- kernels on GPU -----"] + table.to_markdown( + index=False + ).splitlines() + + +_DEFAULT_EXTRACT = _md_tables() + + +_FAILURES = [] + + +def _note_failure(label, why): + """Record a dead case and let the sweep continue past it.""" + _FAILURES.append((label, why)) + print(f"--- {label}: FAILED ({why}), continuing ---", flush=True) + + +@contextlib.contextmanager +def _keep_going(label): + """Op-level net for whatever _run_child cannot catch. + + A child UT that aborts is already handled inside _run_child; this catches + the in-process ops raising Python exceptions. A GPU fault in this process + is not recoverable -- it takes the interpreter with it, and no handler runs. + """ + try: + yield + except Exception as exc: # noqa: BLE001 - a sweep must outlive one bad op + _note_failure(label, f"{type(exc).__name__}: {exc}") + finally: + # A half-finished op can leave allocations behind; the next one should + # not inherit them. + try: + torch.cuda.empty_cache() + except Exception: # noqa: BLE001 - cleanup must not mask the failure + pass + + +def _pin_arch(env): + """Hand the child the arch we already know, so it never runs rocminfo. + + chip_info shells out to rocminfo twice -- once for the arch, once for the + CU count -- and rocminfo takes a per-device rocm_smi mutex on the way in. + One process is fine. A torchrun op starts four ranks at once, and they + contend for that mutex: on 20260901/b45-1 a rank lost it and aborted + ("init_mutex /rocm_smi_renderD128: unlock timed lock", surfacing as + "Allgather operation failed" once the dead rank took the collective with + it), and on b45-2 four rocminfo processes sat in it for minutes, one of + them wedged in D state. Both cost a whole op; the nine single-GPU ops + never noticed, because one process has nobody to contend with. + + GPU_ARCHS covers get_gfx_list, CU_NUM covers get_cu_num -- both are read + from the environment before either shells out. Detected once here, in this + process, where the call is serial. Not forced: an explicit setting from + the caller wins. + """ + env.setdefault("GPU_ARCHS", get_gfx()) + env.setdefault("CU_NUM", str(get_cu_num())) + return env + + +def _run_child(name, cmd, cwd, env=None, extract=None, timeout=None, tail=30, + kernels=True): + """Run a child UT with its output captured and surface only its results. + + Child UTs print their own progress, aiter INFO lines and (with FlyDSL) a + couple of thousand IR-dump lines, which buries the numbers. Capture all of + it, echo what `extract` pulls out, and fall back to the tail of the output + when the child fails or emits nothing recognisable. + """ + extract = extract or _DEFAULT_EXTRACT + # env=None means "inherit ours", which already carries these two. + if env is not None: + _pin_arch(env) + try: + proc = subprocess.run( + cmd, cwd=cwd, env=env, text=True, timeout=timeout, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + ) + except subprocess.TimeoutExpired as exc: + captured = exc.output or "" + print(f"\n===== {name} =====", flush=True) + print(f"--- timed out after {timeout}s, last {tail} lines ---") + print("\n".join(captured.splitlines()[-tail:]), flush=True) + _note_failure(name, f"timed out after {timeout}s") + return + lines = proc.stdout.splitlines() + # `results` decides whether the op reported anything; the kernel digest is + # an annotation and must not stand in for a result table, or an extractor + # that stops matching turns into a silent hole instead of a failure. + results = extract(lines) + rows = list(results) + (_kernel_digest(lines) if kernels else []) + print(f"\n===== {name} =====", flush=True) + print("\n".join(rows) if rows else "(no result rows recognised)", flush=True) + if proc.returncode != 0 or not results: + print(f"--- {name}: exit={proc.returncode}, last {tail} lines ---") + print("\n".join(lines[-tail:]), flush=True) + # Recorded, not raised: one dead shape used to take the rest of the sweep + # with it -- a mega_moe case aborting at tokens/rank=65536 meant the seven + # ops queued behind it never ran at all. + if proc.returncode != 0: + _note_failure(name, f"child exited {proc.returncode}") + elif not results: + _note_failure(name, "no result rows") + + +# --- per-op runners: sweep axes silently, then print one table --- + + +def run_mha(args): + # perf-only fn (no torch ref): sq==sk, hq=64, hk=8(d64)/4(d128), batch=1. + rows = [] + with _silence(): + for init in args.mha_init: + for head_dim, seqlen, causal in _MHA_SHAPES: + hk = 8 if head_dim == 64 else 4 + rows.append( + mha_mod.test_fmha_fwd_with_sink_asm_perf( + head_dim, 64, hk, seqlen, seqlen, 1, causal, init + ) + ) + for row in rows: + if row is not None: + row["dtype"] = "bf16" + _print_table("mha (bf16)", rows, keep=_MHA_KEEP) + + +def run_moe(args): + cfg = _MOE_CONFIG + activation = moe_mod.ActivationType.Silu + rows = [] + data_formats = ["a8w4"] if args.suite == "dsv4" else _MOE_DATA_FORMATS + for tokens, fmt in itertools.product(cfg["tokens"], data_formats): + with _capture() as box: + moe_mod.set_data_format(fmt) + metrics = moe_mod.run_moe( + fmt, + experts=cfg["experts"], + tokens=tokens, + topk=cfg["topk"], + model_dim=cfg["model_dim"], + inter_dim=cfg["inter_dim"], + activation=activation, + use_bias=cfg["use_bias"], + kernel_bench=True, + check_aot_cache=False, + raise_on_fail=False, + ) + # stage1 n = inter_dim*2 (gate+up for silu/swiglu GUGU layout). + aq_bpe, wq_bpe = _MOE_BPE.get(fmt, (1, 1)) + flop1, flop2 = _moe_stage_flops( + tokens, + cfg["topk"], + cfg["model_dim"], + cfg["inter_dim"], + use_g1u1=True, + ) + bytes1, bytes2 = _moe_stage_bytes( + tokens, + cfg["topk"], + cfg["model_dim"], + cfg["inter_dim"], + cfg["experts"], + aq_bpe, + wq_bpe, + use_g1u1=True, + ) + us1, us2 = metrics.get("gemm1_us"), metrics.get("gemm2_us") + total_us = (us1 or 0) + (us2 or 0) if (us1 or us2) else None + bw1, bw2, bwt = ( + _bw(bytes1, us1), + _bw(bytes2, us2), + _bw(bytes1 + bytes2, total_us), + ) + rows.append( + { + "data_format": fmt, + "act": cfg["activation"], + "token": tokens, + "model_dim": cfg["model_dim"], + "inter_dim": cfg["inter_dim"], + "E": cfg["experts"], + "topk": cfg["topk"], + "pass": metrics["passed"], + "gemm1_us": us1, + "gemm1 TFLOPS": _tflops(flop1, us1), + "gemm1 TB/s": bw1, + "gemm2_us": us2, + "gemm2 TFLOPS": _tflops(flop2, us2), + "gemm2 TB/s": bw2, + "total us": round(total_us, 2) if total_us else None, + "total TFLOPS": _tflops(flop1 + flop2, total_us), + "total TB/s": bwt, + "kernel": " + ".join(_kernel_names(box[0].splitlines())) or None, + } + ) + _print_table("flydsl_grouped_gemm (kernel, silu)", rows, keep=_MOE_KEEP) + + +def run_gemm(args): + # Hardware throughput sweep only. Functional/UT mode belongs in the source + # op test and is intentionally not exposed by this performance driver. + init_pairs = [("constant", "constant"), ("uniform", "auto")] + rows = [] + with _silence(): + for (M, N, K), (di, si), intype, outtype in itertools.product( + _GEMM_A4W4_SHAPES, + init_pairs, + ["mxfp4", "nvfp4"], + ["bf16", "fp8"], + ): + rows.append( + gemm_mod.test_gemm( + intype, + M, + N, + K, + 1, + outtype, + di, + si, + mode="perf", + ) + ) + _print_table("gemm_a4w4 (perf)", rows, keep=_GEMM_KEEP) + + +def run_f8gemm(args): + # Generic MXFP8 hardware sweep. DSv4's projection path uses the separate + # a8w8_blockscale runner below, not this F8GEMM kernel family. + rows = [] + with _silence(): + cases = [ + ("hardware", intype, M, N, K, di, si) + for (di, si), intype in itertools.product( + [("constant", "constant"), ("uniform", "auto")], + ["a8w8", "a8w4"], + ) + for M, N, K in _F8GEMM_PERF_SHAPES[intype] + ] + for workload, intype, M, N, K, di, si in cases: + row = f8gemm_mod.test_gemm( + intype, + M, + N, + K, + 1, + data_init=di, + scale_init=si, + mode="perf", + ) + if row is not None: + row["workload"] = workload + rows.append(row) + _print_table(f"mxfp8fp4gemm ({args.suite})", rows, keep=_GEMM_KEEP) + + +def run_a8w8_blockscale(_args): + """Run DSv4 FP8 blockscale linear projections at M=512.""" + # AITER_LOG_MORE=1 is set at module scope for the FlyDSL MoE ops, and a + # child started with env=None inherits this process's whole environ. In this + # UT that turned a clean sweep into an intermittent HSA memory fault, so + # drop it for this child only -- every other op keeps it. + env = os.environ.copy() + env.pop("AITER_LOG_MORE", None) + _run_child( + "gemm_a8w8_blockscale (DSv4)", + [ + sys.executable, + "op_tests/test_gemm_a8w8_blockscale.py", + *( + ["-m", *map(str, _A8W8_BLOCKSCALE_TOKENS)] + if _A8W8_BLOCKSCALE_TOKENS + else [] + ), + "-nk", + "2048,7168", + "7168,16384", + "6144,7168", + "7168,3072", + "65536,1536", + "8192,1536", + "--ck_preshuffle", + "True", + "--flydsl", + ], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + env=env, + ) + + +def run_a16w16(_args): + """Run the DSv4 BF16 linear shapes through the Opus GEMM UT.""" + # test_a16w16 returns only the error; its timing is printed as + # [a16w16] batch=1 M=512 N=64 K=7168 dtype=... | 7.8us | 12.05 TFLOPs | err=0 + # so capture the block and parse that line back out. + batch, K = 1, 7168 + rows = [] + for M, n in itertools.product(_A16W16_MS, _A16W16_NS): + # N=32320/129280 is lm_head (the DeepSeek vocab, whole and TP4-sharded). + # See _A16W16_WIDE_N: this is the one shape rule left here, and it is + # about what DSv4 runs, not about what the kernel can do. + if n > _A16W16_WIDE_N and M > _A16W16_WIDE_N_MAX_M: + rows.append({"batch": batch, "M": M, "N": n, "K": K, + "err_msg": f"skipped: N>{_A16W16_WIDE_N} is lm_head, " + f"capped at M<={_A16W16_WIDE_N_MAX_M}"}) + continue + # No >4 GiB pre-check. opus_dispatch_a16w16_gfx1250 tries the tuned + # table FIRST and returns on a hit; check_shape_4g runs only after that + # misses (opus_gemm_arch_gfx1250.cuh:161), on the way to the split-K + # heuristic kid -- whose launcher is what builds the 32-bit gmem + # descriptors. A tuned 4wave_wl_co winner never reaches it: that + # pipeline addresses gmem through TDM descriptors, which clamp every + # dimension and are not 32-bit bounded. So the limit belongs to one + # fallback path, not to a16w16, and predicting it here would keep + # skipping shapes that tuning has already made runnable. Let the kernel + # raise and record that instead. + try: + with _capture() as box: + err = a16w16_mod.test_a16w16(batch=batch, M=M, N=n, K=K) + except Exception as exc: # noqa: BLE001 - one shape must not end the sweep + rows.append({"batch": batch, "M": M, "N": n, "K": K, + "err_msg": f"{type(exc).__name__}: {exc}"}) + continue + captured = box[0].splitlines() + row = {"batch": batch, "M": M, "N": n, "K": K, "err": err} + # float(): checkAllclose returns a bare 0 for a clean compare but a + # numpy/torch scalar for a mismatch, and only one of those formats. + if err is not None and float(err) > _A16W16_MAX_ERR: + row["err_msg"] = (f"WRONG RESULT: err={float(err):g} " + f"> {_A16W16_MAX_ERR:g}") + _note_failure(f"a16w16 M={M} N={n} K={K}", row["err_msg"]) + for line in captured: + if not line.startswith("[a16w16]"): + continue + fields = [f.strip() for f in line.split("|")] + us = next((f for f in fields if f.endswith("us")), None) + tflops = next((f for f in fields if f.endswith("TFLOPs")), None) + row["us"] = float(us[:-2]) if us else None + row["TFLOPS"] = float(tflops[:-7]) if tflops else None + break + # Which kernel served this shape: a16w16 switches between a splitk pair + # and a 4wave_wl_co variant, and the timing alone does not say which. + row["kernel"] = " + ".join(_kernel_names(captured)) or None + rows.append(row) + _print_table( + "gemm_a16w16_opus (DSv4)", + rows, + keep=["batch", "M", "N", "K", "us", "TFLOPS", "kernel", "err"], + ) + + +def run_mega_moe(_args): + """Run the four-rank DSv4 Mega MoE path vs its base combine, a4w4 and a8w4.""" + # The child ranks need GPU 0 as well. Release any cached allocations held by + # this orchestration process before torchrun starts the four workers. + torch.cuda.empty_cache() + env = os.environ.copy() + # No MORI_SHMEM_HEAP_SIZE default here, for two independent reasons. + # + # Raising it sweep-wide took the machine down: the heap is preallocated per + # rank for every case, not sized per case, so 16 GB became 64 GB reserved on + # every one of them and b45-2 hard-rebooted at tokens/rank=512 -- long + # before the case it was meant to help. + # + # And it would not have helped anyway. The 7.5 GB request at 65536 goes to + # cco's VMM arena, not the shmem heap: "ccoMemAlloc: slot exhausted ... in + # perRankSize=4294967296. Increase perRankVmmSize at ccoCommCreate". That + # size is a ccoCommCreate argument with no environment variable behind it + # (Communicator.DEFAULT_PER_RANK_VMM, 4 GiB), and + # test_mega_moe_gfx1250.py:512 calls Communicator.init() without passing it. + # MORI_SHMEM_HEAP_SIZE is read only in mori/src/shmem/init.cpp and feeds a + # different allocator. The same error also prints "Hint: Increase via + # MORI_SHMEM_HEAP_SIZE" -- that hint is what points the wrong way. + env.update({"MORI_V2_KERNEL_BACKEND": "hip", "MEGA_DISPATCH": "mori"}) + base_cmd = [ + "torchrun", + "--standalone", + "--nproc_per_node=4", + "op_tests/multigpu_tests/test_mega_moe_gfx1250.py", + "-e", + "384", + "-k", + "6", + "-hd", + "7168", + "-id", + "3072", + "--layers", + "61", + "--acc_verify", + "0", + "--profile_table", + "1", + ] + # AITER_FORCE_A8W4 selects the grouped kernel's ACTIVATION dtype (0 -> fp4, + # 1 -> fp8); the weights are mxfp4 either way and -q only picks their layout, + # so the env var and the quant key have to move together. + for tokens, (quant, force_a8w4), (label, combine) in itertools.product( + _MEGA_MOE_TOKENS, + (("a4w4_mxfp4", "0"), ("a8w4_mxfp4", "1")), + (("non-Mega", "base"), ("Mega", "fused")), + ): + _run_child( + f"mega_moe (tokens/rank={tokens}, {quant}, {label}, combine={combine})", + [*base_cmd, "-tpr", str(tokens), "-q", quant, "--combine", combine], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + env={**env, "AITER_FORCE_A8W4": force_a8w4}, + extract=_md_kernel_table, + kernels=False, + ) + + +def run_mhc(_args): + """Run the DSv4 mHC fused-RMSNorm benchmark at M=512, N=7168.""" + _run_child( + "mhc (DSv4, fused RMSNorm)", + [ + sys.executable, + "op_tests/test_mhc.py", + "-n", + "7168", + "-m", + *map(str, _TOKENS), + "--fuse_rmsnorm", + ], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + extract=_md_tables( + (("hip_nofuse_us",), "mhc: fused vs unfused RMSNorm"), + (("unfused_us",), "mhc_post_pre"), + (("hip_us",), "mhc_head"), + ), + ) + + +def run_qk_norm(_args): + """Run DSv4 QK norm + RoPE for prefill and decode token counts.""" + base_cmd = [ + sys.executable, + "op_tests/test_flydsl_qk_norm_rope_quant.py", + "--H", + "128", + "--D", + "512", + "--RD", + "64", + "--no-quant", + "--qweight", + ] + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + _run_child( + "qk_norm", + [*base_cmd, "-T", *map(str, _TOKENS)], + cwd=repo_root, + extract=_md_tables( + (("quant_group_size",), "rope + quant"), + (("rows_written",), "fused SWA write"), + ), + ) + + +def run_score_qk(_args): + """Run DSv4 decode score-QK at batch 512 for short and long CSA KV.""" + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + base_cmd = [ + sys.executable, + "op_tests/op_benchmarks/triton/bench_deepgemm_attention.py", + "--heads", + "64", + "--index_dim", + "128", + "-mtp", + "0", + "--kv_preshuffle", + "--blocksize", + "64", + ] + # None => let the UT pick the batch, so run the KV lengths once each. + for tokens, (label, kv_length) in itertools.product( + _SCORE_QK_TOKENS or (None,), _SCORE_QK_KV_LENGTHS + ): + _run_child( + f"score_qk (decode, B={tokens or 'UT default'}, {label} CSA KV={kv_length})", + [ + *base_cmd, + *(["--batch", str(tokens)] if tokens else []), + "-kv_length", + kv_length, + ], + cwd=repo_root, + extract=_md_from_pandas( + "paged_mqa_logits:", + ("batch", "next_n", "heads", "index_dim", "avg_kv_len", "TFLOPS"), + ), + ) + + +def run_mori_ep(_args): + """Run MORI EPv2 dispatch/combine at the DSv4 MoE shape.""" + # Runs whatever mori the image provides; keeping it current is the image's + # job. Updating it from here moved the measurement target between runs and + # needed a dev ROCm toolchain the pip-wheel images do not ship. + mori = os.environ.get("MORI", "/app/mori") + env = os.environ.copy() + env["PYTHONPATH"] = f"{mori}/python:{mori}" + env["MORI_SOCKET_IFNAME"] = "lo" + env["GLOO_SOCKET_IFNAME"] = "lo" + env["PYTHONUNBUFFERED"] = "1" + backend = env.get("BACKEND", "hip") + env.update( + { + "BACKENDS": backend, + "MORI_V2_KERNEL_BACKEND": backend, + "HIDDEN": env.get("HIDDEN", "7168"), + "TOPK": env.get("TOPK", "6"), + "EPR": env.get("EPR", "96"), + "SWEEP": env.get( + "TOKENS", "64,128,256,512,1024,2048,4096,8192,16384" + ), + "ITERS": env.get("ITERS", "200"), + "WARMUP": "10", + "MODES": env.get("MODES", "eager,graph"), + "COMBINE_IN": env.get("COMBINE_IN", "inplace"), + "CHECK": env.get("CHECK", "1"), + "DBN": "", + "DWPB": "", + "CBN": "", + "CWPB": "", + } + ) + # One child per wire: bench_ep.py reads $DISP once at import and builds the + # transport for that dtype, so the tiers cannot share a process. + for disp in _MORI_EP_DISP: + env["DISP"] = disp + note = " UNCHECKED" if disp in _MORI_EP_UNCHECKED else "" + _run_child( + f"mori_ep (DSv4 dispatch/combine, disp={disp}, combine=bf16{note})", + [ + "torchrun", + "--standalone", + f"--nproc_per_node={env.get('EP', '4')}", + "tests/python/ops/dispatch_combine_v2/bench_ep.py", + ], + cwd=mori, + env=env, + extract=_lines(_quiet), + timeout=3600, + ) + + +def _perf_ratio(num, den): + """triton/asm speed ratio as '1.03x'; 'nanx' when undefined.""" + if num is None or den is None or den == 0: + return "nanx" + return f"{num / den:.2f}x" + + +def _bench_mla_v4_asm_staged(gqa, batch, ctx, split_kv, num_iters, num_warmup): + """Asm kernel (s1) + merge (s2) + total; lives in combo bench only.""" + mod = mla_v4_kargpreld_mod + q_seq = 1 + assert (gqa, q_seq) in mod._SHIPPED_TILE_VARIANTS + if split_kv > 1: + min_split = ctx // split_kv + assert ( + min_split >= 16 + ), f"smallest KV split = floor({ctx}/{split_kv}) = {min_split} < 16" + + device = "cuda" + inputs = mod._build_bf16_inputs( + batch=batch, + kv_seq_lens=ctx, + q_seq_logical=q_seq, + seed=mod._SEED, + gqa_ratio=gqa, + attn_sink=True, + ) + sm_scale = 1.0 / (mod._QUANT_D**0.5) + q_packed, q_rope = mod._native_to_2buff_for_asm(inputs["q_bf16"]) + kv_packed, kv_rope = mod._native_to_2buff_for_asm(inputs["kv_bf16"]) + + total_q = inputs["q_bf16"].size(0) + num_seqs = inputs["qo_indptr"].size(0) - 1 + num_heads = mod.NUM_KV_HEADS * gqa + output_buf = torch.empty( + (total_q, gqa, mod.V_HEAD_DIM), dtype=dtypes.bf16, device=device + ) + split_indptr = torch.tensor( + [i * split_kv for i in range(num_seqs + 1)], + dtype=torch.int32, + device=device, + ) + logits_buf = torch.empty( + (total_q, split_kv, num_heads, mod.V_HEAD_DIM), + dtype=torch.float32, + device=device, + ) + lse_buf = torch.empty( + (total_q, split_kv, num_heads, 1), dtype=torch.float32, device=device + ) + valid_split_count = torch.empty((num_seqs,), dtype=torch.int32, device=device) + + common_kwargs = { + "q": q_packed, + "qrope": q_rope.contiguous(), + "kv_buffer": kv_packed, + "kvrope": kv_rope.contiguous(), + "output": output_buf, + "qo_indptr": inputs["qo_indptr"], + "kv_indptr": inputs["kv_indptr"], + "kv_page_indices": inputs["kv_page_indices"], + "kv_last_page_lens": inputs["kv_last_page_lens"], + "split_indptr": split_indptr, + "max_seqlen_q": inputs["max_seqlen_q"], + "sink": inputs["sink"], + "sm_scale": sm_scale, + "num_kv_splits": split_kv, + "logits": logits_buf, + "attn_lse": lse_buf, + } + perf = {"num_iters": num_iters, "num_warmup": num_warmup, "num_rotate_args": 1} + + _, us_k = run_perftest( + aiter.mla_decode_v4_asm, + q_packed, + q_rope.contiguous(), + kv_packed, + kv_rope.contiguous(), + inputs["qo_indptr"], + inputs["kv_indptr"], + inputs["kv_page_indices"], + split_indptr, + inputs["sink"], + inputs["max_seqlen_q"], + sm_scale, + 0, + split_kv, + logits_buf, + lse_buf, + output_buf, + valid_split_count, + int(split_kv > 1), + inputs["kv_last_page_lens"], + **perf, + ) + _, us_tot = run_perftest( + aiter.mla.mla_decode_fwd_v4_nm, + out_16_nosplit=0, + **common_kwargs, + **perf, + ) + asm_s2 = max(0.0, us_tot - us_k) if split_kv > 1 else 0.0 + return { + "asm_s1": round(us_k, 2), + "asm_s2": round(asm_s2, 2), + "asm_tot": round(us_tot, 2), + } + + +def run_mla_v4_decode(args): + # Side-by-side asm (kargpreld) vs Triton sparse decode on the same shape grid. + iters = args.mla_v4_kargpreld_iters + warmup = args.mla_v4_kargpreld_warmup + mla_v4_triton_mod._PERF["num_iters"] = iters + mla_v4_triton_mod._PERF["num_warmup"] = warmup + default_shapes = ( + _MLA_V4_DSV4_SHAPES + if args.suite == "dsv4" + else _MLA_V4_KARGPRELD_SHAPES + ) + shapes = args.mla_v4_kargpreld_shapes or default_shapes + rows = [] + with _capture() as box: + for gqa, batch, ctx, split_kv in shapes: + row = { + "gqa_ratio": gqa, + "batch": batch, + "kv_seq_lens": ctx, + "num_kv_splits": split_kv, + } + try: + asm = _bench_mla_v4_asm_staged(gqa, batch, ctx, split_kv, iters, warmup) + tri = mla_v4_triton_mod.test_mla_v4_triton_staged( + gqa_ratio=gqa, + batch=batch, + kv_seq_lens=ctx, + num_kv_splits=split_kv, + ) + row.update(asm) + row.update(tri) + row["s1 triton/asm"] = _perf_ratio(row["triton_s1"], row["asm_s1"]) + row["s2 triton/asm"] = _perf_ratio(row["triton_s2"], row["asm_s2"]) + row["tot triton/asm"] = _perf_ratio(row["triton_tot"], row["asm_tot"]) + except (RuntimeError, AssertionError, ValueError) as exc: + msg = str(exc).splitlines()[0] if str(exc) else type(exc).__name__ + row["err_msg"] = msg + rows.append(row) + for row in rows: + row["dtype"] = "bf16" + _print_table( + "mla_v4 decode (bf16, asm vs triton)", + rows, + keep=_MLA_V4_COMPARE_KEEP, + ) + print("\n".join(_kernel_digest(box[0].splitlines())), flush=True) + + +def run_inverse_rope(_args): + """Run DSv4 inverse RoPE + group quant at the tp1 attention-output shape.""" + # -b is (n_local_heads, n_local_groups); 128,16 is V4-Pro at dp/tp1. The UT + # defaults to the two smallest configs instead, which never reach the shape + # the model runs, so name it explicitly. + _run_child( + "inverse_rope_group_quant (DSv4, tp1)", + [ + sys.executable, + "op_tests/test_inverse_rope_group_quant.py", + "-b", + "128,16", + *(["-s", *map(str, _INVERSE_ROPE_TOKENS)] if _INVERSE_ROPE_TOKENS else []), + "-l", + "n32k4", + "--group-size", + "32", + ], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + + +def run_mla_v4_prefill(_args): + """Run DSv4 prefill across two precisions, pools and CSR modes.""" + for tokens in _MLA_PREFILL_TOKENS: + _run_child( + f"mla_v4 prefill (M={tokens}, prec=fp8/bf16, pages=4096/16384)", + [ + sys.executable, + "op_tests/test_pa_sparse_prefill.py", + "-n", + str(tokens), + "--h_q", + "128", + "-d", + "512", + "--total_pages", + "4096", + "16384", + "--total_tokens", + str(tokens), + "--prec", + "fp8", + "bf16", + # bf16 takes the single-tensor Q/K/V/O kernel; only fp8 has an + # asm candidate, so the bf16 rows compare opus against triton + # and leave the asm columns empty. + "--mode", + "dense", + "sparse", + "--no-verify", + ], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + # Not _table_row: the UT has no "latency_us" column (it prints + # "opus us"/"asm us"), so that predicate fell through to its + # starts-with-a-number rule, swallowed the profiler's kernel lines + # as data, and dropped the real header, which starts with "n". + extract=_space_table("total_pages"), + ) + + +OPS = { + "mha": run_mha, + "moe": run_moe, + "gemm": run_gemm, + "f8gemm": run_f8gemm, + "a8w8_blockscale": run_a8w8_blockscale, + "a16w16": run_a16w16, + "mla_v4_decode": run_mla_v4_decode, + "inverse_rope": run_inverse_rope, + "mla_v4_prefill": run_mla_v4_prefill, + "mhc": run_mhc, + "qk_norm": run_qk_norm, + "score_qk": run_score_qk, + "mori_ep": run_mori_ep, + "mega_moe": run_mega_moe, +} +# "gemm" (f4gemm a4w4) stays out: two back-to-back cases differing only in +# outtype (bf16 -> fp8) abort with HSA_STATUS_ERROR_MEMORY_FAULT, while each +# passes in a fresh process -- state leaking across cases, not a kernel bug. +# +# f8gemm has been seen reporting "0 us" / "inf TFLOPS" on part of its sweep. +# run_perftest times through the torch profiler, so a row like that means the +# profiler recorded no GPU work, not that the kernel was fast -- read it +# alongside the kernel digest, which says whether anything reached the GPU. +PERF_OPS = ["mha", "moe", "f8gemm", "mla_v4_decode"] +DSV4_OPS = [ + "mega_moe", + # mori's own EPv2 bench, not an aiter kernel, but it is the dispatch and + # combine either MoE path pays for -- the sweep is incomplete without the + # two all2all legs beside the GEMMs. Reads the mori tree the image ships + # (MORI=/app/mori); keeping that tree current is the image's job. + "mori_ep", + "moe", + # "a8w8_blockscale" stays out of the default sweep, but NOT because the op is + # dead on gfx1250 -- an earlier note here said that and it was wrong. The op + # runs: 20260827, the six DSv4 (n,k) at -m 512 give err=0 at 1210-3564 + # TFLOPS, over ck / asm / flydsl. + # + # What kills it is one line of the UT. An earlier note here blamed #4773's + # gluon tuning rows; that was wrong, and re-tuning would not have helped. + # test_gemm_a8w8_blockscale.py:120 runs an extra "ck strided x_scale" check + # on x_scale.transpose(0,1).contiguous().transpose(0,1) -- the same bytes + # the measured call gets, but stride (1, M) instead of contiguous. #4406 + # added it to cover its own "honor strided x_scale" fix and gated it on + # `if ck_preshuffle:` alone. + # + # Too wide a gate. The fp32 blockscale path does probe stride(0) != 1 to + # learn the layout, so strided coverage means something there. The mxfp8_128 + # path this op runs (--flydsl --ck_preshuffle) does not: it declares the + # layout with is_x_scale_transposed=True and never reads the stride + # (gemm_op_a8w8.py:978-985 states the contract -- x_scale bytes are + # column-major (K//128, M) inside a contiguous (M, K//128) tensor). So the + # strided tensor exercises nothing real here; it only hands triton a + # stride != 1 specialization that dies in make_llir. + # + # A/B 20260828, that line the only variable. The matrix is 162 cases: the + # UT's 27-value default -m, times this op's six (n,k), M outer. As written + # the sweep dies on case 2 (M=2, padded onto the M=16 row); with + # .view(*x_scale.shape) it reaches case 160 -- so every M through 8192, + # M=16 and M=64 among them. Those are exactly the M #4773's rows cover, so + # the gluon kernel compiles and runs once the layout is right. (Case count + # is derived from where the fault lands, not from a per-case log: _run_child + # keeps only the last 30 lines, and a GPU fault is fatal, so reaching M's + # 27th value is itself the proof the first 26 completed.) + # + # The patched sweep still ends in a GPU fault at its last M, but that is a + # separate, older story: m=10240 n=7168 k=3072 run alone passes at err=0, + # 1220 TFLOPS, split-K checks included. Same shape as a fresh process, so + # state carried across cases -- see the f4gemm note above. + # + # The fix is upstream's call: that gate wants to be `ck_preshuffle and not + # use_flydsl_fp8_scale`, matching how the asm/triton block below already + # excludes this path. Narrowing the gate is the right shape of fix, not + # rewriting the line as .view() -- the two calls differ only in stride, so + # .view() would collapse them and drop coverage of the is_x_scale_tranposed + # == False branch that #4406 added the line for. Evidence either way: + # -m 16 -nk 2048,7168 --ck_preshuffle True passes the strided check with + # the line untouched, and only adding --flydsl makes it crash. + # + # Back in the sweep because _A8W8_BLOCKSCALE_TOKENS now starts at 1024, + # which keeps every shape clear of the M that reach those rows. Verified on + # 20260828, rocm/fw-bringup:gfx1250-atom--20260827-ubench: 36/36 cases, + # err=0 on all, 2207-7003 TFLOPS. That run also clears M=10240, the shape + # the earlier sweep faulted on -- more evidence that fault was cross-case + # state and not the shape. + "a8w8_blockscale", + "a16w16", + "mla_v4_decode", + "inverse_rope", + "mla_v4_prefill", + "mhc", + "qk_norm", + "score_qk", +] + + +def main(): + if get_gfx() not in SUPPORTED_GFX: + print( + f"combo bench targets {SUPPORTED_GFX} only; current {get_gfx()} — skipping" + ) + return + # Before any child is spawned: children that inherit our environ (env=None) + # get these too, not just the ones handed an explicit env. See _pin_arch. + _pin_arch(os.environ) + + p = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="combined gfx1250 asm-kernel perf bench (prints only summaries)", + ) + suite = p.add_mutually_exclusive_group(required=True) + suite.add_argument( + "--perf", + action="store_true", + help=f"run hardware-oriented benchmarks (default ops: {', '.join(PERF_OPS)})", + ) + suite.add_argument( + "--dsv4", + action="store_true", + help=( + "run the DeepSeek-V4 fixed-shape suite " + f"(default ops: {', '.join(DSV4_OPS)})" + ), + ) + p.add_argument( + "--ops", + nargs="*", + choices=list(OPS), + default=None, + help=( + "run these ops instead of the suite defaults. Any op is allowed, " + "including ones held out of the defaults because they are broken " + "on this arch (default: suite defaults)" + ), + ) + # mha (SWA fwd asm) — fixed 4-shape grid; init sweep only + p.add_argument( + "--mha-init", + type=str, + nargs="*", + default=["randn", "const0.25"], + choices=["randn", "const0.25"], + ) + # flydsl moe — fixed kernel-bench config (see _MOE_CONFIG) + # mla_v4 (v4 nm kernarg-preload decode) axes + p.add_argument( + "--mla-v4-kargpreld-shapes", + type=_int_quad, + nargs="*", + default=None, + metavar="GQA,BATCH,CTX,SPLIT", + help="Override curated shape grid as gqa,batch,ctx,split tuples " + "(default: suite-specific built-in grid)", + ) + p.add_argument( + "--mla-v4-kargpreld-iters", + type=int, + default=50, + help="mla_v4_kargpreld timed iterations (default: 50)", + ) + p.add_argument( + "--mla-v4-kargpreld-warmup", + type=int, + default=2, + help="mla_v4_kargpreld warmup iterations (default: 2)", + ) + args = p.parse_args() + + args.suite = "dsv4" if args.dsv4 else "perf" + default_ops = DSV4_OPS if args.dsv4 else PERF_OPS + # --ops selects from every op, not just the suite's defaults: an op pulled + # out of the defaults because it is broken on this arch still has to be + # runnable by name to check whether a newer image fixed it. argparse already + # rejects names outside OPS. + selected_ops = args.ops or default_ops + for name in selected_ops: + with _keep_going(name): + OPS[name](args) + + if _FAILURES: + print(f"\n===== {len(_FAILURES)} failed, " + f"{len(selected_ops)} ops selected =====", flush=True) + for label, why in _FAILURES: + print(f" {label}: {why}", flush=True) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/op_tests/triton_tests/attention/test_mla_v4_triton.py b/op_tests/triton_tests/attention/test_mla_v4_triton.py new file mode 100644 index 0000000000..d9c99829d7 --- /dev/null +++ b/op_tests/triton_tests/attention/test_mla_v4_triton.py @@ -0,0 +1,664 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +# from __future__ import annotations + +import pytest +import torch +import triton + +from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse +from aiter.ops.triton.utils._triton import arch_info +from aiter.test_common import benchmark, checkAllclose, run_perftest + +# MLA v4 sparse-decode parity: D=512 heads, page_size=1 unified pool. +_PA_DECODE_SPARSE_D = 512 +_PERF = {"num_iters": 50, "num_warmup": 2} + + +def _sparse_attn_torch(q, kv, attn_sink, topk_idxs, softmax_scale): + """Per-batch sparse multi-head attention with sink in the denominator only. + + Shapes: + q: [B, M, H, D] + kv: [B, N, D] + attn_sink: [H] + topk_idxs: [B, M, K] int32, -1 means skip + Returns: + [B, M, H, D] same dtype as q. + """ + B, M, H, _D = q.shape + K = topk_idxs.shape[-1] + device = q.device + out_dtype = q.dtype + + valid = topk_idxs != -1 + safe_idxs = topk_idxs.clamp(min=0).long() + batch_idx = torch.arange(B, device=device).view(B, 1, 1).expand(B, M, K) + kv_gathered = kv[batch_idx, safe_idxs] # [B, M, K, D] + kv_f32 = kv_gathered.float() + kv_f32 = torch.where( + valid.unsqueeze(-1), kv_f32, torch.zeros((), dtype=kv_f32.dtype, device=device) + ) + + q_f32 = q.float() + scores = torch.einsum("bmhd,bmkd->bmhk", q_f32, kv_f32) * float(softmax_scale) + scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) + + sink = attn_sink.float().view(1, 1, H, 1).expand(B, M, H, 1) + combined = torch.cat([scores, sink], dim=-1) + cmax = combined.amax(dim=-1, keepdim=True) + cmax = torch.where( + cmax == float("-inf"), + torch.zeros((), dtype=cmax.dtype, device=device), + cmax, + ) + weights = (combined - cmax).exp() + denom = weights.sum(dim=-1, keepdim=True) + weights = weights / denom.clamp(min=1e-30) + weights_kv = weights[..., :K] + out = torch.einsum("bmhk,bmkd->bmhd", weights_kv, kv_f32) + return out.to(out_dtype) + + +def pa_decode_sparse_reference( + q, unified_kv, kv_indices, kv_indptr, attn_sink, softmax_scale +): + """Pure-torch reference that materialises per-token KV via gather.""" + T = q.size(0) + indptr = kv_indptr.to(torch.int64) + spans = (indptr[1:] - indptr[:T]).clamp(min=0) + k_dim = int(spans.max().item()) if T > 0 else 1 + if k_dim == 0: + k_dim = 1 + topk_idxs = torch.full((T, k_dim), -1, device=q.device, dtype=torch.int32) + for t in range(T): + s = int(indptr[t].item()) + n = int(spans[t].item()) + if n > 0: + topk_idxs[t, :n] = kv_indices[s : s + n].to(torch.int32) + return _sparse_attn_torch( + q.unsqueeze(0), + unified_kv.unsqueeze(0), + attn_sink, + topk_idxs.unsqueeze(0), + softmax_scale, + ).squeeze(0) + + +# --------------------------------------------------------------------------- +# Input builder +# --------------------------------------------------------------------------- + + +def _make_inputs( + T: int, + H: int, + D: int, + kv_len_per_token: int, + total_pages: int, + dtype=torch.bfloat16, + seed: int = 0, + include_sentinels: bool = False, + variable_len: bool = False, +): + torch.manual_seed(seed) + device = torch.device("cuda") + + q = torch.randn(T, H, D, dtype=dtype, device=device) * 0.5 + unified_kv = torch.randn(total_pages, D, dtype=dtype, device=device) * 0.5 + attn_sink = torch.randn(H, dtype=torch.float32, device=device) * 0.1 + + # Per-token kv_len: fixed or random in [1, kv_len_per_token]. + if variable_len: + kv_lens = torch.randint( + low=1, + high=kv_len_per_token + 1, + size=(T,), + device=device, + dtype=torch.int64, + ) + else: + kv_lens = torch.full((T,), kv_len_per_token, device=device, dtype=torch.int64) + + indptr = torch.zeros(T + 1, device=device, dtype=torch.int64) + indptr[1:] = kv_lens.cumsum(0) + total_indices = int(indptr[-1].item()) + + indices = torch.randint( + low=0, + high=total_pages, + size=(total_indices,), + device=device, + dtype=torch.int32, + ) + if include_sentinels and total_indices > 0: + # Sprinkle a few -1 sentinels. + n_sentinel = max(1, total_indices // 16) + sentinel_pos = torch.randperm(total_indices, device=device)[:n_sentinel] + indices[sentinel_pos] = -1 + + indptr = indptr.to(torch.int32) + softmax_scale = float(D) ** -0.5 + return q, unified_kv, indices, indptr, attn_sink, softmax_scale + + +@benchmark() +def test_mla_v4_triton_staged(gqa_ratio, batch, kv_seq_lens, num_kv_splits): + """Perf-only stage split: main kernel (s1) + reduce (s2) + total.""" + T = batch + H = gqa_ratio + D = _PA_DECODE_SPARSE_D + pages = T * kv_seq_lens + q, unified_kv, indices, indptr, sink, scale = _make_inputs( + T, H, D, kv_seq_lens, pages, variable_len=False + ) + pa_kwargs = { + "has_invalid": False, + "kv_splits": num_kv_splits, + "num_iters": _PERF["num_iters"], + "num_warmup": _PERF["num_warmup"], + "num_rotate_args": 1, + } + _, us_tot = run_perftest( + pa_decode_sparse, + q, + unified_kv, + indices, + indptr, + sink, + scale, + skip_reduce=False, + **pa_kwargs, + ) + if num_kv_splits > 1: + _, us_s1 = run_perftest( + pa_decode_sparse, + q, + unified_kv, + indices, + indptr, + sink, + scale, + skip_reduce=True, + **pa_kwargs, + ) + triton_s2 = max(0.0, us_tot - us_s1) + else: + us_s1 = us_tot + triton_s2 = 0.0 + return { + "triton_s1": round(us_s1, 2), + "triton_s2": round(triton_s2, 2), + "triton_tot": round(us_tot, 2), + } + + +@benchmark() +def test_mla_v4_triton_perf(gqa_ratio, batch, kv_seq_lens, num_kv_splits): + """Perf sweep row for combo bench / gfx1250 Triton sparse MLA v4 decode. + + Shape ids mirror ``test_mla_v4_kargpreld.test_mla_v4_nm``: + T=batch (q_seq=1), H=gqa_ratio, ctx=kv_seq_lens, kv_splits=num_kv_splits, + D=512. + """ + T = batch + H = gqa_ratio + D = _PA_DECODE_SPARSE_D + pages = T * kv_seq_lens + q, unified_kv, indices, indptr, sink, scale = _make_inputs( + T, H, D, kv_seq_lens, pages, variable_len=False + ) + _, us = run_perftest( + pa_decode_sparse, + q, + unified_kv, + indices, + indptr, + sink, + scale, + has_invalid=False, + kv_splits=num_kv_splits, + num_iters=_PERF["num_iters"], + num_warmup=_PERF["num_warmup"], + num_rotate_args=1, + ) + flops = 4 * T * H * kv_seq_lens * D # QK^T + P@V + bpe = q.element_size() + nbytes = (T * H * D + T * kv_seq_lens * D + T * H * D) * bpe + return { + "us": round(us, 2), + "TFLOPS": round(flops / us / 1e6, 2), + "TB/s": round(nbytes / us / 1e6, 3), + } + + +# --------------------------------------------------------------------------- +# skip_reduce: the wrapper hands back the pre-reduce split-K partials and the +# caller is responsible for the log-sum-exp combine + sink fold. This mirrors +# the _pa_decode_sparse_reduce kernel in pure torch so we can validate the +# partials against the dense reference. +# --------------------------------------------------------------------------- + + +def _wrapper_main_kernel_params(T: int, H: int, D: int): + """Reproduce the (use_exp2, block_k) the wrapper picks for the main kernel. + + Must stay in sync with ``pa_decode_sparse``'s USE_EXP2 and block_k logic. + """ + use_gluon = arch_info.get_arch() == "gfx1250" + use_exp2 = True + if use_gluon: + if H >= 128: + block_h = 128 + elif H >= 64: + if T >= 2048: + block_h = 64 + elif T >= 32: + block_h = 32 + else: + block_h = 16 + elif H >= 32: + if T >= 256: + block_h = 32 + else: + block_h = 16 + else: + block_h = triton.next_power_of_2(H) + else: + block_h = triton.next_power_of_2(min(H, 16)) + if use_gluon: + block_k = 16 + if block_h == 128: + block_k = 32 + else: + block_k = 16 if D >= 256 else 32 + return use_exp2, block_k + + +def _reduce_partials_torch( + acc_partial, m_partial, l_partial, attn_sink, kv_indptr, block_k, use_exp2 +): + """Pure-torch port of _pa_decode_sparse_reduce. + + Shapes: + acc_partial: [T, KV_SPLITS, H_padded, D] fp32 + m_partial: [T, KV_SPLITS, H_padded] fp32 + l_partial: [T, KV_SPLITS, H_padded] fp32 + Returns [T, H, D] in attn_sink-implied output dtype (bf16/fp16 caller casts). + """ + T, kv_splits, _, D = acc_partial.shape + H = attn_sink.shape[0] + device = acc_partial.device + + expfn = torch.exp2 if use_exp2 else torch.exp + LOG2E = 1.4426950408889634 + sink_scale = LOG2E if use_exp2 else 1.0 + + indptr = kv_indptr.to(torch.int64) + kv_lens = (indptr[1 : T + 1] - indptr[:T]).clamp(min=0) + seg_ids = torch.arange(kv_splits, device=device) + sink = attn_sink.float() * sink_scale # [H] + + out = torch.empty(T, H, D, dtype=torch.float32, device=device) + for t in range(T): + n = int(kv_lens[t].item()) + # Match the kernel's tiles_per_segment / act_num_segments masking so we + # ignore the stale (uninitialised) partial-buffer slots that the split + # kernel early-returned on. + if n <= 0: + act_num_segments = 0 + else: + tiles_per_segment = triton.cdiv(n, kv_splits * block_k) + act_num_segments = triton.cdiv(n, tiles_per_segment * block_k) + seg_mask = seg_ids < act_num_segments # [KV_SPLITS] + + m_p = m_partial[t, :, :H].clone() # [KV_SPLITS, H] + l_p = l_partial[t, :, :H] + a_p = acc_partial[t, :, :H, :] # [KV_SPLITS, H, D] + m_p = torch.where(seg_mask[:, None], m_p, torch.full_like(m_p, float("-inf"))) + + m_max = m_p.max(dim=0).values # [H] + is_dead = m_p == float("-inf") # [KV_SPLITS, H] + alpha = torch.where(is_dead, torch.zeros_like(m_p), expfn(m_p - m_max[None, :])) + l_comb = torch.where(is_dead, torch.zeros_like(l_p), l_p * alpha).sum(0) # [H] + acc_comb = torch.where( + is_dead[:, :, None], torch.zeros_like(a_p), a_p * alpha[:, :, None] + ).sum( + 0 + ) # [H, D] + + m_final = torch.maximum(m_max, sink) + alpha_kv = expfn(m_max - m_final) + alpha_sink = expfn(sink - m_final) + l_final = l_comb * alpha_kv + alpha_sink + acc_final = acc_comb * alpha_kv[:, None] + denom = l_final.clamp(min=1e-30) + out[t] = torch.where( + l_final[:, None] > 0.0, + acc_final / denom[:, None], + torch.zeros_like(acc_final), + ) + return out + + +@pytest.mark.parametrize("T", [1, 64, 256, 2048]) +@pytest.mark.parametrize("H", [16, 32, 64, 128]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("kv_len", [136, 388, 1024]) +@pytest.mark.parametrize("var_len", [True, False]) +@pytest.mark.parametrize("sentinels", [False]) +@pytest.mark.parametrize("skip_reduce", [False]) +def test_pa_decode_sparse_vs_reference( + T, H, D, kv_len, var_len, sentinels, skip_reduce +): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + pages = T * kv_len + q, ukv, indices, indptr, sink, scale = _make_inputs( + T, + H, + D, + kv_len, + pages, + include_sentinels=sentinels, + variable_len=var_len, + ) + + ref = pa_decode_sparse_reference(q, ukv, indices, indptr, sink, scale) + result = pa_decode_sparse( + q, + ukv, + indices, + indptr, + sink, + scale, + has_invalid=sentinels, + skip_reduce=skip_reduce, + ) + + if isinstance(result, tuple): + # skip_reduce with the split-K path active (kv_splits > 1): the wrapper + # returns raw partials, so do the log-sum-exp combine + sink fold here. + acc_partial, m_partial, l_partial = result + use_exp2, block_k = _wrapper_main_kernel_params(T, H, D) + out = _reduce_partials_torch( + acc_partial, m_partial, l_partial, sink, indptr, block_k, use_exp2 + ).to(q.dtype) + else: + # kv_splits == 1 (skip_reduce is a no-op) or skip_reduce=False: the + # wrapper already returns the final output. + out = result + + tol_err_ratio = 0.01 + assert ( + checkAllclose( + out.to(torch.bfloat16), + ref.to(torch.bfloat16), + atol=5e-3, + rtol=5e-3, + tol_err_ratio=tol_err_ratio, + msg="pa_decode_sparse output", + ) + <= tol_err_ratio + ) + + +# --------------------------------------------------------------------------- +# FP8 KV cache quantization helpers +# --------------------------------------------------------------------------- + +_FP8_GROUP_SIZE = 64 +_FP8_DTYPE = torch.float8_e4m3fnuz + + +def _quantize_kv_fp8(unified_kv, group_size=_FP8_GROUP_SIZE): + """Quantize bf16/fp16 unified_kv to (fp8, scales) with 1xGROUP_SIZE block scaling. + + Returns (kv_fp8, kv_scales) where kv_fp8 is float8_e4m3fnuz and + kv_scales is [total_pages, D // group_size] fp32. + """ + total_pages, D = unified_kv.shape + assert D % group_size == 0 + num_groups = D // group_size + kv_f32 = unified_kv.float().view(total_pages, num_groups, group_size) + amax = kv_f32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + fp8_max = torch.finfo(_FP8_DTYPE).max + scales = (amax / fp8_max).squeeze(-1) # [total_pages, num_groups] + kv_scaled = kv_f32 / amax * fp8_max + kv_fp8 = kv_scaled.view(total_pages, D).to(_FP8_DTYPE) + return kv_fp8, scales.to(torch.float32) + + +def _dequant_kv_fp8(kv_fp8, kv_scales, group_size=_FP8_GROUP_SIZE): + """Dequantize for reference comparison.""" + total_pages, D = kv_fp8.shape + num_groups = D // group_size + kv_f32 = kv_fp8.float().view(total_pages, num_groups, group_size) + scales_expanded = kv_scales.unsqueeze(-1).expand( + total_pages, num_groups, group_size + ) + return (kv_f32 * scales_expanded).view(total_pages, D) + + +@pytest.mark.parametrize("T", [1, 32]) +@pytest.mark.parametrize("H", [16]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("kv_len", [100]) +@pytest.mark.parametrize("var_len", [True, False]) +def test_pa_decode_sparse_fp8_vs_reference(T, H, D, kv_len, var_len): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + pages = T * kv_len + q, ukv_bf16, indices, indptr, sink, scale = _make_inputs( + T, + H, + D, + kv_len, + pages, + variable_len=var_len, + ) + + # Quantize KV to fp8 + scales + kv_fp8, kv_scales = _quantize_kv_fp8(ukv_bf16) + + # Reference: dequant back to bf16, run the torch reference + ukv_deq = _dequant_kv_fp8(kv_fp8, kv_scales).to(q.dtype) + ref = pa_decode_sparse_reference(q, ukv_deq, indices, indptr, sink, scale) + + # Triton kernel with fp8 kv + kv_scales + out = pa_decode_sparse( + q, + kv_fp8, + indices, + indptr, + sink, + scale, + kv_scales=kv_scales, + has_invalid=False, + ) + + tol_err_ratio = 0.01 + assert ( + checkAllclose( + out.to(torch.bfloat16), + ref.to(torch.bfloat16), + atol=1e-2, + rtol=1e-2, + tol_err_ratio=tol_err_ratio, + msg="pa_decode_sparse output", + ) + <= tol_err_ratio + ) + + +def make_packed_cache(num_tokens, D, dtype): + device = "cuda" + rope = 64 # DSv4 RoPE dim, stored bf16 + block = 256 # packed cache page size + nope = D - rope # NoPE dim, stored fp8 e4m3 OCP + nb = triton.cdiv(num_tokens, block) + if dtype == "bf16": + cache = (torch.randn(nb, block, D, device=device) * 0.4).to(torch.bfloat16) + return cache, cache.reshape(nb * block, D).float() + + # per token: [nope fp8 (1B) | rope bf16 (2B) | 8 UE8M0 scale bytes] + data_bytes = nope + rope * 2 + scale_bytes = 8 + row_bytes = data_bytes + scale_bytes + cache = torch.zeros(nb, block, row_bytes, dtype=torch.uint8, device=device) + flat = cache.view(nb, block * row_bytes) + data = flat[:, : block * data_bytes].view(nb, block, data_bytes) + scales_region = flat[:, block * data_bytes :].view(nb, block, scale_bytes) + nope_fp8 = (torch.randn(nb, block, nope, device=device) * 0.4).to( + torch.float8_e4m3fn + ) + data[:, :, :nope] = nope_fp8.view(torch.uint8) + rope_bf16 = (torch.randn(nb, block, rope, device=device) * 0.4).to(torch.bfloat16) + data[:, :, nope:data_bytes] = rope_bf16.view(torch.uint8).view(nb, block, rope * 2) + num_groups = nope // 64 + exps = torch.randint( + 124, 130, (nb, block, num_groups), device=device, dtype=torch.uint8 + ) + scales_region[:, :, :num_groups] = exps + scales = torch.exp2(exps.float() - 127.0).repeat_interleave(64, dim=2) + kv_deq = torch.cat([nope_fp8.float() * scales, rope_bf16.float()], dim=2) + return cache, kv_deq.reshape(nb * block, D) + + +def widen_to_int32_overflow(cache, kv_deq): + """Re-lay ``cache`` as a strided view whose span exceeds a 32-bit offset. + + Same nelement() and same contents, but the dim-0 pitch is stretched so the + last block sits past 2**31 bytes. Only the blocks themselves are written; + the padding between them is left uninitialised, so the pool costs its + address space but not the time to fill it. + """ + nb, block, row = cache.shape + itemsize = cache.element_size() + pitch = triton.cdiv(2**31, max(1, nb - 1) * itemsize) + pitch = max(pitch, block * row) + # the packed fp8 cache is viewed as bfloat16, which needs an even stride + pitch += pitch % 2 + pool = torch.empty( + pitch * (nb - 1) + block * row, dtype=cache.dtype, device=cache.device + ) + view = pool.as_strided((nb, block, row), (pitch, row, 1)) + view.copy_(cache) + assert view.stride(0) * itemsize * (nb - 1) >= 2**31 + return view, kv_deq + + +def two_loop_reference( + q, + main_deq, + main_idx, + main_indptr, + extra_deq, + extra_idx, + extra_indptr, + attn_sink, + softmax_scale, +): + """Reference for the SWA(main) + top-k(extra) two-loop: concatenate the two + dequantized pools, merge the two ragged index sets (extra slots shifted past + the main pool), then reuse ``pa_decode_sparse_reference``. + """ + main_pages = main_deq.shape[0] + combined = torch.cat([main_deq, extra_deq], dim=0).to(q.dtype) + T = main_indptr.numel() - 1 + mi, mp = main_idx.long(), main_indptr.long() + ei, ep = extra_idx.long(), extra_indptr.long() + rows, lens = [], [] + for tok in range(T): + row = torch.cat( + [mi[mp[tok] : mp[tok + 1]], ei[ep[tok] : ep[tok + 1]] + main_pages] + ) + rows.append(row) + lens.append(row.numel()) + combined_idx = torch.cat(rows).to(torch.int32) + combined_indptr = torch.zeros(T + 1, dtype=torch.int32, device=q.device) + combined_indptr[1:] = torch.tensor(lens, device=q.device).cumsum(0) + return pa_decode_sparse_reference( + q, combined, combined_idx, combined_indptr, attn_sink, softmax_scale + ) + + +@pytest.mark.parametrize("T", [1, 32, 128]) +@pytest.mark.parametrize("H", [16]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("main_len", [128]) +@pytest.mark.parametrize("extra_len", [8, 256]) +@pytest.mark.parametrize("dtype", ["bf16", "fp8"]) +@pytest.mark.parametrize("strided_cache", [False, True]) +def test_pa_decode_sparse_two_loop(T, H, D, main_len, extra_len, dtype, strided_cache): + """gfx950 vLLM DSv4 decode path: SWA (main) + top-k (extra) two-loop over + packed caches. fp8 (fp8_ds_mla) is the vLLM production format; bf16 is also + exercised. Skipped off gfx950 (extra_* is a packed-only gluon path).""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + if arch_info.get_arch() != "gfx950": + pytest.skip("two-loop (extra_*) is a gfx950 packed-cache-only path") + if strided_cache: + # The pool has to span >2 GiB for the offsets to overflow, so pin the + # regression to one shape -- the fp8 production format at the largest T + # -- rather than paying it on all 24 combinations. + if dtype != "fp8" or T != 128: + pytest.skip("strided-cache case is pinned to the fp8 T=128 shape") + if torch.cuda.mem_get_info()[0] < 4 * 1024**3: + pytest.skip("needs ~3 GiB free for the >2 GiB strided pool") + + device = "cuda" + torch.manual_seed(0) + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=device) * 0.125 + attn_sink = torch.randn(H, dtype=torch.float32, device=device) * 0.1 + softmax_scale = float(D) ** -0.5 + + # main = contiguous SWA window per query + main_cache, main_deq = make_packed_cache(T * main_len, D, dtype) + query_base = (torch.arange(T, device=device) * main_len)[:, None] + main_idx = ( + (query_base + torch.arange(main_len, device=device)).to(torch.int32).reshape(-1) + ) + main_indptr = torch.arange( + 0, T * main_len + 1, main_len, dtype=torch.int32, device=device + ) + # extra = scattered top-k over a pool + extra_pool = T * extra_len + extra_cache, extra_deq = make_packed_cache(extra_pool, D, dtype) + if strided_cache: + extra_cache, extra_deq = widen_to_int32_overflow(extra_cache, extra_deq) + extra_idx = torch.randint( + 0, extra_pool, (T, extra_len), device=device, dtype=torch.int32 + ).reshape(-1) + extra_indptr = torch.arange( + 0, T * extra_len + 1, extra_len, dtype=torch.int32, device=device + ) + + ref = two_loop_reference( + q, + main_deq, + main_idx, + main_indptr, + extra_deq, + extra_idx, + extra_indptr, + attn_sink, + softmax_scale, + ) + out = pa_decode_sparse( + q, + main_cache, + main_idx, + main_indptr, + attn_sink, + softmax_scale, + extra_cache=extra_cache, + extra_indices=extra_idx, + extra_indptr=extra_indptr, + ) + + tol = 1e-2 if dtype == "fp8" else 5e-3 + torch.testing.assert_close(out, ref, atol=tol, rtol=tol)