← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
generic_api_test.c File Reference

AUTO-GENERATED: model Implementation (IR v6.6 - Explicit Unrolled) More...

#include "ck-kernel-inference.h"
#include "ckernel_engine.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <math.h>
#include "ck_model_api.h"

Go to the source code of this file.

Macros

#define _GNU_SOURCE   /* For MAP_ANONYMOUS, MAP_HUGETLB */
 

Functions

struct __attribute__ ((packed))
 
void * ck_model_create (void)
 
void ck_model_decode (void *model, const int *token, int token_index)
 
void ck_model_forward (void *model, const int *tokens, int num_tokens)
 
void ck_model_free (void *model)
 
void * ck_model_get_base (void *model)
 
const CKModelConfigck_model_get_config (void)
 
float * ck_model_get_logits (void *model)
 
size_t ck_model_get_total_bytes (void *model)
 
void ck_model_precompute_rope (void *model)
 
int ck_model_verify_canaries (void *model)
 
static int model_align_elems (int elems, int elem_bytes, int align_bytes)
 
void model_decode (MODELModel *model, const int *token, int token_index)
 
static void model_decode_token (MODELModel *model, const int *token, int token_index)
 
void model_forward (MODELModel *model, const int *tokens, int num_tokens)
 
static void model_forward_prefill_impl (MODELModel *model, const int *tokens, int num_tokens)
 
static void model_layer_0_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_0_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_10_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_10_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_11_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_11_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_12_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_12_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_13_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_13_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_14_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_14_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_15_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_15_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_16_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_16_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_17_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_17_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_18_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_18_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_19_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_19_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_1_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_1_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_20_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_20_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_21_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_21_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_22_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_22_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_23_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_23_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_2_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_2_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_3_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_3_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_4_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_4_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_5_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_5_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_6_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_6_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_7_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_7_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_8_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_8_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_9_decode (MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
static void model_layer_9_prefill (MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
 
int model_model_allocate (MODELModel *model)
 
void model_model_free (MODELModel *model)
 
void model_precompute_rope (MODELModel *model)
 
static void model_residual_add_token_major (const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)
 
int model_verify_canaries (MODELModel *model)
 

Variables

static CKModelConfig g_model_config
 
 MagicHeader
 

Detailed Description

AUTO-GENERATED: model Implementation (IR v6.6 - Explicit Unrolled)

Generated: 2026-01-12T04:22:25.186752 UTC Total Memory: 3.57 GB Mode: decode Layers: 24 (fully unrolled)

Per-layer quant types: Layer 0: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k Layer 1: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k Layer 2: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k ... (21 more layers)

DO NOT EDIT - Regenerate with build_ir_v6.6.py or codegen_v6.6.py

Definition in file v6.6/test_generated/generic_api_test.c.

Macro Definition Documentation

◆ _GNU_SOURCE

#define _GNU_SOURCE   /* For MAP_ANONYMOUS, MAP_HUGETLB */

Definition at line 19 of file v6.6/test_generated/generic_api_test.c.

Function Documentation

◆ __attribute__()

struct __attribute__ ( (packed)  )

Definition at line 43 of file v6.6/test_generated/generic_api_test.c.

67 {
68 uint32_t magic; /* 0x434B454E */
69 uint32_t version; /* IR version */
70 uint64_t total_bytes;
71 uint64_t weight_bytes;
72 uint64_t activation_bytes;
73 uint32_t num_layers;
74 uint32_t embed_dim;
75 uint32_t num_heads;
76 uint32_t vocab_size;
77 uint32_t max_seq_len;
78 uint32_t canary_count;
79 uint8_t reserved[8]; /* Pad to 64 bytes */
int vocab_size
Definition true_bpe.h:193

◆ ck_model_create()

void * ck_model_create ( void  )

Create and allocate model memory. Returns opaque model pointer, or NULL on failure.

Definition at line 8582 of file v6.6/test_generated/generic_api_test.c.

8582 {
8583 MODELModel *model = malloc(sizeof(MODELModel));
8584 if (!model) return NULL;
8585 if (model_model_allocate(model) != 0) {
8586 free(model);
8587 return NULL;
8588 }
8589 return model;
8590}
int model_model_allocate(MODELModel *model)

References model_model_allocate().

◆ ck_model_decode()

void ck_model_decode ( void *  model,
const int *  token,
int  token_index 
)

Decode single token at position token_index. Used for autoregressive generation.

Definition at line 8606 of file v6.6/test_generated/generic_api_test.c.

8606 {
8607 model_decode((MODELModel *)model, token, token_index);
8608}
const char * token
Definition tokenizer.h:307
void model_decode(MODELModel *model, const int *token, int token_index)

References model_decode(), and token.

◆ ck_model_forward()

void ck_model_forward ( void *  model,
const int *  tokens,
int  num_tokens 
)

Forward pass (prefill) - process multiple tokens. Used for initial prompt processing.

Definition at line 8602 of file v6.6/test_generated/generic_api_test.c.

8602 {
8603 model_forward((MODELModel *)model, tokens, num_tokens);
8604}
void model_forward(MODELModel *model, const int *tokens, int num_tokens)

References model_forward().

◆ ck_model_free()

void ck_model_free ( void *  model)

Free model memory.

Definition at line 8592 of file v6.6/test_generated/generic_api_test.c.

8592 {
8593 if (!model) return;
8594 model_model_free((MODELModel *)model);
8595 free(model);
8596}
void model_model_free(MODELModel *model)

References model_model_free().

◆ ck_model_get_base()

void * ck_model_get_base ( void *  model)

Get model base pointer (for weight loading).

Definition at line 8619 of file v6.6/test_generated/generic_api_test.c.

8619 {
8620 return ((MODELModel *)model)->base;
8621}

◆ ck_model_get_config()

const CKModelConfig * ck_model_get_config ( void  )

Get model configuration (dimensions, sizes, etc.) This is available before allocation.

Definition at line 8578 of file v6.6/test_generated/generic_api_test.c.

8578 {
8579 return &g_model_config;
8580}
static CKModelConfig g_model_config

References g_model_config.

◆ ck_model_get_logits()

float * ck_model_get_logits ( void *  model)

Get pointer to output logits buffer. Size is vocab_size floats.

Definition at line 8610 of file v6.6/test_generated/generic_api_test.c.

8610 {
8611 MODELModel *m = (MODELModel *)model;
8612 return MODEL_PTR(m, MODEL_FOOTER.logits);
8613}

◆ ck_model_get_total_bytes()

size_t ck_model_get_total_bytes ( void *  model)

Get total model size in bytes.

Definition at line 8623 of file v6.6/test_generated/generic_api_test.c.

8623 {
8624 return ((MODELModel *)model)->total_bytes;
8625}

◆ ck_model_precompute_rope()

void ck_model_precompute_rope ( void *  model)

Precompute RoPE cos/sin caches. Call once after allocation, before inference.

Definition at line 8598 of file v6.6/test_generated/generic_api_test.c.

8598 {
8599 model_precompute_rope((MODELModel *)model);
8600}
void model_precompute_rope(MODELModel *model)

References model_precompute_rope().

◆ ck_model_verify_canaries()

int ck_model_verify_canaries ( void *  model)

Verify memory canaries (debug). Returns number of corrupted canaries (0 = OK).

Definition at line 8615 of file v6.6/test_generated/generic_api_test.c.

8615 {
8616 return model_verify_canaries((MODELModel *)model);
8617}
int model_verify_canaries(MODELModel *model)

References model_verify_canaries().

◆ model_align_elems()

static int model_align_elems ( int  elems,
int  elem_bytes,
int  align_bytes 
)
static

Definition at line 176 of file v6.6/test_generated/generic_api_test.c.

176 {
177 int bytes = elems * elem_bytes;
178 int aligned = (bytes + align_bytes - 1) / align_bytes * align_bytes;
179 return aligned / elem_bytes;
180}

◆ model_decode()

void model_decode ( MODELModel *  model,
const int *  token,
int  token_index 
)

Definition at line 8550 of file v6.6/test_generated/generic_api_test.c.

8550 {
8551 model_decode_token(model, token, token_index);
8552}
static void model_decode_token(MODELModel *model, const int *token, int token_index)

References model_decode_token(), and token.

Referenced by ck_model_decode().

◆ model_decode_token()

static void model_decode_token ( MODELModel *  model,
const int *  token,
int  token_index 
)
static

Definition at line 8462 of file v6.6/test_generated/generic_api_test.c.

8466 {
8467 if (!model || !token) return;
8468
8469 const int aligned_embed_dim = 1024;
8470 const int aligned_head_dim = 64;
8471 const int aligned_intermediate_dim = 4864;
8472 const int aligned_context_window = 131072;
8473
8474 if (token_index < 0 || token_index >= aligned_context_window) return;
8475
8476 /* Embedding lookup */
8477 float *embed_out = MODEL_PTR(model, MODEL_HEADER.embedded_input);
8478 const void *embed_weight = (const void *)MODEL_PTR(model, MODEL_HEADER.token_emb);
8479 /* Embedding: Q4_K -> embedding_forward_q4_k */
8480 embedding_forward_q4_k((const int32_t *)token,
8481 1,
8483 embed_weight,
8484 NULL,
8485 embed_out,
8487 aligned_embed_dim,
8488 1,
8489 0);
8490
8491 /* Process each layer explicitly */
8492 model_layer_0_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8493 model_layer_1_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8494 model_layer_2_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8495 model_layer_3_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8496 model_layer_4_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8497 model_layer_5_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8498 model_layer_6_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8499 model_layer_7_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8500 model_layer_8_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8501 model_layer_9_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8502 model_layer_10_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8503 model_layer_11_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8504 model_layer_12_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8505 model_layer_13_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8506 model_layer_14_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8507 model_layer_15_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8508 model_layer_16_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8509 model_layer_17_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8510 model_layer_18_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8511 model_layer_19_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8512 model_layer_20_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8513 model_layer_21_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8514 model_layer_22_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8515 model_layer_23_decode(model, token_index, aligned_embed_dim, aligned_head_dim, aligned_intermediate_dim, aligned_context_window);
8516
8517 /* Final RMSNorm */
8518 float *last_hidden = MODEL_PTR(model, MODEL_LAYERS[23].output);
8519 float *final_ln_weight = MODEL_PTR(model, MODEL_FOOTER.final_ln_weight);
8520 float *final_out = MODEL_PTR(model, MODEL_FOOTER.final_output);
8521 rmsnorm_forward(last_hidden,
8522 final_ln_weight,
8523 final_out,
8524 NULL,
8525 1,
8527 aligned_embed_dim,
8528 1e-06f);
8529
8530 /* LM head projection */
8531 float *logits = MODEL_PTR(model, MODEL_FOOTER.logits);
8532 const void *lm_head = (const void *)MODEL_PTR(model, MODEL_FOOTER.lm_head_weight);
8533 /* LM head: Q4_K -> gemm_nt_q4_k */
8534 gemm_nt_q4_k(final_out, lm_head, NULL, logits, 1, MODEL_VOCAB_SIZE, aligned_embed_dim);
8535}
void gemm_nt_q4_k(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void embedding_forward_q4_k(const int32_t *token_ids, int token_count, int vocab_size, const void *token_embeddings, const float *pos_embeddings, float *output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)
void rmsnorm_forward(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
#define MODEL_VOCAB_SIZE
Definition v6.5_simple.c:30
#define MODEL_EMBED_DIM
Definition v6.5_simple.c:24
static void model_layer_13_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_6_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_3_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_0_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_15_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_22_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_8_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_11_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_16_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_4_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_1_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_20_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_19_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_14_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_7_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_21_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_18_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_5_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_10_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_12_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_17_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_23_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_2_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_9_decode(MODELModel *model, int token_index, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

References embedding_forward_q4_k(), gemm_nt_q4_k(), MODEL_EMBED_DIM, model_layer_0_decode(), model_layer_10_decode(), model_layer_11_decode(), model_layer_12_decode(), model_layer_13_decode(), model_layer_14_decode(), model_layer_15_decode(), model_layer_16_decode(), model_layer_17_decode(), model_layer_18_decode(), model_layer_19_decode(), model_layer_1_decode(), model_layer_20_decode(), model_layer_21_decode(), model_layer_22_decode(), model_layer_23_decode(), model_layer_2_decode(), model_layer_3_decode(), model_layer_4_decode(), model_layer_5_decode(), model_layer_6_decode(), model_layer_7_decode(), model_layer_8_decode(), model_layer_9_decode(), MODEL_VOCAB_SIZE, rmsnorm_forward(), and token.

Referenced by model_decode().

◆ model_forward()

void model_forward ( MODELModel *  model,
const int *  tokens,
int  num_tokens 
)

Definition at line 8541 of file v6.6/test_generated/generic_api_test.c.

8545 {
8546 if (!model || !tokens || num_tokens <= 0) return;
8547 model_forward_prefill_impl(model, tokens, num_tokens);
8548}
static void model_forward_prefill_impl(MODELModel *model, const int *tokens, int num_tokens)

References model_forward_prefill_impl().

Referenced by ck_model_forward().

◆ model_forward_prefill_impl()

static void model_forward_prefill_impl ( MODELModel *  model,
const int *  tokens,
int  num_tokens 
)
static

Definition at line 4148 of file v6.6/test_generated/generic_api_test.c.

4152 {
4153 if (!model || !tokens || num_tokens <= 0) {
4154 return;
4155 }
4156
4157 const int elem_bytes = MODEL_DTYPE_BYTES;
4158 const int aligned_embed_dim = 1024;
4159 const int aligned_head_dim = 64;
4160 const int aligned_intermediate_dim = 4864;
4161 const int aligned_context_window = 131072;
4162
4163 float *embed_out = MODEL_PTR(model, MODEL_HEADER.embedded_input);
4164 const void *embed_weight = (const void *)MODEL_PTR(model, MODEL_HEADER.token_emb);
4165 embedding_forward_q4_k((const int32_t *)tokens,
4166 num_tokens,
4168 embed_weight,
4169 NULL,
4170 embed_out,
4172 aligned_embed_dim,
4173 num_tokens,
4174 0);
4175
4177 model,
4178 num_tokens,
4179 aligned_embed_dim,
4180 aligned_head_dim,
4181 aligned_intermediate_dim,
4182 aligned_context_window);
4183
4185 model,
4186 num_tokens,
4187 aligned_embed_dim,
4188 aligned_head_dim,
4189 aligned_intermediate_dim,
4190 aligned_context_window);
4191
4193 model,
4194 num_tokens,
4195 aligned_embed_dim,
4196 aligned_head_dim,
4197 aligned_intermediate_dim,
4198 aligned_context_window);
4199
4201 model,
4202 num_tokens,
4203 aligned_embed_dim,
4204 aligned_head_dim,
4205 aligned_intermediate_dim,
4206 aligned_context_window);
4207
4209 model,
4210 num_tokens,
4211 aligned_embed_dim,
4212 aligned_head_dim,
4213 aligned_intermediate_dim,
4214 aligned_context_window);
4215
4217 model,
4218 num_tokens,
4219 aligned_embed_dim,
4220 aligned_head_dim,
4221 aligned_intermediate_dim,
4222 aligned_context_window);
4223
4225 model,
4226 num_tokens,
4227 aligned_embed_dim,
4228 aligned_head_dim,
4229 aligned_intermediate_dim,
4230 aligned_context_window);
4231
4233 model,
4234 num_tokens,
4235 aligned_embed_dim,
4236 aligned_head_dim,
4237 aligned_intermediate_dim,
4238 aligned_context_window);
4239
4241 model,
4242 num_tokens,
4243 aligned_embed_dim,
4244 aligned_head_dim,
4245 aligned_intermediate_dim,
4246 aligned_context_window);
4247
4249 model,
4250 num_tokens,
4251 aligned_embed_dim,
4252 aligned_head_dim,
4253 aligned_intermediate_dim,
4254 aligned_context_window);
4255
4257 model,
4258 num_tokens,
4259 aligned_embed_dim,
4260 aligned_head_dim,
4261 aligned_intermediate_dim,
4262 aligned_context_window);
4263
4265 model,
4266 num_tokens,
4267 aligned_embed_dim,
4268 aligned_head_dim,
4269 aligned_intermediate_dim,
4270 aligned_context_window);
4271
4273 model,
4274 num_tokens,
4275 aligned_embed_dim,
4276 aligned_head_dim,
4277 aligned_intermediate_dim,
4278 aligned_context_window);
4279
4281 model,
4282 num_tokens,
4283 aligned_embed_dim,
4284 aligned_head_dim,
4285 aligned_intermediate_dim,
4286 aligned_context_window);
4287
4289 model,
4290 num_tokens,
4291 aligned_embed_dim,
4292 aligned_head_dim,
4293 aligned_intermediate_dim,
4294 aligned_context_window);
4295
4297 model,
4298 num_tokens,
4299 aligned_embed_dim,
4300 aligned_head_dim,
4301 aligned_intermediate_dim,
4302 aligned_context_window);
4303
4305 model,
4306 num_tokens,
4307 aligned_embed_dim,
4308 aligned_head_dim,
4309 aligned_intermediate_dim,
4310 aligned_context_window);
4311
4313 model,
4314 num_tokens,
4315 aligned_embed_dim,
4316 aligned_head_dim,
4317 aligned_intermediate_dim,
4318 aligned_context_window);
4319
4321 model,
4322 num_tokens,
4323 aligned_embed_dim,
4324 aligned_head_dim,
4325 aligned_intermediate_dim,
4326 aligned_context_window);
4327
4329 model,
4330 num_tokens,
4331 aligned_embed_dim,
4332 aligned_head_dim,
4333 aligned_intermediate_dim,
4334 aligned_context_window);
4335
4337 model,
4338 num_tokens,
4339 aligned_embed_dim,
4340 aligned_head_dim,
4341 aligned_intermediate_dim,
4342 aligned_context_window);
4343
4345 model,
4346 num_tokens,
4347 aligned_embed_dim,
4348 aligned_head_dim,
4349 aligned_intermediate_dim,
4350 aligned_context_window);
4351
4353 model,
4354 num_tokens,
4355 aligned_embed_dim,
4356 aligned_head_dim,
4357 aligned_intermediate_dim,
4358 aligned_context_window);
4359
4361 model,
4362 num_tokens,
4363 aligned_embed_dim,
4364 aligned_head_dim,
4365 aligned_intermediate_dim,
4366 aligned_context_window);
4367
4368 float *last_hidden = MODEL_PTR(model, MODEL_LAYERS[MODEL_NUM_LAYERS - 1].output);
4369 float *final_ln_weight = MODEL_PTR(model, MODEL_FOOTER.final_ln_weight);
4370 float *final_out = MODEL_PTR(model, MODEL_FOOTER.final_output);
4371 rmsnorm_forward(last_hidden,
4372 final_ln_weight,
4373 final_out,
4374 NULL,
4375 num_tokens,
4377 aligned_embed_dim,
4378 1e-06f);
4379
4380 float *logits = MODEL_PTR(model, MODEL_FOOTER.logits);
4381 const void *lm_head = (const void *)MODEL_PTR(model, MODEL_FOOTER.lm_head_weight);
4382 const size_t q8_bytes = ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)aligned_embed_dim);
4383 for (int t = 0; t < num_tokens; ++t) {
4384 uint8_t q8_buf[q8_bytes];
4385 const float *row = final_out + (size_t)t * (size_t)aligned_embed_dim;
4386 float *logits_row = logits + (size_t)t * (size_t)MODEL_VOCAB_SIZE;
4387 quantize_row_q8_k(row, q8_buf, aligned_embed_dim);
4388 gemm_nt_q4_k_q8_k(q8_buf,
4389 lm_head,
4390 NULL,
4391 logits_row,
4392 1,
4394 aligned_embed_dim);
4395 }
4396}
@ CK_DT_Q8_K
static size_t ck_dtype_row_bytes(CKDataType dt, size_t n_elements)
Calculate total bytes for n_elements of given dtype.
void gemm_nt_q4_k_q8_k(const void *A_q8, const void *B, const float *bias, float *C, int M, int N, int K)
void quantize_row_q8_k(const float *x, void *y, int k)
#define MODEL_NUM_LAYERS
Definition v6.5_simple.c:25
static void model_layer_13_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_17_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_8_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_22_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_18_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_7_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_4_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_12_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_10_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_21_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_16_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_14_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_3_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_20_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_11_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_23_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_5_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_15_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_0_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_1_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_19_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_2_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_9_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)
static void model_layer_6_prefill(MODELModel *model, int num_tokens, int aligned_embed_dim, int aligned_head_dim, int aligned_intermediate_dim, int aligned_context_window)

References CK_DT_Q8_K, ck_dtype_row_bytes(), embedding_forward_q4_k(), gemm_nt_q4_k_q8_k(), MODEL_EMBED_DIM, model_layer_0_prefill(), model_layer_10_prefill(), model_layer_11_prefill(), model_layer_12_prefill(), model_layer_13_prefill(), model_layer_14_prefill(), model_layer_15_prefill(), model_layer_16_prefill(), model_layer_17_prefill(), model_layer_18_prefill(), model_layer_19_prefill(), model_layer_1_prefill(), model_layer_20_prefill(), model_layer_21_prefill(), model_layer_22_prefill(), model_layer_23_prefill(), model_layer_2_prefill(), model_layer_3_prefill(), model_layer_4_prefill(), model_layer_5_prefill(), model_layer_6_prefill(), model_layer_7_prefill(), model_layer_8_prefill(), model_layer_9_prefill(), MODEL_NUM_LAYERS, MODEL_VOCAB_SIZE, quantize_row_q8_k(), and rmsnorm_forward().

Referenced by model_forward().

◆ model_layer_0_decode()

static void model_layer_0_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 4405 of file v6.6/test_generated/generic_api_test.c.

4412 {
4413 const MODELLayerOffsets *L = &MODEL_LAYERS[0];
4414
4415 float *input = MODEL_PTR(model, MODEL_HEADER.embedded_input);
4416
4417 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
4418 float *ln1_out = MODEL_PTR(model, L->ln1_out);
4419 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
4420 float *ln2_out = MODEL_PTR(model, L->ln2_out);
4421 float *k_cache = MODEL_PTR(model, L->k);
4422 float *v_cache = MODEL_PTR(model, L->v);
4423 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
4424 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
4425 float *residual1 = MODEL_PTR(model, L->residual1);
4426 float *mlp_out = MODEL_PTR(model, L->mlp_out);
4427 float *output = MODEL_PTR(model, L->output);
4428
4429 /* Weights (explicit types for layer 0) */
4430 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
4431 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
4432 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
4433 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
4434 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
4435 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
4436
4437 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
4438 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
4439
4440 const int H = MODEL_NUM_HEADS;
4441 const int H_kv = MODEL_NUM_KV_HEADS;
4442 const int head_dim = MODEL_HEAD_DIM;
4443
4444 float q_token[H * aligned_head_dim];
4445 float k_token[H_kv * aligned_head_dim];
4446 float v_token[H_kv * aligned_head_dim];
4447 float attn_token[H * aligned_head_dim];
4448
4449 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
4450 float fc1_out[2 * aligned_intermediate_dim];
4451 float swiglu_out[aligned_intermediate_dim];
4452
4453 /* Step 1: RMSNorm before attention */
4454 rmsnorm_forward(input,
4455 ln1_gamma,
4456 ln1_out,
4457 NULL,
4458 1,
4460 aligned_embed_dim,
4461 1e-06f);
4462
4463 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
4464
4465 /* Step 2: QKV projection */
4466 /* Q projection: Q4_K -> gemm_nt_q4_k */
4467 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
4468 if (aligned_head_dim > head_dim) {
4469 for (int h = 0; h < H; ++h) {
4470 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
4471 for (int d = head_dim; d < aligned_head_dim; ++d) {
4472 q_head[d] = 0.0f;
4473 }
4474 }
4475 }
4476
4477 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4478 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4479 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
4480 const uint8_t *WK_bytes = (const uint8_t *)WK;
4481 for (int h = 0; h < H_kv; ++h) {
4482 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
4483 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
4484 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4485 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
4486 for (int d = head_dim; d < aligned_head_dim; ++d) {
4487 k_head[d] = 0.0f;
4488 }
4489 }
4490
4491 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4492 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4493 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
4494 const uint8_t *WV_bytes = (const uint8_t *)WV;
4495 for (int h = 0; h < H_kv; ++h) {
4496 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
4497 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
4498 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4499 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
4500 for (int d = head_dim; d < aligned_head_dim; ++d) {
4501 v_head[d] = 0.0f;
4502 }
4503 }
4504
4505 /* Step 3: RoPE */
4506 rope_forward(q_token,
4507 rope_cos,
4508 rope_sin,
4509 H,
4510 1,
4511 head_dim,
4512 aligned_head_dim,
4513 token_index);
4514 for (int h = 0; h < H_kv; ++h) {
4515 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4516 rope_forward(k_head,
4517 rope_cos,
4518 rope_sin,
4519 1,
4520 1,
4521 head_dim,
4522 aligned_head_dim,
4523 token_index);
4524 }
4525
4526 /* Step 4: KV cache write (direct-to-cache) */
4527
4528 /* Step 5: Attention (decode, flash) */
4530 k_cache,
4531 v_cache,
4532 attn_token,
4533 H,
4534 H_kv,
4535 token_index + 1,
4536 aligned_context_window,
4537 head_dim,
4538 aligned_head_dim);
4539
4540 /* Step 6: Output projection */
4541 /* WO projection: Q4_K -> gemm_nt_q4_k */
4542 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
4543
4544 /* Step 7: Residual add */
4545 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
4546
4547 /* Step 8: RMSNorm before MLP */
4548 rmsnorm_forward(residual1,
4549 ln2_gamma,
4550 ln2_out,
4551 NULL,
4552 1,
4554 aligned_embed_dim,
4555 1e-06f);
4556
4557 /* Step 9: MLP (SwiGLU) */
4558 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
4559 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
4560
4561 /* SwiGLU activation */
4562 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
4563
4564 /* Down projection: Q4_K -> gemm_nt_q4_k */
4565 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
4566
4567 /* Step 10: Final residual add */
4568 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
4569}
@ CK_DT_Q4_K
void swiglu_forward(const float *input, float *output, int tokens, int dim)
void rope_forward(float *x, const float *cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)
void attention_forward_decode_head_major_gqa_flash(const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)
#define MODEL_HEAD_DIM
Definition v6.5_simple.c:28
#define MODEL_NUM_KV_HEADS
Definition v6.5_simple.c:27
#define MODEL_NUM_HEADS
Definition v6.5_simple.c:26
static void model_residual_add_token_major(const float *a, const float *b, float *out, int tokens, int aligned_embed_dim)

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_0_prefill()

static void model_layer_0_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 211 of file v6.6/test_generated/generic_api_test.c.

218 {
219 const MODELLayerOffsets *L = &MODEL_LAYERS[0];
220
221 float *input = MODEL_PTR(model, MODEL_HEADER.embedded_input);
222 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
223 float *ln1_out = MODEL_PTR(model, L->ln1_out);
224 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
225 float *ln2_out = MODEL_PTR(model, L->ln2_out);
226 float *q = MODEL_PTR(model, L->q);
227 float *k = MODEL_PTR(model, L->k);
228 float *v = MODEL_PTR(model, L->v);
229 float *attn_out = MODEL_PTR(model, L->attn_out);
230 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
231 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
232 float *residual1 = MODEL_PTR(model, L->residual1);
233 float *fc1_out = MODEL_PTR(model, L->fc1_out);
234 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
235 float *mlp_out = MODEL_PTR(model, L->mlp_out);
236 float *output = MODEL_PTR(model, L->output);
237
238 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
239 const void *WK = (const void *)MODEL_PTR(model, L->wk);
240 const void *WV = (const void *)MODEL_PTR(model, L->wv);
241 const void *WO = (const void *)MODEL_PTR(model, L->wo);
242 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
243 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
244 const float *BQ = NULL;
245 const float *BK = NULL;
246 const float *BV = NULL;
247 const float *BO = NULL;
248 const float *B1 = NULL;
249 const float *B2 = NULL;
250
251 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
252 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
253
254 const int H = MODEL_NUM_HEADS;
255 const int H_kv = MODEL_NUM_KV_HEADS;
256 const int head_dim = MODEL_HEAD_DIM;
257 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
258 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
259 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
260
261 /* RMSNorm before attention */
262 rmsnorm_forward(input,
263 ln1_gamma,
264 ln1_out,
265 NULL,
266 num_tokens,
268 aligned_embed_dim,
269 1e-06f);
270
271 /* Q projection (head-major) */
272 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
273 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
274 for (int h = 0; h < H; ++h) {
275 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
276 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
277 float *q_h = q + (size_t)h * q_head_stride;
278 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
279 }
280
281 /* K projection (head-major) */
282 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
283 const uint8_t *WK_bytes = (const uint8_t *)WK;
284 for (int h = 0; h < H_kv; ++h) {
285 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
286 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
287 float *k_h = k + (size_t)h * kv_head_stride;
288 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
289 }
290
291 /* V projection (head-major) */
292 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
293 const uint8_t *WV_bytes = (const uint8_t *)WV;
294 for (int h = 0; h < H_kv; ++h) {
295 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
296 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
297 float *v_h = v + (size_t)h * kv_head_stride;
298 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
299 }
300
301 /* RoPE */
303 k,
304 rope_cos,
305 rope_sin,
306 H,
307 H_kv,
308 num_tokens,
309 head_dim,
310 aligned_head_dim,
311 0,
312 num_tokens,
313 aligned_context_window);
314
315 /* Attention (prefill, causal) */
317 k,
318 v,
319 attn_out,
320 H,
321 H_kv,
322 num_tokens,
323 head_dim,
324 aligned_head_dim,
325 aligned_context_window);
326
327 /* Output projection (flatten head-major to token-major) */
328 const int K = H * aligned_head_dim;
329 if (K != aligned_embed_dim) {
330 return;
331 }
332 const float *proj_in = attn_out;
333 if (H > 1) {
334 if (!proj_scratch) {
335 return;
336 }
337 for (int t = 0; t < num_tokens; ++t) {
338 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
339 for (int h = 0; h < H; ++h) {
340 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
341 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
342 src,
343 (size_t)aligned_head_dim * sizeof(float));
344 }
345 }
346 proj_in = proj_scratch;
347 }
348 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
349
350 /* Residual add */
351 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
352
353 /* RMSNorm before MLP */
354 rmsnorm_forward(residual1,
355 ln2_gamma,
356 ln2_out,
357 NULL,
358 num_tokens,
360 aligned_embed_dim,
361 1e-06f);
362
363 /* MLP (SwiGLU) */
364 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
365 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
366 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
367
368 /* Final residual add */
369 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
370}
void attention_forward_causal_head_major_gqa_flash_strided(const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)
void rope_forward_qk_strided(float *q, float *k, const float *cos_cache, const float *sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int q_stride_tokens, int k_stride_tokens)

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_10_decode()

