mega_fused_attention_avx.c File Reference

Mega-Fused Attention for AVX (256-bit) and AVX-512 (512-bit) More...

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "ckernel_engine.h"
#include "ckernel_quant.h"
#include "ck_features.h"

Go to the source code of this file.

Macros
#define	MEGA_KV_TILE   32
#define	MEGA_Q_TILE   32
#define	MEGA_REGS   16 /* 16 YMM registers */
#define	MEGA_STACK_MAX   8192
#define	MEGA_VLEN   8 /* 256 / 32 */
#define	REG_K_TILE   "YMM8-YMM11" /* 4 regs for K tile */
#define	REG_O_ACCUM   "Stack+L1" /* O in L1 cache */
#define	REG_Q_ACCUM   "YMM0-YMM7" /* 8 regs for Q tile */
#define	REG_SOFTMAX   "YMM0-YMM1" /* 2 regs for m, l */
#define	REG_TEMP   "YMM2-YMM3" /* 2 regs for temps */
#define	REG_V_TILE   "YMM12-YMM15" /* 4 regs for V tile */

Functions
static float	ck_dot_f32 (const float *a, const float *b, int len)
void	mega_fuse_flash_attention_avx (float *o_out, const float *q, const float *kv_cache_k, const float *kv_cache_v, int num_heads, int num_kv_heads, int seq_len, int cache_capacity, int head_dim, int aligned_head_dim)
	Flash attention with online softmax (AVX version)
static void	mega_fuse_output_proj_residual (const float *attn_token, const float *wo, const float *bo, const float *residual, float *output, int embed_dim, int aligned_embed_dim, int num_heads, int head_dim, int aligned_head_dim)
void	mega_fuse_rmsnorm_qkv_avx (float *q_out, float *k_out, float *v_out, const float *input, const float *gamma, const float *wq, const float *bq, const float *wk, const float *bk, const float *wv, const float *bv, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, float eps)
	Fused RMSNorm + QKV for decode (single token)
void	mega_fuse_rope_inplace_avx (float *q, float *k, const float *rope_cos, const float *rope_sin, int pos, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim)
	Apply RoPE to Q and K (in-place, from L1)
void	mega_fused_attention_decode (float *output, const float *input, const float *residual, const float *ln1_gamma, const float *wq, const float *bq, const float *wk, const float *bk, const float *wv, const float *bv, const float *wo, const float *bo, float *kv_cache_k, float *kv_cache_v, const float *rope_cos, const float *rope_sin, int pos, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int cache_capacity, float eps)
	Mega-fused attention for decode mode (single token)
void	mega_fused_attention_decode_workspace (float *output, const float *input, const float *residual, const float *ln1_gamma, const float *wq, const float *bq, const float *wk, const float *bk, const float *wv, const float *bv, const float *wo, const float *bo, float *kv_cache_k, float *kv_cache_v, const float *rope_cos, const float *rope_sin, int pos, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int cache_capacity, float eps, float *q_output_workspace, size_t q_output_workspace_bytes, float *kv_workspace, size_t kv_workspace_bytes)
	Full mega-fused attention for decode.

Detailed Description

Mega-Fused Attention for AVX (256-bit) and AVX-512 (512-bit)

CK-ENGINE KERNEL RULES:

1. NO malloc/free - memory via bump allocator, pointers passed in
2. NO OpenMP - parallelization at orchestrator/codegen layer
3. NO memcpy for layout - use strided access, not copies
4. API must define: inputs, outputs, workspace, and memory layouts
5. Pure computation - deterministic, no side effects

After changes: make test && make llamacpp-parity-full

The mapped decode provider receives planner-owned workspace. The legacy public wrapper retains one compatibility allocation for direct callers.

Holy grail fusion: RMSNorm → QKV → RoPE → Flash Attention → OutProj + Residual

AVX approach: Keep intermediates in L1 cache (not registers) AVX-512 approach: Keep intermediates in registers

Both achieve the same goal: Eliminate DRAM traffic for intermediates.

Definition in file mega_fused_attention_avx.c.

Macro Definition Documentation

MEGA_KV_TILE

