std :: ml::gpu_mps
Function std::mps_linear_attn_bf16
int mps_linear_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,
int causal,
)Linear attention on token-major bf16, non-causal or causal (causal needs kvLen == seqLen here). Batch-aware like mps_flash_attn_bf16.