static void model_layer_10_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 6095 of file v6.6/test_generated/generic_api_test.c.

6102 {
6103 const MODELLayerOffsets *L = &MODEL_LAYERS[10];
6104
6105 float *input = MODEL_PTR(model, MODEL_LAYERS[9].output);
6106
6107 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
6108 float *ln1_out = MODEL_PTR(model, L->ln1_out);
6109 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
6110 float *ln2_out = MODEL_PTR(model, L->ln2_out);
6111 float *k_cache = MODEL_PTR(model, L->k);
6112 float *v_cache = MODEL_PTR(model, L->v);
6113 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
6114 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
6115 float *residual1 = MODEL_PTR(model, L->residual1);
6116 float *mlp_out = MODEL_PTR(model, L->mlp_out);
6117 float *output = MODEL_PTR(model, L->output);
6118
6119 /* Weights (explicit types for layer 10) */
6120 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
6121 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
6122 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
6123 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
6124 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
6125 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
6126
6127 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
6128 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
6129
6130 const int H = MODEL_NUM_HEADS;
6131 const int H_kv = MODEL_NUM_KV_HEADS;
6132 const int head_dim = MODEL_HEAD_DIM;
6133
6134 float q_token[H * aligned_head_dim];
6135 float k_token[H_kv * aligned_head_dim];
6136 float v_token[H_kv * aligned_head_dim];
6137 float attn_token[H * aligned_head_dim];
6138
6139 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
6140 float fc1_out[2 * aligned_intermediate_dim];
6141 float swiglu_out[aligned_intermediate_dim];
6142
6143 /* Step 1: RMSNorm before attention */
6144 rmsnorm_forward(input,
6145 ln1_gamma,
6146 ln1_out,
6147 NULL,
6148 1,
6150 aligned_embed_dim,
6151 1e-06f);
6152
6153 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
6154
6155 /* Step 2: QKV projection */
6156 /* Q projection: Q4_K -> gemm_nt_q4_k */
6157 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
6158 if (aligned_head_dim > head_dim) {
6159 for (int h = 0; h < H; ++h) {
6160 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
6161 for (int d = head_dim; d < aligned_head_dim; ++d) {
6162 q_head[d] = 0.0f;
6163 }
6164 }
6165 }
6166
6167 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6168 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6169 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
6170 const uint8_t *WK_bytes = (const uint8_t *)WK;
6171 for (int h = 0; h < H_kv; ++h) {
6172 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
6173 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
6174 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6175 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
6176 for (int d = head_dim; d < aligned_head_dim; ++d) {
6177 k_head[d] = 0.0f;
6178 }
6179 }
6180
6181 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6182 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6183 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
6184 const uint8_t *WV_bytes = (const uint8_t *)WV;
6185 for (int h = 0; h < H_kv; ++h) {
6186 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
6187 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
6188 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6189 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
6190 for (int d = head_dim; d < aligned_head_dim; ++d) {
6191 v_head[d] = 0.0f;
6192 }
6193 }
6194
6195 /* Step 3: RoPE */
6196 rope_forward(q_token,
6197 rope_cos,
6198 rope_sin,
6199 H,
6200 1,
6201 head_dim,
6202 aligned_head_dim,
6203 token_index);
6204 for (int h = 0; h < H_kv; ++h) {
6205 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6206 rope_forward(k_head,
6207 rope_cos,
6208 rope_sin,
6209 1,
6210 1,
6211 head_dim,
6212 aligned_head_dim,
6213 token_index);
6214 }
6215
6216 /* Step 4: KV cache write (direct-to-cache) */
6217
6218 /* Step 5: Attention (decode, flash) */
6220 k_cache,
6221 v_cache,
6222 attn_token,
6223 H,
6224 H_kv,
6225 token_index + 1,
6226 aligned_context_window,
6227 head_dim,
6228 aligned_head_dim);
6229
6230 /* Step 6: Output projection */
6231 /* WO projection: Q4_K -> gemm_nt_q4_k */
6232 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
6233
6234 /* Step 7: Residual add */
6235 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
6236
6237 /* Step 8: RMSNorm before MLP */
6238 rmsnorm_forward(residual1,
6239 ln2_gamma,
6240 ln2_out,
6241 NULL,
6242 1,
6244 aligned_embed_dim,
6245 1e-06f);
6246
6247 /* Step 9: MLP (SwiGLU) */
6248 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
6249 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
6250
6251 /* SwiGLU activation */
6252 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
6253
6254 /* Down projection: Q4_K -> gemm_nt_q4_k */
6255 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
6256
6257 /* Step 10: Final residual add */
6258 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
6259}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_10_prefill()

static void model_layer_10_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 1851 of file v6.6/test_generated/generic_api_test.c.

1858 {
1859 const MODELLayerOffsets *L = &MODEL_LAYERS[10];
1860
1861 float *input = MODEL_PTR(model, MODEL_LAYERS[9].output);
1862 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
1863 float *ln1_out = MODEL_PTR(model, L->ln1_out);
1864 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
1865 float *ln2_out = MODEL_PTR(model, L->ln2_out);
1866 float *q = MODEL_PTR(model, L->q);
1867 float *k = MODEL_PTR(model, L->k);
1868 float *v = MODEL_PTR(model, L->v);
1869 float *attn_out = MODEL_PTR(model, L->attn_out);
1870 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
1871 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
1872 float *residual1 = MODEL_PTR(model, L->residual1);
1873 float *fc1_out = MODEL_PTR(model, L->fc1_out);
1874 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
1875 float *mlp_out = MODEL_PTR(model, L->mlp_out);
1876 float *output = MODEL_PTR(model, L->output);
1877
1878 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
1879 const void *WK = (const void *)MODEL_PTR(model, L->wk);
1880 const void *WV = (const void *)MODEL_PTR(model, L->wv);
1881 const void *WO = (const void *)MODEL_PTR(model, L->wo);
1882 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
1883 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
1884 const float *BQ = NULL;
1885 const float *BK = NULL;
1886 const float *BV = NULL;
1887 const float *BO = NULL;
1888 const float *B1 = NULL;
1889 const float *B2 = NULL;
1890
1891 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
1892 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
1893
1894 const int H = MODEL_NUM_HEADS;
1895 const int H_kv = MODEL_NUM_KV_HEADS;
1896 const int head_dim = MODEL_HEAD_DIM;
1897 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
1898 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
1899 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
1900
1901 /* RMSNorm before attention */
1902 rmsnorm_forward(input,
1903 ln1_gamma,
1904 ln1_out,
1905 NULL,
1906 num_tokens,
1908 aligned_embed_dim,
1909 1e-06f);
1910
1911 /* Q projection (head-major) */
1912 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1913 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
1914 for (int h = 0; h < H; ++h) {
1915 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
1916 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
1917 float *q_h = q + (size_t)h * q_head_stride;
1918 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1919 }
1920
1921 /* K projection (head-major) */
1922 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1923 const uint8_t *WK_bytes = (const uint8_t *)WK;
1924 for (int h = 0; h < H_kv; ++h) {
1925 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
1926 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
1927 float *k_h = k + (size_t)h * kv_head_stride;
1928 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1929 }
1930
1931 /* V projection (head-major) */
1932 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1933 const uint8_t *WV_bytes = (const uint8_t *)WV;
1934 for (int h = 0; h < H_kv; ++h) {
1935 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
1936 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
1937 float *v_h = v + (size_t)h * kv_head_stride;
1938 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1939 }
1940
1941 /* RoPE */
1943 k,
1944 rope_cos,
1945 rope_sin,
1946 H,
1947 H_kv,
1948 num_tokens,
1949 head_dim,
1950 aligned_head_dim,
1951 0,
1952 num_tokens,
1953 aligned_context_window);
1954
1955 /* Attention (prefill, causal) */
1957 k,
1958 v,
1959 attn_out,
1960 H,
1961 H_kv,
1962 num_tokens,
1963 head_dim,
1964 aligned_head_dim,
1965 aligned_context_window);
1966
1967 /* Output projection (flatten head-major to token-major) */
1968 const int K = H * aligned_head_dim;
1969 if (K != aligned_embed_dim) {
1970 return;
1971 }
1972 const float *proj_in = attn_out;
1973 if (H > 1) {
1974 if (!proj_scratch) {
1975 return;
1976 }
1977 for (int t = 0; t < num_tokens; ++t) {
1978 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
1979 for (int h = 0; h < H; ++h) {
1980 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
1981 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
1982 src,
1983 (size_t)aligned_head_dim * sizeof(float));
1984 }
1985 }
1986 proj_in = proj_scratch;
1987 }
1988 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1989
1990 /* Residual add */
1991 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1992
1993 /* RMSNorm before MLP */
1994 rmsnorm_forward(residual1,
1995 ln2_gamma,
1996 ln2_out,
1997 NULL,
1998 num_tokens,
2000 aligned_embed_dim,
2001 1e-06f);
2002
2003 /* MLP (SwiGLU) */
2004 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2005 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2006 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2007
2008 /* Final residual add */
2009 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2010}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_11_decode()

static void model_layer_11_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 6264 of file v6.6/test_generated/generic_api_test.c.

6271 {
6272 const MODELLayerOffsets *L = &MODEL_LAYERS[11];
6273
6274 float *input = MODEL_PTR(model, MODEL_LAYERS[10].output);
6275
6276 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
6277 float *ln1_out = MODEL_PTR(model, L->ln1_out);
6278 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
6279 float *ln2_out = MODEL_PTR(model, L->ln2_out);
6280 float *k_cache = MODEL_PTR(model, L->k);
6281 float *v_cache = MODEL_PTR(model, L->v);
6282 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
6283 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
6284 float *residual1 = MODEL_PTR(model, L->residual1);
6285 float *mlp_out = MODEL_PTR(model, L->mlp_out);
6286 float *output = MODEL_PTR(model, L->output);
6287
6288 /* Weights (explicit types for layer 11) */
6289 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
6290 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
6291 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
6292 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
6293 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
6294 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
6295
6296 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
6297 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
6298
6299 const int H = MODEL_NUM_HEADS;
6300 const int H_kv = MODEL_NUM_KV_HEADS;
6301 const int head_dim = MODEL_HEAD_DIM;
6302
6303 float q_token[H * aligned_head_dim];
6304 float k_token[H_kv * aligned_head_dim];
6305 float v_token[H_kv * aligned_head_dim];
6306 float attn_token[H * aligned_head_dim];
6307
6308 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
6309 float fc1_out[2 * aligned_intermediate_dim];
6310 float swiglu_out[aligned_intermediate_dim];
6311
6312 /* Step 1: RMSNorm before attention */
6313 rmsnorm_forward(input,
6314 ln1_gamma,
6315 ln1_out,
6316 NULL,
6317 1,
6319 aligned_embed_dim,
6320 1e-06f);
6321
6322 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
6323
6324 /* Step 2: QKV projection */
6325 /* Q projection: Q4_K -> gemm_nt_q4_k */
6326 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
6327 if (aligned_head_dim > head_dim) {
6328 for (int h = 0; h < H; ++h) {
6329 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
6330 for (int d = head_dim; d < aligned_head_dim; ++d) {
6331 q_head[d] = 0.0f;
6332 }
6333 }
6334 }
6335
6336 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6337 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6338 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
6339 const uint8_t *WK_bytes = (const uint8_t *)WK;
6340 for (int h = 0; h < H_kv; ++h) {
6341 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
6342 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
6343 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6344 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
6345 for (int d = head_dim; d < aligned_head_dim; ++d) {
6346 k_head[d] = 0.0f;
6347 }
6348 }
6349
6350 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6351 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6352 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
6353 const uint8_t *WV_bytes = (const uint8_t *)WV;
6354 for (int h = 0; h < H_kv; ++h) {
6355 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
6356 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
6357 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6358 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
6359 for (int d = head_dim; d < aligned_head_dim; ++d) {
6360 v_head[d] = 0.0f;
6361 }
6362 }
6363
6364 /* Step 3: RoPE */
6365 rope_forward(q_token,
6366 rope_cos,
6367 rope_sin,
6368 H,
6369 1,
6370 head_dim,
6371 aligned_head_dim,
6372 token_index);
6373 for (int h = 0; h < H_kv; ++h) {
6374 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6375 rope_forward(k_head,
6376 rope_cos,
6377 rope_sin,
6378 1,
6379 1,
6380 head_dim,
6381 aligned_head_dim,
6382 token_index);
6383 }
6384
6385 /* Step 4: KV cache write (direct-to-cache) */
6386
6387 /* Step 5: Attention (decode, flash) */
6389 k_cache,
6390 v_cache,
6391 attn_token,
6392 H,
6393 H_kv,
6394 token_index + 1,
6395 aligned_context_window,
6396 head_dim,
6397 aligned_head_dim);
6398
6399 /* Step 6: Output projection */
6400 /* WO projection: Q4_K -> gemm_nt_q4_k */
6401 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
6402
6403 /* Step 7: Residual add */
6404 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
6405
6406 /* Step 8: RMSNorm before MLP */
6407 rmsnorm_forward(residual1,
6408 ln2_gamma,
6409 ln2_out,
6410 NULL,
6411 1,
6413 aligned_embed_dim,
6414 1e-06f);
6415
6416 /* Step 9: MLP (SwiGLU) */
6417 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
6418 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
6419
6420 /* SwiGLU activation */
6421 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
6422
6423 /* Down projection: Q4_K -> gemm_nt_q4_k */
6424 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
6425
6426 /* Step 10: Final residual add */
6427 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
6428}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_11_prefill()

static void model_layer_11_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2015 of file v6.6/test_generated/generic_api_test.c.

2022 {
2023 const MODELLayerOffsets *L = &MODEL_LAYERS[11];
2024
2025 float *input = MODEL_PTR(model, MODEL_LAYERS[10].output);
2026 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
2027 float *ln1_out = MODEL_PTR(model, L->ln1_out);
2028 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
2029 float *ln2_out = MODEL_PTR(model, L->ln2_out);
2030 float *q = MODEL_PTR(model, L->q);
2031 float *k = MODEL_PTR(model, L->k);
2032 float *v = MODEL_PTR(model, L->v);
2033 float *attn_out = MODEL_PTR(model, L->attn_out);
2034 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
2035 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
2036 float *residual1 = MODEL_PTR(model, L->residual1);
2037 float *fc1_out = MODEL_PTR(model, L->fc1_out);
2038 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
2039 float *mlp_out = MODEL_PTR(model, L->mlp_out);
2040 float *output = MODEL_PTR(model, L->output);
2041
2042 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
2043 const void *WK = (const void *)MODEL_PTR(model, L->wk);
2044 const void *WV = (const void *)MODEL_PTR(model, L->wv);
2045 const void *WO = (const void *)MODEL_PTR(model, L->wo);
2046 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
2047 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
2048 const float *BQ = NULL;
2049 const float *BK = NULL;
2050 const float *BV = NULL;
2051 const float *BO = NULL;
2052 const float *B1 = NULL;
2053 const float *B2 = NULL;
2054
2055 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
2056 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
2057
2058 const int H = MODEL_NUM_HEADS;
2059 const int H_kv = MODEL_NUM_KV_HEADS;
2060 const int head_dim = MODEL_HEAD_DIM;
2061 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
2062 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
2063 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
2064
2065 /* RMSNorm before attention */
2066 rmsnorm_forward(input,
2067 ln1_gamma,
2068 ln1_out,
2069 NULL,
2070 num_tokens,
2072 aligned_embed_dim,
2073 1e-06f);
2074
2075 /* Q projection (head-major) */
2076 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2077 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
2078 for (int h = 0; h < H; ++h) {
2079 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
2080 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
2081 float *q_h = q + (size_t)h * q_head_stride;
2082 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2083 }
2084
2085 /* K projection (head-major) */
2086 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2087 const uint8_t *WK_bytes = (const uint8_t *)WK;
2088 for (int h = 0; h < H_kv; ++h) {
2089 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
2090 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
2091 float *k_h = k + (size_t)h * kv_head_stride;
2092 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2093 }
2094
2095 /* V projection (head-major) */
2096 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2097 const uint8_t *WV_bytes = (const uint8_t *)WV;
2098 for (int h = 0; h < H_kv; ++h) {
2099 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
2100 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
2101 float *v_h = v + (size_t)h * kv_head_stride;
2102 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2103 }
2104
2105 /* RoPE */
2107 k,
2108 rope_cos,
2109 rope_sin,
2110 H,
2111 H_kv,
2112 num_tokens,
2113 head_dim,
2114 aligned_head_dim,
2115 0,
2116 num_tokens,
2117 aligned_context_window);
2118
2119 /* Attention (prefill, causal) */
2121 k,
2122 v,
2123 attn_out,
2124 H,
2125 H_kv,
2126 num_tokens,
2127 head_dim,
2128 aligned_head_dim,
2129 aligned_context_window);
2130
2131 /* Output projection (flatten head-major to token-major) */
2132 const int K = H * aligned_head_dim;
2133 if (K != aligned_embed_dim) {
2134 return;
2135 }
2136 const float *proj_in = attn_out;
2137 if (H > 1) {
2138 if (!proj_scratch) {
2139 return;
2140 }
2141 for (int t = 0; t < num_tokens; ++t) {
2142 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
2143 for (int h = 0; h < H; ++h) {
2144 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
2145 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
2146 src,
2147 (size_t)aligned_head_dim * sizeof(float));
2148 }
2149 }
2150 proj_in = proj_scratch;
2151 }
2152 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
2153
2154 /* Residual add */
2155 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
2156
2157 /* RMSNorm before MLP */
2158 rmsnorm_forward(residual1,
2159 ln2_gamma,
2160 ln2_out,
2161 NULL,
2162 num_tokens,
2164 aligned_embed_dim,
2165 1e-06f);
2166
2167 /* MLP (SwiGLU) */
2168 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2169 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2170 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2171
2172 /* Final residual add */
2173 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2174}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_12_decode()

static void model_layer_12_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 6433 of file v6.6/test_generated/generic_api_test.c.

6440 {
6441 const MODELLayerOffsets *L = &MODEL_LAYERS[12];
6442
6443 float *input = MODEL_PTR(model, MODEL_LAYERS[11].output);
6444
6445 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
6446 float *ln1_out = MODEL_PTR(model, L->ln1_out);
6447 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
6448 float *ln2_out = MODEL_PTR(model, L->ln2_out);
6449 float *k_cache = MODEL_PTR(model, L->k);
6450 float *v_cache = MODEL_PTR(model, L->v);
6451 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
6452 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
6453 float *residual1 = MODEL_PTR(model, L->residual1);
6454 float *mlp_out = MODEL_PTR(model, L->mlp_out);
6455 float *output = MODEL_PTR(model, L->output);
6456
6457 /* Weights (explicit types for layer 12) */
6458 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
6459 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
6460 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
6461 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
6462 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
6463 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
6464
6465 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
6466 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
6467
6468 const int H = MODEL_NUM_HEADS;
6469 const int H_kv = MODEL_NUM_KV_HEADS;
6470 const int head_dim = MODEL_HEAD_DIM;
6471
6472 float q_token[H * aligned_head_dim];
6473 float k_token[H_kv * aligned_head_dim];
6474 float v_token[H_kv * aligned_head_dim];
6475 float attn_token[H * aligned_head_dim];
6476
6477 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
6478 float fc1_out[2 * aligned_intermediate_dim];
6479 float swiglu_out[aligned_intermediate_dim];
6480
6481 /* Step 1: RMSNorm before attention */
6482 rmsnorm_forward(input,
6483 ln1_gamma,
6484 ln1_out,
6485 NULL,
6486 1,
6488 aligned_embed_dim,
6489 1e-06f);
6490
6491 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
6492
6493 /* Step 2: QKV projection */
6494 /* Q projection: Q4_K -> gemm_nt_q4_k */
6495 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
6496 if (aligned_head_dim > head_dim) {
6497 for (int h = 0; h < H; ++h) {
6498 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
6499 for (int d = head_dim; d < aligned_head_dim; ++d) {
6500 q_head[d] = 0.0f;
6501 }
6502 }
6503 }
6504
6505 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6506 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6507 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
6508 const uint8_t *WK_bytes = (const uint8_t *)WK;
6509 for (int h = 0; h < H_kv; ++h) {
6510 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
6511 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
6512 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6513 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
6514 for (int d = head_dim; d < aligned_head_dim; ++d) {
6515 k_head[d] = 0.0f;
6516 }
6517 }
6518
6519 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6520 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6521 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
6522 const uint8_t *WV_bytes = (const uint8_t *)WV;
6523 for (int h = 0; h < H_kv; ++h) {
6524 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
6525 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
6526 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6527 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
6528 for (int d = head_dim; d < aligned_head_dim; ++d) {
6529 v_head[d] = 0.0f;
6530 }
6531 }
6532
6533 /* Step 3: RoPE */
6534 rope_forward(q_token,
6535 rope_cos,
6536 rope_sin,
6537 H,
6538 1,
6539 head_dim,
6540 aligned_head_dim,
6541 token_index);
6542 for (int h = 0; h < H_kv; ++h) {
6543 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6544 rope_forward(k_head,
6545 rope_cos,
6546 rope_sin,
6547 1,
6548 1,
6549 head_dim,
6550 aligned_head_dim,
6551 token_index);
6552 }
6553
6554 /* Step 4: KV cache write (direct-to-cache) */
6555
6556 /* Step 5: Attention (decode, flash) */
6558 k_cache,
6559 v_cache,
6560 attn_token,
6561 H,
6562 H_kv,
6563 token_index + 1,
6564 aligned_context_window,
6565 head_dim,
6566 aligned_head_dim);
6567
6568 /* Step 6: Output projection */
6569 /* WO projection: Q4_K -> gemm_nt_q4_k */
6570 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
6571
6572 /* Step 7: Residual add */
6573 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
6574
6575 /* Step 8: RMSNorm before MLP */
6576 rmsnorm_forward(residual1,
6577 ln2_gamma,
6578 ln2_out,
6579 NULL,
6580 1,
6582 aligned_embed_dim,
6583 1e-06f);
6584
6585 /* Step 9: MLP (SwiGLU) */
6586 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
6587 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
6588
6589 /* SwiGLU activation */
6590 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
6591
6592 /* Down projection: Q4_K -> gemm_nt_q4_k */
6593 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
6594
6595 /* Step 10: Final residual add */
6596 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
6597}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_12_prefill()

static void model_layer_12_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2179 of file v6.6/test_generated/generic_api_test.c.

2186 {
2187 const MODELLayerOffsets *L = &MODEL_LAYERS[12];
2188
2189 float *input = MODEL_PTR(model, MODEL_LAYERS[11].output);
2190 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
2191 float *ln1_out = MODEL_PTR(model, L->ln1_out);
2192 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
2193 float *ln2_out = MODEL_PTR(model, L->ln2_out);
2194 float *q = MODEL_PTR(model, L->q);
2195 float *k = MODEL_PTR(model, L->k);
2196 float *v = MODEL_PTR(model, L->v);
2197 float *attn_out = MODEL_PTR(model, L->attn_out);
2198 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
2199 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
2200 float *residual1 = MODEL_PTR(model, L->residual1);
2201 float *fc1_out = MODEL_PTR(model, L->fc1_out);
2202 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
2203 float *mlp_out = MODEL_PTR(model, L->mlp_out);
2204 float *output = MODEL_PTR(model, L->output);
2205
2206 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
2207 const void *WK = (const void *)MODEL_PTR(model, L->wk);
2208 const void *WV = (const void *)MODEL_PTR(model, L->wv);
2209 const void *WO = (const void *)MODEL_PTR(model, L->wo);
2210 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
2211 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
2212 const float *BQ = NULL;
2213 const float *BK = NULL;
2214 const float *BV = NULL;
2215 const float *BO = NULL;
2216 const float *B1 = NULL;
2217 const float *B2 = NULL;
2218
2219 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
2220 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
2221
2222 const int H = MODEL_NUM_HEADS;
2223 const int H_kv = MODEL_NUM_KV_HEADS;
2224 const int head_dim = MODEL_HEAD_DIM;
2225 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
2226 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
2227 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
2228
2229 /* RMSNorm before attention */
2230 rmsnorm_forward(input,
2231 ln1_gamma,
2232 ln1_out,
2233 NULL,
2234 num_tokens,
2236 aligned_embed_dim,
2237 1e-06f);
2238
2239 /* Q projection (head-major) */
2240 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2241 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
2242 for (int h = 0; h < H; ++h) {
2243 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
2244 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
2245 float *q_h = q + (size_t)h * q_head_stride;
2246 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2247 }
2248
2249 /* K projection (head-major) */
2250 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2251 const uint8_t *WK_bytes = (const uint8_t *)WK;
2252 for (int h = 0; h < H_kv; ++h) {
2253 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
2254 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
2255 float *k_h = k + (size_t)h * kv_head_stride;
2256 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2257 }
2258
2259 /* V projection (head-major) */
2260 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2261 const uint8_t *WV_bytes = (const uint8_t *)WV;
2262 for (int h = 0; h < H_kv; ++h) {
2263 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
2264 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
2265 float *v_h = v + (size_t)h * kv_head_stride;
2266 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2267 }
2268
2269 /* RoPE */
2271 k,
2272 rope_cos,
2273 rope_sin,
2274 H,
2275 H_kv,
2276 num_tokens,
2277 head_dim,
2278 aligned_head_dim,
2279 0,
2280 num_tokens,
2281 aligned_context_window);
2282
2283 /* Attention (prefill, causal) */
2285 k,
2286 v,
2287 attn_out,
2288 H,
2289 H_kv,
2290 num_tokens,
2291 head_dim,
2292 aligned_head_dim,
2293 aligned_context_window);
2294
2295 /* Output projection (flatten head-major to token-major) */
2296 const int K = H * aligned_head_dim;
2297 if (K != aligned_embed_dim) {
2298 return;
2299 }
2300 const float *proj_in = attn_out;
2301 if (H > 1) {
2302 if (!proj_scratch) {
2303 return;
2304 }
2305 for (int t = 0; t < num_tokens; ++t) {
2306 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
2307 for (int h = 0; h < H; ++h) {
2308 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
2309 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
2310 src,
2311 (size_t)aligned_head_dim * sizeof(float));
2312 }
2313 }
2314 proj_in = proj_scratch;
2315 }
2316 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
2317
2318 /* Residual add */
2319 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
2320
2321 /* RMSNorm before MLP */
2322 rmsnorm_forward(residual1,
2323 ln2_gamma,
2324 ln2_out,
2325 NULL,
2326 num_tokens,
2328 aligned_embed_dim,
2329 1e-06f);
2330
2331 /* MLP (SwiGLU) */
2332 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2333 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2334 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2335
2336 /* Final residual add */
2337 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2338}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_13_decode()

static void model_layer_13_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 6602 of file v6.6/test_generated/generic_api_test.c.

6609 {
6610 const MODELLayerOffsets *L = &MODEL_LAYERS[13];
6611
6612 float *input = MODEL_PTR(model, MODEL_LAYERS[12].output);
6613
6614 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
6615 float *ln1_out = MODEL_PTR(model, L->ln1_out);
6616 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
6617 float *ln2_out = MODEL_PTR(model, L->ln2_out);
6618 float *k_cache = MODEL_PTR(model, L->k);
6619 float *v_cache = MODEL_PTR(model, L->v);
6620 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
6621 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
6622 float *residual1 = MODEL_PTR(model, L->residual1);
6623 float *mlp_out = MODEL_PTR(model, L->mlp_out);
6624 float *output = MODEL_PTR(model, L->output);
6625
6626 /* Weights (explicit types for layer 13) */
6627 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
6628 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
6629 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
6630 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
6631 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
6632 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
6633
6634 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
6635 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
6636
6637 const int H = MODEL_NUM_HEADS;
6638 const int H_kv = MODEL_NUM_KV_HEADS;
6639 const int head_dim = MODEL_HEAD_DIM;
6640
6641 float q_token[H * aligned_head_dim];
6642 float k_token[H_kv * aligned_head_dim];
6643 float v_token[H_kv * aligned_head_dim];
6644 float attn_token[H * aligned_head_dim];
6645
6646 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
6647 float fc1_out[2 * aligned_intermediate_dim];
6648 float swiglu_out[aligned_intermediate_dim];
6649
6650 /* Step 1: RMSNorm before attention */
6651 rmsnorm_forward(input,
6652 ln1_gamma,
6653 ln1_out,
6654 NULL,
6655 1,
6657 aligned_embed_dim,
6658 1e-06f);
6659
6660 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
6661
6662 /* Step 2: QKV projection */
6663 /* Q projection: Q4_K -> gemm_nt_q4_k */
6664 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
6665 if (aligned_head_dim > head_dim) {
6666 for (int h = 0; h < H; ++h) {
6667 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
6668 for (int d = head_dim; d < aligned_head_dim; ++d) {
6669 q_head[d] = 0.0f;
6670 }
6671 }
6672 }
6673
6674 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6675 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6676 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
6677 const uint8_t *WK_bytes = (const uint8_t *)WK;
6678 for (int h = 0; h < H_kv; ++h) {
6679 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
6680 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
6681 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6682 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
6683 for (int d = head_dim; d < aligned_head_dim; ++d) {
6684 k_head[d] = 0.0f;
6685 }
6686 }
6687
6688 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6689 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6690 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
6691 const uint8_t *WV_bytes = (const uint8_t *)WV;
6692 for (int h = 0; h < H_kv; ++h) {
6693 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
6694 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
6695 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6696 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
6697 for (int d = head_dim; d < aligned_head_dim; ++d) {
6698 v_head[d] = 0.0f;
6699 }
6700 }
6701
6702 /* Step 3: RoPE */
6703 rope_forward(q_token,
6704 rope_cos,
6705 rope_sin,
6706 H,
6707 1,
6708 head_dim,
6709 aligned_head_dim,
6710 token_index);
6711 for (int h = 0; h < H_kv; ++h) {
6712 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6713 rope_forward(k_head,
6714 rope_cos,
6715 rope_sin,
6716 1,
6717 1,
6718 head_dim,
6719 aligned_head_dim,
6720 token_index);
6721 }
6722
6723 /* Step 4: KV cache write (direct-to-cache) */
6724
6725 /* Step 5: Attention (decode, flash) */
6727 k_cache,
6728 v_cache,
6729 attn_token,
6730 H,
6731 H_kv,
6732 token_index + 1,
6733 aligned_context_window,
6734 head_dim,
6735 aligned_head_dim);
6736
6737 /* Step 6: Output projection */
6738 /* WO projection: Q4_K -> gemm_nt_q4_k */
6739 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
6740
6741 /* Step 7: Residual add */
6742 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
6743
6744 /* Step 8: RMSNorm before MLP */
6745 rmsnorm_forward(residual1,
6746 ln2_gamma,
6747 ln2_out,
6748 NULL,
6749 1,
6751 aligned_embed_dim,
6752 1e-06f);
6753
6754 /* Step 9: MLP (SwiGLU) */
6755 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
6756 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
6757
6758 /* SwiGLU activation */
6759 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
6760
6761 /* Down projection: Q4_K -> gemm_nt_q4_k */
6762 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
6763
6764 /* Step 10: Final residual add */
6765 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
6766}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_13_prefill()

