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

Data Structures
struct	CKBufferSpec

struct	CKDimToken

struct	CKKernelSpec

struct	CKPlanBinding

struct	CKPlanStep

struct	CKPlanStepV2

Macros
#define	CKERNEL_MAX_KERNEL_SOURCES 8

Enumerations
enum	CKBufferRole { CK_ROLE_INPUT = 0 , CK_ROLE_OUTPUT , CK_ROLE_ACTIVATION , CK_ROLE_WEIGHT , CK_ROLE_SCRATCH , CK_ROLE_GRAD }

enum	CKBufferScope { CK_SCOPE_LAYER = 0 , CK_SCOPE_GLOBAL }

enum	CKDimKind { CK_DIM_TOKENS = 0 , CK_DIM_EMBED , CK_DIM_ALIGNED_EMBED , CK_DIM_HEAD_DIM , CK_DIM_ALIGNED_HEAD , CK_DIM_NUM_HEADS , CK_DIM_NUM_KV_HEADS , CK_DIM_ALIGNED_CTX , CK_DIM_INTERMEDIATE , CK_DIM_ALIGNED_INTERMEDIATE , CK_DIM_VOCAB , CK_DIM_END }

Variables
const CKPlanStep	ck_decoder_backward_plan []

const size_t	ck_decoder_backward_plan_count

const CKPlanStepV2	ck_decoder_backward_plan_v2 []

const size_t	ck_decoder_backward_plan_v2_count

const size_t	ck_decoder_buffer_count

const CKBufferSpec	ck_decoder_buffers []

const CKPlanStep	ck_decoder_forward_plan []

const size_t	ck_decoder_forward_plan_count

const CKPlanStepV2	ck_decoder_forward_plan_v2 []

const size_t	ck_decoder_forward_plan_v2_count

const size_t	ck_kernel_spec_count

const CKKernelSpec	ck_kernel_specs []

Macro Definition Documentation

◆ CKERNEL_MAX_KERNEL_SOURCES

#define CKERNEL_MAX_KERNEL_SOURCES 8

Definition at line 23 of file ckernel_kernel_specs.h.

Enumeration Type Documentation

◆ CKBufferRole

enum CKBufferRole

Enumerator
CK_ROLE_INPUT
CK_ROLE_OUTPUT
CK_ROLE_ACTIVATION
CK_ROLE_WEIGHT
CK_ROLE_SCRATCH
CK_ROLE_GRAD

Definition at line 36 of file ckernel_kernel_specs.h.

             {
    CK_ROLE_INPUT = 0,
    CK_ROLE_OUTPUT,
    CK_ROLE_ACTIVATION,
    CK_ROLE_WEIGHT,
    CK_ROLE_SCRATCH,
    CK_ROLE_GRAD
} CKBufferRole;

◆ CKBufferScope

enum CKBufferScope

Enumerator
CK_SCOPE_LAYER
CK_SCOPE_GLOBAL

Definition at line 31 of file ckernel_kernel_specs.h.

             {
    CK_SCOPE_LAYER = 0,
    CK_SCOPE_GLOBAL
} CKBufferScope;

◆ CKDimKind

enum CKDimKind

Enumerator
CK_DIM_TOKENS
CK_DIM_EMBED
CK_DIM_ALIGNED_EMBED
CK_DIM_HEAD_DIM
CK_DIM_ALIGNED_HEAD
CK_DIM_NUM_HEADS
CK_DIM_NUM_KV_HEADS
CK_DIM_ALIGNED_CTX
CK_DIM_INTERMEDIATE
CK_DIM_ALIGNED_INTERMEDIATE
CK_DIM_VOCAB
CK_DIM_END

Definition at line 8 of file ckernel_kernel_specs.h.

             {
    CK_DIM_TOKENS = 0,
    CK_DIM_EMBED,
    CK_DIM_ALIGNED_EMBED,
    CK_DIM_HEAD_DIM,
    CK_DIM_ALIGNED_HEAD,
    CK_DIM_NUM_HEADS,
    CK_DIM_NUM_KV_HEADS,
    CK_DIM_ALIGNED_CTX,
    CK_DIM_INTERMEDIATE,
    CK_DIM_ALIGNED_INTERMEDIATE,
    CK_DIM_VOCAB,
    CK_DIM_END
} CKDimKind;

