ckernel_orchestration.h File Reference

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

Go to the source code of this file.

Data Structures
struct	CKLayerBackwardParams
struct	CKLayerForwardParams
struct	CKLayerForwardParamsQ4K

Functions
void	ck_attention_project_head_major (const float *attn_out, const float *wo, const float *bo, float *out, float *scratch, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)
void	ck_attention_project_head_major_backward (const float *d_out, const float *attn_out, const float *wo, float *d_attn_out, float *d_wo, float *d_bo, int tokens, int aligned_embed_dim, int num_heads, int aligned_head_dim)
void	ck_attention_project_head_major_decode_token (const float *attn_token, const float *wo, const float *bo, float *out_token, int embed_dim, int aligned_embed_dim, int num_heads, int aligned_head_dim)
void	ck_gemm_nt_quant (const float *A, const void *B, const float *bias, float *C, int M, int N, int K, CKDataType dtype)
void	ck_layer_backward_rmsnorm_swiglu (const CKLayerBackwardParams *p)
void	ck_layer_forward_rmsnorm_swiglu (const CKLayerForwardParams *p)
void	ck_layer_forward_rmsnorm_swiglu_decode (const CKLayerForwardParams *p, int token_index, int cache_capacity)
void	ck_layer_forward_rmsnorm_swiglu_decode_fused (const CKLayerForwardParams *p, int token_index, int cache_capacity)
void	ck_layer_forward_rmsnorm_swiglu_decode_fused_attn (const CKLayerForwardParams *p, int token_index, int cache_capacity)
void	ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_mlp (const CKLayerForwardParams *p, int token_index, int cache_capacity)
void	ck_layer_forward_rmsnorm_swiglu_decode_q4_k (const CKLayerForwardParamsQ4K *p, int token_index, int cache_capacity)
void	ck_layer_forward_rmsnorm_swiglu_decode_quant (const CKLayerForwardParamsQ4K *p, int token_index, int cache_capacity)
void	ck_layer_forward_rmsnorm_swiglu_q4_k (const CKLayerForwardParamsQ4K *p)
void	ck_layer_forward_rmsnorm_swiglu_quant (const CKLayerForwardParamsQ4K *p)
void	ck_layer_forward_rmsnorm_swiglu_ref (const CKLayerForwardParams *p)
void	ck_mlp_swiglu_forward (const float *input, const float *w1, const float *b1, const float *w2, const float *b2, float *fc1_out, float *swiglu_out, float *output, int tokens, int aligned_embed_dim, int aligned_intermediate_dim)
void	ck_mlp_swiglu_forward_fully_fused_token (const float *input_row, const float *w1, const float *b1, const float *w2, const float *b2, float *output_row, int aligned_embed_dim, int aligned_intermediate_dim)
void	ck_mlp_swiglu_forward_fused_token (const float *input_row, const float *w1, const float *b1, const float *w2, const float *b2, float *swiglu_row, float *output_row, int aligned_embed_dim, int aligned_intermediate_dim)
void	ck_qkv_project_head_major (const float *input, const float *wq, const float *bq, const float *wk, const float *bk, const float *wv, const float *bv, float *q, float *k, float *v, int tokens, int kv_stride_tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)
void	ck_qkv_project_head_major_backward (const float *d_q, const float *d_k, const float *d_v, const float *input, const float *wq, const float *bq, const float *wk, const float *bk, const float *wv, const float *bv, float *d_input, float *d_wq, float *d_bq, float *d_wk, float *d_bk, float *d_wv, float *d_bv, float *scratch, int tokens, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim, int num_threads)
void	ck_qkv_project_head_major_token (const float *input_row, const float *wq, const float *bq, const float *wk, const float *bk, const float *wv, const float *bv, float *q_token, float *k_token, float *v_token, int aligned_embed_dim, int num_heads, int num_kv_heads, int aligned_head_dim)
void	ck_residual_add_backward (const float *d_out, float *d_a, float *d_b, int tokens, int aligned_embed_dim)
void	ck_residual_add_token_major (const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)

Detailed Description

---

LEGACY HEADER - NOT USED IN v6.6

This header declares v6.5 orchestration functions that are NO LONGER USED. v6.6 uses IR Lower 3 + codegen instead of hardcoded orchestration.

v6.6 Architecture (REPLACEMENT):

• Kernel dispatch: version/v6.6/scripts/build_ir_v6_6.py + ckernel_codegen.c
• Memory planning: version/v6.6/scripts/memory_planner_v6_6.py
• Registry: version/v6.6/kernel_maps/KERNEL_REGISTRY.json
• Kernel bindings: version/v6.6/kernel_maps/kernel_bindings.json

Deprecated functions (NOT used in v6.6):

• ck_layer_forward_rmsnorm_swiglu* -> IR Lower 3 + mega_fused_* kernels
• ck_qkv_project_head_major* -> q_proj/k_proj/v_proj ops in IR
• ck_attention_project_head_major* -> out_proj op in IR
• ck_mlp_swiglu_forward* -> mlp_gate_up/mlp_down ops in IR
• ck_gemm_nt_quant -> KERNEL_REGISTRY.json dispatch
• ck_residual_add_token_major -> residual_add op in IR

To remove completely:

1. Delete this header
2. Delete ckernel_orchestration.c
3. Remove from Makefile SRCS list

Last used: v6.5

Deprecated: v6.6 (2026-02)

Definition in file ckernel_orchestration.h.

Function Documentation

ck_attention_project_head_major()

void ck_attention_project_head_major	(	const float *	attn_out,
		const float *	wo,
		const float *	bo,
		float *	out,
		float *	scratch,
		int	tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim
	)

Definition at line 730 of file ckernel_orchestration.c.

{
    if (!attn_out || !wo || !out) {
        return;
    }
    if (num_heads > 1 && !scratch) {
        return;
    }
 
    size_t head_in_stride = (size_t)tokens * (size_t)aligned_head_dim;
    size_t head_weight_stride = (size_t)aligned_embed_dim * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        const float *head_in = attn_out + (size_t)h * head_in_stride;
        const float *wo_h = wo + (size_t)h * head_weight_stride;
 
        if (h == 0) {
            gemm_blocked_serial(head_in, wo_h, bo, out,
                                tokens, aligned_embed_dim, aligned_head_dim);
        } else {
            gemm_blocked_serial(head_in, wo_h, NULL, scratch,
                                tokens, aligned_embed_dim, aligned_head_dim);
            ck_add_inplace(out, scratch, tokens, aligned_embed_dim);
        }
    }
}

References ck_add_inplace(), and gemm_blocked_serial().

Referenced by ck_attention_project_head_major_quant(), and ck_layer_forward_rmsnorm_swiglu().

ck_attention_project_head_major_backward()

void ck_attention_project_head_major_backward	(	const float *	d_out,
		const float *	attn_out,
		const float *	wo,
		float *	d_attn_out,
		float *	d_wo,
		float *	d_bo,
		int	tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim
	)

Definition at line 800 of file ckernel_orchestration.c.

{
    if (!d_out || !attn_out || !wo || !d_attn_out || !d_wo || !d_bo) {
        return;
    }
 
    // Bias gradient: sum over tokens once (bias is applied once in forward).
    for (int d = 0; d < aligned_embed_dim; ++d) {
        d_bo[d] = 0.0f;
    }
    for (int t = 0; t < tokens; ++t) {
        const float *row = d_out + (size_t)t * (size_t)aligned_embed_dim;
        for (int d = 0; d < aligned_embed_dim; ++d) {
            d_bo[d] += row[d];
        }
    }
 
    size_t head_in_stride = (size_t)tokens * (size_t)aligned_head_dim;
    size_t head_weight_stride = (size_t)aligned_embed_dim * (size_t)aligned_head_dim;
 
    float *tmp_b = (float *)calloc((size_t)aligned_embed_dim, sizeof(float));
    if (!tmp_b) {
        return;
    }
 
    for (int h = 0; h < num_heads; ++h) {
        const float *head_in = attn_out + (size_t)h * head_in_stride;
        const float *wo_h = wo + (size_t)h * head_weight_stride;
        float *d_head_in = d_attn_out + (size_t)h * head_in_stride;
        float *d_wo_h = d_wo + (size_t)h * head_weight_stride;
 
        memset(tmp_b, 0, (size_t)aligned_embed_dim * sizeof(float));
        fc2_backward_kernel(d_out,
                            head_in,
                            wo_h,
                            d_head_in,
                            d_wo_h,
                            tmp_b,
                            tokens,
                            aligned_head_dim,
                            aligned_embed_dim,
                            1);
    }
 
    free(tmp_b);
}

References fc2_backward_kernel().

Referenced by ck_layer_backward_rmsnorm_swiglu().

ck_attention_project_head_major_decode_token()

void ck_attention_project_head_major_decode_token	(	const float *	attn_token,
		const float *	wo,
		const float *	bo,
		float *	out_token,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	aligned_head_dim
	)

Definition at line 115 of file attention_decode_fused.c.

{
    const size_t head_in_stride = (size_t)aligned_head_dim;
    const size_t head_weight_stride = (size_t)aligned_embed_dim * (size_t)aligned_head_dim;
 
#pragma omp parallel for schedule(static)
    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 *head_in = attn_token + (size_t)h * head_in_stride;
            const float *wo_row = wo + (size_t)h * head_weight_stride + (size_t)j * (size_t)aligned_head_dim;
            sum += ck_dot_f32(head_in, wo_row, aligned_head_dim);
        }
        out_token[j] = sum;
    }
 
    for (int j = embed_dim; j < aligned_embed_dim; ++j) {
        out_token[j] = 0.0f;
    }
}

References ck_dot_f32().

Referenced by ck_layer_forward_rmsnorm_swiglu_decode(), ck_layer_forward_rmsnorm_swiglu_decode_fused(), and ck_layer_forward_rmsnorm_swiglu_decode_quant().

ck_gemm_nt_quant()

void ck_gemm_nt_quant	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		CKDataType	dtype
	)

Definition at line 335 of file ckernel_orchestration.c.

{
    switch (dtype) {
    case CK_DT_FP32:
        gemm_blocked_serial(A, (const float *)B, bias, C, M, N, K);
        break;
    case CK_DT_Q4_K:
        gemm_nt_q4_k(A, B, bias, C, M, N, K);
        break;
    case CK_DT_Q6_K:
        gemm_nt_q6_k(A, B, bias, C, M, N, K);
        break;
    case CK_DT_Q4_0:
        gemm_nt_q4_0(A, B, bias, C, M, N, K);
        break;
    case CK_DT_Q4_1:
        gemm_nt_q4_1(A, B, bias, C, M, N, K);
        break;
    case CK_DT_Q5_0:
        gemm_nt_q5_0(A, B, bias, C, M, N, K);
        break;
    case CK_DT_Q5_1:
        gemm_nt_q5_1(A, B, bias, C, M, N, K);
        break;
    case CK_DT_Q8_0:
        gemm_nt_q8_0(A, B, bias, C, M, N, K);
        break;
    default:
        break;
    }
}

References C, CK_DT_FP32, CK_DT_Q4_0, CK_DT_Q4_1, CK_DT_Q4_K, CK_DT_Q5_0, CK_DT_Q5_1, CK_DT_Q6_K, CK_DT_Q8_0, gemm_blocked_serial(), gemm_nt_q4_0(), gemm_nt_q4_1(), gemm_nt_q4_k(), gemm_nt_q5_0(), gemm_nt_q5_1(), gemm_nt_q6_k(), and gemm_nt_q8_0().

