std1.0.0-beta-7 SafeC docs

std :: ml::gpu_mps

Function std::mps_attn_batch_bf16

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

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.