gemm_head_major_output.c File Reference

Output projection from head-major attention (NO LAYOUT CONVERSION) More...

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include "ckernel_quant.h"
#include "ckernel_dtype.h"

Go to the source code of this file.

Functions
void	ck_gemm_nt_head_major_q5_0 (const float *attn_out, const void *wo, const float *bias, float *output, int tokens, int embed_dim, int num_heads, int head_dim)
	Output projection from head-major attention (auto-dispatch) More...
void	ck_gemm_nt_head_major_q8_0 (const float *attn_out, const void *wo, const float *bias, float *output, int tokens, int embed_dim, int num_heads, int head_dim)
	Output projection from head-major attention (Q8_0 weights) More...
void	dequant_q5_0_row (const void *src, float *dst, size_t n_elements)
	Dequantize Q5_0 row (multiple blocks) More...
void	gemv_nt_q5_0_head_major_output (float *output, const float *attn_out, const void *wo, const float *bias, int tokens, int embed_dim, int num_heads, int head_dim)
	Output projection reading head-major attention output (Q5_0 weights) More...

Detailed Description

Output projection from head-major attention (NO LAYOUT CONVERSION)

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

PROBLEM THIS SOLVES:

The standard mega_fused_attention_prefill has a bottleneck: attn_out [num_heads, tokens, head_dim] (head-major) → flatten_head_major() - 448 memcpy calls for 32 tokens × 14 heads! → token-major buffer → GEMM output projection

This kernel eliminates the flatten by reading head-major data directly with strided access. The output projection computes:

output[t, n] = bias[n] + sum_h wo[n, h*head_dim:(h+1)*head_dim] @ attn_out[h, t, :]

where wo is Q5_0 quantized [embed_dim, embed_dim] and attn_out is head-major.

Expected speedup: 1.5-2x by eliminating 448 small memcpy calls.

Definition in file gemm_head_major_output.c.

Function Documentation

ck_gemm_nt_head_major_q5_0()

void ck_gemm_nt_head_major_q5_0	(	const float *	attn_out,
		const void *	wo,
		const float *	bias,
		float *	output,
		int	tokens,
		int	embed_dim,
		int	num_heads,
		int	head_dim
	)

Output projection from head-major attention (auto-dispatch)

This replaces flatten_head_major() + ck_gemm_nt_quant() with a single strided-access kernel that reads head-major attention output directly.

Definition at line 328 of file gemm_head_major_output.c.

{
#if defined(__AVX__) && defined(__F16C__)
    gemv_nt_q5_0_head_major_output_avx(output, attn_out, wo, bias,
                                       tokens, embed_dim, num_heads, head_dim);
#else
    gemv_nt_q5_0_head_major_output(output, attn_out, wo, bias,
                                   tokens, embed_dim, num_heads, head_dim);
#endif
}

References gemv_nt_q5_0_head_major_output().

Referenced by mega_fused_attention_prefill().

ck_gemm_nt_head_major_q8_0()

void ck_gemm_nt_head_major_q8_0	(	const float *	attn_out,
		const void *	wo,
		const float *	bias,
		float *	output,
		int	tokens,
		int	embed_dim,
		int	num_heads,
		int	head_dim
	)

Output projection from head-major attention (Q8_0 weights)

Definition at line 353 of file gemm_head_major_output.c.

{
    if (!output || !attn_out || !wo) return;
    if (tokens <= 0 || embed_dim <= 0 || num_heads <= 0 || head_dim <= 0) return;
 
    const int blocks_per_head = head_dim / QK8_0;
    const int blocks_per_row = embed_dim / QK8_0;
    const block_q8_0 *weights = (const block_q8_0 *)wo;
 
    const size_t token_stride = head_dim;
    const size_t head_stride = (size_t)tokens * token_stride;
 
    /* Initialize output */
    if (bias) {
        for (int t = 0; t < tokens; t++) {
            float *out_row = output + (size_t)t * embed_dim;
            for (int n = 0; n < embed_dim; n++) {
                out_row[n] = bias[n];
            }
        }
    } else {
        memset(output, 0, (size_t)tokens * embed_dim * sizeof(float));
    }
 
    /* Accumulate from each head */
    for (int h = 0; h < num_heads; h++) {
        const float *head_data = attn_out + (size_t)h * head_stride;
        const int head_offset = h * blocks_per_head;
 
        for (int n_block = 0; n_block < blocks_per_head; n_block++) {
            for (int n = 0; n < embed_dim; n++) {
                const block_q8_0 *w_row = weights + (size_t)n * blocks_per_row + head_offset + n_block;
                const float d = CK_FP16_TO_FP32(w_row->d);
 
                for (int t = 0; t < tokens; t++) {
                    const float *token_vec = head_data + (size_t)t * token_stride + (size_t)n_block * QK8_0;
                    float sum = 0.0f;
 
                    for (int j = 0; j < QK8_0; j++) {
                        sum += d * (float)w_row->qs[j] * token_vec[j];
                    }
 
                    output[(size_t)t * embed_dim + n] += sum;
                }
            }
        }
    }
}