static void model_layer_13_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2343 of file v6.6/test_generated/generic_api_test.c.

2350 {
2351 const MODELLayerOffsets *L = &MODEL_LAYERS[13];
2352
2353 float *input = MODEL_PTR(model, MODEL_LAYERS[12].output);
2354 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
2355 float *ln1_out = MODEL_PTR(model, L->ln1_out);
2356 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
2357 float *ln2_out = MODEL_PTR(model, L->ln2_out);
2358 float *q = MODEL_PTR(model, L->q);
2359 float *k = MODEL_PTR(model, L->k);
2360 float *v = MODEL_PTR(model, L->v);
2361 float *attn_out = MODEL_PTR(model, L->attn_out);
2362 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
2363 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
2364 float *residual1 = MODEL_PTR(model, L->residual1);
2365 float *fc1_out = MODEL_PTR(model, L->fc1_out);
2366 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
2367 float *mlp_out = MODEL_PTR(model, L->mlp_out);
2368 float *output = MODEL_PTR(model, L->output);
2369
2370 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
2371 const void *WK = (const void *)MODEL_PTR(model, L->wk);
2372 const void *WV = (const void *)MODEL_PTR(model, L->wv);
2373 const void *WO = (const void *)MODEL_PTR(model, L->wo);
2374 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
2375 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
2376 const float *BQ = NULL;
2377 const float *BK = NULL;
2378 const float *BV = NULL;
2379 const float *BO = NULL;
2380 const float *B1 = NULL;
2381 const float *B2 = NULL;
2382
2383 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
2384 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
2385
2386 const int H = MODEL_NUM_HEADS;
2387 const int H_kv = MODEL_NUM_KV_HEADS;
2388 const int head_dim = MODEL_HEAD_DIM;
2389 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
2390 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
2391 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
2392
2393 /* RMSNorm before attention */
2394 rmsnorm_forward(input,
2395 ln1_gamma,
2396 ln1_out,
2397 NULL,
2398 num_tokens,
2400 aligned_embed_dim,
2401 1e-06f);
2402
2403 /* Q projection (head-major) */
2404 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2405 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
2406 for (int h = 0; h < H; ++h) {
2407 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
2408 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
2409 float *q_h = q + (size_t)h * q_head_stride;
2410 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2411 }
2412
2413 /* K projection (head-major) */
2414 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2415 const uint8_t *WK_bytes = (const uint8_t *)WK;
2416 for (int h = 0; h < H_kv; ++h) {
2417 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
2418 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
2419 float *k_h = k + (size_t)h * kv_head_stride;
2420 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2421 }
2422
2423 /* V projection (head-major) */
2424 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2425 const uint8_t *WV_bytes = (const uint8_t *)WV;
2426 for (int h = 0; h < H_kv; ++h) {
2427 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
2428 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
2429 float *v_h = v + (size_t)h * kv_head_stride;
2430 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2431 }
2432
2433 /* RoPE */
2435 k,
2436 rope_cos,
2437 rope_sin,
2438 H,
2439 H_kv,
2440 num_tokens,
2441 head_dim,
2442 aligned_head_dim,
2443 0,
2444 num_tokens,
2445 aligned_context_window);
2446
2447 /* Attention (prefill, causal) */
2449 k,
2450 v,
2451 attn_out,
2452 H,
2453 H_kv,
2454 num_tokens,
2455 head_dim,
2456 aligned_head_dim,
2457 aligned_context_window);
2458
2459 /* Output projection (flatten head-major to token-major) */
2460 const int K = H * aligned_head_dim;
2461 if (K != aligned_embed_dim) {
2462 return;
2463 }
2464 const float *proj_in = attn_out;
2465 if (H > 1) {
2466 if (!proj_scratch) {
2467 return;
2468 }
2469 for (int t = 0; t < num_tokens; ++t) {
2470 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
2471 for (int h = 0; h < H; ++h) {
2472 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
2473 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
2474 src,
2475 (size_t)aligned_head_dim * sizeof(float));
2476 }
2477 }
2478 proj_in = proj_scratch;
2479 }
2480 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
2481
2482 /* Residual add */
2483 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
2484
2485 /* RMSNorm before MLP */
2486 rmsnorm_forward(residual1,
2487 ln2_gamma,
2488 ln2_out,
2489 NULL,
2490 num_tokens,
2492 aligned_embed_dim,
2493 1e-06f);
2494
2495 /* MLP (SwiGLU) */
2496 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2497 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2498 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2499
2500 /* Final residual add */
2501 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2502}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_14_decode()

static void model_layer_14_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 6771 of file v6.6/test_generated/generic_api_test.c.

6778 {
6779 const MODELLayerOffsets *L = &MODEL_LAYERS[14];
6780
6781 float *input = MODEL_PTR(model, MODEL_LAYERS[13].output);
6782
6783 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
6784 float *ln1_out = MODEL_PTR(model, L->ln1_out);
6785 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
6786 float *ln2_out = MODEL_PTR(model, L->ln2_out);
6787 float *k_cache = MODEL_PTR(model, L->k);
6788 float *v_cache = MODEL_PTR(model, L->v);
6789 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
6790 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
6791 float *residual1 = MODEL_PTR(model, L->residual1);
6792 float *mlp_out = MODEL_PTR(model, L->mlp_out);
6793 float *output = MODEL_PTR(model, L->output);
6794
6795 /* Weights (explicit types for layer 14) */
6796 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
6797 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
6798 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
6799 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
6800 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
6801 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
6802
6803 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
6804 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
6805
6806 const int H = MODEL_NUM_HEADS;
6807 const int H_kv = MODEL_NUM_KV_HEADS;
6808 const int head_dim = MODEL_HEAD_DIM;
6809
6810 float q_token[H * aligned_head_dim];
6811 float k_token[H_kv * aligned_head_dim];
6812 float v_token[H_kv * aligned_head_dim];
6813 float attn_token[H * aligned_head_dim];
6814
6815 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
6816 float fc1_out[2 * aligned_intermediate_dim];
6817 float swiglu_out[aligned_intermediate_dim];
6818
6819 /* Step 1: RMSNorm before attention */
6820 rmsnorm_forward(input,
6821 ln1_gamma,
6822 ln1_out,
6823 NULL,
6824 1,
6826 aligned_embed_dim,
6827 1e-06f);
6828
6829 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
6830
6831 /* Step 2: QKV projection */
6832 /* Q projection: Q4_K -> gemm_nt_q4_k */
6833 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
6834 if (aligned_head_dim > head_dim) {
6835 for (int h = 0; h < H; ++h) {
6836 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
6837 for (int d = head_dim; d < aligned_head_dim; ++d) {
6838 q_head[d] = 0.0f;
6839 }
6840 }
6841 }
6842
6843 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6844 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6845 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
6846 const uint8_t *WK_bytes = (const uint8_t *)WK;
6847 for (int h = 0; h < H_kv; ++h) {
6848 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
6849 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
6850 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6851 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
6852 for (int d = head_dim; d < aligned_head_dim; ++d) {
6853 k_head[d] = 0.0f;
6854 }
6855 }
6856
6857 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6858 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6859 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
6860 const uint8_t *WV_bytes = (const uint8_t *)WV;
6861 for (int h = 0; h < H_kv; ++h) {
6862 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
6863 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
6864 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6865 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
6866 for (int d = head_dim; d < aligned_head_dim; ++d) {
6867 v_head[d] = 0.0f;
6868 }
6869 }
6870
6871 /* Step 3: RoPE */
6872 rope_forward(q_token,
6873 rope_cos,
6874 rope_sin,
6875 H,
6876 1,
6877 head_dim,
6878 aligned_head_dim,
6879 token_index);
6880 for (int h = 0; h < H_kv; ++h) {
6881 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6882 rope_forward(k_head,
6883 rope_cos,
6884 rope_sin,
6885 1,
6886 1,
6887 head_dim,
6888 aligned_head_dim,
6889 token_index);
6890 }
6891
6892 /* Step 4: KV cache write (direct-to-cache) */
6893
6894 /* Step 5: Attention (decode, flash) */
6896 k_cache,
6897 v_cache,
6898 attn_token,
6899 H,
6900 H_kv,
6901 token_index + 1,
6902 aligned_context_window,
6903 head_dim,
6904 aligned_head_dim);
6905
6906 /* Step 6: Output projection */
6907 /* WO projection: Q4_K -> gemm_nt_q4_k */
6908 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
6909
6910 /* Step 7: Residual add */
6911 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
6912
6913 /* Step 8: RMSNorm before MLP */
6914 rmsnorm_forward(residual1,
6915 ln2_gamma,
6916 ln2_out,
6917 NULL,
6918 1,
6920 aligned_embed_dim,
6921 1e-06f);
6922
6923 /* Step 9: MLP (SwiGLU) */
6924 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
6925 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
6926
6927 /* SwiGLU activation */
6928 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
6929
6930 /* Down projection: Q4_K -> gemm_nt_q4_k */
6931 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
6932
6933 /* Step 10: Final residual add */
6934 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
6935}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_14_prefill()

static void model_layer_14_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2507 of file v6.6/test_generated/generic_api_test.c.

2514 {
2515 const MODELLayerOffsets *L = &MODEL_LAYERS[14];
2516
2517 float *input = MODEL_PTR(model, MODEL_LAYERS[13].output);
2518 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
2519 float *ln1_out = MODEL_PTR(model, L->ln1_out);
2520 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
2521 float *ln2_out = MODEL_PTR(model, L->ln2_out);
2522 float *q = MODEL_PTR(model, L->q);
2523 float *k = MODEL_PTR(model, L->k);
2524 float *v = MODEL_PTR(model, L->v);
2525 float *attn_out = MODEL_PTR(model, L->attn_out);
2526 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
2527 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
2528 float *residual1 = MODEL_PTR(model, L->residual1);
2529 float *fc1_out = MODEL_PTR(model, L->fc1_out);
2530 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
2531 float *mlp_out = MODEL_PTR(model, L->mlp_out);
2532 float *output = MODEL_PTR(model, L->output);
2533
2534 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
2535 const void *WK = (const void *)MODEL_PTR(model, L->wk);
2536 const void *WV = (const void *)MODEL_PTR(model, L->wv);
2537 const void *WO = (const void *)MODEL_PTR(model, L->wo);
2538 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
2539 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
2540 const float *BQ = NULL;
2541 const float *BK = NULL;
2542 const float *BV = NULL;
2543 const float *BO = NULL;
2544 const float *B1 = NULL;
2545 const float *B2 = NULL;
2546
2547 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
2548 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
2549
2550 const int H = MODEL_NUM_HEADS;
2551 const int H_kv = MODEL_NUM_KV_HEADS;
2552 const int head_dim = MODEL_HEAD_DIM;
2553 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
2554 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
2555 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
2556
2557 /* RMSNorm before attention */
2558 rmsnorm_forward(input,
2559 ln1_gamma,
2560 ln1_out,
2561 NULL,
2562 num_tokens,
2564 aligned_embed_dim,
2565 1e-06f);
2566
2567 /* Q projection (head-major) */
2568 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2569 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
2570 for (int h = 0; h < H; ++h) {
2571 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
2572 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
2573 float *q_h = q + (size_t)h * q_head_stride;
2574 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2575 }
2576
2577 /* K projection (head-major) */
2578 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2579 const uint8_t *WK_bytes = (const uint8_t *)WK;
2580 for (int h = 0; h < H_kv; ++h) {
2581 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
2582 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
2583 float *k_h = k + (size_t)h * kv_head_stride;
2584 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2585 }
2586
2587 /* V projection (head-major) */
2588 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2589 const uint8_t *WV_bytes = (const uint8_t *)WV;
2590 for (int h = 0; h < H_kv; ++h) {
2591 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
2592 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
2593 float *v_h = v + (size_t)h * kv_head_stride;
2594 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2595 }
2596
2597 /* RoPE */
2599 k,
2600 rope_cos,
2601 rope_sin,
2602 H,
2603 H_kv,
2604 num_tokens,
2605 head_dim,
2606 aligned_head_dim,
2607 0,
2608 num_tokens,
2609 aligned_context_window);
2610
2611 /* Attention (prefill, causal) */
2613 k,
2614 v,
2615 attn_out,
2616 H,
2617 H_kv,
2618 num_tokens,
2619 head_dim,
2620 aligned_head_dim,
2621 aligned_context_window);
2622
2623 /* Output projection (flatten head-major to token-major) */
2624 const int K = H * aligned_head_dim;
2625 if (K != aligned_embed_dim) {
2626 return;
2627 }
2628 const float *proj_in = attn_out;
2629 if (H > 1) {
2630 if (!proj_scratch) {
2631 return;
2632 }
2633 for (int t = 0; t < num_tokens; ++t) {
2634 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
2635 for (int h = 0; h < H; ++h) {
2636 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
2637 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
2638 src,
2639 (size_t)aligned_head_dim * sizeof(float));
2640 }
2641 }
2642 proj_in = proj_scratch;
2643 }
2644 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
2645
2646 /* Residual add */
2647 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
2648
2649 /* RMSNorm before MLP */
2650 rmsnorm_forward(residual1,
2651 ln2_gamma,
2652 ln2_out,
2653 NULL,
2654 num_tokens,
2656 aligned_embed_dim,
2657 1e-06f);
2658
2659 /* MLP (SwiGLU) */
2660 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2661 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2662 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2663
2664 /* Final residual add */
2665 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2666}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_15_decode()

static void model_layer_15_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 6940 of file v6.6/test_generated/generic_api_test.c.

6947 {
6948 const MODELLayerOffsets *L = &MODEL_LAYERS[15];
6949
6950 float *input = MODEL_PTR(model, MODEL_LAYERS[14].output);
6951
6952 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
6953 float *ln1_out = MODEL_PTR(model, L->ln1_out);
6954 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
6955 float *ln2_out = MODEL_PTR(model, L->ln2_out);
6956 float *k_cache = MODEL_PTR(model, L->k);
6957 float *v_cache = MODEL_PTR(model, L->v);
6958 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
6959 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
6960 float *residual1 = MODEL_PTR(model, L->residual1);
6961 float *mlp_out = MODEL_PTR(model, L->mlp_out);
6962 float *output = MODEL_PTR(model, L->output);
6963
6964 /* Weights (explicit types for layer 15) */
6965 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
6966 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
6967 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
6968 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
6969 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
6970 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
6971
6972 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
6973 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
6974
6975 const int H = MODEL_NUM_HEADS;
6976 const int H_kv = MODEL_NUM_KV_HEADS;
6977 const int head_dim = MODEL_HEAD_DIM;
6978
6979 float q_token[H * aligned_head_dim];
6980 float k_token[H_kv * aligned_head_dim];
6981 float v_token[H_kv * aligned_head_dim];
6982 float attn_token[H * aligned_head_dim];
6983
6984 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
6985 float fc1_out[2 * aligned_intermediate_dim];
6986 float swiglu_out[aligned_intermediate_dim];
6987
6988 /* Step 1: RMSNorm before attention */
6989 rmsnorm_forward(input,
6990 ln1_gamma,
6991 ln1_out,
6992 NULL,
6993 1,
6995 aligned_embed_dim,
6996 1e-06f);
6997
6998 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
6999
7000 /* Step 2: QKV projection */
7001 /* Q projection: Q4_K -> gemm_nt_q4_k */
7002 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
7003 if (aligned_head_dim > head_dim) {
7004 for (int h = 0; h < H; ++h) {
7005 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
7006 for (int d = head_dim; d < aligned_head_dim; ++d) {
7007 q_head[d] = 0.0f;
7008 }
7009 }
7010 }
7011
7012 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7013 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7014 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
7015 const uint8_t *WK_bytes = (const uint8_t *)WK;
7016 for (int h = 0; h < H_kv; ++h) {
7017 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
7018 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
7019 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7020 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
7021 for (int d = head_dim; d < aligned_head_dim; ++d) {
7022 k_head[d] = 0.0f;
7023 }
7024 }
7025
7026 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7027 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7028 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
7029 const uint8_t *WV_bytes = (const uint8_t *)WV;
7030 for (int h = 0; h < H_kv; ++h) {
7031 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
7032 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
7033 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7034 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
7035 for (int d = head_dim; d < aligned_head_dim; ++d) {
7036 v_head[d] = 0.0f;
7037 }
7038 }
7039
7040 /* Step 3: RoPE */
7041 rope_forward(q_token,
7042 rope_cos,
7043 rope_sin,
7044 H,
7045 1,
7046 head_dim,
7047 aligned_head_dim,
7048 token_index);
7049 for (int h = 0; h < H_kv; ++h) {
7050 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7051 rope_forward(k_head,
7052 rope_cos,
7053 rope_sin,
7054 1,
7055 1,
7056 head_dim,
7057 aligned_head_dim,
7058 token_index);
7059 }
7060
7061 /* Step 4: KV cache write (direct-to-cache) */
7062
7063 /* Step 5: Attention (decode, flash) */
7065 k_cache,
7066 v_cache,
7067 attn_token,
7068 H,
7069 H_kv,
7070 token_index + 1,
7071 aligned_context_window,
7072 head_dim,
7073 aligned_head_dim);
7074
7075 /* Step 6: Output projection */
7076 /* WO projection: Q4_K -> gemm_nt_q4_k */
7077 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
7078
7079 /* Step 7: Residual add */
7080 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
7081
7082 /* Step 8: RMSNorm before MLP */
7083 rmsnorm_forward(residual1,
7084 ln2_gamma,
7085 ln2_out,
7086 NULL,
7087 1,
7089 aligned_embed_dim,
7090 1e-06f);
7091
7092 /* Step 9: MLP (SwiGLU) */
7093 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
7094 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
7095
7096 /* SwiGLU activation */
7097 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
7098
7099 /* Down projection: Q4_K -> gemm_nt_q4_k */
7100 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
7101
7102 /* Step 10: Final residual add */
7103 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
7104}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_15_prefill()

static void model_layer_15_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2671 of file v6.6/test_generated/generic_api_test.c.

2678 {
2679 const MODELLayerOffsets *L = &MODEL_LAYERS[15];
2680
2681 float *input = MODEL_PTR(model, MODEL_LAYERS[14].output);
2682 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
2683 float *ln1_out = MODEL_PTR(model, L->ln1_out);
2684 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
2685 float *ln2_out = MODEL_PTR(model, L->ln2_out);
2686 float *q = MODEL_PTR(model, L->q);
2687 float *k = MODEL_PTR(model, L->k);
2688 float *v = MODEL_PTR(model, L->v);
2689 float *attn_out = MODEL_PTR(model, L->attn_out);
2690 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
2691 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
2692 float *residual1 = MODEL_PTR(model, L->residual1);
2693 float *fc1_out = MODEL_PTR(model, L->fc1_out);
2694 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
2695 float *mlp_out = MODEL_PTR(model, L->mlp_out);
2696 float *output = MODEL_PTR(model, L->output);
2697
2698 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
2699 const void *WK = (const void *)MODEL_PTR(model, L->wk);
2700 const void *WV = (const void *)MODEL_PTR(model, L->wv);
2701 const void *WO = (const void *)MODEL_PTR(model, L->wo);
2702 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
2703 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
2704 const float *BQ = NULL;
2705 const float *BK = NULL;
2706 const float *BV = NULL;
2707 const float *BO = NULL;
2708 const float *B1 = NULL;
2709 const float *B2 = NULL;
2710
2711 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
2712 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
2713
2714 const int H = MODEL_NUM_HEADS;
2715 const int H_kv = MODEL_NUM_KV_HEADS;
2716 const int head_dim = MODEL_HEAD_DIM;
2717 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
2718 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
2719 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
2720
2721 /* RMSNorm before attention */
2722 rmsnorm_forward(input,
2723 ln1_gamma,
2724 ln1_out,
2725 NULL,
2726 num_tokens,
2728 aligned_embed_dim,
2729 1e-06f);
2730
2731 /* Q projection (head-major) */
2732 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2733 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
2734 for (int h = 0; h < H; ++h) {
2735 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
2736 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
2737 float *q_h = q + (size_t)h * q_head_stride;
2738 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2739 }
2740
2741 /* K projection (head-major) */
2742 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2743 const uint8_t *WK_bytes = (const uint8_t *)WK;
2744 for (int h = 0; h < H_kv; ++h) {
2745 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
2746 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
2747 float *k_h = k + (size_t)h * kv_head_stride;
2748 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2749 }
2750
2751 /* V projection (head-major) */
2752 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2753 const uint8_t *WV_bytes = (const uint8_t *)WV;
2754 for (int h = 0; h < H_kv; ++h) {
2755 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
2756 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
2757 float *v_h = v + (size_t)h * kv_head_stride;
2758 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2759 }
2760
2761 /* RoPE */
2763 k,
2764 rope_cos,
2765 rope_sin,
2766 H,
2767 H_kv,
2768 num_tokens,
2769 head_dim,
2770 aligned_head_dim,
2771 0,
2772 num_tokens,
2773 aligned_context_window);
2774
2775 /* Attention (prefill, causal) */
2777 k,
2778 v,
2779 attn_out,
2780 H,
2781 H_kv,
2782 num_tokens,
2783 head_dim,
2784 aligned_head_dim,
2785 aligned_context_window);
2786
2787 /* Output projection (flatten head-major to token-major) */
2788 const int K = H * aligned_head_dim;
2789 if (K != aligned_embed_dim) {
2790 return;
2791 }
2792 const float *proj_in = attn_out;
2793 if (H > 1) {
2794 if (!proj_scratch) {
2795 return;
2796 }
2797 for (int t = 0; t < num_tokens; ++t) {
2798 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
2799 for (int h = 0; h < H; ++h) {
2800 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
2801 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
2802 src,
2803 (size_t)aligned_head_dim * sizeof(float));
2804 }
2805 }
2806 proj_in = proj_scratch;
2807 }
2808 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
2809
2810 /* Residual add */
2811 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
2812
2813 /* RMSNorm before MLP */
2814 rmsnorm_forward(residual1,
2815 ln2_gamma,
2816 ln2_out,
2817 NULL,
2818 num_tokens,
2820 aligned_embed_dim,
2821 1e-06f);
2822
2823 /* MLP (SwiGLU) */
2824 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2825 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2826 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2827
2828 /* Final residual add */
2829 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2830}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_16_decode()

static void model_layer_16_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 7109 of file v6.6/test_generated/generic_api_test.c.

7116 {
7117 const MODELLayerOffsets *L = &MODEL_LAYERS[16];
7118
7119 float *input = MODEL_PTR(model, MODEL_LAYERS[15].output);
7120
7121 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
7122 float *ln1_out = MODEL_PTR(model, L->ln1_out);
7123 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
7124 float *ln2_out = MODEL_PTR(model, L->ln2_out);
7125 float *k_cache = MODEL_PTR(model, L->k);
7126 float *v_cache = MODEL_PTR(model, L->v);
7127 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
7128 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
7129 float *residual1 = MODEL_PTR(model, L->residual1);
7130 float *mlp_out = MODEL_PTR(model, L->mlp_out);
7131 float *output = MODEL_PTR(model, L->output);
7132
7133 /* Weights (explicit types for layer 16) */
7134 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
7135 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
7136 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
7137 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
7138 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
7139 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
7140
7141 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
7142 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
7143
7144 const int H = MODEL_NUM_HEADS;
7145 const int H_kv = MODEL_NUM_KV_HEADS;
7146 const int head_dim = MODEL_HEAD_DIM;
7147
7148 float q_token[H * aligned_head_dim];
7149 float k_token[H_kv * aligned_head_dim];
7150 float v_token[H_kv * aligned_head_dim];
7151 float attn_token[H * aligned_head_dim];
7152
7153 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
7154 float fc1_out[2 * aligned_intermediate_dim];
7155 float swiglu_out[aligned_intermediate_dim];
7156
7157 /* Step 1: RMSNorm before attention */
7158 rmsnorm_forward(input,
7159 ln1_gamma,
7160 ln1_out,
7161 NULL,
7162 1,
7164 aligned_embed_dim,
7165 1e-06f);
7166
7167 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
7168
7169 /* Step 2: QKV projection */
7170 /* Q projection: Q4_K -> gemm_nt_q4_k */
7171 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
7172 if (aligned_head_dim > head_dim) {
7173 for (int h = 0; h < H; ++h) {
7174 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
7175 for (int d = head_dim; d < aligned_head_dim; ++d) {
7176 q_head[d] = 0.0f;
7177 }
7178 }
7179 }
7180
7181 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7182 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7183 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
7184 const uint8_t *WK_bytes = (const uint8_t *)WK;
7185 for (int h = 0; h < H_kv; ++h) {
7186 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
7187 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
7188 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7189 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
7190 for (int d = head_dim; d < aligned_head_dim; ++d) {
7191 k_head[d] = 0.0f;
7192 }
7193 }
7194
7195 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7196 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7197 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
7198 const uint8_t *WV_bytes = (const uint8_t *)WV;
7199 for (int h = 0; h < H_kv; ++h) {
7200 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
7201 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
7202 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7203 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
7204 for (int d = head_dim; d < aligned_head_dim; ++d) {
7205 v_head[d] = 0.0f;
7206 }
7207 }
7208
7209 /* Step 3: RoPE */
7210 rope_forward(q_token,
7211 rope_cos,
7212 rope_sin,
7213 H,
7214 1,
7215 head_dim,
7216 aligned_head_dim,
7217 token_index);
7218 for (int h = 0; h < H_kv; ++h) {
7219 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7220 rope_forward(k_head,
7221 rope_cos,
7222 rope_sin,
7223 1,
7224 1,
7225 head_dim,
7226 aligned_head_dim,
7227 token_index);
7228 }
7229
7230 /* Step 4: KV cache write (direct-to-cache) */
7231
7232 /* Step 5: Attention (decode, flash) */
7234 k_cache,
7235 v_cache,
7236 attn_token,
7237 H,
7238 H_kv,
7239 token_index + 1,
7240 aligned_context_window,
7241 head_dim,
7242 aligned_head_dim);
7243
7244 /* Step 6: Output projection */
7245 /* WO projection: Q4_K -> gemm_nt_q4_k */
7246 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
7247
7248 /* Step 7: Residual add */
7249 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
7250
7251 /* Step 8: RMSNorm before MLP */
7252 rmsnorm_forward(residual1,
7253 ln2_gamma,
7254 ln2_out,
7255 NULL,
7256 1,
7258 aligned_embed_dim,
7259 1e-06f);
7260
7261 /* Step 9: MLP (SwiGLU) */
7262 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
7263 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
7264
7265 /* SwiGLU activation */
7266 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
7267
7268 /* Down projection: Q4_K -> gemm_nt_q4_k */
7269 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
7270
7271 /* Step 10: Final residual add */
7272 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
7273}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_16_prefill()

static void model_layer_16_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2835 of file v6.6/test_generated/generic_api_test.c.

2842 {
2843 const MODELLayerOffsets *L = &MODEL_LAYERS[16];
2844
2845 float *input = MODEL_PTR(model, MODEL_LAYERS[15].output);
2846 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
2847 float *ln1_out = MODEL_PTR(model, L->ln1_out);
2848 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
2849 float *ln2_out = MODEL_PTR(model, L->ln2_out);
2850 float *q = MODEL_PTR(model, L->q);
2851 float *k = MODEL_PTR(model, L->k);
2852 float *v = MODEL_PTR(model, L->v);
2853 float *attn_out = MODEL_PTR(model, L->attn_out);
2854 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
2855 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
2856 float *residual1 = MODEL_PTR(model, L->residual1);
2857 float *fc1_out = MODEL_PTR(model, L->fc1_out);
2858 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
2859 float *mlp_out = MODEL_PTR(model, L->mlp_out);
2860 float *output = MODEL_PTR(model, L->output);
2861
2862 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
2863 const void *WK = (const void *)MODEL_PTR(model, L->wk);
2864 const void *WV = (const void *)MODEL_PTR(model, L->wv);
2865 const void *WO = (const void *)MODEL_PTR(model, L->wo);
2866 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
2867 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
2868 const float *BQ = NULL;
2869 const float *BK = NULL;
2870 const float *BV = NULL;
2871 const float *BO = NULL;
2872 const float *B1 = NULL;
2873 const float *B2 = NULL;
2874
2875 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
2876 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
2877
2878 const int H = MODEL_NUM_HEADS;
2879 const int H_kv = MODEL_NUM_KV_HEADS;
2880 const int head_dim = MODEL_HEAD_DIM;
2881 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
2882 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
2883 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
2884
2885 /* RMSNorm before attention */
2886 rmsnorm_forward(input,
2887 ln1_gamma,
2888 ln1_out,
2889 NULL,
2890 num_tokens,
2892 aligned_embed_dim,
2893 1e-06f);
2894
2895 /* Q projection (head-major) */
2896 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2897 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
2898 for (int h = 0; h < H; ++h) {
2899 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
2900 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
2901 float *q_h = q + (size_t)h * q_head_stride;
2902 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2903 }
2904
2905 /* K projection (head-major) */
2906 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2907 const uint8_t *WK_bytes = (const uint8_t *)WK;
2908 for (int h = 0; h < H_kv; ++h) {
2909 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
2910 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
2911 float *k_h = k + (size_t)h * kv_head_stride;
2912 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2913 }
2914
2915 /* V projection (head-major) */
2916 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
2917 const uint8_t *WV_bytes = (const uint8_t *)WV;
2918 for (int h = 0; h < H_kv; ++h) {
2919 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
2920 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
2921 float *v_h = v + (size_t)h * kv_head_stride;
2922 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
2923 }
2924
2925 /* RoPE */
2927 k,
2928 rope_cos,
2929 rope_sin,
2930 H,
2931 H_kv,
2932 num_tokens,
2933 head_dim,
2934 aligned_head_dim,
2935 0,
2936 num_tokens,
2937 aligned_context_window);
2938
2939 /* Attention (prefill, causal) */
2941 k,
2942 v,
2943 attn_out,
2944 H,
2945 H_kv,
2946 num_tokens,
2947 head_dim,
2948 aligned_head_dim,
2949 aligned_context_window);
2950
2951 /* Output projection (flatten head-major to token-major) */
2952 const int K = H * aligned_head_dim;
2953 if (K != aligned_embed_dim) {
2954 return;
2955 }
2956 const float *proj_in = attn_out;
2957 if (H > 1) {
2958 if (!proj_scratch) {
2959 return;
2960 }
2961 for (int t = 0; t < num_tokens; ++t) {
2962 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
2963 for (int h = 0; h < H; ++h) {
2964 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
2965 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
2966 src,
2967 (size_t)aligned_head_dim * sizeof(float));
2968 }
2969 }
2970 proj_in = proj_scratch;
2971 }
2972 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
2973
2974 /* Residual add */
2975 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
2976
2977 /* RMSNorm before MLP */
2978 rmsnorm_forward(residual1,
2979 ln2_gamma,
2980 ln2_out,
2981 NULL,
2982 num_tokens,
2984 aligned_embed_dim,
2985 1e-06f);
2986
2987 /* MLP (SwiGLU) */
2988 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
2989 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
2990 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
2991
2992 /* Final residual add */
2993 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
2994}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_17_decode()

