std1.0.0-beta-7 SafeC docs

std :: ml::gpu_rocm

Function std::rocm_flash_attention

#include <std/ml/gpu_rocm.h> source ยท definition

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.