[shay-li77]

Summary

Adds a gfx1250 (MI450) ASM forward kernel for FP8 attention: e4m3fn q/k/v with per-tensor fp32 descales, qk head dim 192 / v head dim 128, packed varlen (THD) layout, bf16 output, optional LSE, causal (bottom-right aligned) or not. It is reached through flash_attn_varlen_func (with q/k/v_descale), flash_attn_varlen_fp8_pertensor_func, and a new public wrapper aiter.fmha_fwd_fp8_varlen_asm.

A second commit replaces the per-site eligibility checks of the gfx1250 varlen ASM kernels (the existing bf16 one and the new fp8 one) with a single dispatch gate.

Motivation

• MLA-style prefill (qk 192 / v 128) with many heads and long sequences, e.g. 96 heads with 12K–16K tokens, currently runs on the bf16 D192x128 ASM kernel on gfx1250, or on the existing aiter fp8 baseline kernel. On the target shape (batch 1, 12288 / 14336 tokens, 96 heads), the new kernel is 2.76x / 2.60x faster than the bf16 ASM and 1.28x / 1.39x faster than the fp8 baseline (non-causal / causal). Across the 96-head shapes with 4K or more query tokens, it is 2.2–2.8x and 1.26–1.79x faster (table below).
• The eligibility of the gfx1250 varlen ASM kernels was written out at every site that routes to them: the bf16 kernel twice (the flash_attn_varlen_func pre-router that runs ahead of FlyDSL, and _flash_attn_varlen_forward), and the fp8 kernel would have needed three copies (plus flash_attn_varlen_fp8_pertensor_func with CK disabled). The copies used different argument forms and had already drifted apart.

Changes

• Kernels (hsa/gfx1250/fmha_fwd_fp8_varlen/): non-causal and causal code objects, and their csv.
• C++ entry (csrc/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu, ctypes module module_fmha_fwd_fp8_varlen_asm):
  • Validates the operands and uses strided q/k/v/out views in place.
  • Picks the launch configuration from the problem size and the CU count; AITER_FMHA_FP8_VARLEN_ASM_OPT overrides it.
• Python (aiter/ops/mha.py):
  • Public wrapper fmha_fwd_fp8_varlen_asm(q, k, v, q_descale, k_descale, v_descale, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, softmax_scale=None, is_causal=False, return_lse=False, out=None).
  • Routing from flash_attn_varlen_func (when all three descales are given; checked before FlyDSL, like the bf16 gfx1250 ASM), from flash_attn_varlen_fp8_pertensor_func (with or without CK), and from _flash_attn_varlen_forward. The LSE comes back in the varlen API layout (nheads, total_q).
  • _select_gfx1250_fmha_fwd_varlen_asm decides for both gfx1250 varlen ASM kernels and returns the one that serves the call, or 0. Each routing site calls it once, with the same arguments, behind a cached device check, so other architectures skip it. It costs less host time than the per-site checks it replaces, on gfx1250 and on other architectures.
  • The bf16 "offer to FlyDSL first" policy (non-causal D64/D128, AITER_ENABLE_EXPERIMENTAL) is unchanged and kept apart in _prefer_flydsl_over_gfx1250_bf16_varlen_asm.
  • aiter/jit/optCompilerConfig.json: the new module.
• Op test (op_tests/test_fmha_fwd_with_sink_varlen_asm.py): --dtype fp8 sections and gate tables, described under "Changed op tests".

Routing impact:

• fp8 D192x128 calls with per-tensor fp32 descales through flash_attn_varlen_func or flash_attn_varlen_fp8_pertensor_func on gfx1250 now go to this kernel. Calls it does not accept (other head dims, per-head descales, windows, sinks, dropout, soft-cap, paged KV, padded cu_seqlens, ...) keep their previous route.
• gfx1250 bf16 varlen ASM (second commit): windows are now judged everywhere by the CK varlen convention. A negative window, or one at or beyond max_seqlen_k, is unlimited, and causal ignores the right window.
  • The pre-router now also sends the equivalent spellings of "no window", and causal with a right window, to the ASM. It used to leave those to FlyDSL.
  • _flash_attn_varlen_forward no longer sends a left window of 0, or a non-causal right window of 0, to the ASM. The ASM cannot apply those windows.
  • max_seqlen_q == 0 goes to another backend instead of reaching the bf16 C++ entry, which refuses it.

