Mega-Fused Attention Kernel. More...

#include <stdint.h>
#include "ckernel_dtype.h"

Macros
#define	MEGA_FUSE_KV_TILE 64

#define	MEGA_FUSE_Q_TILE 64

Functions
void	mega_fuse_get_optimal_tiles (int q_tile, int kv_tile, int head_dim)
	Get optimal tile sizes for current CPU. More...

void	mega_fuse_report_stats (int hidden, int num_layers, int seq_len)
	Report memory savings from mega-fusion. More...

void	mega_fuse_rmsnorm_qkv (float q_out, float k_out, float v_out, const float input, const float gamma, const float W_qkv, const float *b_qkv, int hidden, int num_heads, int num_kv_heads, int head_dim, float eps)
	Phase 1: Fused RMSNorm + QKV (intermediates in registers) More...

void	mega_fuse_rmsnorm_qkv_rope (float q_out, float k_out, float v_out, const float input, const float gamma, const float W_qkv, const float b_qkv, const float rope_cos, const float *rope_sin, int pos, int hidden, int num_heads, int num_kv_heads, int head_dim, int max_seq, float eps)
	Phase 2: Fused RMSNorm + QKV + RoPE. More...

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) More...

void	mega_fused_attention_prefill (float output, const float input, const float residual, const float ln1_gamma, const void wq, const float bq, CKDataType wq_dt, const void wk, const float bk, CKDataType wk_dt, const void wv, const float bv, CKDataType wv_dt, const void wo, const float bo, CKDataType wo_dt, float kv_cache_k, float kv_cache_v, const float rope_cos, const float rope_sin, int start_pos, int tokens, int cache_capacity, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, float eps, void *scratch)
	Mega-fused attention for prefill mode (multiple tokens) More...

void	mega_fused_attention_prefill_q8_0 (float output, const float input, const float residual, const float ln1_gamma, const void wq, const float bq, CKDataType wq_dt, const void wk, const float bk, CKDataType wk_dt, const void wv, const float bv, CKDataType wv_dt, const void wo, const float bo, CKDataType wo_dt, float kv_cache_k, float kv_cache_v, const float rope_cos, const float rope_sin, int start_pos, int tokens, int cache_capacity, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, float eps, void *scratch)
	Mega-fused prefill attention kernel (Q8_0 out-proj) More...

size_t	mega_fused_attention_prefill_q8_0_scratch_size (int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)
	Get scratch buffer size for mega_fused_attention_prefill_q8_0. More...

size_t	mega_fused_attention_prefill_scratch_size (int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)
	Get scratch buffer size for mega_fused_attention_prefill. More...

void	mega_fused_outproj_mlp_prefill (float output, const float attn_out, const float residual, const float ln2_gamma, const void wo, const float bo, CKDataType wo_dt, const void w1, const float b1, CKDataType w1_dt, const void w2, const float b2, CKDataType w2_dt, int tokens, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim, int intermediate_dim, int aligned_intermediate_dim, float eps, void *scratch)
	Mega-fused post-attention block (out-proj + RMSNorm2 + MLP) for prefill. More...

size_t	mega_fused_outproj_mlp_prefill_scratch_size (int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim, int aligned_intermediate_dim)
	Get scratch buffer size for mega_fused_outproj_mlp_prefill. More...

Detailed Description

Mega-Fused Attention Kernel.

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

All intermediates stay in registers/L1/L2. Single DRAM round-trip.

Memory Reduction: Before: ~32KB intermediates per layer (stack/heap) After: ~8KB total (input + output only) Reduction: 4-5× per layer, ~100× for full model

Performance Target: Move from memory-bound to compute-bound Expected speedup: 5-10× for attention-heavy workloads

Definition in file mega_fused_attention.h.

Macro Definition Documentation

◆ MEGA_FUSE_KV_TILE

#define MEGA_FUSE_KV_TILE 64

Definition at line 36 of file mega_fused_attention.h.

◆ MEGA_FUSE_Q_TILE

#define MEGA_FUSE_Q_TILE 64

Definition at line 32 of file mega_fused_attention.h.

Function Documentation

