[Ragua1]

Summary

On archs where Triton emulates fp8 in software, the fp8 MHA prefill now converts k and v to fp16 once before the kernel and q once per program, then runs the fp16 dot products. On gfx1101 this is 3.8x to 24.0x faster than today on all 238 measured shapes, median 15.8x, and the error against an fp32 reference is lower on all 238.

Archs with fp8 hardware are not affected: they compile the same attn_fwd as on main.

Motivation

On archs without fp8 WMMA or fp8 conversion instructions, Triton converts fp8 operands to fp16 in software. #6782 lets the Triton MHA fp8 path run on these archs: gfx1100, gfx1101, gfx1102, gfx1150 and gfx1151. In the prefill kernel this conversion runs on every K and V tile of every program, after the tile is read from LDS. The fp8 kernel is therefore many times slower than the fp16 or bf16 kernel on the same data. #6770 shows the same effect in fav3_sage on gfx1100.

Today fp16 or bf16 attention is many times faster on these archs, so the fp8 path matters where the caller cannot switch to it: a model available only in fp8, or a workload whose conversion to fp16 or bf16 is not practical on the target card. With this change the fp8 call runs close to the bf16 kernel on the same data, at a median of 0.95 to 1.22 of its time per head configuration.

Every fp8 value is exact in fp16, so the conversion can skip the descales, and the kernel applies them in fp32.

Changes

• attention_forward_prefill_triton_impl (fwd_prefill.py) takes the new path when q, k and v are fp8, the layout is bshd, and get_arch().supports_sw_emulated_fp8.
• On that path k and v are converted to fp16 by _fp8_to_fp16_bshd, a new Triton kernel: strided bshd input, contiguous output, int64 offsets, because s * nheads * head_dim can exceed 2^31 elements. The conversion is exact for e4m3 and e5m2, so no value is rounded and the fp16 range is never exceeded. The descale shapes are checked before the conversion, as the fp8 branch checks them. A missing descale becomes ones.
• q stays fp8. attn_fwd gets a new constexpr FP8_AS_FP16. When it is set, attn_fwd converts q to fp16 once, right after loading the q block and outside the K/V loop. It loads the three descales with the same indexing as the fp8 branch, passes q_descale * k_descale, taken in fp32, to _attn_fwd_inner, and multiplies the output by v_descale once before the store.
• _attn_fwd_inner gets the same constexpr. When it is set, q_descale * k_descale joins the existing SM_SCALE multiply of the scores, a scalar product per tile. With the constexpr unset, which is every other caller, the generated code is unchanged: a bf16 call through mha_v3.flash_attn_func compiles to the same instructions on main and with this change, and gives the same output.
• The kernel then runs with IS_FP8 off. The output dtype check of the non-fp8 branch accepts this case. The output keeps the dtype the caller allocated, fp32 for flash_attn_3.fwd.
• Unchanged: archs with fp8 hardware, the bhsd and varlen (thd) layouts, the backward pass, decode, and rotary embeddings, which still act on the fp8 q.
• Memory: the converted k and v live for the duration of the call, 2 bytes per element. Converting q in the kernel instead of up front keeps its copy out of memory.
• _fp8_to_fp16_bshd gets a repr from make_kernel_repr.
• op_tests/triton_tests/attention/test_mha_fp8.py: new test_mha_fp8_strided_qkv. q, k and v are views into one packed fp8 qkv tensor, and the output must equal the output for contiguous copies bit for bit.
• op_tests/triton_tests/attention/test_mha_fp8.py: new test_mha_fp8_sw_emulated_path. It forces supports_sw_emulated_fp8, so the new path also runs on archs with fp8 hardware, and compares the output with attention_ref on the dequantized inputs.
• docs/triton_mha_fp8_sw_emulated.md: new note on fp8 MHA inputs on archs with software-emulated fp8, see Documentation.

Alternative considered: bf16 with the descales applied in the conversion

A smaller change converts k and v to bf16 with their descales applied, and q to bf16 with its descale in attn_fwd. It changes only attn_fwd, not _attn_fwd_inner, and needs bf16 because descaled values can exceed the fp16 range. Measured in the same processes as this PR:

	fp8 today	bf16 alternative	this PR (fp16)