static void model_layer_17_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 7278 of file v6.6/test_generated/generic_api_test.c.

7285 {
7286 const MODELLayerOffsets *L = &MODEL_LAYERS[17];
7287
7288 float *input = MODEL_PTR(model, MODEL_LAYERS[16].output);
7289
7290 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
7291 float *ln1_out = MODEL_PTR(model, L->ln1_out);
7292 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
7293 float *ln2_out = MODEL_PTR(model, L->ln2_out);
7294 float *k_cache = MODEL_PTR(model, L->k);
7295 float *v_cache = MODEL_PTR(model, L->v);
7296 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
7297 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
7298 float *residual1 = MODEL_PTR(model, L->residual1);
7299 float *mlp_out = MODEL_PTR(model, L->mlp_out);
7300 float *output = MODEL_PTR(model, L->output);
7301
7302 /* Weights (explicit types for layer 17) */
7303 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
7304 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
7305 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
7306 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
7307 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
7308 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
7309
7310 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
7311 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
7312
7313 const int H = MODEL_NUM_HEADS;
7314 const int H_kv = MODEL_NUM_KV_HEADS;
7315 const int head_dim = MODEL_HEAD_DIM;
7316
7317 float q_token[H * aligned_head_dim];
7318 float k_token[H_kv * aligned_head_dim];
7319 float v_token[H_kv * aligned_head_dim];
7320 float attn_token[H * aligned_head_dim];
7321
7322 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
7323 float fc1_out[2 * aligned_intermediate_dim];
7324 float swiglu_out[aligned_intermediate_dim];
7325
7326 /* Step 1: RMSNorm before attention */
7327 rmsnorm_forward(input,
7328 ln1_gamma,
7329 ln1_out,
7330 NULL,
7331 1,
7333 aligned_embed_dim,
7334 1e-06f);
7335
7336 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
7337
7338 /* Step 2: QKV projection */
7339 /* Q projection: Q4_K -> gemm_nt_q4_k */
7340 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
7341 if (aligned_head_dim > head_dim) {
7342 for (int h = 0; h < H; ++h) {
7343 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
7344 for (int d = head_dim; d < aligned_head_dim; ++d) {
7345 q_head[d] = 0.0f;
7346 }
7347 }
7348 }
7349
7350 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7351 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7352 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
7353 const uint8_t *WK_bytes = (const uint8_t *)WK;
7354 for (int h = 0; h < H_kv; ++h) {
7355 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
7356 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
7357 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7358 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
7359 for (int d = head_dim; d < aligned_head_dim; ++d) {
7360 k_head[d] = 0.0f;
7361 }
7362 }
7363
7364 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7365 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7366 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
7367 const uint8_t *WV_bytes = (const uint8_t *)WV;
7368 for (int h = 0; h < H_kv; ++h) {
7369 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
7370 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
7371 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7372 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
7373 for (int d = head_dim; d < aligned_head_dim; ++d) {
7374 v_head[d] = 0.0f;
7375 }
7376 }
7377
7378 /* Step 3: RoPE */
7379 rope_forward(q_token,
7380 rope_cos,
7381 rope_sin,
7382 H,
7383 1,
7384 head_dim,
7385 aligned_head_dim,
7386 token_index);
7387 for (int h = 0; h < H_kv; ++h) {
7388 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7389 rope_forward(k_head,
7390 rope_cos,
7391 rope_sin,
7392 1,
7393 1,
7394 head_dim,
7395 aligned_head_dim,
7396 token_index);
7397 }
7398
7399 /* Step 4: KV cache write (direct-to-cache) */
7400
7401 /* Step 5: Attention (decode, flash) */
7403 k_cache,
7404 v_cache,
7405 attn_token,
7406 H,
7407 H_kv,
7408 token_index + 1,
7409 aligned_context_window,
7410 head_dim,
7411 aligned_head_dim);
7412
7413 /* Step 6: Output projection */
7414 /* WO projection: Q4_K -> gemm_nt_q4_k */
7415 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
7416
7417 /* Step 7: Residual add */
7418 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
7419
7420 /* Step 8: RMSNorm before MLP */
7421 rmsnorm_forward(residual1,
7422 ln2_gamma,
7423 ln2_out,
7424 NULL,
7425 1,
7427 aligned_embed_dim,
7428 1e-06f);
7429
7430 /* Step 9: MLP (SwiGLU) */
7431 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
7432 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
7433
7434 /* SwiGLU activation */
7435 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
7436
7437 /* Down projection: Q4_K -> gemm_nt_q4_k */
7438 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
7439
7440 /* Step 10: Final residual add */
7441 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
7442}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_17_prefill()

static void model_layer_17_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 2999 of file v6.6/test_generated/generic_api_test.c.

3006 {
3007 const MODELLayerOffsets *L = &MODEL_LAYERS[17];
3008
3009 float *input = MODEL_PTR(model, MODEL_LAYERS[16].output);
3010 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3011 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3012 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3013 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3014 float *q = MODEL_PTR(model, L->q);
3015 float *k = MODEL_PTR(model, L->k);
3016 float *v = MODEL_PTR(model, L->v);
3017 float *attn_out = MODEL_PTR(model, L->attn_out);
3018 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
3019 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
3020 float *residual1 = MODEL_PTR(model, L->residual1);
3021 float *fc1_out = MODEL_PTR(model, L->fc1_out);
3022 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
3023 float *mlp_out = MODEL_PTR(model, L->mlp_out);
3024 float *output = MODEL_PTR(model, L->output);
3025
3026 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
3027 const void *WK = (const void *)MODEL_PTR(model, L->wk);
3028 const void *WV = (const void *)MODEL_PTR(model, L->wv);
3029 const void *WO = (const void *)MODEL_PTR(model, L->wo);
3030 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
3031 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
3032 const float *BQ = NULL;
3033 const float *BK = NULL;
3034 const float *BV = NULL;
3035 const float *BO = NULL;
3036 const float *B1 = NULL;
3037 const float *B2 = NULL;
3038
3039 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
3040 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
3041
3042 const int H = MODEL_NUM_HEADS;
3043 const int H_kv = MODEL_NUM_KV_HEADS;
3044 const int head_dim = MODEL_HEAD_DIM;
3045 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
3046 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
3047 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
3048
3049 /* RMSNorm before attention */
3050 rmsnorm_forward(input,
3051 ln1_gamma,
3052 ln1_out,
3053 NULL,
3054 num_tokens,
3056 aligned_embed_dim,
3057 1e-06f);
3058
3059 /* Q projection (head-major) */
3060 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3061 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
3062 for (int h = 0; h < H; ++h) {
3063 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
3064 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
3065 float *q_h = q + (size_t)h * q_head_stride;
3066 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3067 }
3068
3069 /* K projection (head-major) */
3070 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3071 const uint8_t *WK_bytes = (const uint8_t *)WK;
3072 for (int h = 0; h < H_kv; ++h) {
3073 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
3074 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
3075 float *k_h = k + (size_t)h * kv_head_stride;
3076 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3077 }
3078
3079 /* V projection (head-major) */
3080 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3081 const uint8_t *WV_bytes = (const uint8_t *)WV;
3082 for (int h = 0; h < H_kv; ++h) {
3083 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
3084 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
3085 float *v_h = v + (size_t)h * kv_head_stride;
3086 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3087 }
3088
3089 /* RoPE */
3091 k,
3092 rope_cos,
3093 rope_sin,
3094 H,
3095 H_kv,
3096 num_tokens,
3097 head_dim,
3098 aligned_head_dim,
3099 0,
3100 num_tokens,
3101 aligned_context_window);
3102
3103 /* Attention (prefill, causal) */
3105 k,
3106 v,
3107 attn_out,
3108 H,
3109 H_kv,
3110 num_tokens,
3111 head_dim,
3112 aligned_head_dim,
3113 aligned_context_window);
3114
3115 /* Output projection (flatten head-major to token-major) */
3116 const int K = H * aligned_head_dim;
3117 if (K != aligned_embed_dim) {
3118 return;
3119 }
3120 const float *proj_in = attn_out;
3121 if (H > 1) {
3122 if (!proj_scratch) {
3123 return;
3124 }
3125 for (int t = 0; t < num_tokens; ++t) {
3126 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
3127 for (int h = 0; h < H; ++h) {
3128 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
3129 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
3130 src,
3131 (size_t)aligned_head_dim * sizeof(float));
3132 }
3133 }
3134 proj_in = proj_scratch;
3135 }
3136 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
3137
3138 /* Residual add */
3139 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
3140
3141 /* RMSNorm before MLP */
3142 rmsnorm_forward(residual1,
3143 ln2_gamma,
3144 ln2_out,
3145 NULL,
3146 num_tokens,
3148 aligned_embed_dim,
3149 1e-06f);
3150
3151 /* MLP (SwiGLU) */
3152 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
3153 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
3154 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
3155
3156 /* Final residual add */
3157 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
3158}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_18_decode()

static void model_layer_18_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 7447 of file v6.6/test_generated/generic_api_test.c.

7454 {
7455 const MODELLayerOffsets *L = &MODEL_LAYERS[18];
7456
7457 float *input = MODEL_PTR(model, MODEL_LAYERS[17].output);
7458
7459 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
7460 float *ln1_out = MODEL_PTR(model, L->ln1_out);
7461 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
7462 float *ln2_out = MODEL_PTR(model, L->ln2_out);
7463 float *k_cache = MODEL_PTR(model, L->k);
7464 float *v_cache = MODEL_PTR(model, L->v);
7465 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
7466 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
7467 float *residual1 = MODEL_PTR(model, L->residual1);
7468 float *mlp_out = MODEL_PTR(model, L->mlp_out);
7469 float *output = MODEL_PTR(model, L->output);
7470
7471 /* Weights (explicit types for layer 18) */
7472 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
7473 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
7474 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
7475 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
7476 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
7477 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
7478
7479 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
7480 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
7481
7482 const int H = MODEL_NUM_HEADS;
7483 const int H_kv = MODEL_NUM_KV_HEADS;
7484 const int head_dim = MODEL_HEAD_DIM;
7485
7486 float q_token[H * aligned_head_dim];
7487 float k_token[H_kv * aligned_head_dim];
7488 float v_token[H_kv * aligned_head_dim];
7489 float attn_token[H * aligned_head_dim];
7490
7491 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
7492 float fc1_out[2 * aligned_intermediate_dim];
7493 float swiglu_out[aligned_intermediate_dim];
7494
7495 /* Step 1: RMSNorm before attention */
7496 rmsnorm_forward(input,
7497 ln1_gamma,
7498 ln1_out,
7499 NULL,
7500 1,
7502 aligned_embed_dim,
7503 1e-06f);
7504
7505 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
7506
7507 /* Step 2: QKV projection */
7508 /* Q projection: Q4_K -> gemm_nt_q4_k */
7509 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
7510 if (aligned_head_dim > head_dim) {
7511 for (int h = 0; h < H; ++h) {
7512 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
7513 for (int d = head_dim; d < aligned_head_dim; ++d) {
7514 q_head[d] = 0.0f;
7515 }
7516 }
7517 }
7518
7519 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7520 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7521 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
7522 const uint8_t *WK_bytes = (const uint8_t *)WK;
7523 for (int h = 0; h < H_kv; ++h) {
7524 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
7525 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
7526 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7527 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
7528 for (int d = head_dim; d < aligned_head_dim; ++d) {
7529 k_head[d] = 0.0f;
7530 }
7531 }
7532
7533 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7534 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7535 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
7536 const uint8_t *WV_bytes = (const uint8_t *)WV;
7537 for (int h = 0; h < H_kv; ++h) {
7538 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
7539 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
7540 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7541 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
7542 for (int d = head_dim; d < aligned_head_dim; ++d) {
7543 v_head[d] = 0.0f;
7544 }
7545 }
7546
7547 /* Step 3: RoPE */
7548 rope_forward(q_token,
7549 rope_cos,
7550 rope_sin,
7551 H,
7552 1,
7553 head_dim,
7554 aligned_head_dim,
7555 token_index);
7556 for (int h = 0; h < H_kv; ++h) {
7557 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7558 rope_forward(k_head,
7559 rope_cos,
7560 rope_sin,
7561 1,
7562 1,
7563 head_dim,
7564 aligned_head_dim,
7565 token_index);
7566 }
7567
7568 /* Step 4: KV cache write (direct-to-cache) */
7569
7570 /* Step 5: Attention (decode, flash) */
7572 k_cache,
7573 v_cache,
7574 attn_token,
7575 H,
7576 H_kv,
7577 token_index + 1,
7578 aligned_context_window,
7579 head_dim,
7580 aligned_head_dim);
7581
7582 /* Step 6: Output projection */
7583 /* WO projection: Q4_K -> gemm_nt_q4_k */
7584 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
7585
7586 /* Step 7: Residual add */
7587 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
7588
7589 /* Step 8: RMSNorm before MLP */
7590 rmsnorm_forward(residual1,
7591 ln2_gamma,
7592 ln2_out,
7593 NULL,
7594 1,
7596 aligned_embed_dim,
7597 1e-06f);
7598
7599 /* Step 9: MLP (SwiGLU) */
7600 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
7601 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
7602
7603 /* SwiGLU activation */
7604 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
7605
7606 /* Down projection: Q4_K -> gemm_nt_q4_k */
7607 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
7608
7609 /* Step 10: Final residual add */
7610 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
7611}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_18_prefill()

static void model_layer_18_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 3163 of file v6.6/test_generated/generic_api_test.c.

3170 {
3171 const MODELLayerOffsets *L = &MODEL_LAYERS[18];
3172
3173 float *input = MODEL_PTR(model, MODEL_LAYERS[17].output);
3174 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3175 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3176 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3177 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3178 float *q = MODEL_PTR(model, L->q);
3179 float *k = MODEL_PTR(model, L->k);
3180 float *v = MODEL_PTR(model, L->v);
3181 float *attn_out = MODEL_PTR(model, L->attn_out);
3182 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
3183 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
3184 float *residual1 = MODEL_PTR(model, L->residual1);
3185 float *fc1_out = MODEL_PTR(model, L->fc1_out);
3186 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
3187 float *mlp_out = MODEL_PTR(model, L->mlp_out);
3188 float *output = MODEL_PTR(model, L->output);
3189
3190 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
3191 const void *WK = (const void *)MODEL_PTR(model, L->wk);
3192 const void *WV = (const void *)MODEL_PTR(model, L->wv);
3193 const void *WO = (const void *)MODEL_PTR(model, L->wo);
3194 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
3195 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
3196 const float *BQ = NULL;
3197 const float *BK = NULL;
3198 const float *BV = NULL;
3199 const float *BO = NULL;
3200 const float *B1 = NULL;
3201 const float *B2 = NULL;
3202
3203 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
3204 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
3205
3206 const int H = MODEL_NUM_HEADS;
3207 const int H_kv = MODEL_NUM_KV_HEADS;
3208 const int head_dim = MODEL_HEAD_DIM;
3209 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
3210 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
3211 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
3212
3213 /* RMSNorm before attention */
3214 rmsnorm_forward(input,
3215 ln1_gamma,
3216 ln1_out,
3217 NULL,
3218 num_tokens,
3220 aligned_embed_dim,
3221 1e-06f);
3222
3223 /* Q projection (head-major) */
3224 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3225 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
3226 for (int h = 0; h < H; ++h) {
3227 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
3228 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
3229 float *q_h = q + (size_t)h * q_head_stride;
3230 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3231 }
3232
3233 /* K projection (head-major) */
3234 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3235 const uint8_t *WK_bytes = (const uint8_t *)WK;
3236 for (int h = 0; h < H_kv; ++h) {
3237 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
3238 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
3239 float *k_h = k + (size_t)h * kv_head_stride;
3240 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3241 }
3242
3243 /* V projection (head-major) */
3244 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3245 const uint8_t *WV_bytes = (const uint8_t *)WV;
3246 for (int h = 0; h < H_kv; ++h) {
3247 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
3248 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
3249 float *v_h = v + (size_t)h * kv_head_stride;
3250 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3251 }
3252
3253 /* RoPE */
3255 k,
3256 rope_cos,
3257 rope_sin,
3258 H,
3259 H_kv,
3260 num_tokens,
3261 head_dim,
3262 aligned_head_dim,
3263 0,
3264 num_tokens,
3265 aligned_context_window);
3266
3267 /* Attention (prefill, causal) */
3269 k,
3270 v,
3271 attn_out,
3272 H,
3273 H_kv,
3274 num_tokens,
3275 head_dim,
3276 aligned_head_dim,
3277 aligned_context_window);
3278
3279 /* Output projection (flatten head-major to token-major) */
3280 const int K = H * aligned_head_dim;
3281 if (K != aligned_embed_dim) {
3282 return;
3283 }
3284 const float *proj_in = attn_out;
3285 if (H > 1) {
3286 if (!proj_scratch) {
3287 return;
3288 }
3289 for (int t = 0; t < num_tokens; ++t) {
3290 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
3291 for (int h = 0; h < H; ++h) {
3292 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
3293 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
3294 src,
3295 (size_t)aligned_head_dim * sizeof(float));
3296 }
3297 }
3298 proj_in = proj_scratch;
3299 }
3300 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
3301
3302 /* Residual add */
3303 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
3304
3305 /* RMSNorm before MLP */
3306 rmsnorm_forward(residual1,
3307 ln2_gamma,
3308 ln2_out,
3309 NULL,
3310 num_tokens,
3312 aligned_embed_dim,
3313 1e-06f);
3314
3315 /* MLP (SwiGLU) */
3316 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
3317 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
3318 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
3319
3320 /* Final residual add */
3321 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
3322}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_19_decode()

static void model_layer_19_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 7616 of file v6.6/test_generated/generic_api_test.c.

7623 {
7624 const MODELLayerOffsets *L = &MODEL_LAYERS[19];
7625
7626 float *input = MODEL_PTR(model, MODEL_LAYERS[18].output);
7627
7628 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
7629 float *ln1_out = MODEL_PTR(model, L->ln1_out);
7630 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
7631 float *ln2_out = MODEL_PTR(model, L->ln2_out);
7632 float *k_cache = MODEL_PTR(model, L->k);
7633 float *v_cache = MODEL_PTR(model, L->v);
7634 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
7635 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
7636 float *residual1 = MODEL_PTR(model, L->residual1);
7637 float *mlp_out = MODEL_PTR(model, L->mlp_out);
7638 float *output = MODEL_PTR(model, L->output);
7639
7640 /* Weights (explicit types for layer 19) */
7641 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
7642 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
7643 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
7644 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
7645 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
7646 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
7647
7648 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
7649 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
7650
7651 const int H = MODEL_NUM_HEADS;
7652 const int H_kv = MODEL_NUM_KV_HEADS;
7653 const int head_dim = MODEL_HEAD_DIM;
7654
7655 float q_token[H * aligned_head_dim];
7656 float k_token[H_kv * aligned_head_dim];
7657 float v_token[H_kv * aligned_head_dim];
7658 float attn_token[H * aligned_head_dim];
7659
7660 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
7661 float fc1_out[2 * aligned_intermediate_dim];
7662 float swiglu_out[aligned_intermediate_dim];
7663
7664 /* Step 1: RMSNorm before attention */
7665 rmsnorm_forward(input,
7666 ln1_gamma,
7667 ln1_out,
7668 NULL,
7669 1,
7671 aligned_embed_dim,
7672 1e-06f);
7673
7674 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
7675
7676 /* Step 2: QKV projection */
7677 /* Q projection: Q4_K -> gemm_nt_q4_k */
7678 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
7679 if (aligned_head_dim > head_dim) {
7680 for (int h = 0; h < H; ++h) {
7681 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
7682 for (int d = head_dim; d < aligned_head_dim; ++d) {
7683 q_head[d] = 0.0f;
7684 }
7685 }
7686 }
7687
7688 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7689 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7690 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
7691 const uint8_t *WK_bytes = (const uint8_t *)WK;
7692 for (int h = 0; h < H_kv; ++h) {
7693 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
7694 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
7695 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7696 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
7697 for (int d = head_dim; d < aligned_head_dim; ++d) {
7698 k_head[d] = 0.0f;
7699 }
7700 }
7701
7702 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7703 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7704 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
7705 const uint8_t *WV_bytes = (const uint8_t *)WV;
7706 for (int h = 0; h < H_kv; ++h) {
7707 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
7708 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
7709 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7710 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
7711 for (int d = head_dim; d < aligned_head_dim; ++d) {
7712 v_head[d] = 0.0f;
7713 }
7714 }
7715
7716 /* Step 3: RoPE */
7717 rope_forward(q_token,
7718 rope_cos,
7719 rope_sin,
7720 H,
7721 1,
7722 head_dim,
7723 aligned_head_dim,
7724 token_index);
7725 for (int h = 0; h < H_kv; ++h) {
7726 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7727 rope_forward(k_head,
7728 rope_cos,
7729 rope_sin,
7730 1,
7731 1,
7732 head_dim,
7733 aligned_head_dim,
7734 token_index);
7735 }
7736
7737 /* Step 4: KV cache write (direct-to-cache) */
7738
7739 /* Step 5: Attention (decode, flash) */
7741 k_cache,
7742 v_cache,
7743 attn_token,
7744 H,
7745 H_kv,
7746 token_index + 1,
7747 aligned_context_window,
7748 head_dim,
7749 aligned_head_dim);
7750
7751 /* Step 6: Output projection */
7752 /* WO projection: Q4_K -> gemm_nt_q4_k */
7753 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
7754
7755 /* Step 7: Residual add */
7756 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
7757
7758 /* Step 8: RMSNorm before MLP */
7759 rmsnorm_forward(residual1,
7760 ln2_gamma,
7761 ln2_out,
7762 NULL,
7763 1,
7765 aligned_embed_dim,
7766 1e-06f);
7767
7768 /* Step 9: MLP (SwiGLU) */
7769 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
7770 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
7771
7772 /* SwiGLU activation */
7773 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
7774
7775 /* Down projection: Q4_K -> gemm_nt_q4_k */
7776 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
7777
7778 /* Step 10: Final residual add */
7779 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
7780}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_19_prefill()

static void model_layer_19_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 3327 of file v6.6/test_generated/generic_api_test.c.

3334 {
3335 const MODELLayerOffsets *L = &MODEL_LAYERS[19];
3336
3337 float *input = MODEL_PTR(model, MODEL_LAYERS[18].output);
3338 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3339 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3340 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3341 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3342 float *q = MODEL_PTR(model, L->q);
3343 float *k = MODEL_PTR(model, L->k);
3344 float *v = MODEL_PTR(model, L->v);
3345 float *attn_out = MODEL_PTR(model, L->attn_out);
3346 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
3347 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
3348 float *residual1 = MODEL_PTR(model, L->residual1);
3349 float *fc1_out = MODEL_PTR(model, L->fc1_out);
3350 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
3351 float *mlp_out = MODEL_PTR(model, L->mlp_out);
3352 float *output = MODEL_PTR(model, L->output);
3353
3354 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
3355 const void *WK = (const void *)MODEL_PTR(model, L->wk);
3356 const void *WV = (const void *)MODEL_PTR(model, L->wv);
3357 const void *WO = (const void *)MODEL_PTR(model, L->wo);
3358 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
3359 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
3360 const float *BQ = NULL;
3361 const float *BK = NULL;
3362 const float *BV = NULL;
3363 const float *BO = NULL;
3364 const float *B1 = NULL;
3365 const float *B2 = NULL;
3366
3367 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
3368 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
3369
3370 const int H = MODEL_NUM_HEADS;
3371 const int H_kv = MODEL_NUM_KV_HEADS;
3372 const int head_dim = MODEL_HEAD_DIM;
3373 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
3374 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
3375 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
3376
3377 /* RMSNorm before attention */
3378 rmsnorm_forward(input,
3379 ln1_gamma,
3380 ln1_out,
3381 NULL,
3382 num_tokens,
3384 aligned_embed_dim,
3385 1e-06f);
3386
3387 /* Q projection (head-major) */
3388 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3389 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
3390 for (int h = 0; h < H; ++h) {
3391 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
3392 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
3393 float *q_h = q + (size_t)h * q_head_stride;
3394 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3395 }
3396
3397 /* K projection (head-major) */
3398 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3399 const uint8_t *WK_bytes = (const uint8_t *)WK;
3400 for (int h = 0; h < H_kv; ++h) {
3401 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
3402 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
3403 float *k_h = k + (size_t)h * kv_head_stride;
3404 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3405 }
3406
3407 /* V projection (head-major) */
3408 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3409 const uint8_t *WV_bytes = (const uint8_t *)WV;
3410 for (int h = 0; h < H_kv; ++h) {
3411 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
3412 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
3413 float *v_h = v + (size_t)h * kv_head_stride;
3414 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3415 }
3416
3417 /* RoPE */
3419 k,
3420 rope_cos,
3421 rope_sin,
3422 H,
3423 H_kv,
3424 num_tokens,
3425 head_dim,
3426 aligned_head_dim,
3427 0,
3428 num_tokens,
3429 aligned_context_window);
3430
3431 /* Attention (prefill, causal) */
3433 k,
3434 v,
3435 attn_out,
3436 H,
3437 H_kv,
3438 num_tokens,
3439 head_dim,
3440 aligned_head_dim,
3441 aligned_context_window);
3442
3443 /* Output projection (flatten head-major to token-major) */
3444 const int K = H * aligned_head_dim;
3445 if (K != aligned_embed_dim) {
3446 return;
3447 }
3448 const float *proj_in = attn_out;
3449 if (H > 1) {
3450 if (!proj_scratch) {
3451 return;
3452 }
3453 for (int t = 0; t < num_tokens; ++t) {
3454 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
3455 for (int h = 0; h < H; ++h) {
3456 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
3457 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
3458 src,
3459 (size_t)aligned_head_dim * sizeof(float));
3460 }
3461 }
3462 proj_in = proj_scratch;
3463 }
3464 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
3465
3466 /* Residual add */
3467 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
3468
3469 /* RMSNorm before MLP */
3470 rmsnorm_forward(residual1,
3471 ln2_gamma,
3472 ln2_out,
3473 NULL,
3474 num_tokens,
3476 aligned_embed_dim,
3477 1e-06f);
3478
3479 /* MLP (SwiGLU) */
3480 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
3481 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
3482 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
3483
3484 /* Final residual add */
3485 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
3486}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_1_decode()

static void model_layer_1_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 4574 of file v6.6/test_generated/generic_api_test.c.

4581 {
4582 const MODELLayerOffsets *L = &MODEL_LAYERS[1];
4583
4584 float *input = MODEL_PTR(model, MODEL_LAYERS[0].output);
4585
4586 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
4587 float *ln1_out = MODEL_PTR(model, L->ln1_out);
4588 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
4589 float *ln2_out = MODEL_PTR(model, L->ln2_out);
4590 float *k_cache = MODEL_PTR(model, L->k);
4591 float *v_cache = MODEL_PTR(model, L->v);
4592 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
4593 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
4594 float *residual1 = MODEL_PTR(model, L->residual1);
4595 float *mlp_out = MODEL_PTR(model, L->mlp_out);
4596 float *output = MODEL_PTR(model, L->output);
4597
4598 /* Weights (explicit types for layer 1) */
4599 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
4600 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
4601 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
4602 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
4603 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
4604 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
4605
4606 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
4607 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
4608
4609 const int H = MODEL_NUM_HEADS;
4610 const int H_kv = MODEL_NUM_KV_HEADS;
4611 const int head_dim = MODEL_HEAD_DIM;
4612
4613 float q_token[H * aligned_head_dim];
4614 float k_token[H_kv * aligned_head_dim];
4615 float v_token[H_kv * aligned_head_dim];
4616 float attn_token[H * aligned_head_dim];
4617
4618 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
4619 float fc1_out[2 * aligned_intermediate_dim];
4620 float swiglu_out[aligned_intermediate_dim];
4621
4622 /* Step 1: RMSNorm before attention */
4623 rmsnorm_forward(input,
4624 ln1_gamma,
4625 ln1_out,
4626 NULL,
4627 1,
4629 aligned_embed_dim,
4630 1e-06f);
4631
4632 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
4633
4634 /* Step 2: QKV projection */
4635 /* Q projection: Q4_K -> gemm_nt_q4_k */
4636 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
4637 if (aligned_head_dim > head_dim) {
4638 for (int h = 0; h < H; ++h) {
4639 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
4640 for (int d = head_dim; d < aligned_head_dim; ++d) {
4641 q_head[d] = 0.0f;
4642 }
4643 }
4644 }
4645
4646 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4647 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4648 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
4649 const uint8_t *WK_bytes = (const uint8_t *)WK;
4650 for (int h = 0; h < H_kv; ++h) {
4651 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
4652 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
4653 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4654 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
4655 for (int d = head_dim; d < aligned_head_dim; ++d) {
4656 k_head[d] = 0.0f;
4657 }
4658 }
4659
4660 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4661 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4662 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
4663 const uint8_t *WV_bytes = (const uint8_t *)WV;
4664 for (int h = 0; h < H_kv; ++h) {
4665 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
4666 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
4667 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4668 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
4669 for (int d = head_dim; d < aligned_head_dim; ++d) {
4670 v_head[d] = 0.0f;
4671 }
4672 }
4673
4674 /* Step 3: RoPE */
4675 rope_forward(q_token,
4676 rope_cos,
4677 rope_sin,
4678 H,
4679 1,
4680 head_dim,
4681 aligned_head_dim,
4682 token_index);
4683 for (int h = 0; h < H_kv; ++h) {
4684 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4685 rope_forward(k_head,
4686 rope_cos,
4687 rope_sin,
4688 1,
4689 1,
4690 head_dim,
4691 aligned_head_dim,
4692 token_index);
4693 }
4694
4695 /* Step 4: KV cache write (direct-to-cache) */
4696
4697 /* Step 5: Attention (decode, flash) */
4699 k_cache,
4700 v_cache,
4701 attn_token,
4702 H,
4703 H_kv,
4704 token_index + 1,
4705 aligned_context_window,
4706 head_dim,
4707 aligned_head_dim);
4708
4709 /* Step 6: Output projection */
4710 /* WO projection: Q4_K -> gemm_nt_q4_k */
4711 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
4712
4713 /* Step 7: Residual add */
4714 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
4715
4716 /* Step 8: RMSNorm before MLP */
4717 rmsnorm_forward(residual1,
4718 ln2_gamma,
4719 ln2_out,
4720 NULL,
4721 1,
4723 aligned_embed_dim,
4724 1e-06f);
4725
4726 /* Step 9: MLP (SwiGLU) */
4727 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
4728 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
4729
4730 /* SwiGLU activation */
4731 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
4732
4733 /* Down projection: Q4_K -> gemm_nt_q4_k */
4734 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
4735
4736 /* Step 10: Final residual add */
4737 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
4738}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_1_prefill()

