[Triton] FA v3 API

[micmelesse]

Motivation

Add support for fp8 and paged attention in Flash Attention

Technical Details

Modify existing code so that it confirms to the flash attention v3 api. A user provides fp8 values for q, k and v and their descale values.

Test Plan

update mha tests and bench code

Test Result

Submission Checklist

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

[micmelesse]

You can see examples on how to use the interface for using fp8 with regular and paged attention by looking at the tests like op_tests/triton_tests/test_mha.py

For regular attention fp8, you will see code that looks like this

from aiter.ops.triton.mha_v3 import (
    flash_attn_func as flash_attn_func_v3,
)

# forward
triton_out = flash_attn_func_v3(
    q_fp8,
    k_fp8,
    v_fp8,
    softmax_scale=None,
    causal=CAUSAL,
    q_descale=q_descale,
    k_descale=k_descale,
    v_descale=v_descale,
)

# backward
triton_dq, triton_dk, triton_dv = torch.autograd.grad(
    triton_out, (q_fp8, k_fp8, v_fp8), do.clone()
)

Here is a small model trained on wikitext to test convergence.

for paged attention fp8 which is available in the inference api flash_attn_with_kvcache , you will see code that looks like this

from aiter.ops.triton.mha_v3 import (
    flash_attn_with_kvcache as flash_attn_with_kvcache_v3,
)

# forward
out_kernel = flash_attn_with_kvcache_v3(
            q_fp8,
            k_cache_fp8,
            v_cache_fp8,
            cache_seqlens=cache_seqlens,
            causal=causal,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            page_table=page_table,
        )

[Copilot]

Pull Request Overview

This PR adds support for Flash Attention v3 API with FP8 quantization and paged attention capabilities. The changes modify existing code to conform to the Flash Attention v3 interface while maintaining backward compatibility with v2.

• Updates MHA tests to use v3 API with FP8 support and new quantization utilities
• Introduces new v3 interfaces with proper FP8 descaling and KV cache functionality
• Adds comprehensive FP8 quantization utilities for different tensor layouts

Reviewed Changes

Copilot reviewed 13 out of 15 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
op_tests/triton_tests/test_mha.py	Updates test suite to support both v2 and v3 APIs with FP8 quantization
op_tests/op_benchmarks/triton/bench_mha.py	Updates benchmark code to use v3 FP8 quantization API
aiter/ops/triton/utils/mha_kernel_utils.py	Adds FP8 quantization and dequantization utilities
aiter/ops/triton/mha_v3.py	Implements Flash Attention v3 interface with FP8 support
aiter/ops/triton/mha.py	Removes deprecated FP8 functions and adds v2 KV cache support
Multiple triton kernel files	Adds comprehensive v2/v3 interfaces and FP8 prefill implementation

Comments suppressed due to low confidence (1)

aiter/ops/triton/mha_v3.py:1

• The union syntax int | None should be replaced with Optional[int] for consistent type annotation style and Python version compatibility.

# SPDX-License-Identifier: MIT

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}