speedup over fp8 today, median (min–max)	1	14.9x (4.0–21.8)	15.8x (3.8–24.0)
faster than the bf16 alternative by the criterion under Performance, shapes	–	–	206 of 238
slower than the bf16 alternative by the same criterion, shapes	–	–	19 of 238, all with seqlen ≤ 256, by at most 9 %
no claim either way, shapes	–	–	13 of 238
1 - cos vs fp8 today, median ratio	1	0.831	0.827
lines changed in fwd_prefill.py	–	+121 / −4	+132 / −5

From seqlen 512 up this PR is faster than the bf16 alternative on 154 of 154 shapes. The 19 shapes where the alternative is faster are small calls. 13 of them take under 0.06 ms, and there the alternative saves at most 2.7 µs per call. On the other 6 it saves at most 1.3 % of the call. Summed over all 238 shapes, this PR takes 2948 ms and the alternative 3149 ms. 1 - cos of this PR is lower than that of the bf16 alternative on 238 of 238 shapes, median ratio 0.995.

The fp16 kernel is faster than the bf16 kernel on this arch: plain mha.flash_attn_func in fp16 is faster than in bf16 on 238 of 238 shapes and runs at 0.92 of its bf16 time in the median (0.76–0.96). RDNA3 has f32 to f16 conversion instructions (V_CVT_F16_F32, V_CVT_PK_RTZ_F16_F32) but none to bf16, so the compiler rounds to bf16 with integer instructions. P is converted to the dtype of V on every tile. On 32, 256, 32/8, 128, causal the bf16 alternative's attn_fwd has 4587 instructions and this PR's 4042, and most of the difference is that rounding.

Performance (if applicable)

Before = fp8 today, after = this PR. Both ran in the same process with rotated arm order, mha_v3.flash_attn_func with fp8 q, k, v and descales, default autotune config. Time is the median over 10 fresh processes, GPU time only with the GPU queue filled ahead of each timed window, hot cache.

The timing ran on a measurement tree that adds a runtime switch for the new path to this change. Its attn_fwd compiles to the same instructions as this PR in all 143 compiled variants. Its conversion kernel compiles to 1066 and 1973 instructions in its two variants, against 1066 and 1970 in this PR, and 250 of 250 compared outputs are equal.

Hardware	Workload/configuration	Before	After	Change
gfx1101, Windows 11	B=1, S=1024, Hq/Hk=32/8, d=128, causal	5.954 ms	0.314 ms	18.94x
	B=1, S=4096, Hq/Hk=32/8, d=128, causal	72.935 ms	3.494 ms	20.87x
	B=1, S=16384, Hq/Hk=32/8, d=128, causal	1086.938 ms	50.725 ms	21.43x
	B=8, S=1024, Hq/Hk=32/8, d=128, no mask	69.402 ms	3.426 ms	20.26x
	B=1, S=4096, H=64, d=112, causal	139.896 ms	8.114 ms	17.24x
	B=1, S=4096, H=16, d=64, no mask	20.974 ms	2.032 ms	10.32x
	B=8, Sq=1024, Sk=4096, Hq/Hk=32/8, d=128, no mask	272.411 ms	12.674 ms	21.49x

The full set is 238 shapes: 7 head configurations, seqlen 64 to 16384, causal and no mask, and cross-attention. The criterion was set before the run: drop the lowest and highest of the 10 per-process ratios, then the minimum is above 1 and the median minus 1 is larger than half the range. It holds on 238 of 238 shapes.

The converted k and v cost memory. With 8 GiB of the 16 GiB card already allocated, the largest seqlen that fits is 0.64 to 0.95 of today's, depending on the head configuration.

Measurement conditions, per-configuration results, memory limits

Shapes, bshd: 7 head configurations (q/kv heads and head dim 32/8, 128, 64/8, 128, 128/8, 128, 48/8, 128, 16/16, 64, 16/16, 128, 64/64, 112). Per configuration seqlen 64 to 16384 at batch 1 and at batch·seqlen = 8192 up to seqlen 4096, causal and not, and cross-attention 1024 x 4096 at batch 1 and 8. k and v are quantized once as flash_attn_fp8_func does.