static void model_layer_1_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 375 of file v6.6/test_generated/generic_api_test.c.

382 {
383 const MODELLayerOffsets *L = &MODEL_LAYERS[1];
384
385 float *input = MODEL_PTR(model, MODEL_LAYERS[0].output);
386 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
387 float *ln1_out = MODEL_PTR(model, L->ln1_out);
388 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
389 float *ln2_out = MODEL_PTR(model, L->ln2_out);
390 float *q = MODEL_PTR(model, L->q);
391 float *k = MODEL_PTR(model, L->k);
392 float *v = MODEL_PTR(model, L->v);
393 float *attn_out = MODEL_PTR(model, L->attn_out);
394 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
395 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
396 float *residual1 = MODEL_PTR(model, L->residual1);
397 float *fc1_out = MODEL_PTR(model, L->fc1_out);
398 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
399 float *mlp_out = MODEL_PTR(model, L->mlp_out);
400 float *output = MODEL_PTR(model, L->output);
401
402 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
403 const void *WK = (const void *)MODEL_PTR(model, L->wk);
404 const void *WV = (const void *)MODEL_PTR(model, L->wv);
405 const void *WO = (const void *)MODEL_PTR(model, L->wo);
406 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
407 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
408 const float *BQ = NULL;
409 const float *BK = NULL;
410 const float *BV = NULL;
411 const float *BO = NULL;
412 const float *B1 = NULL;
413 const float *B2 = NULL;
414
415 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
416 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
417
418 const int H = MODEL_NUM_HEADS;
419 const int H_kv = MODEL_NUM_KV_HEADS;
420 const int head_dim = MODEL_HEAD_DIM;
421 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
422 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
423 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
424
425 /* RMSNorm before attention */
426 rmsnorm_forward(input,
427 ln1_gamma,
428 ln1_out,
429 NULL,
430 num_tokens,
432 aligned_embed_dim,
433 1e-06f);
434
435 /* Q projection (head-major) */
436 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
437 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
438 for (int h = 0; h < H; ++h) {
439 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
440 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
441 float *q_h = q + (size_t)h * q_head_stride;
442 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
443 }
444
445 /* K projection (head-major) */
446 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
447 const uint8_t *WK_bytes = (const uint8_t *)WK;
448 for (int h = 0; h < H_kv; ++h) {
449 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
450 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
451 float *k_h = k + (size_t)h * kv_head_stride;
452 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
453 }
454
455 /* V projection (head-major) */
456 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
457 const uint8_t *WV_bytes = (const uint8_t *)WV;
458 for (int h = 0; h < H_kv; ++h) {
459 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
460 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
461 float *v_h = v + (size_t)h * kv_head_stride;
462 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
463 }
464
465 /* RoPE */
467 k,
468 rope_cos,
469 rope_sin,
470 H,
471 H_kv,
472 num_tokens,
473 head_dim,
474 aligned_head_dim,
475 0,
476 num_tokens,
477 aligned_context_window);
478
479 /* Attention (prefill, causal) */
481 k,
482 v,
483 attn_out,
484 H,
485 H_kv,
486 num_tokens,
487 head_dim,
488 aligned_head_dim,
489 aligned_context_window);
490
491 /* Output projection (flatten head-major to token-major) */
492 const int K = H * aligned_head_dim;
493 if (K != aligned_embed_dim) {
494 return;
495 }
496 const float *proj_in = attn_out;
497 if (H > 1) {
498 if (!proj_scratch) {
499 return;
500 }
501 for (int t = 0; t < num_tokens; ++t) {
502 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
503 for (int h = 0; h < H; ++h) {
504 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
505 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
506 src,
507 (size_t)aligned_head_dim * sizeof(float));
508 }
509 }
510 proj_in = proj_scratch;
511 }
512 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
513
514 /* Residual add */
515 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
516
517 /* RMSNorm before MLP */
518 rmsnorm_forward(residual1,
519 ln2_gamma,
520 ln2_out,
521 NULL,
522 num_tokens,
524 aligned_embed_dim,
525 1e-06f);
526
527 /* MLP (SwiGLU) */
528 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
529 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
530 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
531
532 /* Final residual add */
533 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
534}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_20_decode()

static void model_layer_20_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 7785 of file v6.6/test_generated/generic_api_test.c.

7792 {
7793 const MODELLayerOffsets *L = &MODEL_LAYERS[20];
7794
7795 float *input = MODEL_PTR(model, MODEL_LAYERS[19].output);
7796
7797 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
7798 float *ln1_out = MODEL_PTR(model, L->ln1_out);
7799 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
7800 float *ln2_out = MODEL_PTR(model, L->ln2_out);
7801 float *k_cache = MODEL_PTR(model, L->k);
7802 float *v_cache = MODEL_PTR(model, L->v);
7803 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
7804 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
7805 float *residual1 = MODEL_PTR(model, L->residual1);
7806 float *mlp_out = MODEL_PTR(model, L->mlp_out);
7807 float *output = MODEL_PTR(model, L->output);
7808
7809 /* Weights (explicit types for layer 20) */
7810 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
7811 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
7812 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
7813 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
7814 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
7815 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
7816
7817 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
7818 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
7819
7820 const int H = MODEL_NUM_HEADS;
7821 const int H_kv = MODEL_NUM_KV_HEADS;
7822 const int head_dim = MODEL_HEAD_DIM;
7823
7824 float q_token[H * aligned_head_dim];
7825 float k_token[H_kv * aligned_head_dim];
7826 float v_token[H_kv * aligned_head_dim];
7827 float attn_token[H * aligned_head_dim];
7828
7829 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
7830 float fc1_out[2 * aligned_intermediate_dim];
7831 float swiglu_out[aligned_intermediate_dim];
7832
7833 /* Step 1: RMSNorm before attention */
7834 rmsnorm_forward(input,
7835 ln1_gamma,
7836 ln1_out,
7837 NULL,
7838 1,
7840 aligned_embed_dim,
7841 1e-06f);
7842
7843 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
7844
7845 /* Step 2: QKV projection */
7846 /* Q projection: Q4_K -> gemm_nt_q4_k */
7847 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
7848 if (aligned_head_dim > head_dim) {
7849 for (int h = 0; h < H; ++h) {
7850 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
7851 for (int d = head_dim; d < aligned_head_dim; ++d) {
7852 q_head[d] = 0.0f;
7853 }
7854 }
7855 }
7856
7857 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7858 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7859 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
7860 const uint8_t *WK_bytes = (const uint8_t *)WK;
7861 for (int h = 0; h < H_kv; ++h) {
7862 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
7863 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
7864 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7865 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
7866 for (int d = head_dim; d < aligned_head_dim; ++d) {
7867 k_head[d] = 0.0f;
7868 }
7869 }
7870
7871 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
7872 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
7873 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
7874 const uint8_t *WV_bytes = (const uint8_t *)WV;
7875 for (int h = 0; h < H_kv; ++h) {
7876 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
7877 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
7878 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7879 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
7880 for (int d = head_dim; d < aligned_head_dim; ++d) {
7881 v_head[d] = 0.0f;
7882 }
7883 }
7884
7885 /* Step 3: RoPE */
7886 rope_forward(q_token,
7887 rope_cos,
7888 rope_sin,
7889 H,
7890 1,
7891 head_dim,
7892 aligned_head_dim,
7893 token_index);
7894 for (int h = 0; h < H_kv; ++h) {
7895 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
7896 rope_forward(k_head,
7897 rope_cos,
7898 rope_sin,
7899 1,
7900 1,
7901 head_dim,
7902 aligned_head_dim,
7903 token_index);
7904 }
7905
7906 /* Step 4: KV cache write (direct-to-cache) */
7907
7908 /* Step 5: Attention (decode, flash) */
7910 k_cache,
7911 v_cache,
7912 attn_token,
7913 H,
7914 H_kv,
7915 token_index + 1,
7916 aligned_context_window,
7917 head_dim,
7918 aligned_head_dim);
7919
7920 /* Step 6: Output projection */
7921 /* WO projection: Q4_K -> gemm_nt_q4_k */
7922 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
7923
7924 /* Step 7: Residual add */
7925 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
7926
7927 /* Step 8: RMSNorm before MLP */
7928 rmsnorm_forward(residual1,
7929 ln2_gamma,
7930 ln2_out,
7931 NULL,
7932 1,
7934 aligned_embed_dim,
7935 1e-06f);
7936
7937 /* Step 9: MLP (SwiGLU) */
7938 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
7939 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
7940
7941 /* SwiGLU activation */
7942 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
7943
7944 /* Down projection: Q4_K -> gemm_nt_q4_k */
7945 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
7946
7947 /* Step 10: Final residual add */
7948 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
7949}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_20_prefill()

static void model_layer_20_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 3491 of file v6.6/test_generated/generic_api_test.c.

3498 {
3499 const MODELLayerOffsets *L = &MODEL_LAYERS[20];
3500
3501 float *input = MODEL_PTR(model, MODEL_LAYERS[19].output);
3502 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3503 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3504 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3505 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3506 float *q = MODEL_PTR(model, L->q);
3507 float *k = MODEL_PTR(model, L->k);
3508 float *v = MODEL_PTR(model, L->v);
3509 float *attn_out = MODEL_PTR(model, L->attn_out);
3510 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
3511 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
3512 float *residual1 = MODEL_PTR(model, L->residual1);
3513 float *fc1_out = MODEL_PTR(model, L->fc1_out);
3514 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
3515 float *mlp_out = MODEL_PTR(model, L->mlp_out);
3516 float *output = MODEL_PTR(model, L->output);
3517
3518 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
3519 const void *WK = (const void *)MODEL_PTR(model, L->wk);
3520 const void *WV = (const void *)MODEL_PTR(model, L->wv);
3521 const void *WO = (const void *)MODEL_PTR(model, L->wo);
3522 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
3523 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
3524 const float *BQ = NULL;
3525 const float *BK = NULL;
3526 const float *BV = NULL;
3527 const float *BO = NULL;
3528 const float *B1 = NULL;
3529 const float *B2 = NULL;
3530
3531 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
3532 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
3533
3534 const int H = MODEL_NUM_HEADS;
3535 const int H_kv = MODEL_NUM_KV_HEADS;
3536 const int head_dim = MODEL_HEAD_DIM;
3537 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
3538 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
3539 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
3540
3541 /* RMSNorm before attention */
3542 rmsnorm_forward(input,
3543 ln1_gamma,
3544 ln1_out,
3545 NULL,
3546 num_tokens,
3548 aligned_embed_dim,
3549 1e-06f);
3550
3551 /* Q projection (head-major) */
3552 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3553 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
3554 for (int h = 0; h < H; ++h) {
3555 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
3556 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
3557 float *q_h = q + (size_t)h * q_head_stride;
3558 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3559 }
3560
3561 /* K projection (head-major) */
3562 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3563 const uint8_t *WK_bytes = (const uint8_t *)WK;
3564 for (int h = 0; h < H_kv; ++h) {
3565 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
3566 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
3567 float *k_h = k + (size_t)h * kv_head_stride;
3568 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3569 }
3570
3571 /* V projection (head-major) */
3572 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3573 const uint8_t *WV_bytes = (const uint8_t *)WV;
3574 for (int h = 0; h < H_kv; ++h) {
3575 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
3576 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
3577 float *v_h = v + (size_t)h * kv_head_stride;
3578 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3579 }
3580
3581 /* RoPE */
3583 k,
3584 rope_cos,
3585 rope_sin,
3586 H,
3587 H_kv,
3588 num_tokens,
3589 head_dim,
3590 aligned_head_dim,
3591 0,
3592 num_tokens,
3593 aligned_context_window);
3594
3595 /* Attention (prefill, causal) */
3597 k,
3598 v,
3599 attn_out,
3600 H,
3601 H_kv,
3602 num_tokens,
3603 head_dim,
3604 aligned_head_dim,
3605 aligned_context_window);
3606
3607 /* Output projection (flatten head-major to token-major) */
3608 const int K = H * aligned_head_dim;
3609 if (K != aligned_embed_dim) {
3610 return;
3611 }
3612 const float *proj_in = attn_out;
3613 if (H > 1) {
3614 if (!proj_scratch) {
3615 return;
3616 }
3617 for (int t = 0; t < num_tokens; ++t) {
3618 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
3619 for (int h = 0; h < H; ++h) {
3620 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
3621 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
3622 src,
3623 (size_t)aligned_head_dim * sizeof(float));
3624 }
3625 }
3626 proj_in = proj_scratch;
3627 }
3628 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
3629
3630 /* Residual add */
3631 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
3632
3633 /* RMSNorm before MLP */
3634 rmsnorm_forward(residual1,
3635 ln2_gamma,
3636 ln2_out,
3637 NULL,
3638 num_tokens,
3640 aligned_embed_dim,
3641 1e-06f);
3642
3643 /* MLP (SwiGLU) */
3644 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
3645 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
3646 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
3647
3648 /* Final residual add */
3649 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
3650}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_21_decode()

static void model_layer_21_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 7954 of file v6.6/test_generated/generic_api_test.c.

7961 {
7962 const MODELLayerOffsets *L = &MODEL_LAYERS[21];
7963
7964 float *input = MODEL_PTR(model, MODEL_LAYERS[20].output);
7965
7966 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
7967 float *ln1_out = MODEL_PTR(model, L->ln1_out);
7968 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
7969 float *ln2_out = MODEL_PTR(model, L->ln2_out);
7970 float *k_cache = MODEL_PTR(model, L->k);
7971 float *v_cache = MODEL_PTR(model, L->v);
7972 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
7973 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
7974 float *residual1 = MODEL_PTR(model, L->residual1);
7975 float *mlp_out = MODEL_PTR(model, L->mlp_out);
7976 float *output = MODEL_PTR(model, L->output);
7977
7978 /* Weights (explicit types for layer 21) */
7979 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
7980 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
7981 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
7982 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
7983 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
7984 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
7985
7986 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
7987 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
7988
7989 const int H = MODEL_NUM_HEADS;
7990 const int H_kv = MODEL_NUM_KV_HEADS;
7991 const int head_dim = MODEL_HEAD_DIM;
7992
7993 float q_token[H * aligned_head_dim];
7994 float k_token[H_kv * aligned_head_dim];
7995 float v_token[H_kv * aligned_head_dim];
7996 float attn_token[H * aligned_head_dim];
7997
7998 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
7999 float fc1_out[2 * aligned_intermediate_dim];
8000 float swiglu_out[aligned_intermediate_dim];
8001
8002 /* Step 1: RMSNorm before attention */
8003 rmsnorm_forward(input,
8004 ln1_gamma,
8005 ln1_out,
8006 NULL,
8007 1,
8009 aligned_embed_dim,
8010 1e-06f);
8011
8012 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
8013
8014 /* Step 2: QKV projection */
8015 /* Q projection: Q4_K -> gemm_nt_q4_k */
8016 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
8017 if (aligned_head_dim > head_dim) {
8018 for (int h = 0; h < H; ++h) {
8019 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
8020 for (int d = head_dim; d < aligned_head_dim; ++d) {
8021 q_head[d] = 0.0f;
8022 }
8023 }
8024 }
8025
8026 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
8027 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
8028 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
8029 const uint8_t *WK_bytes = (const uint8_t *)WK;
8030 for (int h = 0; h < H_kv; ++h) {
8031 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
8032 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
8033 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8034 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
8035 for (int d = head_dim; d < aligned_head_dim; ++d) {
8036 k_head[d] = 0.0f;
8037 }
8038 }
8039
8040 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
8041 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
8042 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
8043 const uint8_t *WV_bytes = (const uint8_t *)WV;
8044 for (int h = 0; h < H_kv; ++h) {
8045 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
8046 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
8047 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8048 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
8049 for (int d = head_dim; d < aligned_head_dim; ++d) {
8050 v_head[d] = 0.0f;
8051 }
8052 }
8053
8054 /* Step 3: RoPE */
8055 rope_forward(q_token,
8056 rope_cos,
8057 rope_sin,
8058 H,
8059 1,
8060 head_dim,
8061 aligned_head_dim,
8062 token_index);
8063 for (int h = 0; h < H_kv; ++h) {
8064 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8065 rope_forward(k_head,
8066 rope_cos,
8067 rope_sin,
8068 1,
8069 1,
8070 head_dim,
8071 aligned_head_dim,
8072 token_index);
8073 }
8074
8075 /* Step 4: KV cache write (direct-to-cache) */
8076
8077 /* Step 5: Attention (decode, flash) */
8079 k_cache,
8080 v_cache,
8081 attn_token,
8082 H,
8083 H_kv,
8084 token_index + 1,
8085 aligned_context_window,
8086 head_dim,
8087 aligned_head_dim);
8088
8089 /* Step 6: Output projection */
8090 /* WO projection: Q4_K -> gemm_nt_q4_k */
8091 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
8092
8093 /* Step 7: Residual add */
8094 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
8095
8096 /* Step 8: RMSNorm before MLP */
8097 rmsnorm_forward(residual1,
8098 ln2_gamma,
8099 ln2_out,
8100 NULL,
8101 1,
8103 aligned_embed_dim,
8104 1e-06f);
8105
8106 /* Step 9: MLP (SwiGLU) */
8107 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
8108 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
8109
8110 /* SwiGLU activation */
8111 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
8112
8113 /* Down projection: Q4_K -> gemm_nt_q4_k */
8114 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
8115
8116 /* Step 10: Final residual add */
8117 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
8118}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_21_prefill()

static void model_layer_21_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 3655 of file v6.6/test_generated/generic_api_test.c.

3662 {
3663 const MODELLayerOffsets *L = &MODEL_LAYERS[21];
3664
3665 float *input = MODEL_PTR(model, MODEL_LAYERS[20].output);
3666 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3667 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3668 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3669 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3670 float *q = MODEL_PTR(model, L->q);
3671 float *k = MODEL_PTR(model, L->k);
3672 float *v = MODEL_PTR(model, L->v);
3673 float *attn_out = MODEL_PTR(model, L->attn_out);
3674 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
3675 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
3676 float *residual1 = MODEL_PTR(model, L->residual1);
3677 float *fc1_out = MODEL_PTR(model, L->fc1_out);
3678 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
3679 float *mlp_out = MODEL_PTR(model, L->mlp_out);
3680 float *output = MODEL_PTR(model, L->output);
3681
3682 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
3683 const void *WK = (const void *)MODEL_PTR(model, L->wk);
3684 const void *WV = (const void *)MODEL_PTR(model, L->wv);
3685 const void *WO = (const void *)MODEL_PTR(model, L->wo);
3686 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
3687 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
3688 const float *BQ = NULL;
3689 const float *BK = NULL;
3690 const float *BV = NULL;
3691 const float *BO = NULL;
3692 const float *B1 = NULL;
3693 const float *B2 = NULL;
3694
3695 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
3696 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
3697
3698 const int H = MODEL_NUM_HEADS;
3699 const int H_kv = MODEL_NUM_KV_HEADS;
3700 const int head_dim = MODEL_HEAD_DIM;
3701 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
3702 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
3703 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
3704
3705 /* RMSNorm before attention */
3706 rmsnorm_forward(input,
3707 ln1_gamma,
3708 ln1_out,
3709 NULL,
3710 num_tokens,
3712 aligned_embed_dim,
3713 1e-06f);
3714
3715 /* Q projection (head-major) */
3716 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3717 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
3718 for (int h = 0; h < H; ++h) {
3719 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
3720 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
3721 float *q_h = q + (size_t)h * q_head_stride;
3722 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3723 }
3724
3725 /* K projection (head-major) */
3726 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3727 const uint8_t *WK_bytes = (const uint8_t *)WK;
3728 for (int h = 0; h < H_kv; ++h) {
3729 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
3730 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
3731 float *k_h = k + (size_t)h * kv_head_stride;
3732 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3733 }
3734
3735 /* V projection (head-major) */
3736 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3737 const uint8_t *WV_bytes = (const uint8_t *)WV;
3738 for (int h = 0; h < H_kv; ++h) {
3739 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
3740 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
3741 float *v_h = v + (size_t)h * kv_head_stride;
3742 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3743 }
3744
3745 /* RoPE */
3747 k,
3748 rope_cos,
3749 rope_sin,
3750 H,
3751 H_kv,
3752 num_tokens,
3753 head_dim,
3754 aligned_head_dim,
3755 0,
3756 num_tokens,
3757 aligned_context_window);
3758
3759 /* Attention (prefill, causal) */
3761 k,
3762 v,
3763 attn_out,
3764 H,
3765 H_kv,
3766 num_tokens,
3767 head_dim,
3768 aligned_head_dim,
3769 aligned_context_window);
3770
3771 /* Output projection (flatten head-major to token-major) */
3772 const int K = H * aligned_head_dim;
3773 if (K != aligned_embed_dim) {
3774 return;
3775 }
3776 const float *proj_in = attn_out;
3777 if (H > 1) {
3778 if (!proj_scratch) {
3779 return;
3780 }
3781 for (int t = 0; t < num_tokens; ++t) {
3782 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
3783 for (int h = 0; h < H; ++h) {
3784 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
3785 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
3786 src,
3787 (size_t)aligned_head_dim * sizeof(float));
3788 }
3789 }
3790 proj_in = proj_scratch;
3791 }
3792 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
3793
3794 /* Residual add */
3795 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
3796
3797 /* RMSNorm before MLP */
3798 rmsnorm_forward(residual1,
3799 ln2_gamma,
3800 ln2_out,
3801 NULL,
3802 num_tokens,
3804 aligned_embed_dim,
3805 1e-06f);
3806
3807 /* MLP (SwiGLU) */
3808 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
3809 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
3810 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
3811
3812 /* Final residual add */
3813 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
3814}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_22_decode()

static void model_layer_22_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 8123 of file v6.6/test_generated/generic_api_test.c.

8130 {
8131 const MODELLayerOffsets *L = &MODEL_LAYERS[22];
8132
8133 float *input = MODEL_PTR(model, MODEL_LAYERS[21].output);
8134
8135 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
8136 float *ln1_out = MODEL_PTR(model, L->ln1_out);
8137 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
8138 float *ln2_out = MODEL_PTR(model, L->ln2_out);
8139 float *k_cache = MODEL_PTR(model, L->k);
8140 float *v_cache = MODEL_PTR(model, L->v);
8141 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
8142 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
8143 float *residual1 = MODEL_PTR(model, L->residual1);
8144 float *mlp_out = MODEL_PTR(model, L->mlp_out);
8145 float *output = MODEL_PTR(model, L->output);
8146
8147 /* Weights (explicit types for layer 22) */
8148 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
8149 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
8150 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
8151 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
8152 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
8153 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
8154
8155 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
8156 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
8157
8158 const int H = MODEL_NUM_HEADS;
8159 const int H_kv = MODEL_NUM_KV_HEADS;
8160 const int head_dim = MODEL_HEAD_DIM;
8161
8162 float q_token[H * aligned_head_dim];
8163 float k_token[H_kv * aligned_head_dim];
8164 float v_token[H_kv * aligned_head_dim];
8165 float attn_token[H * aligned_head_dim];
8166
8167 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
8168 float fc1_out[2 * aligned_intermediate_dim];
8169 float swiglu_out[aligned_intermediate_dim];
8170
8171 /* Step 1: RMSNorm before attention */
8172 rmsnorm_forward(input,
8173 ln1_gamma,
8174 ln1_out,
8175 NULL,
8176 1,
8178 aligned_embed_dim,
8179 1e-06f);
8180
8181 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
8182
8183 /* Step 2: QKV projection */
8184 /* Q projection: Q4_K -> gemm_nt_q4_k */
8185 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
8186 if (aligned_head_dim > head_dim) {
8187 for (int h = 0; h < H; ++h) {
8188 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
8189 for (int d = head_dim; d < aligned_head_dim; ++d) {
8190 q_head[d] = 0.0f;
8191 }
8192 }
8193 }
8194
8195 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
8196 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
8197 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
8198 const uint8_t *WK_bytes = (const uint8_t *)WK;
8199 for (int h = 0; h < H_kv; ++h) {
8200 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
8201 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
8202 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8203 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
8204 for (int d = head_dim; d < aligned_head_dim; ++d) {
8205 k_head[d] = 0.0f;
8206 }
8207 }
8208
8209 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
8210 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
8211 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
8212 const uint8_t *WV_bytes = (const uint8_t *)WV;
8213 for (int h = 0; h < H_kv; ++h) {
8214 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
8215 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
8216 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8217 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
8218 for (int d = head_dim; d < aligned_head_dim; ++d) {
8219 v_head[d] = 0.0f;
8220 }
8221 }
8222
8223 /* Step 3: RoPE */
8224 rope_forward(q_token,
8225 rope_cos,
8226 rope_sin,
8227 H,
8228 1,
8229 head_dim,
8230 aligned_head_dim,
8231 token_index);
8232 for (int h = 0; h < H_kv; ++h) {
8233 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8234 rope_forward(k_head,
8235 rope_cos,
8236 rope_sin,
8237 1,
8238 1,
8239 head_dim,
8240 aligned_head_dim,
8241 token_index);
8242 }
8243
8244 /* Step 4: KV cache write (direct-to-cache) */
8245
8246 /* Step 5: Attention (decode, flash) */
8248 k_cache,
8249 v_cache,
8250 attn_token,
8251 H,
8252 H_kv,
8253 token_index + 1,
8254 aligned_context_window,
8255 head_dim,
8256 aligned_head_dim);
8257
8258 /* Step 6: Output projection */
8259 /* WO projection: Q4_K -> gemm_nt_q4_k */
8260 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
8261
8262 /* Step 7: Residual add */
8263 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
8264
8265 /* Step 8: RMSNorm before MLP */
8266 rmsnorm_forward(residual1,
8267 ln2_gamma,
8268 ln2_out,
8269 NULL,
8270 1,
8272 aligned_embed_dim,
8273 1e-06f);
8274
8275 /* Step 9: MLP (SwiGLU) */
8276 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
8277 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
8278
8279 /* SwiGLU activation */
8280 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
8281
8282 /* Down projection: Q4_K -> gemm_nt_q4_k */
8283 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
8284
8285 /* Step 10: Final residual add */
8286 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
8287}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_22_prefill()

static void model_layer_22_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 3819 of file v6.6/test_generated/generic_api_test.c.

3826 {
3827 const MODELLayerOffsets *L = &MODEL_LAYERS[22];
3828
3829 float *input = MODEL_PTR(model, MODEL_LAYERS[21].output);
3830 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3831 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3832 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3833 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3834 float *q = MODEL_PTR(model, L->q);
3835 float *k = MODEL_PTR(model, L->k);
3836 float *v = MODEL_PTR(model, L->v);
3837 float *attn_out = MODEL_PTR(model, L->attn_out);
3838 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
3839 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
3840 float *residual1 = MODEL_PTR(model, L->residual1);
3841 float *fc1_out = MODEL_PTR(model, L->fc1_out);
3842 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
3843 float *mlp_out = MODEL_PTR(model, L->mlp_out);
3844 float *output = MODEL_PTR(model, L->output);
3845
3846 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
3847 const void *WK = (const void *)MODEL_PTR(model, L->wk);
3848 const void *WV = (const void *)MODEL_PTR(model, L->wv);
3849 const void *WO = (const void *)MODEL_PTR(model, L->wo);
3850 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
3851 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
3852 const float *BQ = NULL;
3853 const float *BK = NULL;
3854 const float *BV = NULL;
3855 const float *BO = NULL;
3856 const float *B1 = NULL;
3857 const float *B2 = NULL;
3858
3859 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
3860 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
3861
3862 const int H = MODEL_NUM_HEADS;
3863 const int H_kv = MODEL_NUM_KV_HEADS;
3864 const int head_dim = MODEL_HEAD_DIM;
3865 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
3866 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
3867 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
3868
3869 /* RMSNorm before attention */
3870 rmsnorm_forward(input,
3871 ln1_gamma,
3872 ln1_out,
3873 NULL,
3874 num_tokens,
3876 aligned_embed_dim,
3877 1e-06f);
3878
3879 /* Q projection (head-major) */
3880 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3881 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
3882 for (int h = 0; h < H; ++h) {
3883 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
3884 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
3885 float *q_h = q + (size_t)h * q_head_stride;
3886 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3887 }
3888
3889 /* K projection (head-major) */
3890 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3891 const uint8_t *WK_bytes = (const uint8_t *)WK;
3892 for (int h = 0; h < H_kv; ++h) {
3893 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
3894 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
3895 float *k_h = k + (size_t)h * kv_head_stride;
3896 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3897 }
3898
3899 /* V projection (head-major) */
3900 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
3901 const uint8_t *WV_bytes = (const uint8_t *)WV;
3902 for (int h = 0; h < H_kv; ++h) {
3903 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
3904 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
3905 float *v_h = v + (size_t)h * kv_head_stride;
3906 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
3907 }
3908
3909 /* RoPE */
3911 k,
3912 rope_cos,
3913 rope_sin,
3914 H,
3915 H_kv,
3916 num_tokens,
3917 head_dim,
3918 aligned_head_dim,
3919 0,
3920 num_tokens,
3921 aligned_context_window);
3922
3923 /* Attention (prefill, causal) */
3925 k,
3926 v,
3927 attn_out,
3928 H,
3929 H_kv,
3930 num_tokens,
3931 head_dim,
3932 aligned_head_dim,
3933 aligned_context_window);
3934
3935 /* Output projection (flatten head-major to token-major) */
3936 const int K = H * aligned_head_dim;
3937 if (K != aligned_embed_dim) {
3938 return;
3939 }
3940 const float *proj_in = attn_out;
3941 if (H > 1) {
3942 if (!proj_scratch) {
3943 return;
3944 }
3945 for (int t = 0; t < num_tokens; ++t) {
3946 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
3947 for (int h = 0; h < H; ++h) {
3948 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
3949 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
3950 src,
3951 (size_t)aligned_head_dim * sizeof(float));
3952 }
3953 }
3954 proj_in = proj_scratch;
3955 }
3956 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
3957
3958 /* Residual add */
3959 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
3960
3961 /* RMSNorm before MLP */
3962 rmsnorm_forward(residual1,
3963 ln2_gamma,
3964 ln2_out,
3965 NULL,
3966 num_tokens,
3968 aligned_embed_dim,
3969 1e-06f);
3970
3971 /* MLP (SwiGLU) */
3972 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
3973 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
3974 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
3975
3976 /* Final residual add */
3977 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
3978}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_23_decode()

static void model_layer_23_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 8292 of file v6.6/test_generated/generic_api_test.c.