◆ mega_fuse_get_optimal_tiles()

void mega_fuse_get_optimal_tiles	(	int *	q_tile,
		int *	kv_tile,
		int	head_dim
	)

Get optimal tile sizes for current CPU.

◆ mega_fuse_report_stats()

void mega_fuse_report_stats	(	int	hidden,
		int	num_layers,
		int	seq_len
	)

Report memory savings from mega-fusion.

◆ mega_fuse_rmsnorm_qkv()

void mega_fuse_rmsnorm_qkv	(	float *	q_out,
		float *	k_out,
		float *	v_out,
		const float *	input,
		const float *	gamma,
		const float *	W_qkv,
		const float *	b_qkv,
		int	hidden,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		float	eps
	)

Phase 1: Fused RMSNorm + QKV (intermediates in registers)

Simpler step: Just fuse RMSNorm with QKV projection. Q/K/V stay in stack buffers, not DRAM.

◆ mega_fuse_rmsnorm_qkv_rope()

void mega_fuse_rmsnorm_qkv_rope	(	float *	q_out,
		float *	k_out,
		float *	v_out,
		const float *	input,
		const float *	gamma,
		const float *	W_qkv,
		const float *	b_qkv,
		const float *	rope_cos,
		const float *	rope_sin,
		int	pos,
		int	hidden,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq,
		float	eps
	)

Phase 2: Fused RMSNorm + QKV + RoPE.

Q/K stay in output buffers, RoPE applied in-place.

◆ 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()

Mega-fused attention for decode mode (single token)

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) {
         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;
  
     float q_stack[MEGA_STACK_MAX];
     float k_stack[MEGA_STACK_MAX];
     float v_stack[MEGA_STACK_MAX];
     float o_stack[MEGA_STACK_MAX];
  
     float *q = q_stack;
     float *k = k_stack;
     float *v = v_stack;
     float *o = o_stack;
  
     int free_q = 0;
     int free_k = 0;
     int free_v = 0;
     int free_o = 0;
  
     if (q_elems > MEGA_STACK_MAX) {
         q = (float *)malloc(q_elems * sizeof(float));
         if (!q) {
             return;
         }
         free_q = 1;
     }
     if (kv_elems > MEGA_STACK_MAX) {
         k = (float *)malloc(kv_elems * sizeof(float));
         if (!k) {
             if (free_q) free(q);
             return;
         }
         v = (float *)malloc(kv_elems * sizeof(float));
         if (!v) {
             if (free_q) free(q);
             free(k);
             return;
         }
         free_k = 1;
         free_v = 1;
     }
     if (q_elems > MEGA_STACK_MAX) {
         o = (float *)malloc(q_elems * sizeof(float));
         if (!o) {
             if (free_q) free(q);
             if (free_k) free(k);
             if (free_v) free(v);
             return;
         }
         free_o = 1;
     }
  
     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);
  
     if (free_q) free(q);
     if (free_k) free(k);
     if (free_v) free(v);
     if (free_o) free(o);
 }

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

◆ mega_fused_attention_prefill()

void mega_fused_attention_prefill	(	float *	output,
		const float *	input,
		const float *	residual,
		const float *	ln1_gamma,
		const void *	wq,
		const float *	bq,
		CKDataType	wq_dt,
		const void *	wk,
		const float *	bk,
		CKDataType	wk_dt,
		const void *	wv,
		const float *	bv,
		CKDataType	wv_dt,
		const void *	wo,
		const float *	bo,
		CKDataType	wo_dt,
		float *	kv_cache_k,
		float *	kv_cache_v,
		const float *	rope_cos,
		const float *	rope_sin,
		int	start_pos,
		int	tokens,
		int	cache_capacity,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		float	eps,
		void *	scratch
	)

Mega-fused attention for prefill mode (multiple tokens)

Parameters