Variable Documentation

◆ ck_decoder_backward_plan

const CKPlanStep ck_decoder_backward_plan[]

extern

Definition at line 115 of file ckernel_kernel_specs.c.

                                              {
    {"residual_add", NULL},
    {"mlp_down", NULL},
    {"swiglu", NULL},
    {"mlp_up", NULL},
    {"rmsnorm", NULL},
    {"residual_add", NULL},
    {"attn_proj", NULL},
    {"attention", NULL},
    {"rope", "rope_theta>0"},
    {"qkv_project", NULL},
    {"rmsnorm", NULL},
};

Referenced by emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), and emit_kernel_manifest().

◆ ck_decoder_backward_plan_count

const size_t ck_decoder_backward_plan_count

extern

Definition at line 129 of file ckernel_kernel_specs.c.

Referenced by emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), and emit_kernel_manifest().

◆ ck_decoder_backward_plan_v2

const CKPlanStepV2 ck_decoder_backward_plan_v2[]

extern

Definition at line 329 of file ckernel_kernel_specs.c.

                                                   {
    {"residual_add", NULL, ck_decoder_backward_bindings_0, 3},
    {"mlp_down", NULL, ck_decoder_backward_bindings_1, 6},
    {"swiglu", NULL, ck_decoder_backward_bindings_2, 3},
    {"mlp_up", NULL, ck_decoder_backward_bindings_3, 6},
    {"rmsnorm", NULL, ck_decoder_backward_bindings_4, 6},
    {"residual_add", NULL, ck_decoder_backward_bindings_5, 3},
    {"attn_proj", NULL, ck_decoder_backward_bindings_6, 6},
    {"attention", NULL, ck_decoder_backward_bindings_7, 9},
    {"rope", "rope_theta>0", ck_decoder_backward_bindings_8, 6},
    {"qkv_project", NULL, ck_decoder_backward_bindings_9, 14},
    {"rmsnorm", NULL, ck_decoder_backward_bindings_10, 6},
};

Referenced by ck_ir_v2_build_decoder_backward().

◆ ck_decoder_backward_plan_v2_count

const size_t ck_decoder_backward_plan_v2_count

extern

Definition at line 343 of file ckernel_kernel_specs.c.

Referenced by ck_ir_v2_build_decoder_backward().

◆ ck_decoder_buffer_count

const size_t ck_decoder_buffer_count

extern

Definition at line 83 of file ckernel_kernel_specs.c.