Performance (if applicable)

Measured on MI450 (gfx1250 B0, 256 CUs).

• Kernels compared:
  • fp8 ASM (this PR): flash_attn_varlen_fp8_pertensor_func with e4m3fn q/k/v, quantized per tensor from bf16 data (amax / 448), and fp32 descales.
  • bf16 ASM: flash_attn_varlen_func with the same data in bf16, which runs the existing bf16 D192x128 ASM kernel (fmha_bf16_pertokenBf16_hd192x128_128x128[_mask]_varlen).
  • aiter fp8 baseline: the existing aiter fp8 varlen kernel on gfx1250, timed on the same fp8 data.
• Timing:
  • Attention-kernel device time from the torch profiler.
  • Each kernel first runs alone for 1 s to reach steady clocks and power, then 20 timed iterations.
  • Four rounds alternating the kernels; the table shows the median.
• TFLOPS: both GEMMs, counted over the visible (q, k) pairs.
• Inputs: all sequences in a batch have equal length.

batch	seqlen_q / seqlen_k	heads q/kv	causal	fp8 ASM (this PR)	bf16 ASM	speedup vs bf16 ASM	aiter fp8 baseline	speedup vs fp8 baseline
1	12288 / 14336	96/96	no	2158 us (5015 TF)	5956 us (1817 TF)	2.76x	2757 us (3925 TF)	1.28x
1	4096 / 4096	96/96	no	217 us (4741 TF)	568 us (1814 TF)	2.61x	290 us (3560 TF)	1.33x
1	8192 / 8192	96/96	no	846 us (4873 TF)	2283 us (1806 TF)	2.70x	1091 us (3781 TF)	1.29x
1	16384 / 16384	96/96	no	3278 us (5031 TF)	9117 us (1809 TF)	2.78x	4179 us (3946 TF)	1.27x
1	2048 / 14336	96/96	no	373 us (4836 TF)	971 us (1858 TF)	2.60x	474 us (3807 TF)	1.27x
1	8192 / 14336	96/96	no	1470 us (4908 TF)	3979 us (1813 TF)	2.71x	1850 us (3901 TF)	1.26x
2	12288 / 14336	96/96	no	4355 us (4971 TF)	11979 us (1807 TF)	2.75x	5533 us (3912 TF)	1.27x
1	16384 / 16384	32/32	no	1095 us (5019 TF)	2627 us (2092 TF)	2.40x	1387 us (3963 TF)	1.27x
1	16384 / 16384	64/4	no	2161 us (5088 TF)	4903 us (2242 TF)	2.27x	2748 us (4001 TF)	1.27x
4	4096 / 4096	128/128	no	1208 us (4550 TF)	3321 us (1656 TF)	2.75x	1547 us (3554 TF)	1.28x
12	1024 / 1024	128/128	no	301 us (3426 TF)	636 us (1622 TF)	2.11x	406 us (2539 TF)	1.35x
1	12288 / 14336	96/96	yes	1309 us (4725 TF)	3406 us (1816 TF)	2.60x	1821 us (3396 TF)	1.39x
1	4096 / 4096	96/96	yes	130 us (3971 TF)	286 us (1802 TF)	2.20x	232 us (2217 TF)	1.79x
1	8192 / 8192	96/96	yes	466 us (4423 TF)	1140 us (1808 TF)	2.45x	697 us (2958 TF)	1.50x
1	16384 / 16384	96/96	yes	1746 us (4724 TF)	4600 us (1793 TF)	2.64x	2415 us (3415 TF)	1.38x
1	2048 / 14336	96/96	yes	374 us (4476 TF)	884 us (1895 TF)	2.36x	469 us (3572 TF)	1.25x
1	8192 / 14336	96/96	yes	1079 us (4777 TF)	2832 us (1820 TF)	2.63x	1480 us (3483 TF)	1.37x
2	12288 / 14336	96/96	yes	2617 us (4726 TF)	6877 us (1799 TF)	2.63x	3632 us (3406 TF)	1.39x
1	16384 / 16384	32/32	yes	576 us (4776 TF)	1311 us (2096 TF)	2.28x	828 us (3320 TF)	1.44x
1	16384 / 16384	64/4	yes	1128 us (4874 TF)	2504 us (2196 TF)	2.22x	1427 us (3853 TF)	1.27x
4	4096 / 4096	128/128	yes	722 us (3808 TF)	1673 us (1643 TF)	2.32x	1194 us (2303 TF)	1.65x
12	1024 / 1024	128/128	yes	245 us (2105 TF)	425 us (1214 TF)	1.73x	401 us (1287 TF)	1.64x

