[vpietila-amd]

Summary

This PR improves the current FlyDSL FP8 MQA logits kernel on gfx950. The KV-split (grid.y) heuristic now uses the occupancy-aware wave-cost model (same as for gfx942), and a tuned route table replaces the hand-written variant pick. The LDS-pipelined builder gains new head-accumulation, head-reduce, and software-pipelining options, and moves to async copy atoms. The empty/past-the-end KV-window handling is also fixed.

On MI350X, over the 960-case --full sweep, FlyDSL FP8 MQA logits provides a speed-up over the existing Gluon and
the old FlyDSL kernel (geometric mean)

Clean Logits	This PR vs Gluon	This PR vs FlyDSL (main)	FlyDSL (main) vs Gluon (main)
True	1.72x	1.35x	1.28x
False	1.46x	1.43x	1.02x

Motivation

• On main, the gfx950 split policy used a fixed cu_oversub target that ignores how many blocks of a given kernel instance fit on a CU. Small-M / long-KV shapes were badly under- or over-split.
• The hand-written _auto_variant for gfx950 picked from a handful of fixed variants and lost to Gluon on about half the sweep. On main, FlyDSL is 0.91×–1.00× Gluon geomean for H32–H64 at clean_logits=False.
• A row whose window starts past seq_len_kv (past_end), or any empty window, stretched the block's union KV window to seq_len_kv. The block scanned the entire KV for nothing.
• The LDS builder used raw raw_ptr_buffer_load_lds / create_llvm_ptr / buffer_ops addressing and a hand-coded swizzle.

Changes

Host side (aiter/ops/flydsl/fp8_mqa_logits_kernels.py)

• gfx950 _SplitPolicy adopts gfx942's occupancy-aware wave-cost model (block_overhead_tiles=8). A launch that cannot query occupancy keeps main's fixed cu_oversub=4 target.
• gfx950 variants are now a tag grammar built on demand (_variant_builder): mfma<MxNxK>_bkv<B>_r<R>_w<W>[_lds<NB>[_pd<PD>][_swp][_rs][_ls][_ts][_dg]][_h<fma|fma2>][_wpe<N>][_ns<N>]. The gfx942 registry is unchanged.
• The tuned route table _GFX950_ROUTES, keyed by (num_heads class, head_size class) with seq_len / seq_len_kv bounds, replaces the gfx950 _auto_variant heuristic. _auto_variant / _resolve_variant now take head_size, and a head_size the atom's MFMA_K doesn't match is rejected host-side.
• Routes are tuned at head_size 64 and 128. For a larger head_size, _auto_variant skips any route whose LDS slots exceed the device's LDS (_variant_lds_bytes against get_lds_capacity_bytes) and falls back to _GFX950_LARGE_HEAD_FALLBACK if none fits. _variant_supports applies the same LDS check.
• _wpe<N> sets the waves_per_eu compile hint per variant.
• The default gfx950 variant is now mfma32x32x64_bkv64_r1_w1_lds2_swp_rs_hfma.

Kernel (aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py)