Arms in one process: fp8 today with the new path switched off at runtime, this PR, the bf16 alternative, and mha.flash_attn_func on the bf16 inputs the fp8 tensors came from and on the same inputs in fp16. Median of 10 windows, 3 for fp8 today at seqlen 8192 and above. 5 rounds per process with rotated arm order, 10 fresh processes. No row ended its GPU spin before the timed window. The output of every arm was checked against fp32 SDPA in every round.

Stack: TheRock ROCm 10.2 (torch 2.15.0a0+rocm10.2.0a20260926), Triton 3.8.0.post29, AITER_TRITON_ONLY=1, aiter main (457b1151) with this change, display driver 32.0.32015.7029.

Speedup is fp8 today / this PR, median over the 8 processes left per shape. Peak memory is the memory allocated by one call above the inputs. Ranges are over the medians of the shapes.

q/kv heads, head dim	shapes	speedup, median (min–max)	time vs bf16 kernel, median (min–max)	peak memory vs fp8 today, median (min–max)
32/8, 128	34	19.1x (6.5–23.4)	1.07 (0.94–1.87)	1.25 (1.25–1.99)
64/8, 128	34	18.9x (5.9–23.1)	0.98 (0.93–1.62)	1.12 (1.12–1.50)
128/8, 128	34	19.5x (7.5–23.5)	0.95 (0.90–1.51)	1.06 (1.06–1.25)
48/8, 128	34	19.3x (6.0–24.0)	0.99 (0.94–1.80)	1.17 (1.16–1.66)
16/16, 64	34	9.2x (4.3–10.7)	1.07 (0.95–1.32)	1.98 (1.98–4.94)
16/16, 128	34	15.8x (3.8–22.8)	1.21 (1.00–2.12)	1.99 (1.99–4.97)
64/64, 112	34	15.7x (4.7–21.2)	1.22 (0.97–2.04)	1.99 (1.82–4.96)

By seqlen, self-attention shapes of all 7 head configurations:

seqlen	shapes	speedup, median (min–max)	faster / slower than the bf16 alternative, by the criterion
64	28	7.0x (3.8–11.6)	20 / 7
128	28	7.5x (5.1–10.5)	12 / 11
256	28	11.6x (6.6–15.9)	20 / 1
512	28	15.4x (7.5–20.3)	28 / 0
1024	28	18.6x (9.1–21.9)	28 / 0
2048	28	20.1x (9.6–23.2)	28 / 0
4096	28	20.7x (9.9–23.4)	28 / 0
8192	14	21.3x (10.5–23.4)	14 / 0
16384	14	21.5x (10.3–23.2)	14 / 0

The 14 cross-attention shapes 1024 x 4096 give 8.2x to 24.0x.

Accuracy: the dot products see the exact fp8 values, the descales are applied in fp32 to the scores and the output, and P is cast to fp16 instead of fp8. Against an fp32 reference, 1 - cosine similarity is lower than today on 238 of 238 shapes, median ratio 0.83. The maximum absolute error is higher than today on 86 shapes, median ratio 0.99.

Memory limit: the largest seqlen at batch 1 that fits under a torch allocator cap of 90 % of the card, with 8 GiB already allocated. The attention kernel launch was stubbed; allocations and the conversion ran. The bf16 alternative has the same limits as this PR.

q/kv heads, head dim	fp8 today	this PR	this PR / today
32/8, 128	302080	256000	0.85
64/8, 128	157696	144384	0.92
128/8, 128	80896	76800	0.95
48/8, 128	206848	184320	0.89
16/16, 64	947200	604160	0.64
16/16, 128	475136	303104	0.64
64/64, 112	135168	86016	0.64

The largest batch at seqlen 4096 gives the same ratios, 0.64 to 0.95. On an empty card the GQA configurations reach the search's stop at 2^31 elements per tensor before running out of memory, with and without this PR. The extra memory of this PR is 2 bytes × 2 · batch · seqlen_k · kv_heads · head_dim, for the converted k and v. The measured limits match that model within one search step on 252 of 252 searches.

Testing

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