Notes:

• Kernel revision: the table was measured with a slightly earlier revision of this PR's kernel. The PR's kernel was re-measured the same way on the batch 1 and batch 2 12288 / 14336 shapes, and agreed within 0.6%.
• fp8 baseline descales: the aiter fp8 baseline was timed in the version that did not yet take descales. The version that applies them was not faster in a later single run.
• Accuracy: compared with the bf16 kernel's output on the same data, both fp8 kernels have cosine similarity 0.9996–0.9998.

Testing

• Unit tests added or updated
• Performance benchmarks run
• Tested on MI250X
• Tested on MI300X
• Tested on MI450 (gfx1250 B0)

Documentation

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

The new public wrapper documents its contract in its docstring. The existing entry points pick the kernel by themselves, so their usage does not change.

Dependencies

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

Breaking Changes

None. No public signature changes; fmha_fwd_fp8_varlen_asm is new. The routing changes are listed under "Changes".

PR Self-Review

Reviewed revision

1e81335

Behavior and intent

• fp8 calls: on gfx1250, fp8 D192x128 varlen attention with per-tensor descales runs on the new ASM kernel, through the existing varlen entry points and a new public wrapper. It was previously not served by a gfx1250 ASM kernel.
• LSE: the LSE comes from the same fp8-rounded softmax weights that normalize O, so O and LSE stay consistent for chunked-prefill merges. It can differ from the exact log-sum-exp (see "Validation and unverified scope").
• gfx1250 varlen ASM routing: one gate for both kernels replaces the per-site copies. It applies the CK window convention, and fixes the cases listed under "Changes" where the old checks sent calls to the bf16 ASM that it cannot compute, or that its C++ entry refuses.

Validation and unverified scope

• Op test: python3 op_tests/test_fmha_fwd_with_sink_varlen_asm.py --dtype bf16 fp8 on MI450 (gfx1250 B0), container rocm/fw-bringup-mlperf:gfx1250-atom-20261028-dsv41, at the reviewed revision.
  • bf16: 60 correctness rows (D64 / D128 / D192x128) with zero error. Strided q/k/v and strided out are bit-identical to dense. Unsupported strides are refused by the C++ entry and avoided by the public path. The gate table (17 cases) matches.
  • fp8: 80 rows with 0 bad rows and cosine similarity >= 0.999997 against a reference that models the kernel's fp8 rounding. The rows cover dense and strided layouts (kv_split, fused_qkv, padded and odd row strides), GQA, q < k and q > k, empty-KV sequences, causal and not, through all three entry points. Chunked-prefill merge: 6 cases pass. The gate table (18 cases) matches.
• Stress (outside this repo): 800 randomized varlen configurations, each rerun 12 times across launch configurations, entry points and layouts, some under concurrent load; every rerun was bit-identical to the first. The 96-head long-sequence shapes, including the target shapes, were bit-identical across launch configurations.
• Gate refactor (checked on CPU):
  • The new gate makes the same decision as the per-site gates it replaces over about 9.6M argument combinations.
  • Its host cost per call is lower than the original checks, on gfx1250 and on other architectures.