Referenced by ck_find_buffer_spec(), ck_find_buffer_spec(), ck_find_buffer_spec(), ck_find_buffer_spec(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_ir_v2_build_decoder(), ck_ir_v2_find_buffer_spec(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_allocations(), emit_global_allocations(), emit_global_allocations(), emit_global_allocations(), emit_global_offset_fields(), emit_global_offset_fields(), emit_global_offset_fields(), emit_global_offset_fields(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_sgd_update(), emit_sgd_update(), emit_sgd_update(), emit_sgd_update(), emit_zero_grad(), emit_zero_grad(), emit_zero_grad(), and emit_zero_grad().

◆ ck_decoder_buffers

const CKBufferSpec ck_decoder_buffers[]

extern

Definition at line 3 of file ckernel_kernel_specs.c.

                                          {
    {"token_emb", CK_SCOPE_GLOBAL, CK_ROLE_WEIGHT, { { CK_DIM_VOCAB, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"pos_emb", CK_SCOPE_GLOBAL, CK_ROLE_WEIGHT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"embedded_input", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"rope_cos_cache", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_HEAD_DIM, 1, 2 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, "rope_theta", CK_DT_FP32},
    {"rope_sin_cache", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_HEAD_DIM, 1, 2 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, "rope_theta", CK_DT_FP32},
    {"final_ln_weight", CK_SCOPE_GLOBAL, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"final_ln_bias", CK_SCOPE_GLOBAL, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"final_ln_mean", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"final_ln_rstd", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"final_output", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"lm_head_weight", CK_SCOPE_GLOBAL, CK_ROLE_WEIGHT, { { CK_DIM_VOCAB, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, "token_emb", NULL, CK_DT_FP32},
    {"logits", CK_SCOPE_GLOBAL, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_VOCAB, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_token_emb", CK_SCOPE_GLOBAL, CK_ROLE_GRAD, { { CK_DIM_VOCAB, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_pos_emb", CK_SCOPE_GLOBAL, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_final_output", CK_SCOPE_GLOBAL, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_final_input", CK_SCOPE_GLOBAL, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_final_ln_weight", CK_SCOPE_GLOBAL, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_logits", CK_SCOPE_GLOBAL, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_VOCAB, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"input", CK_SCOPE_LAYER, CK_ROLE_INPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"ln1_gamma", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"ln1_out", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"ln1_rstd", CK_SCOPE_LAYER, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 1, NULL, NULL, CK_DT_FP32},
    {"wq", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"bq", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"wk", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"bk", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"wv", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"bv", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"q", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"k", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"v", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"scores", CK_SCOPE_LAYER, CK_ROLE_ACTIVATION, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_CTX, 1, 1 }, { CK_DIM_ALIGNED_CTX, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, "training_enabled", CK_DT_FP32},
    {"attn_out", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"wo", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"bo", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"proj_tmp", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"proj_scratch", CK_SCOPE_LAYER, CK_ROLE_SCRATCH, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"residual1", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"ln2_gamma", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"ln2_out", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"ln2_rstd", CK_SCOPE_LAYER, CK_ROLE_ACTIVATION, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 1, NULL, NULL, CK_DT_FP32},
    {"w1", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_INTERMEDIATE, 2, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"b1", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_INTERMEDIATE, 2, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"fc1_out", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_INTERMEDIATE, 2, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"swiglu_out", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_INTERMEDIATE, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"w2", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_ALIGNED_INTERMEDIATE, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"b2", CK_SCOPE_LAYER, CK_ROLE_WEIGHT, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"mlp_out", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"output", CK_SCOPE_LAYER, CK_ROLE_OUTPUT, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_output", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_residual1", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_mlp_out", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_swiglu_out", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_INTERMEDIATE, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_w2", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_ALIGNED_INTERMEDIATE, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_b2", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_fc1_out", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_INTERMEDIATE, 2, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_ln2_out", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_w1", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_INTERMEDIATE, 2, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_b1", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_INTERMEDIATE, 2, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_ln2_gamma", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_input", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_proj_tmp", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_attn_out", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_wo", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_bo", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_q", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_k", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_v", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_scores", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_CTX, 1, 1 }, { CK_DIM_ALIGNED_CTX, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_ln1_out", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_TOKENS, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_wq", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_bq", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_wk", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_bk", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_wv", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_bv", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_NUM_KV_HEADS, 1, 1 }, { CK_DIM_ALIGNED_HEAD, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
    {"d_ln1_gamma", CK_SCOPE_LAYER, CK_ROLE_GRAD, { { CK_DIM_ALIGNED_EMBED, 1, 1 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 }, { CK_DIM_END, 0, 0 } }, 0, NULL, NULL, CK_DT_FP32},
};

Referenced by ck_find_buffer_spec(), ck_find_buffer_spec(), ck_find_buffer_spec(), ck_find_buffer_spec(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_ir_v2_build_decoder(), ck_ir_v2_find_buffer_spec(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_allocations(), emit_global_allocations(), emit_global_allocations(), emit_global_allocations(), emit_global_offset_fields(), emit_global_offset_fields(), emit_global_offset_fields(), emit_global_offset_fields(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_sgd_update(), emit_sgd_update(), emit_sgd_update(), emit_sgd_update(), emit_zero_grad(), emit_zero_grad(), emit_zero_grad(), and emit_zero_grad().

◆ ck_decoder_forward_plan

const CKPlanStep ck_decoder_forward_plan[]

extern

Definition at line 99 of file ckernel_kernel_specs.c.

                                             {
    {"rmsnorm", NULL},
    {"qkv_project", NULL},
    {"rope", "rope_theta>0"},
    {"attention", NULL},
    {"attn_proj", NULL},
    {"residual_add", NULL},
    {"rmsnorm", NULL},
    {"mlp_up", NULL},
    {"swiglu", NULL},
    {"mlp_down", NULL},
    {"residual_add", NULL},
};

Referenced by emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), and emit_kernel_manifest().

◆ ck_decoder_forward_plan_count

const size_t ck_decoder_forward_plan_count

extern

Definition at line 113 of file ckernel_kernel_specs.c.

Referenced by emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), and emit_kernel_manifest().

◆ ck_decoder_forward_plan_v2

const CKPlanStepV2 ck_decoder_forward_plan_v2[]

extern

Definition at line 313 of file ckernel_kernel_specs.c.

                                                  {
    {"rmsnorm", NULL, ck_decoder_forward_bindings_0, 4},
    {"qkv_project", NULL, ck_decoder_forward_bindings_1, 10},
    {"rope", "rope_theta>0", ck_decoder_forward_bindings_2, 4},
    {"attention", NULL, ck_decoder_forward_bindings_3, 5},
    {"attn_proj", NULL, ck_decoder_forward_bindings_4, 5},
    {"residual_add", NULL, ck_decoder_forward_bindings_5, 3},
    {"rmsnorm", NULL, ck_decoder_forward_bindings_6, 4},
    {"mlp_up", NULL, ck_decoder_forward_bindings_7, 4},
    {"swiglu", NULL, ck_decoder_forward_bindings_8, 2},
    {"mlp_down", NULL, ck_decoder_forward_bindings_9, 4},
    {"residual_add", NULL, ck_decoder_forward_bindings_10, 3},
};

Referenced by ck_ir_v2_build_decoder().

◆ ck_decoder_forward_plan_v2_count

const size_t ck_decoder_forward_plan_v2_count

extern

Definition at line 327 of file ckernel_kernel_specs.c.

Referenced by ck_ir_v2_build_decoder().

◆ ck_kernel_spec_count

const size_t ck_kernel_spec_count

extern

Definition at line 97 of file ckernel_kernel_specs.c.

Referenced by ck_find_kernel_spec(), ck_find_kernel_spec(), ck_find_kernel_spec(), ck_find_kernel_spec(), ck_ir_v2_find_kernel_spec(), emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), and emit_kernel_manifest().

◆ ck_kernel_specs

const CKKernelSpec ck_kernel_specs[]

extern

Definition at line 85 of file ckernel_kernel_specs.c.

                                       {
    {"attention", { "attention_forward_causal_head_major_gqa", "attention_forward_causal_head_major_gqa_bf16", NULL, NULL, NULL }, { "attention_backward_causal_head_major_gqa", "attention_backward_causal_head_major_gqa_bf16", NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32) | CK_DT_MASK(CK_DT_BF16), CK_DT_FP32, { "src/kernels/attention_kernels.c", "src/kernels/softmax_kernels.c", NULL, NULL, NULL, NULL, NULL, NULL }},
    {"attn_proj", { "ck_attention_project_head_major", NULL, NULL, NULL, NULL }, { "ck_attention_project_head_major_backward", NULL, NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32), CK_DT_FP32, { "src/ckernel_orchestration.c", "src/kernels/gemm_kernels.c", "src/kernels/mlp_kernels.c", "src/kernels/gelu_kernels.c", NULL, NULL, NULL, NULL }},
    {"mlp_down", { "gemm_blocked_serial", NULL, NULL, NULL, NULL }, { "fc2_backward_kernel", NULL, NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32), CK_DT_FP32, { "src/kernels/gemm_kernels.c", "src/kernels/mlp_kernels.c", "src/kernels/gelu_kernels.c", NULL, NULL, NULL, NULL, NULL }},
    {"mlp_up", { "gemm_blocked_serial", NULL, NULL, NULL, NULL }, { "fc1_backward_kernel", NULL, NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32), CK_DT_FP32, { "src/kernels/gemm_kernels.c", "src/kernels/mlp_kernels.c", "src/kernels/gelu_kernels.c", NULL, NULL, NULL, NULL, NULL }},
    {"qkv_project", { "ck_qkv_project_head_major", NULL, NULL, NULL, NULL }, { "ck_qkv_project_head_major_backward", NULL, NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32), CK_DT_FP32, { "src/ckernel_orchestration.c", "src/kernels/gemm_kernels.c", "src/kernels/mlp_kernels.c", "src/kernels/gelu_kernels.c", NULL, NULL, NULL, NULL }},
    {"residual_add", { "ck_residual_add_token_major", NULL, NULL, NULL, NULL }, { "ck_residual_add_backward", NULL, NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32), CK_DT_FP32, { "src/ckernel_orchestration.c", NULL, NULL, NULL, NULL, NULL, NULL, NULL }},
    {"rmsnorm", { "rmsnorm_forward", "rmsnorm_forward_bf16", NULL, "rmsnorm_forward_int8", "rmsnorm_forward_int4" }, { "rmsnorm_backward", "rmsnorm_backward_bf16", NULL, "rmsnorm_backward_int8", "rmsnorm_backward_int4" }, CK_DT_MASK(CK_DT_FP32) | CK_DT_MASK(CK_DT_BF16) | CK_DT_MASK(CK_DT_INT8) | CK_DT_MASK(CK_DT_INT4), CK_DT_FP32, { "src/kernels/rmsnorm_kernels.c", "src/kernels/rmsnorm_kernels_bf16.c", "src/kernels/rmsnorm_kernels_int8.c", "src/kernels/rmsnorm_kernels_int4.c", NULL, NULL, NULL, NULL }},
    {"rope", { "rope_forward_qk", "rope_forward_qk_bf16", NULL, NULL, NULL }, { "rope_backward_qk", "rope_backward_qk_bf16", NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32) | CK_DT_MASK(CK_DT_BF16), CK_DT_FP32, { "src/kernels/rope_kernels.c", "src/kernels/rope_kernels_bf16.c", NULL, NULL, NULL, NULL, NULL, NULL }},
    {"swiglu", { "swiglu_forward", "swiglu_forward_bf16", NULL, NULL, NULL }, { "swiglu_backward", "swiglu_backward_bf16", NULL, NULL, NULL }, CK_DT_MASK(CK_DT_FP32) | CK_DT_MASK(CK_DT_BF16), CK_DT_FP32, { "src/kernels/swiglu_kernels.c", "src/kernels/swiglu_kernels_bf16.c", "src/kernels/sigmoid_kernels.c", NULL, NULL, NULL, NULL, NULL }},
};

Referenced by ck_find_kernel_spec(), ck_find_kernel_spec(), ck_find_kernel_spec(), ck_find_kernel_spec(), and ck_ir_v2_find_kernel_spec().

Data Structures

Macros

Enumerations

Variables

Macro Definition Documentation

◆ CKERNEL_MAX_KERNEL_SOURCES

Enumeration Type Documentation

◆ CKBufferRole

◆ CKBufferScope

◆ CKDimKind

Variable Documentation

◆ ck_decoder_backward_plan

◆ ck_decoder_backward_plan_count

◆ ck_decoder_backward_plan_v2

◆ ck_decoder_backward_plan_v2_count

◆ ck_decoder_buffer_count

◆ ck_decoder_buffers

◆ ck_decoder_forward_plan

◆ ck_decoder_forward_plan_count

◆ ck_decoder_forward_plan_v2

◆ ck_decoder_forward_plan_v2_count

◆ ck_kernel_spec_count

◆ ck_kernel_specs