Referenced by ck_attention_project_head_major_quant(), ck_layer_forward_rmsnorm_swiglu_decode_quant(), ck_mlp_swiglu_forward_quant(), ck_qkv_project_head_major_quant(), ck_qkv_project_head_major_token_quant(), and mega_fused_attention_prefill().

ck_layer_backward_rmsnorm_swiglu()

void ck_layer_backward_rmsnorm_swiglu

(

const CKLayerBackwardParams *

p

)

Definition at line 2677 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
 
    int T = p->tokens;
    int aligned_embed = p->aligned_embed_dim;
    int aligned_head = p->aligned_head_dim;
    int aligned_intermediate = p->aligned_intermediate_dim;
    int up_dim = 2 * aligned_intermediate;
    int num_threads = 1;
 
    // 1) Residual add (output = residual1 + mlp_out)
    ck_residual_add_backward(p->d_output, p->d_residual1, p->d_mlp_out, T, aligned_embed);
 
    // 2) MLP down proj backward
    fc2_backward_kernel(p->d_mlp_out,
                        p->swiglu_out,
                        p->w2,
                        p->d_swiglu_out,
                        p->d_w2,
                        p->d_b2,
                        T,
                        aligned_intermediate,
                        aligned_embed,
                        num_threads);
 
    // 3) SwiGLU backward
    swiglu_backward(p->fc1_out, p->d_swiglu_out, p->d_fc1_out, T, aligned_intermediate);
 
    // 4) MLP up proj backward
    fc1_backward_kernel(p->d_fc1_out,
                        p->ln2_out,
                        p->w1,
                        p->d_ln2_out,
                        p->d_w1,
                        p->d_b1,
                        T,
                        aligned_embed,
                        up_dim,
                        num_threads);
 
    // 5) RMSNorm (ln2) backward; reuse d_output as scratch for d_residual1_from_ln2
    rmsnorm_backward(p->d_ln2_out,
                     p->residual1,
                     p->ln2_gamma,
                     p->ln2_rstd,
                     p->d_output,
                     p->d_ln2_gamma,
                     T,
                     p->embed_dim,
                     aligned_embed);
    ck_add_inplace(p->d_residual1, p->d_output, T, aligned_embed);
 
    // 6) Residual add (residual1 = input + proj_tmp)
    ck_residual_add_backward(p->d_residual1, p->d_input, p->d_proj_tmp, T, aligned_embed);
 
    // 7) Attention projection backward
    ck_attention_project_head_major_backward(p->d_proj_tmp,
                                             p->attn_out,
                                             p->wo,
                                             p->d_attn_out,
                                             p->d_wo,
                                             p->d_bo,
                                             T,
                                             aligned_embed,
                                             p->num_heads,
                                             aligned_head);
 
    // 8) Attention backward
    attention_backward_causal_head_major_gqa(p->d_attn_out,
                                             p->q,
                                             p->k,
                                             p->v,
                                             p->scores,
                                             p->d_q,
                                             p->d_k,
                                             p->d_v,
                                             p->d_scores,
                                             p->num_heads,
                                             p->num_kv_heads,
                                             T,
                                             p->head_dim,
                                             aligned_head,
                                             p->aligned_context_window);
 
    // 9) RoPE backward (if enabled)
    if (p->rope_cos && p->rope_sin) {
        rope_backward_qk(p->d_q,
                         p->d_k,
                         p->d_q,
                         p->d_k,
                         p->rope_cos,
                         p->rope_sin,
                         p->num_heads,
                         p->num_kv_heads,
                         T,
                         p->head_dim,
                         aligned_head,
                         p->rope_pos_offset);
    }
 
    // 10) QKV projection backward (scratch uses d_proj_tmp)
    ck_qkv_project_head_major_backward(p->d_q,
                                       p->d_k,
                                       p->d_v,
                                       p->ln1_out,
                                       p->wq,
                                       p->bq,
                                       p->wk,
                                       p->bk,
                                       p->wv,
                                       p->bv,
                                       p->d_ln1_out,
                                       p->d_wq,
                                       p->d_bq,
                                       p->d_wk,
                                       p->d_bk,
                                       p->d_wv,
                                       p->d_bv,
                                       p->d_proj_tmp,
                                       T,
                                       aligned_embed,
                                       p->num_heads,
                                       p->num_kv_heads,
                                       aligned_head,
                                       num_threads);
 
    // 11) RMSNorm (ln1) backward; reuse d_ln1_out as scratch for d_input_from_ln1
    rmsnorm_backward(p->d_ln1_out,
                     p->input,
                     p->ln1_gamma,
                     p->ln1_rstd,
                     p->d_ln1_out,
                     p->d_ln1_gamma,
                     T,
                     p->embed_dim,
                     aligned_embed);
    ck_add_inplace(p->d_input, p->d_ln1_out, T, aligned_embed);
}

References CKLayerBackwardParams::aligned_context_window, CKLayerBackwardParams::aligned_embed_dim, CKLayerBackwardParams::aligned_head_dim, CKLayerBackwardParams::aligned_intermediate_dim, attention_backward_causal_head_major_gqa(), CKLayerBackwardParams::attn_out, CKLayerBackwardParams::bk, CKLayerBackwardParams::bq, CKLayerBackwardParams::bv, ck_add_inplace(), ck_attention_project_head_major_backward(), ck_qkv_project_head_major_backward(), ck_residual_add_backward(), CKLayerBackwardParams::d_attn_out, CKLayerBackwardParams::d_b1, CKLayerBackwardParams::d_b2, CKLayerBackwardParams::d_bk, CKLayerBackwardParams::d_bo, CKLayerBackwardParams::d_bq, CKLayerBackwardParams::d_bv, CKLayerBackwardParams::d_fc1_out, CKLayerBackwardParams::d_input, CKLayerBackwardParams::d_k, CKLayerBackwardParams::d_ln1_gamma, CKLayerBackwardParams::d_ln1_out, CKLayerBackwardParams::d_ln2_gamma, CKLayerBackwardParams::d_ln2_out, CKLayerBackwardParams::d_mlp_out, CKLayerBackwardParams::d_output, CKLayerBackwardParams::d_proj_tmp, CKLayerBackwardParams::d_q, CKLayerBackwardParams::d_residual1, CKLayerBackwardParams::d_scores, CKLayerBackwardParams::d_swiglu_out, CKLayerBackwardParams::d_v, CKLayerBackwardParams::d_w1, CKLayerBackwardParams::d_w2, CKLayerBackwardParams::d_wk, CKLayerBackwardParams::d_wo, CKLayerBackwardParams::d_wq, CKLayerBackwardParams::d_wv, CKLayerBackwardParams::embed_dim, fc1_backward_kernel(), CKLayerBackwardParams::fc1_out, fc2_backward_kernel(), CKLayerBackwardParams::head_dim, CKLayerBackwardParams::input, CKLayerBackwardParams::k, CKLayerBackwardParams::ln1_gamma, CKLayerBackwardParams::ln1_out, CKLayerBackwardParams::ln1_rstd, CKLayerBackwardParams::ln2_gamma, CKLayerBackwardParams::ln2_out, CKLayerBackwardParams::ln2_rstd, CKLayerBackwardParams::num_heads, CKLayerBackwardParams::num_kv_heads, CKLayerBackwardParams::q, CKLayerBackwardParams::residual1, rmsnorm_backward(), rope_backward_qk(), CKLayerBackwardParams::rope_cos, CKLayerBackwardParams::rope_pos_offset, CKLayerBackwardParams::rope_sin, CKLayerBackwardParams::scores, swiglu_backward(), CKLayerBackwardParams::swiglu_out, CKLayerBackwardParams::tokens, CKLayerBackwardParams::v, CKLayerBackwardParams::w1, CKLayerBackwardParams::w2, CKLayerBackwardParams::wk, CKLayerBackwardParams::wo, CKLayerBackwardParams::wq, and CKLayerBackwardParams::wv.

ck_layer_forward_rmsnorm_swiglu()

void ck_layer_forward_rmsnorm_swiglu

(

const CKLayerForwardParams *

p

)

Definition at line 996 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
 
    rmsnorm_forward(p->input,
                    p->ln1_gamma,
                    p->ln1_out,
                    p->ln1_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_qkv_project_head_major(p->ln1_out,
                              p->wq, p->bq,
                              p->wk, p->bk,
                              p->wv, p->bv,
                              p->q, p->k, p->v,
                              p->tokens,
                              p->tokens,
                              p->aligned_embed_dim,
                              p->num_heads,
                              p->num_kv_heads,
                              p->aligned_head_dim);
 
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(p->q,
                        p->k,
                        p->rope_cos,
                        p->rope_sin,
                        p->num_heads,
                        p->num_kv_heads,
                        p->tokens,
                        p->head_dim,
                        p->aligned_head_dim,
                        p->rope_pos_offset);
    }
 
    if (p->scores) {
        attention_forward_causal_head_major_gqa(p->q,
                                               p->k,
                                               p->v,
                                               p->scores,
                                               p->attn_out,
                                               p->num_heads,
                                               p->num_kv_heads,
                                               p->tokens,
                                               p->head_dim,
                                               p->aligned_head_dim,
                                               p->aligned_context_window);
    } else {
        attention_forward_causal_head_major_gqa_flash(p->q,
                                                     p->k,
                                                     p->v,
                                                     p->attn_out,
                                                     p->num_heads,
                                                     p->num_kv_heads,
                                                     p->tokens,
                                                     p->head_dim,
                                                     p->aligned_head_dim);
    }
 
    ck_attention_project_head_major(p->attn_out,
                                    p->wo,
                                    p->bo,
                                    p->proj_tmp,
                                    p->proj_scratch,
                                    p->tokens,
                                    p->aligned_embed_dim,
                                    p->num_heads,
                                    p->aligned_head_dim);
 
    ck_residual_add_token_major(p->input,
                                p->proj_tmp,
                                p->residual1,
                                p->tokens,
                                p->aligned_embed_dim);
 
    rmsnorm_forward(p->residual1,
                    p->ln2_gamma,
                    p->ln2_out,
                    p->ln2_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_mlp_swiglu_forward(p->ln2_out,
                          p->w1,
                          p->b1,
                          p->w2,
                          p->b2,
                          p->fc1_out,
                          p->swiglu_out,
                          p->mlp_out,
                          p->tokens,
                          p->aligned_embed_dim,
                          p->aligned_intermediate_dim);
 
    ck_residual_add_token_major(p->residual1,
                                p->mlp_out,
                                p->output,
                                p->tokens,
                                p->aligned_embed_dim);
}