• Not verified:
  • MI300X / MI350 (gfx942 / gfx950) were not run. Their kernels and routing are untouched; the only new code on their path is the cached device check in front of the gfx1250 gate.
  • gfx1250-strict (A0) was not tested. The device check excludes it, as on the other gfx1250 ASM paths.
  • That the bf16 ASM computed wrong results for a left window of 0, or a non-causal right window of 0, follows from the masks it supports. It was not reproduced on GPU; those calls now go to CK.
  • The LSE difference from the exact log-sum-exp has no fixed bound; it reached about 0.08 in one randomized configuration, where the kernel still matched the reference that models its rounding.
  • One op-test run hit a GPU memory access fault at the start of the bf16 D192x128 strided test. The immediate rerun passed, and the fault has not been reproduced. On this node, the torch reference used by the tests has also faulted intermittently in runs that never called these kernels.

PR-specific question

Changed op tests

op_tests/test_fmha_fwd_with_sink_varlen_asm.py (the existing gfx1250 varlen ASM test; fp8 is selected with --dtype fp8, and bf16 stays the default):

• fp8 accuracy: the layouts, GQA, mask and empty-KV cases above, through the public wrapper, flash_attn_varlen_func and flash_attn_varlen_fp8_pertensor_func. A launch spy checks that each call reaches the ASM kernel once. The reference models the kernel's fp8 rounding, so the check can be tight.
• fp8 chunked prefill: one call over prefix + chunk, compared with two calls merged by their LSE, and with an exact reference for the prefix.
• Gate tables (bf16 and fp8): which calls each kernel must take or refuse. This covers dtype, head dim, sink and descale contracts, strides, windows, sink tokens, dropout, soft-cap and empty Q. These routing decisions cannot be observed from the accuracy tests, because a refused call is still answered correctly by another backend.
• fp8 perf: the fp8 ASM kernel on long-sequence shapes, with reference candidates for comparison.

Not redundant: nothing else tests the fp8 kernel. The gate tables pin decisions that only show up as which backend runs.

