[Triton] FP8 with FA v3 API

[micmelesse]

Motivation

Add support for fp8 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 fp8 by looking at the tests.

for examples of using fp8 with regular and paged attention see the tests at 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,
)

# enable backward for fp8 using dequantized values
set_fp8_dequantized_backward(True)

# 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,
        )

[dhonnappa-amd]

Jenkins CI skipped: Check lint failed. Exiting the entire job...

[micmelesse]

I will reopen in a bit.

[dhonnappa-amd]

Jenkins CI skipped: Required check(s) 'ruff_black' are missing. Exiting the entire job...

[dhonnappa-amd]

Jenkins CI skipped: Check lint failed. Exiting the entire job...

[micmelesse]

Moved here, #1065