output	Output [tokens, aligned_embed_dim] (includes residual add)
input	Input [tokens, aligned_embed_dim]
residual	Residual input [tokens, aligned_embed_dim] (or NULL)
ln1_gamma	RMSNorm gamma [embed_dim]
wq	Q weights [num_heads * aligned_head_dim * aligned_embed_dim]
bq	Q bias [num_heads * aligned_head_dim] (or NULL)
wk	K weights [num_kv_heads * aligned_head_dim * aligned_embed_dim]
bk	K bias [num_kv_heads * aligned_head_dim] (or NULL)
wv	V weights [num_kv_heads * aligned_head_dim * aligned_embed_dim]
bv	V bias [num_kv_heads * aligned_head_dim] (or NULL)
wo	Output projection weights [num_heads * aligned_embed_dim * aligned_head_dim]
bo	Output bias [aligned_embed_dim] (or NULL)
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]
start_pos	Starting position in KV cache
tokens	Number of tokens to process
cache_capacity	KV cache capacity (stride in tokens)
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
head_dim	Head dimension (unpadded)
aligned_head_dim	Aligned head dimension
eps	RMSNorm epsilon

Definition at line 160 of file mega_fused_attention_prefill.c.

 {
     if (!output || !input || !ln1_gamma || !wq || !wk || !wv || !wo ||
         !kv_cache_k || !kv_cache_v || !scratch) {
         return;
     }
     if (tokens <= 0 || cache_capacity <= 0 || embed_dim <= 0 || aligned_embed_dim <= 0 ||
         head_dim <= 0 || aligned_head_dim <= 0 || num_heads <= 0 || num_kv_heads <= 0) {
         return;
     }
     if (aligned_embed_dim < embed_dim || aligned_head_dim < head_dim) {
         return;
     }
     if (start_pos < 0 || start_pos + tokens > cache_capacity) {
         return;
     }
  
     const size_t q_bytes = (size_t)num_heads * (size_t)tokens *
                            (size_t)aligned_head_dim * sizeof(float);
     const size_t attn_bytes = q_bytes;
     const size_t proj_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
     const size_t qkv_scratch_bytes = fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(aligned_embed_dim);
  
     uint8_t *scratch_bytes = (uint8_t *)scratch;
     float *q = (float *)scratch_bytes;
     scratch_bytes += align_up_size(q_bytes, 64);
     float *attn_out = (float *)scratch_bytes;
     scratch_bytes += align_up_size(attn_bytes, 64);
     float *proj_scratch = (float *)scratch_bytes;
     scratch_bytes += align_up_size(proj_bytes, 64);
     void *qkv_scratch = (void *)scratch_bytes;
     (void)qkv_scratch_bytes;
  
     float *k_ptr = kv_cache_k + (size_t)start_pos * (size_t)aligned_head_dim;
     float *v_ptr = kv_cache_v + (size_t)start_pos * (size_t)aligned_head_dim;
  
     if (wq_dt == CK_DT_FP32 && wk_dt == CK_DT_FP32 && wv_dt == CK_DT_FP32) {
         fused_rmsnorm_qkv_prefill_head_major(input,
                                              ln1_gamma,
                                              (const float *)wq, bq,
                                              (const float *)wk, bk,
                                              (const float *)wv, bv,
                                              q,
                                              k_ptr,
                                              v_ptr,
                                              tokens,
                                              embed_dim,
                                              aligned_embed_dim,
                                              num_heads,
                                              num_kv_heads,
                                              head_dim,
                                              aligned_head_dim,
                                              cache_capacity,
                                              eps,
                                              qkv_scratch);
     } else {
         fused_rmsnorm_qkv_prefill_head_major_quant(input,
                                                    ln1_gamma,
                                                    wq, bq, wq_dt,
                                                    wk, bk, wk_dt,
                                                    wv, bv, wv_dt,
                                                    q,
                                                    k_ptr,
                                                    v_ptr,
                                                    tokens,
                                                    embed_dim,
                                                    aligned_embed_dim,
                                                    num_heads,
                                                    num_kv_heads,
                                                    head_dim,
                                                    aligned_head_dim,
                                                    cache_capacity,
                                                    eps,
                                                    qkv_scratch);
     }
  
     if (rope_cos && rope_sin) {
         rope_forward_qk_strided(q,
                                 k_ptr,
                                 rope_cos,
                                 rope_sin,
                                 num_heads,
                                 num_kv_heads,
                                 tokens,
                                 head_dim,
                                 aligned_head_dim,
                                 start_pos,
                                 tokens,
                                 cache_capacity);
     }
  
     if (start_pos == 0) {
         attention_forward_causal_head_major_gqa_flash_strided(q,
                                                              k_ptr,
                                                              v_ptr,
                                                              attn_out,
                                                              num_heads,
                                                              num_kv_heads,
                                                              tokens,
                                                              head_dim,
                                                              aligned_head_dim,
                                                              cache_capacity);
     } else {
         const float scale = 1.0f / sqrtf((float)head_dim);
         const size_t q_head_stride = (size_t)tokens * (size_t)aligned_head_dim;
         const size_t kv_head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
  
         for (int h = 0; h < num_heads; ++h) {
             int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
             const float *k_head = kv_cache_k + (size_t)kv_head * kv_head_stride;
             const float *v_head = kv_cache_v + (size_t)kv_head * kv_head_stride;
  
             for (int i = 0; i < tokens; ++i) {
                 const float *q_vec = q + (size_t)h * q_head_stride + (size_t)i * (size_t)aligned_head_dim;
                 float *out_vec = attn_out + (size_t)h * q_head_stride + (size_t)i * (size_t)aligned_head_dim;
                 attention_flash_decode(out_vec,
                                        q_vec,
                                        k_head,
                                        v_head,
                                        1,
                                        start_pos + i + 1,
                                        1,
                                        aligned_head_dim,
                                        scale);
             }
         }
     }
  
     if ((num_heads * aligned_head_dim) != aligned_embed_dim) {
         return;
     }
  
     if (wo_dt == CK_DT_Q5_0 &&
         ck_q8_0_outproj_enabled() &&
         (aligned_head_dim % QK5_0) == 0 &&
         (aligned_embed_dim % QK5_0) == 0) {
         /* Quantized activations path: Q8_0 attn_out + Q5_0 weights. */
         uint8_t *attn_q8 = (uint8_t *)q;
         quantize_attn_out_head_major_q8_0(attn_out,
                                           attn_q8,
                                           tokens,
                                           num_heads,
                                           aligned_head_dim);
         out_proj_head_major_q5_0_q8_0(attn_q8,
                                       wo,
                                       bo,
                                       output,
                                       tokens,
                                       aligned_embed_dim,
                                       num_heads,
                                       aligned_head_dim);
     } else if (wo_dt == CK_DT_Q5_0 &&
                (aligned_head_dim % QK5_0) == 0 &&
                (aligned_embed_dim % QK5_0) == 0) {
         /* Head-major output projection with Q5_0 weights - no flatten needed */
         ck_gemm_nt_head_major_q5_0(attn_out,
                                     wo,
                                     bo,
                                     output,
                                     tokens,
                                     aligned_embed_dim,
                                     num_heads,
                                     aligned_head_dim);
     } else if (wo_dt == CK_DT_Q8_0 &&
                (aligned_head_dim % QK8_0) == 0 &&
                (aligned_embed_dim % QK8_0) == 0) {
         /* Head-major output projection with Q8_0 weights - no flatten needed */
         ck_gemm_nt_head_major_q8_0(attn_out,
                                     wo,
                                     bo,
                                     output,
                                     tokens,
                                     aligned_embed_dim,
                                     num_heads,
                                     aligned_head_dim);
     } else {
         /* Fallback: flatten then GEMM (slow path) */
         flatten_head_major(attn_out,
                            proj_scratch,
                            tokens,
                            aligned_embed_dim,
                            num_heads,
                            aligned_head_dim);
  
         ck_gemm_nt_quant(proj_scratch,
                          wo,
                          bo,
                          output,
                          tokens,
                          aligned_embed_dim,
                          aligned_embed_dim,
                          wo_dt);
     }
  
     if (residual) {
         ck_residual_add_token_major(residual,
                                     output,
                                     output,
                                     tokens,
                                     aligned_embed_dim);
     }
  
 }

