std :: ml::gpu_mps
Function std::mps_attn_batch_bf16
int mps_attn_batch_bf16(
const &unsigned short q,
const &unsigned short kAll,
const &unsigned short vAll,
const &unsigned int cacheLens,
&unsigned short out,
unsigned long batchSize,
unsigned long numHeads,
unsigned long numKvHeads,
unsigned long headDim,
unsigned long maxKvLen,
int isSliding,
unsigned long slidingWindow,
)Scaled-dot-product attention for one new query token per sequence (a decode step) against each sequence's own K/V history, across the whole batch and every head in one dispatch: grouped-query heads, optional sliding window, softmax, weighted sum over V. No 1/sqrt(headDim) factor is applied -- fold any scaling into q. See gpu_mps_kernels.metal's attn_batch_bf16_kernel for the kernel design.
q: [batchSize, numHeads*headDim] raw bf16 bits.
kAll/vAll: [batchSize, maxKvLen, numKvHeads*headDim] raw bf16 bits --
each sequence's full K/V history including the
just-appended token, gathered contiguously by the caller.
cacheLens: [batchSize] -- each sequence's valid KV length
(<= maxKvLen); rows past it are never read.
out: [batchSize, numHeads*headDim] raw bf16 bits (pre-o_proj).
One-shot dispatch (own commit+wait+readback). Returns 1 on success, 0 if no Metal device is available, pipeline compilation fails, headDim isn't 64, 128, 256 or 512, or numHeads isn't a multiple of numKvHeads -- callers fall back to a CPU path, same contract as every other mps_*_bf16 entry point in this file. There is no KV-length limit.