[sammysun0711]

Motivation

This PR aims to integrate the native gfx950 FP8 paged-prefill attention with asymmetric Q/K and V head dimensions and packed variable-length sequence support from ROCm/FlyDSL#1066 into AITER.

AITER already supports FlyDSL dense FP8 attention through ROCm/aiter#5326, but paged prefill still needs a native path that consumes the serving both kv cache layout used by sglang SGLANG_AITER_KV_CACHE_LAYOUT directly:

• SGLANG_AITER_KV_CACHE_LAYOUT=nhd: Linear KV-cache layout with page size 1
• SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d: Vectorized 5D KV-cache layout with page size 16/64/1024

Technical Details

• Add native causal packed-varlen FP8 paged attention for gfx950 based on FlyDSL 0.3.4.1 runtime.
• Support D128/V128, D192/V128, and D192/V192 with BF16 output.
• Support page-1 linear/linear3d caches and page-16/64/1024 vectorized caches.
• Accept both block-table metadata and native CSR metadata (kv_indptr, kv_page_indices, and kv_last_page_lens).
• Add public API aiter.ops.flydsl.flydsl_flash_attn_paged_prefill_func for downstream farmwork integration.
• Preserve preallocated output, custom stream, graph replay, runtime softmax scale, ragged requests, odd page counts, and large cache addressing contracts, LSE.
• Share common dense/paged FP8 pipeline helpers and split paged GEMM, LDS, softmax, and epilogue helpers into focused modules.
• Add correctness, API, dispatch, metadata, stream, graph, and regression coverage to the standard AITER test shard.

Performance

MI355X (gfx950) with ROCm/HIP 7.2.26015, Torch 2.9.1+rocm7.2.0, B=1, Hq=16, Hkv=1, causal Q:KV=1:2, page size 64, vectorized FP8 K/V, and preallocated BF16 output.

Q / KV	QK / V	FlyDSL paged (ms / TFLOPS)	AITER migrated paged (ms / TFLOPS)	Port delta	AITER public CSR (ms / TFLOPS)	AITER dense (ms / TFLOPS)	Paged/dense
4K / 8K	D128/V128	0.1341 / 1537.25	0.1355 / 1522.09	+1.00%	0.1276 / 1615.36	0.1326 / 1554.60	103.91%
16K / 32K	D128/V128	1.5121 / 2181.42	1.5274 / 2159.60	+1.01%	1.5315 / 2153.78	1.5860 / 2079.82	103.56%
32K / 64K	D128/V128	5.9994 / 2199.26	6.0433 / 2183.28	+0.73%	6.0692 / 2173.98	6.2791 / 2101.29	103.46%
64K / 128K	D128/V128	23.7516 / 2222.03	24.0004 / 2199.00	+1.05%	23.9735 / 2201.46	24.8636 / 2122.65	103.71%
4K / 8K	D192/V128	0.1675 / 1538.61	0.1674 / 1539.25	-0.04%	0.1609 / 1602.08	0.1566 / 1645.70	97.35%
16K / 32K	D192/V128	1.9369 / 2128.83	1.9389 / 2126.59	+0.11%	1.9397 / 2125.67	1.9014 / 2168.56	98.02%
32K / 64K	D192/V128	7.6633 / 2152.20	7.6376 / 2159.43	-0.34%	7.6390 / 2159.04	7.4653 / 2209.26	97.73%
64K / 128K	D192/V128	30.3289 / 2175.18	30.3021 / 2177.11	-0.09%	30.2322 / 2182.14	29.4401 / 2240.86	97.38%
4K / 8K	D192/V192	0.2395 / 1291.12	0.2401 / 1288.26	+0.22%	0.2343 / 1320.05	0.2459 / 1257.66	104.96%
16K / 32K	D192/V192	2.6168 / 1890.80	2.6119 / 1894.37	-0.19%	2.6238 / 1885.79	2.8425 / 1740.67	108.34%
32K / 64K	D192/V192	10.4296 / 1897.61	10.3670 / 1909.08	-0.60%	10.4193 / 1899.50	11.1327 / 1777.78	106.85%
64K / 128K	D192/V192	41.9932 / 1885.19	41.8462 / 1891.81	-0.35%	42.0344 / 1883.35	44.3850 / 1783.60	105.59%