References CKLayerForwardParams::aligned_context_window, CKLayerForwardParams::aligned_embed_dim, CKLayerForwardParams::aligned_head_dim, CKLayerForwardParams::aligned_intermediate_dim, attention_forward_causal_head_major_gqa(), attention_forward_causal_head_major_gqa_flash(), CKLayerForwardParams::attn_out, CKLayerForwardParams::b1, CKLayerForwardParams::b2, CKLayerForwardParams::bk, CKLayerForwardParams::bo, CKLayerForwardParams::bq, CKLayerForwardParams::bv, ck_attention_project_head_major(), ck_mlp_swiglu_forward(), ck_qkv_project_head_major(), ck_residual_add_token_major(), CKLayerForwardParams::embed_dim, CKLayerForwardParams::eps, CKLayerForwardParams::fc1_out, CKLayerForwardParams::head_dim, CKLayerForwardParams::input, CKLayerForwardParams::k, CKLayerForwardParams::ln1_gamma, CKLayerForwardParams::ln1_out, CKLayerForwardParams::ln1_rstd, CKLayerForwardParams::ln2_gamma, CKLayerForwardParams::ln2_out, CKLayerForwardParams::ln2_rstd, CKLayerForwardParams::mlp_out, CKLayerForwardParams::num_heads, CKLayerForwardParams::num_kv_heads, CKLayerForwardParams::output, CKLayerForwardParams::proj_scratch, CKLayerForwardParams::proj_tmp, CKLayerForwardParams::q, CKLayerForwardParams::residual1, rmsnorm_forward(), CKLayerForwardParams::rope_cos, rope_forward_qk(), CKLayerForwardParams::rope_pos_offset, CKLayerForwardParams::rope_sin, CKLayerForwardParams::scores, CKLayerForwardParams::swiglu_out, CKLayerForwardParams::tokens, CKLayerForwardParams::v, CKLayerForwardParams::w1, CKLayerForwardParams::w2, CKLayerForwardParams::wk, CKLayerForwardParams::wo, CKLayerForwardParams::wq, and CKLayerForwardParams::wv.

ck_layer_forward_rmsnorm_swiglu_decode()

void ck_layer_forward_rmsnorm_swiglu_decode	(	const CKLayerForwardParams *	p,
		int	token_index,
		int	cache_capacity
	)

Definition at line 1289 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
    if (!p->input || !p->ln1_gamma || !p->ln2_gamma || !p->ln1_out || !p->ln2_out ||
        !p->wq || !p->wk || !p->wv || !p->wo || !p->w1 || !p->w2 ||
        !p->k || !p->v ||
        !p->proj_tmp || !p->residual1 || !p->fc1_out || !p->swiglu_out || !p->mlp_out || !p->output) {
        return;
    }
    if (token_index < 0 || cache_capacity <= 0 || token_index >= cache_capacity) {
        return;
    }
    if (p->num_heads <= 0 || p->num_kv_heads <= 0 || p->aligned_head_dim <= 0) {
        return;
    }
 
    const int D = p->embed_dim;
    const int aligned_D = p->aligned_embed_dim;
    const int H = p->num_heads;
    const int H_kv = p->num_kv_heads;
    const int hd = p->head_dim;
    const int ad = p->aligned_head_dim;
    const int aligned_intermediate = p->aligned_intermediate_dim;
 
    /* Decode buffers are single-token; token_index only applies to KV cache. */
    const size_t token_slot = 0;
    const float *input_row = p->input + token_slot * (size_t)aligned_D;
    float *ln1_row = p->ln1_out + token_slot * (size_t)aligned_D;
    float *ln2_row = p->ln2_out + token_slot * (size_t)aligned_D;
    float *proj_row = p->proj_tmp + token_slot * (size_t)aligned_D;
    float *residual_row = p->residual1 + token_slot * (size_t)aligned_D;
    float *mlp_row = p->mlp_out + token_slot * (size_t)aligned_D;
    float *out_row = p->output + token_slot * (size_t)aligned_D;
 
    float ln1_rstd_tmp = 0.0f;
    float ln2_rstd_tmp = 0.0f;
    float *ln1_rstd = p->ln1_rstd ? (p->ln1_rstd + token_slot) : &ln1_rstd_tmp;
    float *ln2_rstd = p->ln2_rstd ? (p->ln2_rstd + token_slot) : &ln2_rstd_tmp;
 
    // Scratch for a single token in head-major layout: [head, aligned_head_dim].
    size_t q_elems = (size_t)H * (size_t)ad;
    size_t kv_elems = (size_t)H_kv * (size_t)ad;
    float q_token[q_elems];
    float k_token[kv_elems];
    float v_token[kv_elems];
    float attn_token[q_elems];
 
    // LN1 / RMSNorm.
    rmsnorm_forward(input_row,
                    p->ln1_gamma,
                    ln1_row,
                    ln1_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    // Project Q/K/V for the new token.
    ck_qkv_project_head_major_token(ln1_row,
                                    p->wq, p->bq,
                                    p->wk, p->bk,
                                    p->wv, p->bv,
                                    q_token, k_token, v_token,
                                    aligned_D,
                                    H,
                                    H_kv,
                                    ad);
 
    // RoPE for the new token at absolute position `p->rope_pos_offset`.
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(q_token,
                        k_token,
                        p->rope_cos,
                        p->rope_sin,
                        H,
                        H_kv,
                        /*num_tokens=*/1,
                        hd,
                        ad,
                        p->rope_pos_offset);
    }
 
    // Update KV cache (stores k/v for this token and clears padded lanes).
    kv_cache_write_head_major(k_token,
                              v_token,
                              p->k,
                              p->v,
                              H_kv,
                              token_index,
                              cache_capacity,
                              hd,
                              ad);
 
    // Decode attention for this token using the KV cache.
    ck_attention_flash_decode_wrapper(q_token,
                                   p->k,
                                   p->v,
                                   attn_token,
                                   H,
                                   H_kv,
                                   /*kv_tokens=*/token_index + 1,
                                   cache_capacity,
                                   hd,
                                   ad);
 
    // Output projection (Wo) into token-major buffer (decode-specialized).
    ck_attention_project_head_major_decode_token(attn_token,
                                                 p->wo,
                                                 p->bo,
                                                 proj_row,
                                                 D,
                                                 aligned_D,
                                                 H,
                                                 ad);
 
    // Residual + LN2 / RMSNorm.
    ck_residual_add_token_major(input_row,
                                proj_row,
                                residual_row,
                                /*tokens=*/1,
                                aligned_D);
 
    rmsnorm_forward(residual_row,
                    p->ln2_gamma,
                    ln2_row,
                    ln2_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    // MLP block for this token.
    int up_dim = 2 * aligned_intermediate;
    float *fc1_row = p->fc1_out + token_slot * (size_t)up_dim;
    float *swiglu_row = p->swiglu_out + token_slot * (size_t)aligned_intermediate;
 
    ck_mlp_swiglu_forward(ln2_row,
                          p->w1,
                          p->b1,
                          p->w2,
                          p->b2,
                          fc1_row,
                          swiglu_row,
                          mlp_row,
                          /*tokens=*/1,
                          aligned_D,
                          aligned_intermediate);
 
    // Final residual.
    ck_residual_add_token_major(residual_row,
                                mlp_row,
                                out_row,
                                /*tokens=*/1,
                                aligned_D);
}

References CKLayerForwardParams::aligned_embed_dim, CKLayerForwardParams::aligned_head_dim, CKLayerForwardParams::aligned_intermediate_dim, CKLayerForwardParams::b1, CKLayerForwardParams::b2, CKLayerForwardParams::bk, CKLayerForwardParams::bo, CKLayerForwardParams::bq, CKLayerForwardParams::bv, ck_attention_flash_decode_wrapper(), ck_attention_project_head_major_decode_token(), ck_mlp_swiglu_forward(), ck_qkv_project_head_major_token(), ck_residual_add_token_major(), CKLayerForwardParams::embed_dim, CKLayerForwardParams::eps, CKLayerForwardParams::fc1_out, CKLayerForwardParams::head_dim, CKLayerForwardParams::input, CKLayerForwardParams::k, kv_cache_write_head_major(), CKLayerForwardParams::ln1_gamma, CKLayerForwardParams::ln1_out, CKLayerForwardParams::ln1_rstd, CKLayerForwardParams::ln2_gamma, CKLayerForwardParams::ln2_out, CKLayerForwardParams::ln2_rstd, CKLayerForwardParams::mlp_out, CKLayerForwardParams::num_heads, CKLayerForwardParams::num_kv_heads, CKLayerForwardParams::output, CKLayerForwardParams::proj_tmp, CKLayerForwardParams::residual1, rmsnorm_forward(), CKLayerForwardParams::rope_cos, rope_forward_qk(), CKLayerForwardParams::rope_pos_offset, CKLayerForwardParams::rope_sin, CKLayerForwardParams::swiglu_out, CKLayerForwardParams::v, CKLayerForwardParams::w1, CKLayerForwardParams::w2, CKLayerForwardParams::wk, CKLayerForwardParams::wo, CKLayerForwardParams::wq, and CKLayerForwardParams::wv.

ck_layer_forward_rmsnorm_swiglu_decode_fused()

void ck_layer_forward_rmsnorm_swiglu_decode_fused	(	const CKLayerForwardParams *	p,
		int	token_index,
		int	cache_capacity
	)

Definition at line 1449 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
    if (!p->input || !p->ln1_gamma || !p->ln2_gamma || !p->ln1_out || !p->ln2_out ||
        !p->wq || !p->wk || !p->wv || !p->wo || !p->w1 || !p->w2 ||
        !p->k || !p->v || !p->swiglu_out ||
        !p->proj_tmp || !p->residual1 || !p->mlp_out || !p->output) {
        return;
    }
    if (token_index < 0 || cache_capacity <= 0 || token_index >= cache_capacity) {
        return;
    }
    if (p->num_heads <= 0 || p->num_kv_heads <= 0 || p->aligned_head_dim <= 0) {
        return;
    }
 
    const int D = p->embed_dim;
    const int aligned_D = p->aligned_embed_dim;
    const int H = p->num_heads;
    const int H_kv = p->num_kv_heads;
    const int hd = p->head_dim;
    const int ad = p->aligned_head_dim;
    const int aligned_intermediate = p->aligned_intermediate_dim;
 
    /* Decode buffers are single-token; token_index only applies to KV cache. */
    const size_t token_slot = 0;
    const float *input_row = p->input + token_slot * (size_t)aligned_D;
    float *ln1_row = p->ln1_out + token_slot * (size_t)aligned_D;
    float *ln2_row = p->ln2_out + token_slot * (size_t)aligned_D;
    float *proj_row = p->proj_tmp + token_slot * (size_t)aligned_D;
    float *residual_row = p->residual1 + token_slot * (size_t)aligned_D;
    float *swiglu_row = p->swiglu_out + token_slot * (size_t)aligned_intermediate;
    float *mlp_row = p->mlp_out + token_slot * (size_t)aligned_D;
    float *out_row = p->output + token_slot * (size_t)aligned_D;
 
    float ln1_rstd_tmp = 0.0f;
    float ln2_rstd_tmp = 0.0f;
    float *ln1_rstd = p->ln1_rstd ? (p->ln1_rstd + token_slot) : &ln1_rstd_tmp;
    float *ln2_rstd = p->ln2_rstd ? (p->ln2_rstd + token_slot) : &ln2_rstd_tmp;
 
    // Scratch for a single token in head-major layout: [head, aligned_head_dim].
    size_t q_elems = (size_t)H * (size_t)ad;
    size_t kv_elems = (size_t)H_kv * (size_t)ad;
    float q_token[q_elems];
    float k_token[kv_elems];
    float v_token[kv_elems];
    float attn_token[q_elems];
 
    // LN1 / RMSNorm.
    rmsnorm_forward(input_row,
                    p->ln1_gamma,
                    ln1_row,
                    ln1_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    // Project Q/K/V for the new token.
    ck_qkv_project_head_major_token(ln1_row,
                                    p->wq, p->bq,
                                    p->wk, p->bk,
                                    p->wv, p->bv,
                                    q_token, k_token, v_token,
                                    aligned_D,
                                    H,
                                    H_kv,
                                    ad);
 
    // RoPE for the new token at absolute position `p->rope_pos_offset`.
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(q_token,
                        k_token,
                        p->rope_cos,
                        p->rope_sin,
                        H,
                        H_kv,
                        /*num_tokens=*/1,
                        hd,
                        ad,
                        p->rope_pos_offset);
    }
 
    // Update KV cache (stores k/v for this token and clears padded lanes).
    kv_cache_write_head_major(k_token,
                              v_token,
                              p->k,
                              p->v,
                              H_kv,
                              token_index,
                              cache_capacity,
                              hd,
                              ad);
 
    // Decode attention for this token using the KV cache.
    ck_attention_flash_decode_wrapper(q_token,
                                   p->k,
                                   p->v,
                                   attn_token,
                                   H,
                                   H_kv,
                                   /*kv_tokens=*/token_index + 1,
                                   cache_capacity,
                                   hd,
                                   ad);
 
    // Output projection (Wo) into token-major buffer (decode-specialized).
    ck_attention_project_head_major_decode_token(attn_token,
                                                 p->wo,
                                                 p->bo,
                                                 proj_row,
                                                 D,
                                                 aligned_D,
                                                 H,
                                                 ad);
 
    // Residual + LN2 / RMSNorm.
    ck_residual_add_token_major(input_row,
                                proj_row,
                                residual_row,
                                /*tokens=*/1,
                                aligned_D);
 
    rmsnorm_forward(residual_row,
                    p->ln2_gamma,
                    ln2_row,
                    ln2_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    // MLP block for this token (fully fused - all 3 projections in one pass).
    // Eliminates DRAM round-trip for swiglu intermediate values.
    ck_mlp_swiglu_forward_fully_fused_token(ln2_row,
                                             p->w1,
                                             p->b1,
                                             p->w2,
                                             p->b2,
                                             mlp_row,
                                             aligned_D,
                                             aligned_intermediate);
 
    // Final residual.
    ck_residual_add_token_major(residual_row,
                                mlp_row,
                                out_row,
                                /*tokens=*/1,
                                aligned_D);
}