References align_up_size(), attention_flash_decode(), attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_FP32, CK_DT_Q5_0, CK_DT_Q8_0, ck_gemm_nt_head_major_q5_0(), ck_gemm_nt_head_major_q8_0(), ck_gemm_nt_quant(), ck_q8_0_outproj_enabled(), ck_residual_add_token_major(), flatten_head_major(), fused_rmsnorm_qkv_prefill_head_major(), fused_rmsnorm_qkv_prefill_head_major_quant(), fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(), out_proj_head_major_q5_0_q8_0(), QK5_0, QK8_0, quantize_attn_out_head_major_q8_0(), and rope_forward_qk_strided().

◆ mega_fused_attention_prefill_q8_0()

void mega_fused_attention_prefill_q8_0	(	float *	output,
		const float *	input,
		const float *	residual,
		const float *	ln1_gamma,
		const void *	wq,
		const float *	bq,
		CKDataType	wq_dt,
		const void *	wk,
		const float *	bk,
		CKDataType	wk_dt,
		const void *	wv,
		const float *	bv,
		CKDataType	wv_dt,
		const void *	wo,
		const float *	bo,
		CKDataType	wo_dt,
		float *	kv_cache_k,
		float *	kv_cache_v,
		const float *	rope_cos,
		const float *	rope_sin,
		int	start_pos,
		int	tokens,
		int	cache_capacity,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		float	eps,
		void *	scratch
	)