#define MEGA_KV_TILE   32

Definition at line 118 of file mega_fused_attention_avx.c.

MEGA_Q_TILE

#define MEGA_Q_TILE   32

Definition at line 117 of file mega_fused_attention_avx.c.

MEGA_REGS

#define MEGA_REGS   16 /* 16 YMM registers */

Definition at line 116 of file mega_fused_attention_avx.c.

MEGA_STACK_MAX

#define MEGA_STACK_MAX   8192

Definition at line 119 of file mega_fused_attention_avx.c.

MEGA_VLEN

#define MEGA_VLEN   8 /* 256 / 32 */

Definition at line 115 of file mega_fused_attention_avx.c.

REG_K_TILE

#define REG_K_TILE   "YMM8-YMM11" /* 4 regs for K tile */

Definition at line 123 of file mega_fused_attention_avx.c.

REG_O_ACCUM

#define REG_O_ACCUM   "Stack+L1" /* O in L1 cache */

Definition at line 125 of file mega_fused_attention_avx.c.

REG_Q_ACCUM

#define REG_Q_ACCUM   "YMM0-YMM7" /* 8 regs for Q tile */

Definition at line 122 of file mega_fused_attention_avx.c.

REG_SOFTMAX

#define REG_SOFTMAX   "YMM0-YMM1" /* 2 regs for m, l */

Definition at line 126 of file mega_fused_attention_avx.c.

REG_TEMP

#define REG_TEMP   "YMM2-YMM3" /* 2 regs for temps */

Definition at line 127 of file mega_fused_attention_avx.c.

REG_V_TILE

#define REG_V_TILE   "YMM12-YMM15" /* 4 regs for V tile */

Definition at line 124 of file mega_fused_attention_avx.c.

Function Documentation

ck_dot_f32()

static float ck_dot_f32 ( const float *  a, const float *  b, int  len  )

inlinestatic

Definition at line 54 of file mega_fused_attention_avx.c.

{
#if defined(__AVX512F__)
    __m512 acc = _mm512_setzero_ps();
    int i = 0;
    for (; i <= len - 16; i += 16) {
        __m512 va = _mm512_loadu_ps(a + i);
        __m512 vb = _mm512_loadu_ps(b + i);
        acc = _mm512_fmadd_ps(va, vb, acc);
    }
    float sum = _mm512_reduce_add_ps(acc);
    for (; i < len; ++i) {
        sum += a[i] * b[i];
    }
    return sum;
#elif defined(__AVX__)
    __m256 acc = _mm256_setzero_ps();
    int i = 0;
    for (; i <= len - 8; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        acc = _mm256_add_ps(acc, _mm256_mul_ps(va, vb));
    }
    float sum = ck_hsum256_ps(acc);
    for (; i < len; ++i) {
        sum += a[i] * b[i];
    }
    return sum;
#else
    float sum = 0.0f;
    for (int i = 0; i < len; ++i) {
        sum += a[i] * b[i];
    }
    return sum;
#endif
}

Referenced by mega_fuse_output_proj_residual(), and mega_fuse_rmsnorm_qkv_avx().

mega_fuse_flash_attention_avx()

void mega_fuse_flash_attention_avx	(	float *	o_out,
		const float *	q,
		const float *	kv_cache_k,
		const float *	kv_cache_v,
		int	num_heads,
		int	num_kv_heads,
		int	seq_len,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim
	)

Flash attention with online softmax (AVX version)

Key insight: O, m, l stay in registers throughout! K/V tiles stream from L2 cache.

Parameters

o_out	Output [num_heads * aligned_head_dim] - in registers/L1
q	Q tensor [num_heads * aligned_head_dim] - from L1
kv_cache_k	KV cache K [num_kv_heads * cache_capacity * aligned_head_dim]
kv_cache_v	KV cache V [num_kv_heads * cache_capacity * aligned_head_dim]
num_heads	Number of heads
num_kv_heads	Number of KV heads
seq_len	Current sequence length
cache_capacity	KV cache capacity (head stride)
head_dim	Head dimension
aligned_head_dim	Aligned head dimension