References CKLayerForwardParams::aligned_embed_dim, CKLayerForwardParams::aligned_head_dim, CKLayerForwardParams::aligned_intermediate_dim, CKLayerForwardParams::b1, CKLayerForwardParams::b2, CKLayerForwardParams::bk, CKLayerForwardParams::bo, CKLayerForwardParams::bq, CKLayerForwardParams::bv, ck_attention_flash_decode_wrapper(), ck_attention_project_head_major_decode_token(), ck_mlp_swiglu_forward_fully_fused_token(), ck_qkv_project_head_major_token(), ck_residual_add_token_major(), CKLayerForwardParams::embed_dim, CKLayerForwardParams::eps, CKLayerForwardParams::head_dim, CKLayerForwardParams::input, CKLayerForwardParams::k, kv_cache_write_head_major(), CKLayerForwardParams::ln1_gamma, CKLayerForwardParams::ln1_out, CKLayerForwardParams::ln1_rstd, CKLayerForwardParams::ln2_gamma, CKLayerForwardParams::ln2_out, CKLayerForwardParams::ln2_rstd, CKLayerForwardParams::mlp_out, CKLayerForwardParams::num_heads, CKLayerForwardParams::num_kv_heads, CKLayerForwardParams::output, CKLayerForwardParams::proj_tmp, CKLayerForwardParams::residual1, rmsnorm_forward(), CKLayerForwardParams::rope_cos, rope_forward_qk(), CKLayerForwardParams::rope_pos_offset, CKLayerForwardParams::rope_sin, CKLayerForwardParams::swiglu_out, CKLayerForwardParams::v, CKLayerForwardParams::w1, CKLayerForwardParams::w2, CKLayerForwardParams::wk, CKLayerForwardParams::wo, CKLayerForwardParams::wq, and CKLayerForwardParams::wv.

ck_layer_forward_rmsnorm_swiglu_decode_fused_attn()

void ck_layer_forward_rmsnorm_swiglu_decode_fused_attn	(	const CKLayerForwardParams *	p,
		int	token_index,
		int	cache_capacity
	)

Definition at line 343 of file attention_decode_fused.c.

{
    ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl(p,
                                                           token_index,
                                                           cache_capacity,
                                                           /*fuse_mlp=*/0);
}

References ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl().

ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_mlp()

void ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_mlp	(	const CKLayerForwardParams *	p,
		int	token_index,
		int	cache_capacity
	)

Definition at line 353 of file attention_decode_fused.c.

{
    ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl(p,
                                                           token_index,
                                                           cache_capacity,
                                                           /*fuse_mlp=*/1);
}

References ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl().

ck_layer_forward_rmsnorm_swiglu_decode_q4_k()

void ck_layer_forward_rmsnorm_swiglu_decode_q4_k	(	const CKLayerForwardParamsQ4K *	p,
		int	token_index,
		int	cache_capacity
	)

Definition at line 2117 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
    if (!p->input || !p->ln1_gamma || !p->ln2_gamma || !p->ln1_out || !p->ln2_out ||
        !p->wq || !p->wk || !p->wv || !p->wo || !p->w1 || !p->w2 ||
        !p->k || !p->v ||
        !p->proj_tmp || !p->residual1 || !p->fc1_out || !p->swiglu_out || !p->mlp_out || !p->output) {
        return;
    }
    if (token_index < 0 || cache_capacity <= 0 || token_index >= cache_capacity) {
        return;
    }
 
    const int D = p->embed_dim;
    const int aligned_D = p->aligned_embed_dim;
    const int H = p->num_heads;
    const int H_kv = p->num_kv_heads;
    const int hd = p->head_dim;
    const int ad = p->aligned_head_dim;
    const int aligned_intermediate = p->aligned_intermediate_dim;
    const int K_concat = H * ad;
 
    /* Decode buffers are single-token; token_index only applies to KV cache. */
    const size_t token_slot = 0;
    const float *input_row = p->input + token_slot * (size_t)aligned_D;
    float *ln1_row = p->ln1_out + token_slot * (size_t)aligned_D;
    float *ln2_row = p->ln2_out + token_slot * (size_t)aligned_D;
    float *proj_row = p->proj_tmp + token_slot * (size_t)aligned_D;
    float *residual_row = p->residual1 + token_slot * (size_t)aligned_D;
    float *mlp_row = p->mlp_out + token_slot * (size_t)aligned_D;
    float *out_row = p->output + token_slot * (size_t)aligned_D;
 
    float ln1_rstd_tmp = 0.0f;
    float ln2_rstd_tmp = 0.0f;
    float *ln1_rstd = p->ln1_rstd ? (p->ln1_rstd + token_slot) : &ln1_rstd_tmp;
    float *ln2_rstd = p->ln2_rstd ? (p->ln2_rstd + token_slot) : &ln2_rstd_tmp;
 
    /* Scratch for a single token in head-major layout: [head, aligned_head_dim]. */
    size_t q_elems = (size_t)H * (size_t)ad;
    size_t kv_elems = (size_t)H_kv * (size_t)ad;
    float q_token[q_elems];
    float k_token[kv_elems];
    float v_token[kv_elems];
    float attn_token[q_elems];
 
    /* LN1 / RMSNorm. */
    ck_debug_check_buffer("input_row", input_row, aligned_D);
    rmsnorm_forward(input_row,
                    p->ln1_gamma,
                    ln1_row,
                    ln1_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
    ck_debug_check_buffer("ln1_out (after rmsnorm)", ln1_row, aligned_D);
 
    if (ck_q8k_activations_enabled()) {
        if ((aligned_D % QK_K) == 0 && (aligned_intermediate % QK_K) == 0) {
        const int q8_blocks_embed = aligned_D / QK_K;
        const int q8_blocks_inter = aligned_intermediate / QK_K;
        const int q8_blocks_max = (q8_blocks_embed > q8_blocks_inter) ? q8_blocks_embed : q8_blocks_inter;
        block_q8_K q8_buf[q8_blocks_max];
 
        /* Project Q/K/V with Q8_K activations. */
        quantize_row_q8_k(ln1_row, q8_buf, aligned_D);
        ck_debug_check_q8k("q8_buf (after quantize)", q8_buf, q8_blocks_embed);
        ck_debug_check_q4k_weights("wq weights", p->wq, (aligned_D / QK_K) * (H * ad));
        ck_qkv_project_head_major_token_q4_k_q8_k(q8_buf,
                                                  p->wq, p->bq,
                                                  p->wk, p->bk,
                                                  p->wv, p->bv,
                                                  q_token, k_token, v_token,
                                                  aligned_D,
                                                  H,
                                                  H_kv,
                                                  ad);
        ck_debug_check_buffer("q_token (after QKV proj)", q_token, (int)q_elems);
        ck_debug_check_buffer("k_token (after QKV proj)", k_token, (int)kv_elems);
        ck_debug_check_buffer("v_token (after QKV proj)", v_token, (int)kv_elems);
 
        /* RoPE for the new token at absolute position `p->rope_pos_offset`. */
        if (p->rope_cos && p->rope_sin) {
            rope_forward_qk(q_token,
                            k_token,
                            p->rope_cos,
                            p->rope_sin,
                            H,
                            H_kv,
                            /*num_tokens=*/1,
                            hd,
                            ad,
                            p->rope_pos_offset);
        }
 
        /* Update KV cache. */
        kv_cache_write_head_major(k_token,
                                  v_token,
                                  p->k,
                                  p->v,
                                  H_kv,
                                  token_index,
                                  cache_capacity,
                                  hd,
                                  ad);
 
        /* Decode attention for this token using the KV cache. */
        attention_forward_decode_head_major_gqa_regular(q_token,
                                                      p->k,
                                                      p->v,
                                                      attn_token,
                                                      H,
                                                      H_kv,
                                                      /*kv_tokens=*/token_index + 1,
                                                      cache_capacity,
                                                      hd,
                                                      ad);
        ck_debug_check_buffer("attn_token (after attention)", attn_token, (int)q_elems);
 
        /* Quantized output projection (Wo) with Q8_K activations. */
        quantize_row_q8_k(attn_token, q8_buf, aligned_D);
        gemm_nt_q4_k_q8_k(q8_buf,
                          p->wo,
                          p->bo,
                          proj_row,
                          /*M=*/1,
                          aligned_D,
                          /*K=*/K_concat);
        ck_debug_check_buffer("proj_row (after Wo proj)", proj_row, aligned_D);
 
        for (int j = D; j < aligned_D; ++j) {
            proj_row[j] = 0.0f;
        }
 
        /* Residual + LN2 / RMSNorm. */
        ck_residual_add_token_major(input_row,
                                    proj_row,
                                    residual_row,
                                    /*tokens=*/1,
                                    aligned_D);
 
        rmsnorm_forward(residual_row,
                        p->ln2_gamma,
                        ln2_row,
                        ln2_rstd,
                        /*tokens=*/1,
                        D,
                        aligned_D,
                        p->eps);
 
        /* MLP block for this token (Q8_K activations). */
        int up_dim = 2 * aligned_intermediate;
        float *fc1_row = p->fc1_out + token_slot * (size_t)up_dim;
        float *swiglu_row = p->swiglu_out + token_slot * (size_t)aligned_intermediate;
 
        ck_mlp_swiglu_forward_q4_k_q8_k(ln2_row,
                                        p->w1,
                                        p->b1,
                                        p->w2,
                                        p->b2,
                                        fc1_row,
                                        swiglu_row,
                                        mlp_row,
                                        aligned_D,
                                        aligned_intermediate);
        ck_debug_check_buffer("mlp_row (after MLP)", mlp_row, aligned_D);
 
        /* Final residual. */
        ck_residual_add_token_major(residual_row,
                                    mlp_row,
                                    out_row,
                                    /*tokens=*/1,
                                    aligned_D);
        ck_debug_check_buffer("out_row (final output)", out_row, aligned_D);
        return;
        }
    }
 
    /* Project Q/K/V for the new token (Q4_K weights). */
    ck_qkv_project_head_major_token_q4_k(ln1_row,
                                         p->wq, p->bq,
                                         p->wk, p->bk,
                                         p->wv, p->bv,
                                         q_token, k_token, v_token,
                                         aligned_D,
                                         H,
                                         H_kv,
                                         ad);
 
    /* RoPE for the new token at absolute position `p->rope_pos_offset`. */
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(q_token,
                        k_token,
                        p->rope_cos,
                        p->rope_sin,
                        H,
                        H_kv,
                        /*num_tokens=*/1,
                        hd,
                        ad,
                        p->rope_pos_offset);
    }
 
    /* Update KV cache. */
    kv_cache_write_head_major(k_token,
                              v_token,
                              p->k,
                              p->v,
                              H_kv,
                              token_index,
                              cache_capacity,
                              hd,
                              ad);
 
    /* Decode attention for this token using the KV cache. */
    ck_attention_flash_decode_wrapper(q_token,
                                   p->k,
                                   p->v,
                                   attn_token,
                                   H,
                                   H_kv,
                                   /*kv_tokens=*/token_index + 1,
                                   cache_capacity,
                                   hd,
                                   ad);
 
    /* Quantized output projection: Wo is stored as a flattened Q4_K matrix. */
    gemm_nt_q4_k(attn_token,
                 p->wo,
                 p->bo,
                 proj_row,
                 /*M=*/1,
                 aligned_D,
                 /*K=*/K_concat);
 
    for (int j = D; j < aligned_D; ++j) {
        proj_row[j] = 0.0f;
    }
 
    /* Residual + LN2 / RMSNorm. */
    ck_residual_add_token_major(input_row,
                                proj_row,
                                residual_row,
                                /*tokens=*/1,
                                aligned_D);
 
    rmsnorm_forward(residual_row,
                    p->ln2_gamma,
                    ln2_row,
                    ln2_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    /* MLP block for this token. */
    int up_dim = 2 * aligned_intermediate;
    float *fc1_row = p->fc1_out + token_slot * (size_t)up_dim;
    float *swiglu_row = p->swiglu_out + token_slot * (size_t)aligned_intermediate;
 
    ck_mlp_swiglu_forward_q4_k(ln2_row,
                               p->w1,
                               p->b1,
                               p->w2,
                               p->b2,
                               fc1_row,
                               swiglu_row,
                               mlp_row,
                               /*tokens=*/1,
                               aligned_D,
                               aligned_intermediate);
 
    /* Final residual. */
    ck_residual_add_token_major(residual_row,
                                mlp_row,
                                out_row,
                                /*tokens=*/1,
                                aligned_D);
}

