wehub-resource-sync
338 lines
12 KiB
C++
338 lines
12 KiB
C++
// cider_sdpa_vector.h — Cider v9 optimized SDPA vector kernels
|
|
//
|
|
// Improvements over stock MLX:
|
|
// 1. Contiguous chunk layout (cache-friendly sequential access)
|
|
// 2. FlashInfer-style register tiling (TILE=4 unroll)
|
|
// 3. Adaptive blocks selection per GQA ratio
|
|
//
|
|
// Template parameters:
|
|
// T — data type (float, half, bfloat)
|
|
// D — head dimension (64, 96, 128, 256)
|
|
// BLOCKS — number of blocks for 2-pass (32, 64, 128)
|
|
//
|
|
// Three kernels:
|
|
// cider_sdpa_vector<T, D> — 1-pass (short N)
|
|
// cider_sdpa_vector_2pass_1<T, D, BLOCKS> — 2-pass pass1 (per-block partials)
|
|
// cider_sdpa_vector_2pass_2<T, D, BLOCKS> — 2-pass pass2 (reduce)
|
|
|
|
#pragma once
|
|
|
|
#include <metal_simdgroup>
|
|
|
|
using namespace metal;
|
|
|
|
// ═══════════════════════════════════════════════════════════════
|
|
// 1-pass kernel (N <= threshold, no intermediate buffer)
|
|
// ═══════════════════════════════════════════════════════════════
|
|
template <typename T, int D>
|
|
[[kernel]] void cider_sdpa_vector(
|
|
const device T* queries [[buffer(0)]],
|
|
const device T* keys [[buffer(1)]],
|
|
const device T* values [[buffer(2)]],
|
|
device T* out [[buffer(3)]],
|
|
const constant int& gqa_factor [[buffer(4)]],
|
|
const constant int& N [[buffer(5)]],
|
|
const constant size_t& k_head_stride [[buffer(6)]],
|
|
const constant size_t& k_seq_stride [[buffer(7)]],
|
|
const constant size_t& v_head_stride [[buffer(8)]],
|
|
const constant size_t& v_seq_stride [[buffer(9)]],
|
|
const constant float& scale [[buffer(10)]],
|
|
uint3 tid [[threadgroup_position_in_grid]],
|
|
uint simd_gid [[simdgroup_index_in_threadgroup]],
|
|
uint simd_lid [[thread_index_in_simdgroup]]) {
|
|
|
|
constexpr int BN = 32;
|
|
constexpr int BD = 32;
|
|
constexpr int qk_per_thread = D / BD;
|
|
constexpr int v_per_thread = D / BD;
|
|
|
|
typedef float U;
|
|
|
|
const int inner_k_stride = BN * int(k_seq_stride);
|
|
const int inner_v_stride = BN * int(v_seq_stride);
|
|
|
|
thread U q_reg[qk_per_thread];
|
|
thread U k_reg[qk_per_thread];
|
|
thread U o_reg[v_per_thread];
|
|
|
|
threadgroup U tg_outputs[BN * BD];
|
|
threadgroup U tg_max_scores[BN];
|
|
threadgroup U tg_sum_exp_scores[BN];
|
|
|
|
const int q_batch_head_idx = tid.x;
|
|
const int kv_head_idx = q_batch_head_idx / gqa_factor;
|
|
const int o_offset = q_batch_head_idx;
|
|
|
|
const device T* q_ptr = queries + o_offset * D + simd_lid * qk_per_thread;
|
|
const device T* k_ptr = keys + kv_head_idx * int(k_head_stride)
|
|
+ simd_gid * int(k_seq_stride) + simd_lid * qk_per_thread;
|
|
const device T* v_ptr = values + kv_head_idx * int(v_head_stride)
|
|
+ simd_gid * int(v_seq_stride) + simd_lid * v_per_thread;
|
|
|
|
for (int i = 0; i < qk_per_thread; i++) {
|
|
q_reg[i] = static_cast<U>(scale) * static_cast<U>(q_ptr[i]);
|
|
}
|
|
for (int i = 0; i < v_per_thread; i++) {
|
|
o_reg[i] = 0;
|
|
}
|
|
|
|
U max_score = -1e38f;
|
|
U sum_exp_score = 0;
|
|
|
|
for (int i = simd_gid; i < N; i += BN) {
|
|
for (int j = 0; j < qk_per_thread; j++) {
|
|
k_reg[j] = static_cast<U>(k_ptr[j]);
|
|
}
|
|
U score = 0;
|
|
for (int j = 0; j < qk_per_thread; j++) {
|
|
score += q_reg[j] * k_reg[j];
|
|
}
|
|
score = simd_sum(score);
|
|
|
|
U new_max = max(max_score, score);
|
|
U factor = fast::exp(max_score - new_max);
|
|
U exp_score = fast::exp(score - new_max);
|
|
max_score = new_max;
|
|
sum_exp_score = sum_exp_score * factor + exp_score;
|
|
|
|
for (int j = 0; j < v_per_thread; j++) {
|
|
o_reg[j] = o_reg[j] * factor + exp_score * static_cast<U>(v_ptr[j]);
|
|
}
|
|
|
|
k_ptr += inner_k_stride;
|
|
v_ptr += inner_v_stride;
|
|
}
|
|
|
|
if (simd_lid == 0) {
|
|
tg_max_scores[simd_gid] = max_score;
|
|
tg_sum_exp_scores[simd_gid] = sum_exp_score;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
max_score = tg_max_scores[simd_lid];
|
|
U new_max_final = simd_max(max_score);
|
|
U factor = fast::exp(max_score - new_max_final);
|
|
sum_exp_score = simd_sum(tg_sum_exp_scores[simd_lid] * factor);
|
|
|
|
for (int i = 0; i < v_per_thread; i++) {
|
|
tg_outputs[simd_lid * BD + simd_gid] = o_reg[i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
o_reg[i] = simd_sum(tg_outputs[simd_gid * BD + simd_lid] * factor);
|
|
o_reg[i] = sum_exp_score == 0 ? o_reg[i] : (o_reg[i] / sum_exp_score);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
if (simd_lid == 0) {
|
|
device T* o_ptr = out + o_offset * D + simd_gid * v_per_thread;
|
|
for (int i = 0; i < v_per_thread; i++) {
|
|
o_ptr[i] = static_cast<T>(o_reg[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// ═══════════════════════════════════════════════════════════════
|
|
// 2-pass pass1: per-block partial results
|
|
// v9: contiguous chunks + TILE=4 register tiling
|
|
// ═══════════════════════════════════════════════════════════════
|
|
template <typename T, int D, int BLOCKS>
|
|
[[kernel]] void cider_sdpa_vector_2pass_1(
|
|
const device T* queries [[buffer(0)]],
|
|
const device T* keys [[buffer(1)]],
|
|
const device T* values [[buffer(2)]],
|
|
device T* partials [[buffer(3)]],
|
|
device float* sums [[buffer(4)]],
|
|
device float* maxs [[buffer(5)]],
|
|
const constant int& N [[buffer(7)]],
|
|
const constant size_t& k_head_stride [[buffer(8)]],
|
|
const constant size_t& k_seq_stride [[buffer(9)]],
|
|
const constant size_t& v_head_stride [[buffer(10)]],
|
|
const constant size_t& v_seq_stride [[buffer(11)]],
|
|
const constant float& scale [[buffer(12)]],
|
|
uint3 tptg [[threads_per_threadgroup]],
|
|
uint3 tidtg [[thread_position_in_threadgroup]],
|
|
uint3 tid [[threadgroup_position_in_grid]],
|
|
uint3 tpg [[threadgroups_per_grid]],
|
|
uint simd_lid [[thread_index_in_simdgroup]]) {
|
|
|
|
constexpr int BD = 32;
|
|
constexpr int qk_per_thread = D / BD;
|
|
constexpr int v_per_thread = D / BD;
|
|
constexpr int TILE = 4;
|
|
|
|
typedef float U;
|
|
|
|
thread U q_reg[qk_per_thread];
|
|
thread U o_reg[v_per_thread] = {0};
|
|
|
|
const int kv_head_idx = tid.x;
|
|
const int batch_idx = tid.y;
|
|
const int block_idx = tid.z;
|
|
const int gqa_factor = tptg.y;
|
|
const int q_head_idx = gqa_factor * kv_head_idx + tidtg.y;
|
|
const int num_kv_heads = tpg.x;
|
|
const int num_q_heads = num_kv_heads * gqa_factor;
|
|
const int q_batch_head_idx = batch_idx * num_q_heads + q_head_idx;
|
|
const int o_offset = q_batch_head_idx;
|
|
|
|
queries += o_offset * D + simd_lid * qk_per_thread;
|
|
|
|
const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx;
|
|
|
|
// v9: contiguous chunk layout (not interleaved)
|
|
const int chunk_size = (N + BLOCKS - 1) / BLOCKS;
|
|
const int kv_start = block_idx * chunk_size;
|
|
const int kv_end = min(kv_start + chunk_size, N);
|
|
|
|
const device T* k_ptr = keys + kv_batch_head_idx * int(k_head_stride)
|
|
+ kv_start * int(k_seq_stride) + simd_lid * qk_per_thread;
|
|
const device T* v_ptr = values + kv_batch_head_idx * int(v_head_stride)
|
|
+ kv_start * int(v_seq_stride) + simd_lid * v_per_thread;
|
|
|
|
device T* o_ptr = partials + o_offset * BLOCKS * D + block_idx * D
|
|
+ simd_lid * v_per_thread;
|
|
|
|
// Read query
|
|
for (int i = 0; i < qk_per_thread; i++) {
|
|
q_reg[i] = static_cast<U>(scale) * static_cast<U>(queries[i]);
|
|
}
|
|
|
|
U max_score = -1e38f;
|
|
U sum_exp_score = 0;
|
|
|
|
const int kss = int(k_seq_stride);
|
|
const int vss = int(v_seq_stride);
|
|
|
|
// Main loop with TILE=4 unrolling
|
|
int pos = kv_start;
|
|
const int tiled_end = kv_start + ((kv_end - kv_start) / TILE) * TILE;
|
|
|
|
for (; pos < tiled_end; pos += TILE) {
|
|
U scores[TILE];
|
|
for (int t = 0; t < TILE; t++) {
|
|
U score = 0;
|
|
const device T* kt = k_ptr + t * kss;
|
|
for (int j = 0; j < qk_per_thread; j++) {
|
|
score += q_reg[j] * static_cast<U>(kt[j]);
|
|
}
|
|
scores[t] = simd_sum(score);
|
|
}
|
|
for (int t = 0; t < TILE; t++) {
|
|
U new_max = max(max_score, scores[t]);
|
|
U factor = fast::exp(max_score - new_max);
|
|
U exp_score = fast::exp(scores[t] - new_max);
|
|
max_score = new_max;
|
|
sum_exp_score = sum_exp_score * factor + exp_score;
|
|
|
|
const device T* vt = v_ptr + t * vss;
|
|
for (int j = 0; j < v_per_thread; j++) {
|
|
o_reg[j] = o_reg[j] * factor + exp_score * static_cast<U>(vt[j]);
|
|
}
|
|
}
|
|
k_ptr += TILE * kss;
|
|
v_ptr += TILE * vss;
|
|
}
|
|
// Remainder
|
|
for (; pos < kv_end; pos++) {
|
|
U score = 0;
|
|
for (int j = 0; j < qk_per_thread; j++) {
|
|
score += q_reg[j] * static_cast<U>(k_ptr[j]);
|
|
}
|
|
score = simd_sum(score);
|
|
|
|
U new_max = max(max_score, score);
|
|
U factor = fast::exp(max_score - new_max);
|
|
U exp_score = fast::exp(score - new_max);
|
|
max_score = new_max;
|
|
sum_exp_score = sum_exp_score * factor + exp_score;
|
|
|
|
for (int j = 0; j < v_per_thread; j++) {
|
|
o_reg[j] = o_reg[j] * factor + exp_score * static_cast<U>(v_ptr[j]);
|
|
}
|
|
k_ptr += kss;
|
|
v_ptr += vss;
|
|
}
|
|
|
|
// Write partial results
|
|
if (simd_lid == 0) {
|
|
sums[o_offset * BLOCKS + block_idx] = sum_exp_score;
|
|
maxs[o_offset * BLOCKS + block_idx] = max_score;
|
|
}
|
|
for (int i = 0; i < v_per_thread; i++) {
|
|
o_ptr[i] = static_cast<T>(o_reg[i]);
|
|
}
|
|
}
|
|
|
|
|
|
// ═══════════════════════════════════════════════════════════════
|
|
// 2-pass pass2: reduce partial results across blocks
|
|
// ═══════════════════════════════════════════════════════════════
|
|
template <typename T, int D, int BLOCKS>
|
|
[[kernel]] void cider_sdpa_vector_2pass_2(
|
|
const device T* partials [[buffer(0)]],
|
|
const device float* sums [[buffer(1)]],
|
|
const device float* maxs [[buffer(2)]],
|
|
device T* out [[buffer(3)]],
|
|
uint3 tid [[threadgroup_position_in_grid]],
|
|
uint simd_gid [[simdgroup_index_in_threadgroup]],
|
|
uint simd_lid [[thread_index_in_simdgroup]]) {
|
|
|
|
constexpr int BN = 32;
|
|
constexpr int BD = 32;
|
|
constexpr int elem_per_thread = D / BD;
|
|
|
|
typedef float U;
|
|
|
|
thread U o_reg[elem_per_thread] = {0};
|
|
threadgroup U tg_outputs[BN * BD];
|
|
|
|
const int head_idx = tid.x;
|
|
|
|
const device T* p_ptr = partials + head_idx * BLOCKS * D
|
|
+ simd_gid * D + simd_lid * elem_per_thread;
|
|
const device float* s_ptr = sums + head_idx * BLOCKS;
|
|
const device float* m_ptr = maxs + head_idx * BLOCKS;
|
|
|
|
// Reduce max
|
|
U max_score = -1e38f;
|
|
for (int b = 0; b < BLOCKS / BN; ++b) {
|
|
max_score = max(max_score, m_ptr[simd_lid + BN * b]);
|
|
}
|
|
max_score = simd_max(max_score);
|
|
|
|
// Reduce sum_exp
|
|
U sum_exp_score = 0;
|
|
for (int b = 0; b < BLOCKS / BN; ++b) {
|
|
U factor = fast::exp(m_ptr[simd_lid + BN * b] - max_score);
|
|
sum_exp_score += factor * s_ptr[simd_lid + BN * b];
|
|
}
|
|
sum_exp_score = simd_sum(sum_exp_score);
|
|
|
|
// Reduce partials
|
|
const device float* m_walk = m_ptr;
|
|
for (int b = 0; b < BLOCKS / BN; ++b) {
|
|
U factor = fast::exp(m_walk[simd_gid] - max_score);
|
|
for (int i = 0; i < elem_per_thread; i++) {
|
|
o_reg[i] += factor * static_cast<U>(p_ptr[i]);
|
|
}
|
|
m_walk += BN;
|
|
p_ptr += BN * D;
|
|
}
|
|
|
|
// Transpose + reduce via shared memory
|
|
for (int i = 0; i < elem_per_thread; i++) {
|
|
tg_outputs[simd_lid * BD + simd_gid] = o_reg[i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
o_reg[i] = simd_sum(tg_outputs[simd_gid * BD + simd_lid]);
|
|
o_reg[i] = sum_exp_score == 0 ? o_reg[i] : (o_reg[i] / sum_exp_score);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
if (simd_lid == 0) {
|
|
device T* o_ptr = out + head_idx * D + simd_gid * elem_per_thread;
|
|
for (int i = 0; i < elem_per_thread; i++) {
|
|
o_ptr[i] = static_cast<T>(o_reg[i]);
|
|
}
|
|
}
|
|
}
|