F8 fmha ASM gfx950

[JohnNikolay84]

Motivation

There is no f8 fmha ASM kernel for gfx950, despite a similar kernel exists for gfx942.
Diffusion model spend 60-70% flops just in attention, so f8 can gives even 25% e2e improvement over bf16.

Technical Details

I used MLA ASM kernel for hd192 and reworked it as a pure fmha with hd128. In the end of the day it is not that hard as it is essentially fmha and mfma is 32x32x64 which is perfectly compatible with hd128 too.

Test Plan

• Run video generation prompts with xDiT.
• aiter fmha f8 tests.

Test Result

The improvement is 60->45s. No noticeable artifacts.

python3 op_tests/test_mha_fp8.py -b 1 -n 5 -q 75600 -d 128
    [aiter] import [module_aiter_core] under /workspace/aiter/aiter/jit/module_aiter_core.so
    [aiter] import [module_fmha_v3_fwd] under /workspace/aiter/aiter/jit/module_fmha_v3_fwd.so
    [aiter] type hints mismatch, override to --> fmha_v3_fwd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout_p: float, softmax_scale: float, is_causal: bool, window_size_left: int, window_size_right: int, return_softmax_lse: bool, return_dropout_randval: bool, how_v3_bf16_cvt: int, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, alibi_slopes: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, gen: Optional[torch.Generator] = None) -> List[torch.Tensor]
    [aiter] LoadKernel: _ZN5aiter24fmha_fwd_hd128_fp8_gfx950E hsaco: /workspace/aiter/hsa//gfx950/fmha_v3_fwd/fwd_hd128_fp8.co
    [W424 14:32:50.280727107 collection.cpp:1133] Warning: ROCTracer produced duplicate flow start: 4 (function operator())
    [aiter] LoadKernel: _ZN5aiter19fmha_fwd_hd128_bf16E hsaco: /workspace/aiter/hsa//gfx950/fmha_v3_fwd/fwd_hd128_bf16.co
Output nrms: 0.004180908203125
Output max diff: 0.001953125

[aiter] mha summary:
   quant_fwd_us  quant_fwd_tflops  quant_fwd_gb_per_sec        fwd_us   fwd_tflops  fwd_gb_per_sec
0   7082.736211       2065.772488             27.325033  12004.333737  1218.836624       32.244355

wan22_i2v_14b_fp8_aiter_mla-mha_fp8_gemms_hybrid_aiter_2steps_staging.mp4

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 2911 --add-label <label>

[jcaraban]

@JohnNikolay84 while this was developed in the context of Diffusion Inference on xDiT, I think the testing can be limited to AITER. Your change to aiter/ops/mha.py makes it so your ASM gets selected in the MHA fp8 path, right? In that case op_tests/test_mha_fp8.py is now expected to jump from 1200 TFlops to 2000 TFlops on mi355. If you can paste those logs, I'd consider that enough.

[JohnNikolay84]

@jcaraban

python3 op_tests/test_mha_fp8.py -b 1 -n 5 -q 75600 -d 128
[aiter] import [module_aiter_core] under /workspace/aiter/aiter/jit/module_aiter_core.so
[aiter] import [module_fmha_v3_fwd] under /workspace/aiter/aiter/jit/module_fmha_v3_fwd.so
[aiter] type hints mismatch, override to --> fmha_v3_fwd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout_p: float, softmax_scale: float, is_causal: bool, window_size_left: int, window_size_right: int, return_softmax_lse: bool, return_dropout_randval: bool, how_v3_bf16_cvt: int, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, alibi_slopes: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, gen: Optional[torch.Generator] = None) -> List[torch.Tensor]
[aiter] LoadKernel: _ZN5aiter24fmha_fwd_hd128_fp8_gfx950E hsaco: /workspace/aiter/hsa//gfx950/fmha_v3_fwd/fwd_hd128_fp8.co
[W424 14:32:50.280727107 collection.cpp:1133] Warning: ROCTracer produced duplicate flow start: 4 (function operator())
[aiter] LoadKernel: _ZN5aiter19fmha_fwd_hd128_bf16E hsaco: /workspace/aiter/hsa//gfx950/fmha_v3_fwd/fwd_hd128_bf16.co
Output nrms: 0.004180908203125
Output max diff: 0.001953125

