std :: ml::gpu_rocm
Function std::rocm_flash_attention
int rocm_flash_attention(
struct RocmBuf q,
struct RocmBuf k,
struct RocmBuf v,
struct RocmBuf out,
unsigned long heads,
unsigned long seqLen,
unsigned long kvLen,
unsigned long headDim,
float scale,
int causal,
)fp32, batched over heads: q and out are [heads][seqLen][headDim], k and v [heads][kvLen][headDim], all row-major; headDim <= 128. 1 on success, 0 on a bad shape or failure. Forward only, queued on the stream like every op.
On GPUs with matrix cores (CDNA: MI100 / MI200 / MI300) both ops run on them (fp32 MFMA); elsewhere on CUDA-core kernels with the same results up to rounding.
Flash attention: softmax(q k^T * scale) v without forming the score matrix -- K/V tiles in LDS and an online softmax per query row (Dao et al.), the same result as std/ml/attention.h's attention_forward up to rounding. scale <= 0 means 1/sqrt(headDim). causal: query i sees keys j <= i + (kvLen - seqLen) (so kvLen >= seqLen; with a KV cache the new queries are the last seqLen positions). A query with no visible key gets zeros.