[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 7375 --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 e6dfeae1f21f · 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. Which architecture, dtype and shape boundaries select the changed kernel path, and what invariant prevents unsupported inputs from reaching it?
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
• FAIL — PR-specific question has no substantive answer
• FAIL — changed op test not explained in Changed op tests: op_tests/test_mha_optional_flydsl.py

_{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

FP8 routing can silently drop the requested diagonal-only mask, and the dispatch check can reject supported operands by using the build target instead of the runtime GPU architecture.

2 open issues

• Zero-width windows incorrectly routed to ASM without mask bounds · New
• Build target override misroutes runtime dispatch on gfx1250 · New

What changed in this PR

Adds a gfx1250 assembly-backed FP8 variable-length attention path for 192-dimensional queries/keys and 128-dimensional values, integrated into AITER’s existing attention APIs.

Changes:

• Adds a HIP launcher, kernel manifest, and JIT configuration for per-tensor FP8 descales and BF16 output.
• Routes supported calls through the new ASM path.
• Extends correctness and performance tests with strided inputs, chunked prefill, and unsupported-operand checks.

File	Description
op_tests/​test_fmha_fwd_with_sink_varlen_asm.py	Adds FP8 references, correctness cases, benchmarks, and dtype selection.
hsa/​gfx1250/​fmha_fwd_fp8_varlen/​fmha_fwd_fp8_varlen.csv	Registers causal and non-causal FP8 kernels.
csrc/​py_itfs_cu/​asm_fmha_fwd_fp8_varlen.cu	Implements operand validation, kernel arguments, and launch selection.
aiter/​ops/​mha.py	Adds the FP8 wrapper and public API dispatch gates.
aiter/​jit/​optCompilerConfig.json	Registers compilation and manifest code generation for the new module.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Incorrect softmax LSE computation can corrupt attention-state merging, while unchecked CPU descale pointers and zero-head inputs risk invalid memory access and undefined behavior.

4 open issues

• Compute softmax LSE from unrounded FP32 probabilities · New
• Reject descales on a different device from q · New
• Reject non-positive head counts before modulo · New
• Require negative infinity for non-finite reference LSE rows · New

2 resolved since last review

• Zero-width windows incorrectly routed to ASM without mask bounds
• Build target override misroutes runtime dispatch on gfx1250

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Unresolved device-validation gaps risk GPU memory faults, while FP8 probability underflow in the default route can severely distort attention outputs.

No open issues 🎉

2 previously missed issues

In code that hasn't changed since last review

Keep FP8 attention opt-in until probability underflow is bounded

aiter/​ops/​mha.py:3421

[inferred] Could this default route remain opt-in until FP8 probability underflow is bounded against exact attention? For non-causal Sq=64, Sk=16384, H=1, unit descales and the default scale, set Q[:,0,0]=8, K[0,0,0]=14, all other Q/K coordinates to zero, and V to +1 for the first key and -1 elsewhere: the documented rounding model yields O=+1 instead of exact O≈-0.6699 because every tail probability rounds to zero. Reproduce this on MI450, add an exact-reference regression, and preserve small probabilities before enabling automatic dispatch; the same loss also makes chunked merges partition-dependent.

Validate operand devices before launching the GPU kernel

csrc/​py_itfs_cu/​asm_fmha_fwd_fp8_varlen.cu:206

[inferred] Could this entry reject non-GPU and mismatched-device operands before constructing the device guard? CPU int32 cu_seqlens_q/k retain their device through the public wrapper and pass the current checks, so their host pointers reach the GPU kernel and may cause a memory fault rather than a synchronous argument error. Require Q to be on a GPU and K/V/out/LSE/cu_seqlens to share that device, as already required for the descales.

4 resolved since last review

• Reject descales on a different device from q
• Compute softmax LSE from unrounded FP32 probabilities
• Require negative infinity for non-finite reference LSE rows
• Reject non-positive head counts before modulo

_{Review effort: Balanced}

[shay-li77]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision e6dfeae1f21f · 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. Which architecture, dtype and shape boundaries select the changed kernel path, and what invariant prevents unsupported inputs from reaching it?
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
• FAIL — PR-specific question has no substantive answer
• FAIL — changed op test not explained in Changed op tests: op_tests/test_mha_optional_flydsl.py

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

• Architecture: the running GPU is gfx1250 B0 or later. _IS_GFX1250_ASM_DEVICE requires get_gfx_runtime() to be gfx1250 / gfx1250-strict and is_gfx1250_asm_supported() to be true, which excludes A0. On every other architecture the sites stop at this constant.
• fp8 kernel (new):
  • q/k/v are float8_e4m3fn, 3-D packed THD.
  • Head dims: q/k 192, v 128.
  • The q head count is a positive multiple of the kv head count.
  • q/k/v_descale are each a single fp32 element on q's device.
  • out is None or bf16; no sink logits.
• bf16 kernel (existing, gate merged):
  • q is bf16, with no descales.
  • D64 / D128 (hdim_v == hdim_q) or D192x128.
  • D64 requires sink logits; D128 / D192x128 refuse them.
• Both kernels:
  • max_seqlen_q > 0.
  • Plain or causal (bottom-right) attention only. There is no window: a negative window, or one at or beyond max_seqlen_k, counts as none, and causal ignores the right window.
  • No sink-token span, dropout, soft-cap, bias, ALiBi, paged KV (block_table), padded cu_seqlens or returned softmax.
• D192x128 strides: q/k/v/out are used in place and need a unit last-dim stride, a token stride in (0, 2^24) B and a head stride in (0, 2^31) B. Other views go to another backend.

The pre-router policy is unchanged: bf16 non-causal D64/D128, and D64/D128/D192x128 with AITER_ENABLE_EXPERIMENTAL, are offered to FlyDSL first.

What invariant blocks unsupported inputs?

• One gate, same arguments at every site. The pre-router only short-circuits to FlashAttnVarlenFunc, which calls _flash_attn_varlen_forward. That function asks the same gate again after maybe_contiguous, so a call the pre-router accepts is served by that ASM kernel. A call the gate rejects keeps its previous route (FlyDSL, CK or Triton).
• The C++ entry validates again. It checks dtypes, head dims, head counts (positive, divisible), descale dtype / size / device, out and lse shapes, and strides, and raises instead of launching. require_gfx1250_asm raises on A0 for direct calls to the public wrapper.
• Tests: the op test pins this with the gate tables (bf16 18 cases, fp8 20 cases), the ops-refused cases (CPU descale, zero kv heads) and the unsupported-stride cases.