8299 {
8300 const MODELLayerOffsets *L = &MODEL_LAYERS[23];
8301
8302 float *input = MODEL_PTR(model, MODEL_LAYERS[22].output);
8303
8304 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
8305 float *ln1_out = MODEL_PTR(model, L->ln1_out);
8306 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
8307 float *ln2_out = MODEL_PTR(model, L->ln2_out);
8308 float *k_cache = MODEL_PTR(model, L->k);
8309 float *v_cache = MODEL_PTR(model, L->v);
8310 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
8311 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
8312 float *residual1 = MODEL_PTR(model, L->residual1);
8313 float *mlp_out = MODEL_PTR(model, L->mlp_out);
8314 float *output = MODEL_PTR(model, L->output);
8315
8316 /* Weights (explicit types for layer 23) */
8317 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
8318 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
8319 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
8320 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
8321 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
8322 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
8323
8324 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
8325 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
8326
8327 const int H = MODEL_NUM_HEADS;
8328 const int H_kv = MODEL_NUM_KV_HEADS;
8329 const int head_dim = MODEL_HEAD_DIM;
8330
8331 float q_token[H * aligned_head_dim];
8332 float k_token[H_kv * aligned_head_dim];
8333 float v_token[H_kv * aligned_head_dim];
8334 float attn_token[H * aligned_head_dim];
8335
8336 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
8337 float fc1_out[2 * aligned_intermediate_dim];
8338 float swiglu_out[aligned_intermediate_dim];
8339
8340 /* Step 1: RMSNorm before attention */
8341 rmsnorm_forward(input,
8342 ln1_gamma,
8343 ln1_out,
8344 NULL,
8345 1,
8347 aligned_embed_dim,
8348 1e-06f);
8349
8350 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
8351
8352 /* Step 2: QKV projection */
8353 /* Q projection: Q4_K -> gemm_nt_q4_k */
8354 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
8355 if (aligned_head_dim > head_dim) {
8356 for (int h = 0; h < H; ++h) {
8357 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
8358 for (int d = head_dim; d < aligned_head_dim; ++d) {
8359 q_head[d] = 0.0f;
8360 }
8361 }
8362 }
8363
8364 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
8365 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
8366 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
8367 const uint8_t *WK_bytes = (const uint8_t *)WK;
8368 for (int h = 0; h < H_kv; ++h) {
8369 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
8370 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
8371 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8372 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
8373 for (int d = head_dim; d < aligned_head_dim; ++d) {
8374 k_head[d] = 0.0f;
8375 }
8376 }
8377
8378 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
8379 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
8380 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
8381 const uint8_t *WV_bytes = (const uint8_t *)WV;
8382 for (int h = 0; h < H_kv; ++h) {
8383 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
8384 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
8385 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8386 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
8387 for (int d = head_dim; d < aligned_head_dim; ++d) {
8388 v_head[d] = 0.0f;
8389 }
8390 }
8391
8392 /* Step 3: RoPE */
8393 rope_forward(q_token,
8394 rope_cos,
8395 rope_sin,
8396 H,
8397 1,
8398 head_dim,
8399 aligned_head_dim,
8400 token_index);
8401 for (int h = 0; h < H_kv; ++h) {
8402 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
8403 rope_forward(k_head,
8404 rope_cos,
8405 rope_sin,
8406 1,
8407 1,
8408 head_dim,
8409 aligned_head_dim,
8410 token_index);
8411 }
8412
8413 /* Step 4: KV cache write (direct-to-cache) */
8414
8415 /* Step 5: Attention (decode, flash) */
8417 k_cache,
8418 v_cache,
8419 attn_token,
8420 H,
8421 H_kv,
8422 token_index + 1,
8423 aligned_context_window,
8424 head_dim,
8425 aligned_head_dim);
8426
8427 /* Step 6: Output projection */
8428 /* WO projection: Q4_K -> gemm_nt_q4_k */
8429 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
8430
8431 /* Step 7: Residual add */
8432 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
8433
8434 /* Step 8: RMSNorm before MLP */
8435 rmsnorm_forward(residual1,
8436 ln2_gamma,
8437 ln2_out,
8438 NULL,
8439 1,
8441 aligned_embed_dim,
8442 1e-06f);
8443
8444 /* Step 9: MLP (SwiGLU) */
8445 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
8446 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
8447
8448 /* SwiGLU activation */
8449 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
8450
8451 /* Down projection: Q4_K -> gemm_nt_q4_k */
8452 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
8453
8454 /* Step 10: Final residual add */
8455 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
8456}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_23_prefill()

static void model_layer_23_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 3983 of file v6.6/test_generated/generic_api_test.c.

3990 {
3991 const MODELLayerOffsets *L = &MODEL_LAYERS[23];
3992
3993 float *input = MODEL_PTR(model, MODEL_LAYERS[22].output);
3994 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
3995 float *ln1_out = MODEL_PTR(model, L->ln1_out);
3996 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
3997 float *ln2_out = MODEL_PTR(model, L->ln2_out);
3998 float *q = MODEL_PTR(model, L->q);
3999 float *k = MODEL_PTR(model, L->k);
4000 float *v = MODEL_PTR(model, L->v);
4001 float *attn_out = MODEL_PTR(model, L->attn_out);
4002 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
4003 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
4004 float *residual1 = MODEL_PTR(model, L->residual1);
4005 float *fc1_out = MODEL_PTR(model, L->fc1_out);
4006 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
4007 float *mlp_out = MODEL_PTR(model, L->mlp_out);
4008 float *output = MODEL_PTR(model, L->output);
4009
4010 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
4011 const void *WK = (const void *)MODEL_PTR(model, L->wk);
4012 const void *WV = (const void *)MODEL_PTR(model, L->wv);
4013 const void *WO = (const void *)MODEL_PTR(model, L->wo);
4014 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
4015 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
4016 const float *BQ = NULL;
4017 const float *BK = NULL;
4018 const float *BV = NULL;
4019 const float *BO = NULL;
4020 const float *B1 = NULL;
4021 const float *B2 = NULL;
4022
4023 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
4024 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
4025
4026 const int H = MODEL_NUM_HEADS;
4027 const int H_kv = MODEL_NUM_KV_HEADS;
4028 const int head_dim = MODEL_HEAD_DIM;
4029 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4030 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
4031 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
4032
4033 /* RMSNorm before attention */
4034 rmsnorm_forward(input,
4035 ln1_gamma,
4036 ln1_out,
4037 NULL,
4038 num_tokens,
4040 aligned_embed_dim,
4041 1e-06f);
4042
4043 /* Q projection (head-major) */
4044 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
4045 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
4046 for (int h = 0; h < H; ++h) {
4047 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
4048 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
4049 float *q_h = q + (size_t)h * q_head_stride;
4050 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
4051 }
4052
4053 /* K projection (head-major) */
4054 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
4055 const uint8_t *WK_bytes = (const uint8_t *)WK;
4056 for (int h = 0; h < H_kv; ++h) {
4057 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
4058 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
4059 float *k_h = k + (size_t)h * kv_head_stride;
4060 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
4061 }
4062
4063 /* V projection (head-major) */
4064 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
4065 const uint8_t *WV_bytes = (const uint8_t *)WV;
4066 for (int h = 0; h < H_kv; ++h) {
4067 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
4068 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
4069 float *v_h = v + (size_t)h * kv_head_stride;
4070 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
4071 }
4072
4073 /* RoPE */
4075 k,
4076 rope_cos,
4077 rope_sin,
4078 H,
4079 H_kv,
4080 num_tokens,
4081 head_dim,
4082 aligned_head_dim,
4083 0,
4084 num_tokens,
4085 aligned_context_window);
4086
4087 /* Attention (prefill, causal) */
4089 k,
4090 v,
4091 attn_out,
4092 H,
4093 H_kv,
4094 num_tokens,
4095 head_dim,
4096 aligned_head_dim,
4097 aligned_context_window);
4098
4099 /* Output projection (flatten head-major to token-major) */
4100 const int K = H * aligned_head_dim;
4101 if (K != aligned_embed_dim) {
4102 return;
4103 }
4104 const float *proj_in = attn_out;
4105 if (H > 1) {
4106 if (!proj_scratch) {
4107 return;
4108 }
4109 for (int t = 0; t < num_tokens; ++t) {
4110 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
4111 for (int h = 0; h < H; ++h) {
4112 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
4113 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
4114 src,
4115 (size_t)aligned_head_dim * sizeof(float));
4116 }
4117 }
4118 proj_in = proj_scratch;
4119 }
4120 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
4121
4122 /* Residual add */
4123 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
4124
4125 /* RMSNorm before MLP */
4126 rmsnorm_forward(residual1,
4127 ln2_gamma,
4128 ln2_out,
4129 NULL,
4130 num_tokens,
4132 aligned_embed_dim,
4133 1e-06f);
4134
4135 /* MLP (SwiGLU) */
4136 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
4137 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
4138 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
4139
4140 /* Final residual add */
4141 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
4142}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_2_decode()

static void model_layer_2_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 4743 of file v6.6/test_generated/generic_api_test.c.

4750 {
4751 const MODELLayerOffsets *L = &MODEL_LAYERS[2];
4752
4753 float *input = MODEL_PTR(model, MODEL_LAYERS[1].output);
4754
4755 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
4756 float *ln1_out = MODEL_PTR(model, L->ln1_out);
4757 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
4758 float *ln2_out = MODEL_PTR(model, L->ln2_out);
4759 float *k_cache = MODEL_PTR(model, L->k);
4760 float *v_cache = MODEL_PTR(model, L->v);
4761 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
4762 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
4763 float *residual1 = MODEL_PTR(model, L->residual1);
4764 float *mlp_out = MODEL_PTR(model, L->mlp_out);
4765 float *output = MODEL_PTR(model, L->output);
4766
4767 /* Weights (explicit types for layer 2) */
4768 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
4769 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
4770 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
4771 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
4772 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
4773 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
4774
4775 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
4776 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
4777
4778 const int H = MODEL_NUM_HEADS;
4779 const int H_kv = MODEL_NUM_KV_HEADS;
4780 const int head_dim = MODEL_HEAD_DIM;
4781
4782 float q_token[H * aligned_head_dim];
4783 float k_token[H_kv * aligned_head_dim];
4784 float v_token[H_kv * aligned_head_dim];
4785 float attn_token[H * aligned_head_dim];
4786
4787 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
4788 float fc1_out[2 * aligned_intermediate_dim];
4789 float swiglu_out[aligned_intermediate_dim];
4790
4791 /* Step 1: RMSNorm before attention */
4792 rmsnorm_forward(input,
4793 ln1_gamma,
4794 ln1_out,
4795 NULL,
4796 1,
4798 aligned_embed_dim,
4799 1e-06f);
4800
4801 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
4802
4803 /* Step 2: QKV projection */
4804 /* Q projection: Q4_K -> gemm_nt_q4_k */
4805 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
4806 if (aligned_head_dim > head_dim) {
4807 for (int h = 0; h < H; ++h) {
4808 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
4809 for (int d = head_dim; d < aligned_head_dim; ++d) {
4810 q_head[d] = 0.0f;
4811 }
4812 }
4813 }
4814
4815 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4816 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4817 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
4818 const uint8_t *WK_bytes = (const uint8_t *)WK;
4819 for (int h = 0; h < H_kv; ++h) {
4820 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
4821 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
4822 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4823 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
4824 for (int d = head_dim; d < aligned_head_dim; ++d) {
4825 k_head[d] = 0.0f;
4826 }
4827 }
4828
4829 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4830 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4831 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
4832 const uint8_t *WV_bytes = (const uint8_t *)WV;
4833 for (int h = 0; h < H_kv; ++h) {
4834 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
4835 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
4836 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4837 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
4838 for (int d = head_dim; d < aligned_head_dim; ++d) {
4839 v_head[d] = 0.0f;
4840 }
4841 }
4842
4843 /* Step 3: RoPE */
4844 rope_forward(q_token,
4845 rope_cos,
4846 rope_sin,
4847 H,
4848 1,
4849 head_dim,
4850 aligned_head_dim,
4851 token_index);
4852 for (int h = 0; h < H_kv; ++h) {
4853 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4854 rope_forward(k_head,
4855 rope_cos,
4856 rope_sin,
4857 1,
4858 1,
4859 head_dim,
4860 aligned_head_dim,
4861 token_index);
4862 }
4863
4864 /* Step 4: KV cache write (direct-to-cache) */
4865
4866 /* Step 5: Attention (decode, flash) */
4868 k_cache,
4869 v_cache,
4870 attn_token,
4871 H,
4872 H_kv,
4873 token_index + 1,
4874 aligned_context_window,
4875 head_dim,
4876 aligned_head_dim);
4877
4878 /* Step 6: Output projection */
4879 /* WO projection: Q4_K -> gemm_nt_q4_k */
4880 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
4881
4882 /* Step 7: Residual add */
4883 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
4884
4885 /* Step 8: RMSNorm before MLP */
4886 rmsnorm_forward(residual1,
4887 ln2_gamma,
4888 ln2_out,
4889 NULL,
4890 1,
4892 aligned_embed_dim,
4893 1e-06f);
4894
4895 /* Step 9: MLP (SwiGLU) */
4896 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
4897 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
4898
4899 /* SwiGLU activation */
4900 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
4901
4902 /* Down projection: Q4_K -> gemm_nt_q4_k */
4903 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
4904
4905 /* Step 10: Final residual add */
4906 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
4907}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_2_prefill()

static void model_layer_2_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 539 of file v6.6/test_generated/generic_api_test.c.

546 {
547 const MODELLayerOffsets *L = &MODEL_LAYERS[2];
548
549 float *input = MODEL_PTR(model, MODEL_LAYERS[1].output);
550 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
551 float *ln1_out = MODEL_PTR(model, L->ln1_out);
552 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
553 float *ln2_out = MODEL_PTR(model, L->ln2_out);
554 float *q = MODEL_PTR(model, L->q);
555 float *k = MODEL_PTR(model, L->k);
556 float *v = MODEL_PTR(model, L->v);
557 float *attn_out = MODEL_PTR(model, L->attn_out);
558 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
559 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
560 float *residual1 = MODEL_PTR(model, L->residual1);
561 float *fc1_out = MODEL_PTR(model, L->fc1_out);
562 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
563 float *mlp_out = MODEL_PTR(model, L->mlp_out);
564 float *output = MODEL_PTR(model, L->output);
565
566 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
567 const void *WK = (const void *)MODEL_PTR(model, L->wk);
568 const void *WV = (const void *)MODEL_PTR(model, L->wv);
569 const void *WO = (const void *)MODEL_PTR(model, L->wo);
570 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
571 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
572 const float *BQ = NULL;
573 const float *BK = NULL;
574 const float *BV = NULL;
575 const float *BO = NULL;
576 const float *B1 = NULL;
577 const float *B2 = NULL;
578
579 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
580 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
581
582 const int H = MODEL_NUM_HEADS;
583 const int H_kv = MODEL_NUM_KV_HEADS;
584 const int head_dim = MODEL_HEAD_DIM;
585 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
586 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
587 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
588
589 /* RMSNorm before attention */
590 rmsnorm_forward(input,
591 ln1_gamma,
592 ln1_out,
593 NULL,
594 num_tokens,
596 aligned_embed_dim,
597 1e-06f);
598
599 /* Q projection (head-major) */
600 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
601 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
602 for (int h = 0; h < H; ++h) {
603 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
604 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
605 float *q_h = q + (size_t)h * q_head_stride;
606 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
607 }
608
609 /* K projection (head-major) */
610 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
611 const uint8_t *WK_bytes = (const uint8_t *)WK;
612 for (int h = 0; h < H_kv; ++h) {
613 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
614 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
615 float *k_h = k + (size_t)h * kv_head_stride;
616 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
617 }
618
619 /* V projection (head-major) */
620 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
621 const uint8_t *WV_bytes = (const uint8_t *)WV;
622 for (int h = 0; h < H_kv; ++h) {
623 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
624 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
625 float *v_h = v + (size_t)h * kv_head_stride;
626 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
627 }
628
629 /* RoPE */
631 k,
632 rope_cos,
633 rope_sin,
634 H,
635 H_kv,
636 num_tokens,
637 head_dim,
638 aligned_head_dim,
639 0,
640 num_tokens,
641 aligned_context_window);
642
643 /* Attention (prefill, causal) */
645 k,
646 v,
647 attn_out,
648 H,
649 H_kv,
650 num_tokens,
651 head_dim,
652 aligned_head_dim,
653 aligned_context_window);
654
655 /* Output projection (flatten head-major to token-major) */
656 const int K = H * aligned_head_dim;
657 if (K != aligned_embed_dim) {
658 return;
659 }
660 const float *proj_in = attn_out;
661 if (H > 1) {
662 if (!proj_scratch) {
663 return;
664 }
665 for (int t = 0; t < num_tokens; ++t) {
666 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
667 for (int h = 0; h < H; ++h) {
668 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
669 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
670 src,
671 (size_t)aligned_head_dim * sizeof(float));
672 }
673 }
674 proj_in = proj_scratch;
675 }
676 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
677
678 /* Residual add */
679 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
680
681 /* RMSNorm before MLP */
682 rmsnorm_forward(residual1,
683 ln2_gamma,
684 ln2_out,
685 NULL,
686 num_tokens,
688 aligned_embed_dim,
689 1e-06f);
690
691 /* MLP (SwiGLU) */
692 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
693 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
694 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
695
696 /* Final residual add */
697 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
698}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_3_decode()

static void model_layer_3_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 4912 of file v6.6/test_generated/generic_api_test.c.

4919 {
4920 const MODELLayerOffsets *L = &MODEL_LAYERS[3];
4921
4922 float *input = MODEL_PTR(model, MODEL_LAYERS[2].output);
4923
4924 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
4925 float *ln1_out = MODEL_PTR(model, L->ln1_out);
4926 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
4927 float *ln2_out = MODEL_PTR(model, L->ln2_out);
4928 float *k_cache = MODEL_PTR(model, L->k);
4929 float *v_cache = MODEL_PTR(model, L->v);
4930 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
4931 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
4932 float *residual1 = MODEL_PTR(model, L->residual1);
4933 float *mlp_out = MODEL_PTR(model, L->mlp_out);
4934 float *output = MODEL_PTR(model, L->output);
4935
4936 /* Weights (explicit types for layer 3) */
4937 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
4938 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
4939 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
4940 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
4941 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
4942 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
4943
4944 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
4945 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
4946
4947 const int H = MODEL_NUM_HEADS;
4948 const int H_kv = MODEL_NUM_KV_HEADS;
4949 const int head_dim = MODEL_HEAD_DIM;
4950
4951 float q_token[H * aligned_head_dim];
4952 float k_token[H_kv * aligned_head_dim];
4953 float v_token[H_kv * aligned_head_dim];
4954 float attn_token[H * aligned_head_dim];
4955
4956 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
4957 float fc1_out[2 * aligned_intermediate_dim];
4958 float swiglu_out[aligned_intermediate_dim];
4959
4960 /* Step 1: RMSNorm before attention */
4961 rmsnorm_forward(input,
4962 ln1_gamma,
4963 ln1_out,
4964 NULL,
4965 1,
4967 aligned_embed_dim,
4968 1e-06f);
4969
4970 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
4971
4972 /* Step 2: QKV projection */
4973 /* Q projection: Q4_K -> gemm_nt_q4_k */
4974 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
4975 if (aligned_head_dim > head_dim) {
4976 for (int h = 0; h < H; ++h) {
4977 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
4978 for (int d = head_dim; d < aligned_head_dim; ++d) {
4979 q_head[d] = 0.0f;
4980 }
4981 }
4982 }
4983
4984 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4985 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
4986 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
4987 const uint8_t *WK_bytes = (const uint8_t *)WK;
4988 for (int h = 0; h < H_kv; ++h) {
4989 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
4990 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
4991 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
4992 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
4993 for (int d = head_dim; d < aligned_head_dim; ++d) {
4994 k_head[d] = 0.0f;
4995 }
4996 }
4997
4998 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
4999 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5000 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
5001 const uint8_t *WV_bytes = (const uint8_t *)WV;
5002 for (int h = 0; h < H_kv; ++h) {
5003 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
5004 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
5005 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5006 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
5007 for (int d = head_dim; d < aligned_head_dim; ++d) {
5008 v_head[d] = 0.0f;
5009 }
5010 }
5011
5012 /* Step 3: RoPE */
5013 rope_forward(q_token,
5014 rope_cos,
5015 rope_sin,
5016 H,
5017 1,
5018 head_dim,
5019 aligned_head_dim,
5020 token_index);
5021 for (int h = 0; h < H_kv; ++h) {
5022 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5023 rope_forward(k_head,
5024 rope_cos,
5025 rope_sin,
5026 1,
5027 1,
5028 head_dim,
5029 aligned_head_dim,
5030 token_index);
5031 }
5032
5033 /* Step 4: KV cache write (direct-to-cache) */
5034
5035 /* Step 5: Attention (decode, flash) */
5037 k_cache,
5038 v_cache,
5039 attn_token,
5040 H,
5041 H_kv,
5042 token_index + 1,
5043 aligned_context_window,
5044 head_dim,
5045 aligned_head_dim);
5046
5047 /* Step 6: Output projection */
5048 /* WO projection: Q4_K -> gemm_nt_q4_k */
5049 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
5050
5051 /* Step 7: Residual add */
5052 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
5053
5054 /* Step 8: RMSNorm before MLP */
5055 rmsnorm_forward(residual1,
5056 ln2_gamma,
5057 ln2_out,
5058 NULL,
5059 1,
5061 aligned_embed_dim,
5062 1e-06f);
5063
5064 /* Step 9: MLP (SwiGLU) */
5065 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
5066 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
5067
5068 /* SwiGLU activation */
5069 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
5070
5071 /* Down projection: Q4_K -> gemm_nt_q4_k */
5072 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
5073
5074 /* Step 10: Final residual add */
5075 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
5076}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_3_prefill()

static void model_layer_3_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 703 of file v6.6/test_generated/generic_api_test.c.

710 {
711 const MODELLayerOffsets *L = &MODEL_LAYERS[3];
712
713 float *input = MODEL_PTR(model, MODEL_LAYERS[2].output);
714 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
715 float *ln1_out = MODEL_PTR(model, L->ln1_out);
716 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
717 float *ln2_out = MODEL_PTR(model, L->ln2_out);
718 float *q = MODEL_PTR(model, L->q);
719 float *k = MODEL_PTR(model, L->k);
720 float *v = MODEL_PTR(model, L->v);
721 float *attn_out = MODEL_PTR(model, L->attn_out);
722 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
723 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
724 float *residual1 = MODEL_PTR(model, L->residual1);
725 float *fc1_out = MODEL_PTR(model, L->fc1_out);
726 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
727 float *mlp_out = MODEL_PTR(model, L->mlp_out);
728 float *output = MODEL_PTR(model, L->output);
729
730 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
731 const void *WK = (const void *)MODEL_PTR(model, L->wk);
732 const void *WV = (const void *)MODEL_PTR(model, L->wv);
733 const void *WO = (const void *)MODEL_PTR(model, L->wo);
734 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
735 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
736 const float *BQ = NULL;
737 const float *BK = NULL;
738 const float *BV = NULL;
739 const float *BO = NULL;
740 const float *B1 = NULL;
741 const float *B2 = NULL;
742
743 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
744 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
745
746 const int H = MODEL_NUM_HEADS;
747 const int H_kv = MODEL_NUM_KV_HEADS;
748 const int head_dim = MODEL_HEAD_DIM;
749 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
750 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
751 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
752
753 /* RMSNorm before attention */
754 rmsnorm_forward(input,
755 ln1_gamma,
756 ln1_out,
757 NULL,
758 num_tokens,
760 aligned_embed_dim,
761 1e-06f);
762
763 /* Q projection (head-major) */
764 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
765 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
766 for (int h = 0; h < H; ++h) {
767 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
768 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
769 float *q_h = q + (size_t)h * q_head_stride;
770 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
771 }
772
773 /* K projection (head-major) */
774 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
775 const uint8_t *WK_bytes = (const uint8_t *)WK;
776 for (int h = 0; h < H_kv; ++h) {
777 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
778 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
779 float *k_h = k + (size_t)h * kv_head_stride;
780 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
781 }
782
783 /* V projection (head-major) */
784 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
785 const uint8_t *WV_bytes = (const uint8_t *)WV;
786 for (int h = 0; h < H_kv; ++h) {
787 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
788 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
789 float *v_h = v + (size_t)h * kv_head_stride;
790 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
791 }
792
793 /* RoPE */
795 k,
796 rope_cos,
797 rope_sin,
798 H,
799 H_kv,
800 num_tokens,
801 head_dim,
802 aligned_head_dim,
803 0,
804 num_tokens,
805 aligned_context_window);
806
807 /* Attention (prefill, causal) */
809 k,
810 v,
811 attn_out,
812 H,
813 H_kv,
814 num_tokens,
815 head_dim,
816 aligned_head_dim,
817 aligned_context_window);
818
819 /* Output projection (flatten head-major to token-major) */
820 const int K = H * aligned_head_dim;
821 if (K != aligned_embed_dim) {
822 return;
823 }
824 const float *proj_in = attn_out;
825 if (H > 1) {
826 if (!proj_scratch) {
827 return;
828 }
829 for (int t = 0; t < num_tokens; ++t) {
830 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
831 for (int h = 0; h < H; ++h) {
832 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
833 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
834 src,
835 (size_t)aligned_head_dim * sizeof(float));
836 }
837 }
838 proj_in = proj_scratch;
839 }
840 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
841
842 /* Residual add */
843 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
844
845 /* RMSNorm before MLP */
846 rmsnorm_forward(residual1,
847 ln2_gamma,
848 ln2_out,
849 NULL,
850 num_tokens,
852 aligned_embed_dim,
853 1e-06f);
854
855 /* MLP (SwiGLU) */
856 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
857 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
858 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
859
860 /* Final residual add */
861 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
862}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_4_decode()

static void model_layer_4_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 5081 of file v6.6/test_generated/generic_api_test.c.

5088 {
5089 const MODELLayerOffsets *L = &MODEL_LAYERS[4];
5090
5091 float *input = MODEL_PTR(model, MODEL_LAYERS[3].output);
5092
5093 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
5094 float *ln1_out = MODEL_PTR(model, L->ln1_out);
5095 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
5096 float *ln2_out = MODEL_PTR(model, L->ln2_out);
5097 float *k_cache = MODEL_PTR(model, L->k);
5098 float *v_cache = MODEL_PTR(model, L->v);
5099 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
5100 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
5101 float *residual1 = MODEL_PTR(model, L->residual1);
5102 float *mlp_out = MODEL_PTR(model, L->mlp_out);
5103 float *output = MODEL_PTR(model, L->output);
5104
5105 /* Weights (explicit types for layer 4) */
5106 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
5107 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
5108 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
5109 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
5110 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
5111 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
5112
5113 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
5114 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
5115
5116 const int H = MODEL_NUM_HEADS;
5117 const int H_kv = MODEL_NUM_KV_HEADS;
5118 const int head_dim = MODEL_HEAD_DIM;
5119
5120 float q_token[H * aligned_head_dim];
5121 float k_token[H_kv * aligned_head_dim];
5122 float v_token[H_kv * aligned_head_dim];
5123 float attn_token[H * aligned_head_dim];
5124
5125 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
5126 float fc1_out[2 * aligned_intermediate_dim];
5127 float swiglu_out[aligned_intermediate_dim];
5128
5129 /* Step 1: RMSNorm before attention */
5130 rmsnorm_forward(input,
5131 ln1_gamma,
5132 ln1_out,
5133 NULL,
5134 1,
5136 aligned_embed_dim,
5137 1e-06f);
5138
5139 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
5140
5141 /* Step 2: QKV projection */
5142 /* Q projection: Q4_K -> gemm_nt_q4_k */
5143 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
5144 if (aligned_head_dim > head_dim) {
5145 for (int h = 0; h < H; ++h) {
5146 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
5147 for (int d = head_dim; d < aligned_head_dim; ++d) {
5148 q_head[d] = 0.0f;
5149 }
5150 }
5151 }
5152
5153 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5154 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5155 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
5156 const uint8_t *WK_bytes = (const uint8_t *)WK;
5157 for (int h = 0; h < H_kv; ++h) {
5158 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
5159 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
5160 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5161 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
5162 for (int d = head_dim; d < aligned_head_dim; ++d) {
5163 k_head[d] = 0.0f;
5164 }
5165 }
5166
5167 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5168 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5169 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
5170 const uint8_t *WV_bytes = (const uint8_t *)WV;
5171 for (int h = 0; h < H_kv; ++h) {
5172 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
5173 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
5174 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5175 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
5176 for (int d = head_dim; d < aligned_head_dim; ++d) {
5177 v_head[d] = 0.0f;
5178 }
5179 }
5180
5181 /* Step 3: RoPE */
5182 rope_forward(q_token,
5183 rope_cos,
5184 rope_sin,
5185 H,
5186 1,
5187 head_dim,
5188 aligned_head_dim,
5189 token_index);
5190 for (int h = 0; h < H_kv; ++h) {
5191 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5192 rope_forward(k_head,
5193 rope_cos,
5194 rope_sin,
5195 1,
5196 1,
5197 head_dim,
5198 aligned_head_dim,
5199 token_index);
5200 }
5201
5202 /* Step 4: KV cache write (direct-to-cache) */
5203
5204 /* Step 5: Attention (decode, flash) */
5206 k_cache,
5207 v_cache,
5208 attn_token,
5209 H,
5210 H_kv,
5211 token_index + 1,
5212 aligned_context_window,
5213 head_dim,
5214 aligned_head_dim);
5215
5216 /* Step 6: Output projection */
5217 /* WO projection: Q4_K -> gemm_nt_q4_k */
5218 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
5219
5220 /* Step 7: Residual add */
5221 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
5222
5223 /* Step 8: RMSNorm before MLP */
5224 rmsnorm_forward(residual1,
5225 ln2_gamma,
5226 ln2_out,
5227 NULL,
5228 1,
5230 aligned_embed_dim,
5231 1e-06f);
5232
5233 /* Step 9: MLP (SwiGLU) */
5234 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
5235 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
5236
5237 /* SwiGLU activation */
5238 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
5239
5240 /* Down projection: Q4_K -> gemm_nt_q4_k */
5241 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
5242
5243 /* Step 10: Final residual add */
5244 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
5245}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_4_prefill()

static void model_layer_4_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 867 of file v6.6/test_generated/generic_api_test.c.