References CKLayerForwardParamsQ4K::aligned_embed_dim, CKLayerForwardParamsQ4K::aligned_head_dim, CKLayerForwardParamsQ4K::aligned_intermediate_dim, attention_forward_decode_head_major_gqa_regular(), CKLayerForwardParamsQ4K::b1, CKLayerForwardParamsQ4K::b2, CKLayerForwardParamsQ4K::bk, CKLayerForwardParamsQ4K::bo, CKLayerForwardParamsQ4K::bq, CKLayerForwardParamsQ4K::bv, ck_attention_flash_decode_wrapper(), ck_debug_check_buffer(), ck_debug_check_q4k_weights(), ck_debug_check_q8k(), ck_mlp_swiglu_forward_q4_k(), ck_mlp_swiglu_forward_q4_k_q8_k(), ck_q8k_activations_enabled(), ck_qkv_project_head_major_token_q4_k(), ck_qkv_project_head_major_token_q4_k_q8_k(), ck_residual_add_token_major(), CKLayerForwardParamsQ4K::embed_dim, CKLayerForwardParamsQ4K::eps, CKLayerForwardParamsQ4K::fc1_out, gemm_nt_q4_k(), gemm_nt_q4_k_q8_k(), CKLayerForwardParamsQ4K::head_dim, CKLayerForwardParamsQ4K::input, CKLayerForwardParamsQ4K::k, kv_cache_write_head_major(), CKLayerForwardParamsQ4K::ln1_gamma, CKLayerForwardParamsQ4K::ln1_out, CKLayerForwardParamsQ4K::ln1_rstd, CKLayerForwardParamsQ4K::ln2_gamma, CKLayerForwardParamsQ4K::ln2_out, CKLayerForwardParamsQ4K::ln2_rstd, CKLayerForwardParamsQ4K::mlp_out, CKLayerForwardParamsQ4K::num_heads, CKLayerForwardParamsQ4K::num_kv_heads, CKLayerForwardParamsQ4K::output, CKLayerForwardParamsQ4K::proj_tmp, QK_K, quantize_row_q8_k(), CKLayerForwardParamsQ4K::residual1, rmsnorm_forward(), CKLayerForwardParamsQ4K::rope_cos, rope_forward_qk(), CKLayerForwardParamsQ4K::rope_pos_offset, CKLayerForwardParamsQ4K::rope_sin, CKLayerForwardParamsQ4K::swiglu_out, CKLayerForwardParamsQ4K::v, CKLayerForwardParamsQ4K::w1, CKLayerForwardParamsQ4K::w2, CKLayerForwardParamsQ4K::wk, CKLayerForwardParamsQ4K::wo, CKLayerForwardParamsQ4K::wq, and CKLayerForwardParamsQ4K::wv.

ck_layer_forward_rmsnorm_swiglu_decode_quant()

void ck_layer_forward_rmsnorm_swiglu_decode_quant	(	const CKLayerForwardParamsQ4K *	p,
		int	token_index,
		int	cache_capacity
	)

Definition at line 2512 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
    if (!p->input || !p->ln1_gamma || !p->ln2_gamma || !p->ln1_out || !p->ln2_out ||
        !p->wq || !p->wk || !p->wv || !p->wo || !p->w1 || !p->w2 ||
        !p->k || !p->v ||
        !p->proj_tmp || !p->proj_scratch || !p->residual1 || !p->fc1_out || !p->swiglu_out || !p->mlp_out || !p->output) {
        return;
    }
    if (token_index < 0 || cache_capacity <= 0 || token_index >= cache_capacity) {
        return;
    }
 
    const int D = p->embed_dim;
    const int aligned_D = p->aligned_embed_dim;
    const int H = p->num_heads;
    const int H_kv = p->num_kv_heads;
    const int hd = p->head_dim;
    const int ad = p->aligned_head_dim;
    const int aligned_intermediate = p->aligned_intermediate_dim;
    const int K_concat = H * ad;
 
    /* Decode buffers are single-token; token_index only applies to KV cache. */
    const size_t token_slot = 0;
    const float *input_row = p->input + token_slot * (size_t)aligned_D;
    float *ln1_row = p->ln1_out + token_slot * (size_t)aligned_D;
    float *ln2_row = p->ln2_out + token_slot * (size_t)aligned_D;
    float *proj_row = p->proj_tmp + token_slot * (size_t)aligned_D;
    float *residual_row = p->residual1 + token_slot * (size_t)aligned_D;
    float *mlp_row = p->mlp_out + token_slot * (size_t)aligned_D;
    float *out_row = p->output + token_slot * (size_t)aligned_D;
 
    float ln1_rstd_tmp = 0.0f;
    float ln2_rstd_tmp = 0.0f;
    float *ln1_rstd = p->ln1_rstd ? (p->ln1_rstd + token_slot) : &ln1_rstd_tmp;
    float *ln2_rstd = p->ln2_rstd ? (p->ln2_rstd + token_slot) : &ln2_rstd_tmp;
 
    size_t q_elems = (size_t)H * (size_t)ad;
    size_t kv_elems = (size_t)H_kv * (size_t)ad;
    float q_token[q_elems];
    float k_token[kv_elems];
    float v_token[kv_elems];
    float attn_token[q_elems];
 
    rmsnorm_forward(input_row,
                    p->ln1_gamma,
                    ln1_row,
                    ln1_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    ck_qkv_project_head_major_token_quant(ln1_row,
                                          p->wq, p->bq, p->wq_dtype,
                                          p->wk, p->bk, p->wk_dtype,
                                          p->wv, p->bv, p->wv_dtype,
                                          q_token, k_token, v_token,
                                          aligned_D,
                                          H,
                                          H_kv,
                                          ad);
 
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(q_token,
                        k_token,
                        p->rope_cos,
                        p->rope_sin,
                        H,
                        H_kv,
                        /*num_tokens=*/1,
                        hd,
                        ad,
                        p->rope_pos_offset);
    }
 
    kv_cache_write_head_major(k_token,
                              v_token,
                              p->k,
                              p->v,
                              H_kv,
                              token_index,
                              cache_capacity,
                              hd,
                              ad);
 
    ck_attention_flash_decode_wrapper(q_token,
                                   p->k,
                                   p->v,
                                   attn_token,
                                   H,
                                   H_kv,
                                   /*kv_tokens=*/token_index + 1,
                                   cache_capacity,
                                   hd,
                                   ad);
 
    if (p->wo_dtype == CK_DT_FP32) {
        ck_attention_project_head_major_decode_token(attn_token,
                                                     (const float *)p->wo,
                                                     p->bo,
                                                     proj_row,
                                                     D,
                                                     aligned_D,
                                                     H,
                                                     ad);
    } else {
        /* Quantized attention output projection - handle all quant types */
        ck_gemm_nt_quant(attn_token,
                         p->wo,
                         p->bo,
                         proj_row,
                         /*M=*/1,
                         aligned_D,
                         /*K=*/K_concat,
                         p->wo_dtype);
        for (int j = D; j < aligned_D; ++j) {
            proj_row[j] = 0.0f;
        }
    }
 
    ck_residual_add_token_major(input_row,
                                proj_row,
                                residual_row,
                                /*tokens=*/1,
                                aligned_D);
 
    rmsnorm_forward(residual_row,
                    p->ln2_gamma,
                    ln2_row,
                    ln2_rstd,
                    /*tokens=*/1,
                    D,
                    aligned_D,
                    p->eps);
 
    int up_dim = 2 * aligned_intermediate;
    float *fc1_row = p->fc1_out + token_slot * (size_t)up_dim;
    float *swiglu_row = p->swiglu_out + token_slot * (size_t)aligned_intermediate;
 
    ck_mlp_swiglu_forward_quant(ln2_row,
                                p->w1,
                                p->b1,
                                p->w1_dtype,
                                p->w2,
                                p->b2,
                                p->w2_dtype,
                                fc1_row,
                                swiglu_row,
                                mlp_row,
                                /*tokens=*/1,
                                aligned_D,
                                aligned_intermediate);
 
    ck_residual_add_token_major(residual_row,
                                mlp_row,
                                out_row,
                                /*tokens=*/1,
                                aligned_D);
}