• head_acc modes for the per-lane head sum: add (previous behavior), fma (one non-reassociable FMA chain), and fma2 (two packed v_pk_fma_f32 chains).
• rs_head: the 64 // MFMA_N n-tiles of a group are head-reduced together with permlane swaps (warp_reduce_scatter_strided), and all 64 lanes store one contiguous run. Before, each n-tile needed a shuffle butterfly and an MFMA_N-lane store.
• sw_pipe: a depth-2 software pipeline that issues item k+1's MFMAs before consuming item k's accumulators, with iglp_opt(0) for WPB ≤ 2.
• lds_scales (kv_scales staged into LDS by the same async DMA), tile_skip (per-wave skip of tiles outside its rows' windows) and dma_guard (no prefetch past the last tile). Also a prefetch_depth tag.
• KV staging now uses the fx.rocdl.cdna4.BufferLoadAsyncLDS{128b,32b} copy atoms. The bank-conflict swizzle is a SwizzleType composed layout. It is applied on both the DMA source and the read view, so it must be its own inverse (S >= B). Its width and row-bit position follow the head size, so it also holds for head_size > 256 (verified at 512). The builder asserts this, and asserts that its LDS slots fit the device. A and B fragments are tiled copies derived from the atom's tiled MMA (make_tiled_copy_A/B). This removes buffer_ops and create_llvm_ptr from this kernel.
• Empty-window fix in both builders: a non-live row (past the end, or cu_ends <= cu_starts after clamping) takes the window [seq_len_kv, 0), so it no longer stretches the block's union window.

Shared helpers (aiter/ops/flydsl/kernels/kernels_common.py)

• warp_reduce_strided and warp_reduce_scatter_strided are MFMA-accumulator head reductions with fx.coop-style signatures plus stride.
• FlyDSL 0.4.1's fx.coop cannot express them. All of its warp collectives (warp_reduce, warp_reduce_batched[_to_blocked|_to_striped], and the ROCm DPP override of warp_reduce) fold contiguous lane groups (XOR 1..width/2), while the head reduce must fold lane bits at and above MFMA_N. permlane{16,32}_swap is also exposed only as the raw ROCDL op. The helpers keep fx.coop signatures so they can be swapped once fx.coop gains a lane-stride option.

Test (op_tests/test_flydsl_fp8_mqa_logits.py)

• verify also pins every routed gfx950 variant (and gfx942's auto variants) on a small odd shape across num_heads, head_dim, clean_logits and the causal/past_end windows. These are the production-variant pins.
• New long-context indexer shapes: 512×32768, 1024×65536, 2048×131072.
• Bench rows record the resolved variant.
• New --shapes-csv option: sweep shapes from a CSV. Per-row num_heads/head_dim/window/batch columns pin those axes for the row; otherwise the row sweeps the CLI axes.
• New -o/--output option: also write the summary tables and the speedup line to a markdown file.
• The production-variant pins also run at head_dim 512 (_PRODUCTION_EXTRA_HEAD_DIMS), skipping variants whose atom cannot tile the problem or whose LDS slots do not fit.
• verify_auto_routed_head_dims (gfx950): at head_dim 512, for 5 shapes on every side of the route bounds and every num_heads, _auto_variant must pick a route that tiles the problem and fits LDS, and the unforced op must match the reference. On bfe265224 all 15 cases fail grading (calc_diff 0.33-0.43, or an occupancy-query error from LDS overflow); they pass here.

Performance

Hardware: AMD Instinct MI350X (gfx950, 256 CUs), ROCm 10.0, torch 2.13.0, Triton 3.8.0, FlyDSL 0.4.1.

Workload: op_tests/test_flydsl_fp8_mqa_logits.py --scenario bench --full with the PR branch's default shape list, run on both trees:

• Shapes (s_q×s_kv): 1×1, 1×16, 1×113, 17×76, 61×113, 61×1024, 128×1024, 1024×1024, 1024×1560, 64×2048, 64×8192, 256×8192, 257×8192, 128×64, 1024×1000, 512×32768, 1024×65536, 2048×131072.
• Other axes: num_heads {32, 64, 128} × head_dim {64, 128} × windows {causal, cp, misaligned, empty, past_end} × clean_logits {0, 1}, fp8 e4m3fn. That is 960 cases per tree.

Measurement:

• Each case is graded against the fp32 reference, then timed with HIP-graph replay. Every replay is preceded by an LLC flush. Each time is the median over 20 samples of the mean of 50 replays.
• The sweep was split into 6 shards by (num_heads, head_dim), one per GPU. Main used the same shard-to-GPU mapping.
• PR: commit bfe265224. Main: 720b36ce0.
• "vs Gluon" uses the Gluon times from the PR run.
• All 960 cases passed grading on both trees.

Each table gives the geomean speedup (>1 means the left-hand kernel is faster).

clean_logits=True

num_heads	head_dim	cases	PR FlyDSL vs Gluon	PR FlyDSL vs main FlyDSL	main FlyDSL vs Gluon (ref.)
32	64	80	1.51	1.30	1.16
32	128	80	1.47	1.32	1.11
64	64	80	1.42	1.36	1.05
64	128	80	1.46	1.34	1.09
128	64	80	1.74	1.41	1.23
128	128	80	3.24	1.37	2.37
all	all	480	1.72	1.35	1.28

clean_logits=False

num_heads	head_dim	cases	PR FlyDSL vs Gluon	PR FlyDSL vs main FlyDSL	main FlyDSL vs Gluon (ref.)
32	64	80	1.28	1.38	0.93
32	128	80	1.23	1.39	0.88
64	64	80	1.14	1.46	0.78
64	128	80	1.17	1.41	0.83
128	64	80	1.47	1.51	0.97
128	128	80	3.18	1.44	2.20
all	all	480	1.46	1.43	1.02

Notes:

• Excluding past_end, where main hits the union-window bug, PR vs main is 1.21× (clean_logits=True) and 1.23× (clean_logits=False) over 372 cases each.

• Long-context shapes (s_kv ≥ 32k, 90 cases per value):

  clean_logits	PR vs Gluon	PR vs main FlyDSL	main FlyDSL vs Gluon
  True	1.41	2.11	0.67
  False	1.50	2.94	0.51

• Coverage of wins:

  • The PR is faster than Gluon on 473 of 480 cases (clean_logits=True) and 375 of 480 (clean_logits=False).
  • It is faster than main on 458 of 480 and 464 of 480.
  • The worst case against main is 0.93×: 1024×65536 H128 D128 empty window, clean_logits=True, 79.1 µs vs 73.3 µs, dominated by the -inf fill. It is one of 3 cases below 0.97×.
  • The worst cases against Gluon are 0.65–0.81×. All are long-KV misaligned windows (1024×65536, at both clean_logits values) or clean_logits=False misaligned/empty windows at 1024×65536, 2048×131072 and 1024×1560.

• Per-window breakdown (geomean):

  window	cl=True vs Gluon	cl=True vs main	cl=False vs Gluon	cl=False vs main
  causal	1.90	1.24	1.60	1.24
  cp	1.94	1.40	1.71	1.44
  empty	1.57	1.09	1.30	1.12
  misaligned	1.65	1.24	1.34	1.25
  past_end	1.68	1.94	1.54	2.42

Testing

• Unit tests added or updated
• Performance benchmarks run
• Tested on MI250X
• Tested on MI300X

Tested on MI350X (gfx950) only; see "Validation and unverified scope".

Documentation

• Docstrings updated
• User guide updated
• Performance guide updated
• No documentation change is needed, with the reason stated below

The variant tag grammar, the route table and every new builder option are documented in docstrings and comments next to the code. There is no user-guide page for this op.

Dependencies

• No new third-party dependencies added
• New dependencies are listed and justified below

Breaking Changes

None for the public op (flydsl_fp8_mqa_logits). For gfx950 only:

• The fixed gfx950 variant registry is gone, so the old tags such as mfma32x32x64_bkv64_r1_w2_lds3 now resolve through the tag grammar. Any tag in that grammar still builds.
• KERNEL_VARIANTS lists the routed tags.
• The swizzle= builder argument is removed; the swizzle is always on.

PR Self-Review

Reviewed revision

067ae36

Behavior and intent

• gfx950 dispatch:
  • Variant selection moves from a hand-written heuristic to a route table tuned by search over the builder knob space.
  • num_splits is now sized by the occupancy-aware wave-cost model.
  • Main's heuristic lost to Gluon on about half the sweep at clean_logits=False.
• New kernel options: head_acc, rs_head, sw_pipe, lds_scales, tile_skip, dma_guard and prefetch_depth. They reduce VALU work per head term and the cost of the head reduce, and they overlap the MFMAs with the epilogue.
• Empty and past-the-end windows: they no longer widen a block's union KV window. On main, such rows made the whole block scan up to seq_len_kv for nothing (up to ~100× slower at s_kv = 128k), although the output was still correct.
• gfx942: the empty-window fix and the head_acc parameter (default add) also touch the gfx942 builder. The gfx942 routes and split policy are unchanged.

Known limitation: cold first launch inside graph capture

This carries over the limitation already stated in #5921 for gfx942; it now applies to gfx950 as well.

• num_splits comes from the kernel's resident blocks per CU, read once per launcher with FlyDSL's max_blocks_per_cu() and cached.
• That query loads a scratch copy of the kernel module, queries it, and unloads it. HIP rejects hipModuleUnload during stream capture (hipErrorStreamCaptureUnsupported), and the error invalidates the capture.
• Compiling the kernel and loading its module are both allowed during capture. A cold launch inside capture therefore works; only the occupancy query is skipped.
• When a launcher's first call happens inside graph capture, the split falls back to the fixed cu_oversub=4 target (main's gfx950 behavior), and the captured graph keeps that split. Output is unaffected; only the split, and therefore performance, can differ from the tuned one.
• To get the tuned split in a captured graph, run the op once eagerly at the same (num_heads, head_size, clean_logits, variant) before capture; a warmup run before capture is enough.
• A fix without the fallback needs FlyDSL to answer the occupancy query without unloading a module during capture, e.g. from the module the launcher already loaded.

Validation and unverified scope

Validated on MI350X (gfx950), FlyDSL 0.4.1:

• FLYDSL_RUNTIME_ENABLE_CACHE=0 python op_tests/test_flydsl_fp8_mqa_logits.py --scenario verify: all 815 case runs passed. That is 80 reduced-set cases, 720 production-variant pins (including 204 at head_dim 512) and 15 auto-routed head-dim cases.
• python op_tests/test_flydsl_fp8_mqa_logits.py --scenario bench --full -s <18 shapes> --num-heads N --head-dim D on the PR (bfe265224) and on main 720b36ce0: 960 of 960 cases passed grading on both. The timings are in the Performance section.
• black --check and ruff check pass on all changed files.

Unverified:

• gfx942 (MI300X/MI325X) was not run in this validation, although the gfx942 builder received the empty-window fix and the head_acc parameter (default unchanged). CI's gfx942 op-test lane covers it.
• Routes are tuned at each class's own (num_heads, head_size). Other sizes that map to a class run untuned routes. head_size 512 is verified for correctness only, not tuned or timed; head sizes above 512 are not tested.
• For one route, mfma32x32x64_bkv64_r1_w1_lds4_pd1_swp_rs_hfma_wpe3 at H64 D128, the compiler warns that waves_per_eu=3 is not reached (occupancy 1). It was still the measured best for its window.

PR-specific question

What would each changed op test fail on before this PR, and why is that coverage not already provided?

• Production-variant pins (verify_production_variants):
  • They force every variant that _auto_variant can route to onto a small odd shape (67×1500: a clamped tail tile and a partial row block). They run across num_heads, head_dim, clean_logits and the causal/past_end windows.
  • Before this PR the routed variants (_rs, _swp, _ls, _ts, _dg, _h*, _wpe*) did not exist. Main has no counterpart.
  • Without the pins, the shape-dependent router reaches a given variant only by chance in the reduced/full sweeps, so a broken epilogue in a rarely routed variant would go unnoticed.
• Long-context shapes (512×32768, 1024×65536, 2048×131072): these are the only shapes reaching the long-KV grid.y split range and the routes tuned for s_kv > 4096.
  • On main they pass grading but expose the union-window bug in the bench: past_end at 2048×131072 is 1590–5815 µs on main versus 19–61 µs here.
  • No existing shape is long enough to show it.
• variant column, --shapes-csv, -o: reporting and input plumbing; they do not change coverage.

Changed op tests

op_tests/test_flydsl_fp8_mqa_logits.py remains the single canonical FlyDSL fp8_mqa_logits correctness and perf test:

• The production-variant verify table, which pins every routed variant.
• Three long-context shapes in the default sweep.
• A variant column in the bench table.
• --shapes-csv, to bench production shape dumps through the same grading and timing path instead of a second script.
• -o/--output, to save the summary tables.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 7412 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[aiter-bot]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 067ae36d288e · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

- **gfx950 dispatch:**
  - Variant selection moves from a hand-written heuristic to a route table tuned by search over the builder knob space.
  - `num_splits` is now sized by the occupancy-aware wave-cost model.
  - Main's heuristic lost to Gluon on about half the sweep at clean_logits=False.
- **New kernel options:** `head_acc`, `rs_head`, `sw_pipe`, `lds_scales`, `tile_skip`, `dma_guard` and `prefetch_depth`. They reduce VALU work per head term and the cost of the head reduce, and they overlap the MFMAs with the epilogue.
- **Empty and past-the-end windows:** they no longer widen a block's union KV window. On main, such rows made the whole block scan up to `seq_len_kv` for nothing (up to ~100× slower at s_kv = 128k), although the output was still correct.
- **gfx942:** the empty-window fix and the `head_acc` parameter (default `add`) also touch the gfx942 builder. The gfx942 routes and split policy are unchanged.

Validation and unverified scope

Validated on MI350X (gfx950), FlyDSL 0.4.1:
- `FLYDSL_RUNTIME_ENABLE_CACHE=0 python op_tests/test_flydsl_fp8_mqa_logits.py --scenario verify`: all 815 case runs passed. That is 80 reduced-set cases, 720 production-variant pins (including 204 at head_dim 512) and 15 auto-routed head-dim cases.
- `python op_tests/test_flydsl_fp8_mqa_logits.py --scenario bench --full -s <18 shapes> --num-heads N --head-dim D` on the PR (`bfe265224`) and on main `720b36ce0`: 960 of 960 cases passed grading on both. The timings are in the Performance section.
- `black --check` and `ruff check` pass on all changed files.

Unverified:
- gfx942 (MI300X/MI325X) was not run in this validation, although the gfx942 builder received the empty-window fix and the `head_acc` parameter (default unchanged). CI's gfx942 op-test lane covers it.
- Routes are tuned at each class's own (num_heads, head_size). Other sizes that map to a class run untuned routes. head_size 512 is verified for correctness only, not tuned or timed; head sizes above 512 are not tested.
- For one route, `mfma32x32x64_bkv64_r1_w1_lds4_pd1_swp_rs_hfma_wpe3` at H64 D128, the compiler warns that `waves_per_eu=3` is not reached (occupancy 1). It was still the measured best for its window.

PR-specific question

What would each changed op test fail on before this PR, and why is that coverage not already provided?

- **Production-variant pins (`verify_production_variants`):**
  - They force every variant that `_auto_variant` can route to onto a small odd shape (67×1500: a clamped tail tile and a partial row block). They run across num_heads, head_dim, clean_logits and the causal/past_end windows.
  - Before this PR the routed variants (`_rs`, `_swp`, `_ls`, `_ts`, `_dg`, `_h*`, `_wpe*`) did not exist. Main has no counterpart.
  - Without the pins, the shape-dependent router reaches a given variant only by chance in the reduced/full sweeps, so a broken epilogue in a rarely routed variant would go unnoticed.
- **Long-context shapes (512×32768, 1024×65536, 2048×131072):** these are the only shapes reaching the long-KV grid.y split range and the routes tuned for s_kv > 4096.
  - On main they pass grading but expose the union-window bug in the bench: past_end at 2048×131072 is 1590–5815 µs on main versus 19–61 µs here.
  - No existing shape is long enough to show it.
- **`variant` column, `--shapes-csv`, `-o`:** reporting and input plumbing; they do not change coverage.

Changed op tests

`op_tests/test_flydsl_fp8_mqa_logits.py` remains the single canonical FlyDSL fp8_mqa_logits correctness and perf test:
- The production-variant verify table, which pins every routed variant.
- Three long-context shapes in the default sweep.
- A `variant` column in the bench table.
- `--shapes-csv`, to bench production shape dumps through the same grading and timing path instead of a second script.
- `-o/--output`, to save the summary tables.

PR state

• WAIT — pull request is still a draft

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot did not review the current head SHA

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• PASS — Aiter Test Gate: passed

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}