i8fp8 fmha gfx950/942 asm

[jcaraban]

Motivation

fp8 FMHA delivers large speedups over bf16, but its precision is insufficient for diffusion models — the quantization error shows up as visible artifacts in generated images and videos.

This PR adds a Sage-like i8fp8 FMHA (int8 Q/K, fp8 V, bf16 output) as a middle point between bf16 and fp8. It is more precise than fp8 (enough to prevent artifacts) while running faster than bf16, about 95% of fp8 performance.

Technical Details

New asm FMHA forward kernel for head-dim 128 on gfx950 (MI350) and gfx942 (MI300):

• Quantization scheme: int8 Q and K (QK product accumulated in int32), fp8 (e4m3) V, bf16 output. Per-tensor or per-(batch, kv-head) q/k/v_descale (float32) applied on the descale path.
• Dispatch: routed implicitly by tensor dtype (q,k = int8 and v = fp8) — no new API surface or override flag is required. The new data_type string "i8fp8bf16" threads through the existing v3 dispatch:
  • csrc/py_itfs_cu/asm_mha_fwd.cu — detects is_i8fp8, validates q/k(int8) and v(fp8) dtypes and the descale tensors, and forces a bf16 output.
  • csrc/cpp_itfs/mha_fwd.cu — sets in_bpe = 1, wires the descale pointers, and admits i8fp8bf16 through the v3 support guard.
  • hsa/gfx950/.../fmha_fwd.csv and hsa/gfx942/.../fmha_fwd.csv — one lookup row each pointing at the new fwd_hd128_i8fp8.co kernel.
• Python entry point: flash_attn_i8fp8_pertensor_func(q, k, v, q_descale, k_descale, v_descale, ...) in aiter/ops/mha.py, plus an is_fmha_v3_i8fp8() gate in _flash_attn_forward.
• Constraints: head-dim 128, non-causal full attention, bf16 output only, descales required.

Note: causal and grouped variants could be generated if useful for othe rusers

Test Plan

Benchmarked and accuracy-checked via op_tests/op_benchmarks/triton/bench_sage.py, which was extended with an aiter_i8fp8 runner, a transformer-like input distribution, QK rotation/clipping for better quantized quality, and a refactored accuracy-metrics report (cosine similarity against the bf16 and Sage references):

  python op_tests/op_benchmarks/triton/bench_sage.py --b 1 --hq 16 --sq 16384 --d 128 --kernel all

Test Result

• Functional: i8fp8 kernel runs on gfx950 (MI350) and gfx942 (MI300) for hd128.
• Accuracy: cosine similarity 0.998198 vs BF16 — above FP8 although below Triton Sage, but sufficient
• Performance: ~1860 TFLOPS (B1, H16, S16384, D128) for full attention on gfx950

bench_sage --kernel=all  (b=1 hq=5 sq=75600 sk=75600 d=128 input=transformer):
kernel             time(ms)     TFLOPS          MAE         MaxE       Cosine
------------------------------------------------------------------------------
sage_fp8             9.6745    1512.37    5.198e-04    4.297e-02     0.998552
sage_mxfp4           7.9293    1845.21    2.386e-03    2.988e-01     0.969782
aiter_fp8            7.1052    2059.24    7.851e-04    8.008e-02     0.996940
aiter_i8fp8          7.8580    1861.96    6.900e-04    4.492e-02     0.997761
aiter_bf16          12.2508    1194.32    0.000e+00    0.000e+00     1.000000

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: 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 3186 --add-label <label>

[valarLip]

why we need new python api..

[jcaraban]

I tried to minimize changes to existing logic on this PR, and leave those for a holistic refactor later, but yes in this case flash_attn_fp8_pertensor_func could have been reused with minimal intrusion

[jcaraban]

but even then, I find the duplication less severe than making the fp8 wrapper dispatch Sage-like i8f8, when the later doesn't fall back to flash_attn_func_v3_triton().

[JohnNikolay84]

That is a lot of code duplication. Why not having a single function that takes tensor types to check against as an argument ?

[jcaraban]

The whole AITER mha API is long due for a refactor, precisely so different Q/K/V/O dtypes combinations can go through with minimal code duplication. However in this PR I tried to keep my changes to a minimum, thus introducing temporal duplication. My proposal is to refactor after the mxfp4 and mxfp6 attention variants have also merged, to minimize the chances of breaking something (like e.g here )