References CKLayerForwardParamsQ4K::aligned_embed_dim, CKLayerForwardParamsQ4K::aligned_head_dim, CKLayerForwardParamsQ4K::aligned_intermediate_dim, CKLayerForwardParamsQ4K::b1, CKLayerForwardParamsQ4K::b2, CKLayerForwardParamsQ4K::bk, CKLayerForwardParamsQ4K::bo, CKLayerForwardParamsQ4K::bq, CKLayerForwardParamsQ4K::bv, ck_attention_flash_decode_wrapper(), ck_attention_project_head_major_decode_token(), CK_DT_FP32, ck_gemm_nt_quant(), ck_mlp_swiglu_forward_quant(), ck_qkv_project_head_major_token_quant(), ck_residual_add_token_major(), CKLayerForwardParamsQ4K::embed_dim, CKLayerForwardParamsQ4K::eps, CKLayerForwardParamsQ4K::fc1_out, CKLayerForwardParamsQ4K::head_dim, CKLayerForwardParamsQ4K::input, CKLayerForwardParamsQ4K::k, kv_cache_write_head_major(), CKLayerForwardParamsQ4K::ln1_gamma, CKLayerForwardParamsQ4K::ln1_out, CKLayerForwardParamsQ4K::ln1_rstd, CKLayerForwardParamsQ4K::ln2_gamma, CKLayerForwardParamsQ4K::ln2_out, CKLayerForwardParamsQ4K::ln2_rstd, CKLayerForwardParamsQ4K::mlp_out, CKLayerForwardParamsQ4K::num_heads, CKLayerForwardParamsQ4K::num_kv_heads, CKLayerForwardParamsQ4K::output, CKLayerForwardParamsQ4K::proj_scratch, CKLayerForwardParamsQ4K::proj_tmp, CKLayerForwardParamsQ4K::residual1, rmsnorm_forward(), CKLayerForwardParamsQ4K::rope_cos, rope_forward_qk(), CKLayerForwardParamsQ4K::rope_pos_offset, CKLayerForwardParamsQ4K::rope_sin, CKLayerForwardParamsQ4K::swiglu_out, CKLayerForwardParamsQ4K::v, CKLayerForwardParamsQ4K::w1, CKLayerForwardParamsQ4K::w1_dtype, CKLayerForwardParamsQ4K::w2, CKLayerForwardParamsQ4K::w2_dtype, CKLayerForwardParamsQ4K::wk, CKLayerForwardParamsQ4K::wk_dtype, CKLayerForwardParamsQ4K::wo, CKLayerForwardParamsQ4K::wo_dtype, CKLayerForwardParamsQ4K::wq, CKLayerForwardParamsQ4K::wq_dtype, CKLayerForwardParamsQ4K::wv, and CKLayerForwardParamsQ4K::wv_dtype.

ck_layer_forward_rmsnorm_swiglu_q4_k()

void ck_layer_forward_rmsnorm_swiglu_q4_k

(

const CKLayerForwardParamsQ4K *

p

)

Definition at line 1910 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
 
    const int aligned_D = p->aligned_embed_dim;
    const int aligned_intermediate = p->aligned_intermediate_dim;
 
    rmsnorm_forward(p->input,
                    p->ln1_gamma,
                    p->ln1_out,
                    p->ln1_rstd,
                    p->tokens,
                    p->embed_dim,
                    aligned_D,
                    p->eps);
 
    if (ck_q8k_activations_enabled()) {
        if ((aligned_D % QK_K) == 0 && (aligned_intermediate % QK_K) == 0) {
            ck_qkv_project_head_major_q4_k_q8_k(p->ln1_out,
                                                p->wq, p->bq,
                                                p->wk, p->bk,
                                                p->wv, p->bv,
                                                p->q, p->k, p->v,
                                                p->tokens,
                                                p->tokens,
                                                aligned_D,
                                                p->num_heads,
                                                p->num_kv_heads,
                                                p->aligned_head_dim);
 
            if (p->rope_cos && p->rope_sin) {
                rope_forward_qk(p->q,
                                p->k,
                                p->rope_cos,
                                p->rope_sin,
                                p->num_heads,
                                p->num_kv_heads,
                                p->tokens,
                                p->head_dim,
                                p->aligned_head_dim,
                                p->rope_pos_offset);
            }
 
            if (p->scores) {
                attention_forward_causal_head_major_gqa(p->q,
                                                       p->k,
                                                       p->v,
                                                       p->scores,
                                                       p->attn_out,
                                                       p->num_heads,
                                                       p->num_kv_heads,
                                                       p->tokens,
                                                       p->head_dim,
                                                       p->aligned_head_dim,
                                                       p->aligned_context_window);
            } else {
                attention_forward_causal_head_major_gqa_flash(p->q,
                                                             p->k,
                                                             p->v,
                                                             p->attn_out,
                                                             p->num_heads,
                                                             p->num_kv_heads,
                                                             p->tokens,
                                                             p->head_dim,
                                                             p->aligned_head_dim);
            }
 
            ck_attention_project_head_major_q4_k_q8_k(p->attn_out,
                                                      p->wo,
                                                      p->bo,
                                                      p->proj_tmp,
                                                      p->tokens,
                                                      aligned_D,
                                                      p->num_heads,
                                                      p->aligned_head_dim);
 
            ck_residual_add_token_major(p->input,
                                        p->proj_tmp,
                                        p->residual1,
                                        p->tokens,
                                        aligned_D);
 
            rmsnorm_forward(p->residual1,
                            p->ln2_gamma,
                            p->ln2_out,
                            p->ln2_rstd,
                            p->tokens,
                            p->embed_dim,
                            aligned_D,
                            p->eps);
 
            ck_mlp_swiglu_forward_q4_k_q8_k_prefill(p->ln2_out,
                                                    p->w1,
                                                    p->b1,
                                                    p->w2,
                                                    p->b2,
                                                    p->fc1_out,
                                                    p->swiglu_out,
                                                    p->mlp_out,
                                                    p->tokens,
                                                    aligned_D,
                                                    aligned_intermediate);
 
            ck_residual_add_token_major(p->residual1,
                                        p->mlp_out,
                                        p->output,
                                        p->tokens,
                                        aligned_D);
            return;
        }
    }
 
    ck_qkv_project_head_major_q4_k(p->ln1_out,
                                  p->wq, p->bq,
                                  p->wk, p->bk,
                                  p->wv, p->bv,
                                  p->q, p->k, p->v,
                                  p->tokens,
                                  p->tokens,
                                  aligned_D,
                                  p->num_heads,
                                  p->num_kv_heads,
                                  p->aligned_head_dim);
 
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(p->q,
                        p->k,
                        p->rope_cos,
                        p->rope_sin,
                        p->num_heads,
                        p->num_kv_heads,
                        p->tokens,
                        p->head_dim,
                        p->aligned_head_dim,
                        p->rope_pos_offset);
    }
 
    if (p->scores) {
        attention_forward_causal_head_major_gqa(p->q,
                                               p->k,
                                               p->v,
                                               p->scores,
                                               p->attn_out,
                                               p->num_heads,
                                               p->num_kv_heads,
                                               p->tokens,
                                               p->head_dim,
                                               p->aligned_head_dim,
                                               p->aligned_context_window);
    } else {
        attention_forward_causal_head_major_gqa_flash(p->q,
                                                     p->k,
                                                     p->v,
                                                     p->attn_out,
                                                     p->num_heads,
                                                     p->num_kv_heads,
                                                     p->tokens,
                                                     p->head_dim,
                                                     p->aligned_head_dim);
    }
 
    ck_attention_project_head_major_q4_k(p->attn_out,
                                         p->wo,
                                         p->bo,
                                         p->proj_tmp,
                                         p->proj_scratch,
                                         p->tokens,
                                         p->aligned_embed_dim,
                                         p->num_heads,
                                         p->aligned_head_dim);
 
    ck_residual_add_token_major(p->input,
                                p->proj_tmp,
                                p->residual1,
                                p->tokens,
                                p->aligned_embed_dim);
 
    rmsnorm_forward(p->residual1,
                    p->ln2_gamma,
                    p->ln2_out,
                    p->ln2_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_mlp_swiglu_forward_q4_k(p->ln2_out,
                               p->w1,
                               p->b1,
                               p->w2,
                               p->b2,
                               p->fc1_out,
                               p->swiglu_out,
                               p->mlp_out,
                               p->tokens,
                               p->aligned_embed_dim,
                               p->aligned_intermediate_dim);
 
    ck_residual_add_token_major(p->residual1,
                                p->mlp_out,
                                p->output,
                                p->tokens,
                                p->aligned_embed_dim);
}

References CKLayerForwardParamsQ4K::aligned_context_window, CKLayerForwardParamsQ4K::aligned_embed_dim, CKLayerForwardParamsQ4K::aligned_head_dim, CKLayerForwardParamsQ4K::aligned_intermediate_dim, attention_forward_causal_head_major_gqa(), attention_forward_causal_head_major_gqa_flash(), CKLayerForwardParamsQ4K::attn_out, CKLayerForwardParamsQ4K::b1, CKLayerForwardParamsQ4K::b2, CKLayerForwardParamsQ4K::bk, CKLayerForwardParamsQ4K::bo, CKLayerForwardParamsQ4K::bq, CKLayerForwardParamsQ4K::bv, ck_attention_project_head_major_q4_k(), ck_attention_project_head_major_q4_k_q8_k(), ck_mlp_swiglu_forward_q4_k(), ck_mlp_swiglu_forward_q4_k_q8_k_prefill(), ck_q8k_activations_enabled(), ck_qkv_project_head_major_q4_k(), ck_qkv_project_head_major_q4_k_q8_k(), ck_residual_add_token_major(), CKLayerForwardParamsQ4K::embed_dim, CKLayerForwardParamsQ4K::eps, CKLayerForwardParamsQ4K::fc1_out, CKLayerForwardParamsQ4K::head_dim, CKLayerForwardParamsQ4K::input, CKLayerForwardParamsQ4K::k, CKLayerForwardParamsQ4K::ln1_gamma, CKLayerForwardParamsQ4K::ln1_out, CKLayerForwardParamsQ4K::ln1_rstd, CKLayerForwardParamsQ4K::ln2_gamma, CKLayerForwardParamsQ4K::ln2_out, CKLayerForwardParamsQ4K::ln2_rstd, CKLayerForwardParamsQ4K::mlp_out, CKLayerForwardParamsQ4K::num_heads, CKLayerForwardParamsQ4K::num_kv_heads, CKLayerForwardParamsQ4K::output, CKLayerForwardParamsQ4K::proj_scratch, CKLayerForwardParamsQ4K::proj_tmp, CKLayerForwardParamsQ4K::q, QK_K, CKLayerForwardParamsQ4K::residual1, rmsnorm_forward(), CKLayerForwardParamsQ4K::rope_cos, rope_forward_qk(), CKLayerForwardParamsQ4K::rope_pos_offset, CKLayerForwardParamsQ4K::rope_sin, CKLayerForwardParamsQ4K::scores, CKLayerForwardParamsQ4K::swiglu_out, CKLayerForwardParamsQ4K::tokens, CKLayerForwardParamsQ4K::v, CKLayerForwardParamsQ4K::w1, CKLayerForwardParamsQ4K::w2, CKLayerForwardParamsQ4K::wk, CKLayerForwardParamsQ4K::wo, CKLayerForwardParamsQ4K::wq, and CKLayerForwardParamsQ4K::wv.

ck_layer_forward_rmsnorm_swiglu_quant()

void ck_layer_forward_rmsnorm_swiglu_quant

(

const CKLayerForwardParamsQ4K *

p

)