874 {
875 const MODELLayerOffsets *L = &MODEL_LAYERS[4];
876
877 float *input = MODEL_PTR(model, MODEL_LAYERS[3].output);
878 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
879 float *ln1_out = MODEL_PTR(model, L->ln1_out);
880 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
881 float *ln2_out = MODEL_PTR(model, L->ln2_out);
882 float *q = MODEL_PTR(model, L->q);
883 float *k = MODEL_PTR(model, L->k);
884 float *v = MODEL_PTR(model, L->v);
885 float *attn_out = MODEL_PTR(model, L->attn_out);
886 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
887 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
888 float *residual1 = MODEL_PTR(model, L->residual1);
889 float *fc1_out = MODEL_PTR(model, L->fc1_out);
890 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
891 float *mlp_out = MODEL_PTR(model, L->mlp_out);
892 float *output = MODEL_PTR(model, L->output);
893
894 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
895 const void *WK = (const void *)MODEL_PTR(model, L->wk);
896 const void *WV = (const void *)MODEL_PTR(model, L->wv);
897 const void *WO = (const void *)MODEL_PTR(model, L->wo);
898 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
899 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
900 const float *BQ = NULL;
901 const float *BK = NULL;
902 const float *BV = NULL;
903 const float *BO = NULL;
904 const float *B1 = NULL;
905 const float *B2 = NULL;
906
907 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
908 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
909
910 const int H = MODEL_NUM_HEADS;
911 const int H_kv = MODEL_NUM_KV_HEADS;
912 const int head_dim = MODEL_HEAD_DIM;
913 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
914 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
915 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
916
917 /* RMSNorm before attention */
918 rmsnorm_forward(input,
919 ln1_gamma,
920 ln1_out,
921 NULL,
922 num_tokens,
924 aligned_embed_dim,
925 1e-06f);
926
927 /* Q projection (head-major) */
928 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
929 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
930 for (int h = 0; h < H; ++h) {
931 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
932 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
933 float *q_h = q + (size_t)h * q_head_stride;
934 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
935 }
936
937 /* K projection (head-major) */
938 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
939 const uint8_t *WK_bytes = (const uint8_t *)WK;
940 for (int h = 0; h < H_kv; ++h) {
941 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
942 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
943 float *k_h = k + (size_t)h * kv_head_stride;
944 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
945 }
946
947 /* V projection (head-major) */
948 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
949 const uint8_t *WV_bytes = (const uint8_t *)WV;
950 for (int h = 0; h < H_kv; ++h) {
951 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
952 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
953 float *v_h = v + (size_t)h * kv_head_stride;
954 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
955 }
956
957 /* RoPE */
959 k,
960 rope_cos,
961 rope_sin,
962 H,
963 H_kv,
964 num_tokens,
965 head_dim,
966 aligned_head_dim,
967 0,
968 num_tokens,
969 aligned_context_window);
970
971 /* Attention (prefill, causal) */
973 k,
974 v,
975 attn_out,
976 H,
977 H_kv,
978 num_tokens,
979 head_dim,
980 aligned_head_dim,
981 aligned_context_window);
982
983 /* Output projection (flatten head-major to token-major) */
984 const int K = H * aligned_head_dim;
985 if (K != aligned_embed_dim) {
986 return;
987 }
988 const float *proj_in = attn_out;
989 if (H > 1) {
990 if (!proj_scratch) {
991 return;
992 }
993 for (int t = 0; t < num_tokens; ++t) {
994 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
995 for (int h = 0; h < H; ++h) {
996 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
997 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
998 src,
999 (size_t)aligned_head_dim * sizeof(float));
1000 }
1001 }
1002 proj_in = proj_scratch;
1003 }
1004 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1005
1006 /* Residual add */
1007 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1008
1009 /* RMSNorm before MLP */
1010 rmsnorm_forward(residual1,
1011 ln2_gamma,
1012 ln2_out,
1013 NULL,
1014 num_tokens,
1016 aligned_embed_dim,
1017 1e-06f);
1018
1019 /* MLP (SwiGLU) */
1020 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
1021 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
1022 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
1023
1024 /* Final residual add */
1025 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
1026}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_5_decode()

static void model_layer_5_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 5250 of file v6.6/test_generated/generic_api_test.c.

5257 {
5258 const MODELLayerOffsets *L = &MODEL_LAYERS[5];
5259
5260 float *input = MODEL_PTR(model, MODEL_LAYERS[4].output);
5261
5262 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
5263 float *ln1_out = MODEL_PTR(model, L->ln1_out);
5264 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
5265 float *ln2_out = MODEL_PTR(model, L->ln2_out);
5266 float *k_cache = MODEL_PTR(model, L->k);
5267 float *v_cache = MODEL_PTR(model, L->v);
5268 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
5269 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
5270 float *residual1 = MODEL_PTR(model, L->residual1);
5271 float *mlp_out = MODEL_PTR(model, L->mlp_out);
5272 float *output = MODEL_PTR(model, L->output);
5273
5274 /* Weights (explicit types for layer 5) */
5275 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
5276 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
5277 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
5278 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
5279 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
5280 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
5281
5282 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
5283 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
5284
5285 const int H = MODEL_NUM_HEADS;
5286 const int H_kv = MODEL_NUM_KV_HEADS;
5287 const int head_dim = MODEL_HEAD_DIM;
5288
5289 float q_token[H * aligned_head_dim];
5290 float k_token[H_kv * aligned_head_dim];
5291 float v_token[H_kv * aligned_head_dim];
5292 float attn_token[H * aligned_head_dim];
5293
5294 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
5295 float fc1_out[2 * aligned_intermediate_dim];
5296 float swiglu_out[aligned_intermediate_dim];
5297
5298 /* Step 1: RMSNorm before attention */
5299 rmsnorm_forward(input,
5300 ln1_gamma,
5301 ln1_out,
5302 NULL,
5303 1,
5305 aligned_embed_dim,
5306 1e-06f);
5307
5308 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
5309
5310 /* Step 2: QKV projection */
5311 /* Q projection: Q4_K -> gemm_nt_q4_k */
5312 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
5313 if (aligned_head_dim > head_dim) {
5314 for (int h = 0; h < H; ++h) {
5315 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
5316 for (int d = head_dim; d < aligned_head_dim; ++d) {
5317 q_head[d] = 0.0f;
5318 }
5319 }
5320 }
5321
5322 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5323 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5324 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
5325 const uint8_t *WK_bytes = (const uint8_t *)WK;
5326 for (int h = 0; h < H_kv; ++h) {
5327 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
5328 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
5329 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5330 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
5331 for (int d = head_dim; d < aligned_head_dim; ++d) {
5332 k_head[d] = 0.0f;
5333 }
5334 }
5335
5336 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5337 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5338 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
5339 const uint8_t *WV_bytes = (const uint8_t *)WV;
5340 for (int h = 0; h < H_kv; ++h) {
5341 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
5342 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
5343 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5344 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
5345 for (int d = head_dim; d < aligned_head_dim; ++d) {
5346 v_head[d] = 0.0f;
5347 }
5348 }
5349
5350 /* Step 3: RoPE */
5351 rope_forward(q_token,
5352 rope_cos,
5353 rope_sin,
5354 H,
5355 1,
5356 head_dim,
5357 aligned_head_dim,
5358 token_index);
5359 for (int h = 0; h < H_kv; ++h) {
5360 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5361 rope_forward(k_head,
5362 rope_cos,
5363 rope_sin,
5364 1,
5365 1,
5366 head_dim,
5367 aligned_head_dim,
5368 token_index);
5369 }
5370
5371 /* Step 4: KV cache write (direct-to-cache) */
5372
5373 /* Step 5: Attention (decode, flash) */
5375 k_cache,
5376 v_cache,
5377 attn_token,
5378 H,
5379 H_kv,
5380 token_index + 1,
5381 aligned_context_window,
5382 head_dim,
5383 aligned_head_dim);
5384
5385 /* Step 6: Output projection */
5386 /* WO projection: Q4_K -> gemm_nt_q4_k */
5387 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
5388
5389 /* Step 7: Residual add */
5390 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
5391
5392 /* Step 8: RMSNorm before MLP */
5393 rmsnorm_forward(residual1,
5394 ln2_gamma,
5395 ln2_out,
5396 NULL,
5397 1,
5399 aligned_embed_dim,
5400 1e-06f);
5401
5402 /* Step 9: MLP (SwiGLU) */
5403 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
5404 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
5405
5406 /* SwiGLU activation */
5407 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
5408
5409 /* Down projection: Q4_K -> gemm_nt_q4_k */
5410 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
5411
5412 /* Step 10: Final residual add */
5413 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
5414}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_5_prefill()

static void model_layer_5_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 1031 of file v6.6/test_generated/generic_api_test.c.

1038 {
1039 const MODELLayerOffsets *L = &MODEL_LAYERS[5];
1040
1041 float *input = MODEL_PTR(model, MODEL_LAYERS[4].output);
1042 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
1043 float *ln1_out = MODEL_PTR(model, L->ln1_out);
1044 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
1045 float *ln2_out = MODEL_PTR(model, L->ln2_out);
1046 float *q = MODEL_PTR(model, L->q);
1047 float *k = MODEL_PTR(model, L->k);
1048 float *v = MODEL_PTR(model, L->v);
1049 float *attn_out = MODEL_PTR(model, L->attn_out);
1050 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
1051 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
1052 float *residual1 = MODEL_PTR(model, L->residual1);
1053 float *fc1_out = MODEL_PTR(model, L->fc1_out);
1054 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
1055 float *mlp_out = MODEL_PTR(model, L->mlp_out);
1056 float *output = MODEL_PTR(model, L->output);
1057
1058 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
1059 const void *WK = (const void *)MODEL_PTR(model, L->wk);
1060 const void *WV = (const void *)MODEL_PTR(model, L->wv);
1061 const void *WO = (const void *)MODEL_PTR(model, L->wo);
1062 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
1063 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
1064 const float *BQ = NULL;
1065 const float *BK = NULL;
1066 const float *BV = NULL;
1067 const float *BO = NULL;
1068 const float *B1 = NULL;
1069 const float *B2 = NULL;
1070
1071 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
1072 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
1073
1074 const int H = MODEL_NUM_HEADS;
1075 const int H_kv = MODEL_NUM_KV_HEADS;
1076 const int head_dim = MODEL_HEAD_DIM;
1077 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
1078 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
1079 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
1080
1081 /* RMSNorm before attention */
1082 rmsnorm_forward(input,
1083 ln1_gamma,
1084 ln1_out,
1085 NULL,
1086 num_tokens,
1088 aligned_embed_dim,
1089 1e-06f);
1090
1091 /* Q projection (head-major) */
1092 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1093 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
1094 for (int h = 0; h < H; ++h) {
1095 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
1096 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
1097 float *q_h = q + (size_t)h * q_head_stride;
1098 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1099 }
1100
1101 /* K projection (head-major) */
1102 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1103 const uint8_t *WK_bytes = (const uint8_t *)WK;
1104 for (int h = 0; h < H_kv; ++h) {
1105 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
1106 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
1107 float *k_h = k + (size_t)h * kv_head_stride;
1108 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1109 }
1110
1111 /* V projection (head-major) */
1112 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1113 const uint8_t *WV_bytes = (const uint8_t *)WV;
1114 for (int h = 0; h < H_kv; ++h) {
1115 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
1116 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
1117 float *v_h = v + (size_t)h * kv_head_stride;
1118 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1119 }
1120
1121 /* RoPE */
1123 k,
1124 rope_cos,
1125 rope_sin,
1126 H,
1127 H_kv,
1128 num_tokens,
1129 head_dim,
1130 aligned_head_dim,
1131 0,
1132 num_tokens,
1133 aligned_context_window);
1134
1135 /* Attention (prefill, causal) */
1137 k,
1138 v,
1139 attn_out,
1140 H,
1141 H_kv,
1142 num_tokens,
1143 head_dim,
1144 aligned_head_dim,
1145 aligned_context_window);
1146
1147 /* Output projection (flatten head-major to token-major) */
1148 const int K = H * aligned_head_dim;
1149 if (K != aligned_embed_dim) {
1150 return;
1151 }
1152 const float *proj_in = attn_out;
1153 if (H > 1) {
1154 if (!proj_scratch) {
1155 return;
1156 }
1157 for (int t = 0; t < num_tokens; ++t) {
1158 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
1159 for (int h = 0; h < H; ++h) {
1160 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
1161 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
1162 src,
1163 (size_t)aligned_head_dim * sizeof(float));
1164 }
1165 }
1166 proj_in = proj_scratch;
1167 }
1168 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1169
1170 /* Residual add */
1171 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1172
1173 /* RMSNorm before MLP */
1174 rmsnorm_forward(residual1,
1175 ln2_gamma,
1176 ln2_out,
1177 NULL,
1178 num_tokens,
1180 aligned_embed_dim,
1181 1e-06f);
1182
1183 /* MLP (SwiGLU) */
1184 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
1185 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
1186 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
1187
1188 /* Final residual add */
1189 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
1190}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_6_decode()

static void model_layer_6_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 5419 of file v6.6/test_generated/generic_api_test.c.

5426 {
5427 const MODELLayerOffsets *L = &MODEL_LAYERS[6];
5428
5429 float *input = MODEL_PTR(model, MODEL_LAYERS[5].output);
5430
5431 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
5432 float *ln1_out = MODEL_PTR(model, L->ln1_out);
5433 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
5434 float *ln2_out = MODEL_PTR(model, L->ln2_out);
5435 float *k_cache = MODEL_PTR(model, L->k);
5436 float *v_cache = MODEL_PTR(model, L->v);
5437 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
5438 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
5439 float *residual1 = MODEL_PTR(model, L->residual1);
5440 float *mlp_out = MODEL_PTR(model, L->mlp_out);
5441 float *output = MODEL_PTR(model, L->output);
5442
5443 /* Weights (explicit types for layer 6) */
5444 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
5445 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
5446 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
5447 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
5448 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
5449 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
5450
5451 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
5452 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
5453
5454 const int H = MODEL_NUM_HEADS;
5455 const int H_kv = MODEL_NUM_KV_HEADS;
5456 const int head_dim = MODEL_HEAD_DIM;
5457
5458 float q_token[H * aligned_head_dim];
5459 float k_token[H_kv * aligned_head_dim];
5460 float v_token[H_kv * aligned_head_dim];
5461 float attn_token[H * aligned_head_dim];
5462
5463 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
5464 float fc1_out[2 * aligned_intermediate_dim];
5465 float swiglu_out[aligned_intermediate_dim];
5466
5467 /* Step 1: RMSNorm before attention */
5468 rmsnorm_forward(input,
5469 ln1_gamma,
5470 ln1_out,
5471 NULL,
5472 1,
5474 aligned_embed_dim,
5475 1e-06f);
5476
5477 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
5478
5479 /* Step 2: QKV projection */
5480 /* Q projection: Q4_K -> gemm_nt_q4_k */
5481 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
5482 if (aligned_head_dim > head_dim) {
5483 for (int h = 0; h < H; ++h) {
5484 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
5485 for (int d = head_dim; d < aligned_head_dim; ++d) {
5486 q_head[d] = 0.0f;
5487 }
5488 }
5489 }
5490
5491 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5492 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5493 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
5494 const uint8_t *WK_bytes = (const uint8_t *)WK;
5495 for (int h = 0; h < H_kv; ++h) {
5496 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
5497 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
5498 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5499 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
5500 for (int d = head_dim; d < aligned_head_dim; ++d) {
5501 k_head[d] = 0.0f;
5502 }
5503 }
5504
5505 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5506 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5507 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
5508 const uint8_t *WV_bytes = (const uint8_t *)WV;
5509 for (int h = 0; h < H_kv; ++h) {
5510 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
5511 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
5512 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5513 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
5514 for (int d = head_dim; d < aligned_head_dim; ++d) {
5515 v_head[d] = 0.0f;
5516 }
5517 }
5518
5519 /* Step 3: RoPE */
5520 rope_forward(q_token,
5521 rope_cos,
5522 rope_sin,
5523 H,
5524 1,
5525 head_dim,
5526 aligned_head_dim,
5527 token_index);
5528 for (int h = 0; h < H_kv; ++h) {
5529 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5530 rope_forward(k_head,
5531 rope_cos,
5532 rope_sin,
5533 1,
5534 1,
5535 head_dim,
5536 aligned_head_dim,
5537 token_index);
5538 }
5539
5540 /* Step 4: KV cache write (direct-to-cache) */
5541
5542 /* Step 5: Attention (decode, flash) */
5544 k_cache,
5545 v_cache,
5546 attn_token,
5547 H,
5548 H_kv,
5549 token_index + 1,
5550 aligned_context_window,
5551 head_dim,
5552 aligned_head_dim);
5553
5554 /* Step 6: Output projection */
5555 /* WO projection: Q4_K -> gemm_nt_q4_k */
5556 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
5557
5558 /* Step 7: Residual add */
5559 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
5560
5561 /* Step 8: RMSNorm before MLP */
5562 rmsnorm_forward(residual1,
5563 ln2_gamma,
5564 ln2_out,
5565 NULL,
5566 1,
5568 aligned_embed_dim,
5569 1e-06f);
5570
5571 /* Step 9: MLP (SwiGLU) */
5572 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
5573 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
5574
5575 /* SwiGLU activation */
5576 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
5577
5578 /* Down projection: Q4_K -> gemm_nt_q4_k */
5579 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
5580
5581 /* Step 10: Final residual add */
5582 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
5583}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_6_prefill()

static void model_layer_6_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 1195 of file v6.6/test_generated/generic_api_test.c.

1202 {
1203 const MODELLayerOffsets *L = &MODEL_LAYERS[6];
1204
1205 float *input = MODEL_PTR(model, MODEL_LAYERS[5].output);
1206 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
1207 float *ln1_out = MODEL_PTR(model, L->ln1_out);
1208 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
1209 float *ln2_out = MODEL_PTR(model, L->ln2_out);
1210 float *q = MODEL_PTR(model, L->q);
1211 float *k = MODEL_PTR(model, L->k);
1212 float *v = MODEL_PTR(model, L->v);
1213 float *attn_out = MODEL_PTR(model, L->attn_out);
1214 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
1215 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
1216 float *residual1 = MODEL_PTR(model, L->residual1);
1217 float *fc1_out = MODEL_PTR(model, L->fc1_out);
1218 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
1219 float *mlp_out = MODEL_PTR(model, L->mlp_out);
1220 float *output = MODEL_PTR(model, L->output);
1221
1222 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
1223 const void *WK = (const void *)MODEL_PTR(model, L->wk);
1224 const void *WV = (const void *)MODEL_PTR(model, L->wv);
1225 const void *WO = (const void *)MODEL_PTR(model, L->wo);
1226 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
1227 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
1228 const float *BQ = NULL;
1229 const float *BK = NULL;
1230 const float *BV = NULL;
1231 const float *BO = NULL;
1232 const float *B1 = NULL;
1233 const float *B2 = NULL;
1234
1235 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
1236 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
1237
1238 const int H = MODEL_NUM_HEADS;
1239 const int H_kv = MODEL_NUM_KV_HEADS;
1240 const int head_dim = MODEL_HEAD_DIM;
1241 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
1242 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
1243 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
1244
1245 /* RMSNorm before attention */
1246 rmsnorm_forward(input,
1247 ln1_gamma,
1248 ln1_out,
1249 NULL,
1250 num_tokens,
1252 aligned_embed_dim,
1253 1e-06f);
1254
1255 /* Q projection (head-major) */
1256 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1257 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
1258 for (int h = 0; h < H; ++h) {
1259 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
1260 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
1261 float *q_h = q + (size_t)h * q_head_stride;
1262 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1263 }
1264
1265 /* K projection (head-major) */
1266 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1267 const uint8_t *WK_bytes = (const uint8_t *)WK;
1268 for (int h = 0; h < H_kv; ++h) {
1269 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
1270 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
1271 float *k_h = k + (size_t)h * kv_head_stride;
1272 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1273 }
1274
1275 /* V projection (head-major) */
1276 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1277 const uint8_t *WV_bytes = (const uint8_t *)WV;
1278 for (int h = 0; h < H_kv; ++h) {
1279 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
1280 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
1281 float *v_h = v + (size_t)h * kv_head_stride;
1282 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1283 }
1284
1285 /* RoPE */
1287 k,
1288 rope_cos,
1289 rope_sin,
1290 H,
1291 H_kv,
1292 num_tokens,
1293 head_dim,
1294 aligned_head_dim,
1295 0,
1296 num_tokens,
1297 aligned_context_window);
1298
1299 /* Attention (prefill, causal) */
1301 k,
1302 v,
1303 attn_out,
1304 H,
1305 H_kv,
1306 num_tokens,
1307 head_dim,
1308 aligned_head_dim,
1309 aligned_context_window);
1310
1311 /* Output projection (flatten head-major to token-major) */
1312 const int K = H * aligned_head_dim;
1313 if (K != aligned_embed_dim) {
1314 return;
1315 }
1316 const float *proj_in = attn_out;
1317 if (H > 1) {
1318 if (!proj_scratch) {
1319 return;
1320 }
1321 for (int t = 0; t < num_tokens; ++t) {
1322 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
1323 for (int h = 0; h < H; ++h) {
1324 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
1325 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
1326 src,
1327 (size_t)aligned_head_dim * sizeof(float));
1328 }
1329 }
1330 proj_in = proj_scratch;
1331 }
1332 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1333
1334 /* Residual add */
1335 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1336
1337 /* RMSNorm before MLP */
1338 rmsnorm_forward(residual1,
1339 ln2_gamma,
1340 ln2_out,
1341 NULL,
1342 num_tokens,
1344 aligned_embed_dim,
1345 1e-06f);
1346
1347 /* MLP (SwiGLU) */
1348 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
1349 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
1350 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
1351
1352 /* Final residual add */
1353 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
1354}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_7_decode()

static void model_layer_7_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 5588 of file v6.6/test_generated/generic_api_test.c.

5595 {
5596 const MODELLayerOffsets *L = &MODEL_LAYERS[7];
5597
5598 float *input = MODEL_PTR(model, MODEL_LAYERS[6].output);
5599
5600 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
5601 float *ln1_out = MODEL_PTR(model, L->ln1_out);
5602 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
5603 float *ln2_out = MODEL_PTR(model, L->ln2_out);
5604 float *k_cache = MODEL_PTR(model, L->k);
5605 float *v_cache = MODEL_PTR(model, L->v);
5606 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
5607 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
5608 float *residual1 = MODEL_PTR(model, L->residual1);
5609 float *mlp_out = MODEL_PTR(model, L->mlp_out);
5610 float *output = MODEL_PTR(model, L->output);
5611
5612 /* Weights (explicit types for layer 7) */
5613 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
5614 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
5615 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
5616 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
5617 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
5618 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
5619
5620 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
5621 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
5622
5623 const int H = MODEL_NUM_HEADS;
5624 const int H_kv = MODEL_NUM_KV_HEADS;
5625 const int head_dim = MODEL_HEAD_DIM;
5626
5627 float q_token[H * aligned_head_dim];
5628 float k_token[H_kv * aligned_head_dim];
5629 float v_token[H_kv * aligned_head_dim];
5630 float attn_token[H * aligned_head_dim];
5631
5632 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
5633 float fc1_out[2 * aligned_intermediate_dim];
5634 float swiglu_out[aligned_intermediate_dim];
5635
5636 /* Step 1: RMSNorm before attention */
5637 rmsnorm_forward(input,
5638 ln1_gamma,
5639 ln1_out,
5640 NULL,
5641 1,
5643 aligned_embed_dim,
5644 1e-06f);
5645
5646 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
5647
5648 /* Step 2: QKV projection */
5649 /* Q projection: Q4_K -> gemm_nt_q4_k */
5650 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
5651 if (aligned_head_dim > head_dim) {
5652 for (int h = 0; h < H; ++h) {
5653 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
5654 for (int d = head_dim; d < aligned_head_dim; ++d) {
5655 q_head[d] = 0.0f;
5656 }
5657 }
5658 }
5659
5660 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5661 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5662 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
5663 const uint8_t *WK_bytes = (const uint8_t *)WK;
5664 for (int h = 0; h < H_kv; ++h) {
5665 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
5666 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
5667 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5668 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
5669 for (int d = head_dim; d < aligned_head_dim; ++d) {
5670 k_head[d] = 0.0f;
5671 }
5672 }
5673
5674 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5675 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5676 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
5677 const uint8_t *WV_bytes = (const uint8_t *)WV;
5678 for (int h = 0; h < H_kv; ++h) {
5679 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
5680 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
5681 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5682 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
5683 for (int d = head_dim; d < aligned_head_dim; ++d) {
5684 v_head[d] = 0.0f;
5685 }
5686 }
5687
5688 /* Step 3: RoPE */
5689 rope_forward(q_token,
5690 rope_cos,
5691 rope_sin,
5692 H,
5693 1,
5694 head_dim,
5695 aligned_head_dim,
5696 token_index);
5697 for (int h = 0; h < H_kv; ++h) {
5698 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5699 rope_forward(k_head,
5700 rope_cos,
5701 rope_sin,
5702 1,
5703 1,
5704 head_dim,
5705 aligned_head_dim,
5706 token_index);
5707 }
5708
5709 /* Step 4: KV cache write (direct-to-cache) */
5710
5711 /* Step 5: Attention (decode, flash) */
5713 k_cache,
5714 v_cache,
5715 attn_token,
5716 H,
5717 H_kv,
5718 token_index + 1,
5719 aligned_context_window,
5720 head_dim,
5721 aligned_head_dim);
5722
5723 /* Step 6: Output projection */
5724 /* WO projection: Q4_K -> gemm_nt_q4_k */
5725 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
5726
5727 /* Step 7: Residual add */
5728 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
5729
5730 /* Step 8: RMSNorm before MLP */
5731 rmsnorm_forward(residual1,
5732 ln2_gamma,
5733 ln2_out,
5734 NULL,
5735 1,
5737 aligned_embed_dim,
5738 1e-06f);
5739
5740 /* Step 9: MLP (SwiGLU) */
5741 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
5742 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
5743
5744 /* SwiGLU activation */
5745 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
5746
5747 /* Down projection: Q4_K -> gemm_nt_q4_k */
5748 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
5749
5750 /* Step 10: Final residual add */
5751 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
5752}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_7_prefill()

static void model_layer_7_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 1359 of file v6.6/test_generated/generic_api_test.c.

1366 {
1367 const MODELLayerOffsets *L = &MODEL_LAYERS[7];
1368
1369 float *input = MODEL_PTR(model, MODEL_LAYERS[6].output);
1370 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
1371 float *ln1_out = MODEL_PTR(model, L->ln1_out);
1372 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
1373 float *ln2_out = MODEL_PTR(model, L->ln2_out);
1374 float *q = MODEL_PTR(model, L->q);
1375 float *k = MODEL_PTR(model, L->k);
1376 float *v = MODEL_PTR(model, L->v);
1377 float *attn_out = MODEL_PTR(model, L->attn_out);
1378 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
1379 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
1380 float *residual1 = MODEL_PTR(model, L->residual1);
1381 float *fc1_out = MODEL_PTR(model, L->fc1_out);
1382 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
1383 float *mlp_out = MODEL_PTR(model, L->mlp_out);
1384 float *output = MODEL_PTR(model, L->output);
1385
1386 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
1387 const void *WK = (const void *)MODEL_PTR(model, L->wk);
1388 const void *WV = (const void *)MODEL_PTR(model, L->wv);
1389 const void *WO = (const void *)MODEL_PTR(model, L->wo);
1390 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
1391 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
1392 const float *BQ = NULL;
1393 const float *BK = NULL;
1394 const float *BV = NULL;
1395 const float *BO = NULL;
1396 const float *B1 = NULL;
1397 const float *B2 = NULL;
1398
1399 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
1400 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
1401
1402 const int H = MODEL_NUM_HEADS;
1403 const int H_kv = MODEL_NUM_KV_HEADS;
1404 const int head_dim = MODEL_HEAD_DIM;
1405 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
1406 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
1407 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
1408
1409 /* RMSNorm before attention */
1410 rmsnorm_forward(input,
1411 ln1_gamma,
1412 ln1_out,
1413 NULL,
1414 num_tokens,
1416 aligned_embed_dim,
1417 1e-06f);
1418
1419 /* Q projection (head-major) */
1420 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1421 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
1422 for (int h = 0; h < H; ++h) {
1423 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
1424 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
1425 float *q_h = q + (size_t)h * q_head_stride;
1426 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1427 }
1428
1429 /* K projection (head-major) */
1430 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1431 const uint8_t *WK_bytes = (const uint8_t *)WK;
1432 for (int h = 0; h < H_kv; ++h) {
1433 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
1434 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
1435 float *k_h = k + (size_t)h * kv_head_stride;
1436 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1437 }
1438
1439 /* V projection (head-major) */
1440 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1441 const uint8_t *WV_bytes = (const uint8_t *)WV;
1442 for (int h = 0; h < H_kv; ++h) {
1443 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
1444 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
1445 float *v_h = v + (size_t)h * kv_head_stride;
1446 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1447 }
1448
1449 /* RoPE */
1451 k,
1452 rope_cos,
1453 rope_sin,
1454 H,
1455 H_kv,
1456 num_tokens,
1457 head_dim,
1458 aligned_head_dim,
1459 0,
1460 num_tokens,
1461 aligned_context_window);
1462
1463 /* Attention (prefill, causal) */
1465 k,
1466 v,
1467 attn_out,
1468 H,
1469 H_kv,
1470 num_tokens,
1471 head_dim,
1472 aligned_head_dim,
1473 aligned_context_window);
1474
1475 /* Output projection (flatten head-major to token-major) */
1476 const int K = H * aligned_head_dim;
1477 if (K != aligned_embed_dim) {
1478 return;
1479 }
1480 const float *proj_in = attn_out;
1481 if (H > 1) {
1482 if (!proj_scratch) {
1483 return;
1484 }
1485 for (int t = 0; t < num_tokens; ++t) {
1486 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
1487 for (int h = 0; h < H; ++h) {
1488 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
1489 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
1490 src,
1491 (size_t)aligned_head_dim * sizeof(float));
1492 }
1493 }
1494 proj_in = proj_scratch;
1495 }
1496 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1497
1498 /* Residual add */
1499 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1500
1501 /* RMSNorm before MLP */
1502 rmsnorm_forward(residual1,
1503 ln2_gamma,
1504 ln2_out,
1505 NULL,
1506 num_tokens,
1508 aligned_embed_dim,
1509 1e-06f);
1510
1511 /* MLP (SwiGLU) */
1512 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
1513 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
1514 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
1515
1516 /* Final residual add */
1517 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
1518}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_8_decode()

static void model_layer_8_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 5757 of file v6.6/test_generated/generic_api_test.c.

