std :: ml::gpu_mps
Struct std::MpsAttnShape
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
query heads
key/value heads; divides heads
query length
key/value length
32, 64, 96, 128, 256 or 512
flash attention sliding window; 0 = none