Mega-fused prefill attention kernel (Q8_0 out-proj)

Same layout and scratch requirements as mega_fused_attention_prefill.

Definition at line 105 of file mega_fused_attention_prefill_q8_0.c.

 {
     if (!output || !input || !ln1_gamma || !wq || !wk || !wv || !wo ||
         !kv_cache_k || !kv_cache_v || !scratch) {
         return;
     }
     if (tokens <= 0 || cache_capacity <= 0 || embed_dim <= 0 || aligned_embed_dim <= 0 ||
         head_dim <= 0 || aligned_head_dim <= 0 || num_heads <= 0 || num_kv_heads <= 0) {
         return;
     }
     if (aligned_embed_dim < embed_dim || aligned_head_dim < head_dim) {
         return;
     }
     if (start_pos < 0 || start_pos + tokens > cache_capacity) {
         return;
     }
     if (wo_dt != CK_DT_Q8_0) {
         return;
     }
     if ((aligned_head_dim % QK8_0) != 0 || (aligned_embed_dim % QK8_0) != 0) {
         return;
     }
  
     const size_t q_bytes = (size_t)num_heads * (size_t)tokens *
                            (size_t)aligned_head_dim * sizeof(float);
     const size_t attn_bytes = q_bytes;
     const size_t proj_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
     const size_t qkv_scratch_bytes = fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(aligned_embed_dim);
  
     uint8_t *scratch_bytes = (uint8_t *)scratch;
     float *q = (float *)scratch_bytes;
     scratch_bytes += align_up_size(q_bytes, 64);
     float *attn_out = (float *)scratch_bytes;
     scratch_bytes += align_up_size(attn_bytes, 64);
     float *proj_scratch = (float *)scratch_bytes;
     scratch_bytes += align_up_size(proj_bytes, 64);
     void *qkv_scratch = (void *)scratch_bytes;
     (void)qkv_scratch_bytes;
     (void)proj_scratch;
  
     float *k_ptr = kv_cache_k + (size_t)start_pos * (size_t)aligned_head_dim;
     float *v_ptr = kv_cache_v + (size_t)start_pos * (size_t)aligned_head_dim;
  
     if (wq_dt == CK_DT_FP32 && wk_dt == CK_DT_FP32 && wv_dt == CK_DT_FP32) {
         fused_rmsnorm_qkv_prefill_head_major(input,
                                              ln1_gamma,
                                              (const float *)wq, bq,
                                              (const float *)wk, bk,
                                              (const float *)wv, bv,
                                              q,
                                              k_ptr,
                                              v_ptr,
                                              tokens,
                                              embed_dim,
                                              aligned_embed_dim,
                                              num_heads,
                                              num_kv_heads,
                                              head_dim,
                                              aligned_head_dim,
                                              cache_capacity,
                                              eps,
                                              qkv_scratch);
     } else {
         fused_rmsnorm_qkv_prefill_head_major_quant(input,
                                                    ln1_gamma,
                                                    wq, bq, wq_dt,
                                                    wk, bk, wk_dt,
                                                    wv, bv, wv_dt,
                                                    q,
                                                    k_ptr,
                                                    v_ptr,
                                                    tokens,
                                                    embed_dim,
                                                    aligned_embed_dim,
                                                    num_heads,
                                                    num_kv_heads,
                                                    head_dim,
                                                    aligned_head_dim,
                                                    cache_capacity,
                                                    eps,
                                                    qkv_scratch);
     }
  
     if (rope_cos && rope_sin) {
         rope_forward_qk_strided(q,
                                 k_ptr,
                                 rope_cos,
                                 rope_sin,
                                 num_heads,
                                 num_kv_heads,
                                 tokens,
                                 head_dim,
                                 aligned_head_dim,
                                 start_pos,
                                 tokens,
                                 cache_capacity);
     }
  
     if (start_pos == 0) {
         attention_forward_causal_head_major_gqa_flash_strided(q,
                                                              k_ptr,
                                                              v_ptr,
                                                              attn_out,
                                                              num_heads,
                                                              num_kv_heads,
                                                              tokens,
                                                              head_dim,
                                                              aligned_head_dim,
                                                              cache_capacity);
     } else {
         const float scale = 1.0f / sqrtf((float)head_dim);
         const size_t q_head_stride = (size_t)tokens * (size_t)aligned_head_dim;
         const size_t kv_head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
  
         for (int h = 0; h < num_heads; ++h) {
             int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
             const float *k_head = kv_cache_k + (size_t)kv_head * kv_head_stride;
             const float *v_head = kv_cache_v + (size_t)kv_head * kv_head_stride;
  
             for (int i = 0; i < tokens; ++i) {
                 const float *q_vec = q + (size_t)h * q_head_stride + (size_t)i * (size_t)aligned_head_dim;
                 float *out_vec = attn_out + (size_t)h * q_head_stride + (size_t)i * (size_t)aligned_head_dim;
                 attention_flash_decode(out_vec,
                                        q_vec,
                                        k_head,
                                        v_head,
                                        1,
                                        start_pos + i + 1,
                                        1,
                                        aligned_head_dim,
                                        scale);
             }
         }
     }
  
     if ((num_heads * aligned_head_dim) != aligned_embed_dim) {
         return;
     }
  
     /* Quantized activations path: Q8_0 attn_out + Q8_0 weights. */
     {
         uint8_t *attn_q8 = (uint8_t *)q;
         quantize_attn_out_head_major_q8_0(attn_out,
                                           attn_q8,
                                           tokens,
                                           num_heads,
                                           aligned_head_dim);
         out_proj_head_major_q8_0_q8_0(attn_q8,
                                       wo,
                                       bo,
                                       output,
                                       tokens,
                                       aligned_embed_dim,
                                       num_heads,
                                       aligned_head_dim);
     }
  
     if (residual) {
         ck_residual_add_token_major(residual,
                                     output,
                                     output,
                                     tokens,
                                     aligned_embed_dim);
     }
 }