Definition at line 2401 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
 
    rmsnorm_forward(p->input,
                    p->ln1_gamma,
                    p->ln1_out,
                    p->ln1_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_qkv_project_head_major_quant(p->ln1_out,
                                   p->wq, p->bq, p->wq_dtype,
                                   p->wk, p->bk, p->wk_dtype,
                                   p->wv, p->bv, p->wv_dtype,
                                   p->q, p->k, p->v,
                                   p->tokens,
                                   p->tokens,
                                   p->aligned_embed_dim,
                                   p->num_heads,
                                   p->num_kv_heads,
                                   p->aligned_head_dim);
 
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(p->q,
                        p->k,
                        p->rope_cos,
                        p->rope_sin,
                        p->num_heads,
                        p->num_kv_heads,
                        p->tokens,
                        p->head_dim,
                        p->aligned_head_dim,
                        p->rope_pos_offset);
    }
 
    if (p->scores) {
        attention_forward_causal_head_major_gqa(p->q,
                                               p->k,
                                               p->v,
                                               p->scores,
                                               p->attn_out,
                                               p->num_heads,
                                               p->num_kv_heads,
                                               p->tokens,
                                               p->head_dim,
                                               p->aligned_head_dim,
                                               p->aligned_context_window);
    } else {
        attention_forward_causal_head_major_gqa_flash(p->q,
                                                     p->k,
                                                     p->v,
                                                     p->attn_out,
                                                     p->num_heads,
                                                     p->num_kv_heads,
                                                     p->tokens,
                                                     p->head_dim,
                                                     p->aligned_head_dim);
    }
 
    ck_attention_project_head_major_quant(p->attn_out,
                                          p->wo,
                                          p->bo,
                                          p->proj_tmp,
                                          p->proj_scratch,
                                          p->tokens,
                                          p->aligned_embed_dim,
                                          p->num_heads,
                                          p->aligned_head_dim,
                                          p->wo_dtype);
 
    ck_residual_add_token_major(p->input,
                                p->proj_tmp,
                                p->residual1,
                                p->tokens,
                                p->aligned_embed_dim);
 
    rmsnorm_forward(p->residual1,
                    p->ln2_gamma,
                    p->ln2_out,
                    p->ln2_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_mlp_swiglu_forward_quant(p->ln2_out,
                                p->w1,
                                p->b1,
                                p->w1_dtype,
                                p->w2,
                                p->b2,
                                p->w2_dtype,
                                p->fc1_out,
                                p->swiglu_out,
                                p->mlp_out,
                                p->tokens,
                                p->aligned_embed_dim,
                                p->aligned_intermediate_dim);
 
    ck_residual_add_token_major(p->residual1,
                                p->mlp_out,
                                p->output,
                                p->tokens,
                                p->aligned_embed_dim);
}

References CKLayerForwardParamsQ4K::aligned_context_window, CKLayerForwardParamsQ4K::aligned_embed_dim, CKLayerForwardParamsQ4K::aligned_head_dim, CKLayerForwardParamsQ4K::aligned_intermediate_dim, attention_forward_causal_head_major_gqa(), attention_forward_causal_head_major_gqa_flash(), CKLayerForwardParamsQ4K::attn_out, CKLayerForwardParamsQ4K::b1, CKLayerForwardParamsQ4K::b2, CKLayerForwardParamsQ4K::bk, CKLayerForwardParamsQ4K::bo, CKLayerForwardParamsQ4K::bq, CKLayerForwardParamsQ4K::bv, ck_attention_project_head_major_quant(), ck_mlp_swiglu_forward_quant(), ck_qkv_project_head_major_quant(), ck_residual_add_token_major(), CKLayerForwardParamsQ4K::embed_dim, CKLayerForwardParamsQ4K::eps, CKLayerForwardParamsQ4K::fc1_out, CKLayerForwardParamsQ4K::head_dim, CKLayerForwardParamsQ4K::input, CKLayerForwardParamsQ4K::k, CKLayerForwardParamsQ4K::ln1_gamma, CKLayerForwardParamsQ4K::ln1_out, CKLayerForwardParamsQ4K::ln1_rstd, CKLayerForwardParamsQ4K::ln2_gamma, CKLayerForwardParamsQ4K::ln2_out, CKLayerForwardParamsQ4K::ln2_rstd, CKLayerForwardParamsQ4K::mlp_out, CKLayerForwardParamsQ4K::num_heads, CKLayerForwardParamsQ4K::num_kv_heads, CKLayerForwardParamsQ4K::output, CKLayerForwardParamsQ4K::proj_scratch, CKLayerForwardParamsQ4K::proj_tmp, CKLayerForwardParamsQ4K::q, CKLayerForwardParamsQ4K::residual1, rmsnorm_forward(), CKLayerForwardParamsQ4K::rope_cos, rope_forward_qk(), CKLayerForwardParamsQ4K::rope_pos_offset, CKLayerForwardParamsQ4K::rope_sin, CKLayerForwardParamsQ4K::scores, CKLayerForwardParamsQ4K::swiglu_out, CKLayerForwardParamsQ4K::tokens, CKLayerForwardParamsQ4K::v, CKLayerForwardParamsQ4K::w1, CKLayerForwardParamsQ4K::w1_dtype, CKLayerForwardParamsQ4K::w2, CKLayerForwardParamsQ4K::w2_dtype, CKLayerForwardParamsQ4K::wk, CKLayerForwardParamsQ4K::wk_dtype, CKLayerForwardParamsQ4K::wo, CKLayerForwardParamsQ4K::wo_dtype, CKLayerForwardParamsQ4K::wq, CKLayerForwardParamsQ4K::wq_dtype, CKLayerForwardParamsQ4K::wv, and CKLayerForwardParamsQ4K::wv_dtype.

ck_layer_forward_rmsnorm_swiglu_ref()

void ck_layer_forward_rmsnorm_swiglu_ref

(

const CKLayerForwardParams *

p

)

Definition at line 1104 of file ckernel_orchestration.c.

{
    if (!p) {
        return;
    }
 
    rmsnorm_forward(p->input,
                    p->ln1_gamma,
                    p->ln1_out,
                    p->ln1_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_qkv_project_head_major_ref(p->ln1_out,
                                  p->wq, p->bq,
                                  p->wk, p->bk,
                                  p->wv, p->bv,
                                  p->q, p->k, p->v,
                                  p->tokens,
                                  p->tokens,
                                  p->aligned_embed_dim,
                                  p->num_heads,
                                  p->num_kv_heads,
                                  p->aligned_head_dim);
 
    if (p->rope_cos && p->rope_sin) {
        rope_forward_qk(p->q,
                        p->k,
                        p->rope_cos,
                        p->rope_sin,
                        p->num_heads,
                        p->num_kv_heads,
                        p->tokens,
                        p->head_dim,
                        p->aligned_head_dim,
                        p->rope_pos_offset);
    }
 
    if (p->scores) {
        attention_forward_causal_head_major_gqa(p->q,
                                               p->k,
                                               p->v,
                                               p->scores,
                                               p->attn_out,
                                               p->num_heads,
                                               p->num_kv_heads,
                                               p->tokens,
                                               p->head_dim,
                                               p->aligned_head_dim,
                                               p->aligned_context_window);
    } else {
        attention_forward_causal_head_major_gqa_flash(p->q,
                                                     p->k,
                                                     p->v,
                                                     p->attn_out,
                                                     p->num_heads,
                                                     p->num_kv_heads,
                                                     p->tokens,
                                                     p->head_dim,
                                                     p->aligned_head_dim);
    }
 
    ck_attention_project_head_major_ref(p->attn_out,
                                        p->wo,
                                        p->bo,
                                        p->proj_tmp,
                                        p->proj_scratch,
                                        p->tokens,
                                        p->aligned_embed_dim,
                                        p->num_heads,
                                        p->aligned_head_dim);
 
    ck_residual_add_token_major(p->input,
                                p->proj_tmp,
                                p->residual1,
                                p->tokens,
                                p->aligned_embed_dim);
 
    rmsnorm_forward(p->residual1,
                    p->ln2_gamma,
                    p->ln2_out,
                    p->ln2_rstd,
                    p->tokens,
                    p->embed_dim,
                    p->aligned_embed_dim,
                    p->eps);
 
    ck_mlp_swiglu_forward_ref(p->ln2_out,
                              p->w1,
                              p->b1,
                              p->w2,
                              p->b2,
                              p->fc1_out,
                              p->swiglu_out,
                              p->mlp_out,
                              p->tokens,
                              p->aligned_embed_dim,
                              p->aligned_intermediate_dim);
 
    ck_residual_add_token_major(p->residual1,
                                p->mlp_out,
                                p->output,
                                p->tokens,
                                p->aligned_embed_dim);
}

References CKLayerForwardParams::aligned_context_window, CKLayerForwardParams::aligned_embed_dim, CKLayerForwardParams::aligned_head_dim, CKLayerForwardParams::aligned_intermediate_dim, attention_forward_causal_head_major_gqa(), attention_forward_causal_head_major_gqa_flash(), CKLayerForwardParams::attn_out, CKLayerForwardParams::b1, CKLayerForwardParams::b2, CKLayerForwardParams::bk, CKLayerForwardParams::bo, CKLayerForwardParams::bq, CKLayerForwardParams::bv, ck_attention_project_head_major_ref(), ck_mlp_swiglu_forward_ref(), ck_qkv_project_head_major_ref(), ck_residual_add_token_major(), CKLayerForwardParams::embed_dim, CKLayerForwardParams::eps, CKLayerForwardParams::fc1_out, CKLayerForwardParams::head_dim, CKLayerForwardParams::input, CKLayerForwardParams::k, CKLayerForwardParams::ln1_gamma, CKLayerForwardParams::ln1_out, CKLayerForwardParams::ln1_rstd, CKLayerForwardParams::ln2_gamma, CKLayerForwardParams::ln2_out, CKLayerForwardParams::ln2_rstd, CKLayerForwardParams::mlp_out, CKLayerForwardParams::num_heads, CKLayerForwardParams::num_kv_heads, CKLayerForwardParams::output, CKLayerForwardParams::proj_scratch, CKLayerForwardParams::proj_tmp, CKLayerForwardParams::q, CKLayerForwardParams::residual1, rmsnorm_forward(), CKLayerForwardParams::rope_cos, rope_forward_qk(), CKLayerForwardParams::rope_pos_offset, CKLayerForwardParams::rope_sin, CKLayerForwardParams::scores, CKLayerForwardParams::swiglu_out, CKLayerForwardParams::tokens, CKLayerForwardParams::v, CKLayerForwardParams::w1, CKLayerForwardParams::w2, CKLayerForwardParams::wk, CKLayerForwardParams::wo, CKLayerForwardParams::wq, and CKLayerForwardParams::wv.

ck_mlp_swiglu_forward()

void ck_mlp_swiglu_forward	(	const float *	input,
		const float *	w1,
		const float *	b1,
		const float *	w2,
		const float *	b2,
		float *	fc1_out,
		float *	swiglu_out,
		float *	output,
		int	tokens,
		int	aligned_embed_dim,
		int	aligned_intermediate_dim
	)

Definition at line 952 of file ckernel_orchestration.c.

{
    int up_dim = 2 * aligned_intermediate_dim;
    gemm_blocked_serial(input, w1, b1, fc1_out,
                        tokens, up_dim, aligned_embed_dim);
 
    swiglu_forward(fc1_out, swiglu_out, tokens, aligned_intermediate_dim);
 
    gemm_blocked_serial(swiglu_out, w2, b2, output,
                        tokens, aligned_embed_dim, aligned_intermediate_dim);
}

References gemm_blocked_serial(), and swiglu_forward().

Referenced by ck_layer_forward_rmsnorm_swiglu(), ck_layer_forward_rmsnorm_swiglu_decode(), and ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl().

ck_mlp_swiglu_forward_fully_fused_token()

void ck_mlp_swiglu_forward_fully_fused_token	(	const float *	input_row,
		const float *	w1,
		const float *	b1,
		const float *	w2,
		const float *	b2,
		float *	output_row,
		int	aligned_embed_dim,
		int	aligned_intermediate_dim
	)

Definition at line 1247 of file ckernel_orchestration.c.

{
    if (!input_row || !w1 || !w2 || !output_row) {
        return;
    }
 
    // Split w1 into gate and up projections
    // w1 layout: [2 * aligned_intermediate_dim, aligned_embed_dim]
    //   First half: W_gate [aligned_intermediate_dim, aligned_embed_dim]
    //   Second half: W_up [aligned_intermediate_dim, aligned_embed_dim]
    const float *w_gate = w1;
    const float *w_up = w1 + (size_t)aligned_intermediate_dim * (size_t)aligned_embed_dim;
 
    // Split b1 into gate and up biases (if present)
    const float *b_gate = b1;
    const float *b_up = b1 ? (b1 + aligned_intermediate_dim) : NULL;
 
    // w2 is W_down: [aligned_embed_dim, aligned_intermediate_dim]
    const float *w_down = w2;
    const float *b_down = b2;
 
    // Call the fully fused kernel - eliminates DRAM round-trip for swiglu
    // Uses aligned dimensions since weights are stored with alignment padding
    fused_mlp_swiglu_decode_v2(input_row,
                                w_gate,
                                w_up,
                                w_down,
                                b_gate,
                                b_up,
                                b_down,
                                output_row,
                                aligned_embed_dim,
                                aligned_intermediate_dim);
}

References fused_mlp_swiglu_decode_v2().

Referenced by ck_layer_forward_rmsnorm_swiglu_decode_fused(), and ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl().

ck_mlp_swiglu_forward_fused_token()

void ck_mlp_swiglu_forward_fused_token	(	const float *	input_row,
		const float *	w1,
		const float *	b1,
		const float *	w2,
		const float *	b2,
		float *	swiglu_row,
		float *	output_row,
		int	aligned_embed_dim,
		int	aligned_intermediate_dim
	)

Definition at line 1212 of file ckernel_orchestration.c.

{
    if (!input_row || !w1 || !w2 || !swiglu_row || !output_row) {
        return;
    }
 
    const float *w_gate = w1;
    const float *w_up = w1 + (size_t)aligned_intermediate_dim * (size_t)aligned_embed_dim;
    const float *b_gate = b1;
    const float *b_up = b1 ? (b1 + aligned_intermediate_dim) : NULL;
 
    gemm_swiglu_fused(input_row,
                      w_gate,
                      w_up,
                      b_gate,
                      b_up,
                      swiglu_row,
                      /*M=*/1,
                      /*N=*/aligned_intermediate_dim,
                      /*K=*/aligned_embed_dim);
 
    gemm_blocked_serial(swiglu_row, w2, b2, output_row,
                        /*M=*/1,
                        /*N=*/aligned_embed_dim,
                        /*K=*/aligned_intermediate_dim);
}

References gemm_blocked_serial(), and gemm_swiglu_fused().

ck_qkv_project_head_major()

void ck_qkv_project_head_major	(	const float *	input,
		const float *	wq,
		const float *	bq,
		const float *	wk,
		const float *	bk,
		const float *	wv,
		const float *	bv,
		float *	q,
		float *	k,
		float *	v,
		int	tokens,
		int	kv_stride_tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	aligned_head_dim
	)

Definition at line 168 of file ckernel_orchestration.c.

{
    if (!input || !wq || !wk || !wv || !q || !k || !v) {
        return;
    }
    if (kv_stride_tokens < tokens) {
        return;
    }
 
    size_t head_weight_stride = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
    size_t q_head_stride = (size_t)tokens * (size_t)aligned_head_dim;
    size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        const float *wq_h = wq + (size_t)h * head_weight_stride;
        const float *bq_h = bq ? (bq + (size_t)h * (size_t)aligned_head_dim) : NULL;
        float *q_h = q + (size_t)h * q_head_stride;
 
        gemm_blocked_serial(input, wq_h, bq_h, q_h,
                            tokens, aligned_head_dim, aligned_embed_dim);
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *wk_h = wk + (size_t)h * head_weight_stride;
        const float *wv_h = wv + (size_t)h * head_weight_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 + (size_t)h * kv_head_stride;
        float *v_h = v + (size_t)h * kv_head_stride;
 
        gemm_blocked_serial(input, wk_h, bk_h, k_h,
                            tokens, aligned_head_dim, aligned_embed_dim);
        gemm_blocked_serial(input, wv_h, bv_h, v_h,
                            tokens, aligned_head_dim, aligned_embed_dim);
    }
}