Not tested on MI250X or MI300X. The new path is taken only on the archs in SW_EMULATED_FP8_ARCHS. On other archs the host code takes the existing branches, and attn_fwd compiles to the same TTIR as on main for bf16, fp16 and the IS_FP8 path, checked by compiling on gfx1101. test_mha_fp8_sw_emulated_path forces the new path on these archs as well. Test results are under "Validation and unverified scope".

Documentation

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

New docs/triton_mha_fp8_sw_emulated.md: why fp16 or bf16 inputs are the better choice on archs with software-emulated fp8, what memory an fp8 call takes, and that the forward output for fp8 inputs is fp32. The public API and the output dtype do not change.

Dependencies

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

Breaking Changes

None.

Possible follow-up: output dtype for fp8 inputs

Not part of this PR. For fp8 q, k and v, flash_attn_3.fwd allocates the output in fp32 (interface_v3.py:226). On gfx1101 with this change, an fp16 output is faster than fp32 on 6 of 6 shapes and fits a 1.22x to 1.59x longer seqlen. #7232 adds an out_dtype argument to _flash_attn_forward in mha.py; interface_v3.py is a separate entry point.

Measurements and the case for fp16 or bf16

The comment above interface_v3.py:226 says that bfloat16 caused mismatches on a small set of tests. The public mha_v3.flash_attn_func takes no out, so a caller cannot choose the output dtype. The choice does not depend on the arch, so it applies to archs with fp8 hardware as well. For comparison, FlashAttention-3 returns bf16 for fp8 inputs (hopper/flash_api.cpp).

Measured on gfx1101 with this change, with the output preallocated through flash_attn_3.fwd(out=...):

output	largest seqlen vs fp32 output, 8 GiB allocated	largest batch vs fp32 output, 8 GiB allocated	time vs fp32 output, 6 shapes	1 - cos vs fp32 output
fp16	1.22–1.59	1.19–1.61	0.79–0.98, faster on 6 of 6	1.000
bf16	1.22–1.59	1.19–1.61	0.90–1.03, faster on 5 of 6	1.002–1.003

The bf16 output is slower than the fp16 output on 6 of 6 shapes, by 0.6 % to 13.7 %. Storing a bf16 value takes 5 instructions on RDNA3, an fp16 value one v_fma_mixlo_f16, which also absorbs the last multiply. With interface_v3.py:226 changed to fp16 on main (457b1151), the forward tests of test_mha_fp8.py (-k "not backward") give 44 passed and 4 skipped, the same as without the change. With bf16 the result is the same.

The risk of fp16 is its range. The output is a weighted average of V after the descale, so it overflows only where the unquantized V exceeds 65504. How often that happens in real models was not measured. bf16 has the fp32 range and the same memory saving, at a smaller speedup and a slightly higher error.

PR Self-Review

Reviewed revision

a002007

Behavior and intent

On gfx1100, gfx1101, gfx1102, gfx1150 and gfx1151, an fp8 bshd prefill through fwd_prefill.py converts k and v to fp16 in a separate kernel and q in attn_fwd, and runs the fp16 dot products with the descales applied in fp32. Today the same call converts every K and V tile in software inside the loop, which makes it 3.8x to 24.0x slower on gfx1101. Archs with fp8 hardware, other layouts, backward and decode keep the current code; with the new constexpr unset, _attn_fwd_inner compiles to the same instructions as on main.

The change adds no tuning values. The arch condition uses the existing get_arch().supports_sw_emulated_fp8 from #6782.

Validation and unverified scope

gfx1101 (60 CU), Windows 11, TheRock ROCm 10.2 (torch 2.15.0a0+rocm10.2.0a20260926), Triton 3.8.0.post29, AITER_TRITON_ONLY=1, aiter main @ 457b1151. Tests ran as pytest op_tests/triton_tests/attention/test_mha_fp8.py -q from the aiter root, through a local wrapper that binds aiter.dtypes under AITER_TRITON_ONLY and caps the PyTorch allocator at 90 % of memory. The wrapper does not change kernels, configs or tests.

tree	tests	passed	failed	skipped	failures that pass in a fresh process
main (457b1151)	176	104	20	52	15 of 15
this PR (8805a4f8)	188	123	13	52	8 of 8