Definition at line 444 of file mega_fused_attention_avx.c.

{
    const int hd = head_dim;
    const float scale = 1.0f / sqrtf((float)hd);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; h++) {
        const float *q_h = q + (size_t)h * (size_t)aligned_head_dim;
        const int kv_idx = h % num_kv_heads;
        const float *k_cache = kv_cache_k + (size_t)kv_idx * head_stride;
        const float *v_cache = kv_cache_v + (size_t)kv_idx * head_stride;
 
        /* O, m, l in registers for this head */
        float o_h[aligned_head_dim]; /* in L1 */
        float m = -INFINITY; /* running max */
        float l = 0.0f;      /* running sum */
 
        /* Initialize O to zeros */
        memset(o_h, 0, (size_t)aligned_head_dim * sizeof(float));
 
        /* Iterate over KV cache tiles */
        for (int t = 0; t < seq_len; t += MEGA_KV_TILE) {
            int tile_end = t + MEGA_KV_TILE;
            if (tile_end > seq_len) tile_end = seq_len;
            int tile_size = tile_end - t;
 
            /* Load K tile from L2 cache */
            float k_tile[MEGA_KV_TILE * hd];
            for (int i = 0; i < tile_size; i++) {
                memcpy(k_tile + (size_t)i * (size_t)hd,
                       k_cache + (size_t)(t + i) * (size_t)aligned_head_dim,
                       (size_t)hd * sizeof(float));
            }
 
            /* S_ij = Q @ K_tile.T / sqrt(d) - in registers */
            float s_row[MEGA_KV_TILE];
            for (int j = 0; j < tile_size; j++) {
                s_row[j] = 0.0f;
                for (int i = 0; i < hd; i++) {
                    s_row[j] += q_h[i] * k_tile[j * hd + i];
                }
                s_row[j] *= scale;
            }
 
            /* Online softmax update */
            float m_new = m;
            for (int j = 0; j < tile_size; j++) {
                if (s_row[j] > m_new) m_new = s_row[j];
            }
 
            float l_new = 0.0f;
            for (int j = 0; j < tile_size; j++) {
                float p = expf(s_row[j] - m_new);
                s_row[j] = p;
                l_new += p;
            }
 
            /* Scale O by exp(m - m_new) and add P @ V */
            float exp_m_diff = expf(m - m_new);
            for (int i = 0; i < hd; i++) {
                o_h[i] *= exp_m_diff;
            }
 
            /* Load V tile and accumulate */
            for (int j = 0; j < tile_size; j++) {
                float p = s_row[j];
                for (int i = 0; i < hd; i++) {
                    o_h[i] += p * v_cache[(size_t)(t + j) * (size_t)aligned_head_dim + (size_t)i];
                }
            }
 
            l = l * exp_m_diff + l_new;
            m = m_new;
        }
 
        /* Normalize by l */
        for (int i = 0; i < hd; i++) {
            o_h[i] /= l;
        }
        for (int i = hd; i < aligned_head_dim; ++i) {
            o_h[i] = 0.0f;
        }
 
        /* Store O - still in L1, goes to output projection */
        memcpy(o_out + (size_t)h * (size_t)aligned_head_dim,
               o_h,
               (size_t)aligned_head_dim * sizeof(float));
    }
}

References MEGA_KV_TILE.

Referenced by mega_fused_attention_decode_workspace().

mega_fuse_output_proj_residual()

static void mega_fuse_output_proj_residual ( const float *  attn_token, const float *  wo, const float *  bo, const float *  residual, float *  output, int  embed_dim, int  aligned_embed_dim, int  num_heads, int  head_dim, int  aligned_head_dim  )

static

Definition at line 551 of file mega_fused_attention_avx.c.