5764 {
5765 const MODELLayerOffsets *L = &MODEL_LAYERS[8];
5766
5767 float *input = MODEL_PTR(model, MODEL_LAYERS[7].output);
5768
5769 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
5770 float *ln1_out = MODEL_PTR(model, L->ln1_out);
5771 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
5772 float *ln2_out = MODEL_PTR(model, L->ln2_out);
5773 float *k_cache = MODEL_PTR(model, L->k);
5774 float *v_cache = MODEL_PTR(model, L->v);
5775 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
5776 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
5777 float *residual1 = MODEL_PTR(model, L->residual1);
5778 float *mlp_out = MODEL_PTR(model, L->mlp_out);
5779 float *output = MODEL_PTR(model, L->output);
5780
5781 /* Weights (explicit types for layer 8) */
5782 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
5783 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
5784 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
5785 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
5786 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
5787 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
5788
5789 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
5790 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
5791
5792 const int H = MODEL_NUM_HEADS;
5793 const int H_kv = MODEL_NUM_KV_HEADS;
5794 const int head_dim = MODEL_HEAD_DIM;
5795
5796 float q_token[H * aligned_head_dim];
5797 float k_token[H_kv * aligned_head_dim];
5798 float v_token[H_kv * aligned_head_dim];
5799 float attn_token[H * aligned_head_dim];
5800
5801 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
5802 float fc1_out[2 * aligned_intermediate_dim];
5803 float swiglu_out[aligned_intermediate_dim];
5804
5805 /* Step 1: RMSNorm before attention */
5806 rmsnorm_forward(input,
5807 ln1_gamma,
5808 ln1_out,
5809 NULL,
5810 1,
5812 aligned_embed_dim,
5813 1e-06f);
5814
5815 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
5816
5817 /* Step 2: QKV projection */
5818 /* Q projection: Q4_K -> gemm_nt_q4_k */
5819 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
5820 if (aligned_head_dim > head_dim) {
5821 for (int h = 0; h < H; ++h) {
5822 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
5823 for (int d = head_dim; d < aligned_head_dim; ++d) {
5824 q_head[d] = 0.0f;
5825 }
5826 }
5827 }
5828
5829 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5830 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5831 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
5832 const uint8_t *WK_bytes = (const uint8_t *)WK;
5833 for (int h = 0; h < H_kv; ++h) {
5834 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
5835 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
5836 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5837 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
5838 for (int d = head_dim; d < aligned_head_dim; ++d) {
5839 k_head[d] = 0.0f;
5840 }
5841 }
5842
5843 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5844 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
5845 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
5846 const uint8_t *WV_bytes = (const uint8_t *)WV;
5847 for (int h = 0; h < H_kv; ++h) {
5848 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
5849 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
5850 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5851 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
5852 for (int d = head_dim; d < aligned_head_dim; ++d) {
5853 v_head[d] = 0.0f;
5854 }
5855 }
5856
5857 /* Step 3: RoPE */
5858 rope_forward(q_token,
5859 rope_cos,
5860 rope_sin,
5861 H,
5862 1,
5863 head_dim,
5864 aligned_head_dim,
5865 token_index);
5866 for (int h = 0; h < H_kv; ++h) {
5867 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
5868 rope_forward(k_head,
5869 rope_cos,
5870 rope_sin,
5871 1,
5872 1,
5873 head_dim,
5874 aligned_head_dim,
5875 token_index);
5876 }
5877
5878 /* Step 4: KV cache write (direct-to-cache) */
5879
5880 /* Step 5: Attention (decode, flash) */
5882 k_cache,
5883 v_cache,
5884 attn_token,
5885 H,
5886 H_kv,
5887 token_index + 1,
5888 aligned_context_window,
5889 head_dim,
5890 aligned_head_dim);
5891
5892 /* Step 6: Output projection */
5893 /* WO projection: Q4_K -> gemm_nt_q4_k */
5894 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
5895
5896 /* Step 7: Residual add */
5897 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
5898
5899 /* Step 8: RMSNorm before MLP */
5900 rmsnorm_forward(residual1,
5901 ln2_gamma,
5902 ln2_out,
5903 NULL,
5904 1,
5906 aligned_embed_dim,
5907 1e-06f);
5908
5909 /* Step 9: MLP (SwiGLU) */
5910 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
5911 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
5912
5913 /* SwiGLU activation */
5914 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
5915
5916 /* Down projection: Q4_K -> gemm_nt_q4_k */
5917 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
5918
5919 /* Step 10: Final residual add */
5920 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
5921}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_8_prefill()

static void model_layer_8_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 1523 of file v6.6/test_generated/generic_api_test.c.

1530 {
1531 const MODELLayerOffsets *L = &MODEL_LAYERS[8];
1532
1533 float *input = MODEL_PTR(model, MODEL_LAYERS[7].output);
1534 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
1535 float *ln1_out = MODEL_PTR(model, L->ln1_out);
1536 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
1537 float *ln2_out = MODEL_PTR(model, L->ln2_out);
1538 float *q = MODEL_PTR(model, L->q);
1539 float *k = MODEL_PTR(model, L->k);
1540 float *v = MODEL_PTR(model, L->v);
1541 float *attn_out = MODEL_PTR(model, L->attn_out);
1542 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
1543 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
1544 float *residual1 = MODEL_PTR(model, L->residual1);
1545 float *fc1_out = MODEL_PTR(model, L->fc1_out);
1546 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
1547 float *mlp_out = MODEL_PTR(model, L->mlp_out);
1548 float *output = MODEL_PTR(model, L->output);
1549
1550 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
1551 const void *WK = (const void *)MODEL_PTR(model, L->wk);
1552 const void *WV = (const void *)MODEL_PTR(model, L->wv);
1553 const void *WO = (const void *)MODEL_PTR(model, L->wo);
1554 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
1555 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
1556 const float *BQ = NULL;
1557 const float *BK = NULL;
1558 const float *BV = NULL;
1559 const float *BO = NULL;
1560 const float *B1 = NULL;
1561 const float *B2 = NULL;
1562
1563 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
1564 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
1565
1566 const int H = MODEL_NUM_HEADS;
1567 const int H_kv = MODEL_NUM_KV_HEADS;
1568 const int head_dim = MODEL_HEAD_DIM;
1569 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
1570 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
1571 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
1572
1573 /* RMSNorm before attention */
1574 rmsnorm_forward(input,
1575 ln1_gamma,
1576 ln1_out,
1577 NULL,
1578 num_tokens,
1580 aligned_embed_dim,
1581 1e-06f);
1582
1583 /* Q projection (head-major) */
1584 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1585 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
1586 for (int h = 0; h < H; ++h) {
1587 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
1588 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
1589 float *q_h = q + (size_t)h * q_head_stride;
1590 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1591 }
1592
1593 /* K projection (head-major) */
1594 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1595 const uint8_t *WK_bytes = (const uint8_t *)WK;
1596 for (int h = 0; h < H_kv; ++h) {
1597 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
1598 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
1599 float *k_h = k + (size_t)h * kv_head_stride;
1600 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1601 }
1602
1603 /* V projection (head-major) */
1604 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1605 const uint8_t *WV_bytes = (const uint8_t *)WV;
1606 for (int h = 0; h < H_kv; ++h) {
1607 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
1608 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
1609 float *v_h = v + (size_t)h * kv_head_stride;
1610 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1611 }
1612
1613 /* RoPE */
1615 k,
1616 rope_cos,
1617 rope_sin,
1618 H,
1619 H_kv,
1620 num_tokens,
1621 head_dim,
1622 aligned_head_dim,
1623 0,
1624 num_tokens,
1625 aligned_context_window);
1626
1627 /* Attention (prefill, causal) */
1629 k,
1630 v,
1631 attn_out,
1632 H,
1633 H_kv,
1634 num_tokens,
1635 head_dim,
1636 aligned_head_dim,
1637 aligned_context_window);
1638
1639 /* Output projection (flatten head-major to token-major) */
1640 const int K = H * aligned_head_dim;
1641 if (K != aligned_embed_dim) {
1642 return;
1643 }
1644 const float *proj_in = attn_out;
1645 if (H > 1) {
1646 if (!proj_scratch) {
1647 return;
1648 }
1649 for (int t = 0; t < num_tokens; ++t) {
1650 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
1651 for (int h = 0; h < H; ++h) {
1652 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
1653 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
1654 src,
1655 (size_t)aligned_head_dim * sizeof(float));
1656 }
1657 }
1658 proj_in = proj_scratch;
1659 }
1660 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1661
1662 /* Residual add */
1663 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1664
1665 /* RMSNorm before MLP */
1666 rmsnorm_forward(residual1,
1667 ln2_gamma,
1668 ln2_out,
1669 NULL,
1670 num_tokens,
1672 aligned_embed_dim,
1673 1e-06f);
1674
1675 /* MLP (SwiGLU) */
1676 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
1677 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
1678 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
1679
1680 /* Final residual add */
1681 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
1682}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_layer_9_decode()

static void model_layer_9_decode ( MODELModel *  model,
int  token_index,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 5926 of file v6.6/test_generated/generic_api_test.c.

5933 {
5934 const MODELLayerOffsets *L = &MODEL_LAYERS[9];
5935
5936 float *input = MODEL_PTR(model, MODEL_LAYERS[8].output);
5937
5938 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
5939 float *ln1_out = MODEL_PTR(model, L->ln1_out);
5940 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
5941 float *ln2_out = MODEL_PTR(model, L->ln2_out);
5942 float *k_cache = MODEL_PTR(model, L->k);
5943 float *v_cache = MODEL_PTR(model, L->v);
5944 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
5945 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
5946 float *residual1 = MODEL_PTR(model, L->residual1);
5947 float *mlp_out = MODEL_PTR(model, L->mlp_out);
5948 float *output = MODEL_PTR(model, L->output);
5949
5950 /* Weights (explicit types for layer 9) */
5951 const void *WQ = (const void *)MODEL_PTR(model, L->wq); /* Q4_K */
5952 const void *WK = (const void *)MODEL_PTR(model, L->wk); /* Q4_K */
5953 const void *WV = (const void *)MODEL_PTR(model, L->wv); /* Q4_K */
5954 const void *WO = (const void *)MODEL_PTR(model, L->wo); /* Q4_K */
5955 const void *W1 = (const void *)MODEL_PTR(model, L->w1); /* Q4_K (gate+up) */
5956 const void *W2 = (const void *)MODEL_PTR(model, L->w2); /* Q4_K (down) */
5957
5958 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
5959 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
5960
5961 const int H = MODEL_NUM_HEADS;
5962 const int H_kv = MODEL_NUM_KV_HEADS;
5963 const int head_dim = MODEL_HEAD_DIM;
5964
5965 float q_token[H * aligned_head_dim];
5966 float k_token[H_kv * aligned_head_dim];
5967 float v_token[H_kv * aligned_head_dim];
5968 float attn_token[H * aligned_head_dim];
5969
5970 /* Local MLP buffers (avoid layout dependencies for intermediate values) */
5971 float fc1_out[2 * aligned_intermediate_dim];
5972 float swiglu_out[aligned_intermediate_dim];
5973
5974 /* Step 1: RMSNorm before attention */
5975 rmsnorm_forward(input,
5976 ln1_gamma,
5977 ln1_out,
5978 NULL,
5979 1,
5981 aligned_embed_dim,
5982 1e-06f);
5983
5984 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
5985
5986 /* Step 2: QKV projection */
5987 /* Q projection: Q4_K -> gemm_nt_q4_k */
5988 gemm_nt_q4_k(ln1_out, WQ, NULL, q_token, 1, H * head_dim, aligned_embed_dim);
5989 if (aligned_head_dim > head_dim) {
5990 for (int h = 0; h < H; ++h) {
5991 float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
5992 for (int d = head_dim; d < aligned_head_dim; ++d) {
5993 q_head[d] = 0.0f;
5994 }
5995 }
5996 }
5997
5998 /* K projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
5999 const size_t wk_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6000 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wk_head_elems);
6001 const uint8_t *WK_bytes = (const uint8_t *)WK;
6002 for (int h = 0; h < H_kv; ++h) {
6003 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
6004 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
6005 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6006 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_head, 1, head_dim, aligned_embed_dim);
6007 for (int d = head_dim; d < aligned_head_dim; ++d) {
6008 k_head[d] = 0.0f;
6009 }
6010 }
6011
6012 /* V projection: Q4_K -> gemm_nt_q4_k (direct-to-cache) */
6013 const size_t wv_head_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
6014 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, wv_head_elems);
6015 const uint8_t *WV_bytes = (const uint8_t *)WV;
6016 for (int h = 0; h < H_kv; ++h) {
6017 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
6018 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
6019 float *v_head = v_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6020 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_head, 1, head_dim, aligned_embed_dim);
6021 for (int d = head_dim; d < aligned_head_dim; ++d) {
6022 v_head[d] = 0.0f;
6023 }
6024 }
6025
6026 /* Step 3: RoPE */
6027 rope_forward(q_token,
6028 rope_cos,
6029 rope_sin,
6030 H,
6031 1,
6032 head_dim,
6033 aligned_head_dim,
6034 token_index);
6035 for (int h = 0; h < H_kv; ++h) {
6036 float *k_head = k_cache + (size_t)h * kv_head_stride + (size_t)token_index * (size_t)aligned_head_dim;
6037 rope_forward(k_head,
6038 rope_cos,
6039 rope_sin,
6040 1,
6041 1,
6042 head_dim,
6043 aligned_head_dim,
6044 token_index);
6045 }
6046
6047 /* Step 4: KV cache write (direct-to-cache) */
6048
6049 /* Step 5: Attention (decode, flash) */
6051 k_cache,
6052 v_cache,
6053 attn_token,
6054 H,
6055 H_kv,
6056 token_index + 1,
6057 aligned_context_window,
6058 head_dim,
6059 aligned_head_dim);
6060
6061 /* Step 6: Output projection */
6062 /* WO projection: Q4_K -> gemm_nt_q4_k */
6063 gemm_nt_q4_k(attn_token, WO, NULL, proj_tmp, 1, aligned_embed_dim, H * head_dim);
6064
6065 /* Step 7: Residual add */
6066 model_residual_add_token_major(input, proj_tmp, residual1, 1, aligned_embed_dim);
6067
6068 /* Step 8: RMSNorm before MLP */
6069 rmsnorm_forward(residual1,
6070 ln2_gamma,
6071 ln2_out,
6072 NULL,
6073 1,
6075 aligned_embed_dim,
6076 1e-06f);
6077
6078 /* Step 9: MLP (SwiGLU) */
6079 /* Gate+Up projection: Q4_K -> gemm_nt_q4_k */
6080 gemm_nt_q4_k(ln2_out, W1, NULL, fc1_out, 1, 2 * aligned_intermediate_dim, aligned_embed_dim);
6081
6082 /* SwiGLU activation */
6083 swiglu_forward(fc1_out, swiglu_out, 1, aligned_intermediate_dim);
6084
6085 /* Down projection: Q4_K -> gemm_nt_q4_k */
6086 gemm_nt_q4_k(swiglu_out, W2, NULL, mlp_out, 1, aligned_embed_dim, aligned_intermediate_dim);
6087
6088 /* Step 10: Final residual add */
6089 model_residual_add_token_major(residual1, mlp_out, output, 1, aligned_embed_dim);
6090}

References attention_forward_decode_head_major_gqa_flash(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward(), and swiglu_forward().

Referenced by model_decode_token().

◆ model_layer_9_prefill()

static void model_layer_9_prefill ( MODELModel *  model,
int  num_tokens,
int  aligned_embed_dim,
int  aligned_head_dim,
int  aligned_intermediate_dim,
int  aligned_context_window 
)
static

Definition at line 1687 of file v6.6/test_generated/generic_api_test.c.

1694 {
1695 const MODELLayerOffsets *L = &MODEL_LAYERS[9];
1696
1697 float *input = MODEL_PTR(model, MODEL_LAYERS[8].output);
1698 float *ln1_gamma = MODEL_PTR(model, L->ln1_gamma);
1699 float *ln1_out = MODEL_PTR(model, L->ln1_out);
1700 float *ln2_gamma = MODEL_PTR(model, L->ln2_gamma);
1701 float *ln2_out = MODEL_PTR(model, L->ln2_out);
1702 float *q = MODEL_PTR(model, L->q);
1703 float *k = MODEL_PTR(model, L->k);
1704 float *v = MODEL_PTR(model, L->v);
1705 float *attn_out = MODEL_PTR(model, L->attn_out);
1706 float *proj_tmp = MODEL_PTR(model, L->proj_tmp);
1707 float *proj_scratch = MODEL_PTR(model, L->proj_scratch);
1708 float *residual1 = MODEL_PTR(model, L->residual1);
1709 float *fc1_out = MODEL_PTR(model, L->fc1_out);
1710 float *swiglu_out = MODEL_PTR(model, L->swiglu_out);
1711 float *mlp_out = MODEL_PTR(model, L->mlp_out);
1712 float *output = MODEL_PTR(model, L->output);
1713
1714 const void *WQ = (const void *)MODEL_PTR(model, L->wq);
1715 const void *WK = (const void *)MODEL_PTR(model, L->wk);
1716 const void *WV = (const void *)MODEL_PTR(model, L->wv);
1717 const void *WO = (const void *)MODEL_PTR(model, L->wo);
1718 const void *W1 = (const void *)MODEL_PTR(model, L->w1);
1719 const void *W2 = (const void *)MODEL_PTR(model, L->w2);
1720 const float *BQ = NULL;
1721 const float *BK = NULL;
1722 const float *BV = NULL;
1723 const float *BO = NULL;
1724 const float *B1 = NULL;
1725 const float *B2 = NULL;
1726
1727 float *rope_cos = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
1728 float *rope_sin = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
1729
1730 const int H = MODEL_NUM_HEADS;
1731 const int H_kv = MODEL_NUM_KV_HEADS;
1732 const int head_dim = MODEL_HEAD_DIM;
1733 const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
1734 const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
1735 const size_t kv_head_stride = (size_t)aligned_context_window * (size_t)aligned_head_dim;
1736
1737 /* RMSNorm before attention */
1738 rmsnorm_forward(input,
1739 ln1_gamma,
1740 ln1_out,
1741 NULL,
1742 num_tokens,
1744 aligned_embed_dim,
1745 1e-06f);
1746
1747 /* Q projection (head-major) */
1748 const size_t wq_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1749 const uint8_t *WQ_bytes = (const uint8_t *)WQ;
1750 for (int h = 0; h < H; ++h) {
1751 const void *wq_h = (const void *)(WQ_bytes + (size_t)h * wq_head_bytes);
1752 const float *bq_h = BQ ? (BQ + (size_t)h * (size_t)aligned_head_dim) : NULL;
1753 float *q_h = q + (size_t)h * q_head_stride;
1754 gemm_nt_q4_k(ln1_out, wq_h, bq_h, q_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1755 }
1756
1757 /* K projection (head-major) */
1758 const size_t wk_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1759 const uint8_t *WK_bytes = (const uint8_t *)WK;
1760 for (int h = 0; h < H_kv; ++h) {
1761 const void *wk_h = (const void *)(WK_bytes + (size_t)h * wk_head_bytes);
1762 const float *bk_h = BK ? (BK + (size_t)h * (size_t)aligned_head_dim) : NULL;
1763 float *k_h = k + (size_t)h * kv_head_stride;
1764 gemm_nt_q4_k(ln1_out, wk_h, bk_h, k_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1765 }
1766
1767 /* V projection (head-major) */
1768 const size_t wv_head_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, head_w_elems);
1769 const uint8_t *WV_bytes = (const uint8_t *)WV;
1770 for (int h = 0; h < H_kv; ++h) {
1771 const void *wv_h = (const void *)(WV_bytes + (size_t)h * wv_head_bytes);
1772 const float *bv_h = BV ? (BV + (size_t)h * (size_t)aligned_head_dim) : NULL;
1773 float *v_h = v + (size_t)h * kv_head_stride;
1774 gemm_nt_q4_k(ln1_out, wv_h, bv_h, v_h, num_tokens, aligned_head_dim, aligned_embed_dim);
1775 }
1776
1777 /* RoPE */
1779 k,
1780 rope_cos,
1781 rope_sin,
1782 H,
1783 H_kv,
1784 num_tokens,
1785 head_dim,
1786 aligned_head_dim,
1787 0,
1788 num_tokens,
1789 aligned_context_window);
1790
1791 /* Attention (prefill, causal) */
1793 k,
1794 v,
1795 attn_out,
1796 H,
1797 H_kv,
1798 num_tokens,
1799 head_dim,
1800 aligned_head_dim,
1801 aligned_context_window);
1802
1803 /* Output projection (flatten head-major to token-major) */
1804 const int K = H * aligned_head_dim;
1805 if (K != aligned_embed_dim) {
1806 return;
1807 }
1808 const float *proj_in = attn_out;
1809 if (H > 1) {
1810 if (!proj_scratch) {
1811 return;
1812 }
1813 for (int t = 0; t < num_tokens; ++t) {
1814 float *dst = proj_scratch + (size_t)t * (size_t)aligned_embed_dim;
1815 for (int h = 0; h < H; ++h) {
1816 const float *src = attn_out + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
1817 memcpy(dst + (size_t)h * (size_t)aligned_head_dim,
1818 src,
1819 (size_t)aligned_head_dim * sizeof(float));
1820 }
1821 }
1822 proj_in = proj_scratch;
1823 }
1824 gemm_nt_q4_k(proj_in, WO, BO, proj_tmp, num_tokens, aligned_embed_dim, K);
1825
1826 /* Residual add */
1827 model_residual_add_token_major(input, proj_tmp, residual1, num_tokens, aligned_embed_dim);
1828
1829 /* RMSNorm before MLP */
1830 rmsnorm_forward(residual1,
1831 ln2_gamma,
1832 ln2_out,
1833 NULL,
1834 num_tokens,
1836 aligned_embed_dim,
1837 1e-06f);
1838
1839 /* MLP (SwiGLU) */
1840 gemm_nt_q4_k(ln2_out, W1, B1, fc1_out, num_tokens, 2 * aligned_intermediate_dim, aligned_embed_dim);
1841 swiglu_forward(fc1_out, swiglu_out, num_tokens, aligned_intermediate_dim);
1842 gemm_nt_q4_k(swiglu_out, W2, B2, mlp_out, num_tokens, aligned_embed_dim, aligned_intermediate_dim);
1843
1844 /* Final residual add */
1845 model_residual_add_token_major(residual1, mlp_out, output, num_tokens, aligned_embed_dim);
1846}

References attention_forward_causal_head_major_gqa_flash_strided(), CK_DT_Q4_K, ck_dtype_row_bytes(), gemm_nt_q4_k(), MODEL_EMBED_DIM, MODEL_HEAD_DIM, MODEL_NUM_HEADS, MODEL_NUM_KV_HEADS, model_residual_add_token_major(), rmsnorm_forward(), rope_forward_qk_strided(), and swiglu_forward().

Referenced by model_forward_prefill_impl().

◆ model_model_allocate()

int model_model_allocate ( MODELModel *  model)

Definition at line 88 of file v6.6/test_generated/generic_api_test.c.

88 {
89 size_t total = MODEL_TOTAL_BYTES;
90
91#ifdef __linux__
92 model->base = mmap(NULL, total,
93 PROT_READ | PROT_WRITE,
94 MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB,
95 -1, 0);
96 if (model->base == MAP_FAILED) {
97 model->base = mmap(NULL, total,
98 PROT_READ | PROT_WRITE,
99 MAP_PRIVATE | MAP_ANONYMOUS,
100 -1, 0);
101 }
102 if (model->base == MAP_FAILED) {
103 perror("mmap failed");
104 return -1;
105 }
106#else
107 model->base = aligned_alloc(64, total);
108 if (!model->base) {
109 perror("aligned_alloc failed");
110 return -1;
111 }
112#endif
113
114 model->total_bytes = total;
115
116 /* Initialize magic header */
117 MagicHeader *header = (MagicHeader *)model->base;
118 header->magic = MODEL_MAGIC;
119 header->version = 5;
120 header->total_bytes = MODEL_TOTAL_BYTES;
121 header->weight_bytes = MODEL_WEIGHT_BYTES;
122 header->activation_bytes = MODEL_ACTIVATION_BYTES;
123 header->num_layers = MODEL_NUM_LAYERS;
124 header->embed_dim = MODEL_EMBED_DIM;
125 header->num_heads = MODEL_NUM_HEADS;
126 header->vocab_size = MODEL_VOCAB_SIZE;
127 header->max_seq_len = MODEL_MAX_SEQ_LEN;
128 header->canary_count = MODEL_CANARY_COUNT;
129
130 /* Initialize canary guards */
131 for (int i = 0; i < MODEL_CANARY_COUNT; i++) {
132 uint32_t *ptr = (uint32_t*)((char*)model->base + MODEL_CANARIES[i].offset);
133 for (int j = 0; j < (MODEL_CANARY_SIZE / 4); j++) {
134 ptr[j] = MODEL_CANARY_VALUE;
135 }
136 }
137
138 return 0;
139}
#define MODEL_MAX_SEQ_LEN
Definition v6.5_simple.c:31

References MagicHeader, MODEL_EMBED_DIM, MODEL_MAX_SEQ_LEN, MODEL_NUM_HEADS, MODEL_NUM_LAYERS, and MODEL_VOCAB_SIZE.

Referenced by ck_model_create().

◆ model_model_free()

void model_model_free ( MODELModel *  model)

Definition at line 141 of file v6.6/test_generated/generic_api_test.c.

141 {
142 if (!model || !model->base) return;
143#ifdef __linux__
144 munmap(model->base, model->total_bytes);
145#else
146 free(model->base);
147#endif
148 model->base = NULL;
149 model->total_bytes = 0;
150}

Referenced by ck_model_free().

◆ model_precompute_rope()

void model_precompute_rope ( MODELModel *  model)

Definition at line 186 of file v6.6/test_generated/generic_api_test.c.

186 {
187 const int T = MODEL_MAX_SEQ_LEN;
188 const int D = MODEL_HEAD_DIM / 2;
189 const float theta = 1000000.0f;
190
191 float *cos_ptr = MODEL_PTR(model, MODEL_GLOBALS.rope_cos_cache);
192 float *sin_ptr = MODEL_PTR(model, MODEL_GLOBALS.rope_sin_cache);
193
194 for (int pos = 0; pos < T; pos++) {
195 for (int i = 0; i < D; i++) {
196 float freq = 1.0f / powf(theta, (float)(2 * i) / (float)(D * 2));
197 float angle = (float)pos * freq;
198 cos_ptr[pos * D + i] = cosf(angle);
199 sin_ptr[pos * D + i] = sinf(angle);
200 }
201 }
202}

References MODEL_HEAD_DIM, and MODEL_MAX_SEQ_LEN.

Referenced by ck_model_precompute_rope().

◆ model_residual_add_token_major()

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

Definition at line 43 of file v6.6/test_generated/generic_api_test.c.

49 {
50 if (!a || !b || !out) {
51 return;
52 }
53 for (int t = 0; t < tokens; ++t) {
54 const float *pa = a + (size_t)t * (size_t)aligned_embed_dim;
55 const float *pb = b + (size_t)t * (size_t)aligned_embed_dim;
56 float *pc = out + (size_t)t * (size_t)aligned_embed_dim;
57 for (int d = 0; d < aligned_embed_dim; ++d) {
58 pc[d] = pa[d] + pb[d];
59 }
60 }
61}

Referenced by model_layer_0_decode(), model_layer_0_prefill(), model_layer_10_decode(), model_layer_10_prefill(), model_layer_11_decode(), model_layer_11_prefill(), model_layer_12_decode(), model_layer_12_prefill(), model_layer_13_decode(), model_layer_13_prefill(), model_layer_14_decode(), model_layer_14_prefill(), model_layer_15_decode(), model_layer_15_prefill(), model_layer_16_decode(), model_layer_16_prefill(), model_layer_17_decode(), model_layer_17_prefill(), model_layer_18_decode(), model_layer_18_prefill(), model_layer_19_decode(), model_layer_19_prefill(), model_layer_1_decode(), model_layer_1_prefill(), model_layer_20_decode(), model_layer_20_prefill(), model_layer_21_decode(), model_layer_21_prefill(), model_layer_22_decode(), model_layer_22_prefill(), model_layer_23_decode(), model_layer_23_prefill(), model_layer_2_decode(), model_layer_2_prefill(), model_layer_3_decode(), model_layer_3_prefill(), model_layer_4_decode(), model_layer_4_prefill(), model_layer_5_decode(), model_layer_5_prefill(), model_layer_6_decode(), model_layer_6_prefill(), model_layer_7_decode(), model_layer_7_prefill(), model_layer_8_decode(), model_layer_8_prefill(), model_layer_9_decode(), and model_layer_9_prefill().

◆ model_verify_canaries()

int model_verify_canaries ( MODELModel *  model)

Definition at line 152 of file v6.6/test_generated/generic_api_test.c.

152 {
153 int errors = 0;
154 uint32_t *ptr;
155
156 for (int i = 0; i < MODEL_CANARY_COUNT; i++) {
157 ptr = (uint32_t*)((char*)model->base + MODEL_CANARIES[i].offset);
158 for (int j = 0; j < 4; j++) {
159 if (ptr[j] != MODEL_CANARY_VALUE) {
160 fprintf(stderr, "CANARY CORRUPTION: %s at offset 0x%lX\n",
161 MODEL_CANARIES[i].name,
162 MODEL_CANARIES[i].offset);
163 errors++;
164 break;
165 }
166 }
167 }
168
169 return errors;
170}

Referenced by ck_model_verify_canaries().

Variable Documentation

◆ g_model_config

CKModelConfig g_model_config
static
Initial value:
= {
.embed_dim = MODEL_EMBED_DIM,
.num_heads = MODEL_NUM_HEADS,
.num_kv_heads = MODEL_NUM_KV_HEADS,
.head_dim = MODEL_HEAD_DIM,
.intermediate_size = MODEL_INTERMEDIATE,
.num_layers = MODEL_NUM_LAYERS,
.vocab_size = MODEL_VOCAB_SIZE,
.max_seq_len = MODEL_MAX_SEQ_LEN,
.total_bytes = MODEL_TOTAL_BYTES,
.weight_bytes = MODEL_WEIGHT_BYTES,
.activation_bytes = MODEL_ACTIVATION_BYTES,
.model_name = "model",
.model_family = "model",
}
#define MODEL_INTERMEDIATE
Definition v6.5_simple.c:36

Definition at line 8562 of file v6.6/test_generated/generic_api_test.c.

8562 {
8563 .embed_dim = MODEL_EMBED_DIM,
8564 .num_heads = MODEL_NUM_HEADS,
8565 .num_kv_heads = MODEL_NUM_KV_HEADS,
8566 .head_dim = MODEL_HEAD_DIM,
8567 .intermediate_size = MODEL_INTERMEDIATE,
8568 .num_layers = MODEL_NUM_LAYERS,
8569 .vocab_size = MODEL_VOCAB_SIZE,
8570 .max_seq_len = MODEL_MAX_SEQ_LEN,
8571 .total_bytes = MODEL_TOTAL_BYTES,
8572 .weight_bytes = MODEL_WEIGHT_BYTES,
8573 .activation_bytes = MODEL_ACTIVATION_BYTES,
8574 .model_name = "model",
8575 .model_family = "model",
8576};

Referenced by ck_model_get_config().

◆ MagicHeader

MagicHeader

Definition at line 80 of file v6.6/test_generated/generic_api_test.c.

Referenced by model_model_allocate().