References CK_FP16_TO_FP32, block_q8_0::d, QK8_0, and block_q8_0::qs.

Referenced by mega_fused_attention_prefill().

dequant_q5_0_row()

void dequant_q5_0_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q5_0 row (multiple blocks)

Definition at line 196 of file dequant_kernels.c.

{
    const block_q5_0 *blocks = (const block_q5_0 *)src;
    const size_t n_blocks = n_elements / QK5_0;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q5_0_block(&blocks[b], &dst[b * QK5_0]);
    }
}

gemv_nt_q5_0_head_major_output()

void gemv_nt_q5_0_head_major_output	(	float *	output,
		const float *	attn_out,
		const void *	wo,
		const float *	bias,
		int	tokens,
		int	embed_dim,
		int	num_heads,
		int	head_dim
	)

Output projection reading head-major attention output (Q5_0 weights)

Parameters

output	Output [tokens, embed_dim] (token-major, written contiguously)
attn_out	Attention output [num_heads, tokens, head_dim] (head-major, strided)
wo	Output weights in Q5_0 format [embed_dim, embed_dim]
bias	Optional bias [embed_dim]
tokens	Number of tokens
embed_dim	Output embedding dimension
num_heads	Number of attention heads
head_dim	Head dimension (must be multiple of 32 for Q5_0)

Definition at line 62 of file gemm_head_major_output.c.

{
    if (!output || !attn_out || !wo) return;
    if (tokens <= 0 || embed_dim <= 0 || num_heads <= 0 || head_dim <= 0) return;
 
    const int blocks_per_head = head_dim / QK5_0;
    const int blocks_per_row = embed_dim / QK5_0;
    const block_q5_0 *weights = (const block_q5_0 *)wo;
 
    /* Strides for head-major layout */
    const size_t token_stride = head_dim;           /* attn_out[h][t] offset */
    const size_t head_stride = (size_t)tokens * token_stride;  /* attn_out[h] offset */
 
    /* Initialize output with bias (if provided) */
    if (bias) {
        for (int t = 0; t < tokens; t++) {
            float *out_row = output + (size_t)t * embed_dim;
            for (int n = 0; n < embed_dim; n++) {
                out_row[n] = bias[n];
            }
        }
    } else {
        memset(output, 0, (size_t)tokens * embed_dim * sizeof(float));
    }
 
    /* Accumulate contributions from each head */
    for (int h = 0; h < num_heads; h++) {
        const float *head_data = attn_out + (size_t)h * head_stride;
 
        /* For each output row (n) corresponding to this head's slice */
        const int head_offset = h * blocks_per_head;
 
        for (int n_block = 0; n_block < blocks_per_head; n_block++) {
            for (int n = 0; n < embed_dim; n++) {
                const block_q5_0 *w_row = weights + (size_t)n * blocks_per_row + head_offset + n_block;
                const float d = CK_FP16_TO_FP32(w_row->d);
 
                /* Get high bits */
                uint32_t qh;
                memcpy(&qh, w_row->qh, sizeof(qh));
 
                /* Accumulate for all tokens at once (better cache reuse) */
                for (int t = 0; t < tokens; t++) {
                    const float *token_vec = head_data + (size_t)t * token_stride + (size_t)n_block * QK5_0;
                    float sum = 0.0f;
 
                    /* Q5_0 dot product for this block */
                    for (int j = 0; j < QK5_0 / 2; j++) {
                        const uint8_t packed = w_row->qs[j];
                        const int lo = (packed & 0x0F);
                        const int hi = (packed >> 4);
                        const int xh_0 = ((qh >> (j + 0)) << 4) & 0x10;
                        const int xh_1 = ((qh >> (j + 12))) & 0x10;
                        const int q0 = (lo | xh_0) - 16;
                        const int q1 = (hi | xh_1) - 16;
 
                        sum += d * (float)q0 * token_vec[j];
                        sum += d * (float)q1 * token_vec[j + 16];
                    }
 
                    output[(size_t)t * embed_dim + n] += sum;
                }
            }
        }
    }
}

References CK_FP16_TO_FP32, block_q5_0::d, block_q5_0::qh, QK5_0, and block_q5_0::qs.

Referenced by ck_gemm_nt_head_major_q5_0().