580 lines
28 KiB
Mathematica
580 lines
28 KiB
Mathematica
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_attn.m
|
|
* @brief Host dispatch and MSL source for the attn kernel family: fp32 KV
|
|
* append, split-K flash decode plus its combine, and tiled causal GQA
|
|
* flash prefill.
|
|
*
|
|
* @status All four kernels implemented and tested (metal/tests/test_metal_attn.c).
|
|
* @issues The tile loop is shared by decode and prefill through the MSL helper
|
|
* fmattn_accumulate_range, which takes the per-thread `acc` and `qreg`
|
|
* arrays as thread-address-space pointers. Both loops over them are bounded
|
|
* by function constants, so after inlining the indices are constants and the
|
|
* arrays promote to registers — but that promotion is the compiler's choice,
|
|
* not something the source guarantees. If this ever shows up as a perf cliff,
|
|
* the fix is to duplicate the body into each kernel, not to add a runtime
|
|
* bound.
|
|
* @todo
|
|
* - [ ] fm_attn_blk — per-KV-block mask classification (fully masked → skip).
|
|
*
|
|
* The CUDA sibling (cuda_kernels/attention.cu) reduces each q.k dot across a
|
|
* 16-lane sub-warp with __shfl_down_sync(..., 16). There is no 16-wide
|
|
* sub-group on a 32-wide Metal SIMD group, so the mapping is different here:
|
|
* one SIMD GROUP owns one query head of the GQA group and reduces a whole
|
|
* 256-element dot with a single simd_sum, and the query row lives in
|
|
* REGISTERS (8 floats per lane) instead of threadgroup memory. Same math,
|
|
* different decomposition.
|
|
*/
|
|
|
|
#include "metal/kernels/fm_kernel_attn.h"
|
|
|
|
#include "metal/fm_metal_internal.h"
|
|
#include "metal/kernels/fm_kernel_registry.h"
|
|
|
|
#include <string.h>
|
|
|
|
/* The MSL below sizes its threadgroup arrays with literals, because
|
|
* fm_metal's dispatch helpers never call setThreadgroupMemoryLength and a
|
|
* threadgroup array therefore has to be statically sized. These assertions
|
|
* are what keeps those literals honest. */
|
|
_Static_assert(FM_ATTN_TILE_KEYS == 16, "MSL k_tile/score_s/p_s assume 16");
|
|
_Static_assert(FM_ATTN_MAX_HEAD_DIM == 256, "MSL k_tile assumes 256");
|
|
_Static_assert(FM_ATTN_MAX_GROUPS == 8, "MSL score_s/p_s/m_s assume 8");
|
|
_Static_assert(FM_ATTN_THREADS == 256, "MSL assumes 8 SIMD groups of 32");
|
|
|
|
/* ================================================================
|
|
* Host dispatch
|
|
* ================================================================ */
|
|
|
|
/* Shapes the flash kernels accept. head_dim must be a multiple of the SIMD
|
|
* width because the QK dot splits one head across 32 lanes with no tail, and
|
|
* it must not exceed FM_ATTN_MAX_HEAD_DIM because one thread owns one output
|
|
* dimension and the workgroup is FM_ATTN_THREADS wide. */
|
|
static bool fm_attn_shape_is_valid(int kv_groups, int head_dim)
|
|
{
|
|
if (kv_groups < 1 || kv_groups > FM_ATTN_MAX_GROUPS) return false;
|
|
if (head_dim < 32 || head_dim > FM_ATTN_MAX_HEAD_DIM) return false;
|
|
return head_dim % 32 == 0;
|
|
}
|
|
|
|
/* kv_groups and head_dim are function constants, not kernel arguments: they
|
|
* bound the per-thread acc[] and qreg[] loops, and a runtime bound would
|
|
* leave those arrays in thread-local memory instead of registers. */
|
|
static int fm_attn_specialize(fm_fc_values *fc, int kv_groups, int head_dim)
|
|
{
|
|
fm_fc_reset(fc);
|
|
if (fm_fc_set_int(fc, FM_ATTN_FC_KV_GROUPS, kv_groups) != 0) return -1;
|
|
if (fm_fc_set_int(fc, FM_ATTN_FC_HEAD_DIM, head_dim) != 0) return -1;
|
|
return 0;
|
|
}
|
|
|
|
/* The dispatch helpers carry uint32 arguments only, so a float rides across
|
|
* as its bit pattern and the shader reads it back with as_type<float>. */
|
|
static uint32_t fm_attn_float_bits(float value)
|
|
{
|
|
uint32_t bits = 0;
|
|
memcpy(&bits, &value, sizeof(bits));
|
|
return bits;
|
|
}
|
|
|
|
int fm_attn_decode_plan(int seq_len, int n_head_kv, int target_workgroups,
|
|
int *split_len_out)
|
|
{
|
|
if (split_len_out) *split_len_out = 0;
|
|
if (seq_len < 1 || n_head_kv < 1) return 0;
|
|
if (target_workgroups <= 0)
|
|
target_workgroups = FM_ATTN_DECODE_TARGET_WORKGROUPS;
|
|
|
|
int n_splits = target_workgroups / n_head_kv;
|
|
if (n_splits < 1) n_splits = 1;
|
|
|
|
/* Round the slice up to a whole tile so every workgroup starts on a tile
|
|
* boundary, then recompute the count from the rounded length — otherwise
|
|
* the last split would be empty whenever the rounding overshoots. This
|
|
* is also what keeps a short context from fanning out into slivers: at
|
|
* seq_len 128 the requested 16 splits round to 8 tiles of 16 keys. */
|
|
int split_len = (seq_len + n_splits - 1) / n_splits;
|
|
split_len = ((split_len + FM_ATTN_TILE_KEYS - 1) / FM_ATTN_TILE_KEYS) *
|
|
FM_ATTN_TILE_KEYS;
|
|
if (split_len_out) *split_len_out = split_len;
|
|
return (seq_len + split_len - 1) / split_len;
|
|
}
|
|
|
|
void fm_attn_decode_scratch_floats(int n_head, int head_dim, int n_splits,
|
|
size_t *m_floats, size_t *acc_floats)
|
|
{
|
|
size_t partials = (size_t)(n_head > 0 ? n_head : 0) *
|
|
(size_t)(n_splits > 0 ? n_splits : 0);
|
|
if (m_floats) *m_floats = partials;
|
|
if (acc_floats) *acc_floats = partials * (size_t)(head_dim > 0 ? head_dim : 0);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!fm_metal_is_initialized()) return;
|
|
if (!k_new.buffer || !v_new.buffer || !k_cache.buffer || !v_cache.buffer) return;
|
|
if (n <= 0 || n_head_kv <= 0 || head_dim <= 0 || max_seq <= 0 || pos0 < 0) return;
|
|
|
|
int total = n * n_head_kv * head_dim;
|
|
fm_bufref refs[] = { k_new, v_new, k_cache, v_cache };
|
|
uint32_t consts[] = { (uint32_t)n_head_kv, (uint32_t)head_dim,
|
|
(uint32_t)max_seq, (uint32_t)pos0, (uint32_t)total };
|
|
fm_dispatch_1d_refs(@"fmattn_kv_append", NULL, refs, 4, consts, 5, total);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!k_new || !v_new || !k_cache || !v_cache) return;
|
|
fm_attn_kv_append_at(fm_bufref_at(k_new, 0), fm_bufref_at(v_new, 0),
|
|
fm_bufref_at(k_cache, 0), fm_bufref_at(v_cache, 0),
|
|
n, n_head_kv, head_dim, max_seq, pos0);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!fm_metal_is_initialized()) return;
|
|
if (!q.buffer || !k_cache.buffer || !v_cache.buffer) return;
|
|
if (!m_partial.buffer || !l_partial.buffer || !acc_partial.buffer) return;
|
|
if (!fm_attn_shape_is_valid(kv_groups, head_dim)) return;
|
|
if (n_head_kv <= 0 || max_seq <= 0 || seq_len < 1) return;
|
|
if (n_splits < 1 || split_len < 1) return;
|
|
|
|
fm_fc_values fc;
|
|
if (fm_attn_specialize(&fc, kv_groups, head_dim) != 0) return;
|
|
|
|
fm_bufref refs[] = { q, k_cache, v_cache, m_partial, l_partial, acc_partial };
|
|
uint32_t consts[] = { (uint32_t)max_seq, (uint32_t)seq_len,
|
|
(uint32_t)n_splits, (uint32_t)split_len,
|
|
fm_attn_float_bits(scale) };
|
|
/* One workgroup per (kv head, split): the split-K fan-out. At batch 1 a
|
|
* single workgroup per kv head leaves a 10-core M4 almost idle. */
|
|
fm_dispatch_threadgroups_refs(@"fmattn_flash_decode", &fc, refs, 6,
|
|
consts, 5, n_head_kv * n_splits,
|
|
FM_ATTN_THREADS);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!q || !k_cache || !v_cache || !m_partial || !l_partial || !acc_partial)
|
|
return;
|
|
fm_attn_flash_decode_at(fm_bufref_at(q, 0), fm_bufref_at(k_cache, 0),
|
|
fm_bufref_at(v_cache, 0),
|
|
fm_bufref_at(m_partial, 0),
|
|
fm_bufref_at(l_partial, 0),
|
|
fm_bufref_at(acc_partial, 0),
|
|
kv_groups, n_head_kv, head_dim, max_seq, seq_len,
|
|
n_splits, split_len, scale);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!fm_metal_is_initialized()) return;
|
|
if (!m_partial.buffer || !l_partial.buffer || !acc_partial.buffer || !out.buffer)
|
|
return;
|
|
if (n_head <= 0 || head_dim <= 0 || n_splits < 1) return;
|
|
|
|
fm_bufref refs[] = { m_partial, l_partial, acc_partial, out };
|
|
uint32_t consts[] = { (uint32_t)head_dim, (uint32_t)n_splits };
|
|
fm_dispatch_threadgroups_refs(@"fmattn_flash_decode_reduce", NULL, refs, 4,
|
|
consts, 2, n_head, head_dim);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!m_partial || !l_partial || !acc_partial || !out) return;
|
|
fm_attn_flash_decode_reduce_at(fm_bufref_at(m_partial, 0),
|
|
fm_bufref_at(l_partial, 0),
|
|
fm_bufref_at(acc_partial, 0),
|
|
fm_bufref_at(out, 0),
|
|
n_head, head_dim, n_splits);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!fm_metal_is_initialized()) return;
|
|
if (!q.buffer || !k_cache.buffer || !v_cache.buffer || !out.buffer) return;
|
|
if (!fm_attn_shape_is_valid(kv_groups, head_dim)) return;
|
|
if (n <= 0 || n_head_kv <= 0 || max_seq <= 0 || pos0 < 0) return;
|
|
if (pos0 + n > max_seq) return; /* would wrap inside the attended window */
|
|
|
|
fm_fc_values fc;
|
|
if (fm_attn_specialize(&fc, kv_groups, head_dim) != 0) return;
|
|
|
|
fm_bufref refs[] = { q, k_cache, v_cache, out };
|
|
uint32_t consts[] = { (uint32_t)n_head_kv, (uint32_t)max_seq,
|
|
(uint32_t)pos0, fm_attn_float_bits(scale) };
|
|
/* One workgroup per (query row, kv head). The whole GQA group shares one
|
|
* workgroup so each staged K tile is read once and scored by all of it. */
|
|
fm_dispatch_threadgroups_refs(@"fmattn_flash_prefill", &fc, refs, 4,
|
|
consts, 4, n * n_head_kv, FM_ATTN_THREADS);
|
|
}
|
|
|
|
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)
|
|
{
|
|
if (!q || !k_cache || !v_cache || !out) return;
|
|
fm_attn_flash_prefill_at(fm_bufref_at(q, 0), fm_bufref_at(k_cache, 0),
|
|
fm_bufref_at(v_cache, 0), fm_bufref_at(out, 0),
|
|
n, kv_groups, n_head_kv, head_dim, max_seq, pos0,
|
|
scale);
|
|
}
|
|
|
|
/* ================================================================
|
|
* MSL fragment
|
|
* ================================================================ */
|
|
|
|
NSString *fm_kernel_msl_attn(void)
|
|
{
|
|
return @
|
|
"/* ============================================================\n"
|
|
" * Attention: causal GQA over an fp32 ring KV cache.\n"
|
|
" *\n"
|
|
" * Layouts (all fp32, all row-major):\n"
|
|
" * q [rows][n_head][head_dim], global head = kvh*G + g\n"
|
|
" * k_cache [n_head_kv][max_seq][head_dim]\n"
|
|
" * v_cache same as k_cache\n"
|
|
" *\n"
|
|
" * Decomposition. A workgroup is 8 SIMD groups of 32 lanes and owns ONE\n"
|
|
" * kv head. SIMD group g scores query head g of the GQA group; its 32\n"
|
|
" * lanes each hold head_dim/32 elements of that query row in REGISTERS\n"
|
|
" * for the whole kernel, and one simd_sum turns a lane-partial into a\n"
|
|
" * whole q.k dot. Thread tid owns output dimension tid during the PV\n"
|
|
" * accumulate. Only K is staged in threadgroup memory, because it is the\n"
|
|
" * one operand every SIMD group of the workgroup re-reads.\n"
|
|
" *\n"
|
|
" * Online softmax. Per tile: fold the tile's scores into the running max\n"
|
|
" * m, rescale the running sum l and the output accumulator by\n"
|
|
" * exp(m_old - m_new), add the tile. The [seq][seq] score matrix is never\n"
|
|
" * materialized, and the result differs from a single-pass softmax only\n"
|
|
" * by reassociation.\n"
|
|
" *\n"
|
|
" * Causality is a key COUNT, not a -inf fill: the loop stops at the\n"
|
|
" * query's own position and the final tile is clamped to the live key\n"
|
|
" * count, so a future key is never read, never scored and never weighted.\n"
|
|
" * ============================================================ */\n"
|
|
"\n"
|
|
"constant int fmattn_tile_keys = 16;\n"
|
|
"constant int fmattn_max_head_dim = 256;\n"
|
|
"constant int fmattn_max_groups = 8;\n"
|
|
"constant int fmattn_threads = 256;\n"
|
|
"constant int fmattn_simd_width = 32;\n"
|
|
"\n"
|
|
"/* 'No key has been folded in yet.' NOT -INFINITY: fm_init compiles this\n"
|
|
" * library with MTLMathModeFast, under which an operand may be assumed\n"
|
|
" * finite, so a real infinity here is undefined behaviour. A finite floor\n"
|
|
" * behaves identically for every use it has — a real score can never reach\n"
|
|
" * it, max() picks the real score, and exp(floor - anything_real)\n"
|
|
" * underflows to exactly 0, which is the weight an empty range needs. */\n"
|
|
"constant float fmattn_score_floor = -1.0e30f;\n"
|
|
"\n"
|
|
"/* Shapes as function constants so every loop over acc[] and qreg[] has a\n"
|
|
" * compile-time bound and both arrays stay in registers. */\n"
|
|
"constant int fmattn_kv_groups_in [[function_constant(260)]];\n"
|
|
"constant int fmattn_kv_groups =\n"
|
|
" is_function_constant_defined(fmattn_kv_groups_in) ? fmattn_kv_groups_in\n"
|
|
" : fmattn_max_groups;\n"
|
|
"constant int fmattn_head_dim_in [[function_constant(261)]];\n"
|
|
"constant int fmattn_head_dim =\n"
|
|
" is_function_constant_defined(fmattn_head_dim_in) ? fmattn_head_dim_in\n"
|
|
" : fmattn_max_head_dim;\n"
|
|
"constant int fmattn_q_per_lane = fmattn_head_dim / fmattn_simd_width;\n"
|
|
"\n"
|
|
"/* Append one post-RoPE K/V element to the ring cache. One thread per\n"
|
|
" * element of k_new; row t of the batch lands at (pos0 + t) % max_seq. */\n"
|
|
"kernel void fmattn_kv_append(\n"
|
|
" device const float* k_new [[buffer(0)]],\n"
|
|
" device const float* v_new [[buffer(1)]],\n"
|
|
" device float* k_cache [[buffer(2)]],\n"
|
|
" device float* v_cache [[buffer(3)]],\n"
|
|
" constant uint& n_head_kv [[buffer(4)]],\n"
|
|
" constant uint& head_dim [[buffer(5)]],\n"
|
|
" constant uint& max_seq [[buffer(6)]],\n"
|
|
" constant uint& pos0 [[buffer(7)]],\n"
|
|
" constant uint& total [[buffer(8)]],\n"
|
|
" uint i [[thread_position_in_grid]]\n"
|
|
") {\n"
|
|
" if (i >= total) return;\n"
|
|
" uint d = i % head_dim;\n"
|
|
" uint th = i / head_dim; // flattened (row, kv head)\n"
|
|
" uint h = th % n_head_kv;\n"
|
|
" uint t = th / n_head_kv;\n"
|
|
" uint cache_pos = (pos0 + t) % max_seq;\n"
|
|
" ulong dst = ((ulong)h * max_seq + cache_pos) * head_dim + d;\n"
|
|
" k_cache[dst] = k_new[i];\n"
|
|
" v_cache[dst] = v_new[i];\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* Online-softmax pass over keys [k_lo, k_hi) for the whole GQA group.\n"
|
|
" *\n"
|
|
" * On entry m_s/l_s hold the running state (-inf / 0 for a fresh pass) and\n"
|
|
" * acc[] the running output; on exit all three have absorbed the range.\n"
|
|
" * Shared verbatim by decode (which passes one split's slice) and prefill\n"
|
|
" * (which passes [0, query position]), because they differ only in which\n"
|
|
" * range they hand over and what they do with the result. */\n"
|
|
"static void fmattn_accumulate_range(\n"
|
|
" device const float* k_head, device const float* v_head,\n"
|
|
" thread float* qreg, thread float* acc,\n"
|
|
" threadgroup float* k_tile, threadgroup float* score_s,\n"
|
|
" threadgroup float* p_s, threadgroup float* m_s,\n"
|
|
" threadgroup float* l_s, threadgroup float* corr_s,\n"
|
|
" uint k_lo, uint k_hi, float scale, uint tid, uint sgid, uint lane)\n"
|
|
"{\n"
|
|
" const uint HD = uint(fmattn_head_dim);\n"
|
|
" const uint G = uint(fmattn_kv_groups);\n"
|
|
" const uint TK = uint(fmattn_tile_keys);\n"
|
|
" const uint MG = uint(fmattn_max_groups);\n"
|
|
"\n"
|
|
" for (uint j0 = k_lo; j0 < k_hi; j0 += TK) {\n"
|
|
" uint tk = min(TK, k_hi - j0);\n"
|
|
"\n"
|
|
" /* Stage the tile's K rows once for the whole group. Flattened so\n"
|
|
" * consecutive threads read consecutive addresses whatever HD is. */\n"
|
|
" for (uint e = tid; e < tk * HD; e += uint(fmattn_threads))\n"
|
|
" k_tile[e] = k_head[(ulong)(j0 + e / HD) * HD + (e % HD)];\n"
|
|
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
|
"\n"
|
|
" /* q.k for every (head, key) in the tile. sgid and tk are uniform\n"
|
|
" * across the SIMD group, so simd_sum sees all 32 lanes. */\n"
|
|
" if (sgid < G) {\n"
|
|
" for (uint j = 0; j < tk; j++) {\n"
|
|
" float part = 0.0f;\n"
|
|
" threadgroup const float* krow = k_tile + (ulong)j * HD;\n"
|
|
" for (int i = 0; i < fmattn_q_per_lane; i++)\n"
|
|
" part += qreg[i] * krow[lane + uint(i) * 32u];\n"
|
|
" float dot = simd_sum(part);\n"
|
|
" if (lane == 0) score_s[sgid * TK + j] = dot * scale;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
|
"\n"
|
|
" /* Fold the tile into (m, l), one thread per head. p_s is stored\n"
|
|
" * key-major so the PV loop below reads a contiguous run. */\n"
|
|
" if (tid < G) {\n"
|
|
" threadgroup const float* sc = score_s + tid * TK;\n"
|
|
" float m_new = m_s[tid];\n"
|
|
" for (uint j = 0; j < tk; j++) m_new = max(m_new, sc[j]);\n"
|
|
" float corr = exp(m_s[tid] - m_new);\n"
|
|
" float sum = 0.0f;\n"
|
|
" for (uint j = 0; j < tk; j++) {\n"
|
|
" float pj = exp(sc[j] - m_new);\n"
|
|
" p_s[j * MG + tid] = pj;\n"
|
|
" sum += pj;\n"
|
|
" }\n"
|
|
" l_s[tid] = l_s[tid] * corr + sum;\n"
|
|
" m_s[tid] = m_new;\n"
|
|
" corr_s[tid] = corr;\n"
|
|
" }\n"
|
|
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
|
"\n"
|
|
" /* Rescale to the tile's max, then accumulate PV. Thread tid owns\n"
|
|
" * output dim tid, so the V read is coalesced across the group. */\n"
|
|
" if (tid < HD) {\n"
|
|
" for (int g = 0; g < fmattn_kv_groups; g++) acc[g] *= corr_s[g];\n"
|
|
" for (uint j = 0; j < tk; j++) {\n"
|
|
" float vv = v_head[(ulong)(j0 + j) * HD + tid];\n"
|
|
" threadgroup const float* pj = p_s + j * MG;\n"
|
|
" for (int g = 0; g < fmattn_kv_groups; g++)\n"
|
|
" acc[g] += pj[g] * vv;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
|
" }\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* Loads this SIMD group's query head into registers: lane holds elements\n"
|
|
" * lane, lane+32, ... of head `sgid` of the group based at q_base. */\n"
|
|
"static void fmattn_load_query(device const float* q_base, thread float* qreg,\n"
|
|
" uint sgid, uint lane)\n"
|
|
"{\n"
|
|
" for (int i = 0; i < fmattn_q_per_lane; i++) qreg[i] = 0.0f;\n"
|
|
" if (sgid >= uint(fmattn_kv_groups)) return;\n"
|
|
" device const float* head = q_base + (ulong)sgid * uint(fmattn_head_dim);\n"
|
|
" for (int i = 0; i < fmattn_q_per_lane; i++)\n"
|
|
" qreg[i] = head[lane + uint(i) * 32u];\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* One split of the decode key range for one kv head.\n"
|
|
" *\n"
|
|
" * Writes the split's UNNORMALIZED accumulator plus its (m, l) so\n"
|
|
" * fmattn_flash_decode_reduce can put every split on a common scale. The\n"
|
|
" * partial index is (kvh*G + g)*n_splits + split — head-major, so the\n"
|
|
" * reduce reads one head's splits contiguously.\n"
|
|
" *\n"
|
|
" * A split whose slice starts past seq_len writes m = -inf, l = 0 and a\n"
|
|
" * zero accumulator; the reduce weights that by exp(-inf - M) = 0. */\n"
|
|
"kernel void fmattn_flash_decode(\n"
|
|
" device const float* q [[buffer(0)]],\n"
|
|
" device const float* k_cache [[buffer(1)]],\n"
|
|
" device const float* v_cache [[buffer(2)]],\n"
|
|
" device float* m_out [[buffer(3)]],\n"
|
|
" device float* l_out [[buffer(4)]],\n"
|
|
" device float* acc_out [[buffer(5)]],\n"
|
|
" constant uint& max_seq [[buffer(6)]],\n"
|
|
" constant uint& seq_len [[buffer(7)]],\n"
|
|
" constant uint& n_splits [[buffer(8)]],\n"
|
|
" constant uint& split_len [[buffer(9)]],\n"
|
|
" constant uint& scale_bits [[buffer(10)]],\n"
|
|
" uint tgid [[threadgroup_position_in_grid]],\n"
|
|
" uint tid [[thread_position_in_threadgroup]],\n"
|
|
" uint sgid [[simdgroup_index_in_threadgroup]],\n"
|
|
" uint lane [[thread_index_in_simdgroup]]\n"
|
|
") {\n"
|
|
" threadgroup float k_tile[16 * 256];\n"
|
|
" threadgroup float score_s[8 * 16];\n"
|
|
" threadgroup float p_s[16 * 8];\n"
|
|
" threadgroup float m_s[8];\n"
|
|
" threadgroup float l_s[8];\n"
|
|
" threadgroup float corr_s[8];\n"
|
|
"\n"
|
|
" const uint HD = uint(fmattn_head_dim);\n"
|
|
" const uint G = uint(fmattn_kv_groups);\n"
|
|
" const uint kvh = tgid / n_splits;\n"
|
|
" const uint split = tgid % n_splits;\n"
|
|
"\n"
|
|
" uint k_lo = split * split_len;\n"
|
|
" uint k_hi = k_lo + split_len;\n"
|
|
" if (k_lo > seq_len) k_lo = seq_len;\n"
|
|
" if (k_hi > seq_len) k_hi = seq_len;\n"
|
|
"\n"
|
|
" float qreg[8];\n"
|
|
" float acc[8];\n"
|
|
" fmattn_load_query(q + (ulong)kvh * G * HD, qreg, sgid, lane);\n"
|
|
" for (int g = 0; g < fmattn_kv_groups; g++) acc[g] = 0.0f;\n"
|
|
" if (tid < G) { m_s[tid] = fmattn_score_floor; l_s[tid] = 0.0f; }\n"
|
|
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
|
"\n"
|
|
" fmattn_accumulate_range(k_cache + (ulong)kvh * max_seq * HD,\n"
|
|
" v_cache + (ulong)kvh * max_seq * HD,\n"
|
|
" qreg, acc, k_tile, score_s, p_s, m_s, l_s,\n"
|
|
" corr_s, k_lo, k_hi, as_type<float>(scale_bits),\n"
|
|
" tid, sgid, lane);\n"
|
|
"\n"
|
|
" if (tid < HD) {\n"
|
|
" for (int g = 0; g < fmattn_kv_groups; g++) {\n"
|
|
" ulong po = (ulong)(kvh * G + uint(g)) * n_splits + split;\n"
|
|
" acc_out[po * HD + tid] = acc[g];\n"
|
|
" }\n"
|
|
" }\n"
|
|
" if (tid < G) {\n"
|
|
" ulong po = (ulong)(kvh * G + tid) * n_splits + split;\n"
|
|
" m_out[po] = m_s[tid];\n"
|
|
" l_out[po] = l_s[tid];\n"
|
|
" }\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* Combine the per-split partials of one query head.\n"
|
|
" *\n"
|
|
" * Take M = max of the split maxima, weight each split by exp(m_s - M),\n"
|
|
" * sum the weighted l's and accumulators, divide. Identical to a\n"
|
|
" * single-pass softmax over the union of the slices. One workgroup per\n"
|
|
" * head, thread tid owns output dim tid. */\n"
|
|
"kernel void fmattn_flash_decode_reduce(\n"
|
|
" device const float* m_in [[buffer(0)]],\n"
|
|
" device const float* l_in [[buffer(1)]],\n"
|
|
" device const float* acc_in [[buffer(2)]],\n"
|
|
" device float* out [[buffer(3)]],\n"
|
|
" constant uint& head_dim [[buffer(4)]],\n"
|
|
" constant uint& n_splits [[buffer(5)]],\n"
|
|
" uint head [[threadgroup_position_in_grid]],\n"
|
|
" uint tid [[thread_position_in_threadgroup]]\n"
|
|
") {\n"
|
|
" if (tid >= head_dim) return;\n"
|
|
" device const float* m_h = m_in + (ulong)head * n_splits;\n"
|
|
" device const float* l_h = l_in + (ulong)head * n_splits;\n"
|
|
" device const float* a_h = acc_in + (ulong)head * n_splits * head_dim;\n"
|
|
"\n"
|
|
" float peak = fmattn_score_floor;\n"
|
|
" for (uint s = 0; s < n_splits; s++) peak = max(peak, m_h[s]);\n"
|
|
"\n"
|
|
" float o = 0.0f;\n"
|
|
" float total = 0.0f;\n"
|
|
" for (uint s = 0; s < n_splits; s++) {\n"
|
|
" float w = exp(m_h[s] - peak);\n"
|
|
" total += l_h[s] * w;\n"
|
|
" o += a_h[(ulong)s * head_dim + tid] * w;\n"
|
|
" }\n"
|
|
" out[(ulong)head * head_dim + tid] = o / total;\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* Causal GQA attention for one query row against the cache.\n"
|
|
" *\n"
|
|
" * Workgroup (t, kvh): query row t sits at absolute position pos0 + t and\n"
|
|
" * attends cache rows [0, pos0 + t]. Unlike decode this normalizes in\n"
|
|
" * place — there is no split to combine. */\n"
|
|
"kernel void fmattn_flash_prefill(\n"
|
|
" device const float* q [[buffer(0)]],\n"
|
|
" device const float* k_cache [[buffer(1)]],\n"
|
|
" device const float* v_cache [[buffer(2)]],\n"
|
|
" device float* out [[buffer(3)]],\n"
|
|
" constant uint& n_head_kv [[buffer(4)]],\n"
|
|
" constant uint& max_seq [[buffer(5)]],\n"
|
|
" constant uint& pos0 [[buffer(6)]],\n"
|
|
" constant uint& scale_bits [[buffer(7)]],\n"
|
|
" uint tgid [[threadgroup_position_in_grid]],\n"
|
|
" uint tid [[thread_position_in_threadgroup]],\n"
|
|
" uint sgid [[simdgroup_index_in_threadgroup]],\n"
|
|
" uint lane [[thread_index_in_simdgroup]]\n"
|
|
") {\n"
|
|
" threadgroup float k_tile[16 * 256];\n"
|
|
" threadgroup float score_s[8 * 16];\n"
|
|
" threadgroup float p_s[16 * 8];\n"
|
|
" threadgroup float m_s[8];\n"
|
|
" threadgroup float l_s[8];\n"
|
|
" threadgroup float corr_s[8];\n"
|
|
"\n"
|
|
" const uint HD = uint(fmattn_head_dim);\n"
|
|
" const uint G = uint(fmattn_kv_groups);\n"
|
|
" const uint t = tgid / n_head_kv;\n"
|
|
" const uint kvh = tgid % n_head_kv;\n"
|
|
" const uint n_head = n_head_kv * G;\n"
|
|
" const ulong row_base = ((ulong)t * n_head + (ulong)kvh * G) * HD;\n"
|
|
"\n"
|
|
" float qreg[8];\n"
|
|
" float acc[8];\n"
|
|
" fmattn_load_query(q + row_base, qreg, sgid, lane);\n"
|
|
" for (int g = 0; g < fmattn_kv_groups; g++) acc[g] = 0.0f;\n"
|
|
" if (tid < G) { m_s[tid] = fmattn_score_floor; l_s[tid] = 0.0f; }\n"
|
|
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
|
|
"\n"
|
|
" fmattn_accumulate_range(k_cache + (ulong)kvh * max_seq * HD,\n"
|
|
" v_cache + (ulong)kvh * max_seq * HD,\n"
|
|
" qreg, acc, k_tile, score_s, p_s, m_s, l_s,\n"
|
|
" corr_s, 0u, pos0 + t + 1u,\n"
|
|
" as_type<float>(scale_bits), tid, sgid, lane);\n"
|
|
"\n"
|
|
" if (tid < HD) {\n"
|
|
" device float* o = out + row_base;\n"
|
|
" for (int g = 0; g < fmattn_kv_groups; g++)\n"
|
|
" o[(ulong)uint(g) * HD + tid] = acc[g] / l_s[g];\n"
|
|
" }\n"
|
|
"}\n"
|
|
"\n"
|
|
;
|
|
}
|