319 lines
14 KiB
C
319 lines
14 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_attn.h
|
|
* @brief Attention family: fp32 KV cache append, flash decode (the batch-1
|
|
* vector path, split-K across workgroups plus a combine kernel), and flash
|
|
* prefill (tiled causal GQA with online softmax).
|
|
*
|
|
* All three kernels compute exactly one thing — softmax(scale * qK^T) V over
|
|
* a causal key range — and differ only in how that range is carved up:
|
|
*
|
|
* prefill one workgroup per (query row, kv head); keys [0, pos0+t]
|
|
* decode one workgroup per (split, kv head); keys [split_lo, split_hi)
|
|
* of the single query row, partial (m, l, acc) per split
|
|
* reduce one workgroup per query head; folds the splits back together
|
|
*
|
|
* The score matrix is never materialized: each workgroup carries a running
|
|
* max `m` and running sum `l` and rescales its output accumulator whenever
|
|
* the max moves (online softmax). Tile order therefore differs from a
|
|
* single-pass CPU softmax and the results are NOT bit-exact — the bar is
|
|
* ~1e-3 with a stable argmax (spec §11).
|
|
*
|
|
* @status kv_append, flash_decode + flash_decode_reduce and flash_prefill are
|
|
* implemented and tested (metal/tests/test_metal_attn.c, 14 cases) against
|
|
* an independent double-precision CPU reference at context depths 1, 2, 16,
|
|
* 17, 513 and 4097, including forced multi-split decode, empty splits, both
|
|
* causality probes and a constant-V normalization check. Worst drift
|
|
* 1.19e-07 absolute, 8.9e-06 of the output RMS.
|
|
* @issues None
|
|
* @todo
|
|
* - [ ] fm_attn_blk — per-KV-block mask classification so fully-masked
|
|
* blocks are skipped. Long-context lever, not needed for correctness.
|
|
*
|
|
* FUNCTION-CONSTANT INDICES OWNED BY THIS FAMILY: 260-279.
|
|
* Function constants share ONE index space across the whole assembled
|
|
* shader library, so each family reserves a block; two families colliding
|
|
* on an index fails the compile of every kernel in the backend, not just
|
|
* theirs. Allocated so far: 0-2 probe / 30 Q3_K / 50 Q5_K / 60 Q6_K / 70 Q8_0 / 80 Q4_K /
|
|
* 100-109 rope / 110-119 norm / 120-129 elementwise / 200-219 ssm /
|
|
* 220-239 ssm_chunked / 240-259 moe / 260-279 attn.
|
|
* In use here: 260 kv_groups, 261 head_dim. 262-279 free.
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_KERNEL_ATTN_H
|
|
#define FOUNDRY_METAL_KERNEL_ATTN_H
|
|
|
|
#include "metal/fm_metal.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* ================================================================
|
|
* Shape limits
|
|
*
|
|
* The threadgroup staging buffers are declared at these maxima because
|
|
* fm_metal's dispatch helpers never call setThreadgroupMemoryLength — a
|
|
* kernel's threadgroup arrays have to be statically sized. The staged K
|
|
* tile is FM_ATTN_TILE_KEYS * FM_ATTN_MAX_HEAD_DIM floats (16 KiB at the
|
|
* limits below), which is what caps them.
|
|
* ================================================================ */
|
|
|
|
/* Keys staged per tile. One tile is the unit the online softmax folds. */
|
|
#define FM_ATTN_TILE_KEYS 16
|
|
|
|
/* Largest head_dim any entry point accepts. Must be a multiple of 32
|
|
* (the SIMD width the QK dot is split across). */
|
|
#define FM_ATTN_MAX_HEAD_DIM 256
|
|
|
|
/* Largest GQA group (query heads per kv head) any entry point accepts.
|
|
* One SIMD group of the workgroup scores one query head, and a workgroup
|
|
* is 8 SIMD groups wide. */
|
|
#define FM_ATTN_MAX_GROUPS 8
|
|
|
|
/* Threads per workgroup for the flash kernels: 8 SIMD groups of 32. */
|
|
#define FM_ATTN_THREADS 256
|
|
|
|
/* Function-constant indices this family binds. Both flash kernels are
|
|
* specialized on them so the per-thread accumulator arrays have
|
|
* compile-time loop bounds and stay in registers. */
|
|
#define FM_ATTN_FC_KV_GROUPS 260
|
|
#define FM_ATTN_FC_HEAD_DIM 261
|
|
|
|
/* Workgroups the default decode plan aims for, in total across all kv heads.
|
|
*
|
|
* MEASURED by metal/tests/bench_attn_split.c on M4 / 10 GPU cores,
|
|
* n_head=16 n_head_kv=2 head_dim=256. Speedup over the unsplit kernel
|
|
* (2 workgroups, one per kv head):
|
|
*
|
|
* workgroups 2 4 8 16 32 64
|
|
* seq 128 1.00x 1.42x 1.66x --- tile rounding caps this length ---
|
|
* seq 512 1.00x 1.34x 1.73x 2.28x 2.19x 2.10x
|
|
* seq 4096 1.00x 2.57x 4.54x 6.68x 9.37x 8.69x
|
|
*
|
|
* Absolute milliseconds are not quoted because they moved 2x between runs on
|
|
* a shared machine; the ratios within a run were stable across five runs and
|
|
* are what the choice rests on. 32 is the smallest count at which the curve
|
|
* has flattened at every length tested, and 64 was never better than 32
|
|
* outside the noise. At long context the split is not a micro-optimization
|
|
* — it is most of the decode path's speed.
|
|
*
|
|
* Note the SHAPE of that table: the best split_len moves 16 -> 32 -> 256
|
|
* across the three lengths while the best workgroup count does not move at
|
|
* all. What the GPU wants held constant is the number of workgroups, which
|
|
* is why this plan is parameterized by workgroups rather than by the
|
|
* keys-per-split knob the CUDA decode path uses. */
|
|
#define FM_ATTN_DECODE_TARGET_WORKGROUPS 32
|
|
|
|
/* ================================================================
|
|
* Split-K planning
|
|
* ================================================================ */
|
|
|
|
/**
|
|
* Plans the decode split: how many workgroups share the key range and how
|
|
* many keys each one takes.
|
|
*
|
|
* `split_len` is rounded up to FM_ATTN_TILE_KEYS so every workgroup starts
|
|
* on a tile boundary, and the returned split count is recomputed from the
|
|
* rounded length — so no planned split is ever empty, and a short context
|
|
* naturally uses fewer workgroups than asked for rather than fanning out
|
|
* into slivers.
|
|
*
|
|
* @param seq_len Live keys in the cache (>= 1)
|
|
* @param n_head_kv Key/value heads; the dispatch runs
|
|
* n_head_kv * splits workgroups
|
|
* @param target_workgroups Total workgroups to aim for
|
|
* (0 selects FM_ATTN_DECODE_TARGET_WORKGROUPS)
|
|
* @param split_len_out Receives the keys per workgroup
|
|
* @return Number of splits, or 0 if the arguments are invalid
|
|
*
|
|
* Thread-safe: yes (pure function)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
int fm_attn_decode_plan(int seq_len, int n_head_kv, int target_workgroups,
|
|
int *split_len_out);
|
|
|
|
/**
|
|
* Floats of scratch the split-K partials need, for each of the three
|
|
* buffers fm_attn_flash_decode_at writes.
|
|
*
|
|
* @param n_head Total query heads (n_head_kv * kv_groups)
|
|
* @param head_dim Per-head dimension
|
|
* @param n_splits Splits the decode will be dispatched with
|
|
* @param m_floats Receives the element count for the m and l buffers
|
|
* @param acc_floats Receives the element count for the acc buffer
|
|
*
|
|
* Thread-safe: yes (pure function)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
void fm_attn_decode_scratch_floats(int n_head, int head_dim, int n_splits,
|
|
size_t *m_floats, size_t *acc_floats);
|
|
|
|
/* ================================================================
|
|
* KV cache append
|
|
* ================================================================ */
|
|
|
|
/**
|
|
* Appends `n` rows of post-RoPE K and V into the ring KV cache.
|
|
*
|
|
* Row t of the batch lands at cache position (pos0 + t) % max_seq, which is
|
|
* the same ring rule the host cache uses, so a host<->device sync stays a
|
|
* contiguous per-head row range.
|
|
*
|
|
* @param k_new [n][n_head_kv][head_dim] fp32
|
|
* @param v_new [n][n_head_kv][head_dim] fp32
|
|
* @param k_cache [n_head_kv][max_seq][head_dim] fp32 (written)
|
|
* @param v_cache [n_head_kv][max_seq][head_dim] fp32 (written)
|
|
* @param n Rows to append (0 is a no-op, not an error)
|
|
* @param n_head_kv Key/value heads
|
|
* @param head_dim Per-head dimension
|
|
* @param max_seq Ring capacity in rows
|
|
* @param pos0 Absolute position of row 0
|
|
*
|
|
* Q4_0 KV quantization is explicitly out of scope (spec §12).
|
|
*
|
|
* Thread-safe: no (shares the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_attn_kv_append_at(fm_bufref k_new, fm_bufref v_new,
|
|
fm_bufref k_cache, fm_bufref v_cache,
|
|
int n, int n_head_kv, int head_dim, int max_seq,
|
|
int pos0);
|
|
|
|
/* fm_attn_kv_append_at with every operand bound at offset 0. */
|
|
void fm_attn_kv_append(fm_buffer *k_new, fm_buffer *v_new,
|
|
fm_buffer *k_cache, fm_buffer *v_cache,
|
|
int n, int n_head_kv, int head_dim, int max_seq,
|
|
int pos0);
|
|
|
|
/* ================================================================
|
|
* Flash decode (split-K vector path)
|
|
* ================================================================ */
|
|
|
|
/**
|
|
* Scores one query row against key slice [split*split_len, ...) of the
|
|
* cache and writes that slice's partial online-softmax state.
|
|
*
|
|
* This is half of the decode path: the partials it writes are meaningless
|
|
* until fm_attn_flash_decode_reduce_at folds them. Both take `n_splits`, and
|
|
* it must be the same value in both calls — it is the stride of the partial
|
|
* arrays, not just a workgroup count.
|
|
*
|
|
* The caller owns the split plan (see fm_attn_decode_plan). Passing a plan
|
|
* whose splits run past seq_len is safe: an empty split writes m = -inf,
|
|
* l = 0 and a zero accumulator, which the reduce weights to zero.
|
|
*
|
|
* @param q [n_head][head_dim] fp32, the single query row
|
|
* @param k_cache [n_head_kv][max_seq][head_dim] fp32
|
|
* @param v_cache [n_head_kv][max_seq][head_dim] fp32
|
|
* @param m_partial [n_head][n_splits] fp32 (written) running max
|
|
* @param l_partial [n_head][n_splits] fp32 (written) running sum
|
|
* @param acc_partial [n_head][n_splits][head_dim] fp32 (written) UNnormalized
|
|
* @param kv_groups Query heads per kv head, 1..FM_ATTN_MAX_GROUPS
|
|
* @param n_head_kv Key/value heads
|
|
* @param head_dim Per-head dimension, multiple of 32, <= FM_ATTN_MAX_HEAD_DIM
|
|
* @param max_seq Ring capacity in rows
|
|
* @param seq_len Live keys, >= 1; the query attends [0, seq_len)
|
|
* @param n_splits Workgroups per kv head, >= 1
|
|
* @param split_len Keys per split; the caller's plan, not derived here
|
|
* @param scale 1/sqrt(head_dim)
|
|
*
|
|
* Thread-safe: no (shares the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_attn_flash_decode_at(fm_bufref q, fm_bufref k_cache, fm_bufref v_cache,
|
|
fm_bufref m_partial, fm_bufref l_partial,
|
|
fm_bufref acc_partial,
|
|
int kv_groups, int n_head_kv, int head_dim,
|
|
int max_seq, int seq_len, int n_splits,
|
|
int split_len, float scale);
|
|
|
|
/* fm_attn_flash_decode_at with every operand bound at offset 0. */
|
|
void fm_attn_flash_decode(fm_buffer *q, fm_buffer *k_cache, fm_buffer *v_cache,
|
|
fm_buffer *m_partial, fm_buffer *l_partial,
|
|
fm_buffer *acc_partial,
|
|
int kv_groups, int n_head_kv, int head_dim,
|
|
int max_seq, int seq_len, int n_splits,
|
|
int split_len, float scale);
|
|
|
|
/**
|
|
* Folds the per-split partials into the final attention output.
|
|
*
|
|
* Takes the max of the per-split maxima, reweights each split's sum and
|
|
* accumulator by exp(m_split - M), and divides. Algebraically identical to
|
|
* a single-pass softmax over the whole key range; a split that saw no keys
|
|
* carries m = -inf and drops out with weight zero.
|
|
*
|
|
* @param m_partial [n_head][n_splits] fp32 from fm_attn_flash_decode_at
|
|
* @param l_partial [n_head][n_splits] fp32
|
|
* @param acc_partial [n_head][n_splits][head_dim] fp32
|
|
* @param out [n_head][head_dim] fp32 (written)
|
|
* @param n_head Total query heads
|
|
* @param head_dim Per-head dimension
|
|
* @param n_splits Same value passed to fm_attn_flash_decode_at
|
|
*
|
|
* Thread-safe: no (shares the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_attn_flash_decode_reduce_at(fm_bufref m_partial, fm_bufref l_partial,
|
|
fm_bufref acc_partial, fm_bufref out,
|
|
int n_head, int head_dim, int n_splits);
|
|
|
|
/* fm_attn_flash_decode_reduce_at with every operand bound at offset 0. */
|
|
void fm_attn_flash_decode_reduce(fm_buffer *m_partial, fm_buffer *l_partial,
|
|
fm_buffer *acc_partial, fm_buffer *out,
|
|
int n_head, int head_dim, int n_splits);
|
|
|
|
/* ================================================================
|
|
* Flash prefill (tiled causal GQA)
|
|
* ================================================================ */
|
|
|
|
/**
|
|
* Causal GQA attention for a batch of query rows against the cache.
|
|
*
|
|
* Query row t sits at absolute position pos0 + t and attends cache rows
|
|
* [0, pos0 + t] — causality is enforced by the key loop bound and by
|
|
* clamping the final tile to the live key count, so no future key is ever
|
|
* read, scored or weighted. The caller must have appended rows [0, n) to
|
|
* the cache first, and must guarantee pos0 + n <= max_seq (no ring wrap
|
|
* inside the attended window).
|
|
*
|
|
* @param q [n][n_head][head_dim] fp32, post QK-norm and RoPE
|
|
* @param k_cache [n_head_kv][max_seq][head_dim] fp32
|
|
* @param v_cache [n_head_kv][max_seq][head_dim] fp32
|
|
* @param out [n][n_head][head_dim] fp32 (written)
|
|
* @param n Query rows (0 is a no-op, not an error)
|
|
* @param kv_groups Query heads per kv head, 1..FM_ATTN_MAX_GROUPS
|
|
* @param n_head_kv Key/value heads
|
|
* @param head_dim Per-head dimension, multiple of 32, <= FM_ATTN_MAX_HEAD_DIM
|
|
* @param max_seq Ring capacity in rows
|
|
* @param pos0 Absolute position of query row 0
|
|
* @param scale 1/sqrt(head_dim)
|
|
*
|
|
* Thread-safe: no (shares the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_attn_flash_prefill_at(fm_bufref q, fm_bufref k_cache, fm_bufref v_cache,
|
|
fm_bufref out, int n, int kv_groups,
|
|
int n_head_kv, int head_dim, int max_seq,
|
|
int pos0, float scale);
|
|
|
|
/* fm_attn_flash_prefill_at with every operand bound at offset 0. */
|
|
void fm_attn_flash_prefill(fm_buffer *q, fm_buffer *k_cache, fm_buffer *v_cache,
|
|
fm_buffer *out, int n, int kv_groups,
|
|
int n_head_kv, int head_dim, int max_seq,
|
|
int pos0, float scale);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_METAL_KERNEL_ATTN_H */
|