From a7d8d99abef4f61dcbe6784ed99125c63f9382af Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Mon, 24 Aug 2026 07:40:09 +0000 Subject: [PATCH 1/7] [FlyDSL] gfx942 fp8_mqa_logits: let _auto_variant choose rows_per_block _auto_variant returned f"mfma_r2_w{wpb}", so rows_per_block was pinned at 2 and seven of the nine registered variants -- including every member of the r4 family -- could never be selected. r4 amortizes each KV tile load over twice as many query rows and is faster from seq_len 8 upward. The gap is widest exactly where it costs most: vLLM chunks indexer prefill to fit VLLM_SPARSE_INDEXER_MAX_LOGITS_MB (512 MB), which caps seq_len at 1024 when seq_len_kv is 131072, so the existing seq_len >= 2048 branch cannot fire at long context and every such call took mfma_r2_w4 -- the median of the nine by speed, with the best 1.65x faster. Measured on MI325X (gfx942), seq_len_kv 131072, best variant vs the r2 pick: seq_len 1 r2 25.2 us r4 79.1 us r4 3.1x worse seq_len 4 r2 24.3 us r4 31.5 us r4 1.3x worse seq_len 8 r2 34.0 us r4 32.9 us r4 1.03x better seq_len 16 r2 56.5 us r4 47.5 us r4 1.19x better seq_len 1024 r2 2564.3 us r4 1520.3 us r4 1.69x better Below seq_len 8 the host padding of seq_len up to a multiple of RPB dominates -- at seq_len 1 an r4 kernel computes 4 rows to obtain 1 -- so r2 is kept there and behaviour for those shapes is unchanged. Logits are bitwise identical across all variants at every shape tested, so this is purely a blocking/occupancy change. End to end on 8x MI325X, TP8, GLM-5.2-FP8, 131072 in / 1024 out, concurrency 8: median TPOT improves 7.22% and output throughput 6.60%. This kernel is 16.8% of GPU time at that point. Signed-off-by: Jin Tao --- .../kernels/mqa_logits/fp8_mqa_logits.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py index 0397c48418d..ccb70fffe67 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py @@ -387,12 +387,33 @@ def _mk_builder(rpb, wpb): DEFAULT_VARIANT = "mfma_r2_w4" +# Smallest seq_len at which RPB=4 pays off. Below it the host padding to a +# multiple of RPB dominates: at seq_len=1 an RPB=4 kernel computes 4 rows to +# obtain 1. Measured crossover on gfx942 at seq_len_kv=131072 (us, best variant +# vs the RPB=2 pick this replaces): +# seq_len 1 -> r2 25.2 wins, r4 79.1 (3.1x worse) +# seq_len 4 -> r2 24.3 wins, r4 31.5 (1.3x worse) +# seq_len 8 -> r4 32.9 wins by 1.03x +# seq_len 16 -> r4 47.5 wins by 1.19x +# seq_len 1024 -> r4 1520.3 wins by 1.69x +_RPB4_MIN_SEQ_LEN = 8 + + def _auto_variant(seq_len, seq_len_kv): - """Pick (RPB, WPB) from the problem shape: 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.""" + """Pick (RPB, WPB) from the problem shape. + + RPB=4 amortizes each KV tile load over twice as many query rows as RPB=2 and + wins from ``_RPB4_MIN_SEQ_LEN`` upward, by 1.65x at the long-context prefill + shapes vLLM issues (seq_len 1024, seq_len_kv 131072). Below that threshold + the padding of seq_len up to a multiple of RPB costs more than the extra + reuse buys, so RPB=2 stays. + + 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. + """ + rpb = 4 if seq_len >= _RPB4_MIN_SEQ_LEN else 2 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}" def _resolve_variant(variant, seq_len, seq_len_kv): From 5691e17cf6970cc2f314eb2558f1eb6622770951 Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Mon, 24 Aug 2026 08:06:10 +0000 Subject: [PATCH 2/7] [FlyDSL] gfx942 fp8_mqa_logits: pick RPB on element count, and keep it a divisor Refines the previous commit's rule after a 2-D sweep. That rule keyed RPB off seq_len alone with a crossover measured only at seq_len_kv=131072; sweeping the other contexts shows the crossover is not a seq_len threshold at all, and that a second effect was being read as one. RPB tracks the logits element count. Over seq_len 1..8192 x seq_len_kv 1024..262144 on MI325X, the boundaries land on the same element count at every context: RPB=1 wins below 2**19 elements (27/27 shapes), RPB=2 at 2**19 (6/6), RPB=4 from 2**21 up (38/38), with 2**20 a transition band split 3/3. Keying off seq_len instead put the previous rule on the wrong side at low context: at seq_len 16, seq_len_kv 1024 it chose RPB=4 and ran 1.26x slower than RPB=1. RPB must also divide seq_len. When it does not, the launcher pads with four torch.cat calls; that is a flat ~44 us of host-side overhead, independent of seq_len_kv, and it is the whole of the "small seq_len" penalty the previous commit attributed to wasted rows. At seq_len 1, seq_len_kv 131072: RPB=1 23.1 us, RPB=2 67.8 us, of which the four cats are 44.1 us and pre-padding by hand recovers all of it (21.9 us). So the penalty is not proportional to the padding -- 1 wasted row of 2 costs the same as 3 of 4 -- and it applies to every odd seq_len, which the old rule sent to RPB=2 unconditionally. Stepping down to a divisor is only right while the kernel is cheap relative to that fixed cost, so it is gated to the same 2**21 elements: at seq_len 1025, seq_len_kv 131072 the dividing RPB=1 takes 3880 us against 2601 us for a padded RPB=2. Measured on MI325X, no FLYDSL_FP8_MQA_LOGITS_VARIANT set, old pick vs new: seq_len seq_len_kv old new speedup 1024 131072 r2_w4 2598.7 r4_w4 1623.5 1.60x 1025 131072 r2_w4 2609.3 r4_w4 1594.0 1.64x 512 131072 r2_w4 1190.1 r4_w4 795.3 1.50x 700 50000 r2_w4 661.6 r4_w4 440.3 1.50x 333 12000 r2_w4 118.4 r4_w4 97.1 1.22x 16 131072 r2_w4 58.3 r4_w4 47.4 1.23x 3 131072 r2_w4 66.4 r1_w4 28.1 2.36x 1 131072 r2_w4 69.4 r1_w4 23.2 2.99x 1 1024 r2_w4 66.2 r1_w4 22.2 2.98x No shape measured regresses; the smallest gain is 1.04x. Against the best of the nine variants at each shape, pooled over held-out data (non-power-of-two shapes, a fine seq_len sweep, and head counts 16 and 64), the geometric mean cost falls from 1.45x to 1.03x and the worst case from 3.17x to 1.41x. Logits are bitwise identical across all nine variants at all 180 shapes swept (1620 timings), so this remains purely a blocking/occupancy change. WPB is deliberately left alone. It is worth a few percent at most here, and unlike RPB its optimum moves with the head count -- at 64 heads the current WPB rule costs 1.61x worst case where a fixed WPB=4 costs 1.07x -- so it needs its own sweep rather than a change fitted to one head count. Signed-off-by: Jin Tao Co-authored-by: Cursor --- .../kernels/mqa_logits/fp8_mqa_logits.py | 50 ++++++++++++------- 1 file changed, 32 insertions(+), 18 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py index ccb70fffe67..3490376991d 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py @@ -387,31 +387,45 @@ def _mk_builder(rpb, wpb): DEFAULT_VARIANT = "mfma_r2_w4" -# Smallest seq_len at which RPB=4 pays off. Below it the host padding to a -# multiple of RPB dominates: at seq_len=1 an RPB=4 kernel computes 4 rows to -# obtain 1. Measured crossover on gfx942 at seq_len_kv=131072 (us, best variant -# vs the RPB=2 pick this replaces): -# seq_len 1 -> r2 25.2 wins, r4 79.1 (3.1x worse) -# seq_len 4 -> r2 24.3 wins, r4 31.5 (1.3x worse) -# seq_len 8 -> r4 32.9 wins by 1.03x -# seq_len 16 -> r4 47.5 wins by 1.19x -# seq_len 1024 -> r4 1520.3 wins by 1.69x -_RPB4_MIN_SEQ_LEN = 8 +# RPB crossovers track the logits element count seq_len * seq_len_kv, not +# seq_len alone: a larger RPB amortizes each KV tile load over more query rows, +# and that pays off once there is enough total work. Measured on gfx942 +# (MI325X) over seq_len 1..8192 x seq_len_kv 1024..262144, the boundaries land +# on the same element count at every context: +# +# elements < 2**19 RPB=1 fastest at all 27 shapes measured +# elements = 2**19 RPB=2 fastest at all 6 +# elements >= 2**21 RPB=4 fastest at all 38 +# +# (2**20 is the transition band, split 3/3 between RPB=2 and RPB=4.) +_RPB2_MIN_ELEMS = 2**19 +_RPB4_MIN_ELEMS = 2**21 def _auto_variant(seq_len, seq_len_kv): """Pick (RPB, WPB) from the problem shape. - RPB=4 amortizes each KV tile load over twice as many query rows as RPB=2 and - wins from ``_RPB4_MIN_SEQ_LEN`` upward, by 1.65x at the long-context prefill - shapes vLLM issues (seq_len 1024, seq_len_kv 131072). Below that threshold - the padding of seq_len up to a multiple of RPB costs more than the extra - reuse buys, so RPB=2 stays. + RPB was previously pinned to 2, which left the whole r4 family unreachable + and cost 1.65x at the long-context prefill shapes vLLM issues (seq_len 1024, + seq_len_kv 131072, i.e. 2**27 elements). - 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. + RPB must also divide seq_len. When it does not, the launcher pads seq_len up + with four ``torch.cat`` calls, a flat ~44 us of host-side overhead that does + not scale with the work -- so it swamps small shapes and is noise on large + ones. Stepping down to a divisor is therefore right below + ``_RPB4_MIN_ELEMS`` and wrong above it: at seq_len 1025, seq_len_kv 131072, + RPB=1 divides but runs 3880 us against 2601 us for a padded RPB=2. + + WPB is left as it was. It is worth a few percent at most here, and unlike + RPB its optimum moves with the head count, so it needs its own sweep. """ - rpb = 4 if seq_len >= _RPB4_MIN_SEQ_LEN else 2 + 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_r{rpb}_w{wpb}" From d863e4dcc428c2df646ee9a40a9238347287c8f7 Mon Sep 17 00:00:00 2001 From: "jin.tao@amd.com" Date: Tue, 25 Aug 2026 07:38:13 +0000 Subject: [PATCH 3/7] [FlyDSL] gfx942 fp8_mqa_logits: trim _auto_variant comments per review Move RPB element-count thresholds into _auto_variant and shorten the docstring; logic unchanged. Signed-off-by: Jin Tao Co-authored-by: Cursor --- .../kernels/mqa_logits/fp8_mqa_logits.py | 36 ++++--------------- 1 file changed, 6 insertions(+), 30 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py index 3490376991d..1f9ebce1927 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py @@ -387,42 +387,18 @@ def _mk_builder(rpb, wpb): DEFAULT_VARIANT = "mfma_r2_w4" -# RPB crossovers track the logits element count seq_len * seq_len_kv, not -# seq_len alone: a larger RPB amortizes each KV tile load over more query rows, -# and that pays off once there is enough total work. Measured on gfx942 -# (MI325X) over seq_len 1..8192 x seq_len_kv 1024..262144, the boundaries land -# on the same element count at every context: -# -# elements < 2**19 RPB=1 fastest at all 27 shapes measured -# elements = 2**19 RPB=2 fastest at all 6 -# elements >= 2**21 RPB=4 fastest at all 38 -# -# (2**20 is the transition band, split 3/3 between RPB=2 and RPB=4.) -_RPB2_MIN_ELEMS = 2**19 -_RPB4_MIN_ELEMS = 2**21 - - def _auto_variant(seq_len, seq_len_kv): """Pick (RPB, WPB) from the problem shape. - RPB was previously pinned to 2, which left the whole r4 family unreachable - and cost 1.65x at the long-context prefill shapes vLLM issues (seq_len 1024, - seq_len_kv 131072, i.e. 2**27 elements). - - RPB must also divide seq_len. When it does not, the launcher pads seq_len up - with four ``torch.cat`` calls, a flat ~44 us of host-side overhead that does - not scale with the work -- so it swamps small shapes and is noise on large - ones. Stepping down to a divisor is therefore right below - ``_RPB4_MIN_ELEMS`` and wrong above it: at seq_len 1025, seq_len_kv 131072, - RPB=1 divides but runs 3880 us against 2601 us for a padded RPB=2. - - WPB is left as it was. It is worth a few percent at most here, and unlike - RPB its optimum moves with the head count, so it needs its own sweep. + RPB from ``seq_len * seq_len_kv`` thresholds; step down to a divisor of + ``seq_len`` when padding overhead would dominate. WPB unchanged. """ + rpb2_min_elems = 2**19 + rpb4_min_elems = 2**21 elems = seq_len * seq_len_kv - if elems < _RPB2_MIN_ELEMS: + if elems < rpb2_min_elems: rpb = 1 - elif elems < _RPB4_MIN_ELEMS: + elif elems < rpb4_min_elems: rpb = 2 if seq_len % 2 == 0 else 1 else: rpb = 4 From 326a4227e8765b3d47f3db7b908cfbfa66bfd627 Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Wed, 16 Sep 2026 09:19:11 +0000 Subject: [PATCH 4/7] [FlyDSL] gfx942 fp8_mqa_logits: scope the _auto_variant step-down note The docstring stated the divisor step-down as an unconditional rule, but it only applies in the middle band: above the top threshold RPB stays 4 and the padding is accepted. Say which band it applies to, and why the top band is exempt. No logic change. Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py index ed9d62894f9..6539e6d1fae 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py @@ -376,8 +376,9 @@ def _mk_builder(rpb, wpb): def _auto_variant(seq_len, seq_len_kv): """Pick (RPB, WPB) from the problem shape. - RPB from ``seq_len * seq_len_kv`` thresholds; step down to a divisor of - ``seq_len`` when padding overhead would dominate. WPB unchanged. + 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 unchanged. """ rpb2_min_elems = 2**19 rpb4_min_elems = 2**21 From 3437167c6f49f724a39c3075240bfe2c1e4b5ec8 Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Wed, 16 Sep 2026 09:19:11 +0000 Subject: [PATCH 5/7] [FlyDSL] gfx942 fp8_mqa_logits: unit-test the variant selector The shape sweep in test_flydsl_fp8_mqa_logits.py never reaches r4 -- its largest default shape is 1024 x 1560, below the 2**21 threshold -- and nothing asserted _auto_variant directly, so a regression could re-pin RPB to 2 with every correctness test still green. Cover the RPB bands at both thresholds, that the edges track seq_len * seq_len_kv rather than seq_len alone, the middle-band step-down on odd seq_len (and that it stops above the top threshold), the production 1024 x 131072 shape, the unchanged WPB rule, and _resolve_variant precedence so the auto path is confirmed to be the default. Pure shape arithmetic, so no kernel launch. Co-authored-by: Cursor --- .../test_flydsl_fp8_mqa_logits_variant.py | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py b/op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py new file mode 100644 index 00000000000..5e1b731dc21 --- /dev/null +++ b/op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Selector coverage for the FlyDSL fp8_mqa_logits kernel variants. + +``_auto_variant`` is pure shape arithmetic, so these launch no kernel: they pin +the RPB bands, the middle-band step-down, and the unchanged WPB rule, none of +which the shape sweep in ``test_flydsl_fp8_mqa_logits.py`` reaches (its largest +shape is below the top RPB threshold, so ``r4`` never gets selected there). +""" + +import pytest + +from aiter.ops.flydsl.kernels.mqa_logits.fp8_mqa_logits import ( + KERNEL_VARIANTS, + _auto_variant, + _resolve_variant, +) + +RPB2_MIN_ELEMS = 2**19 +RPB4_MIN_ELEMS = 2**21 + + +def _rpb_wpb(seq_len, seq_len_kv): + tag = _auto_variant(seq_len, seq_len_kv) + assert tag in KERNEL_VARIANTS, f"{tag} is not a registered variant" + _, rpb, wpb = tag.split("_") + return int(rpb[1:]), int(wpb[1:]) + + +@pytest.mark.parametrize( + "seq_len, seq_len_kv, expected_rpb", + [ + (1, 1024, 1), # below RPB2_MIN_ELEMS + (64, 4096, 1), + (256, 1024, 1), + (64, 8192, 2), # exactly RPB2_MIN_ELEMS + (512, 1024, 2), + (1024, 1560, 2), # largest shape the existing sweep covers + (513, 1024, 1), # middle band, odd seq_len: step down + (1, 2**20, 1), + (2048, 1024, 4), # exactly RPB4_MIN_ELEMS + (1024, 131072, 4), # vLLM long-context indexer prefill + (1025, 131072, 4), # odd seq_len: padding accepted up here + ], +) +def test_auto_variant_rpb_bands(seq_len, seq_len_kv, expected_rpb): + rpb, _ = _rpb_wpb(seq_len, seq_len_kv) + assert rpb == expected_rpb + + +@pytest.mark.parametrize("seq_len_kv", [1024, 8192, 131072]) +@pytest.mark.parametrize( + "threshold, below, at", [(RPB2_MIN_ELEMS, 1, 2), (RPB4_MIN_ELEMS, 2, 4)] +) +def test_auto_variant_rpb_thresholds_track_element_count( + seq_len_kv, threshold, below, at +): + """Band edges land on ``seq_len * seq_len_kv``, not on ``seq_len`` alone.""" + seq_len = threshold // seq_len_kv + assert seq_len % 2 == 0, "keep the step-down out of this assertion" + assert _rpb_wpb(seq_len - 2, seq_len_kv)[0] == below + assert _rpb_wpb(seq_len, seq_len_kv)[0] == at + + +@pytest.mark.parametrize( + "seq_len_kv, odd_mid, even_mid, odd_above", + [(1024, 513, 514, 2049), (8192, 65, 66, 257)], +) +def test_auto_variant_steps_down_only_in_middle_band( + seq_len_kv, odd_mid, even_mid, odd_above +): + """An indivisible seq_len costs four host-side cats, a fixed price the + middle band cannot absorb and the top band can.""" + assert _rpb_wpb(odd_mid, seq_len_kv)[0] == 1 + assert _rpb_wpb(even_mid, seq_len_kv)[0] == 2 + assert _rpb_wpb(odd_above, seq_len_kv)[0] == 4 + + +def test_auto_variant_reaches_every_rpb_family(): + seen = { + _rpb_wpb(seq_len, seq_len_kv)[0] + for seq_len in (1, 3, 512, 513, 1024, 2048) + for seq_len_kv in (1024, 8192, 131072) + } + assert seen == {1, 2, 4} + + +@pytest.mark.parametrize( + "seq_len, seq_len_kv, expected_wpb", + [ + (2048, 8192, 2), + (4096, 131072, 2), + (2047, 8192, 4), + (2048, 4096, 4), + (1024, 131072, 4), + (1, 1024, 4), + ], +) +def test_auto_variant_wpb_rule_unchanged(seq_len, seq_len_kv, expected_wpb): + _, wpb = _rpb_wpb(seq_len, seq_len_kv) + assert wpb == expected_wpb + + +def test_resolve_variant_precedence(monkeypatch): + """Explicit > env > shape-adaptive, so the auto path is what runs by default.""" + monkeypatch.delenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", raising=False) + assert _resolve_variant(None, 1024, 131072) == "mfma_r4_w4" + + monkeypatch.setenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", "mfma_r1_w1") + assert _resolve_variant(None, 1024, 131072) == "mfma_r1_w1" + assert _resolve_variant("mfma_r2_w2", 1024, 131072) == "mfma_r2_w2" + + +def test_resolve_variant_rejects_unknown_tag(monkeypatch): + monkeypatch.delenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", raising=False) + with pytest.raises(ValueError, match="unknown fp8_mqa_logits variant"): + _resolve_variant("mfma_r3_w4", 1024, 131072) From 095113b2f6f26571a775d2b0a0749f2c0b478bdb Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Fri, 25 Sep 2026 07:09:12 +0000 Subject: [PATCH 6/7] [FlyDSL] gfx942 fp8_mqa_logits: put selector tests on the CI path CI only discovers op_tests/test_*.py, so the r1/r2/r4 selector coverage never ran. Move it there, run pytest from __main__, and add two r4 shapes to the GPU sweep so auto-selected padding is launched. Signed-off-by: Jin Tao Co-authored-by: Cursor --- op_tests/test_flydsl_fp8_mqa_logits.py | 6 +++++- .../test_flydsl_fp8_mqa_logits_variant.py | 13 ++++++++----- 2 files changed, 13 insertions(+), 6 deletions(-) rename op_tests/{flydsl_tests => }/test_flydsl_fp8_mqa_logits_variant.py (91%) diff --git a/op_tests/test_flydsl_fp8_mqa_logits.py b/op_tests/test_flydsl_fp8_mqa_logits.py index c390868c663..805ef488755 100644 --- a/op_tests/test_flydsl_fp8_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_mqa_logits.py @@ -648,7 +648,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 +750,10 @@ def main(): # enough that most blocks end up owning an empty column range. (64, 2048), (64, 8192), + # r4 auto-select: 256*8192 == 2**21. 257 is odd, so the launcher + # pads seq_len up to a multiple of RPB=4. + (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), diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py b/op_tests/test_flydsl_fp8_mqa_logits_variant.py similarity index 91% rename from op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py rename to op_tests/test_flydsl_fp8_mqa_logits_variant.py index 11c260af7bc..a91a6ea8c3a 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits_variant.py +++ b/op_tests/test_flydsl_fp8_mqa_logits_variant.py @@ -2,10 +2,9 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. """Selector coverage for the FlyDSL fp8_mqa_logits kernel variants. -``_auto_variant`` is pure shape arithmetic, so these launch no kernel: they pin -the RPB bands, the middle-band step-down, and the unchanged WPB rule, none of -which the shape sweep in ``test_flydsl_fp8_mqa_logits.py`` reaches (its largest -shape is below the top RPB threshold, so ``r4`` never gets selected there). +CI discovers ``op_tests/test_*.py``. These pin the RPB bands, the middle-band +step-down, and the WPB rule; they launch no kernel. The GPU sweep covers the +auto-selected ``r4`` path. """ import pytest @@ -42,7 +41,7 @@ def _rpb_wpb(seq_len, seq_len_kv): (256, 1024, 1), (64, 8192, 2), # exactly RPB2_MIN_ELEMS (512, 1024, 2), - (1024, 1560, 2), # largest shape the existing sweep covers + (1024, 1560, 2), (513, 1024, 1), # middle band, odd seq_len: step down (1, 2**20, 1), (2048, 1024, 4), # exactly RPB4_MIN_ELEMS @@ -122,3 +121,7 @@ def test_resolve_variant_rejects_unknown_tag(monkeypatch): monkeypatch.delenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", raising=False) with pytest.raises(ValueError, match="unknown fp8_mqa_logits variant"): _resolve_variant("mfma_r3_w4", 1024, 131072, NUM_HEADS) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) From feb833c3512172041aae3b2c0dfcf39e5b1afc18 Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Fri, 25 Sep 2026 10:13:26 +0000 Subject: [PATCH 7/7] [FlyDSL] gfx942 fp8_mqa_logits: fold the selector pins into the op test aiter op tests are plain scripts, not pytest, so drop test_flydsl_fp8_mqa_logits_variant.py and pin the gfx942 auto-selected variant in test_flydsl_fp8_mqa_logits.py instead: a host-only verify_auto_variant table at both RPB band edges, either side of the odd-seq_len step-down, the long-context prefill shape and the WPB switch. It runs in the default verify scenario on gfx942, and its failures count toward the same exit code as the kernel sweep. The selector import sits in the file's existing ImportError guard. Signed-off-by: Jin Tao Co-authored-by: Cursor --- op_tests/test_flydsl_fp8_mqa_logits.py | 65 +++++++-- .../test_flydsl_fp8_mqa_logits_variant.py | 127 ------------------ 2 files changed, 56 insertions(+), 136 deletions(-) delete mode 100644 op_tests/test_flydsl_fp8_mqa_logits_variant.py diff --git a/op_tests/test_flydsl_fp8_mqa_logits.py b/op_tests/test_flydsl_fp8_mqa_logits.py index 805ef488755..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 @@ -750,8 +780,8 @@ def main(): # enough that most blocks end up owning an empty column range. (64, 2048), (64, 8192), - # r4 auto-select: 256*8192 == 2**21. 257 is odd, so the launcher - # pads seq_len up to a multiple of RPB=4. + # 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 @@ -820,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 = [] @@ -838,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( @@ -849,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) diff --git a/op_tests/test_flydsl_fp8_mqa_logits_variant.py b/op_tests/test_flydsl_fp8_mqa_logits_variant.py deleted file mode 100644 index a91a6ea8c3a..00000000000 --- a/op_tests/test_flydsl_fp8_mqa_logits_variant.py +++ /dev/null @@ -1,127 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Selector coverage for the FlyDSL fp8_mqa_logits kernel variants. - -CI discovers ``op_tests/test_*.py``. These pin the RPB bands, the middle-band -step-down, and the WPB rule; they launch no kernel. The GPU sweep covers the -auto-selected ``r4`` path. -""" - -import pytest - -from aiter.ops.flydsl.fp8_mqa_logits_kernels import ( - _ARCH, - KERNEL_VARIANTS, - _auto_variant, - _resolve_variant, -) - -pytestmark = pytest.mark.skipif( - _ARCH != "gfx942", reason="gfx942 variant selector and registry" -) - -RPB2_MIN_ELEMS = 2**19 -RPB4_MIN_ELEMS = 2**21 -# The gfx942 rule ignores num_heads; it only has to be a multiple of MFMA_M=16. -NUM_HEADS = 32 - - -def _rpb_wpb(seq_len, seq_len_kv): - tag = _auto_variant(seq_len, seq_len_kv, NUM_HEADS) - assert tag in KERNEL_VARIANTS, f"{tag} is not a registered variant" - _, rpb, wpb = tag.split("_") - return int(rpb[1:]), int(wpb[1:]) - - -@pytest.mark.parametrize( - "seq_len, seq_len_kv, expected_rpb", - [ - (1, 1024, 1), # below RPB2_MIN_ELEMS - (64, 4096, 1), - (256, 1024, 1), - (64, 8192, 2), # exactly RPB2_MIN_ELEMS - (512, 1024, 2), - (1024, 1560, 2), - (513, 1024, 1), # middle band, odd seq_len: step down - (1, 2**20, 1), - (2048, 1024, 4), # exactly RPB4_MIN_ELEMS - (1024, 131072, 4), # vLLM long-context indexer prefill - (1025, 131072, 4), # odd seq_len: padding accepted up here - ], -) -def test_auto_variant_rpb_bands(seq_len, seq_len_kv, expected_rpb): - rpb, _ = _rpb_wpb(seq_len, seq_len_kv) - assert rpb == expected_rpb - - -@pytest.mark.parametrize("seq_len_kv", [1024, 8192, 131072]) -@pytest.mark.parametrize( - "threshold, below, at", [(RPB2_MIN_ELEMS, 1, 2), (RPB4_MIN_ELEMS, 2, 4)] -) -def test_auto_variant_rpb_thresholds_track_element_count( - seq_len_kv, threshold, below, at -): - """Band edges land on ``seq_len * seq_len_kv``, not on ``seq_len`` alone.""" - seq_len = threshold // seq_len_kv - assert seq_len % 2 == 0, "keep the step-down out of this assertion" - assert _rpb_wpb(seq_len - 2, seq_len_kv)[0] == below - assert _rpb_wpb(seq_len, seq_len_kv)[0] == at - - -@pytest.mark.parametrize( - "seq_len_kv, odd_mid, even_mid, odd_above", - [(1024, 513, 514, 2049), (8192, 65, 66, 257)], -) -def test_auto_variant_steps_down_only_in_middle_band( - seq_len_kv, odd_mid, even_mid, odd_above -): - """An indivisible seq_len costs four host-side cats, a fixed price the - middle band cannot absorb and the top band can.""" - assert _rpb_wpb(odd_mid, seq_len_kv)[0] == 1 - assert _rpb_wpb(even_mid, seq_len_kv)[0] == 2 - assert _rpb_wpb(odd_above, seq_len_kv)[0] == 4 - - -def test_auto_variant_reaches_every_rpb_family(): - seen = { - _rpb_wpb(seq_len, seq_len_kv)[0] - for seq_len in (1, 3, 512, 513, 1024, 2048) - for seq_len_kv in (1024, 8192, 131072) - } - assert seen == {1, 2, 4} - - -@pytest.mark.parametrize( - "seq_len, seq_len_kv, expected_wpb", - [ - (2048, 8192, 2), - (4096, 131072, 2), - (2047, 8192, 4), - (2048, 4096, 4), - (1024, 131072, 4), - (1, 1024, 4), - ], -) -def test_auto_variant_wpb_rule_unchanged(seq_len, seq_len_kv, expected_wpb): - _, wpb = _rpb_wpb(seq_len, seq_len_kv) - assert wpb == expected_wpb - - -def test_resolve_variant_precedence(monkeypatch): - """Explicit > env > shape-adaptive, so the auto path is what runs by default.""" - monkeypatch.delenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", raising=False) - assert _resolve_variant(None, 1024, 131072, NUM_HEADS) == "mfma_r4_w4" - - monkeypatch.setenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", "mfma_r1_w1") - assert _resolve_variant(None, 1024, 131072, NUM_HEADS) == "mfma_r1_w1" - assert _resolve_variant("mfma_r2_w2", 1024, 131072, NUM_HEADS) == "mfma_r2_w2" - - -def test_resolve_variant_rejects_unknown_tag(monkeypatch): - monkeypatch.delenv("FLYDSL_FP8_MQA_LOGITS_VARIANT", raising=False) - with pytest.raises(ValueError, match="unknown fp8_mqa_logits variant"): - _resolve_variant("mfma_r3_w4", 1024, 131072, NUM_HEADS) - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__, "-v"]))