diff --git a/aiter/ops/flydsl/fp8_mqa_logits_kernels.py b/aiter/ops/flydsl/fp8_mqa_logits_kernels.py index c6569a6ea99..f9029b3efc3 100644 --- a/aiter/ops/flydsl/fp8_mqa_logits_kernels.py +++ b/aiter/ops/flydsl/fp8_mqa_logits_kernels.py @@ -347,9 +347,11 @@ def _parse_variant(tag): def _auto_variant(seq_len, seq_len_kv, num_heads): """Pick a variant from the problem shape. - gfx942 (unchanged): RPB=2 always; WPB=2 packs more column tiles per wave - when M and N are both large, else WPB=4 for more wavefronts on small-M / - short-window shapes. + gfx942: RPB from ``seq_len * seq_len_kv`` thresholds; in the middle band an + odd ``seq_len`` drops RPB to 1, since padding costs a fixed host-side + overhead that only the largest shapes outgrow. WPB=2 packs more column + tiles per wave when M and N are both large, else WPB=4 for more + wavefronts on small-M / short-window shapes. gfx950 H>=128: mfma32x32x64 at r=1 always -- ample compute, more blocks. @@ -364,8 +366,17 @@ def _auto_variant(seq_len, seq_len_kv, num_heads): shapes (KV pressure high), r=1 otherwise. """ if _ARCH == "gfx942": + rpb2_min_elems = 2**19 + rpb4_min_elems = 2**21 + elems = seq_len * seq_len_kv + if elems < rpb2_min_elems: + rpb = 1 + elif elems < rpb4_min_elems: + rpb = 2 if seq_len % 2 == 0 else 1 + else: + rpb = 4 wpb = 2 if (seq_len >= 2048 and seq_len_kv >= 8192) else 4 - return f"mfma_r2_w{wpb}" + return f"mfma_r{rpb}_w{wpb}" if _ARCH == "gfx950": if num_heads >= 128: return "mfma32x32x64_bkv64_r1_w2_lds3" diff --git a/op_tests/test_flydsl_fp8_mqa_logits.py b/op_tests/test_flydsl_fp8_mqa_logits.py index c390868c663..cc186a501e8 100644 --- a/op_tests/test_flydsl_fp8_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_mqa_logits.py @@ -68,8 +68,9 @@ try: from aiter.ops.flydsl import flydsl_fp8_mqa_logits + from aiter.ops.flydsl.fp8_mqa_logits_kernels import _auto_variant except ImportError: - flydsl_fp8_mqa_logits = None + flydsl_fp8_mqa_logits = _auto_variant = None # Bench timing knobs, set from argv in main(), read by _time_us. BENCH_WARMUP = 10 @@ -405,6 +406,35 @@ def verify_fp8_mqa_logits( return ret +# The variant gfx942 auto-selects, pinned at both RPB band edges in s_q * s_k, +# on either side of the odd-s_q step-down, and at the WPB switch. +_GFX942_AUTO_VARIANTS = { + (1, 1024): "mfma_r1_w4", + (62, 8192): "mfma_r1_w4", + (64, 8192): "mfma_r2_w4", # 2**19 + (512, 1024): "mfma_r2_w4", # 2**19 + (65, 8192): "mfma_r1_w4", # odd, below 2**21: steps down + (254, 8192): "mfma_r2_w4", + (256, 8192): "mfma_r4_w4", # 2**21 + (257, 8192): "mfma_r4_w4", # odd, from 2**21 up: pads + (1024, 131072): "mfma_r4_w4", # long-context indexer prefill + (2048, 8192): "mfma_r4_w2", +} + + +@benchmark() +def verify_auto_variant(s_q, s_k, num_heads): + """Check the gfx942 auto-selected variant against its pin. Launches nothing.""" + variant = _auto_variant(s_q, s_k, num_heads) + expected = _GFX942_AUTO_VARIANTS[(s_q, s_k)] + if variant != expected: + raise AssertionError( + f"auto-selected {variant}, expected {expected} " + f"[s_q={s_q} s_k={s_k} nh={num_heads}]" + ) + return {"gfx": get_gfx(), "variant": variant, "status": "ok"} + + _FLUSH_CACHE = None @@ -648,7 +678,7 @@ def _full_set(args): def _reduced_set(args): """One case per (shape, window) pair, with the remaining axes rotated. - 56 cases on the defaults, which is what the CI lane runs. Shape and window + 65 cases on the defaults, which is what the CI lane runs. Shape and window are the axes that reach genuinely distinct kernel paths -- the grid.y split, the negative/empty window collapse, the cu_starts clamp -- so they are covered exhaustively. num_heads, head_dim, clean_logits and the operand @@ -750,6 +780,10 @@ def main(): # enough that most blocks end up owning an empty column range. (64, 2048), (64, 8192), + # gfx942 auto-selects r4 from 256*8192 == 2**21 up (pinned in + # _GFX942_AUTO_VARIANTS). 257 is odd, so the launcher pads it. + (256, 8192), + (257, 8192), # s_kv < s_q. A causal mask then puts cu_ends below zero on the # leading rows, so these cover the negative-window path end to end. (128, 64), @@ -816,6 +850,27 @@ def main(): ) failures = [] + total = len(cases) * len(scenarios) + if "verify" in scenarios and get_gfx() == "gfx942": + rows = [] + for (s_q, s_k), nh in itertools.product(_GFX942_AUTO_VARIANTS, args.num_heads): + pin = {"s_q": s_q, "s_k": s_k, "num_heads": nh} + try: + rows.append(verify_auto_variant(**pin)) + except AssertionError as exc: + aiter.logger.error("FAILED auto_variant case: %s", exc) + failures.append(("auto_variant", f"s_q={s_q} s_k={s_k} nh={nh}", exc)) + rows.append({**pin, "gfx": get_gfx(), "status": "FAIL"}) + total += len(rows) + aiter.logger.info( + "fp8_mqa_logits auto_variant summary (markdown):\n%s", + pd.DataFrame(rows).to_markdown(index=False), + ) + elif "verify" in scenarios: + aiter.logger.warning( + "%s: skipping the auto_variant pins, which are gfx942-only", get_gfx() + ) + for scenario in scenarios: run = verify_fp8_mqa_logits if scenario == "verify" else bench_fp8_mqa_logits rows = [] @@ -834,7 +889,7 @@ def main(): exc, exc_info=not isinstance(exc, AssertionError), ) - failures.append((scenario, case, exc)) + failures.append((scenario, _case_tag(*case), exc)) rows.append({**case._asdict(), "gfx": get_gfx(), "status": "FAIL"}) df = pd.DataFrame(rows) aiter.logger.info( @@ -845,16 +900,12 @@ def main(): if "speedup" in df: _log_speedup(df["speedup"]) - total = len(cases) * len(scenarios) if failures: aiter.logger.error( "fp8_mqa_logits: %d of %d case runs FAILED\n%s", len(failures), total, - "\n".join( - f" [{sc}] {_case_tag(*case)}\n {exc}" - for sc, case, exc in failures - ), + "\n".join(f" [{sc}] {tag}\n {exc}" for sc, tag, exc in failures), ) sys.exit(1) aiter.logger.info("fp8_mqa_logits: all %d case runs passed", total)