{
    if (!attn_token || !wo || !output) {
        return;
    }
 
    const size_t head_weight_stride = (size_t)aligned_embed_dim * (size_t)aligned_head_dim;
 
    for (int j = 0; j < embed_dim; ++j) {
        float sum = bo ? bo[j] : 0.0f;
        for (int h = 0; h < num_heads; ++h) {
            const float *o_h = attn_token + (size_t)h * (size_t)aligned_head_dim;
            const float *wo_row = wo + (size_t)h * head_weight_stride + (size_t)j * (size_t)aligned_head_dim;
            sum += ck_dot_f32(o_h, wo_row, head_dim);
        }
        output[j] = sum + (residual ? residual[j] : 0.0f);
    }
 
    for (int j = embed_dim; j < aligned_embed_dim; ++j) {
        output[j] = 0.0f;
    }
}

References ck_dot_f32().

Referenced by mega_fused_attention_decode_workspace().

mega_fuse_rmsnorm_qkv_avx()

void mega_fuse_rmsnorm_qkv_avx	(	float *	q_out,
		float *	k_out,
		float *	v_out,
		const float *	input,
		const float *	gamma,
		const float *	wq,
		const float *	bq,
		const float *	wk,
		const float *	bk,
		const float *	wv,
		const float *	bv,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		float	eps
	)

Fused RMSNorm + QKV for decode (single token)

Intermediates stay in L1/L2. Output buffers are head-major.

Definition at line 143 of file mega_fused_attention_avx.c.

{
    if (!q_out || !k_out || !v_out || !input || !wq || !wk || !wv) {
        return;
    }
    if (embed_dim <= 0 || aligned_embed_dim <= 0 || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    float ln1_row[aligned_embed_dim];
    float sum_sq = 0.0f;
 
#if defined(__AVX512F__)
    __m512 sum_vec = _mm512_setzero_ps();
    int i = 0;
    for (; i + 16 <= embed_dim; i += 16) {
        __m512 xv = _mm512_loadu_ps(input + i);
        sum_vec = _mm512_fmadd_ps(xv, xv, sum_vec);
    }
    sum_sq = _mm512_reduce_add_ps(sum_vec);
    for (; i < embed_dim; ++i) {
        sum_sq += input[i] * input[i];
    }
#elif defined(__AVX__)
    __m256 sum_vec = _mm256_setzero_ps();
    int i = 0;
    for (; i + 8 <= embed_dim; i += 8) {
        __m256 xv = _mm256_loadu_ps(input + i);
        sum_vec = _mm256_add_ps(sum_vec, _mm256_mul_ps(xv, xv));
    }
    sum_sq = ck_hsum256_ps(sum_vec);
    for (; i < embed_dim; ++i) {
        sum_sq += input[i] * input[i];
    }
#else
    for (int i = 0; i < embed_dim; ++i) {
        sum_sq += input[i] * input[i];
    }
#endif
 
    float rstd = 1.0f / sqrtf(sum_sq / (float)embed_dim + eps);
 
#if defined(__AVX512F__)
    if (gamma) {
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        int j = 0;
        for (; j + 16 <= embed_dim; j += 16) {
            __m512 xv = _mm512_loadu_ps(input + j);
            __m512 gv = _mm512_loadu_ps(gamma + j);
            __m512 yv = _mm512_mul_ps(_mm512_mul_ps(xv, rstd_vec), gv);
            _mm512_storeu_ps(ln1_row + j, yv);
        }
        for (; j < embed_dim; ++j) {
            ln1_row[j] = input[j] * rstd * gamma[j];
        }
    } else {
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        int j = 0;
        for (; j + 16 <= embed_dim; j += 16) {
            __m512 xv = _mm512_loadu_ps(input + j);
            __m512 yv = _mm512_mul_ps(xv, rstd_vec);
            _mm512_storeu_ps(ln1_row + j, yv);
        }
        for (; j < embed_dim; ++j) {
            ln1_row[j] = input[j] * rstd;
        }
    }
#elif defined(__AVX__)
    if (gamma) {
        __m256 rstd_vec = _mm256_set1_ps(rstd);
        int j = 0;
        for (; j + 8 <= embed_dim; j += 8) {
            __m256 xv = _mm256_loadu_ps(input + j);
            __m256 gv = _mm256_loadu_ps(gamma + j);
            __m256 yv = _mm256_mul_ps(_mm256_mul_ps(xv, rstd_vec), gv);
            _mm256_storeu_ps(ln1_row + j, yv);
        }
        for (; j < embed_dim; ++j) {
            ln1_row[j] = input[j] * rstd * gamma[j];
        }
    } else {
        __m256 rstd_vec = _mm256_set1_ps(rstd);
        int j = 0;
        for (; j + 8 <= embed_dim; j += 8) {
            __m256 xv = _mm256_loadu_ps(input + j);
            __m256 yv = _mm256_mul_ps(xv, rstd_vec);
            _mm256_storeu_ps(ln1_row + j, yv);
        }
        for (; j < embed_dim; ++j) {
            ln1_row[j] = input[j] * rstd;
        }
    }
#else
    for (int j = 0; j < embed_dim; ++j) {
        ln1_row[j] = input[j] * rstd * (gamma ? gamma[j] : 1.0f);
    }
#endif
 
    for (int j = embed_dim; j < aligned_embed_dim; ++j) {
        ln1_row[j] = 0.0f;
    }
 
    const size_t head_w_stride = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        const float *wq_h = wq + (size_t)h * head_w_stride;
        const float *bq_h = bq ? (bq + (size_t)h * (size_t)aligned_head_dim) : NULL;
        float *q_h = q_out + (size_t)h * (size_t)aligned_head_dim;
        for (int d = 0; d < head_dim; ++d) {
            const float *row = wq_h + (size_t)d * (size_t)aligned_embed_dim;
            float sum = ck_dot_f32(ln1_row, row, aligned_embed_dim);
            q_h[d] = sum + (bq_h ? bq_h[d] : 0.0f);
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) {
            q_h[d] = 0.0f;
        }
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *wk_h = wk + (size_t)h * head_w_stride;
        const float *wv_h = wv + (size_t)h * head_w_stride;
        const float *bk_h = bk ? (bk + (size_t)h * (size_t)aligned_head_dim) : NULL;
        const float *bv_h = bv ? (bv + (size_t)h * (size_t)aligned_head_dim) : NULL;
        float *k_h = k_out + (size_t)h * (size_t)aligned_head_dim;
        float *v_h = v_out + (size_t)h * (size_t)aligned_head_dim;
        for (int d = 0; d < head_dim; ++d) {
            const float *wk_row = wk_h + (size_t)d * (size_t)aligned_embed_dim;
            const float *wv_row = wv_h + (size_t)d * (size_t)aligned_embed_dim;
            float k_sum = ck_dot_f32(ln1_row, wk_row, aligned_embed_dim);
            float v_sum = ck_dot_f32(ln1_row, wv_row, aligned_embed_dim);
            k_h[d] = k_sum + (bk_h ? bk_h[d] : 0.0f);
            v_h[d] = v_sum + (bv_h ? bv_h[d] : 0.0f);
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) {
            k_h[d] = 0.0f;
            v_h[d] = 0.0f;
        }
    }
}