The FlyDSL source column uses PR #1066 commit 8a662b0 with its locally built FlyDSL 0.3.3 runtime. The AITER migrated columns use commit c933927e with FlyDSL 0.3.2.

Port delta is AITER paged attention with direct block-table latency versus FlyDSL source paged attention latency. Paged/dense uses the AITER public CSR page metadata path (used by aiter.ops.mha.mha_batch_prefill_func) versus AITER FlyDSL dense attention; values above 100% mean paged is faster.

Key results:

• The migrated paged kernel is within -0.60% to +1.05% latency of the current FlyDSL source implementation across all 12 shapes.
• AITER public CSR dispatch is within 0.45% of the direct migrated CSR kernel.
• Paged attention is 3.46% to 3.91% faster than dense attention for D128/V128 and 4.96% to 8.34% faster for D192/V192.
• For D192/V128, page attention is 1.98% to 2.65% behind the Flydsl dense attention in AITER.

Test Plan

python3 -m pytest -q op_tests/test_flydsl_paged_fmha.py

Performance coverage:

• Compare current FlyDSL PR source with the direct AITER port using block-table metadata.
• Compare direct AITER block-table and CSR paths.
• Compare the AITER public CSR dispatcher with the direct CSR kernel.
• Compare AITER paged CSR with AITER FlyDSL dense attention using identical logical Q/K/V tensors.
• Cross-check source and port output in separate processes to avoid copied JIT-kernel cache collisions.

Test Result

.................................................................................................................... [ 54%]
...............................................................................................                      [100%]
===================================================== warnings summary =====================================================
../../../../opt/venv/lib/python3.10/site-packages/Cython/Distutils/old_build_ext.py:15
../../../../opt/venv/lib/python3.10/site-packages/Cython/Distutils/old_build_ext.py:15
  /opt/venv/lib/python3.10/site-packages/Cython/Distutils/old_build_ext.py:15: DeprecationWarning: dep_util is Deprecated. Use functions from setuptools instead.
    from distutils.dep_util import newer, newer_group

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
211 passed, 2 warnings in 141.51s (0:02:21)

• Direct AITER block-table, direct CSR, and public CSR outputs are bit-identical for all 12 benchmark shapes.
• Paged versus dense passes rtol=0.02, atol=0.02; maximum absolute difference is 0.001953125, and minimum row cosine is 0.9999994.
• Flydsl FP8 paged attention vs AITER migrated paged attention checks at 4K/8K pass for all supported dimension pairs. D128/V128 is bit-identical; D192 maximum absolute difference is 0.001953125, with minimum row cosine at least 0.99999958.

Reviewed revision

fdd594f

Behavior and intent

Add the explicit flydsl_flash_attn_paged_prefill_func API for gfx950 causal
FP8 prefill with packed queries, BF16 output and optional FP32 LSE.

• Support page-1 linear or page-16/64/1024 vectorized KV caches through rectangular or CSR metadata
• Supporting Q/K–V pairs 128/128, 192/128 and 192/192.
• Avoids mandatory dense KV materialization for the existing dense/varlen interface.

Validation and unverified scope

MI350X/gfx950, 256 CUs, ROCm 7.2, PyTorch 2.9.1, FlyDSL 0.4.1:
python -m pytest -q op_tests/test_flydsl_paged_fmha.py: 140 passed, no failures or skips.

Unverified for this revision: downstream model integration.

PR-specific question

The CI change adds one literal filename to the existing pytest route in
.github/scripts/aiter_test.sh and a 240-second sharding estimate in
.github/split-test-times/aiter.tsv. It adds no workflow event, token permission, credential or runner configuration.

Changed op tests

op_tests/test_flydsl_paged_fmha.py is the only added op-test file. Existing
dense FMHA tests cannot exercise shuffled page IDs, CSR ranges, page tails or
independent paged V addressing. All categories below belong to this PR's new
test file; each parameterized case is counted once.

