FlashAttention3 Forward Compatibility

This document is a quick reference for the current FlashAttention-3-compatible forward coverage provided by MATE on MUSA.

Compatibility Overview

If you already use FlashAttention-3 style Python APIs, start here to see what works on MATE and how to try it.

1. Install flash_attn_3.

2. Import flash_attn_interface.

3. Call the usual forward APIs.

4. Check the matrix below for supported modes, dtypes, and limits.

This page covers how to get started with forward compatibility on MUSA. It does not cover kernel internals, backward, or autograd behavior.

Getting Started

Install the compatibility wrapper from the MUSA wheel source:

python -m pip install flash_attn_3 \
  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple

Minimal forward example:

import torch
from flash_attn_interface import flash_attn_varlen_func

device = "musa"
dtype = torch.bfloat16

q = torch.randn((96, 32, 128), device=device, dtype=dtype)
k = torch.randn((160, 8, 128), device=device, dtype=dtype)
v = torch.randn((160, 8, 128), device=device, dtype=dtype)

cu_seqlens_q = torch.tensor([0, 32, 96], device=device, dtype=torch.int32)
cu_seqlens_k = torch.tensor([0, 64, 160], device=device, dtype=torch.int32)

out = flash_attn_varlen_func(
    q=q,
    k=k,
    v=v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=64,
    max_seqlen_k=96,
    causal=False,
)

If your integration uses KV cache, paged KV, scheduler metadata, FP8, or only_qv, check the compatibility matrix below before porting more code.

Minimal KV-cache example:

import torch
from flash_attn_interface import flash_attn_with_kvcache

device = "musa"
dtype = torch.bfloat16

q = torch.randn((1, 4, 32, 128), device=device, dtype=dtype)
k_cache = torch.randn((8, 16, 8, 128), device=device, dtype=dtype)
v_cache = torch.randn((8, 16, 8, 128), device=device, dtype=dtype)
cache_seqlens = torch.tensor([32], device=device, dtype=torch.int32)
block_table = torch.tensor([[0, 1]], device=device, dtype=torch.int32)

out, softmax_lse = flash_attn_with_kvcache(
    q=q,
    k_cache=k_cache,
    v_cache=v_cache,
    cache_seqlens=cache_seqlens,
    block_table=block_table,
    causal=True,
)

At a Glance

Area	Status	Notes
Q Mode	✅ Supported	Normal, Ragged, Padded
KV Mode	✅ Supported	Normal, Ragged, Padded, Paged
Append New KV	✅ Supported	flash_attn_with_kvcache appends via k / v; packed new KV is supported via cu_seqlens_k_new
RoPE Input	✅ Supported	flash_attn_with_kvcache supports rotary_cos / rotary_sin, Interleaved, and Non-interleaved
Cache Index Options	✅ Supported	cache_batch_idx, cache_leftpad
Mask Mode	✅ Supported	None, Causal, Local, Local + attention_chunk
Score Mode	✅ Supported	Standard softmax and softcap
Page Size	✅ Supported	1, 16, 64, and arbitrary page sizes
Dtype	✅ Supported	bf16, fp16, torch.float8_e4m3fn, and torch.float8_e5m2 forward inputs; FP8 uses q_descale / k_descale / v_descale scaling
QV Input	✅ Supported	The forward path supports an optional qv input, including FP8 inputs with qv
HeadDim	✅ Supported	headdim <= 512, subject to dtype- and qv-specific alignment rules below
Optimization	✅ Supported	SplitKV, PackGQA, SchedulerMetadata
Output	✅ Supported	out, softmax_lse

MATE Extensions

Extension	Status	Notes
Context Parallel	✅ Supported	cp_world_size, cp_rank, cp_tot_seqused_k
Learnable Sink	✅ Supported	Supported on the local-attention path

Head-Dimension Alignment

The forward path supports headdim <= 512, but the supported range does not mean that arbitrary head dimensions are accepted. The alignment rules are:

Configuration	Q/K head dimension	V head dimension
FP16 or BF16 without qv	Must be divisible by 2	Must be divisible by 2
FP8 without qv	Must be divisible by 4	Must be divisible by 4
With qv	Must be divisible by 8	Must be divisible by 8

These are alignment gates in addition to the other API and kernel limits. For example, FP16/BF16 D=130 and FP8 D=132 are alignment-valid without qv, while FP16/BF16 D=129 and FP8 D=130 are invalid. With qv, Q/K D=128 and V D=256 are alignment-valid, while either dimension equal to 130 is invalid. The qv path also requires V D <= 512.

Notes

• This page summarizes the compatibility surface, not every internal kernel detail.

• FP8 forward support includes torch.float8_e4m3fn and torch.float8_e5m2; pass optional q_descale, k_descale, and v_descale tensors with shape (batch_size, num_heads_kv) when scale factors are required.

• When both q and the optional qv input are FP8, q_descale applies to both query tensors; k_descale and v_descale still apply to the KV inputs.

• RoPE is supported only when appending new KV through k / v; rotary_dim must be <= headdim and divisible by 16.

• Chunked forward is optimized for large headdim, including Gemma 512-512 attention.

• Local + attention_chunk requires MUSA SDK >= 5.1.0.

• FP8 attention works on the forward path today. For best performance, use MUSA SDK 5.2.0 or newer when available.

• For wrapper-level usage, see the FlashAttention wrapper page.