References ck_dot_f32().

Referenced by mega_fused_attention_decode_workspace().

mega_fuse_rope_inplace_avx()

void mega_fuse_rope_inplace_avx	(	float *	q,
		float *	k,
		const float *	rope_cos,
		const float *	rope_sin,
		int	pos,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim
	)

Apply RoPE to Q and K (in-place, from L1)

Q and K are already in L1 from QKV projection. Just apply rotation in-place.

Definition at line 311 of file mega_fused_attention_avx.c.

{
    if (!q || !k || !rope_cos || !rope_sin || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
    if ((head_dim & 1) != 0) {
        return;
    }
 
    int half = head_dim / 2;
    const float *cos_ptr = rope_cos + (size_t)pos * (size_t)half;
    const float *sin_ptr = rope_sin + (size_t)pos * (size_t)half;
 
    for (int h = 0; h < num_heads; ++h) {
        float *q_h = q + (size_t)h * (size_t)aligned_head_dim;
        int i = 0;
#if defined(__AVX512F__)
        for (; i + 16 <= half; i += 16) {
            __m512 q0 = _mm512_loadu_ps(q_h + i);
            __m512 q1 = _mm512_loadu_ps(q_h + i + half);
            __m512 cos = _mm512_loadu_ps(cos_ptr + i);
            __m512 sin = _mm512_loadu_ps(sin_ptr + i);
 
            __m512 q_rot0 = _mm512_sub_ps(_mm512_mul_ps(q0, cos), _mm512_mul_ps(q1, sin));
            __m512 q_rot1 = _mm512_add_ps(_mm512_mul_ps(q0, sin), _mm512_mul_ps(q1, cos));
 
            _mm512_storeu_ps(q_h + i, q_rot0);
            _mm512_storeu_ps(q_h + i + half, q_rot1);
        }
#elif defined(__AVX__)
        for (; i + 8 <= half; i += 8) {
            __m256 q0 = _mm256_loadu_ps(q_h + i);
            __m256 q1 = _mm256_loadu_ps(q_h + i + half);
            __m256 cos = _mm256_loadu_ps(cos_ptr + i);
            __m256 sin = _mm256_loadu_ps(sin_ptr + i);
 
            __m256 q_rot0 = _mm256_sub_ps(_mm256_mul_ps(q0, cos), _mm256_mul_ps(q1, sin));
            __m256 q_rot1 = _mm256_add_ps(_mm256_mul_ps(q0, sin), _mm256_mul_ps(q1, cos));
 
            _mm256_storeu_ps(q_h + i, q_rot0);
            _mm256_storeu_ps(q_h + i + half, q_rot1);
        }
#endif
        for (; i < half; ++i) {
            float q0 = q_h[i];
            float q1 = q_h[i + half];
            float c = cos_ptr[i];
            float s = sin_ptr[i];
            q_h[i] = q0 * c - q1 * s;
            q_h[i + half] = q0 * s + q1 * c;
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) {
            q_h[d] = 0.0f;
        }
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        float *k_h = k + (size_t)h * (size_t)aligned_head_dim;
        int i = 0;
#if defined(__AVX512F__)
        for (; i + 16 <= half; i += 16) {
            __m512 k0 = _mm512_loadu_ps(k_h + i);
            __m512 k1 = _mm512_loadu_ps(k_h + i + half);
            __m512 cos = _mm512_loadu_ps(cos_ptr + i);
            __m512 sin = _mm512_loadu_ps(sin_ptr + i);
 
            __m512 k_rot0 = _mm512_sub_ps(_mm512_mul_ps(k0, cos), _mm512_mul_ps(k1, sin));
            __m512 k_rot1 = _mm512_add_ps(_mm512_mul_ps(k0, sin), _mm512_mul_ps(k1, cos));
 
            _mm512_storeu_ps(k_h + i, k_rot0);
            _mm512_storeu_ps(k_h + i + half, k_rot1);
        }
#elif defined(__AVX__)
        for (; i + 8 <= half; i += 8) {
            __m256 k0 = _mm256_loadu_ps(k_h + i);
            __m256 k1 = _mm256_loadu_ps(k_h + i + half);
            __m256 cos = _mm256_loadu_ps(cos_ptr + i);
            __m256 sin = _mm256_loadu_ps(sin_ptr + i);
 
            __m256 k_rot0 = _mm256_sub_ps(_mm256_mul_ps(k0, cos), _mm256_mul_ps(k1, sin));
            __m256 k_rot1 = _mm256_add_ps(_mm256_mul_ps(k0, sin), _mm256_mul_ps(k1, cos));
 
            _mm256_storeu_ps(k_h + i, k_rot0);
            _mm256_storeu_ps(k_h + i + half, k_rot1);
        }
#endif
        for (; i < half; ++i) {
            float k0 = k_h[i];
            float k1 = k_h[i + half];
            float c = cos_ptr[i];
            float s = sin_ptr[i];
            k_h[i] = k0 * c - k1 * s;
            k_h[i + half] = k0 * s + k1 * c;
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) {
            k_h[d] = 0.0f;
        }
    }
}

Referenced by mega_fused_attention_decode_workspace().

mega_fused_attention_decode()

void mega_fused_attention_decode	(	float *	output,
		const float *	input,
		const float *	residual,
		const float *	ln1_gamma,
		const float *	wq,
		const float *	bq,
		const float *	wk,
		const float *	bk,
		const float *	wv,
		const float *	bv,
		const float *	wo,
		const float *	bo,
		float *	kv_cache_k,
		float *	kv_cache_v,
		const float *	rope_cos,
		const float *	rope_sin,
		int	pos,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		int	cache_capacity,
		float	eps
	)

Mega-fused attention for decode mode (single token)

This is the "holy grail" - all operations fused, no intermediates to DRAM.

Parameters

output	Output [aligned_embed_dim] (includes residual add)
input	Input [aligned_embed_dim]
residual	Residual input [aligned_embed_dim] (or NULL)
ln1_gamma	RMSNorm gamma [embed_dim]
wq	Q weights (quantized) [num_heads * aligned_head_dim * aligned_embed_dim]
bq	Q bias [num_heads * aligned_head_dim] (or NULL)
wq_dt	Q weight dtype (CK_DT_Q5_0/CK_DT_Q8_0/CK_DT_FP32)
wk	K weights (quantized) [num_kv_heads * aligned_head_dim * aligned_embed_dim]
bk	K bias [num_kv_heads * aligned_head_dim] (or NULL)
wk_dt	K weight dtype (CK_DT_Q5_0/CK_DT_Q8_0/CK_DT_FP32)
wv	V weights (quantized) [num_kv_heads * aligned_head_dim * aligned_embed_dim]
bv	V bias [num_kv_heads * aligned_head_dim] (or NULL)
wv_dt	V weight dtype (CK_DT_Q5_0/CK_DT_Q8_0/CK_DT_FP32)
wo	Output projection weights (quantized) [aligned_embed_dim * aligned_embed_dim]
bo	Output bias [aligned_embed_dim] (or NULL)
wo_dt	Output weight dtype (CK_DT_Q5_0/CK_DT_FP32)
kv_cache_k	KV cache for K [num_kv_heads * cache_capacity * aligned_head_dim]
kv_cache_v	KV cache for V [num_kv_heads * cache_capacity * aligned_head_dim]
rope_cos	RoPE cos [max_seq, head_dim/2]
rope_sin	RoPE sin [max_seq, head_dim/2]
pos	Current position in sequence
embed_dim	Model hidden dimension (unpadded)
aligned_embed_dim	Aligned hidden dimension
num_heads	Number of attention heads
num_kv_heads	Number of KV heads (for GQA)
head_dim	Head dimension (unpadded)
aligned_head_dim	Aligned head dimension
cache_capacity	KV cache capacity (stride in tokens)
eps	RMSNorm epsilon
scratch	Scratch buffer from mega_fused_attention_prefill_scratch_size()

Definition at line 678 of file mega_fused_attention_avx.c.

{
    if (num_heads <= 0 || num_kv_heads <= 0 || aligned_head_dim <= 0) {
        return;
    }
    const size_t q_elems = (size_t)num_heads * (size_t)aligned_head_dim;
    const size_t kv_elems = (size_t)num_kv_heads * (size_t)aligned_head_dim;
    if (q_elems > SIZE_MAX / (2 * sizeof(float)) ||
        kv_elems > SIZE_MAX / (2 * sizeof(float)) ||
        q_elems > SIZE_MAX / 2 - kv_elems) {
        return;
    }
    const size_t total_elems = 2 * (q_elems + kv_elems);
    if (total_elems > SIZE_MAX / sizeof(float)) {
        return;
    }
 
    float stack_workspace[4 * MEGA_STACK_MAX];
    float *workspace = stack_workspace;
    if (total_elems > 4 * MEGA_STACK_MAX) {
        workspace = (float *)malloc(total_elems * sizeof(float));
        if (!workspace) {
            return;
        }
    }
 
    const size_t q_output_bytes = 2 * q_elems * sizeof(float);
    const size_t kv_bytes = 2 * kv_elems * sizeof(float);
    mega_fused_attention_decode_workspace(
        output, input, residual, ln1_gamma,
        wq, bq, wk, bk, wv, bv, wo, bo,
        kv_cache_k, kv_cache_v, rope_cos, rope_sin,
        pos, embed_dim, aligned_embed_dim, num_heads, num_kv_heads,
        head_dim, aligned_head_dim, cache_capacity, eps,
        workspace, q_output_bytes,
        workspace + 2 * q_elems, kv_bytes);
 
    if (workspace != stack_workspace) {
        free(workspace);
    }
}

References mega_fused_attention_decode_workspace(), and MEGA_STACK_MAX.

mega_fused_attention_decode_workspace()

void mega_fused_attention_decode_workspace	(	float *	output,
		const float *	input,
		const float *	residual,
		const float *	ln1_gamma,
		const float *	wq,
		const float *	bq,
		const float *	wk,
		const float *	bk,
		const float *	wv,
		const float *	bv,
		const float *	wo,
		const float *	bo,
		float *	kv_cache_k,
		float *	kv_cache_v,
		const float *	rope_cos,
		const float *	rope_sin,
		int	pos,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		int	cache_capacity,
		float	eps,
		float *	q_output_workspace,
		size_t	q_output_workspace_bytes,
		float *	kv_workspace,
		size_t	kv_workspace_bytes
	)

Full mega-fused attention for decode.

RMSNorm → QKV → RoPE → Flash Attn → OutProj + Residual

Definition at line 589 of file mega_fused_attention_avx.c.

{
    if (!output || !input || !ln1_gamma || !wq || !wk || !wv || !wo ||
        !kv_cache_k || !kv_cache_v || !q_output_workspace || !kv_workspace) {
        return;
    }
    if (embed_dim <= 0 || aligned_embed_dim <= 0 || head_dim <= 0 || aligned_head_dim <= 0 ||
        num_heads <= 0 || num_kv_heads <= 0 || cache_capacity <= 0) {
        return;
    }
    if (pos < 0 || pos >= cache_capacity) {
        return;
    }
    if (aligned_embed_dim < embed_dim || aligned_head_dim < head_dim) {
        return;
    }
 
    const size_t q_elems = (size_t)num_heads * (size_t)aligned_head_dim;
    const size_t kv_elems = (size_t)num_kv_heads * (size_t)aligned_head_dim;
    if (q_elems > SIZE_MAX / (2 * sizeof(float)) ||
        kv_elems > SIZE_MAX / (2 * sizeof(float))) {
        return;
    }
    const size_t q_output_required = 2 * q_elems * sizeof(float);
    const size_t kv_required = 2 * kv_elems * sizeof(float);
    if (q_output_workspace_bytes < q_output_required ||
        kv_workspace_bytes < kv_required) {
        return;
    }
 
    float *q = q_output_workspace;
    float *o = q + q_elems;
    float *k = kv_workspace;
    float *v = k + kv_elems;
 
    mega_fuse_rmsnorm_qkv_avx(q, k, v, input, ln1_gamma,
                              wq, bq, wk, bk, wv, bv,
                              embed_dim, aligned_embed_dim,
                              num_heads, num_kv_heads,
                              head_dim, aligned_head_dim, eps);
 
    if (rope_cos && rope_sin) {
        mega_fuse_rope_inplace_avx(q, k, rope_cos, rope_sin, pos,
                                   num_heads, num_kv_heads,
                                   head_dim, aligned_head_dim);
    }
 
    kv_cache_write_head_major(k, v,
                              kv_cache_k, kv_cache_v,
                              num_kv_heads, pos,
                              cache_capacity,
                              head_dim, aligned_head_dim);
 
    mega_fuse_flash_attention_avx(o, q, kv_cache_k, kv_cache_v,
                                  num_heads, num_kv_heads,
                                  pos + 1, cache_capacity,
                                  head_dim, aligned_head_dim);
 
    mega_fuse_output_proj_residual(o, wo, bo, residual, output,
                                   embed_dim, aligned_embed_dim,
                                   num_heads, head_dim, aligned_head_dim);
 
}

References kv_cache_write_head_major(), mega_fuse_flash_attention_avx(), mega_fuse_output_proj_residual(), mega_fuse_rmsnorm_qkv_avx(), and mega_fuse_rope_inplace_avx().

Referenced by mega_fused_attention_decode().