std1.0.0-beta-7 SafeC docs

std :: ml::gpu_mps

Struct std::MpsAttnShape

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

struct MpsAttnShape {
    unsigned long batch;
    unsigned long heads;
    unsigned long kvHeads;
    unsigned long seqLen;
    unsigned long kvLen;
    unsigned long headDim;
    unsigned long window;
}

Batched multi-head versions of attention.h's flash_attention_forward and linear_attention_forward on the GPU, in fp32, fp16 or bf16 (all math in fp32). Two calling forms share the same kernels:

  • mps_flash_attention / mps_linear_attention (and the host-array f32 versions) take the same arguments as gpu_rocm.h's and gpu_cuda_attention.h's rocm/cuda_ ops, so code moves between AMD, NVIDIA and Apple GPUs unchanged: fp32 device buffers (from MpsBuf::upload), q and out [heads][seqLen][headDim], k and v [heads][kvLen][headDim] (fold a batch into heads).
  • The _ex versions take an MpsAttnShape (a batch, grouped-query K/V with kvHeads dividing heads -- query head h reads kv head h / (heads / kvHeads) -- and a sliding window) and a dtype: 0 = fp32, MPS_FP16, MPS_BF16 (raw 16-bit values, as from MpsBuf::upload_f16/upload_bf16). Q and out are then [batch][heads][seqLen][headDim], K and V [batch][kvHeads][kvLen][headDim]. All buffers are row-major and contiguous.

Flash attention: softmax(q k^T * scale) v without forming the score matrix -- FlashAttention-2 on the GPU's simdgroup matrix units for head dims up to 256, a simdgroup-reduction kernel for 512. 'scale' <= 0 means 1/sqrt(headDim). Linear attention: phi(x) = elu(x)+1, out_i = phi(q_i) S / phi(q_i) z with S = sum_j phi(k_j) v_j^T and z = sum_j phi(k_j) -- no scale, since the division cancels it.

Masking: query i sits at position pos = i + (kvLen - seqLen) -- the last seqLen keys are the queries' own tokens, any before them a KV cache. causal hides keys j > pos (needs kvLen >= seqLen); a window w > 0 (flash only) also hides keys j <= pos - w, i.e. the query sees its last w positions. A query with no visible key gets zeros.

Head dims: 32, 64, 96, 128, 256 or 512 (mps_attention_head_dim_ok). The device-buffer forms encode onto the open batch (results valid after mps_batch_end()) or, outside one, run and wait. 1 = done/encoded, 0 = bad shape or no Metal device -- callers fall back to a CPU path.

Fields

§unsigned long batchsource
§unsigned long headssource

query heads

§unsigned long kvHeadssource

key/value heads; divides heads

§unsigned long seqLensource

query length

§unsigned long kvLensource

key/value length

§unsigned long headDimsource

32, 64, 96, 128, 256 or 512

§unsigned long windowsource

flash attention sliding window; 0 = none