• Every failure is torch.OutOfMemoryError in test_mha_backward or test_mha_backward_varlen, which this PR does not change. Which backward tests fail depends on the memory left by earlier tests in the same process.
• The 5 cases 64-2048-2048-4 fail on both trees, also in a fresh process, and are not counted in the last column.
• The forward tests pass on both trees. The new test_mha_fp8_strided_qkv passes 8 of 8 with this change and 8 of 8 on main. The timing run checked the output of every arm against fp32 SDPA on 238 shapes, 0 failures.
• The new test_mha_fp8_sw_emulated_path passes 4 of 4 with this change. On the previous head c2067660, before the review fixes, its 2 fp16_large cases fail with NaN in the output.
• Unchanged paths: the TTIR of attn_fwd from this PR equals main for bf16, fp16 and the IS_FP8 path that archs with fp8 hardware take, compiled on gfx1101 with location info removed. The IS_FP8 arm used fp8 e4m3fn inputs with supports_sw_emulated_fp8 forced to False.
• gfx1150, Windows 11: test_mha_fp8.py on this change before the repr, other=0.0, the new tests and the review fixes gives 124 passed, 52 skipped, 0 failed. Timed with Triton 3.8.0.post28 and the gfx1151 attn_fwd config entry instead of the default one, this change is faster than today on 97 of 102 shapes by the same criterion and slower on none, median speedup 1.46x to 60.41x.
• gfx1101, Linux under WSL2, same torch build and wrapper, Triton 3.8.0, HSA_OVERRIDE_GFX_VERSION=11.0.1 because the WSL driver reports the card as gfx1102: this PR (8805a4f8) gives 79 passed, 57 failed, 52 skipped. All forward tests pass, including the two new tests. All 57 failures are torch.OutOfMemoryError in the backward tests; 52 pass in a fresh process, the other 5 are the 64-2048-2048-4 cases above.
• Unverified: bare-metal Linux, and gfx1100, gfx1102 and gfx1151, which take the same path by construction. CI runners have fp8 hardware, so there only test_mha_fp8_sw_emulated_path runs the new path, and it was run only on gfx1101, where the forced flag is already true. The varlen and bhsd layouts keep the in-kernel conversion and are not covered.

PR-specific question

test_mha_fp8_strided_qkv passes on main as well, 8 of 8 on gfx1101. It does not catch a defect that exists before this PR. It guards the new conversion kernel: on the archs in SW_EMULATED_FP8_ARCHS it is the only test that feeds non-contiguous fp8 k and v to _fp8_to_fp16_bshd, whose strides differ from the contiguous copy it writes. The existing fp8 tests in test_mha_fp8.py and test_mha_v3.py build q, k and v with torch.randn, so their bshd inputs are contiguous. A stride or mask error in the conversion would change the output for packed inputs only, and the bit-exact comparison catches it.

test_mha_fp8_sw_emulated_path is the test that runs the new path on archs with fp8 hardware, which the CI runners have. Its fp16_large cases pass fp16 descales of 256 for q and k, whose product exceeds the fp16 range. They fail on the previous head c2067660, where the product was taken in the descale dtype, and pass with the fp32 product.

The main boundary of the change is the condition in attention_forward_prefill_triton_impl: fp8 q, k and v, bshd layout and supports_sw_emulated_fp8. Every other call keeps the current kernel and the same generated code. The converted k and v add 2 bytes per element for the call, which lowers the largest seqlen to 0.64–0.95 of today's with 8 GiB allocated. The conversion kernel computes offsets in int64, because one tensor can exceed 2^31 elements.

Changed op tests

op_tests/triton_tests/attention/test_mha_fp8.py: adds test_mha_fp8_strided_qkv, 8 cases over seqlen 64 and 300, head dim 64 and 128, causal and not. It covers fp8 q, k and v that are views into one packed qkv tensor, with descales, through the public mha_v3.flash_attn_func. It is not redundant: no existing fp8 test passes non-contiguous bshd inputs, and seqlen 300 exercises the partial last block of the conversion. The test has no arch condition of its own, so it runs on every arch in ANY_FP8_ARCHS. On archs with hardware fp8 the conversion does not run, and there the test covers the existing attn_fwd with strided fp8 inputs. It was run only on gfx1101.