References gemm_blocked_serial().

Referenced by ck_layer_forward_rmsnorm_swiglu().

ck_qkv_project_head_major_backward()

void ck_qkv_project_head_major_backward	(	const float *	d_q,
		const float *	d_k,
		const float *	d_v,
		const float *	input,
		const float *	wq,
		const float *	bq,
		const float *	wk,
		const float *	bk,
		const float *	wv,
		const float *	bv,
		float *	d_input,
		float *	d_wq,
		float *	d_bq,
		float *	d_wk,
		float *	d_bk,
		float *	d_wv,
		float *	d_bv,
		float *	scratch,
		int	tokens,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	aligned_head_dim,
		int	num_threads
	)

Definition at line 856 of file ckernel_orchestration.c.

{
    if (!d_q || !d_k || !d_v || !input || !wq || !wk || !wv ||
        !d_input || !d_wq || !d_bq || !d_wk || !d_bk || !d_wv || !d_bv || !scratch) {
        return;
    }
 
    size_t total_in = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t i = 0; i < total_in; ++i) {
        d_input[i] = 0.0f;
    }
 
    size_t head_weight_stride = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
    size_t head_out_stride = (size_t)tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        const float *d_q_h = d_q + (size_t)h * head_out_stride;
        const float *wq_h = wq + (size_t)h * head_weight_stride;
        float *d_wq_h = d_wq + (size_t)h * head_weight_stride;
        float *d_bq_h = d_bq + (size_t)h * (size_t)aligned_head_dim;
 
        fc2_backward_kernel(d_q_h,
                            input,
                            wq_h,
                            scratch,
                            d_wq_h,
                            d_bq_h,
                            tokens,
                            aligned_embed_dim,
                            aligned_head_dim,
                            num_threads);
        ck_add_inplace(d_input, scratch, tokens, aligned_embed_dim);
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *d_k_h = d_k + (size_t)h * head_out_stride;
        const float *d_v_h = d_v + (size_t)h * head_out_stride;
 
        const float *wk_h = wk + (size_t)h * head_weight_stride;
        const float *wv_h = wv + (size_t)h * head_weight_stride;
 
        float *d_wk_h = d_wk + (size_t)h * head_weight_stride;
        float *d_wv_h = d_wv + (size_t)h * head_weight_stride;
 
        float *d_bk_h = d_bk + (size_t)h * (size_t)aligned_head_dim;
        float *d_bv_h = d_bv + (size_t)h * (size_t)aligned_head_dim;
 
        fc2_backward_kernel(d_k_h,
                            input,
                            wk_h,
                            scratch,
                            d_wk_h,
                            d_bk_h,
                            tokens,
                            aligned_embed_dim,
                            aligned_head_dim,
                            num_threads);
        ck_add_inplace(d_input, scratch, tokens, aligned_embed_dim);
 
        fc2_backward_kernel(d_v_h,
                            input,
                            wv_h,
                            scratch,
                            d_wv_h,
                            d_bv_h,
                            tokens,
                            aligned_embed_dim,
                            aligned_head_dim,
                            num_threads);
        ck_add_inplace(d_input, scratch, tokens, aligned_embed_dim);
    }
}

References ck_add_inplace(), and fc2_backward_kernel().

Referenced by ck_layer_backward_rmsnorm_swiglu().

ck_qkv_project_head_major_token()

void ck_qkv_project_head_major_token	(	const float *	input_row,
		const float *	wq,
		const float *	bq,
		const float *	wk,
		const float *	bk,
		const float *	wv,
		const float *	bv,
		float *	q_token,
		float *	k_token,
		float *	v_token,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	aligned_head_dim
	)

Definition at line 78 of file attention_decode_fused.c.

{
    if (!input_row || !wq || !wk || !wv || !q_token || !k_token || !v_token) {
        return;
    }
 
    const int q_out = num_heads * aligned_head_dim;
    gemm_blocked_serial(input_row, wq, bq, q_token,
                        /*tokens=*/1, q_out, aligned_embed_dim);
 
    size_t head_weight_stride = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
#pragma omp parallel for schedule(static) if(num_kv_heads > 1)
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *wk_h = wk + (size_t)h * head_weight_stride;
        const float *wv_h = wv + (size_t)h * head_weight_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_token + (size_t)h * (size_t)aligned_head_dim;
        float *v_h = v_token + (size_t)h * (size_t)aligned_head_dim;
 
        gemm_blocked_serial(input_row, wk_h, bk_h, k_h,
                            /*tokens=*/1, aligned_head_dim, aligned_embed_dim);
        gemm_blocked_serial(input_row, wv_h, bv_h, v_h,
                            /*tokens=*/1, aligned_head_dim, aligned_embed_dim);
    }
}

References gemm_blocked_serial().

Referenced by ck_layer_forward_rmsnorm_swiglu_decode(), ck_layer_forward_rmsnorm_swiglu_decode_fused(), and ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl().

ck_residual_add_backward()

void ck_residual_add_backward	(	const float *	d_out,
		float *	d_a,
		float *	d_b,
		int	tokens,
		int	aligned_embed_dim
	)

Definition at line 151 of file ckernel_orchestration.c.

{
    if (!d_out || !d_a || !d_b) {
        return;
    }
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t i = 0; i < total; ++i) {
        float v = d_out[i];
        d_a[i] = v;
        d_b[i] = v;
    }
}

Referenced by ck_layer_backward_rmsnorm_swiglu().

ck_residual_add_token_major()

void ck_residual_add_token_major	(	const float *	a,
		const float *	b,
		float *	out,
		int	tokens,
		int	aligned_embed_dim
	)

Definition at line 139 of file ckernel_orchestration.c.

{
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t i = 0; i < total; ++i) {
        out[i] = a[i] + b[i];
    }
}

Referenced by ck_layer_forward_rmsnorm_swiglu(), ck_layer_forward_rmsnorm_swiglu_decode(), ck_layer_forward_rmsnorm_swiglu_decode_fused(), ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl(), ck_layer_forward_rmsnorm_swiglu_decode_q4_k(), ck_layer_forward_rmsnorm_swiglu_decode_quant(), ck_layer_forward_rmsnorm_swiglu_q4_k(), ck_layer_forward_rmsnorm_swiglu_quant(), ck_layer_forward_rmsnorm_swiglu_ref(), mega_fused_attention_prefill(), and mega_fused_attention_prefill_q8_0().