← Back to C-Kernel-Engine Docs Doxygen Source Documentation
v6.5/test_generated/generic_api_test.c
Go to the documentation of this file.
1 /**
2  * @file generic_api_test.c
3  * @brief AUTO-GENERATED: model Implementation (IR v6.5 - Explicit Unrolled)
4  *
5  * Generated: 2026-01-12T04:22:25.186752 UTC
6  * Total Memory: 3.57 GB
7  * Mode: decode
8  * Layers: 24 (fully unrolled)
9  *
10  * Per-layer quant types:
11  * Layer 0: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
12  * Layer 1: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
13  * Layer 2: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
14  * ... (21 more layers)
15  *
16  * DO NOT EDIT - Regenerate with build_ir_v6.5.py or codegen_v6.5.py
17  */
18 
19 #define _GNU_SOURCE /* For MAP_ANONYMOUS, MAP_HUGETLB */
20 
21 #include "ck-kernel-inference.h"
22 
23 #include "ckernel_engine.h"
24 
25 #include <stdio.h>
26 #include <stdlib.h>
27 #include <string.h>
28 #include <stdint.h>
29 #include <math.h>
30 
31 #ifdef __linux__
32 #include <sys/mman.h>
33 #endif
34 
35 #if MODEL_DTYPE_BYTES != 4
36 #error "model: v6.5 codegen currently supports fp32 only. Use --dtype=fp32."
37 #endif
38 
39 /* ============================================================================
40  * LOCAL HELPERS (no orchestration dependency)
41  * ============================================================================ */
42 
43 static void model_residual_add_token_major(
44  const float *a,
45  const float *b,
46  float *out,
47  int tokens,
48  int aligned_embed_dim
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 }
62 
63 /* ============================================================================
64  * MAGIC HEADER
65  * ============================================================================ */
66 
67 typedef struct __attribute__((packed)) {
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 */
80 } MagicHeader;
81 
82 _Static_assert(sizeof(MagicHeader) == 64, "MagicHeader must be 64 bytes");
83 
84 /* ============================================================================
85  * ALLOCATION
86  * ============================================================================ */
87 
88 int model_model_allocate(MODELModel *model) {
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 }
140 
141 void model_model_free(MODELModel *model) {
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 }
151 
152 int model_verify_canaries(MODELModel *model) {
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 }
171 
172 /* ============================================================================
173  * ALIGNMENT HELPERS
174  * ============================================================================ */
175 
176 static int model_align_elems(int elems, int elem_bytes, int align_bytes) {
177  int bytes = elems * elem_bytes;
178  int aligned = (bytes + align_bytes - 1) / align_bytes * align_bytes;
179  return aligned / elem_bytes;
180 }
181 
182 /* ============================================================================
183  * ROPE PRECOMPUTE
184  * ============================================================================ */
185 
186 void model_precompute_rope(MODELModel *model) {
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 }
203 
204 /* ============================================================================
205  * EXPLICIT PER-LAYER PREFILL FUNCTIONS (v6.5 unrolled)
206  * ============================================================================ */
207 
208 /*
209  * Layer 0: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
210  */
211 static void model_layer_0_prefill(
212  MODELModel *model,
213  int num_tokens,
214  int aligned_embed_dim,
215  int aligned_head_dim,
216  int aligned_intermediate_dim,
217  int aligned_context_window
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,
267  MODEL_EMBED_DIM,
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 */
302  rope_forward_qk_strided(q,
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) */
316  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
359  MODEL_EMBED_DIM,
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 }
371 
372 /*
373  * Layer 1: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
374  */
375 static void model_layer_1_prefill(
376  MODELModel *model,
377  int num_tokens,
378  int aligned_embed_dim,
379  int aligned_head_dim,
380  int aligned_intermediate_dim,
381  int aligned_context_window
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,
431  MODEL_EMBED_DIM,
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 */
466  rope_forward_qk_strided(q,
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) */
480  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
523  MODEL_EMBED_DIM,
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 }
535 
536 /*
537  * Layer 2: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
538  */
539 static void model_layer_2_prefill(
540  MODELModel *model,
541  int num_tokens,
542  int aligned_embed_dim,
543  int aligned_head_dim,
544  int aligned_intermediate_dim,
545  int aligned_context_window
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,
595  MODEL_EMBED_DIM,
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 */
630  rope_forward_qk_strided(q,
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) */
644  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
687  MODEL_EMBED_DIM,
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 }
699 
700 /*
701  * Layer 3: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
702  */
703 static void model_layer_3_prefill(
704  MODELModel *model,
705  int num_tokens,
706  int aligned_embed_dim,
707  int aligned_head_dim,
708  int aligned_intermediate_dim,
709  int aligned_context_window
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,
759  MODEL_EMBED_DIM,
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 */
794  rope_forward_qk_strided(q,
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) */
808  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
851  MODEL_EMBED_DIM,
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 }
863 
864 /*
865  * Layer 4: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
866  */
867 static void model_layer_4_prefill(
868  MODELModel *model,
869  int num_tokens,
870  int aligned_embed_dim,
871  int aligned_head_dim,
872  int aligned_intermediate_dim,
873  int aligned_context_window
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,
923  MODEL_EMBED_DIM,
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 */
958  rope_forward_qk_strided(q,
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) */
972  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1015  MODEL_EMBED_DIM,
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 }
1027 
1028 /*
1029  * Layer 5: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
1030  */
1031 static void model_layer_5_prefill(
1032  MODELModel *model,
1033  int num_tokens,
1034  int aligned_embed_dim,
1035  int aligned_head_dim,
1036  int aligned_intermediate_dim,
1037  int aligned_context_window
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,
1087  MODEL_EMBED_DIM,
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 */
1122  rope_forward_qk_strided(q,
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) */
1136  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1179  MODEL_EMBED_DIM,
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 }
1191 
1192 /*
1193  * Layer 6: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
1194  */
1195 static void model_layer_6_prefill(
1196  MODELModel *model,
1197  int num_tokens,
1198  int aligned_embed_dim,
1199  int aligned_head_dim,
1200  int aligned_intermediate_dim,
1201  int aligned_context_window
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,
1251  MODEL_EMBED_DIM,
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 */
1286  rope_forward_qk_strided(q,
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) */
1300  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1343  MODEL_EMBED_DIM,
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 }
1355 
1356 /*
1357  * Layer 7: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
1358  */
1359 static void model_layer_7_prefill(
1360  MODELModel *model,
1361  int num_tokens,
1362  int aligned_embed_dim,
1363  int aligned_head_dim,
1364  int aligned_intermediate_dim,
1365  int aligned_context_window
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,
1415  MODEL_EMBED_DIM,
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 */
1450  rope_forward_qk_strided(q,
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) */
1464  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1507  MODEL_EMBED_DIM,
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 }
1519 
1520 /*
1521  * Layer 8: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
1522  */
1523 static void model_layer_8_prefill(
1524  MODELModel *model,
1525  int num_tokens,
1526  int aligned_embed_dim,
1527  int aligned_head_dim,
1528  int aligned_intermediate_dim,
1529  int aligned_context_window
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,
1579  MODEL_EMBED_DIM,
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 */
1614  rope_forward_qk_strided(q,
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) */
1628  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1671  MODEL_EMBED_DIM,
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 }
1683 
1684 /*
1685  * Layer 9: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
1686  */
1687 static void model_layer_9_prefill(
1688  MODELModel *model,
1689  int num_tokens,
1690  int aligned_embed_dim,
1691  int aligned_head_dim,
1692  int aligned_intermediate_dim,
1693  int aligned_context_window
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,
1743  MODEL_EMBED_DIM,
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 */
1778  rope_forward_qk_strided(q,
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) */
1792  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1835  MODEL_EMBED_DIM,
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 }
1847 
1848 /*
1849  * Layer 10: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
1850  */
1851 static void model_layer_10_prefill(
1852  MODELModel *model,
1853  int num_tokens,
1854  int aligned_embed_dim,
1855  int aligned_head_dim,
1856  int aligned_intermediate_dim,
1857  int aligned_context_window
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,
1907  MODEL_EMBED_DIM,
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 */
1942  rope_forward_qk_strided(q,
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) */
1956  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
1999  MODEL_EMBED_DIM,
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 }
2011 
2012 /*
2013  * Layer 11: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2014  */
2015 static void model_layer_11_prefill(
2016  MODELModel *model,
2017  int num_tokens,
2018  int aligned_embed_dim,
2019  int aligned_head_dim,
2020  int aligned_intermediate_dim,
2021  int aligned_context_window
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,
2071  MODEL_EMBED_DIM,
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 */
2106  rope_forward_qk_strided(q,
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) */
2120  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
2163  MODEL_EMBED_DIM,
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 }
2175 
2176 /*
2177  * Layer 12: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2178  */
2179 static void model_layer_12_prefill(
2180  MODELModel *model,
2181  int num_tokens,
2182  int aligned_embed_dim,
2183  int aligned_head_dim,
2184  int aligned_intermediate_dim,
2185  int aligned_context_window
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,
2235  MODEL_EMBED_DIM,
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 */
2270  rope_forward_qk_strided(q,
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) */
2284  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
2327  MODEL_EMBED_DIM,
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 }
2339 
2340 /*
2341  * Layer 13: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2342  */
2343 static void model_layer_13_prefill(
2344  MODELModel *model,
2345  int num_tokens,
2346  int aligned_embed_dim,
2347  int aligned_head_dim,
2348  int aligned_intermediate_dim,
2349  int aligned_context_window
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,
2399  MODEL_EMBED_DIM,
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 */
2434  rope_forward_qk_strided(q,
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) */
2448  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
2491  MODEL_EMBED_DIM,
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 }
2503 
2504 /*
2505  * Layer 14: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2506  */
2507 static void model_layer_14_prefill(
2508  MODELModel *model,
2509  int num_tokens,
2510  int aligned_embed_dim,
2511  int aligned_head_dim,
2512  int aligned_intermediate_dim,
2513  int aligned_context_window
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,
2563  MODEL_EMBED_DIM,
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 */
2598  rope_forward_qk_strided(q,
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) */
2612  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
2655  MODEL_EMBED_DIM,
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 }
2667 
2668 /*
2669  * Layer 15: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2670  */
2671 static void model_layer_15_prefill(
2672  MODELModel *model,
2673  int num_tokens,
2674  int aligned_embed_dim,
2675  int aligned_head_dim,
2676  int aligned_intermediate_dim,
2677  int aligned_context_window
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,
2727  MODEL_EMBED_DIM,
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 */
2762  rope_forward_qk_strided(q,
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) */
2776  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
2819  MODEL_EMBED_DIM,
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 }
2831 
2832 /*
2833  * Layer 16: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2834  */
2835 static void model_layer_16_prefill(
2836  MODELModel *model,
2837  int num_tokens,
2838  int aligned_embed_dim,
2839  int aligned_head_dim,
2840  int aligned_intermediate_dim,
2841  int aligned_context_window
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,
2891  MODEL_EMBED_DIM,
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 */
2926  rope_forward_qk_strided(q,
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) */
2940  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
2983  MODEL_EMBED_DIM,
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 }
2995 
2996 /*
2997  * Layer 17: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
2998  */
2999 static void model_layer_17_prefill(
3000  MODELModel *model,
3001  int num_tokens,
3002  int aligned_embed_dim,
3003  int aligned_head_dim,
3004  int aligned_intermediate_dim,
3005  int aligned_context_window
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,
3055  MODEL_EMBED_DIM,
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 */
3090  rope_forward_qk_strided(q,
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) */
3104  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
3147  MODEL_EMBED_DIM,
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 }
3159 
3160 /*
3161  * Layer 18: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
3162  */
3163 static void model_layer_18_prefill(
3164  MODELModel *model,
3165  int num_tokens,
3166  int aligned_embed_dim,
3167  int aligned_head_dim,
3168  int aligned_intermediate_dim,
3169  int aligned_context_window
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,
3219  MODEL_EMBED_DIM,
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 */
3254  rope_forward_qk_strided(q,
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) */
3268  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
3311  MODEL_EMBED_DIM,
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 }
3323 
3324 /*
3325  * Layer 19: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
3326  */
3327 static void model_layer_19_prefill(
3328  MODELModel *model,
3329  int num_tokens,
3330  int aligned_embed_dim,
3331  int aligned_head_dim,
3332  int aligned_intermediate_dim,
3333  int aligned_context_window
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,
3383  MODEL_EMBED_DIM,
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 */
3418  rope_forward_qk_strided(q,
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) */
3432  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
3475  MODEL_EMBED_DIM,
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 }
3487 
3488 /*
3489  * Layer 20: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
3490  */
3491 static void model_layer_20_prefill(
3492  MODELModel *model,
3493  int num_tokens,
3494  int aligned_embed_dim,
3495  int aligned_head_dim,
3496  int aligned_intermediate_dim,
3497  int aligned_context_window
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,
3547  MODEL_EMBED_DIM,
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 */
3582  rope_forward_qk_strided(q,
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) */
3596  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
3639  MODEL_EMBED_DIM,
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 }
3651 
3652 /*
3653  * Layer 21: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
3654  */
3655 static void model_layer_21_prefill(
3656  MODELModel *model,
3657  int num_tokens,
3658  int aligned_embed_dim,
3659  int aligned_head_dim,
3660  int aligned_intermediate_dim,
3661  int aligned_context_window
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,
3711  MODEL_EMBED_DIM,
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 */
3746  rope_forward_qk_strided(q,
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) */
3760  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
3803  MODEL_EMBED_DIM,
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 }
3815 
3816 /*
3817  * Layer 22: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
3818  */
3819 static void model_layer_22_prefill(
3820  MODELModel *model,
3821  int num_tokens,
3822  int aligned_embed_dim,
3823  int aligned_head_dim,
3824  int aligned_intermediate_dim,
3825  int aligned_context_window
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,
3875  MODEL_EMBED_DIM,
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 */
3910  rope_forward_qk_strided(q,
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) */
3924  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
3967  MODEL_EMBED_DIM,
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 }
3979 
3980 /*
3981  * Layer 23: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
3982  */
3983 static void model_layer_23_prefill(
3984  MODELModel *model,
3985  int num_tokens,
3986  int aligned_embed_dim,
3987  int aligned_head_dim,
3988  int aligned_intermediate_dim,
3989  int aligned_context_window
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,
4039  MODEL_EMBED_DIM,
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 */
4074  rope_forward_qk_strided(q,
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) */
4088  attention_forward_causal_head_major_gqa_flash_strided(q,
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,
4131  MODEL_EMBED_DIM,
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 }
4143 
4144 /* ============================================================================
4145  * FORWARD PASS (PREFILL)
4146  * ============================================================================ */
4147 
4148 static void model_forward_prefill_impl(
4149  MODELModel *model,
4150  const int *tokens,
4151  int num_tokens
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,
4167  MODEL_VOCAB_SIZE,
4168  embed_weight,
4169  NULL,
4170  embed_out,
4171  MODEL_EMBED_DIM,
4172  aligned_embed_dim,
4173  num_tokens,
4174  0);
4175 
4176  model_layer_0_prefill(
4177  model,
4178  num_tokens,
4179  aligned_embed_dim,
4180  aligned_head_dim,
4181  aligned_intermediate_dim,
4182  aligned_context_window);
4183 
4184  model_layer_1_prefill(
4185  model,
4186  num_tokens,
4187  aligned_embed_dim,
4188  aligned_head_dim,
4189  aligned_intermediate_dim,
4190  aligned_context_window);
4191 
4192  model_layer_2_prefill(
4193  model,
4194  num_tokens,
4195  aligned_embed_dim,
4196  aligned_head_dim,
4197  aligned_intermediate_dim,
4198  aligned_context_window);
4199 
4200  model_layer_3_prefill(
4201  model,
4202  num_tokens,
4203  aligned_embed_dim,
4204  aligned_head_dim,
4205  aligned_intermediate_dim,
4206  aligned_context_window);
4207 
4208  model_layer_4_prefill(
4209  model,
4210  num_tokens,
4211  aligned_embed_dim,
4212  aligned_head_dim,
4213  aligned_intermediate_dim,
4214  aligned_context_window);
4215 
4216  model_layer_5_prefill(
4217  model,
4218  num_tokens,
4219  aligned_embed_dim,
4220  aligned_head_dim,
4221  aligned_intermediate_dim,
4222  aligned_context_window);
4223 
4224  model_layer_6_prefill(
4225  model,
4226  num_tokens,
4227  aligned_embed_dim,
4228  aligned_head_dim,
4229  aligned_intermediate_dim,
4230  aligned_context_window);
4231 
4232  model_layer_7_prefill(
4233  model,
4234  num_tokens,
4235  aligned_embed_dim,
4236  aligned_head_dim,
4237  aligned_intermediate_dim,
4238  aligned_context_window);
4239 
4240  model_layer_8_prefill(
4241  model,
4242  num_tokens,
4243  aligned_embed_dim,
4244  aligned_head_dim,
4245  aligned_intermediate_dim,
4246  aligned_context_window);
4247 
4248  model_layer_9_prefill(
4249  model,
4250  num_tokens,
4251  aligned_embed_dim,
4252  aligned_head_dim,
4253  aligned_intermediate_dim,
4254  aligned_context_window);
4255 
4256  model_layer_10_prefill(
4257  model,
4258  num_tokens,
4259  aligned_embed_dim,
4260  aligned_head_dim,
4261  aligned_intermediate_dim,
4262  aligned_context_window);
4263 
4264  model_layer_11_prefill(
4265  model,
4266  num_tokens,
4267  aligned_embed_dim,
4268  aligned_head_dim,
4269  aligned_intermediate_dim,
4270  aligned_context_window);
4271 
4272  model_layer_12_prefill(
4273  model,
4274  num_tokens,
4275  aligned_embed_dim,
4276  aligned_head_dim,
4277  aligned_intermediate_dim,
4278  aligned_context_window);
4279 
4280  model_layer_13_prefill(
4281  model,
4282  num_tokens,
4283  aligned_embed_dim,
4284  aligned_head_dim,
4285  aligned_intermediate_dim,
4286  aligned_context_window);
4287 
4288  model_layer_14_prefill(
4289  model,
4290  num_tokens,
4291  aligned_embed_dim,
4292  aligned_head_dim,
4293  aligned_intermediate_dim,
4294  aligned_context_window);
4295 
4296  model_layer_15_prefill(
4297  model,
4298  num_tokens,
4299  aligned_embed_dim,
4300  aligned_head_dim,
4301  aligned_intermediate_dim,
4302  aligned_context_window);
4303 
4304  model_layer_16_prefill(
4305  model,
4306  num_tokens,
4307  aligned_embed_dim,
4308  aligned_head_dim,
4309  aligned_intermediate_dim,
4310  aligned_context_window);
4311 
4312  model_layer_17_prefill(
4313  model,
4314  num_tokens,
4315  aligned_embed_dim,
4316  aligned_head_dim,
4317  aligned_intermediate_dim,
4318  aligned_context_window);
4319 
4320  model_layer_18_prefill(
4321  model,
4322  num_tokens,
4323  aligned_embed_dim,
4324  aligned_head_dim,
4325  aligned_intermediate_dim,
4326  aligned_context_window);
4327 
4328  model_layer_19_prefill(
4329  model,
4330  num_tokens,
4331  aligned_embed_dim,
4332  aligned_head_dim,
4333  aligned_intermediate_dim,
4334  aligned_context_window);
4335 
4336  model_layer_20_prefill(
4337  model,
4338  num_tokens,
4339  aligned_embed_dim,
4340  aligned_head_dim,
4341  aligned_intermediate_dim,
4342  aligned_context_window);
4343 
4344  model_layer_21_prefill(
4345  model,
4346  num_tokens,
4347  aligned_embed_dim,
4348  aligned_head_dim,
4349  aligned_intermediate_dim,
4350  aligned_context_window);
4351 
4352  model_layer_22_prefill(
4353  model,
4354  num_tokens,
4355  aligned_embed_dim,
4356  aligned_head_dim,
4357  aligned_intermediate_dim,
4358  aligned_context_window);
4359 
4360  model_layer_23_prefill(
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,
4376  MODEL_EMBED_DIM,
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,
4393  MODEL_VOCAB_SIZE,
4394  aligned_embed_dim);
4395  }
4396 }
4397 
4398 /* ============================================================================
4399  * EXPLICIT PER-LAYER DECODE FUNCTIONS (v6.5 unrolled)
4400  * ============================================================================ */
4401 
4402 /*
4403  * Layer 0: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
4404  */
4405 static void model_layer_0_decode(
4406  MODELModel *model,
4407  int token_index,
4408  int aligned_embed_dim,
4409  int aligned_head_dim,
4410  int aligned_intermediate_dim,
4411  int aligned_context_window
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,
4459  MODEL_EMBED_DIM,
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) */
4529  attention_forward_decode_head_major_gqa_flash(q_token,
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,
4553  MODEL_EMBED_DIM,
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 }
4570 
4571 /*
4572  * Layer 1: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
4573  */
4574 static void model_layer_1_decode(
4575  MODELModel *model,
4576  int token_index,
4577  int aligned_embed_dim,
4578  int aligned_head_dim,
4579  int aligned_intermediate_dim,
4580  int aligned_context_window
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,
4628  MODEL_EMBED_DIM,
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) */
4698  attention_forward_decode_head_major_gqa_flash(q_token,
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,
4722  MODEL_EMBED_DIM,
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 }
4739 
4740 /*
4741  * Layer 2: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
4742  */
4743 static void model_layer_2_decode(
4744  MODELModel *model,
4745  int token_index,
4746  int aligned_embed_dim,
4747  int aligned_head_dim,
4748  int aligned_intermediate_dim,
4749  int aligned_context_window
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,
4797  MODEL_EMBED_DIM,
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) */
4867  attention_forward_decode_head_major_gqa_flash(q_token,
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,
4891  MODEL_EMBED_DIM,
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 }
4908 
4909 /*
4910  * Layer 3: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
4911  */
4912 static void model_layer_3_decode(
4913  MODELModel *model,
4914  int token_index,
4915  int aligned_embed_dim,
4916  int aligned_head_dim,
4917  int aligned_intermediate_dim,
4918  int aligned_context_window
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,
4966  MODEL_EMBED_DIM,
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) */
5036  attention_forward_decode_head_major_gqa_flash(q_token,
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,
5060  MODEL_EMBED_DIM,
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 }
5077 
5078 /*
5079  * Layer 4: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
5080  */
5081 static void model_layer_4_decode(
5082  MODELModel *model,
5083  int token_index,
5084  int aligned_embed_dim,
5085  int aligned_head_dim,
5086  int aligned_intermediate_dim,
5087  int aligned_context_window
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,
5135  MODEL_EMBED_DIM,
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) */
5205  attention_forward_decode_head_major_gqa_flash(q_token,
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,
5229  MODEL_EMBED_DIM,
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 }
5246 
5247 /*
5248  * Layer 5: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
5249  */
5250 static void model_layer_5_decode(
5251  MODELModel *model,
5252  int token_index,
5253  int aligned_embed_dim,
5254  int aligned_head_dim,
5255  int aligned_intermediate_dim,
5256  int aligned_context_window
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,
5304  MODEL_EMBED_DIM,
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) */
5374  attention_forward_decode_head_major_gqa_flash(q_token,
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,
5398  MODEL_EMBED_DIM,
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 }
5415 
5416 /*
5417  * Layer 6: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
5418  */
5419 static void model_layer_6_decode(
5420  MODELModel *model,
5421  int token_index,
5422  int aligned_embed_dim,
5423  int aligned_head_dim,
5424  int aligned_intermediate_dim,
5425  int aligned_context_window
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,
5473  MODEL_EMBED_DIM,
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) */
5543  attention_forward_decode_head_major_gqa_flash(q_token,
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,
5567  MODEL_EMBED_DIM,
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 }
5584 
5585 /*
5586  * Layer 7: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
5587  */
5588 static void model_layer_7_decode(
5589  MODELModel *model,
5590  int token_index,
5591  int aligned_embed_dim,
5592  int aligned_head_dim,
5593  int aligned_intermediate_dim,
5594  int aligned_context_window
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,
5642  MODEL_EMBED_DIM,
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) */
5712  attention_forward_decode_head_major_gqa_flash(q_token,
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,
5736  MODEL_EMBED_DIM,
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 }
5753 
5754 /*
5755  * Layer 8: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
5756  */
5757 static void model_layer_8_decode(
5758  MODELModel *model,
5759  int token_index,
5760  int aligned_embed_dim,
5761  int aligned_head_dim,
5762  int aligned_intermediate_dim,
5763  int aligned_context_window
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,
5811  MODEL_EMBED_DIM,
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) */
5881  attention_forward_decode_head_major_gqa_flash(q_token,
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,
5905  MODEL_EMBED_DIM,
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 }
5922 
5923 /*
5924  * Layer 9: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
5925  */
5926 static void model_layer_9_decode(
5927  MODELModel *model,
5928  int token_index,
5929  int aligned_embed_dim,
5930  int aligned_head_dim,
5931  int aligned_intermediate_dim,
5932  int aligned_context_window
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,
5980  MODEL_EMBED_DIM,
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) */
6050  attention_forward_decode_head_major_gqa_flash(q_token,
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,
6074  MODEL_EMBED_DIM,
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 }
6091 
6092 /*
6093  * Layer 10: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
6094  */
6095 static void model_layer_10_decode(
6096  MODELModel *model,
6097  int token_index,
6098  int aligned_embed_dim,
6099  int aligned_head_dim,
6100  int aligned_intermediate_dim,
6101  int aligned_context_window
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,
6149  MODEL_EMBED_DIM,
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) */
6219  attention_forward_decode_head_major_gqa_flash(q_token,
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,
6243  MODEL_EMBED_DIM,
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 }
6260 
6261 /*
6262  * Layer 11: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
6263  */
6264 static void model_layer_11_decode(
6265  MODELModel *model,
6266  int token_index,
6267  int aligned_embed_dim,
6268  int aligned_head_dim,
6269  int aligned_intermediate_dim,
6270  int aligned_context_window
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,
6318  MODEL_EMBED_DIM,
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) */
6388  attention_forward_decode_head_major_gqa_flash(q_token,
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,
6412  MODEL_EMBED_DIM,
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 }
6429 
6430 /*
6431  * Layer 12: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
6432  */
6433 static void model_layer_12_decode(
6434  MODELModel *model,
6435  int token_index,
6436  int aligned_embed_dim,
6437  int aligned_head_dim,
6438  int aligned_intermediate_dim,
6439  int aligned_context_window
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,
6487  MODEL_EMBED_DIM,
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) */
6557  attention_forward_decode_head_major_gqa_flash(q_token,
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,
6581  MODEL_EMBED_DIM,
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 }
6598 
6599 /*
6600  * Layer 13: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
6601  */
6602 static void model_layer_13_decode(
6603  MODELModel *model,
6604  int token_index,
6605  int aligned_embed_dim,
6606  int aligned_head_dim,
6607  int aligned_intermediate_dim,
6608  int aligned_context_window
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,
6656  MODEL_EMBED_DIM,
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) */
6726  attention_forward_decode_head_major_gqa_flash(q_token,
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,
6750  MODEL_EMBED_DIM,
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 }
6767 
6768 /*
6769  * Layer 14: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
6770  */
6771 static void model_layer_14_decode(
6772  MODELModel *model,
6773  int token_index,
6774  int aligned_embed_dim,
6775  int aligned_head_dim,
6776  int aligned_intermediate_dim,
6777  int aligned_context_window
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,
6825  MODEL_EMBED_DIM,
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) */
6895  attention_forward_decode_head_major_gqa_flash(q_token,
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,
6919  MODEL_EMBED_DIM,
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 }
6936 
6937 /*
6938  * Layer 15: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
6939  */
6940 static void model_layer_15_decode(
6941  MODELModel *model,
6942  int token_index,
6943  int aligned_embed_dim,
6944  int aligned_head_dim,
6945  int aligned_intermediate_dim,
6946  int aligned_context_window
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,
6994  MODEL_EMBED_DIM,
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) */
7064  attention_forward_decode_head_major_gqa_flash(q_token,
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,
7088  MODEL_EMBED_DIM,
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 }
7105 
7106 /*
7107  * Layer 16: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
7108  */
7109 static void model_layer_16_decode(
7110  MODELModel *model,
7111  int token_index,
7112  int aligned_embed_dim,
7113  int aligned_head_dim,
7114  int aligned_intermediate_dim,
7115  int aligned_context_window
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,
7163  MODEL_EMBED_DIM,
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) */
7233  attention_forward_decode_head_major_gqa_flash(q_token,
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,
7257  MODEL_EMBED_DIM,
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 }
7274 
7275 /*
7276  * Layer 17: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
7277  */
7278 static void model_layer_17_decode(
7279  MODELModel *model,
7280  int token_index,
7281  int aligned_embed_dim,
7282  int aligned_head_dim,
7283  int aligned_intermediate_dim,
7284  int aligned_context_window
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,
7332  MODEL_EMBED_DIM,
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) */
7402  attention_forward_decode_head_major_gqa_flash(q_token,
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,
7426  MODEL_EMBED_DIM,
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 }
7443 
7444 /*
7445  * Layer 18: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
7446  */
7447 static void model_layer_18_decode(
7448  MODELModel *model,
7449  int token_index,
7450  int aligned_embed_dim,
7451  int aligned_head_dim,
7452  int aligned_intermediate_dim,
7453  int aligned_context_window
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,
7501  MODEL_EMBED_DIM,
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) */
7571  attention_forward_decode_head_major_gqa_flash(q_token,
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,
7595  MODEL_EMBED_DIM,
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 }
7612 
7613 /*
7614  * Layer 19: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
7615  */
7616 static void model_layer_19_decode(
7617  MODELModel *model,
7618  int token_index,
7619  int aligned_embed_dim,
7620  int aligned_head_dim,
7621  int aligned_intermediate_dim,
7622  int aligned_context_window
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,
7670  MODEL_EMBED_DIM,
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) */
7740  attention_forward_decode_head_major_gqa_flash(q_token,
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,
7764  MODEL_EMBED_DIM,
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 }
7781 
7782 /*
7783  * Layer 20: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
7784  */
7785 static void model_layer_20_decode(
7786  MODELModel *model,
7787  int token_index,
7788  int aligned_embed_dim,
7789  int aligned_head_dim,
7790  int aligned_intermediate_dim,
7791  int aligned_context_window
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,
7839  MODEL_EMBED_DIM,
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) */
7909  attention_forward_decode_head_major_gqa_flash(q_token,
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,
7933  MODEL_EMBED_DIM,
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 }
7950 
7951 /*
7952  * Layer 21: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
7953  */
7954 static void model_layer_21_decode(
7955  MODELModel *model,
7956  int token_index,
7957  int aligned_embed_dim,
7958  int aligned_head_dim,
7959  int aligned_intermediate_dim,
7960  int aligned_context_window
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,
8008  MODEL_EMBED_DIM,
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) */
8078  attention_forward_decode_head_major_gqa_flash(q_token,
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,
8102  MODEL_EMBED_DIM,
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 }
8119 
8120 /*
8121  * Layer 22: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
8122  */
8123 static void model_layer_22_decode(
8124  MODELModel *model,
8125  int token_index,
8126  int aligned_embed_dim,
8127  int aligned_head_dim,
8128  int aligned_intermediate_dim,
8129  int aligned_context_window
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,
8177  MODEL_EMBED_DIM,
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) */
8247  attention_forward_decode_head_major_gqa_flash(q_token,
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,
8271  MODEL_EMBED_DIM,
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 }
8288 
8289 /*
8290  * Layer 23: wq=q4_k wk=q4_k wv=q4_k wo=q4_k w1=q4_k w2=q4_k
8291  */
8292 static void model_layer_23_decode(
8293  MODELModel *model,
8294  int token_index,
8295  int aligned_embed_dim,
8296  int aligned_head_dim,
8297  int aligned_intermediate_dim,
8298  int aligned_context_window
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,
8346  MODEL_EMBED_DIM,
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) */
8416  attention_forward_decode_head_major_gqa_flash(q_token,
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,
8440  MODEL_EMBED_DIM,
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 }
8457 
8458 /* ============================================================================
8459  * DECODE TOKEN (calls each layer explicitly)
8460  * ============================================================================ */
8461 
8462 static void model_decode_token(
8463  MODELModel *model,
8464  const int *token,
8465  int token_index
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,
8482  MODEL_VOCAB_SIZE,
8483  embed_weight,
8484  NULL,
8485  embed_out,
8486  MODEL_EMBED_DIM,
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,
8526  MODEL_EMBED_DIM,
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 }
8536 
8537 /* ============================================================================
8538  * PUBLIC API
8539  * ============================================================================ */
8540 
8541 void model_forward(
8542  MODELModel *model,
8543  const int *tokens,
8544  int num_tokens
8545 ) {
8546  if (!model || !tokens || num_tokens <= 0) return;
8547  model_forward_prefill_impl(model, tokens, num_tokens);
8548 }
8549 
8550 void model_decode(MODELModel *model, const int *token, int token_index) {
8551  model_decode_token(model, token, token_index);
8552 }
8553 
8554 /* ============================================================================
8555  * GENERIC MODEL API - Model-agnostic interface
8556  * These functions have the SAME names for all models.
8557  * CLI code uses these, so it works with any model.
8558  * ============================================================================ */
8559 
8560 #include "ck_model_api.h"
8561 
8562 static CKModelConfig g_model_config = {
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 };
8577 
8578 const CKModelConfig *ck_model_get_config(void) {
8579  return &g_model_config;
8580 }
8581 
8582 void *ck_model_create(void) {
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 }
8591 
8592 void ck_model_free(void *model) {
8593  if (!model) return;
8594  model_model_free((MODELModel *)model);
8595  free(model);
8596 }
8597 
8598 void ck_model_precompute_rope(void *model) {
8599  model_precompute_rope((MODELModel *)model);
8600 }
8601 
8602 void ck_model_forward(void *model, const int *tokens, int num_tokens) {
8603  model_forward((MODELModel *)model, tokens, num_tokens);
8604 }
8605 
8606 void ck_model_decode(void *model, const int *token, int token_index) {
8607  model_decode((MODELModel *)model, token, token_index);
8608 }
8609 
8610 float *ck_model_get_logits(void *model) {
8611  MODELModel *m = (MODELModel *)model;
8612  return MODEL_PTR(m, MODEL_FOOTER.logits);
8613 }
8614 
8615 int ck_model_verify_canaries(void *model) {
8616  return model_verify_canaries((MODELModel *)model);
8617 }
8618 
8619 void *ck_model_get_base(void *model) {
8620  return ((MODELModel *)model)->base;
8621 }
8622 
8623 size_t ck_model_get_total_bytes(void *model) {
8624  return ((MODELModel *)model)->total_bytes;
8625 }
ck_model_api.h
Generic Model API - Model-agnostic interface for CK-Engine.
CK_DT_Q4_K
@ CK_DT_Q4_K
Definition: ckernel_dtype.h:40
CK_DT_Q8_K
@ CK_DT_Q8_K
Definition: ckernel_dtype.h:43
ck_dtype_row_bytes
static size_t ck_dtype_row_bytes(CKDataType dt, size_t n_elements)
Calculate total bytes for n_elements of given dtype.
Definition: ckernel_dtype.h:138
ckernel_engine.h
attention_forward_causal_head_major_gqa_flash_strided
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)
Definition: attention_kernels.c:859
rope_forward_qk_strided
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)
Definition: rope_kernels.c:472
swiglu_forward
void swiglu_forward(const float *input, float *output, int tokens, int dim)
Definition: swiglu_kernels.c:131
gemm_nt_q4_k
void gemm_nt_q4_k(const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
Definition: gemm_kernels_q4k.c:683
rope_forward
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)
Definition: rope_kernels.c:180
embedding_forward_q4_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)
Definition: embedding_kernels.c:76
gemm_nt_q4_k_q8_k
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)
Definition: gemm_kernels_q4k_q8k.c:295
attention_forward_decode_head_major_gqa_flash
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)
Definition: attention_kernels.c:1467
quantize_row_q8_k
void quantize_row_q8_k(const float *x, void *y, int k)
Definition: gemm_kernels_q4k_q8k.c:107
rmsnorm_forward
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)
Definition: rmsnorm_kernels.c:50
token
const char * token
Definition: tokenizer.h:306
vocab_size
int vocab_size
Definition: true_bpe.h:185
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)
Definition: v6.5/test_generated/generic_api_test.c:2343
ck_model_get_config
const CKModelConfig * ck_model_get_config(void)
Definition: v6.5/test_generated/generic_api_test.c:8578
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)
Definition: v6.5/test_generated/generic_api_test.c:2999
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)
Definition: v6.5/test_generated/generic_api_test.c:6602
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)
Definition: v6.5/test_generated/generic_api_test.c:5419
model_model_allocate
int model_model_allocate(MODELModel *model)
Definition: v6.5/test_generated/generic_api_test.c:88
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)
Definition: v6.5/test_generated/generic_api_test.c:4912
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)
Definition: v6.5/test_generated/generic_api_test.c:4405
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)
Definition: v6.5/test_generated/generic_api_test.c:6940
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)
Definition: v6.5/test_generated/generic_api_test.c:1523
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)
Definition: v6.5/test_generated/generic_api_test.c:8123
MagicHeader
MagicHeader
Definition: v6.5/test_generated/generic_api_test.c:80
model_forward_prefill_impl
static void model_forward_prefill_impl(MODELModel *model, const int *tokens, int num_tokens)
Definition: v6.5/test_generated/generic_api_test.c:4148
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)
Definition: v6.5/test_generated/generic_api_test.c:3819
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)
Definition: v6.5/test_generated/generic_api_test.c:3163
ck_model_create
void * ck_model_create(void)
Definition: v6.5/test_generated/generic_api_test.c:8582
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)
Definition: v6.5/test_generated/generic_api_test.c:5757
g_model_config
static CKModelConfig g_model_config
Definition: v6.5/test_generated/generic_api_test.c:8562
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)
Definition: v6.5/test_generated/generic_api_test.c:1359
model_decode
void model_decode(MODELModel *model, const int *token, int token_index)
Definition: v6.5/test_generated/generic_api_test.c:8550
ck_model_precompute_rope
void ck_model_precompute_rope(void *model)
Definition: v6.5/test_generated/generic_api_test.c:8598
model_precompute_rope
void model_precompute_rope(MODELModel *model)
Definition: v6.5/test_generated/generic_api_test.c:186
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)
Definition: v6.5/test_generated/generic_api_test.c:867
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)
Definition: v6.5/test_generated/generic_api_test.c:6264
ck_model_forward
void ck_model_forward(void *model, const int *tokens, int num_tokens)
Definition: v6.5/test_generated/generic_api_test.c:8602
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)
Definition: v6.5/test_generated/generic_api_test.c:2179
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)
Definition: v6.5/test_generated/generic_api_test.c:1851
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)
Definition: v6.5/test_generated/generic_api_test.c:7109
model_forward
void model_forward(MODELModel *model, const int *tokens, int num_tokens)
Definition: v6.5/test_generated/generic_api_test.c:8541
ck_model_verify_canaries
int ck_model_verify_canaries(void *model)
Definition: v6.5/test_generated/generic_api_test.c:8615
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)
Definition: v6.5/test_generated/generic_api_test.c:3655
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)
Definition: v6.5/test_generated/generic_api_test.c:2835
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)
Definition: v6.5/test_generated/generic_api_test.c:5081
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)
Definition: v6.5/test_generated/generic_api_test.c:4574
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)
Definition: v6.5/test_generated/generic_api_test.c:2507
model_verify_canaries
int model_verify_canaries(MODELModel *model)
Definition: v6.5/test_generated/generic_api_test.c:152
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)
Definition: v6.5/test_generated/generic_api_test.c:7785
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)
Definition: v6.5/test_generated/generic_api_test.c:703
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)
Definition: v6.5/test_generated/generic_api_test.c:7616
ck_model_free
void ck_model_free(void *model)
Definition: v6.5/test_generated/generic_api_test.c:8592
ck_model_get_base
void * ck_model_get_base(void *model)
Definition: v6.5/test_generated/generic_api_test.c:8619
ck_model_decode
void ck_model_decode(void *model, const int *token, int token_index)
Definition: v6.5/test_generated/generic_api_test.c:8606
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)
Definition: v6.5/test_generated/generic_api_test.c:6771
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)
Definition: v6.5/test_generated/generic_api_test.c:43
model_model_free
void model_model_free(MODELModel *model)
Definition: v6.5/test_generated/generic_api_test.c:141
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)
Definition: v6.5/test_generated/generic_api_test.c:3491
ck_model_get_total_bytes
size_t ck_model_get_total_bytes(void *model)
Definition: v6.5/test_generated/generic_api_test.c:8623
ck_model_get_logits
float * ck_model_get_logits(void *model)
Definition: v6.5/test_generated/generic_api_test.c:8610
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)
Definition: v6.5/test_generated/generic_api_test.c:2015
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)
Definition: v6.5/test_generated/generic_api_test.c:5588
__attribute__
struct __attribute__((packed))
Definition: v6.5/test_generated/generic_api_test.c:67
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)
Definition: v6.5/test_generated/generic_api_test.c:3983
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)
Definition: v6.5/test_generated/generic_api_test.c:7954
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)
Definition: v6.5/test_generated/generic_api_test.c:1031
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)
Definition: v6.5/test_generated/generic_api_test.c:7447
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)
Definition: v6.5/test_generated/generic_api_test.c:5250
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)
Definition: v6.5/test_generated/generic_api_test.c:2671
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)
Definition: v6.5/test_generated/generic_api_test.c:6095
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)
Definition: v6.5/test_generated/generic_api_test.c:6433
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)
Definition: v6.5/test_generated/generic_api_test.c:7278
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)
Definition: v6.5/test_generated/generic_api_test.c:211
model_align_elems
static int model_align_elems(int elems, int elem_bytes, int align_bytes)
Definition: v6.5/test_generated/generic_api_test.c:176
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)
Definition: v6.5/test_generated/generic_api_test.c:375
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)
Definition: v6.5/test_generated/generic_api_test.c:8292
model_decode_token
static void model_decode_token(MODELModel *model, const int *token, int token_index)
Definition: v6.5/test_generated/generic_api_test.c:8462
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)
Definition: v6.5/test_generated/generic_api_test.c:3327
_Static_assert
_Static_assert(sizeof(MagicHeader)==64, "MagicHeader must be 64 bytes")
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)
Definition: v6.5/test_generated/generic_api_test.c:4743
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)
Definition: v6.5/test_generated/generic_api_test.c:539
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)
Definition: v6.5/test_generated/generic_api_test.c:1687
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)
Definition: v6.5/test_generated/generic_api_test.c:1195
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)
Definition: v6.5/test_generated/generic_api_test.c:5926
ck-kernel-inference.h
AUTO-GENERATED: qwen2_0.5b_decode Memory Layout.
MODEL_VOCAB_SIZE
#define MODEL_VOCAB_SIZE
Definition: v6_simple.c:30
MODEL_INTERMEDIATE
#define MODEL_INTERMEDIATE
Definition: v6_simple.c:36
MODEL_NUM_LAYERS
#define MODEL_NUM_LAYERS
Definition: v6_simple.c:25
MODEL_HEAD_DIM
#define MODEL_HEAD_DIM
Definition: v6_simple.c:28
MODEL_NUM_KV_HEADS
#define MODEL_NUM_KV_HEADS
Definition: v6_simple.c:27
MODEL_EMBED_DIM
#define MODEL_EMBED_DIM
Definition: v6_simple.c:24
MODEL_MAX_SEQ_LEN
#define MODEL_MAX_SEQ_LEN
Definition: v6_simple.c:31
MODEL_NUM_HEADS
#define MODEL_NUM_HEADS
Definition: v6_simple.c:26