From 28d9a5be5bc1477d841f027e73b844b4adb1bf9b Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Thu, 3 Sep 2026 22:58:46 +0800 Subject: [PATCH 1/3] perf(gfx1250): complete DSv4 data init controls --- op_tests/bench_gfx1250_combo.py | 71 ++++-- op_tests/bench_init.py | 241 ++++++++++++++++++ .../triton/bench_deepgemm_attention.py | 55 ++-- op_tests/test_flydsl_grouped_gemm_gfx1250.py | 2 +- op_tests/test_flydsl_qk_norm_rope_quant.py | 19 +- op_tests/test_gemm_a8w8_blockscale.py | 2 +- op_tests/test_inverse_rope_group_quant.py | 99 ++++++- op_tests/test_mhc.py | 2 +- op_tests/test_mla_v4_kargpreld.py | 75 ++++-- op_tests/test_pa_sparse_prefill.py | 91 +++++-- .../attention/test_mla_v4_triton.py | 66 ++++- 11 files changed, 631 insertions(+), 92 deletions(-) create mode 100644 op_tests/bench_init.py diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index d01500c7d1..be1c13df55 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -102,14 +102,13 @@ The current passthrough matrix is: DATA + SCALE + seed moe, gemm, f8gemm, a8w8_blockscale - DATA + seed a16w16, mega_moe, mhc, qk_norm + DATA + seed a16w16, mega_moe, mhc, qk_norm, inverse_rope, + score_qk, mla_v4_decode, mla_v4_prefill DATA mapping only mha (norm -> randn, constant -> const0.25) - native init only mla_v4_decode, mla_v4_prefill, score_qk, - inverse_rope, mori_ep + native init only mori_ep ``--scale-init`` is reported as not applicable for operators without a scale -operand. ``mla_v4_prefill`` still receives ``--seed`` because its native data -generator exposes that control even though it does not expose a distribution. +operand. 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 @@ -683,6 +682,8 @@ def _moe_stage_bytes( ] _MLA_V4_COMPARE_KEEP = [ "dtype", + "data_init", + "seed", "gqa_ratio", "batch", "kv_seq_lens", @@ -1533,7 +1534,7 @@ def run_qk_norm(args): def run_score_qk(args): """Run DSv4 decode score-QK at batch 512 for short and long CSA KV.""" - _unsupported_init(args, "score_qk") + _unused_scale_init(args, "score_qk") repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) base_cmd = [ sys.executable, @@ -1549,16 +1550,23 @@ def run_score_qk(args): "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 + for tokens, (label, kv_length), data_init in itertools.product( + _SCORE_QK_TOKENS or (None,), + _SCORE_QK_KV_LENGTHS, + args.data_init or ["norm"], ): _run_child( - f"score_qk (decode, B={tokens or 'UT default'}, {label} CSA KV={kv_length})", + f"score_qk (decode, B={tokens or 'UT default'}, {label} " + f"CSA KV={kv_length}, init={data_init}, seed={args.seed})", [ *base_cmd, *(["--batch", str(tokens)] if tokens else []), "-kv_length", kv_length, + "--data-init", + data_init, + "--seed", + str(args.seed), ], cwd=repo_root, extract=_md_from_pandas( @@ -1629,7 +1637,9 @@ def _perf_ratio(num, den): return f"{num / den:.2f}x" -def _bench_mla_v4_asm_staged(gqa, batch, ctx, split_kv, num_iters, num_warmup): +def _bench_mla_v4_asm_staged( + gqa, batch, ctx, split_kv, num_iters, num_warmup, data_init, seed +): """Asm kernel (s1) + merge (s2) + total; lives in combo bench only.""" mod = mla_v4_kargpreld_mod q_seq = 1 @@ -1645,7 +1655,8 @@ def _bench_mla_v4_asm_staged(gqa, batch, ctx, split_kv, num_iters, num_warmup): batch=batch, kv_seq_lens=ctx, q_seq_logical=q_seq, - seed=mod._SEED, + seed=seed, + data_init=data_init, gqa_ratio=gqa, attn_sink=True, ) @@ -1733,7 +1744,7 @@ def _bench_mla_v4_asm_staged(gqa, batch, ctx, split_kv, num_iters, num_warmup): def run_mla_v4_decode(args): # Side-by-side asm (kargpreld) vs Triton sparse decode on the same shape grid. - _unsupported_init(args, "mla_v4_decode") + _unused_scale_init(args, "mla_v4_decode") iters = args.mla_v4_kargpreld_iters warmup = args.mla_v4_kargpreld_warmup mla_v4_triton_mod._PERF["num_iters"] = iters @@ -1744,22 +1755,38 @@ def run_mla_v4_decode(args): else _MLA_V4_KARGPRELD_SHAPES ) shapes = args.mla_v4_kargpreld_shapes or default_shapes + data_inits = args.data_init or ["norm"] rows = [] with _capture() as box: - for gqa, batch, ctx, split_kv in shapes: + for (gqa, batch, ctx, split_kv), data_init in itertools.product( + shapes, data_inits + ): row = { + "data_init": data_init, + "seed": args.seed, "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) + asm = _bench_mla_v4_asm_staged( + gqa, + batch, + ctx, + split_kv, + iters, + warmup, + data_init, + args.seed, + ) tri = mla_v4_triton_mod.test_mla_v4_triton_staged( gqa_ratio=gqa, batch=batch, kv_seq_lens=ctx, num_kv_splits=split_kv, + data_init=data_init, + seed=args.seed, ) row.update(asm) row.update(tri) @@ -1782,7 +1809,7 @@ def run_mla_v4_decode(args): def run_inverse_rope(args): """Run DSv4 inverse RoPE + group quant at the tp1 attention-output shape.""" - _unsupported_init(args, "inverse_rope") + _unused_scale_init(args, "inverse_rope") # -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. @@ -1798,6 +1825,13 @@ def run_inverse_rope(args): "n32k4", "--group-size", "32", + *( + ["--data-init", *args.data_init] + if args.data_init is not None + else [] + ), + "--seed", + str(args.seed), ], cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ) @@ -1805,7 +1839,7 @@ def run_inverse_rope(args): def run_mla_v4_prefill(args): """Run DSv4 prefill across two precisions, pools and CSR modes.""" - _unsupported_init(args, "mla_v4_prefill") + _unused_scale_init(args, "mla_v4_prefill") for tokens in _MLA_PREFILL_TOKENS: _run_child( f"mla_v4 prefill (M={tokens}, prec=fp8/bf16, pages=4096/16384)", @@ -1835,6 +1869,11 @@ def run_mla_v4_prefill(args): "--no-verify", "--seed", str(args.seed), + *( + ["--data-init", *args.data_init] + if args.data_init is not None + else [] + ), ], cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), # Not _table_row: the UT has no "latency_us" column (it prints diff --git a/op_tests/bench_init.py b/op_tests/bench_init.py new file mode 100644 index 0000000000..5106133e66 --- /dev/null +++ b/op_tests/bench_init.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Bench-style static DATA/SCALE init for gfx1250 FP4/FP8 GEMM tests. + +Two *independent* axes: data and scale are sampled independently, **not** "make +the data then derive the scale by quantizing it". + +DATA (per element, iid) -- ``--data-init``:: + + uniform : FP4 -> U(-3,3), FP8 -> U(-6,6) [default] + gaussian : N(0,1) [norm-dist / LLM-like] + trig : trig_float in [-2,2] [optimistic pattern] + random : pure random on-wire codes [overly pessimistic baseline] + +then round-to-nearest into the on-wire low-precision format (FP4 e2m1 packed +2/byte, FP8 e4m3). + +SCALE (per block) -- ``--scale-init``:: + + pow2_binomial : E8M0 value = 2^(Binomial(2n+1,0.5) - (n+1)), n=10 + gaussian : N(0.34375, 0.08) -> E4M3 + random : random on-wire byte in a modest range + auto : format-recommended default (E8M0 -> pow2_binomial, + E4M3 -> gaussian(0.34375, 0.08)) [default] + +Every sampler takes a ``torch.Generator`` so a fixed ``--seed`` reproduces the +buffers bit-for-bit. ``constant`` init is operator-specific (fixed representable +bytes) and stays in the individual test files. +""" + +import math + +import torch + +FP8_E4M3 = torch.float8_e4m3fn # gfx1250 fp8 data & E4M3 block-scale encoding + +# Selectable distributions (kept in sync with the test-file argparse choices). +DATA_DISTS = ("constant", "uniform", "gaussian", "trig", "random") +E8M0_SCALE_DISTS = ("constant", "pow2_binomial", "random", "auto") +E4M3_SCALE_DISTS = ("constant", "gaussian", "random", "auto") + +# Per-format uniform ranges for scaled low-precision DATA. +FP4_UNIFORM = (-3.0, 3.0) # E2M1 max is 6.0; keep headroom +FP8_UNIFORM = (-6.0, 6.0) # a touch wider than FP4 + +# E4M3 block-scale gaussian: recognisable non-2^ center + spread. +E4M3_SCALE_MEAN, E4M3_SCALE_STD = 0.34375, 0.08 + +# pow2_binomial exponent recenter: e = Binomial(2n+1, 0.5) - (n+1). +# n=10 -> Binomial(21,0.5)-11, exponents centred at -0.5. Lower n narrows the +# scale dynamic range (useful if a correctness ref hits the allclose tolerance). +POW2_BINOMIAL_N = 10 + +# e4m3fn NaN encodings (OCP): S.1111.111 -> 0x7F / 0xFF. Everything else is a +# finite representable code, so "random codes" only needs to avoid these two. +_E4M3_NAN_POS, _E4M3_NAN_NEG = 0x7F, 0xFF + + +def make_generator(seed, device="cuda"): + """Seeded ``torch.Generator`` -- same seed => bit-identical buffers.""" + return torch.Generator(device=device).manual_seed(int(seed)) + + +# --------------------------------------------------------------------------- # +# DATA +# --------------------------------------------------------------------------- # +# Cap on the f32 staging buffer. The low-precision buffers are built by sampling +# f32 and narrowing, so a perf shape like (1048576, 16384) would ask for a single +# 64 GiB intermediate before anything shrinks. Generate it in row chunks instead. +_STAGE_ELEMS = 1 << 28 # 256M f32 = 1 GiB per chunk + + +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 _trig_phase(dist, gen, device): + """Draw trig's phase once, outside the chunk loop. + + Sampling it per chunk would both consume the generator a chunk-count- + dependent number of times and make the pattern depend on how the rows were + split -- same seed, different buffer. + """ + if dist != "trig": + return None + return torch.rand(1, generator=gen, device=device).item() * (2.0 * math.pi) + + +def _sample_f32(shape, dist, gen, *, lo, hi, device, phase=None, idx_offset=0): + """Sample an f32 tensor for the requested continuous DATA distribution. + + ``idx_offset`` is the flat index this chunk starts at, so a row-chunked + caller reproduces the same trig pattern as a single full-size call. + """ + if dist == "uniform": + return torch.empty(shape, dtype=torch.float32, device=device).uniform_( + lo, hi, generator=gen + ) + if dist == "gaussian": + return torch.empty(shape, dtype=torch.float32, device=device).normal_( + 0.0, 1.0, generator=gen + ) + if dist == "trig": + # Deterministic trig_float in [-2,2]; the generator only jitters the + # phase so --seed still varies the pattern without breaking repro. + n = 1 + for s in shape: + n *= s + if phase is None: + phase = _trig_phase(dist, gen, device) + idx = torch.arange( + idx_offset, idx_offset + n, dtype=torch.float32, device=device + ) + return (2.0 * torch.sin(0.017 * idx + phase)).reshape(shape) + raise ValueError(f"data dist {dist!r} is not continuous; use fill_* dispatch") + + +def fill_fp4(shape, dist, gen, *, uniform=FP4_UNIFORM, device="cuda"): + """Return packed e2m1 ``uint8`` of shape ``(rows, cols // 2)``. + + ``dist`` in {uniform, gaussian, trig} -> sample f32 then round to e2m1. + ``dist == "random"`` -> uniform over all e2m1 codes (every byte is a valid + pair of e2m1 nibbles). + """ + rows, cols = shape + assert cols % 2 == 0, f"FP4 needs even columns, got {cols}" + if dist == "random": + return torch.randint( + 0, 256, (rows, cols // 2), dtype=torch.uint8, device=device, generator=gen + ) + # Local import: fp4_utils pulls in triton; keep module import cheap. + from aiter.utility import fp4_utils + + out = torch.empty((rows, cols // 2), dtype=torch.uint8, device=device) + phase = _trig_phase(dist, gen, device) + for r0, r1 in _row_chunks(rows, cols): + v = _sample_f32( + (r1 - r0, cols), + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + phase=phase, + idx_offset=r0 * cols, + ) + 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"): + """Return an e4m3 tensor of ``shape``. + + ``dist`` in {uniform, gaussian, trig} -> sample f32 then cast to e4m3. + ``dist == "random"`` -> uniform over finite e4m3 codes (NaN bytes remapped). + """ + if dist == "random": + b = torch.randint( + 0, 256, shape, dtype=torch.uint8, device=device, generator=gen + ) + b[b == _E4M3_NAN_POS] = 0x00 # +NaN -> +0 + b[b == _E4M3_NAN_NEG] = 0x80 # -NaN -> -0 + return b.view(FP8_E4M3) + if len(shape) != 2: + v = _sample_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) + phase = _trig_phase(dist, gen, device) + for r0, r1 in _row_chunks(rows, cols): + v = _sample_f32( + (r1 - r0, cols), + dist, + gen, + lo=uniform[0], + hi=uniform[1], + device=device, + phase=phase, + idx_offset=r0 * cols, + ) + out[r0:r1] = v.to(FP8_E4M3) + del v + return out + + +# --------------------------------------------------------------------------- # +# SCALE +# --------------------------------------------------------------------------- # +def fill_scale_e8m0(shape, dist="auto", gen=None, *, device="cuda", n=POW2_BINOMIAL_N): + """Return E8M0 on-wire ``uint8`` (biased exponent, bias 127). + + ``auto``/``pow2_binomial`` -> value = 2^(Binomial(2n+1, 0.5) - (n+1)). + ``random`` -> uniform exponent in [-2, 2] (modest, ref-friendly). + """ + if dist not in E8M0_SCALE_DISTS: + raise ValueError(f"E8M0 scale dist {dist!r}; choose from {E8M0_SCALE_DISTS}") + if dist == "random": + return torch.randint( + 125, 130, shape, dtype=torch.uint8, device=device, generator=gen + ) + # value = 2^(Binomial(2n+1, 0.5) - (n+1)); Binomial(k, 0.5) == popcount of a + # uniform k-bit int, so one randint + popcount (vs 2n+1 rand kernels). + 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 + 127).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_scale_e4m3(shape, dist="auto", gen=None, *, device="cuda"): + """Return E4M3 on-wire ``uint8``. + + ``auto``/``gaussian`` -> N(0.34375, 0.08), clamped non-negative, cast e4m3. + ``random`` -> uniform over e4m3 bytes in [0x20, 0x50) (legacy NVFP4 range). + """ + if dist not in E4M3_SCALE_DISTS: + raise ValueError(f"E4M3 scale dist {dist!r}; choose from {E4M3_SCALE_DISTS}") + if dist == "random": + return torch.randint( + 0x20, 0x50, shape, dtype=torch.uint8, device=device, generator=gen + ) + v = torch.empty(shape, dtype=torch.float32, device=device).normal_( + E4M3_SCALE_MEAN, E4M3_SCALE_STD, generator=gen + ) + v.clamp_(min=0.0) # block scales are non-negative + return v.to(FP8_E4M3).view(torch.uint8) diff --git a/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py b/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py index 7efa3fea0d..7724016354 100644 --- a/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py +++ b/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py @@ -15,7 +15,7 @@ ) from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.test_common import run_perftest +from aiter.test_common import DATA_DISTS, fill, make_generator, run_perftest def cdiv(x: int, y: int) -> int: @@ -189,7 +189,7 @@ def create_paged_mqa_logits_configs(args: argparse.Namespace): return configs -def run_benchmark(args: argparse.Namespace): +def run_benchmark(args: argparse.Namespace, data_init: str = "norm"): ChunkK = 128 WavePerEU = 5 @@ -197,8 +197,9 @@ def run_benchmark(args: argparse.Namespace): def test_deepgemm_fp8_paged_mqa_logits( batch_size, next_n, heads, index_dim, avg_kv_length, kv_storage_kind ): - torch.manual_seed(0) - random.seed(0) + torch.manual_seed(args.seed) + random.seed(args.seed) + gen = make_generator(args.seed) max_model_len = 2 * avg_kv_length blocksize = args.blocksize if args.kv_preshuffle else 1 @@ -225,19 +226,22 @@ def test_deepgemm_fp8_paged_mqa_logits( ) prefix_sum_context_lens[1:] = torch.cumsum(context_lens, dim=0) - q = torch.randn( - (batch_size, next_n, heads, index_dim), - device="cuda", + q = fill( + (batch_size * next_n * heads, index_dim), + data_init, + gen, dtype=torch.bfloat16, - ) - kv_cache = torch.randn( - (num_blocks, blocksize, 1, index_dim), - device="cuda", + ).view(batch_size, next_n, heads, index_dim) + kv_cache = fill( + (num_blocks * blocksize, index_dim), + data_init, + gen, dtype=torch.bfloat16, - ) - weights = torch.randn( + ).view(num_blocks, blocksize, 1, index_dim) + weights = fill( (batch_size * next_n, heads), - device="cuda", + data_init, + gen, dtype=torch.float32, ) @@ -273,7 +277,7 @@ def test_deepgemm_fp8_paged_mqa_logits( for i in range(batch_size): ctx_len = int(context_lens[i].item()) kv_indices[prefix_sum_context_lens[i] : prefix_sum_context_lens[i + 1]] = ( - torch.randperm(max_model_len, device="cuda")[:ctx_len] + torch.randperm(max_model_len, device="cuda", generator=gen)[:ctx_len] ) if kv_storage_kind == "non_ragged_k": @@ -369,6 +373,10 @@ def test_deepgemm_fp8_paged_mqa_logits( def calc_diff(x: torch.Tensor, y: torch.Tensor): x, y = x.double(), y.double() denominator = (x * x + y * y).sum() + # zero-init makes both logits tensors exactly zero. Treat that + # exact match as zero error instead of reporting 0/0 -> NaN. + if denominator == 0: + return torch.zeros_like(denominator) sim = 2 * (x * y).sum() / denominator return 1 - sim @@ -474,6 +482,21 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): action="store_true", help="Disable varctx schedule (only applies with --kv_preshuffle)", ) + parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) for Q, KV and weights", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for input data and generated index tables (default: 0)", + ) args = parser.parse_args() - run_benchmark(args) + for data_init in args.data_init: + print(f"data_init={data_init} seed={args.seed}") + run_benchmark(args, data_init=data_init) diff --git a/op_tests/test_flydsl_grouped_gemm_gfx1250.py b/op_tests/test_flydsl_grouped_gemm_gfx1250.py index 66a5d00e48..bb5f671165 100644 --- a/op_tests/test_flydsl_grouped_gemm_gfx1250.py +++ b/op_tests/test_flydsl_grouped_gemm_gfx1250.py @@ -36,6 +36,7 @@ import torch from aiter import ActivationType, QuantType, logger +from aiter import test_common as bench_init from aiter.aot.flydsl.common import run_only_env from aiter.fused_moe import ( fused_moe, @@ -47,7 +48,6 @@ from aiter.ops.quant import per_1x32_f4_quant from aiter.ops.shuffle import moe_shuffle_scale, moe_shuffle_weight from aiter.utility import dtypes, fp4_utils -from op_tests import bench_init # Build every tensor straight on the device (like op_tests/test_moe_2stage.py) so # the test body has no `.cuda()` / `.float().cuda()` plumbing. diff --git a/op_tests/test_flydsl_qk_norm_rope_quant.py b/op_tests/test_flydsl_qk_norm_rope_quant.py index 0da5565deb..6896f85204 100755 --- a/op_tests/test_flydsl_qk_norm_rope_quant.py +++ b/op_tests/test_flydsl_qk_norm_rope_quant.py @@ -588,7 +588,14 @@ def test_flydsl_swa_write( # a write; they change no in-pool byte, so only a dirtied guard row can show # that one regressed. G = _SWA_GUARD_ROWS - pool = torch.zeros(G + num_rows + G, D, dtype=torch.bfloat16, device=device) + # A zero-filled pool cannot distinguish "the kernel wrote a zero row" from + # "the kernel skipped this row" when --init zero. Seed every row with an + # unreachable sentinel so the write-mask check remains valid for all data + # distributions. + sentinel = torch.finfo(torch.bfloat16).max + pool = torch.full( + (G + num_rows + G, D), sentinel, dtype=torch.bfloat16, device=device + ) swa_kv = pool[G : G + num_rows] mode_kw = ( {"swa_dest_rows": index_t} @@ -636,7 +643,7 @@ def test_flydsl_swa_write( # The scatter is a verbatim copy of kv_out, so the reference IS kv_out # gathered onto the rows the addressing mode selects. - expected = torch.zeros_like(swa_kv) + expected = torch.full_like(swa_kv, sentinel) n_written = 0 for t in range(T): if dest[t] < 0: @@ -658,10 +665,12 @@ def test_flydsl_swa_write( ) # A skipped token must reach NO row, not merely the right one. assert ( - int((swa_kv != 0).any(dim=1).sum()) == n_written + int((swa_kv != sentinel).any(dim=1).sum()) == n_written ), f"{mode}: a skipped token still reached the pool" - assert not pool[:G].any(), f"{mode}: scatter wrote BEFORE the pool" - assert not pool[G + num_rows :].any(), f"{mode}: scatter wrote PAST the pool" + assert (pool[:G] == sentinel).all(), f"{mode}: scatter wrote BEFORE the pool" + assert (pool[G + num_rows :] == sentinel).all(), ( + f"{mode}: scatter wrote PAST the pool" + ) # The scatter must not perturb the primary outputs. ref_q, ref_kv, _, _ = flydsl_qk_norm_rope_quant( diff --git a/op_tests/test_gemm_a8w8_blockscale.py b/op_tests/test_gemm_a8w8_blockscale.py index bf35b2fc80..7f954d98fb 100644 --- a/op_tests/test_gemm_a8w8_blockscale.py +++ b/op_tests/test_gemm_a8w8_blockscale.py @@ -16,11 +16,11 @@ import aiter from aiter import dtypes +from aiter import test_common as bench_init from aiter.ops.gemm_op_a8w8 import gemm_a8w8_blockscale_ck, gemm_a8w8_blockscale_cktile from aiter.ops.shuffle import shuffle_weight from aiter.test_common import benchmark, checkAllclose, perftest from aiter.utility import fp4_utils -from op_tests import bench_init block_shape = (128, 128) TEST_NUM_ITERS = 100 diff --git a/op_tests/test_inverse_rope_group_quant.py b/op_tests/test_inverse_rope_group_quant.py index f1272cbe79..16600d9cbc 100644 --- a/op_tests/test_inverse_rope_group_quant.py +++ b/op_tests/test_inverse_rope_group_quant.py @@ -32,8 +32,11 @@ from aiter.ops.quant import dynamic_per_group_scaled_quant from aiter.ops.triton.rope.rope import RotateStyle, _rope_cached_bwd from aiter.test_common import ( + DATA_DISTS, benchmark, checkAllclose, + fill, + make_generator, run_perftest, ) @@ -295,7 +298,7 @@ def _check_scale_layout(scale, s, g, ks, scale_layout, group_size, name): ), f"{name}: {scale_layout} scale should be {expect}, got {tuple(scale.shape)}" -def _make_inputs(s, h, head_dim, rd, dtype, seed=0): +def _make_inputs(s, h, head_dim, rd, dtype, data_init="norm", seed=0): """Build (o, positions, cos, sin) for one config. cos/sin are the 2D [max_pos, rd//2] the op takes. A model holding the @@ -305,12 +308,16 @@ def _make_inputs(s, h, head_dim, rd, dtype, seed=0): own call site, the way run_inverse_rope_inplace does for the triton rope. Shared by the sweep and the graph check so the two cannot drift. """ - torch.manual_seed(seed) + gen = make_generator(seed) positions = torch.arange(s, dtype=dtypes.i64) % MAX_POS # /10 keeps a group's amax away from fp8 saturation, like a real # post-softmax attention output. - o = torch.randn((s, h, head_dim), dtype=dtype) / 10 - theta = torch.randn((MAX_POS, rd // 2), dtype=dtypes.fp32) + o = ( + fill((s * h, head_dim), data_init, gen, dtype=dtype) + .view(s, h, head_dim) + .div_(10) + ) + theta = fill((MAX_POS, rd // 2), data_init, gen, dtype=dtypes.fp32) cos = torch.cos(theta).to(dtype).contiguous() sin = torch.sin(theta).to(dtype).contiguous() return o, positions, cos, sin @@ -440,12 +447,23 @@ def run_unfused(x, positions, cos, sin, num_groups, quant_group_size, rd, out): @benchmark() def test_inverse_rope_group_quant( - s, h, g, head_dim, rd, group_size, dtype, scale_layout + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init="norm", + seed=0, ): d = h * head_dim // g scale_n = d // group_size - o, positions, cos, sin = _make_inputs(s, h, head_dim, rd, dtype) + o, positions, cos, sin = _make_inputs( + s, h, head_dim, rd, dtype, data_init=data_init, seed=seed + ) ref = run_torch(o, positions, cos, sin, g, group_size, rd) ref_rt = run_torch(o, positions, cos, sin, g, group_size, rd, roundtrip=True) @@ -546,14 +564,27 @@ def unfused_once(): return ret -def check_graph(s, h, g, head_dim, rd, group_size, dtype, scale_layout): +def check_graph( + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init="norm", + seed=0, +): """Capture the op in a HIP graph, replay on fresh data, compare against eager. Not part of the perf table: this is a pass/fail check that the host-side dispatch tier and the pre-allocated buffers survive capture/replay. """ d = h * head_dim // g - o, positions, cos, sin = _make_inputs(s, h, head_dim, rd, dtype) + o, positions, cos, sin = _make_inputs( + s, h, head_dim, rd, dtype, data_init=data_init, seed=seed + ) x_fp8, x_scale = _alloc_outputs(s, g, d, group_size, scale_layout=scale_layout) kwargs = { "num_groups": g, @@ -575,7 +606,9 @@ def check_graph(s, h, g, head_dim, rd, group_size, dtype, scale_layout): inverse_rope_group_quant_cpp(o, positions, cos, sin, **kwargs) # Replay on new data, then compare against an eager run on the same data. - o2, positions2, cos2, sin2 = _make_inputs(s, h, head_dim, rd, dtype, seed=7) + o2, positions2, cos2, sin2 = _make_inputs( + s, h, head_dim, rd, dtype, data_init=data_init, seed=seed + 7 + ) o.copy_(o2) positions.copy_(positions2) cos.copy_(cos2) @@ -740,6 +773,19 @@ def main(): help="""Also run the HIP-graph capture/replay check over the same sweep. e.g.: --graph -s 1 4 32 128 300 512 700 2048""", ) + parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) (default: norm)", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for o and the RoPE cache source (default: 0)", + ) parser.add_argument( "--opus-tree", default=os.environ.get("AITER_OPUS_TREE"), @@ -756,13 +802,22 @@ def main(): for dtype in args.dtype: df = [] - for (h, g), s, head_dim, rd, group_size, scale_layout in itertools.product( + for ( + (h, g), + s, + head_dim, + rd, + group_size, + scale_layout, + data_init, + ) in itertools.product( args.hg, args.tokens, args.head_dim, args.rope_dim, args.group_size, args.scale_layout, + args.data_init, ): # n32k4 only exists at group 32: its four packed k groups are one # WMMA-K=128 step, so 4 * group_size has to be 128. The op rejects @@ -770,11 +825,31 @@ def main(): if scale_layout == "n32k4" and group_size != 32: continue ret = test_inverse_rope_group_quant( - s, h, g, head_dim, rd, group_size, dtype, scale_layout + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init=data_init, + seed=args.seed, ) df.append(ret) if args.graph: - check_graph(s, h, g, head_dim, rd, group_size, dtype, scale_layout) + check_graph( + s, + h, + g, + head_dim, + rd, + group_size, + dtype, + scale_layout, + data_init=data_init, + seed=args.seed, + ) df = pd.DataFrame(df) aiter.logger.info( "inverse_rope_group_quant summary (markdown):\n%s", diff --git a/op_tests/test_mhc.py b/op_tests/test_mhc.py index ede26320ea..f688b63e8a 100644 --- a/op_tests/test_mhc.py +++ b/op_tests/test_mhc.py @@ -1039,7 +1039,7 @@ def test_mhc_post_pre( choices=[dtypes.d_dtypes["fp16"], dtypes.d_dtypes["bf16"]], nargs="*", metavar="{fp16, bf16}", - default=["bf16"], + default=[dtypes.bf16], help="""Data type. e.g.: -d bf16""", ) diff --git a/op_tests/test_mla_v4_kargpreld.py b/op_tests/test_mla_v4_kargpreld.py index ed0324125b..e88343ecbb 100644 --- a/op_tests/test_mla_v4_kargpreld.py +++ b/op_tests/test_mla_v4_kargpreld.py @@ -39,7 +39,14 @@ import aiter.mla # main no longer auto-imports submodules; need explicit from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx -from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.test_common import ( + DATA_DISTS, + benchmark, + checkAllclose, + fill, + make_generator, + run_perftest, +) torch.set_default_device("cuda") @@ -62,7 +69,6 @@ # Perf iteration counts (kept out of the @benchmark signature so they don't # become table columns). main() overrides these from --iters / --warmup. _PERF = {"num_iters": 2, "num_warmup": 1} -_SEED = 0 # --------------------------------------------------------------------------- @@ -297,6 +303,7 @@ def _build_bf16_inputs( kv_seq_lens=64, q_seq_logical=4, seed=0, + data_init="norm", device="cuda", gqa_ratio=GQA_RATIO, attn_sink=True, @@ -311,19 +318,24 @@ def _build_bf16_inputs( mismatch shows up as an err blowup, not a silent pass. False -> per-head -inf ("no sink" no-op: exp(-inf - max) = 0). """ - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) + gen = make_generator(seed, device=device) total_q = batch * q_seq_logical num_page = batch * (kv_seq_lens // PAGE_SIZE) - q_bf16 = torch.randn( - (total_q, gqa_ratio, _QUANT_D), dtype=dtypes.bf16, device=device - ) - kv_bf16 = torch.randn( - (num_page, PAGE_SIZE, NUM_KV_HEADS, _QUANT_D), + q_bf16 = fill( + (total_q * gqa_ratio, _QUANT_D), + data_init, + gen, dtype=dtypes.bf16, device=device, - ) + ).view(total_q, gqa_ratio, _QUANT_D) + kv_bf16 = fill( + (num_page * PAGE_SIZE * NUM_KV_HEADS, _QUANT_D), + data_init, + gen, + dtype=dtypes.bf16, + device=device, + ).view(num_page, PAGE_SIZE, NUM_KV_HEADS, _QUANT_D) qo_indptr = ( torch.arange(0, batch + 1, dtype=torch.int32, device=device) * q_seq_logical @@ -345,7 +357,16 @@ def _build_bf16_inputs( if attn_sink: # randn*10 so the sink contributes materially (~15%) to the softmax; # well above tolerance, so a dropped/mis-scaled sink is a hard mismatch. - sink = torch.randn(num_heads, dtype=torch.float32, device=device) * 10.0 + sink = ( + fill( + (num_heads,), + data_init, + gen, + dtype=torch.float32, + device=device, + ) + * 10.0 + ) else: sink = torch.full( (num_heads,), float("-inf"), dtype=torch.float32, device=device @@ -377,6 +398,8 @@ def test_mla_v4_nm( num_kv_splits=1, gqa_ratio=GQA_RATIO, attn_sink=True, + data_init="norm", + seed=0, ): """Time each v4 nm kernel candidate, check it against the torch fp8-dequant reference, and return per-candidate `us` / `TFLOPS` / `TB/s` / `err`. @@ -413,7 +436,8 @@ def test_mla_v4_nm( batch=batch, kv_seq_lens=kv_seq_lens, q_seq_logical=q_seq_logical, - seed=_SEED, + seed=seed, + data_init=data_init, gqa_ratio=gqa_ratio, attn_sink=attn_sink, ) @@ -733,13 +757,18 @@ def main(): default=[True], help="attn sink value(s) to sweep. e.g. --attn-sink True False", ) + parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) for Q, KV and attention sink", + ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--iters", type=int, default=50, help="Perf timed iterations") parser.add_argument("--warmup", type=int, default=2, help="Perf warmup iterations") args = parser.parse_args() - global _SEED - _SEED = args.seed _PERF["num_iters"] = args.iters _PERF["num_warmup"] = args.warmup @@ -752,8 +781,20 @@ def main(): ] df = [] - for (nhead, decode_qlen), batch, kv_seq_lens, split_kv, sink in itertools.product( - nhead_combos, args.batch, args.kv_seq_lens, args.split_kv, args.attn_sink + for ( + (nhead, decode_qlen), + batch, + kv_seq_lens, + split_kv, + sink, + data_init, + ) in itertools.product( + nhead_combos, + args.batch, + args.kv_seq_lens, + args.split_kv, + args.attn_sink, + args.data_init, ): try: df.append( @@ -764,6 +805,8 @@ def main(): num_kv_splits=split_kv, gqa_ratio=nhead, attn_sink=sink, + data_init=data_init, + seed=args.seed, ) ) except (RuntimeError, AssertionError) as exc: diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index afce1a712d..c427f8ef51 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -44,7 +44,14 @@ pa_sparse_prefill_fp8_opus, pa_sparse_prefill_opus, ) -from aiter.test_common import benchmark, checkAllclose, perftest +from aiter.test_common import ( + DATA_DISTS, + benchmark, + checkAllclose, + fill, + make_generator, + perftest, +) try: from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse @@ -307,19 +314,32 @@ def _make_inputs( mode: str = "sparse", device: torch.device | str = "cuda", seed: int = 0, + data_init: str = "norm", ) -> dict: assert mode in _MODES - torch.manual_seed(seed) device = torch.device(device) + gen = make_generator(seed, device=device) - q = (torch.randn(n, h, d, device=device, dtype=torch.float32) * 0.5).to(dtype) + q = ( + fill((n * h, d), data_init, gen, dtype=torch.float32, device=device) + .view(n, h, d) + .mul_(0.5) + ).to(dtype) unified_kv = ( - torch.randn(total_pages, d, device=device, dtype=torch.float32) * 0.5 + fill( + (total_pages, d), data_init, gen, dtype=torch.float32, device=device + ) + * 0.5 ).to(dtype) - kv = (torch.randn(total_tokens, d, device=device, dtype=torch.float32) * 0.5).to( - dtype + kv = ( + fill( + (total_tokens, d), data_init, gen, dtype=torch.float32, device=device + ) + * 0.5 + ).to(dtype) + attn_sink = ( + fill((h,), data_init, gen, dtype=torch.float32, device=device) * 0.25 ) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 def _csr(total_rows: int, seed_offset: int): if mode == "sparse": @@ -357,19 +377,37 @@ def _make_inputs_fp8( mode: str = "sparse", device: torch.device | str = "cuda", seed: int = 0, + data_init: str = "norm", ) -> dict: """Returns ``{"kernel": ..., "ref": ...}``: the split fp8/bf16 tensors the kernels take, and the dequantized fp32 rows the reference takes. """ assert mode in _MODES - torch.manual_seed(seed) device = torch.device(device) + gen = make_generator(seed, device=device) def _streams(rows: int): nope_fp8, deq = _quantize_nope( - torch.randn(rows, _FP8_D_NOPE, device=device) * 0.5 + fill( + (rows, _FP8_D_NOPE), + data_init, + gen, + dtype=torch.float32, + device=device, + ) + * 0.5 ) - rope = (torch.randn(rows, _FP8_D_ROPE, device=device) * 0.5).to(torch.bfloat16) + rope = ( + fill( + (rows, _FP8_D_ROPE), + data_init, + gen, + dtype=torch.float32, + device=device, + ) + * 0.5 + ) + rope = rope.to(torch.bfloat16) row_fp32 = torch.cat([deq, rope.to(torch.float32)], dim=1) # [rows, 512] return nope_fp8, rope, row_fp32 @@ -380,7 +418,9 @@ def _streams(rows: int): ukn, ukr, ukv_fp32 = _streams(total_pages) kn, kr, kv_fp32 = _streams(total_tokens) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 + attn_sink = ( + fill((h,), data_init, gen, dtype=torch.float32, device=device) * 0.25 + ) def _csr(total_rows: int, seed_offset: int): if mode == "sparse": @@ -488,6 +528,7 @@ def run_pa_sparse_prefill( mode: str = "sparse", backends: tuple = _BACKENDS, seed: int = 0, + data_init: str = "norm", verify: bool = True, bench: bool = True, ) -> dict | None: @@ -498,7 +539,7 @@ def run_pa_sparse_prefill( softmax_scale = 1.0 / math.sqrt(d) msg = ( f"[N={n} H={h} D={d} total_pages={total_pages} total_tokens={total_tokens} " - f"prec={prec} mode={mode}]" + f"prec={prec} mode={mode} data_init={data_init} seed={seed}]" ) wanted = [b for b in _PREC_BACKENDS[prec] if b in backends] @@ -506,7 +547,15 @@ def run_pa_sparse_prefill( candidates: list = [] if prec == "fp8": - data = _make_inputs_fp8(n, h, total_pages, total_tokens, mode=mode, seed=seed) + data = _make_inputs_fp8( + n, + h, + total_pages, + total_tokens, + mode=mode, + seed=seed, + data_init=data_init, + ) kernel_inputs = data["kernel"] ref_fn, ref_inputs = _ref_pa_sparse_prefill_fp8, data["ref"] if "opus" in wanted: @@ -537,6 +586,7 @@ def run_pa_sparse_prefill( _PREC_TO_DTYPE[prec], mode=mode, seed=seed, + data_init=data_init, ) ref_fn, ref_inputs = _ref_pa_sparse_prefill_opus, kernel_inputs if "opus" in wanted: @@ -744,6 +794,13 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): default=0, help="RNG seed for input + CSR generation", ) +parser.add_argument( + "--data-init", + nargs="+", + choices=list(DATA_DISTS), + default=["norm"], + help="DATA initialization distribution(s) for Q, KV and attention sink", +) if __name__ == "__main__": @@ -751,12 +808,13 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): rows = [] # product varies its last argument fastest -> this is also the row order. - for prec, mode, h, n, pages_arg in itertools.product( + for prec, mode, h, n, pages_arg, data_init in itertools.product( args.prec, args.mode, args.h_q, args.n_tokens, args.total_pages, + args.data_init, ): total_pages = pages_arg if pages_arg > 0 else n # 0 is "mirror -n" total_tokens = args.total_tokens if args.total_tokens is not None else n @@ -770,6 +828,7 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): mode=mode, backends=tuple(args.backend), seed=args.seed, + data_init=data_init, verify=not args.no_verify, bench=not args.no_bench, ) @@ -784,7 +843,9 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): if drop_cols: df = df.drop(columns=drop_cols) # Column order otherwise follows whichever row first ran a backend. - lead = [c for c in ("prec", "mode", "h", "n") if c in df.columns] + lead = [ + c for c in ("prec", "mode", "data_init", "h", "n") if c in df.columns + ] rest = [c for c in df.columns if c not in lead] metrics = [c for b in _BACKENDS for c in rest if c.startswith(f"{b} ")] df = df[lead + [c for c in rest if c not in metrics] + metrics] diff --git a/op_tests/triton_tests/attention/test_mla_v4_triton.py b/op_tests/triton_tests/attention/test_mla_v4_triton.py index d9c99829d7..e4033afd66 100644 --- a/op_tests/triton_tests/attention/test_mla_v4_triton.py +++ b/op_tests/triton_tests/attention/test_mla_v4_triton.py @@ -9,7 +9,13 @@ 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 +from aiter.test_common import ( + benchmark, + checkAllclose, + fill, + make_generator, + run_perftest, +) # MLA v4 sparse-decode parity: D=512 heads, page_size=1 unified pool. _PA_DECODE_SPARSE_D = 512 @@ -99,15 +105,25 @@ def _make_inputs( total_pages: int, dtype=torch.bfloat16, seed: int = 0, + data_init: str = "norm", include_sentinels: bool = False, variable_len: bool = False, ): torch.manual_seed(seed) device = torch.device("cuda") + gen = make_generator(seed) - 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 + q = ( + fill((T * H, D), data_init, gen, dtype=dtype, device=device) + .view(T, H, D) + .mul_(0.5) + ) + unified_kv = ( + fill((total_pages, D), data_init, gen, dtype=dtype, device=device) * 0.5 + ) + attn_sink = ( + fill((H,), data_init, gen, dtype=torch.float32, device=device) * 0.1 + ) # Per-token kv_len: fixed or random in [1, kv_len_per_token]. if variable_len: @@ -117,6 +133,7 @@ def _make_inputs( size=(T,), device=device, dtype=torch.int64, + generator=gen, ) else: kv_lens = torch.full((T,), kv_len_per_token, device=device, dtype=torch.int64) @@ -131,11 +148,14 @@ def _make_inputs( size=(total_indices,), device=device, dtype=torch.int32, + generator=gen, ) 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] + sentinel_pos = torch.randperm(total_indices, device=device, generator=gen)[ + :n_sentinel + ] indices[sentinel_pos] = -1 indptr = indptr.to(torch.int32) @@ -144,14 +164,28 @@ def _make_inputs( @benchmark() -def test_mla_v4_triton_staged(gqa_ratio, batch, kv_seq_lens, num_kv_splits): +def test_mla_v4_triton_staged( + gqa_ratio, + batch, + kv_seq_lens, + num_kv_splits, + data_init="norm", + seed=0, +): """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 + T, + H, + D, + kv_seq_lens, + pages, + variable_len=False, + data_init=data_init, + seed=seed, ) pa_kwargs = { "has_invalid": False, @@ -195,7 +229,14 @@ def test_mla_v4_triton_staged(gqa_ratio, batch, kv_seq_lens, num_kv_splits): @benchmark() -def test_mla_v4_triton_perf(gqa_ratio, batch, kv_seq_lens, num_kv_splits): +def test_mla_v4_triton_perf( + gqa_ratio, + batch, + kv_seq_lens, + num_kv_splits, + data_init="norm", + seed=0, +): """Perf sweep row for combo bench / gfx1250 Triton sparse MLA v4 decode. Shape ids mirror ``test_mla_v4_kargpreld.test_mla_v4_nm``: @@ -207,7 +248,14 @@ def test_mla_v4_triton_perf(gqa_ratio, batch, kv_seq_lens, num_kv_splits): 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 + T, + H, + D, + kv_seq_lens, + pages, + variable_len=False, + data_init=data_init, + seed=seed, ) _, us = run_perftest( pa_decode_sparse, From 3e62f8462be6ab15015dbc749c8f0cfa7cd91699 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Fri, 4 Sep 2026 10:54:49 +0800 Subject: [PATCH 2/3] perf(gfx1250): report per-case SMI telemetry --- aiter/test_common.py | 48 +++ op_tests/bench_gfx1250_combo.py | 686 ++++++++++++++++++-------------- op_tests/smi_monitor.py | 76 ++++ 3 files changed, 520 insertions(+), 290 deletions(-) diff --git a/aiter/test_common.py b/aiter/test_common.py index 4f523e23e2..929dc38291 100644 --- a/aiter/test_common.py +++ b/aiter/test_common.py @@ -12,6 +12,7 @@ from aiter import logger pd.set_option("display.max_rows", 200) +_SMI_LABEL_COUNTS = {} ## debug ## # pd.set_option("display.max_rows", None) # pd.set_option("display.max_columns", None) @@ -124,6 +125,53 @@ def wrapper(*args, **kwargs): avg = get_trace_perf(prof, num_iters) logger.info(f"avg: {avg} us/iter with hipgraph") + if os.environ.get("AITER_SMI_MONITOR", "0") == "1": + fn_name = getattr(func, "__name__", "kernel") + skipped = { + name.strip() + for name in os.environ.get("AITER_SMI_SKIP_FUNCTIONS", "").split(",") + if name.strip() + } + if fn_name in skipped: + return data, avg + # Import lazily: normal library/test use has no amdsmi dependency. + # Combo and its child UTs run from the repository root, where the + # standalone op_tests monitor module is importable. + try: + from op_tests.smi_monitor import replay_with_smi + except ModuleNotFoundError: + # Direct ``python op_tests/foo.py`` puts op_tests itself, + # rather than the repository root, at sys.path[0]. + from smi_monitor import replay_with_smi + + if testGraph: + replay = graph.replay + # One replay contains num_iters calls captured above. + replay_us = avg * num_iters + else: + replay_index = 0 + + def replay(): + nonlocal replay_index + replay_args, replay_kwargs = rotate_args[ + replay_index % len(rotate_args) + ] + replay_index += 1 + return func(*replay_args, **replay_kwargs) + + replay_us = avg + + case_label = os.environ.get("AITER_SMI_LABEL", "benchmark_case") + label_key = (case_label, fn_name) + occurrence = _SMI_LABEL_COUNTS.get(label_key, 0) + 1 + _SMI_LABEL_COUNTS[label_key] = occurrence + replay_with_smi( + replay, + label=f"{case_label}/{fn_name}#{occurrence}", + synchronize=torch.cuda.synchronize, + estimated_us=replay_us, + ) + return data, avg return wrapper diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index be1c13df55..eff973c894 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -6,11 +6,11 @@ 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. +Output discipline: combo-owned pandas summaries are printed as record-oriented +JSON. Existing child-UT summaries are extracted without changing those UTs. 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. Run from the aiter repo root so `op_tests/` siblings import cleanly: @@ -81,15 +81,18 @@ lose it outright -- see _pin_arch. Set either yourself and yours wins. -Optional GPU telemetry wraps the whole selected-op sweep, never an individual -kernel timing region: +Optional GPU telemetry replays each already-prepared benchmark case in its own +sampling window, after its normal latency measurement: python op_tests/bench_gfx1250_combo.py --dsv4 \ - --smi-monitor --smi-device 0 --smi-interval 0.05 + --smi-monitor --smi-device 0 --smi-interval 0.05 --smi-duration 1.0 -The monitor uses the Python ``amdsmi`` package and prints min/mean/median/max -for the collected clocks, power, temperature, activity and VRAM metrics after -the sweep. +Input initialization, compilation, correctness and warmup are outside the SMI +window. The monitor uses the Python ``amdsmi`` package and prints a case-tagged +min/mean/median/max table for clocks, power, temperature, activity and VRAM. +The single-GPU replay hook is intentionally disabled for ``mega_moe`` and +``mori_ep``; their multi-rank communication loops need a separate rank/device +telemetry design. Supported operator inputs can be overridden consistently with: @@ -254,12 +257,13 @@ import argparse import contextlib import itertools +import json import subprocess import sys import tempfile import warnings -from smi_monitor import monitor_gpu +from smi_monitor import SMI_RESULT_PREFIX warnings.filterwarnings("ignore") @@ -316,6 +320,31 @@ def _silence(): ) SUPPORTED_GFX = ["gfx1250"] +_SMI_ROWS = [] + + +@contextlib.contextmanager +def _smi_case(label): + """Set the case label consumed by the common perftest hook.""" + old = os.environ.get("AITER_SMI_LABEL") + if os.environ.get("AITER_SMI_MONITOR") == "1": + os.environ["AITER_SMI_LABEL"] = label + try: + yield + finally: + if old is None: + os.environ.pop("AITER_SMI_LABEL", None) + else: + os.environ["AITER_SMI_LABEL"] = old + + +def _without_smi(env): + """Return a child environment with single-GPU telemetry disabled.""" + clean = env.copy() + for key in tuple(clean): + if key.startswith("AITER_SMI_"): + clean.pop(key) + return clean # 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) @@ -487,75 +516,6 @@ def _int_quad(s): 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 = [ @@ -593,13 +553,10 @@ def _moe_stage_bytes( "pass", "gemm1_us", "gemm1 TFLOPS", - "gemm1 TB/s", + "gemm1 GB/s", "gemm2_us", "gemm2 TFLOPS", - "gemm2 TB/s", - "total us", - "total TFLOPS", - "total TB/s", + "gemm2 GB/s", "kernel", ] # Fixed kernel-bench config (mirrors test_flydsl_grouped_gemm_gfx1250.py --scenario kernel). @@ -726,9 +683,36 @@ def _capture(): os.close(old2) tmp.seek(0) box.append(tmp.read()) + _collect_smi_rows(box[0].splitlines()) + + +def _collect_smi_rows(lines): + """Collect structured per-case SMI records emitted by this or a child UT.""" + for line in lines: + marker = line.find(SMI_RESULT_PREFIX) + if marker < 0: + continue + try: + record = json.loads(line[marker + len(SMI_RESULT_PREFIX) :]) + except json.JSONDecodeError: + continue + base = { + "case": record.get("label"), + "device": record.get("device"), + "duration_s": round(record.get("duration_s", 0.0), 3), + "launches": record.get("launches"), + "samples": record.get("samples"), + "sample_status": record.get("sample_status"), + } + metrics = record.get("metrics", {}) + if not metrics: + _SMI_ROWS.append({**base, "metric": "(no metrics)"}) + for metric, stats in metrics.items(): + _SMI_ROWS.append({**base, "metric": metric, **stats}) def _print_table(name, rows, keep=None): + """Print one named DataFrame as a JSON object with record-oriented rows.""" 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`. @@ -740,8 +724,8 @@ def _print_table(name, rows, keep=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)) + records = json.loads(df.to_json(orient="records")) + print(json.dumps({"name": name, "rows": records}, indent=2), flush=True) # Compiler / logger / IR-dump chatter the child UTs interleave with results. @@ -1044,7 +1028,7 @@ def _pin_arch(env): def _run_child(name, cmd, cwd, env=None, extract=None, timeout=None, tail=30, - kernels=True): + kernels=True, smi=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 @@ -1053,6 +1037,15 @@ def _run_child(name, cmd, cwd, env=None, extract=None, timeout=None, tail=30, when the child fails or emits nothing recognisable. """ extract = extract or _DEFAULT_EXTRACT + # Give every child invocation its combo-owned SMI case label. UTs remain + # unaware of telemetry; the common perftest hook reads this environment. + if smi and os.environ.get("AITER_SMI_MONITOR") == "1": + env = os.environ.copy() if env is None else env.copy() + env["AITER_SMI_LABEL"] = name + old_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + f"{cwd}{os.pathsep}{old_pythonpath}" if old_pythonpath else str(cwd) + ) # env=None means "inherit ours", which already carries these two. if env is not None: _pin_arch(env) @@ -1069,6 +1062,7 @@ def _run_child(name, cmd, cwd, env=None, extract=None, timeout=None, tail=30, _note_failure(name, f"timed out after {timeout}s") return lines = proc.stdout.splitlines() + _collect_smi_rows(lines) # `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. @@ -1109,12 +1103,17 @@ def run_mha(args): with _silence(): for init in inits: 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 + with _smi_case( + f"mha/batch=1/hq=64/hk={8 if head_dim == 64 else 4}/" + f"sq={seqlen}/sk={seqlen}/d={head_dim}/causal={int(causal)}/" + f"data={init}" + ): + 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" @@ -1130,7 +1129,12 @@ def run_moe(args): for tokens, fmt, (data_init, scale_init) in itertools.product( cfg["tokens"], data_formats, init_pairs ): - with _capture() as box: + label = ( + f"moe/fmt={fmt}/tokens={tokens}/experts={cfg['experts']}/" + f"topk={cfg['topk']}/hd={cfg['model_dim']}/id={cfg['inter_dim']}/" + f"data={data_init}/scale={scale_init}/seed={args.seed}" + ) + with _smi_case(label), _capture() as box: moe_mod.set_data_format(fmt) metrics = moe_mod.run_moe( fmt, @@ -1148,32 +1152,7 @@ def run_moe(args): 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, @@ -1188,14 +1167,11 @@ def run_moe(args): "seed": args.seed, "pass": metrics["passed"], "gemm1_us": us1, - "gemm1 TFLOPS": _tflops(flop1, us1), - "gemm1 TB/s": bw1, + "gemm1 TFLOPS": metrics.get("gemm1_tflops"), + "gemm1 GB/s": metrics.get("gemm1_bandwidth_gbs"), "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, + "gemm2 TFLOPS": metrics.get("gemm2_tflops"), + "gemm2 GB/s": metrics.get("gemm2_bandwidth_gbs"), "kernel": " + ".join(_kernel_names(box[0].splitlines())) or None, } ) @@ -1216,20 +1192,24 @@ def run_gemm(args): ["mxfp4", "nvfp4"], ["bf16", "fp8"], ): - rows.append( - gemm_mod.test_gemm( - intype, - M, - N, - K, - 1, - outtype, - di, - si, - seed=args.seed, - mode="perf", + with _smi_case( + f"gemm_a4w4/intype={intype}/out={outtype}/M={M}/N={N}/K={K}/" + f"data={di}/scale={si}/seed={args.seed}" + ): + rows.append( + gemm_mod.test_gemm( + intype, + M, + N, + K, + 1, + outtype, + di, + si, + seed=args.seed, + mode="perf", + ) ) - ) _print_table("gemm_a4w4 (perf)", rows, keep=_GEMM_KEEP) @@ -1250,17 +1230,21 @@ def run_f8gemm(args): 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, - seed=args.seed, - mode="perf", - ) + with _smi_case( + f"mxfp8fp4gemm/intype={intype}/M={M}/N={N}/K={K}/" + f"data={di}/scale={si}/seed={args.seed}" + ): + row = f8gemm_mod.test_gemm( + intype, + M, + N, + K, + 1, + data_init=di, + scale_init=si, + seed=args.seed, + mode="perf", + ) if row is not None: row["workload"] = workload rows.append(row) @@ -1275,42 +1259,61 @@ def run_a8w8_blockscale(args): # 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", - *( - ["--data-init", *args.data_init] - if args.data_init is not None - else [] - ), - *( - ["--scale-init", *args.scale_init] - if args.scale_init is not None - else [] - ), - "--seed", - str(args.seed), - ], - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - env=env, + nk_shapes = ( + (2048, 7168), + (7168, 16384), + (6144, 7168), + (7168, 3072), + (65536, 1536), + (8192, 1536), + ) + + def run_case(tokens, shapes, init_pairs, label): + _run_child( + label, + [ + sys.executable, + "op_tests/test_gemm_a8w8_blockscale.py", + "-m", + *map(str, tokens), + "-nk", + *(f"{n},{k}" for n, k in shapes), + "--ck_preshuffle", + "True", + "--flydsl", + "--data-init", + *(data for data, _ in init_pairs), + "--scale-init", + *(scale for _, scale in init_pairs), + "--seed", + str(args.seed), + ], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + env=env, + ) + + init_pairs = _init_pairs( + args, defaults=(("constant", "constant"), ("uniform", "auto")) ) + if args.smi_monitor: + for m, (n, k), pair in itertools.product( + _A8W8_BLOCKSCALE_TOKENS, nk_shapes, init_pairs + ): + data_init, scale_init = pair + run_case( + (m,), + ((n, k),), + (pair,), + f"a8w8_blockscale/M={m}/N={n}/K={k}/data={data_init}/" + f"scale={scale_init}/seed={args.seed}", + ) + else: + run_case( + _A8W8_BLOCKSCALE_TOKENS, + nk_shapes, + init_pairs, + "gemm_a8w8_blockscale (DSv4)", + ) def run_a16w16(args): @@ -1346,7 +1349,10 @@ def run_a16w16(args): # skipping shapes that tuning has already made runnable. Let the kernel # raise and record that instead. try: - with _capture() as box: + with _smi_case( + f"a16w16/batch={batch}/M={M}/N={n}/K={K}/" + f"data={data_init}/seed={args.seed}" + ), _capture() as box: err = a16w16_mod.test_a16w16( batch=batch, M=M, @@ -1394,7 +1400,7 @@ def run_mega_moe(args): # 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() + env = _without_smi(os.environ) # No MORI_SHMEM_HEAP_SIZE default here, for two independent reasons. # # Raising it sweep-wide took the machine down: the heap is preallocated per @@ -1462,38 +1468,48 @@ def run_mega_moe(args): cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), env={**env, "AITER_FORCE_A8W4": force_a8w4}, extract=_md_kernel_table, - kernels=False, + kernels=False, + smi=False, ) def run_mhc(args): """Run the DSv4 mHC fused-RMSNorm benchmark at M=512, N=7168.""" _unused_scale_init(args, "mhc") - _run_child( - "mhc (DSv4, fused RMSNorm)", - [ - sys.executable, - "op_tests/test_mhc.py", - "-n", - "7168", - "-m", - *map(str, _TOKENS), - "--fuse_rmsnorm", - *( - ["--data-init", *args.data_init] - if args.data_init is not None - else [] + def run_case(tokens, data_inits, label): + _run_child( + label, + [ + sys.executable, + "op_tests/test_mhc.py", + "-n", + "7168", + "-m", + *map(str, tokens), + "--fuse_rmsnorm", + "--data-init", + *data_inits, + "--seed", + str(args.seed), + ], + 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"), ), - "--seed", - str(args.seed), - ], - 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"), - ), - ) + ) + + data_inits = args.data_init or ["norm"] + if args.smi_monitor: + for m, data_init in itertools.product(_TOKENS, data_inits): + run_case( + (m,), + (data_init,), + f"mhc/M={m}/N=7168/fuse_rmsnorm=1/data={data_init}/seed={args.seed}", + ) + else: + run_case(_TOKENS, data_inits, "mhc (DSv4, fused RMSNorm)") def run_qk_norm(args): @@ -1512,14 +1528,18 @@ def run_qk_norm(args): ] repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) _unused_scale_init(args, "qk_norm") - for data_init in args.data_init or [None]: + data_inits = args.data_init or [None] + + def run_case(tokens, data_init): qk_init = "normal" if data_init == "norm" else data_init _run_child( - f"qk_norm (init={qk_init or 'native-default'}, seed={args.seed})", + f"qk_norm/T={','.join(map(str, tokens))}/H=128/D=512/RD=64/" + f"qweight=both/swa=direct,paged/init={qk_init or 'native-default'}/" + f"seed={args.seed}", [ *base_cmd, "-T", - *map(str, _TOKENS), + *map(str, tokens), *(["--init", qk_init] if qk_init else []), "--seed", str(args.seed), @@ -1531,6 +1551,13 @@ def run_qk_norm(args): ), ) + if args.smi_monitor: + for token, data_init in itertools.product(_TOKENS, data_inits): + run_case((token,), data_init) + else: + for data_init in data_inits: + run_case(_TOKENS, data_init) + def run_score_qk(args): """Run DSv4 decode score-QK at batch 512 for short and long CSA KV.""" @@ -1583,7 +1610,7 @@ def run_mori_ep(args): # 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 = _without_smi(os.environ) env["PYTHONPATH"] = f"{mori}/python:{mori}" env["MORI_SOCKET_IFNAME"] = "lo" env["GLOO_SOCKET_IFNAME"] = "lo" @@ -1627,6 +1654,7 @@ def run_mori_ep(args): env=env, extract=_lines(_quiet), timeout=3600, + smi=False, ) @@ -1770,24 +1798,28 @@ def run_mla_v4_decode(args): "num_kv_splits": split_kv, } try: - asm = _bench_mla_v4_asm_staged( - gqa, - batch, - ctx, - split_kv, - iters, - warmup, - data_init, - args.seed, - ) - tri = mla_v4_triton_mod.test_mla_v4_triton_staged( - gqa_ratio=gqa, - batch=batch, - kv_seq_lens=ctx, - num_kv_splits=split_kv, - data_init=data_init, - seed=args.seed, - ) + with _smi_case( + f"mla_v4_decode/gqa={gqa}/batch={batch}/ctx={ctx}/" + f"split={split_kv}/data={data_init}/seed={args.seed}" + ): + asm = _bench_mla_v4_asm_staged( + gqa, + batch, + ctx, + split_kv, + iters, + warmup, + data_init, + args.seed, + ) + tri = mla_v4_triton_mod.test_mla_v4_triton_staged( + gqa_ratio=gqa, + batch=batch, + kv_seq_lens=ctx, + num_kv_splits=split_kv, + data_init=data_init, + seed=args.seed, + ) row.update(asm) row.update(tri) row["s1 triton/asm"] = _perf_ratio(row["triton_s1"], row["asm_s1"]) @@ -1813,36 +1845,53 @@ def run_inverse_rope(args): # -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", - *( - ["--data-init", *args.data_init] - if args.data_init is not None - else [] - ), - "--seed", - str(args.seed), - ], - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - ) + def run_case(tokens, data_inits, label): + _run_child( + label, + [ + sys.executable, + "op_tests/test_inverse_rope_group_quant.py", + "-b", + "128,16", + "-s", + *map(str, tokens), + "-l", + "n32k4", + "--group-size", + "32", + "--data-init", + *data_inits, + "--seed", + str(args.seed), + ], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + + data_inits = args.data_init or ["norm"] + if args.smi_monitor: + for tokens, data_init in itertools.product(_INVERSE_ROPE_TOKENS, data_inits): + run_case( + (tokens,), + (data_init,), + f"inverse_rope/s={tokens}/heads=128/groups=16/layout=n32k4/" + f"group_size=32/data={data_init}/seed={args.seed}", + ) + else: + run_case( + _INVERSE_ROPE_TOKENS, + data_inits, + "inverse_rope_group_quant (DSv4, tp1)", + ) def run_mla_v4_prefill(args): """Run DSv4 prefill across two precisions, pools and CSR modes.""" _unused_scale_init(args, "mla_v4_prefill") - for tokens in _MLA_PREFILL_TOKENS: + data_inits = args.data_init or ["norm"] + + def run_case(tokens, pages, precs, modes, backends, init_values, label): _run_child( - f"mla_v4 prefill (M={tokens}, prec=fp8/bf16, pages=4096/16384)", + label, [ sys.executable, "op_tests/test_pa_sparse_prefill.py", @@ -1853,27 +1902,23 @@ def run_mla_v4_prefill(args): "-d", "512", "--total_pages", - "4096", - "16384", + *map(str, pages), "--total_tokens", str(tokens), "--prec", - "fp8", - "bf16", + *precs, # 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", + *modes, + "--backend", + *backends, "--no-verify", "--seed", str(args.seed), - *( - ["--data-init", *args.data_init] - if args.data_init is not None - else [] - ), + "--data-init", + *init_values, ], cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), # Not _table_row: the UT has no "latency_us" column (it prints @@ -1883,6 +1928,39 @@ def run_mla_v4_prefill(args): extract=_space_table("total_pages"), ) + if args.smi_monitor: + backend_by_prec = {"fp8": ("opus", "asm"), "bf16": ("opus", "triton")} + for tokens, pages, prec, mode, data_init in itertools.product( + _MLA_PREFILL_TOKENS, + (4096, 16384), + ("fp8", "bf16"), + ("dense", "sparse"), + data_inits, + ): + for backend in backend_by_prec[prec]: + run_case( + tokens, + (pages,), + (prec,), + (mode,), + (backend,), + (data_init,), + f"mla_v4_prefill/M={tokens}/H=128/D=512/pages={pages}/" + f"total_tokens={tokens}/prec={prec}/mode={mode}/backend={backend}/" + f"data={data_init}/seed={args.seed}", + ) + else: + for tokens in _MLA_PREFILL_TOKENS: + run_case( + tokens, + (4096, 16384), + ("fp8", "bf16"), + ("dense", "sparse"), + ("opus", "asm", "triton"), + data_inits, + f"mla_v4 prefill (M={tokens}, prec=fp8/bf16, pages=4096/16384)", + ) + OPS = { "mha": run_mha, @@ -2042,7 +2120,7 @@ def main(): p.add_argument( "--smi-monitor", action="store_true", - help="sample power/clocks/temperature/utilization for the whole sweep", + help="replay and sample each timed benchmark case after latency measurement", ) p.add_argument( "--smi-device", @@ -2056,6 +2134,12 @@ def main(): default=0.05, help="amdsmi sampling interval in seconds (default: 0.05)", ) + p.add_argument( + "--smi-duration", + type=float, + default=1.0, + help="minimum replay window in seconds for each benchmark case (default: 1.0)", + ) # mha (SWA fwd asm) — fixed 4-shape grid; init sweep only p.add_argument( "--mha-init", @@ -2092,6 +2176,29 @@ def main(): p.error("--smi-device must be non-negative") if args.smi_interval <= 0: p.error("--smi-interval must be positive") + if args.smi_duration <= 0: + p.error("--smi-duration must be positive") + + if args.smi_monitor: + smi_output = tempfile.NamedTemporaryFile( + prefix="aiter_smi_", suffix=".jsonl", delete=False + ) + smi_output.close() + os.environ.update( + { + "AITER_SMI_MONITOR": "1", + "AITER_SMI_DEVICE": str(args.smi_device), + "AITER_SMI_INTERVAL": str(args.smi_interval), + "AITER_SMI_DURATION": str(args.smi_duration), + "AITER_SMI_OUTPUT_PATH": smi_output.name, + # Reference implementations timed by a few legacy UTs are not + # hardware candidates and must not produce telemetry rows. + "AITER_SMI_SKIP_FUNCTIONS": "run_torch,run_torch2", + } + ) + else: + smi_output = None + os.environ.pop("AITER_SMI_MONITOR", None) args.suite = "dsv4" if args.dsv4 else "perf" default_ops = DSV4_OPS if args.dsv4 else PERF_OPS @@ -2100,26 +2207,25 @@ def main(): # runnable by name to check whether a newer image fixed it. argparse already # rejects names outside OPS. selected_ops = args.ops or default_ops - monitor_context = ( - monitor_gpu(device_index=args.smi_device, interval_s=args.smi_interval) - if args.smi_monitor - else contextlib.nullcontext(None) - ) - with monitor_context as monitor: - for name in selected_ops: - with _keep_going(name): - OPS[name](args) - - if monitor is not None: - rows = [ - {"metric": metric, **stats} - for metric, stats in monitor.summary().items() - ] + for name in selected_ops: + with _keep_going(name): + OPS[name](args) + + if args.smi_monitor: + try: + with open(smi_output.name, encoding="utf-8") as output: + _collect_smi_rows(output) + finally: + os.unlink(smi_output.name) _print_table( - f"amdsmi (device={args.smi_device}, interval={args.smi_interval}s, " - f"samples={len(monitor.samples)})", - rows, - keep=["metric", "min", "mean", "median", "max", "n"], + f"amdsmi per benchmark case (device={args.smi_device}, " + f"interval={args.smi_interval}s, min_duration={args.smi_duration}s)", + _SMI_ROWS, + keep=[ + "case", "device", "duration_s", "launches", "samples", + "sample_status", + "metric", "min", "mean", "median", "max", "n", + ], ) if _FAILURES: diff --git a/op_tests/smi_monitor.py b/op_tests/smi_monitor.py index 9d51b9b55d..6cbb82a4ef 100644 --- a/op_tests/smi_monitor.py +++ b/op_tests/smi_monitor.py @@ -19,11 +19,15 @@ from __future__ import annotations import ctypes +import json +import os import threading import time from contextlib import contextmanager from typing import Generator +SMI_RESULT_PREFIX = "AITER_SMI_RESULT " + try: import amdsmi @@ -262,3 +266,75 @@ def monitor_gpu( mon = GpuMonitor(device_index=device_index, interval_s=interval_s) with mon: yield mon + + +def smi_replay_enabled() -> bool: + """Whether the benchmark requested an isolated SMI replay window.""" + return os.environ.get("AITER_SMI_MONITOR", "0") == "1" + + +def replay_with_smi( + fn, + *, + label: str, + synchronize, + estimated_us: float | None = None, +) -> dict | None: + """Repeat one already-prepared benchmark case under the GPU monitor. + + Input creation, compilation, correctness and the latency measurement happen + before this function is called. Batching launches between synchronizations + keeps short kernels busy while still checking the wall-clock deadline often + enough for slow kernels. + """ + if not smi_replay_enabled(): + return None + + device = int(os.environ.get("AITER_SMI_DEVICE", "0")) + interval_s = float(os.environ.get("AITER_SMI_INTERVAL", "0.05")) + duration_s = float(os.environ.get("AITER_SMI_DURATION", "1.0")) + if interval_s <= 0 or duration_s <= 0: + raise ValueError("AITER_SMI_INTERVAL and AITER_SMI_DURATION must be positive") + + # Aim for roughly one synchronization per monitor tick. The cap prevents a + # near-zero/invalid latency estimate from enqueueing an unbounded amount of + # work, while a slow case is synchronized after every launch. + if estimated_us is not None and estimated_us > 0: + batch_iters = max(1, min(1024, int(interval_s * 1e6 / estimated_us))) + else: + batch_iters = 1 + + synchronize() + launches = 0 + start = time.perf_counter() + with monitor_gpu(device_index=device, interval_s=interval_s) as monitor: + while launches == 0 or time.perf_counter() - start < duration_s: + for _ in range(batch_iters): + fn() + launches += batch_iters + synchronize() + elapsed_s = time.perf_counter() - start + + result = { + "label": label, + "device": device, + "interval_s": interval_s, + "duration_s": elapsed_s, + "launches": launches, + "samples": len(monitor.samples), + "metrics": monitor.summary(), + } + expected_samples = max(1, int(duration_s / interval_s)) + result["sample_status"] = ( + "ok" if len(monitor.samples) >= max(2, expected_samples // 2) else "insufficient" + ) + # A shared JSONL sink survives fd silencing and child processes. Standalone + # UT runs without a sink still get a machine-readable stdout record. + line = SMI_RESULT_PREFIX + json.dumps(result, sort_keys=True) + output_path = os.environ.get("AITER_SMI_OUTPUT_PATH") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(line + "\n") + else: + print(line, flush=True) + return result From f5d5c026e461390e779ea4f702c44dfc76f45e7d Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Fri, 4 Sep 2026 11:31:23 +0800 Subject: [PATCH 3/3] perf(gfx1250): monitor Mega MoE per rank --- op_tests/bench_gfx1250_combo.py | 12 +-- .../multigpu_tests/test_mega_moe_gfx1250.py | 73 ++++++++++++++++++ op_tests/smi_monitor.py | 77 +++++++++++++++---- 3 files changed, 139 insertions(+), 23 deletions(-) diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index eff973c894..182f4af9f4 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -90,9 +90,9 @@ Input initialization, compilation, correctness and warmup are outside the SMI window. The monitor uses the Python ``amdsmi`` package and prints a case-tagged min/mean/median/max table for clocks, power, temperature, activity and VRAM. -The single-GPU replay hook is intentionally disabled for ``mega_moe`` and -``mori_ep``; their multi-rank communication loops need a separate rank/device -telemetry design. +``mega_moe`` has every rank monitor its local GPU around synchronized graph +replays, then gathers the four summaries to rank 0; ``mori_ep`` remains disabled +until its dispatch/combine loop exposes an aligned telemetry window. Supported operator inputs can be overridden consistently with: @@ -698,6 +698,7 @@ def _collect_smi_rows(lines): continue base = { "case": record.get("label"), + "rank": record.get("rank"), "device": record.get("device"), "duration_s": round(record.get("duration_s", 0.0), 3), "launches": record.get("launches"), @@ -1400,7 +1401,7 @@ def run_mega_moe(args): # 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 = _without_smi(os.environ) + 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 @@ -1469,7 +1470,6 @@ def run_mega_moe(args): env={**env, "AITER_FORCE_A8W4": force_a8w4}, extract=_md_kernel_table, kernels=False, - smi=False, ) @@ -2222,7 +2222,7 @@ def main(): f"interval={args.smi_interval}s, min_duration={args.smi_duration}s)", _SMI_ROWS, keep=[ - "case", "device", "duration_s", "launches", "samples", + "case", "rank", "device", "duration_s", "launches", "samples", "sample_status", "metric", "min", "mean", "median", "max", "n", ], diff --git a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py index c09afe3eec..47994b7db3 100644 --- a/op_tests/multigpu_tests/test_mega_moe_gfx1250.py +++ b/op_tests/multigpu_tests/test_mega_moe_gfx1250.py @@ -37,7 +37,9 @@ """ import argparse +import math import os +import time import torch import torch.distributed as dist @@ -786,6 +788,76 @@ def _event_device_us(e): return 0.0 +def _run_distributed_smi_replay(pipe, dist_ctx, median_us, n_layers): + """Replay the Mega graph while every rank monitors its local GPU.""" + if os.environ.get("AITER_SMI_MONITOR", "0") != "1": + return + + from op_tests.smi_monitor import GpuMonitor, emit_smi_result + + interval_s = float(os.environ.get("AITER_SMI_INTERVAL", "0.05")) + duration_s = float(os.environ.get("AITER_SMI_DURATION", "1.0")) + if interval_s <= 0 or duration_s <= 0 or median_us <= 0: + raise ValueError( + "Mega MoE SMI interval, duration and measured median must be positive" + ) + replay_count = max(1, math.ceil(duration_s * 1e6 / median_us)) + + monitor = None + monitor_error = None + try: + monitor = GpuMonitor( + device_index=torch.cuda.current_device(), interval_s=interval_s + ) + monitor.start() + except Exception as error: # noqa: BLE001 - propagate to every rank + monitor_error = f"rank {dist_ctx.rank}: {type(error).__name__}: {error}" + + monitor_errors = dist_ctx.gather_objects(monitor_error) + failed = [error for error in monitor_errors if error is not None] + if failed: + if monitor is not None: + monitor.stop() + raise RuntimeError("Mega MoE SMI monitor failed to start: " + "; ".join(failed)) + + # Gloo barriers align CPU submission without adding a GPU collective to the + # measured Mega graph window. Every rank executes exactly the same replay + # count; a local duration loop would diverge and deadlock the collectives. + dist.barrier() + window_start = time.perf_counter() + for _ in range(replay_count): + pipe.graph.replay() + torch.cuda.synchronize() + window_end = time.perf_counter() + monitor.stop() + dist.barrier() + + samples = [ + sample + for sample in monitor.samples + if window_start <= sample["timestamp_s"] <= window_end + ] + expected_samples = max(1, int(duration_s / interval_s)) + base_label = os.environ.get("AITER_SMI_LABEL", "mega_moe") + local_result = { + "label": f"{base_label}/mega_graph_{n_layers}_layers", + "device": dist_ctx.local_rank, + "rank": dist_ctx.rank, + "interval_s": interval_s, + "duration_s": window_end - window_start, + "launches": replay_count, + "samples": len(samples), + "sample_status": ( + "ok" if len(samples) >= max(2, expected_samples // 2) else "insufficient" + ), + "metrics": monitor.summary(start_s=window_start, end_s=window_end), + } + results = dist_ctx.gather_objects(local_result) + if dist_ctx.rank == 0: + for result in results: + emit_smi_result(result) + + def _aggregate_prof_table(prof, dist_ctx, per_layer_denom=1.0, row_limit=200): """Collect the torch.profiler per-kernel table ACROSS ranks (collective; call on every rank). Each rank contributes {name: (self_device_us_total, count)}; @@ -960,6 +1032,7 @@ def main(): stats = {k: dist_ctx.allreduce_avg_float(v) for k, v in stats.items()} per_layer_us = stats["median"] / n_layers prof_us = dist_ctx.allreduce_avg_float(prof_us) + _run_distributed_smi_replay(pipe, dist_ctx, stats["median"], n_layers) tbl = None if args.profile_table: tbl = _aggregate_prof_table( diff --git a/op_tests/smi_monitor.py b/op_tests/smi_monitor.py index 6cbb82a4ef..062c7da694 100644 --- a/op_tests/smi_monitor.py +++ b/op_tests/smi_monitor.py @@ -160,6 +160,8 @@ def __init__(self, device_index: int = 0, interval_s: float = 0.05) -> None: self._samples: list[dict] = [] self._thread: threading.Thread | None = None self._stop_event = threading.Event() + self._ready_event = threading.Event() + self._error: BaseException | None = None # ------------------------------------------------------------------ # Public API @@ -170,9 +172,18 @@ def start(self) -> None: if self._thread is not None and self._thread.is_alive(): raise RuntimeError("GpuMonitor is already running") self._samples = [] + self._error = None self._stop_event.clear() + self._ready_event.clear() self._thread = threading.Thread(target=self._poll_loop, daemon=True) self._thread.start() + if not self._ready_event.wait(timeout=10.0): + self._stop_event.set() + raise RuntimeError("timed out while initializing amdsmi monitor") + if self._error is not None: + error = self._error + self.stop() + raise RuntimeError(f"failed to initialize amdsmi monitor: {error}") def stop(self) -> None: """Stop background polling and wait for the thread to finish.""" @@ -186,16 +197,24 @@ def samples(self) -> list[dict]: """Collected samples; each is a dict with 'timestamp_s' plus metric keys.""" return list(self._samples) - def summary(self) -> dict: - """Return min/mean/median/max for every numeric metric across all samples.""" - if not self._samples: + def summary( + self, *, start_s: float | None = None, end_s: float | None = None + ) -> dict: + """Return metric summaries, optionally restricted to a timestamp window.""" + samples = [ + sample + for sample in self._samples + if (start_s is None or sample["timestamp_s"] >= start_s) + and (end_s is None or sample["timestamp_s"] <= end_s) + ] + if not samples: return {} - keys = [k for k in self._samples[0] if k != "timestamp_s"] + keys = {key for sample in samples for key in sample if key != "timestamp_s"} result: dict = {} - for key in keys: + for key in sorted(keys): vals = sorted( s[key] - for s in self._samples + for s in samples if s.get(key) is not None and s[key] != "N/A" ) if not vals: @@ -228,13 +247,27 @@ def __exit__(self, *_) -> None: # ------------------------------------------------------------------ def _poll_loop(self) -> None: - amdsmi.amdsmi_init() + initialized = False try: + target_bdf = None if isinstance(self._device_index, int): - devices = amdsmi.amdsmi_get_processor_handles() - handle = devices[self._device_index] + target_bdf = _hip_device_bdf(self._device_index) + amdsmi.amdsmi_init() + initialized = True + if target_bdf is not None: + handles = { + _amdsmi_bdf_str(handle): handle + for handle in amdsmi.amdsmi_get_processor_handles() + } + if target_bdf not in handles: + raise RuntimeError( + f"no amdsmi handle for HIP device {self._device_index} " + f"({target_bdf})" + ) + handle = handles[target_bdf] else: handle = self._device_index + self._ready_event.set() while not self._stop_event.is_set(): t0 = time.perf_counter() self._samples.append(_collect_sample(handle)) @@ -242,8 +275,13 @@ def _poll_loop(self) -> None: remaining = self._interval_s - elapsed if remaining > 0: self._stop_event.wait(timeout=remaining) + except BaseException as error: + self._error = error + self._ready_event.set() finally: - amdsmi.amdsmi_shut_down() + if initialized: + amdsmi.amdsmi_shut_down() + self._ready_event.set() # ------------------------------------------------------------------ @@ -273,6 +311,17 @@ def smi_replay_enabled() -> bool: return os.environ.get("AITER_SMI_MONITOR", "0") == "1" +def emit_smi_result(result: dict) -> None: + """Write one structured result to the combo JSONL sink or stdout.""" + line = SMI_RESULT_PREFIX + json.dumps(result, sort_keys=True) + output_path = os.environ.get("AITER_SMI_OUTPUT_PATH") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(line + "\n") + else: + print(line, flush=True) + + def replay_with_smi( fn, *, @@ -330,11 +379,5 @@ def replay_with_smi( ) # A shared JSONL sink survives fd silencing and child processes. Standalone # UT runs without a sink still get a machine-readable stdout record. - line = SMI_RESULT_PREFIX + json.dumps(result, sort_keys=True) - output_path = os.environ.get("AITER_SMI_OUTPUT_PATH") - if output_path: - with open(output_path, "a", encoding="utf-8") as output: - output.write(line + "\n") - else: - print(line, flush=True) + emit_smi_result(result) return result