op_tests/triton_tests/attention/test_mha_fp8.py: adds test_mha_fp8_sw_emulated_path, 4 cases over causal and not and two descale sets: fp32 descales between 0.5 and 1.5, and fp16 descales of 256 for q and k with q and k scaled down by 2^-6. It sets GpuArch.supports_sw_emulated_fp8 to true with monkeypatch, so the new path runs on every arch in ANY_FP8_ARCHS while the arch name and its configs stay real. It checks that _fp8_to_fp16 ran twice, for k and v, and compares the output with attention_ref on the dequantized fp32 inputs. It is not redundant: no other test runs the new path on archs with fp8 hardware. It was run only on gfx1101.

Related

#6782, #6770, #7232

[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 7277 --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]

WAITING ON AUTHOR

PR Self-Review

Revision a00200757452 · 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 behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
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.}

[Ragua1]

@brunomazzottiamd, this PR follows up on the emulated fp8 performance you raised in your review of #6782. In my reply there (#6782 (comment)) I proposed converting to fp16 once and running the fp16 kernel on archs without fp8 hardware. This PR implements that for the bshd prefill. On gfx1101 the fp8 prefill is 3.8x to 24.0x faster than on main across 238 shapes, median 15.8x. Archs with fp8 hardware compile the same attn_fwd as before. The fav3_sage path from #6770 is not changed.

[Copilot]

🟡 Changes recommended

Descale multiplication and int32 stride arithmetic can overflow, producing NaNs or invalid memory reads, while the software FP8 emulation path lacks targeted test coverage.

4 open issues

• Combined descales can overflow and produce NaN scores · New
• Use int64 for the stride multiplication · New
• Add a test for software FP8 emulation on hardware-FP8 GPUs · New
• Allocate ones only when a descale is missing · New

What changed in this PR

This PR speeds up Triton MHA fp8 prefill on GPUs that emulate fp8 in software. It uses fp16 dot products for the bshd path while preserving existing dispatch for hardware-fp8 GPUs and other layouts.

Changes:

• Convert K and V to fp16 once per call and Q once per program, with descales applied during attention.
• Add a bit-exact comparison of packed, strided fp8 inputs against contiguous copies.
• Document conversion costs, memory trade-offs, and output dtype.

File	Description
op_tests/​triton_tests/​attention/​test_mha_fp8.py	Adds strided fp8 Q/K/V coverage.
docs/​triton_mha_fp8_sw_emulated.md	Explains software-emulated fp8 behavior and trade-offs.
aiter/​ops/​triton/​_triton_kernels/​flash_attn_triton_amd/​fwd_prefill.py	Adds the conversion kernel and fp16 prefill path.

_{Review effort: Balanced}

[Ragua1]

Merged main at 075695e8 into the branch. The only conflict was the import block of fwd_prefill.py (#5505 added get_num_xcds), and the change itself is unchanged. After the merge, test_mha_fp8.py on gfx1101 passes on Triton 3.8 and 3.9, except the five 64-2048-2048-4 backward cases that run out of memory on main as well. The other backward failures were out-of-memory errors and passed when re-run in a fresh process.

The change also holds on Triton 3.9. Measured on gfx1101 with the PR head, same torch and ROCm build, only the Triton version differs; main is the same tree with the new path switched off. Median of 5 fresh processes per version, fp32 output. The last two columns are the time on 3.8 divided by the time on 3.9, above 1 is faster on 3.9:

(b, s, hq, hk, d, causal)	main / PR, Triton 3.8	main / PR, Triton 3.9	main, 3.8 / 3.9	PR, 3.8 / 3.9
2, 1024, 16, 16, 128, no	19.7x	19.3x	1.01	0.99
2, 1024, 16, 16, 128, yes	16.9x	16.2x	1.07	1.01
2, 1024, 16, 16, 64, no	10.0x	11.8x	0.92	1.04
1, 4096, 24, 24, 128, no	23.0x	22.2x	1.02	0.98
1, 4096, 32, 8, 128, yes	21.2x	18.7x	1.15	1.01

[micmelesse]

@Ragua1 Can you rebase your pr?

[Ragua1]

@Ragua1 Can you rebase your pr?

done