[aiter] mha summary:
   quant_fwd_us  quant_fwd_tflops  quant_fwd_gb_per_sec        fwd_us   fwd_tflops  fwd_gb_per_sec
0   7082.736211       2065.772488             27.325033  12004.333737  1218.836624       32.244355

[jcaraban]

Thanks! Here is a plot of the long-context TFlops (per head) for @valarLip et al

[yichiche]

@JohnNikolay84 Thanks for the effort. To clarify: is this kernel limited to MHA with hd128, or might it be extended to other head dimensions in the future?

[jcaraban]

@yichiche what other head dimensions do you have in mind? hd192 might be trivial since Sergey derived his MHA from the persistent fp8 MLA. However our focus is diffusion models where hd128 is typically enough

[yichiche]

@jcaraban While running Wan 2.2 end-to-end, I hit a segmentation fault. In this model, both self-attention (where Q/K/V share the same sequence length) and cross-attention go through the MHA path. The crash occurs in cross-attention with seq_q = 90000, seq_k = seq_v = 512.

Minimal repro (CLI)

cd /sgl-workspace/aiter_mha_fp8/aiter
python3 op_tests/test_mha_fp8.py -b 1 -n 40 -q 90000 -k 512 -d 128

[aiter] import [module_aiter_core] under /sgl-workspace/aiter_mha_fp8/aiter/aiter/jit/module_aiter_core.so
[aiter] import [module_fmha_v3_fwd] under /sgl-workspace/aiter_mha_fp8/aiter/aiter/jit/module_fmha_v3_fwd.so
[aiter] type hints mismatch, override to --> fmha_v3_fwd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout_p: float, softmax_scale: float, is_causal: bool, window_size_left: int, window_size_right: int, return_softmax_lse: bool, return_dropout_randval: bool, how_v3_bf16_cvt: int, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, alibi_slopes: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, gen: Optional[torch.Generator] = None) -> List[torch.Tensor]
[aiter] hipModuleLoad: /sgl-workspace/aiter_mha_fp8/aiter/hsa//gfx950/fmha_v3_fwd/fwd_hd128_fp8.co GetFunction: _ZN5aiter24fmha_fwd_hd128_fp8_gfx950E
[aiter] hipModuleGetFunction: _ZN5aiter24fmha_fwd_hd128_fp8_gfx950E Success
Memory access fault by GPU node-2 (Agent handle: 0x5592cecab380) on address 0x7f1c01000000. Reason: Unknown.

[jcaraban]

Thanks for the minimal repro 👍 The crash seems constrained to cross attention (K != Q) and Heads 2, 4, 8, 10, 20, 40... however it doesn't crash for H 1, 3, 5, 6, 7, 9...

[JohnNikolay84]

@yichiche thank you for the report. Great catch. The issue should be fixed now.

[yichiche]

@yichiche thank you for the report. Great catch. The issue should be fixed now.

Hi, may I follow up the status of this PR? When will merge it?

[JohnNikolay84]

@valarLip do you have any ,comments/objections about this PR?

[valarLip]

looks really good

[JohnNikolay84]

I've asked Claude to compare CK vs asm for the causal kernel too

[carlushuang]

This PR regession all the MI355 BF16 attention causal masking attention, will increase execution to 2x because of

[carlushuang]

#3390 => This PR fixed this regression, we don't globally set tg_div = 1 (hence will have all the bf16 mi355 attn regression), only check the FP8 attention introduced in this PR to set tgdiv = 1

cc @JohnNikolay84 @yichiche @jcaraban
cc @valarLip @danyao12 @shengnxu