References align_up_size(), attention_flash_decode(), attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_FP32, CK_DT_Q8_0, ck_residual_add_token_major(), fused_rmsnorm_qkv_prefill_head_major(), fused_rmsnorm_qkv_prefill_head_major_quant(), fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(), out_proj_head_major_q8_0_q8_0(), QK8_0, quantize_attn_out_head_major_q8_0(), and rope_forward_qk_strided().

◆ mega_fused_attention_prefill_q8_0_scratch_size()

size_t mega_fused_attention_prefill_q8_0_scratch_size	(	int	tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim
	)

Get scratch buffer size for mega_fused_attention_prefill_q8_0.

Definition at line 84 of file mega_fused_attention_prefill_q8_0.c.

 {
     if (tokens <= 0 || aligned_embed_dim <= 0 || num_heads <= 0 || aligned_head_dim <= 0) {
         return 0;
     }
  
     const size_t q_bytes = (size_t)num_heads * (size_t)tokens *
                            (size_t)aligned_head_dim * sizeof(float);
     const size_t attn_bytes = q_bytes;
     const size_t proj_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
     const size_t qkv_scratch = fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(aligned_embed_dim);
  
     return align_up_size(q_bytes, 64) +
            align_up_size(attn_bytes, 64) +
            align_up_size(proj_bytes, 64) +
            align_up_size(qkv_scratch, 64);
 }

