std :: ml::gpu_mps
Function std::mps_flash_attn_bf16
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.