std1.0.0-beta-7 SafeC docs

std :: ml::gpu_mps

Function std::mps_flash_attn_bf16

#include <std/ml/gpu_mps.h> source · definition

int mps_flash_attn_bf16(
    const &unsigned short q,
    const &unsigned short k,
    const &unsigned short v,
    &unsigned short out,
    unsigned long seqLen,
    unsigned long kvLen,
    unsigned long numHeads,
    unsigned long numKvHeads,
    unsigned long headDim,
    float scale,
    int causal,
    int isSliding,
    unsigned long slidingWindow,
)

The same flash and linear attention on one sequence laid out the way a transformer's projections produce it -- q/out [seqLen, numHeadsheadDim], k/v [kvLen, numKvHeadsheadDim], raw bf16 bits -- plus a recurrent linear attention whose state stays on the GPU. Use mps_attn_batch_bf16 above for one-query-per-sequence decode steps across a batch.

mps_flash_attn_bf16: 'scale' as above (pass 1.0 for models that pre-scale q); isSliding/slidingWindow as the window w. Batch-aware like mps_sub_f32: inside an open batch the result lands in 'out' at mps_batch_end(); otherwise on return. Returns 0 on the conditions above.