References align_up_size(), and fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size().

◆ mega_fused_attention_prefill_scratch_size()

size_t mega_fused_attention_prefill_scratch_size	(	int	tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim
	)

Get scratch buffer size for mega_fused_attention_prefill.

Definition at line 139 of file mega_fused_attention_prefill.c.

 {
     if (tokens <= 0 || aligned_embed_dim <= 0 || num_heads <= 0 || aligned_head_dim <= 0) {
         return 0;
     }
  
     const size_t q_bytes = (size_t)num_heads * (size_t)tokens *
                            (size_t)aligned_head_dim * sizeof(float);
     const size_t attn_bytes = q_bytes;
     const size_t proj_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
     const size_t qkv_scratch = fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size(aligned_embed_dim);
  
     return align_up_size(q_bytes, 64) +
            align_up_size(attn_bytes, 64) +
            align_up_size(proj_bytes, 64) +
            align_up_size(qkv_scratch, 64);
 }

References align_up_size(), and fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size().

◆ mega_fused_outproj_mlp_prefill()

void mega_fused_outproj_mlp_prefill	(	float *	output,
		const float *	attn_out,
		const float *	residual,
		const float *	ln2_gamma,
		const void *	wo,
		const float *	bo,
		CKDataType	wo_dt,
		const void *	w1,
		const float *	b1,
		CKDataType	w1_dt,
		const void *	w2,
		const float *	b2,
		CKDataType	w2_dt,
		int	tokens,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim,
		int	intermediate_dim,
		int	aligned_intermediate_dim,
		float	eps,
		void *	scratch
	)

Mega-fused post-attention block (out-proj + RMSNorm2 + MLP) for prefill.

Uses head-major attention output and quantized out-proj (Q5_0/Q8_0 weights).