Added test category	Cases	Purpose
Core paged layouts and ragged requests	15	Check all supported cache layouts and Q/K–V width pairs against FP32 attention.
LSE and caller-owned buffers	18	Check LSE, empty requests, buffer validation and output canaries.
Metadata contracts and normalization	10	Check optional page-one lengths, strided metadata, scalar descales and rectangular-table precedence.
Addressing modes and page-ID bounds	14	Compare bounded/wide pointers, byte-offset bases and scalar/paired page lookup.
Graph replay and stream lifetime	18	Check updated metadata on replay, stream ordering and copy-source lifetime.
Probability rescaling and runtime scale	18	Check lazy/eager parity, rescaling boundaries and scale changes without rebuilding.
Poisoned tails, shared pages and output guards	30	Check inactive NaN isolation, shared pages and writes bounded to valid output rows.
Large addresses and integer overflow	11	Exercise >4 GiB offsets, signed-int32 edges and output-address wraparound.
Explicit compile/preload	6	Verify compilation does not execute the kernel or modify output buffers.
Total	140	All cases passed.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[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; main branch always runs both MI35X and MI300X
multigpu	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: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
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 5556 --add-label <label>

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 opt this PR out.

[Copilot]

🟡 Changes recommended

The unresolved FlyDSL import breaks AITER_TRITON_ONLY=1 fallback behavior.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds native gfx950 FlyDSL FP8 paged-prefill attention with asymmetric head dimensions, variable-length support, multiple cache layouts, and expanded testing.

Changes:

• Adds paged FP8 kernels, metadata handling, GEMMs, softmax, and output epilogues.
• Integrates FlyDSL dispatch and public APIs.
• Adds correctness, stream, graph, regression, and CI coverage.

Critical finding: aiter/ops/mha.py unconditionally imports FlyDSL, causing ModuleNotFoundError in AITER_TRITON_ONLY=1 builds instead of using the existing fallback.

File summaries

File	Summary
op_tests/test_flydsl_paged_fmha.py	Adds comprehensive paged FP8 correctness and regression tests.
aiter/ops/mha.py	Integrates public FP8 dispatch.
aiter/ops/flydsl/kernels/fmha_gfx950/pipeline.py	Refactors shared dense and paged pipeline helpers.
aiter/ops/flydsl/kernels/fmha_gfx950/paged_pipeline.py	Adds paged kernel context and metadata handling.
aiter/ops/flydsl/kernels/fmha_gfx950/paged_op_softmax.py	Implements paged masking and softmax.
aiter/ops/flydsl/kernels/fmha_gfx950/paged_op_lds.py	Implements KV-cache staging and LDS movement.
aiter/ops/flydsl/kernels/fmha_gfx950/paged_op_gemm.py	Implements paged FP8 GEMMs.
aiter/ops/flydsl/kernels/fmha_gfx950/paged_op_epilog.py	Implements BF16 output stores.
aiter/ops/flydsl/kernels/fmha_gfx950/flash_attn_paged_fp8_gfx950.py	Builds the native gfx950 paged kernel.
aiter/ops/flydsl/kernels/fmha_gfx950/common.py	Provides shared memory and buffer helpers.
aiter/ops/flydsl/kernels/flash_attn_paged_fp8_func_gfx950.py	Provides public validation and launch wrapping.
aiter/ops/flydsl/fmha_kernels.py	Adds capability gating and paged-attention dispatch.
aiter/ops/flydsl/__init__.py	Exports the paged API.
.github/scripts/split_tests.sh	Registers test shard timing metadata.
.github/scripts/aiter_test.sh	Configures execution of the new pytest coverage.

Review details

• Files reviewed: 15/15 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Four unresolved review findings remain, including three critical issues and a FlyDSL version-gating issue.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

aiter/ops/flydsl/init.py:46

• This new lazy export is backed by code explicitly tied to FlyDSL 0.3.2 (requirements.txt pins that release), but the package-level gate still accepts FlyDSL 0.2.4 and later. On a 0.2.4–0.3.1 installation, accessing this attribute can fail inside the paged module instead of falling back; raise _MIN_FLYDSL_VERSION to 0.3.2 or version-gate this path.

    "flydsl_flash_attn_paged_fp8_func": (
        ".kernels.flash_attn_paged_fp8_func_gfx950",
        "flydsl_flash_attn_paged_fp8_func",

• Files reviewed: 15/15 changed files
• Comments generated: 3
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

A verified critical correctness issue and additional review findings require follow-up.

Review details

Suppressed comments (2)

aiter/ops/flydsl/kernels/fmha_gfx950/paged_op_lds.py:172

• For the supported D192/V128 specialization, NUM_DMA_K=2 and _load_vectorized() still runs on all 8 waves, although a 64×192 tile has only (64*192)/(2*64*16)=6 valid wave slots. Waves 6–7 therefore issue zero/OOB-padded K loads that are never consumed, wasting about 25% of K-staging DMA work on every tile. Apply the same copy_waves guard used by the V192 path to skip these loads while keeping all waves in the rendezvous schedule.

            if const_expr(traits.HEAD_DIM_V == 192):

                @flyc.jit
                def _load_compact():
                    copy_waves = (traits.BLOCK_N * traits.HEAD_DIM) // (

aiter/ops/mha.py:4106

• [inferred] This unconditional branch makes the new kernel the default for every supported mha_batch_prefill_func call, but the performance table compares the port with the FlyDSL source rather than with the pre-change AITER/CK path. That leaves a possible regression in the production dispatcher unmeasured even though the new test passes. Author must add a reproducible base-versus-head benchmark for the public dispatcher (or explicitly limit the performance claim to source-port parity).

    if q.dtype == torch.float8_e4m3fn:
        try:
            from .flydsl.fmha_kernels import flydsl_flash_attn_batch_prefill_func
        except ModuleNotFoundError as exc:
            if exc.name != "flydsl":
                raise
        else:
            flydsl_result = flydsl_flash_attn_batch_prefill_func(

• Files reviewed: 15/15 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Three unresolved findings remain: two moderate fallback issues and one nit involving D192/V128 staging.

Review details

Suppressed comments (3)

aiter/ops/flydsl/fmha_kernels.py:317

• This dispatcher is called before the existing CK fallback, but it does not gate the native wrapper's CSR precondition for page sizes greater than one: flydsl_flash_attn_paged_fp8_func rejects kv_last_page_lens=None for those layouts. A public mha_batch_prefill_func call without that optional metadata therefore raises from the native path instead of returning None and falling back. Check this condition before dispatch (while retaining the native path when last-page lengths are supplied).

    try:
        _cache_geometry(q, k, v)
    except NotImplementedError:
        return None

aiter/ops/flydsl/kernels/fmha_gfx950/paged_op_lds.py:170

• [verified] In the supported D192/V128 specialisation, HEAD_DIM_V == 192 is false, so this loader runs all eight waves and two DMA bands even though D192 has only 12 valid 16-byte D-groups. The valid mask turns groups 12–15 into zero/sentinel loads, but those four groups still consume roughly 25% extra K staging traffic on every tile; only D192/V192 takes the six-wave compact path. Author must either compact K staging based on HEAD_DIM == 192 as well or provide a measured reason this extra traffic is required.

        if const_expr(traits.HEAD_DIM_V == 192):

            @flyc.jit
            def _load_compact():
                copy_waves = (traits.BLOCK_N * traits.HEAD_DIM) // (
                    self.NUM_DMA_K * traits.WARP_SIZE * traits.KV_VEC_SIZE
                )
                if self.wave_id_uni < fx.Int64(copy_waves):
                    _load_vectorized()

            _load_compact()

aiter/ops/mha.py:4104

• The optional-backend guard only handles a missing flydsl module. Importing .flydsl.fmha_kernels also executes aiter/ops/flydsl/__init__.py, which raises ImportError for an unavailable/unsupported FlyDSL version (__init__.py:15-24); in that environment every FP8 batch-prefill call now raises here instead of reaching CK/Triton. Handle that version-check failure as an unavailable optional backend while preserving unrelated import errors.

        except ModuleNotFoundError as exc:
            if exc.name != "flydsl":
                raise

• Files reviewed: 15/15 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Five review findings remain, including runtime architecture dispatch consistency and a duplicate CI shard entry.

3 open findings

• Use shared get_gfx_runtime helper for architecture dispatch
• Use get_gfx for test architecture selection
• Remove duplicate test timing entry

1 resolved since last review

• Duplicate case arm makes aiter_test.sh syntactically invalid

Previously missed (1)

In code that hasn't changed since last review

Duplicate test row causes duplicate sharding and execution

.github/​split-test-times/​aiter.tsv:39

This added row duplicates the existing op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py entry at line 41. The shard script indexes rows rather than deduplicating paths, so that test is assigned and run twice; remove the duplicate row to avoid inflating CI time and distorting shard timing.

🧠 Review effort: Lite

[Copilot]

🟡 Changes recommended

The paged pipeline has a critical unsafe block-table load, and runtime architecture probing is inconsistent with the shared dispatch contract.

2 open findings

• Guard block-table loads for odd or empty page counts · New
• Use get_gfx for test architecture selection

2 resolved since last review

• Use shared get_gfx_runtime helper for architecture dispatch
• Remove duplicate test timing entry

🧠 Review effort: Lite

[Copilot]

🟡 Changes recommended

Align the kernel builder’s runtime architecture check with AITER’s shared get_gfx_runtime() contract.

2 open issues

• Use AITER's get_gfx_runtime() for runtime architecture gating · New
• Use get_gfx for test architecture selection

1 resolved since last review

• Guard block-table loads for odd or empty page counts

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

A critical current-stream propagation issue remains unresolved in the paged kernel launcher.

1 open issue

• Pass resolved current stream to FlyDSL launcher · New

2 resolved since last review

• Use AITER's get_gfx_runtime() for runtime architecture gating
• Use get_gfx for test architecture selection

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

Approval should be deferred because this large gfx950-specific FlyDSL kernel depends on MI355X/ROCm runtime validation of its LDS/MFMA scheduling, addressing, and graph/stream lifetime behavior, which is unavailable in the review environment.

No open issues 🎉

1 resolved since last review

• Pass resolved current stream to FlyDSL launcher

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

The gfx950 FlyDSL kernel’s low-level LDS, DMA, MFMA, pointer-boundary, graph/stream, and performance behavior requires independent execution-based validation on the target MI355X/ROCm environment.

No open issues 🎉

_{Review effort: Lite}

[aiter-bot]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision 8b0cede41e59 · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

1. What behavior does this PR change, and why is the existing behavior insufficient?
2. What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
3. What permissions and untrusted-input boundary does this CI change have, and can fork PR code or text reach a privileged execution path?
4. For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

• FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA

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

[Copilot]

🟡 Changes recommended

The new LDS DMA path uses the undocumented literal address-space value 2 instead of the established fx.AddressSpace.Shared enum, creating an avoidable API and backend-mapping risk.

1 open issue

• Use fx.AddressSpace.Shared instead of literal address-space value 2 · New

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

Human assessment is needed because a critical verified issue allows BF16 output storage to exceed the kernel descriptor’s 2 GiB byte limit despite the element-count check, potentially causing incorrect memory addressing in both wrapper paths.

No open issues 🎉

1 resolved since last review

• Use fx.AddressSpace.Shared instead of literal address-space value 2

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

A critical active issue remains because the fixed-base output descriptor can wrap at 2 GiB and silently misaddress valid BF16 outputs, while the current guard is too high and the element-count check does not prevent it.

1 open issue

• BF16 output descriptor can wrap above 2 GiB · New

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

A critical masking defect remains in the grouped K staging path, allowing padded tokens from partial or zero-length pages to affect softmax normalization and LSE.

1 open issue

• Mask invalid tokens in grouped K staging · New

1 resolved since last review

• BF16 output descriptor can wrap above 2 GiB

_{Review effort: Lite}