Definition at line 184 of file mega_fused_outproj_mlp_prefill.c.

 {
     if (!output || !attn_out || !residual || !ln2_gamma ||
         !wo || !w1 || !w2 || !scratch) {
         return;
     }
     if (tokens <= 0 || embed_dim <= 0 || aligned_embed_dim <= 0 ||
         num_heads <= 0 || aligned_head_dim <= 0 ||
         intermediate_dim <= 0 || aligned_intermediate_dim <= 0) {
         return;
     }
     if (aligned_embed_dim < embed_dim || aligned_head_dim <= 0 ||
         aligned_intermediate_dim < intermediate_dim) {
         return;
     }
     if (aligned_embed_dim != num_heads * aligned_head_dim) {
         return;
     }
     if ((aligned_embed_dim % 32) != 0 || (aligned_head_dim % 32) != 0) {
         return;
     }
     if ((aligned_intermediate_dim % QK_K) != 0) {
         return;
     }
     if (wo_dt != CK_DT_Q5_0 && wo_dt != CK_DT_Q8_0) {
         return;
     }
     if (w1_dt != CK_DT_Q5_0 && w1_dt != CK_DT_Q8_0) {
         return;
     }
     if (w2_dt != CK_DT_Q4_K && w2_dt != CK_DT_Q6_K) {
         return;
     }
  
     const size_t q8_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0,
                                                    (size_t)aligned_head_dim);
     const size_t attn_q8_bytes = (size_t)num_heads * (size_t)tokens * q8_row_bytes;
     const size_t h1_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
     const size_t ln2_bytes = h1_bytes;
  
     uint8_t *scratch_bytes = (uint8_t *)scratch;
     uint8_t *attn_q8 = scratch_bytes;
     scratch_bytes += align_up_size(attn_q8_bytes, 64);
     float *h1 = (float *)scratch_bytes;
     scratch_bytes += align_up_size(h1_bytes, 64);
     float *ln2_out = (float *)scratch_bytes;
     scratch_bytes += align_up_size(ln2_bytes, 64);
     void *mlp_scratch = (void *)scratch_bytes;
  
     quantize_attn_out_head_major_q8_0(attn_out,
                                       attn_q8,
                                       tokens,
                                       num_heads,
                                       aligned_head_dim);
  
     if (wo_dt == CK_DT_Q8_0) {
         out_proj_head_major_q8_0_q8_0(attn_q8,
                                       wo,
                                       bo,
                                       h1,
                                       tokens,
                                       aligned_embed_dim,
                                       num_heads,
                                       aligned_head_dim);
     } else {
         out_proj_head_major_q5_0_q8_0(attn_q8,
                                       wo,
                                       bo,
                                       h1,
                                       tokens,
                                       aligned_embed_dim,
                                       num_heads,
                                       aligned_head_dim);
     }
  
     for (int t = 0; t < tokens; ++t) {
         const float *res_row = residual + (size_t)t * (size_t)aligned_embed_dim;
         float *h1_row = h1 + (size_t)t * (size_t)aligned_embed_dim;
         add_inplace_f32(h1_row, res_row, aligned_embed_dim);
     }
  
     rmsnorm_forward(h1,
                     ln2_gamma,
                     ln2_out,
                     NULL,
                     tokens,
                     embed_dim,
                     aligned_embed_dim,
                     eps);
  
     fused_mlp_swiglu_prefill_w1w2_quant(ln2_out,
                                         w1,
                                         b1,
                                         w1_dt,
                                         w2,
                                         b2,
                                         w2_dt,
                                         output,
                                         tokens,
                                         embed_dim,
                                         aligned_embed_dim,
                                         intermediate_dim,
                                         aligned_intermediate_dim,
                                         mlp_scratch);
  
     for (int t = 0; t < tokens; ++t) {
         const float *h1_row = h1 + (size_t)t * (size_t)aligned_embed_dim;
         float *out_row = output + (size_t)t * (size_t)aligned_embed_dim;
         add_inplace_f32(out_row, h1_row, aligned_embed_dim);
     }
 }

References add_inplace_f32(), align_up_size(), CK_DT_Q4_K, CK_DT_Q5_0, CK_DT_Q6_K, CK_DT_Q8_0, ck_dtype_row_bytes(), fused_mlp_swiglu_prefill_w1w2_quant(), out_proj_head_major_q5_0_q8_0(), out_proj_head_major_q8_0_q8_0(), QK_K, quantize_attn_out_head_major_q8_0(), and rmsnorm_forward().

◆ mega_fused_outproj_mlp_prefill_scratch_size()

size_t mega_fused_outproj_mlp_prefill_scratch_size	(	int	tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim,
		int	aligned_intermediate_dim
	)

Get scratch buffer size for mega_fused_outproj_mlp_prefill.

Definition at line 159 of file mega_fused_outproj_mlp_prefill.c.

 {
     if (tokens <= 0 || aligned_embed_dim <= 0 || num_heads <= 0 ||
         aligned_head_dim <= 0 || aligned_intermediate_dim <= 0) {
         return 0;
     }
  
     const size_t q8_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0,
                                                    (size_t)aligned_head_dim);
     const size_t attn_q8_bytes = (size_t)num_heads * (size_t)tokens * q8_row_bytes;
     const size_t h1_bytes = (size_t)tokens * (size_t)aligned_embed_dim * sizeof(float);
     const size_t ln2_bytes = h1_bytes;
     const size_t mlp_scratch = fused_mlp_swiglu_prefill_w1w2_quant_scratch_size(
         aligned_embed_dim, aligned_intermediate_dim);
  
     return align_up_size(attn_q8_bytes, 64) +
            align_up_size(h1_bytes, 64) +
            align_up_size(ln2_bytes, 64) +
            align_up_size(mlp_scratch, 64);
 }

References align_up_size(), CK_DT_Q8_0, ck_dtype_row_bytes(), and fused_mlp_swiglu_prefill_w1w2_quant_scratch_size().

Macros

Functions