← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
ckernel_codegen.c
Go to the documentation of this file.
1/**
2 * @file ckernel_codegen.c
3 *
4 * ===========================================================================
5 * LEGACY CODE - NOT USED IN v6.6
6 * ===========================================================================
7 *
8 * This file contains the older C-side kernel/codegen orchestration path.
9 * It is kept for reference and potential fallback use but is not part of the
10 * v6.6 generated-model inference flow.
11 *
12 * v6.6 Architecture:
13 * - IR Lower 1/2/3 drives execution planning
14 * - version/v6.6/scripts/codegen_v6_6.py emits model_v6_6.c directly
15 * - Kernel dispatch is selected from kernel_maps + lowered call IR
16 *
17 * To remove completely:
18 * 1. Delete this file
19 * 2. Remove ckernel_codegen.* references from legacy build paths
20 *
21 * Last used: v6.5
22 * Deprecated: v6.6 (2026-02)
23 * ===========================================================================
24 */
25
26#include "ckernel_codegen.h"
27#include "ckernel_registry.h"
28#include "ckernel_kernel_specs.h"
29
30#include <errno.h>
31#include <stdio.h>
32#include <stdlib.h>
33#include <string.h>
34
35static const char *op_name(CKOpType op)
36{
37 switch (op) {
38 case CK_OP_RMSNORM: return "RMSNORM";
39 case CK_OP_LINEAR_QKV: return "LINEAR_QKV";
40 case CK_OP_ATTENTION: return "ATTENTION";
41 case CK_OP_ADD: return "ADD";
42 case CK_OP_LINEAR: return "LINEAR";
43 case CK_OP_SPLIT: return "SPLIT";
44 case CK_OP_SWIGLU: return "SWIGLU";
45 case CK_OP_RMSNORM_BWD: return "RMSNORM_BWD";
46 case CK_OP_LINEAR_QKV_BWD: return "LINEAR_QKV_BWD";
47 case CK_OP_ATTENTION_BWD: return "ATTENTION_BWD";
48 case CK_OP_ADD_BWD: return "ADD_BWD";
49 case CK_OP_LINEAR_BWD: return "LINEAR_BWD";
50 case CK_OP_SPLIT_BWD: return "SPLIT_BWD";
51 case CK_OP_SWIGLU_BWD: return "SWIGLU_BWD";
52 default: return "UNKNOWN";
53 }
54}
55
56static const CKBufferSpec *ck_find_buffer_spec(const char *name)
57{
58 if (!name) {
59 return NULL;
60 }
61 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
62 if (strcmp(ck_decoder_buffers[i].name, name) == 0) {
63 return &ck_decoder_buffers[i];
64 }
65 }
66 return NULL;
67}
68
69static const CKKernelSpec *ck_find_kernel_spec(const char *name)
70{
71 if (!name) {
72 return NULL;
73 }
74 for (size_t i = 0; i < ck_kernel_spec_count; ++i) {
75 if (strcmp(ck_kernel_specs[i].name, name) == 0) {
76 return &ck_kernel_specs[i];
77 }
78 }
79 return NULL;
80}
81
82static int ck_buffer_should_alloc(const CKBufferSpec *spec)
83{
84 return spec && spec->role != CK_ROLE_INPUT;
85}
86
87static int ck_buffer_uses_weight_dtype(const CKBufferSpec *spec)
88{
89 if (!spec || spec->role != CK_ROLE_WEIGHT || !spec->name) {
90 return 0;
91 }
92
93 /* Weight-only quantization targets: big GEMM weights and tied embeddings.
94 * Small vectors (norm scales, biases) remain fp32 for simplicity. */
95 return (strcmp(spec->name, "token_emb") == 0) ||
96 (strcmp(spec->name, "wq") == 0) ||
97 (strcmp(spec->name, "wk") == 0) ||
98 (strcmp(spec->name, "wv") == 0) ||
99 (strcmp(spec->name, "wo") == 0) ||
100 (strcmp(spec->name, "w1") == 0) ||
101 (strcmp(spec->name, "w2") == 0) ||
102 (strcmp(spec->name, "lm_head_weight") == 0);
103}
104
105static const char *ck_weight_dtype_expr(const CKBufferSpec *spec)
106{
107 if (!spec || spec->role != CK_ROLE_WEIGHT || !spec->name) {
108 return NULL;
109 }
110 if (spec->scope == CK_SCOPE_GLOBAL) {
111 if (strcmp(spec->name, "token_emb") == 0) {
112 return "m->token_emb_dtype";
113 }
114 if (strcmp(spec->name, "lm_head_weight") == 0) {
115 return "m->lm_head_weight_dtype";
116 }
117 return NULL;
118 }
119 if (spec->scope == CK_SCOPE_LAYER) {
120 if (strcmp(spec->name, "wq") == 0) return "L->wq_dtype";
121 if (strcmp(spec->name, "wk") == 0) return "L->wk_dtype";
122 if (strcmp(spec->name, "wv") == 0) return "L->wv_dtype";
123 if (strcmp(spec->name, "wo") == 0) return "L->wo_dtype";
124 if (strcmp(spec->name, "w1") == 0) return "L->w1_dtype";
125 if (strcmp(spec->name, "w2") == 0) return "L->w2_dtype";
126 return NULL;
127 }
128 return NULL;
129}
130
131static void emit_offset_field(FILE *out, const char *name)
132{
133 fprintf(out, " size_t %s_offset;\n", name);
134}
135
136static void emit_layer_offsets_struct(FILE *out)
137{
138 fprintf(out, "typedef struct {\n");
139 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
140 const CKBufferSpec *spec = &ck_decoder_buffers[i];
141 if (spec->scope != CK_SCOPE_LAYER) {
142 continue;
143 }
144 if (!ck_buffer_should_alloc(spec)) {
145 continue;
146 }
147 emit_offset_field(out, spec->name);
148 }
149 fprintf(out,
150 " CKDataType wq_dtype;\n"
151 " CKDataType wk_dtype;\n"
152 " CKDataType wv_dtype;\n"
153 " CKDataType wo_dtype;\n"
154 " CKDataType w1_dtype;\n"
155 " CKDataType w2_dtype;\n");
156 fprintf(out, "} LayerOffsets;\n\n");
157}
158
159static void emit_global_offset_fields(FILE *out)
160{
161 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
162 const CKBufferSpec *spec = &ck_decoder_buffers[i];
163 if (spec->scope != CK_SCOPE_GLOBAL) {
164 continue;
165 }
166 if (!ck_buffer_should_alloc(spec)) {
167 continue;
168 }
169 emit_offset_field(out, spec->name);
170 }
171}
172
173static void emit_model_struct(FILE *out)
174{
175 fprintf(out,
176 "typedef LayerOffsets TrulyOptimalLayer;\n\n"
177 "typedef struct {\n"
178 " char magic[8];\n"
179 " uint32_t version;\n"
180 " uint32_t model_type;\n"
181 "\n"
182 " int num_layers;\n"
183 " int vocab_size;\n"
184 " int embed_dim;\n"
185 " int context_window;\n"
186 " int intermediate_size;\n"
187 "\n"
188 " size_t aligned_embed_dim;\n"
189 " size_t aligned_head_dim;\n"
190 " size_t aligned_attn_context_window;\n"
191 "\n"
192 " int num_cores;\n"
193 " int tokens_per_core;\n"
194 " int num_attention_heads;\n"
195 " int num_kv_heads;\n"
196 " int head_dim;\n"
197 " float rms_norm_eps;\n"
198 " float rope_theta;\n"
199 "\n"
200 " uint8_t *memory_base;\n"
201 " size_t total_bytes;\n"
202 " size_t elem_bytes;\n"
203 " CKDataType weight_dtype;\n"
204 " CKDataType token_emb_dtype;\n"
205 " CKDataType pos_emb_dtype;\n"
206 " CKDataType lm_head_weight_dtype;\n"
207 " bool weights_mixed;\n"
208 " bool weights_quantized;\n"
209 " size_t layer_stride;\n"
210 "\n"
211 " size_t layers_start_offset;\n");
212
213 emit_global_offset_fields(out);
214
215 fprintf(out,
216 "\n"
217 " TrulyOptimalLayer *layers;\n"
218 "\n"
219 " GradientStorage gradients;\n"
220 " bool training_enabled;\n"
221 " float learning_rate;\n"
222 " int lr_warmup_steps;\n"
223 " float lr_warmup_init;\n"
224 " float grad_clip;\n"
225 " size_t training_cache_samples;\n"
226 " int active_tokens;\n"
227 " TaskType task_type;\n"
228 " OptimizerType optimizer;\n"
229 " uint64_t optimizer_step;\n"
230 " float adam_beta1;\n"
231 " float adam_beta2;\n"
232 " float adam_eps;\n"
233 " float weight_decay;\n"
234 " bool ema_enabled;\n"
235 " float ema_decay;\n"
236 " bool optimizer_state_initialized;\n"
237 "\n"
238 " bool seq_cls_enabled;\n"
239 " int seq_cls_num_classes;\n"
240 " int seq_cls_pooling;\n"
241 " size_t seq_cls_weight_offset;\n"
242 " size_t seq_cls_bias_offset;\n"
243 "\n"
244 " bool kv_cache_enabled;\n"
245 " int kv_cache_capacity;\n"
246 " int kv_cache_tokens;\n"
247 "\n"
248 " long *training_data_buffer;\n"
249 " long num_training_tokens;\n"
250 "\n"
251 " uint8_t checksum[32];\n"
252 " uint8_t reserved[32];\n"
253 "} TransformerModel;\n\n");
254}
255
256static void emit_dim_expr(FILE *out, CKDimKind dim)
257{
258 switch (dim) {
259 case CK_DIM_TOKENS: fprintf(out, "(size_t)m->context_window"); break;
260 case CK_DIM_EMBED: fprintf(out, "(size_t)m->embed_dim"); break;
261 case CK_DIM_ALIGNED_EMBED: fprintf(out, "m->aligned_embed_dim"); break;
262 case CK_DIM_HEAD_DIM: fprintf(out, "(size_t)m->head_dim"); break;
263 case CK_DIM_ALIGNED_HEAD: fprintf(out, "m->aligned_head_dim"); break;
264 case CK_DIM_NUM_HEADS: fprintf(out, "(size_t)m->num_attention_heads"); break;
265 case CK_DIM_NUM_KV_HEADS: fprintf(out, "(size_t)m->num_kv_heads"); break;
266 case CK_DIM_ALIGNED_CTX: fprintf(out, "m->aligned_attn_context_window"); break;
267 case CK_DIM_INTERMEDIATE: fprintf(out, "(size_t)m->intermediate_size"); break;
268 case CK_DIM_ALIGNED_INTERMEDIATE:fprintf(out, "aligned_intermediate_dim"); break;
269 case CK_DIM_VOCAB: fprintf(out, "(size_t)m->vocab_size"); break;
270 case CK_DIM_END: fprintf(out, "0"); break;
271 }
272}
273
274static void emit_shape_expr(FILE *out, const CKDimToken *shape)
275{
276 int first = 1;
277 for (int i = 0; i < 4; ++i) {
278 if (shape[i].dim == CK_DIM_END) {
279 break;
280 }
281 if (!first) {
282 fprintf(out, " * ");
283 }
284 fprintf(out, "(");
285 emit_dim_expr(out, shape[i].dim);
286 if (shape[i].mult != 1) {
287 fprintf(out, " * %d", shape[i].mult);
288 }
289 if (shape[i].div != 1) {
290 fprintf(out, " / %d", shape[i].div);
291 }
292 fprintf(out, ")");
293 first = 0;
294 }
295 if (first) {
296 fprintf(out, "0");
297 }
298}
299
300static void emit_bump_bytes_assignment(FILE *out,
301 const char *indent,
302 const char *struct_prefix,
303 const char *name,
304 const CKDimToken *shape)
305{
306 fprintf(out, "%s%s%s_offset = bump_bytes(&off, (", indent, struct_prefix, name);
307 emit_shape_expr(out, shape);
308 fprintf(out, ") * elem_bytes, CACHELINE_BYTES);\n");
309}
310
311static void emit_bump_bytes_assignment_weight_dtype(FILE *out,
312 const char *indent,
313 const char *struct_prefix,
314 const char *name,
315 const CKDimToken *shape,
316 const char *dtype_expr)
317{
318 fprintf(out, "%s%s%s_offset = bump_bytes(&off, ck_dtype_row_bytes(%s, (",
319 indent, struct_prefix, name, dtype_expr);
320 emit_shape_expr(out, shape);
321 fprintf(out, ")), CACHELINE_BYTES);\n");
322}
323
324static void emit_training_conditional_assignment(FILE *out,
325 const char *indent,
326 const char *struct_prefix,
327 const char *name,
328 const CKDimToken *shape)
329{
330 /* Allocate gradient buffers only if training is enabled at init time */
331 fprintf(out, "%s%s%s_offset = m->training_enabled ? bump_bytes(&off, (", indent, struct_prefix, name);
332 emit_shape_expr(out, shape);
333 fprintf(out, ") * elem_bytes, CACHELINE_BYTES) : 0;\n");
334}
335
336static void emit_global_allocations(FILE *out)
337{
338 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
339 const CKBufferSpec *spec = &ck_decoder_buffers[i];
340 if (spec->scope != CK_SCOPE_GLOBAL) {
341 continue;
342 }
343 if (!ck_buffer_should_alloc(spec)) {
344 continue;
345 }
346 if (spec->role == CK_ROLE_GRAD) {
347 emit_training_conditional_assignment(out, " ", "m->", spec->name, spec->shape);
348 continue;
349 }
350 if (spec->alias_of) {
351 const CKBufferSpec *alias = ck_find_buffer_spec(spec->alias_of);
352 if (alias && alias->scope == CK_SCOPE_GLOBAL) {
353 fprintf(out, " m->%s_offset = m->%s_offset;\n", spec->name, spec->alias_of);
354 }
355 continue;
356 }
357 if (spec->condition && strcmp(spec->condition, "rope_theta") == 0) {
358 fprintf(out, " if (m->rope_theta > 0.0f) {\n");
359 fprintf(out, " m->%s_offset = bump_bytes(&off, (", spec->name);
360 emit_shape_expr(out, spec->shape);
361 fprintf(out, ") * elem_bytes, CACHELINE_BYTES);\n");
362 fprintf(out, " } else {\n");
363 fprintf(out, " m->%s_offset = 0;\n", spec->name);
364 fprintf(out, " }\n");
365 continue;
366 }
367 if (spec->condition && strcmp(spec->condition, "training_enabled") == 0) {
368 fprintf(out, " if (m->training_enabled) {\n");
369 fprintf(out, " m->%s_offset = bump_bytes(&off, (", spec->name);
370 emit_shape_expr(out, spec->shape);
371 fprintf(out, ") * elem_bytes, CACHELINE_BYTES);\n");
372 fprintf(out, " } else {\n");
373 fprintf(out, " m->%s_offset = 0;\n", spec->name);
374 fprintf(out, " }\n");
375 continue;
376 }
377 if (ck_buffer_uses_weight_dtype(spec)) {
378 const char *dtype_expr = ck_weight_dtype_expr(spec);
379 if (dtype_expr) {
380 emit_bump_bytes_assignment_weight_dtype(out, " ", "m->", spec->name, spec->shape, dtype_expr);
381 continue;
382 }
383 }
384 emit_bump_bytes_assignment(out, " ", "m->", spec->name, spec->shape);
385 }
386}
387
388static void emit_layer_allocations(FILE *out)
389{
390 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
391 const CKBufferSpec *spec = &ck_decoder_buffers[i];
392 if (spec->scope != CK_SCOPE_LAYER) {
393 continue;
394 }
395 if (!ck_buffer_should_alloc(spec)) {
396 continue;
397 }
398 if (spec->role == CK_ROLE_GRAD) {
399 emit_training_conditional_assignment(out, " ", "L->", spec->name, spec->shape);
400 continue;
401 }
402 if (spec->condition && strcmp(spec->condition, "training_enabled") == 0) {
403 fprintf(out, " if (m->training_enabled) {\n");
404 fprintf(out, " L->%s_offset = bump_bytes(&off, (", spec->name);
405 emit_shape_expr(out, spec->shape);
406 fprintf(out, ") * elem_bytes, CACHELINE_BYTES);\n");
407 fprintf(out, " } else {\n");
408 fprintf(out, " L->%s_offset = 0;\n", spec->name);
409 fprintf(out, " }\n");
410 continue;
411 }
412 if (ck_buffer_uses_weight_dtype(spec)) {
413 const char *dtype_expr = ck_weight_dtype_expr(spec);
414 if (dtype_expr) {
415 emit_bump_bytes_assignment_weight_dtype(out, " ", "L->", spec->name, spec->shape, dtype_expr);
416 continue;
417 }
418 }
419 emit_bump_bytes_assignment(out, " ", "L->", spec->name, spec->shape);
420 }
421}
422
423static void emit_global_aliases_to_layer(FILE *out)
424{
425 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
426 const CKBufferSpec *spec = &ck_decoder_buffers[i];
427 if (spec->scope != CK_SCOPE_GLOBAL || !spec->alias_of) {
428 continue;
429 }
430 const CKBufferSpec *alias = ck_find_buffer_spec(spec->alias_of);
431 if (!alias || alias->scope != CK_SCOPE_LAYER) {
432 continue;
433 }
434 fprintf(out,
435 " if (m->num_layers > 0) {\n"
436 " m->%s_offset = m->layers[m->num_layers - 1].%s_offset;\n"
437 " } else {\n"
438 " m->%s_offset = 0;\n"
439 " }\n",
440 spec->name, spec->alias_of, spec->name);
441 }
442}
443
444static void emit_zero_grad(FILE *out)
445{
446 fprintf(out,
447 "static void zero_grad(TransformerModel *m)\n"
448 "{\n"
449 " if (!m || !m->training_enabled) return;\n"
450 " uint8_t *base = m->memory_base;\n"
451 " size_t aligned_intermediate_dim = align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n");
452
453 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
454 const CKBufferSpec *spec = &ck_decoder_buffers[i];
455 if (spec->role != CK_ROLE_GRAD || spec->scope != CK_SCOPE_GLOBAL) {
456 continue;
457 }
458 fprintf(out, " if (m->%s_offset) {\n", spec->name);
459 fprintf(out, " memset(base + m->%s_offset, 0, (", spec->name);
460 emit_shape_expr(out, spec->shape);
461 fprintf(out, ") * m->elem_bytes);\n");
462 fprintf(out, " }\n");
463 }
464
465 fprintf(out,
466 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
467 " TrulyOptimalLayer *L = &m->layers[layer];\n");
468 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
469 const CKBufferSpec *spec = &ck_decoder_buffers[i];
470 if (spec->role != CK_ROLE_GRAD || spec->scope != CK_SCOPE_LAYER) {
471 continue;
472 }
473 fprintf(out, " if (L->%s_offset) {\n", spec->name);
474 fprintf(out, " memset(base + L->%s_offset, 0, (", spec->name);
475 emit_shape_expr(out, spec->shape);
476 fprintf(out, ") * m->elem_bytes);\n");
477 fprintf(out, " }\n");
478 }
479 fprintf(out,
480 " }\n"
481 "}\n\n");
482}
483
484static void emit_sgd_update(FILE *out)
485{
486 fprintf(out,
487 "static void sgd_update(TransformerModel *m, float lr)\n"
488 "{\n"
489 " if (!m || !m->training_enabled || lr == 0.0f) return;\n"
490 " uint8_t *base = m->memory_base;\n"
491 " size_t aligned_intermediate_dim = align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n");
492
493 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
494 const CKBufferSpec *spec = &ck_decoder_buffers[i];
495 if (spec->role != CK_ROLE_WEIGHT || spec->scope != CK_SCOPE_GLOBAL) {
496 continue;
497 }
498 if (spec->alias_of) {
499 continue;
500 }
501 char grad_name[128];
502 snprintf(grad_name, sizeof(grad_name), "d_%s", spec->name);
503 const CKBufferSpec *grad = ck_find_buffer_spec(grad_name);
504 if (!grad || grad->scope != CK_SCOPE_GLOBAL) {
505 continue;
506 }
507 fprintf(out,
508 " if (m->%s_offset && m->%s_offset) {\n"
509 " float *w = ptr_f32(base, m->%s_offset);\n"
510 " float *g = ptr_f32(base, m->%s_offset);\n"
511 " size_t count = (",
512 spec->name, grad_name, spec->name, grad_name);
513 emit_shape_expr(out, spec->shape);
514 fprintf(out,
515 ");\n"
516 " for (size_t i = 0; i < count; ++i) {\n"
517 " w[i] -= lr * g[i];\n"
518 " }\n"
519 " }\n");
520 }
521
522 fprintf(out,
523 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
524 " TrulyOptimalLayer *L = &m->layers[layer];\n");
525 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
526 const CKBufferSpec *spec = &ck_decoder_buffers[i];
527 if (spec->role != CK_ROLE_WEIGHT || spec->scope != CK_SCOPE_LAYER) {
528 continue;
529 }
530 char grad_name[128];
531 snprintf(grad_name, sizeof(grad_name), "d_%s", spec->name);
532 const CKBufferSpec *grad = ck_find_buffer_spec(grad_name);
533 if (!grad || grad->scope != CK_SCOPE_LAYER) {
534 continue;
535 }
536 fprintf(out,
537 " if (L->%s_offset && L->%s_offset) {\n"
538 " float *w = ptr_f32(base, L->%s_offset);\n"
539 " float *g = ptr_f32(base, L->%s_offset);\n"
540 " size_t count = (",
541 spec->name, grad_name, spec->name, grad_name);
542 emit_shape_expr(out, spec->shape);
543 fprintf(out,
544 ");\n"
545 " for (size_t i = 0; i < count; ++i) {\n"
546 " w[i] -= lr * g[i];\n"
547 " }\n"
548 " }\n");
549 }
550 fprintf(out,
551 " }\n"
552 "}\n\n");
553}
554
555static int emit_unique_source(FILE *f,
556 const char *path,
557 const char **seen,
558 size_t *seen_count,
559 size_t seen_cap)
560{
561 if (!path || !path[0]) {
562 return 0;
563 }
564 for (size_t i = 0; i < *seen_count; ++i) {
565 if (strcmp(seen[i], path) == 0) {
566 return 0;
567 }
568 }
569 if (*seen_count >= seen_cap) {
570 return -1;
571 }
572 fputs(path, f);
573 fputc('\n', f);
574 seen[*seen_count] = path;
575 (*seen_count)++;
576 return 0;
577}
578
579static int ck_plan_step_enabled(const CKPlanStep *step, const CKIRGraph *cfg)
580{
581 if (!step || !step->condition || !cfg) {
582 return 1;
583 }
584 if (strcmp(step->condition, "rope_theta") == 0) {
585 return cfg->config.rope_theta > 0.0f;
586 }
587 if (strcmp(step->condition, "rope_theta>0") == 0) {
588 return cfg->config.rope_theta > 0.0f;
589 }
590 return 1;
591}
592
593static int emit_plan_sources(FILE *f,
594 const CKPlanStep *plan,
595 size_t plan_count,
596 const CKIRGraph *cfg,
597 const char **seen,
598 size_t *seen_count,
599 size_t seen_cap)
600{
601 for (size_t i = 0; i < plan_count; ++i) {
602 const CKPlanStep *step = &plan[i];
603 if (!ck_plan_step_enabled(step, cfg)) {
604 continue;
605 }
606 const CKKernelSpec *spec = ck_find_kernel_spec(step->kernel);
607 if (!spec) {
608 continue;
609 }
610 for (size_t s = 0; s < CKERNEL_MAX_KERNEL_SOURCES; ++s) {
611 const char *src = spec->sources[s];
612 if (!src) {
613 continue;
614 }
615 if (emit_unique_source(f, src, seen, seen_count, seen_cap) != 0) {
616 return -1;
617 }
618 }
619 }
620 return 0;
621}
622
623static const char *ck_first_layer_buffer_name(void)
624{
625 for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
626 const CKBufferSpec *spec = &ck_decoder_buffers[i];
627 if (spec->scope != CK_SCOPE_LAYER) {
628 continue;
629 }
630 if (!ck_buffer_should_alloc(spec)) {
631 continue;
632 }
633 return spec->name;
634 }
635 return "ln1_gamma";
636}
637
638void ck_codegen_c_skeleton(const CKIRGraph *forward,
639 const CKIRGraph *backward,
640 FILE *out)
641{
642 if (!forward || !out) {
643 return;
644 }
645
646 fprintf(out,
647 "/* Auto-generated skeleton from CKIRGraph.\n"
648 " * This file sketches the structure of the forward and backward\n"
649 " * execution for a decoder-only transformer. It is NOT yet a\n"
650 " * complete, runnable implementation. You can use it as a\n"
651 " * starting point to wire buffers, kernel calls, and memory layout.\n"
652 " */\n\n");
653
654 fprintf(out, "#include \"ckernel_engine.h\"\n");
655 fprintf(out, "#include \"ckernel_model.h\"\n");
656 fprintf(out, "#include \"ckernel_alloc.h\"\n\n");
657
658 /* Forward function */
659 fprintf(out,
660 "void run_decoder_forward(TransformerModel *model /*, inputs, etc. */)\n"
661 "{\n"
662 " for (int layer = 0; layer < model->cfg.num_layers; ++layer) {\n"
663 " /* Forward pass for layer */\n");
664
665 int nodes_per_layer = 0;
666 if (forward->num_nodes > 0) {
667 int l0 = forward->nodes[0].id.layer;
668 for (int i = 0; i < forward->num_nodes; ++i) {
669 if (forward->nodes[i].id.layer != l0) {
670 break;
671 }
672 nodes_per_layer++;
673 }
674 }
675
676 if (nodes_per_layer <= 0) {
677 nodes_per_layer = forward->num_nodes;
678 }
679
680 fprintf(out, " /* This layer has %d IR nodes */\n", nodes_per_layer);
681
682 for (int i = 0; i < nodes_per_layer; ++i) {
683 const CKIRNode *n = &forward->nodes[i];
684 fprintf(out, " // L%%d: %s\n", op_name(n->op));
685 fprintf(out,
686 " // outputs: [");
687 for (int o = 0; o < n->n_outputs; ++o) {
688 if (o > 0) fprintf(out, ", ");
689 fprintf(out, "L%%d:N%d:%d", n->id.node, o);
690 }
691 fprintf(out, "]\n");
692 fprintf(out, " // inputs : [");
693 for (int j = 0; j < n->n_inputs; ++j) {
694 const CKInputRef *inp = &n->inputs[j];
695 if (j > 0) fprintf(out, ", ");
696 if (inp->producer.node == 0xFFFFu) {
697 fprintf(out, "IN");
698 } else {
699 fprintf(out, "L%%d:N%u:%u",
700 (unsigned)inp->producer.node,
701 (unsigned)inp->out_index);
702 }
703 }
704 fprintf(out, "]\n");
705 fprintf(out,
706 " // TODO: bind buffers/weights and call %s kernel here\n\n",
707 op_name(n->op));
708 }
709
710 fprintf(out,
711 " } /* end for layer */\n"
712 "}\n\n");
713
714 /* Backward skeleton */
715 if (backward && backward->nodes && backward->num_nodes > 0) {
716 fprintf(out,
717 "void run_decoder_backward(TransformerModel *model /*, grads, etc. */)\n"
718 "{\n"
719 " for (int layer = model->cfg.num_layers - 1; layer >= 0; --layer) {\n"
720 " /* Backward pass for layer */\n");
721
722 int bwd_per_layer = 0;
723 int l0 = backward->nodes[0].id.layer;
724 for (int i = 0; i < backward->num_nodes; ++i) {
725 if (backward->nodes[i].id.layer != l0) break;
726 bwd_per_layer++;
727 }
728 if (bwd_per_layer <= 0) bwd_per_layer = backward->num_nodes;
729
730 fprintf(out, " /* This layer has %d backward IR nodes */\n", bwd_per_layer);
731
732 for (int i = 0; i < bwd_per_layer; ++i) {
733 const CKIRNode *n = &backward->nodes[i];
734 fprintf(out, " // L%%d: %s\n", op_name(n->op));
735 fprintf(out,
736 " // TODO: wire gradient tensors and call %s kernel here\n\n",
737 op_name(n->op));
738 }
739
740 fprintf(out,
741 " } /* end for layer */\n"
742 "}\n\n");
743 }
744
745 fprintf(out,
746 "int main(int argc, char **argv)\n"
747 "{\n"
748 " (void)argc; (void)argv;\n"
749 " TransformerModel model = {0};\n"
750 " model.cfg.num_layers = %d;\n"
751 " model.cfg.hidden_size = %d;\n"
752 " model.cfg.intermediate_size = %d;\n"
753 " model.cfg.num_heads = %d;\n"
754 " model.cfg.num_kv_heads = %d;\n"
755 " model.cfg.vocab_size = %d;\n"
756 " model.cfg.context_window = %d;\n"
757 " model.cfg.rms_norm_eps = %.9g;\n"
758 " model.cfg.rope_theta = %.9g;\n"
759 " layout_transformer_from_ir(&model, NULL); /* TODO: pass IR if needed */\n"
760 " size_t bytes = model.total_bytes;\n"
761 " model.memory_base = (uint8_t *)ck_huge_alloc(bytes);\n"
762 " if (!model.memory_base) {\n"
763 " fprintf(stderr, \"Failed to allocate %%zu bytes for model\\n\", bytes);\n"
764 " return 1;\n"
765 " }\n"
766 " // TODO: load weights into model.memory_base based on offsets\n"
767 " run_decoder_forward(&model);\n"
768 " // TODO: run_decoder_backward(&model) when training\n"
769 " ck_huge_free(model.memory_base, bytes);\n"
770 " return 0;\n"
771 "}\n",
772 forward->config.num_layers,
773 forward->config.hidden_size,
774 forward->config.intermediate_size,
775 forward->config.num_heads,
776 forward->config.num_kv_heads,
777 forward->config.vocab_size,
778 forward->config.context_window,
779 forward->config.rms_norm_eps,
780 forward->config.rope_theta);
781}
782
783static int emit_runtime_preamble(FILE *out)
784{
785 fprintf(out,
786 "/* Auto-generated runtime from CKIRGraph.\n"
787 " * This file wires the existing C-Kernel-Engine kernels into a\n"
788 " * decoder-only transformer forward pass.\n"
789 " *\n"
790 " * Compile (scalar): gcc -O2 generated_model.c $(cat generated_model.c.kernels) -Iinclude -lm -o generated_model\n"
791 " * Compile (AVX-512): gcc -O3 -mavx512f -mfma generated_model.c $(cat generated_model.c.kernels) -Iinclude -lm -o generated_model\n"
792 " */\n\n");
793
794 fprintf(out,
795 "#define _GNU_SOURCE\n"
796 "#include <stddef.h>\n"
797 "#include <stdint.h>\n"
798 "#include <stdbool.h>\n"
799 "#include <stdio.h>\n"
800 "#include <stdlib.h>\n"
801 "#include <string.h>\n"
802 "#include <math.h>\n"
803 "#include <errno.h>\n"
804 "#include <sys/types.h>\n"
805 "#include <unistd.h>\n"
806 "#include \"ckernel_engine.h\"\n"
807 "#include \"ckernel_dtype.h\"\n"
808 "#include \"ckernel_orchestration.h\"\n"
809 "#include \"ckernel_alloc.h\"\n\n");
810
811 fprintf(out,
812 "#define CACHELINE_BYTES 64\n"
813 "static size_t align_up_bytes(size_t n, size_t align) {\n"
814 " if (align == 0) return n;\n"
815 " return (n + align - 1) & ~(align - 1);\n"
816 "}\n\n"
817 "static size_t align_up_elems(size_t elems, size_t elem_bytes, size_t align) {\n"
818 " size_t bytes = elems * elem_bytes;\n"
819 " bytes = align_up_bytes(bytes, align);\n"
820 " return bytes / elem_bytes;\n"
821 "}\n\n"
822 "static size_t bump_bytes(size_t *off, size_t bytes, size_t align) {\n"
823 " size_t start = align_up_bytes(*off, align);\n"
824 " *off = start + bytes;\n"
825 " return start;\n"
826 "}\n\n"
827 "static inline float *ptr_f32(uint8_t *base, size_t offset) {\n"
828 " return (float *)(base + offset);\n"
829 "}\n"
830 "static inline const float *cptr_f32(const uint8_t *base, size_t offset) {\n"
831 " return (const float *)(base + offset);\n"
832 "}\n\n");
833
834 fprintf(out,
835 "static inline uint8_t *ptr_u8(uint8_t *base, size_t offset) {\n"
836 " return base + offset;\n"
837 "}\n"
838 "static inline const void *cptr_void(const uint8_t *base, size_t offset) {\n"
839 " return (const void *)(base + offset);\n"
840 "}\n\n");
841
842 return 0;
843}
844
845static int emit_kernel_manifest(const CKIRGraph *forward, const char *runtime_path)
846{
847 if (!forward || !runtime_path) {
848 return -1;
849 }
850
851 const char *suffix = ".kernels";
852 size_t len = strlen(runtime_path) + strlen(suffix) + 1;
853 char *path = (char *)malloc(len);
854 if (!path) {
855 return -1;
856 }
857 snprintf(path, len, "%s%s", runtime_path, suffix);
858
859 FILE *f = fopen(path, "wb");
860 if (!f) {
861 fprintf(stderr, "ck_codegen_emit_runtime: failed to open %s: %s\n",
862 path, strerror(errno));
863 free(path);
864 return -1;
865 }
866
867 size_t seen_cap = ck_kernel_spec_count * CKERNEL_MAX_KERNEL_SOURCES + 8;
868 const char **seen = (const char **)calloc(seen_cap, sizeof(char *));
869 if (!seen) {
870 fclose(f);
871 free(path);
872 return -1;
873 }
874 size_t seen_count = 0;
875
876 emit_unique_source(f, "src/ckernel_alloc.c", seen, &seen_count, seen_cap);
877 emit_unique_source(f, "src/ckernel_strict.c", seen, &seen_count, seen_cap);
878 emit_unique_source(f, "src/cpu_features.c", seen, &seen_count, seen_cap);
879 emit_unique_source(f, "src/kernels/embedding_kernels.c", seen, &seen_count, seen_cap);
880 /* Quantized inference support (K-quants: Q4_K_M / Q4_K / Q6_K). */
881 emit_unique_source(f, "src/kernels/dequant_kernels.c", seen, &seen_count, seen_cap);
882 emit_unique_source(f, "src/kernels/gemm_kernels_q4k.c", seen, &seen_count, seen_cap);
883 emit_unique_source(f, "src/kernels/gemm_kernels_q4k_q8k.c", seen, &seen_count, seen_cap);
884 emit_unique_source(f, "src/kernels/gemm_kernels_q6k.c", seen, &seen_count, seen_cap);
885 /* Legacy quant support (Q4_0 / Q4_1 / Q5_0 / Q5_1 / Q8_0). */
886 emit_unique_source(f, "src/kernels/gemm_kernels_q4_0.c", seen, &seen_count, seen_cap);
887 emit_unique_source(f, "src/kernels/gemm_kernels_q4_1.c", seen, &seen_count, seen_cap);
888 emit_unique_source(f, "src/kernels/gemm_kernels_q5_0.c", seen, &seen_count, seen_cap);
889 emit_unique_source(f, "src/kernels/gemm_kernels_q5_1.c", seen, &seen_count, seen_cap);
890 emit_unique_source(f, "src/kernels/gemm_kernels_q8_0.c", seen, &seen_count, seen_cap);
891 /* SSE/AVX/AVX2/VNNI fallback implementations for quantized kernels. */
892 emit_unique_source(f, "src/kernels/gemm_kernels_q4k_sse.c", seen, &seen_count, seen_cap);
893 emit_unique_source(f, "src/kernels/gemm_kernels_q4k_avx.c", seen, &seen_count, seen_cap);
894 emit_unique_source(f, "src/kernels/gemm_kernels_q4k_q8k_avx2.c", seen, &seen_count, seen_cap);
895 emit_unique_source(f, "src/kernels/gemm_kernels_q4k_q8k_vnni.c", seen, &seen_count, seen_cap);
896 emit_unique_source(f, "src/kernels/gemm_kernels_q5_0_sse.c", seen, &seen_count, seen_cap);
897 emit_unique_source(f, "src/kernels/gemm_kernels_q5_0_sse_v2.c", seen, &seen_count, seen_cap);
898 emit_unique_source(f, "src/kernels/gemm_kernels_q6k_sse.c", seen, &seen_count, seen_cap);
899 emit_unique_source(f, "src/kernels/quantize_row_q8_k_sse.c", seen, &seen_count, seen_cap);
900 emit_unique_source(f, "src/kernels/quantize_row_q8_k_avx.c", seen, &seen_count, seen_cap);
901 emit_unique_source(f, "src/kernels/quantize_row_q8_k_avx2.c", seen, &seen_count, seen_cap);
902 emit_unique_source(f, "src/kernels/quantize_row_q8_k_avx512.c", seen, &seen_count, seen_cap);
903 emit_unique_source(f, "src/kernels/rope_kernels.c", seen, &seen_count, seen_cap);
904 emit_unique_source(f, "src/kernels/loss_kernels.c", seen, &seen_count, seen_cap);
905 emit_unique_source(f, "src/kernels/kv_cache_kernels.c", seen, &seen_count, seen_cap);
906 /* Fused kernels used by orchestration layer. */
907 emit_unique_source(f, "src/kernels/gemm_fused_kernels.c", seen, &seen_count, seen_cap);
908 emit_unique_source(f, "src/kernels/mlp_fused_decode.c", seen, &seen_count, seen_cap);
909 emit_unique_source(f, "src/kernels/gemm_microkernel.c", seen, &seen_count, seen_cap);
910 emit_unique_source(f, "src/kernels/attention_decode_fused.c", seen, &seen_count, seen_cap);
911 emit_unique_source(f, "src/kernels/attention_flash_true.c", seen, &seen_count, seen_cap);
912 if (emit_plan_sources(f,
913 ck_decoder_forward_plan,
914 ck_decoder_forward_plan_count,
915 forward,
916 seen,
917 &seen_count,
918 seen_cap) != 0) {
919 free(seen);
920 fclose(f);
921 free(path);
922 return -1;
923 }
924 if (emit_plan_sources(f,
925 ck_decoder_backward_plan,
926 ck_decoder_backward_plan_count,
927 forward,
928 seen,
929 &seen_count,
930 seen_cap) != 0) {
931 free(seen);
932 fclose(f);
933 free(path);
934 return -1;
935 }
936 free(seen);
937
938 fclose(f);
939 fprintf(stderr, "[ck_codegen] kernels manifest written to %s\n", path);
940 free(path);
941 return 0;
942}
943
944/* Library API functions emitted when mode == CK_EMIT_LIBRARY */
945static void emit_library_api(FILE *out, const CKIRGraph *forward)
946{
947 fprintf(out,
948 "\n/* ═══════════════════════════════════════════════════════════════\n"
949 " * C-Kernel-Engine Library API (for dlopen)\n"
950 " * ═══════════════════════════════════════════════════════════════ */\n\n"
951 "#ifdef _WIN32\n"
952 "#define CK_EXPORT __declspec(dllexport)\n"
953 "#else\n"
954 "#define CK_EXPORT __attribute__((visibility(\"default\")))\n"
955 "#endif\n\n"
956 "typedef struct {\n"
957 " int num_layers;\n"
958 " int hidden_size;\n"
959 " int intermediate_size;\n"
960 " int num_heads;\n"
961 " int num_kv_heads;\n"
962 " int vocab_size;\n"
963 " int context_window;\n"
964 " float rms_norm_eps;\n"
965 " float rope_theta;\n"
966 "} CKModelInfo;\n\n"
967 "static TransformerModel g_model = {0};\n"
968 "static int g_initialized = 0;\n"
969 "static int g_fuse_swiglu_decode = -2;\n"
970 "static int g_fuse_attn_decode = -2;\n\n"
971 "static int ck_fuse_swiglu_decode_mode(void)\n"
972 "{\n"
973 " if (g_fuse_swiglu_decode != -2) return g_fuse_swiglu_decode;\n"
974 " const char *env = getenv(\"CK_FUSE_SWIGLU_DECODE\");\n"
975 " if (!env || !env[0]) {\n"
976 " g_fuse_swiglu_decode = -1; /* auto */\n"
977 " return g_fuse_swiglu_decode;\n"
978 " }\n"
979 " if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || env[0] == 'F') {\n"
980 " g_fuse_swiglu_decode = 0;\n"
981 " } else {\n"
982 " g_fuse_swiglu_decode = 1;\n"
983 " }\n"
984 " return g_fuse_swiglu_decode;\n"
985 "}\n\n");
986
987 fprintf(out,
988 "static int ck_fuse_attn_decode_mode(void)\n"
989 "{\n"
990 " if (g_fuse_attn_decode != -2) return g_fuse_attn_decode;\n"
991 " const char *env = getenv(\"CK_FUSE_ATTN_DECODE\");\n"
992 " if (!env || !env[0]) {\n"
993 " g_fuse_attn_decode = -1; /* auto */\n"
994 " return g_fuse_attn_decode;\n"
995 " }\n"
996 " if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || env[0] == 'F') {\n"
997 " g_fuse_attn_decode = 0;\n"
998 " } else {\n"
999 " g_fuse_attn_decode = 1;\n"
1000 " }\n"
1001 " return g_fuse_attn_decode;\n"
1002 "}\n\n");
1003
1004 fprintf(out,
1005 "static int run_model_decode(TransformerModel *m, int32_t token)\n"
1006 "{\n"
1007 " if (!m || !m->memory_base) return -1;\n"
1008 " /* KV-cache decode is an inference-only fast path; training uses the full forward/backward graph. */\n"
1009 " if (m->training_enabled) return -4;\n"
1010 " if (!m->kv_cache_enabled) return -2;\n"
1011 "\n"
1012 " int cache_cap = m->kv_cache_capacity > 0 ? m->kv_cache_capacity : m->context_window;\n"
1013 " if (cache_cap > m->context_window) cache_cap = m->context_window;\n"
1014 " int t = m->kv_cache_tokens;\n"
1015 " if (t < 0) t = 0;\n"
1016 " if (t >= cache_cap) return -3;\n"
1017 "\n"
1018 " embed_token_at(m, token, t);\n"
1019 "\n"
1020 " uint8_t *base = m->memory_base;\n"
1021 " float *current = ptr_f32(base, m->embedded_input_offset);\n"
1022 " int aligned_intermediate_dim = (int)align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n"
1023 " int fuse_swiglu_mode = ck_fuse_swiglu_decode_mode();\n"
1024 " int use_fused_swiglu = 0;\n"
1025 " if (fuse_swiglu_mode > 0) {\n"
1026 " use_fused_swiglu = 1;\n"
1027 " } else if (fuse_swiglu_mode == 0) {\n"
1028 " use_fused_swiglu = 0;\n"
1029 " } else if (!ck_strict_parity_enabled()) {\n"
1030 " use_fused_swiglu = 1;\n"
1031 " }\n"
1032 " int fuse_attn_mode = ck_fuse_attn_decode_mode();\n"
1033 " int use_fused_attn = 0;\n"
1034 " if (fuse_attn_mode > 0) {\n"
1035 " use_fused_attn = 1;\n"
1036 " } else if (fuse_attn_mode == 0) {\n"
1037 " use_fused_attn = 0;\n"
1038 " } else if (!ck_strict_parity_enabled()) {\n"
1039 " use_fused_attn = 1;\n"
1040 " }\n"
1041 " if (m->weights_quantized) {\n"
1042 " use_fused_swiglu = 0;\n"
1043 " use_fused_attn = 0;\n"
1044 " }\n"
1045 "\n"
1046 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1047 " TrulyOptimalLayer *L = &m->layers[layer];\n"
1048 " if (!m->weights_mixed && m->weight_dtype == CK_DT_Q4_K) {\n"
1049 " CKLayerForwardParamsQ4K p = {0};\n"
1050 " p.tokens = cache_cap;\n"
1051 " p.embed_dim = m->embed_dim;\n"
1052 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
1053 " p.num_heads = m->num_attention_heads;\n"
1054 " p.num_kv_heads = m->num_kv_heads;\n"
1055 " p.head_dim = m->head_dim;\n"
1056 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
1057 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
1058 " p.intermediate_dim = m->intermediate_size;\n"
1059 " p.aligned_intermediate_dim = aligned_intermediate_dim;\n"
1060 " p.eps = m->rms_norm_eps;\n"
1061 " p.rope_pos_offset = t;\n"
1062 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
1063 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
1064 " p.input = current;\n"
1065 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
1066 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
1067 " p.wq = cptr_void(base, L->wq_offset);\n"
1068 " p.bq = cptr_f32(base, L->bq_offset);\n"
1069 " p.wk = cptr_void(base, L->wk_offset);\n"
1070 " p.bk = cptr_f32(base, L->bk_offset);\n"
1071 " p.wv = cptr_void(base, L->wv_offset);\n"
1072 " p.bv = cptr_f32(base, L->bv_offset);\n"
1073 " p.wo = cptr_void(base, L->wo_offset);\n"
1074 " p.bo = cptr_f32(base, L->bo_offset);\n"
1075 " p.w1 = cptr_void(base, L->w1_offset);\n"
1076 " p.b1 = cptr_f32(base, L->b1_offset);\n"
1077 " p.w2 = cptr_void(base, L->w2_offset);\n"
1078 " p.b2 = cptr_f32(base, L->b2_offset);\n"
1079 " p.ln1_out = ptr_f32(base, L->ln1_out_offset);\n"
1080 " p.ln1_rstd = ptr_f32(base, L->ln1_rstd_offset);\n"
1081 " p.k = ptr_f32(base, L->k_offset);\n"
1082 " p.v = ptr_f32(base, L->v_offset);\n"
1083 " p.proj_tmp = ptr_f32(base, L->proj_tmp_offset);\n"
1084 " p.proj_scratch = ptr_f32(base, L->proj_scratch_offset);\n"
1085 " p.residual1 = ptr_f32(base, L->residual1_offset);\n"
1086 " p.ln2_out = ptr_f32(base, L->ln2_out_offset);\n"
1087 " p.ln2_rstd = ptr_f32(base, L->ln2_rstd_offset);\n"
1088 " p.fc1_out = ptr_f32(base, L->fc1_out_offset);\n"
1089 " p.swiglu_out = ptr_f32(base, L->swiglu_out_offset);\n"
1090 " p.mlp_out = ptr_f32(base, L->mlp_out_offset);\n"
1091 " p.output = ptr_f32(base, L->output_offset);\n"
1092 "\n"
1093 " ck_layer_forward_rmsnorm_swiglu_decode_q4_k(&p, t, cache_cap);\n"
1094 " } else if (m->weights_quantized) {\n"
1095 " CKLayerForwardParamsQ4K p = {0};\n"
1096 " p.tokens = cache_cap;\n"
1097 " p.embed_dim = m->embed_dim;\n"
1098 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
1099 " p.num_heads = m->num_attention_heads;\n"
1100 " p.num_kv_heads = m->num_kv_heads;\n"
1101 " p.head_dim = m->head_dim;\n"
1102 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
1103 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
1104 " p.intermediate_dim = m->intermediate_size;\n"
1105 " p.aligned_intermediate_dim = aligned_intermediate_dim;\n"
1106 " p.eps = m->rms_norm_eps;\n"
1107 " p.rope_pos_offset = t;\n"
1108 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
1109 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
1110 " p.input = current;\n"
1111 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
1112 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
1113 " p.wq = cptr_void(base, L->wq_offset);\n"
1114 " p.bq = cptr_f32(base, L->bq_offset);\n"
1115 " p.wk = cptr_void(base, L->wk_offset);\n"
1116 " p.bk = cptr_f32(base, L->bk_offset);\n"
1117 " p.wv = cptr_void(base, L->wv_offset);\n"
1118 " p.bv = cptr_f32(base, L->bv_offset);\n"
1119 " p.wo = cptr_void(base, L->wo_offset);\n"
1120 " p.bo = cptr_f32(base, L->bo_offset);\n"
1121 " p.w1 = cptr_void(base, L->w1_offset);\n"
1122 " p.b1 = cptr_f32(base, L->b1_offset);\n"
1123 " p.w2 = cptr_void(base, L->w2_offset);\n"
1124 " p.b2 = cptr_f32(base, L->b2_offset);\n"
1125 " p.ln1_out = ptr_f32(base, L->ln1_out_offset);\n"
1126 " p.ln1_rstd = ptr_f32(base, L->ln1_rstd_offset);\n"
1127 " p.k = ptr_f32(base, L->k_offset);\n"
1128 " p.v = ptr_f32(base, L->v_offset);\n"
1129 " p.proj_tmp = ptr_f32(base, L->proj_tmp_offset);\n"
1130 " p.proj_scratch = ptr_f32(base, L->proj_scratch_offset);\n"
1131 " p.residual1 = ptr_f32(base, L->residual1_offset);\n"
1132 " p.ln2_out = ptr_f32(base, L->ln2_out_offset);\n"
1133 " p.ln2_rstd = ptr_f32(base, L->ln2_rstd_offset);\n"
1134 " p.fc1_out = ptr_f32(base, L->fc1_out_offset);\n"
1135 " p.swiglu_out = ptr_f32(base, L->swiglu_out_offset);\n"
1136 " p.mlp_out = ptr_f32(base, L->mlp_out_offset);\n"
1137 " p.output = ptr_f32(base, L->output_offset);\n"
1138 " p.wq_dtype = L->wq_dtype;\n"
1139 " p.wk_dtype = L->wk_dtype;\n"
1140 " p.wv_dtype = L->wv_dtype;\n"
1141 " p.wo_dtype = L->wo_dtype;\n"
1142 " p.w1_dtype = L->w1_dtype;\n"
1143 " p.w2_dtype = L->w2_dtype;\n"
1144 "\n"
1145 " ck_layer_forward_rmsnorm_swiglu_decode_quant(&p, t, cache_cap);\n"
1146 " } else {\n"
1147 " CKLayerForwardParams p = {0};\n"
1148 " p.tokens = cache_cap;\n"
1149 " p.embed_dim = m->embed_dim;\n"
1150 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
1151 " p.num_heads = m->num_attention_heads;\n"
1152 " p.num_kv_heads = m->num_kv_heads;\n"
1153 " p.head_dim = m->head_dim;\n"
1154 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
1155 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
1156 " p.intermediate_dim = m->intermediate_size;\n"
1157 " p.aligned_intermediate_dim = aligned_intermediate_dim;\n"
1158 " p.eps = m->rms_norm_eps;\n"
1159 " p.rope_pos_offset = t;\n"
1160 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
1161 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
1162 " p.input = current;\n"
1163 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
1164 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
1165 " p.wq = cptr_f32(base, L->wq_offset);\n"
1166 " p.bq = cptr_f32(base, L->bq_offset);\n"
1167 " p.wk = cptr_f32(base, L->wk_offset);\n"
1168 " p.bk = cptr_f32(base, L->bk_offset);\n"
1169 " p.wv = cptr_f32(base, L->wv_offset);\n"
1170 " p.bv = cptr_f32(base, L->bv_offset);\n"
1171 " p.wo = cptr_f32(base, L->wo_offset);\n"
1172 " p.bo = cptr_f32(base, L->bo_offset);\n"
1173 " p.w1 = cptr_f32(base, L->w1_offset);\n"
1174 " p.b1 = cptr_f32(base, L->b1_offset);\n"
1175 " p.w2 = cptr_f32(base, L->w2_offset);\n"
1176 " p.b2 = cptr_f32(base, L->b2_offset);\n"
1177 " p.ln1_out = ptr_f32(base, L->ln1_out_offset);\n"
1178 " p.ln1_rstd = ptr_f32(base, L->ln1_rstd_offset);\n"
1179 " p.k = ptr_f32(base, L->k_offset);\n"
1180 " p.v = ptr_f32(base, L->v_offset);\n"
1181 " p.proj_tmp = ptr_f32(base, L->proj_tmp_offset);\n"
1182 " p.proj_scratch = ptr_f32(base, L->proj_scratch_offset);\n"
1183 " p.residual1 = ptr_f32(base, L->residual1_offset);\n"
1184 " p.ln2_out = ptr_f32(base, L->ln2_out_offset);\n"
1185 " p.ln2_rstd = ptr_f32(base, L->ln2_rstd_offset);\n"
1186 " p.fc1_out = ptr_f32(base, L->fc1_out_offset);\n"
1187 " p.swiglu_out = ptr_f32(base, L->swiglu_out_offset);\n"
1188 " p.mlp_out = ptr_f32(base, L->mlp_out_offset);\n"
1189 " p.output = ptr_f32(base, L->output_offset);\n"
1190 "\n"
1191 " if (use_fused_attn) {\n"
1192 " if (use_fused_swiglu) {\n"
1193 " ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_mlp(&p, t, cache_cap);\n"
1194 " } else {\n"
1195 " ck_layer_forward_rmsnorm_swiglu_decode_fused_attn(&p, t, cache_cap);\n"
1196 " }\n"
1197 " } else if (use_fused_swiglu) {\n"
1198 " ck_layer_forward_rmsnorm_swiglu_decode_fused(&p, t, cache_cap);\n"
1199 " } else {\n"
1200 " ck_layer_forward_rmsnorm_swiglu_decode(&p, t, cache_cap);\n"
1201 " }\n"
1202 " }\n"
1203 " current = ptr_f32(base, L->output_offset);\n"
1204 " }\n"
1205 "\n"
1206 " int V = m->vocab_size;\n"
1207 " int D = m->embed_dim;\n"
1208 " int aligned_D = (int)m->aligned_embed_dim;\n"
1209 " float *final_in = current + (size_t)t * aligned_D;\n"
1210 " float *final_out = ptr_f32(base, m->final_output_offset) + (size_t)t * aligned_D;\n"
1211 " float *final_rstd = ptr_f32(base, m->final_ln_rstd_offset) + (size_t)t;\n"
1212 "\n"
1213 " rmsnorm_forward(final_in,\n"
1214 " cptr_f32(base, m->final_ln_weight_offset),\n"
1215 " final_out,\n"
1216 " final_rstd,\n"
1217 " 1,\n"
1218 " D,\n"
1219 " aligned_D,\n"
1220 " m->rms_norm_eps);\n"
1221 " if (V > 0) {\n"
1222 " float *logits_row = ptr_f32(base, m->logits_offset) + (size_t)t * (size_t)V;\n"
1223 " if (m->lm_head_weight_dtype == CK_DT_Q4_K) {\n"
1224 " gemm_nt_q4_k(final_out,\n"
1225 " cptr_void(base, m->lm_head_weight_offset),\n"
1226 " NULL,\n"
1227 " logits_row,\n"
1228 " 1,\n"
1229 " V,\n"
1230 " aligned_D);\n"
1231 " } else if (m->lm_head_weight_dtype == CK_DT_Q6_K) {\n"
1232 " gemm_nt_q6_k(final_out,\n"
1233 " cptr_void(base, m->lm_head_weight_offset),\n"
1234 " NULL,\n"
1235 " logits_row,\n"
1236 " 1,\n"
1237 " V,\n"
1238 " aligned_D);\n"
1239 " } else {\n"
1240 " lm_head_forward(final_out,\n"
1241 " cptr_f32(base, m->lm_head_weight_offset),\n"
1242 " logits_row,\n"
1243 " 1,\n"
1244 " V,\n"
1245 " D,\n"
1246 " aligned_D);\n"
1247 " }\n"
1248 " }\n"
1249 "\n"
1250 " m->kv_cache_tokens = t + 1;\n"
1251 " m->active_tokens = m->kv_cache_tokens;\n"
1252 " return 0;\n"
1253 "}\n\n");
1254
1255 /* ck_model_init */
1256 fprintf(out,
1257 "CK_EXPORT int ck_model_init(const char *weights_path)\n"
1258 "{\n"
1259 " if (g_initialized) return 0;\n"
1260 " memcpy(g_model.magic, \"BUMPWGT3\", 8);\n"
1261 " g_model.version = 3;\n"
1262 " g_model.model_type = 0;\n"
1263 " g_model.num_layers = %d;\n"
1264 " g_model.embed_dim = %d;\n"
1265 " g_model.intermediate_size = %d;\n"
1266 " g_model.num_attention_heads = %d;\n"
1267 " g_model.num_kv_heads = %d;\n"
1268 " g_model.vocab_size = %d;\n"
1269 " g_model.context_window = %d;\n"
1270 " g_model.rms_norm_eps = (float)%.9g;\n"
1271 " g_model.rope_theta = (float)%.9g;\n"
1272 " g_model.num_cores = 1;\n"
1273 " g_model.task_type = TASK_LM;\n"
1274 " g_model.weight_dtype = CK_DT_FP32;\n"
1275 " const char *wd = getenv(\"CK_WEIGHT_DTYPE\");\n"
1276 " if (wd) {\n"
1277 " if (strcmp(wd, \"q4_k\") == 0 || strcmp(wd, \"q4_k_m\") == 0 ||\n"
1278 " strcmp(wd, \"Q4_K\") == 0 || strcmp(wd, \"Q4_K_M\") == 0) {\n"
1279 " g_model.weight_dtype = CK_DT_Q4_K;\n"
1280 " } else if (strcmp(wd, \"q6_k\") == 0 || strcmp(wd, \"q6_k_l\") == 0 ||\n"
1281 " strcmp(wd, \"Q6_K\") == 0 || strcmp(wd, \"Q6_K_L\") == 0) {\n"
1282 " g_model.weight_dtype = CK_DT_Q6_K;\n"
1283 " }\n"
1284 " }\n"
1285 " init_weight_dtypes_uniform(&g_model, g_model.weight_dtype);\n"
1286 " refresh_weight_flags(&g_model);\n"
1287 " /* Check env var to pre-allocate gradient buffers for training */\n"
1288 " const char *train_env = getenv(\"CK_ENABLE_TRAINING\");\n"
1289 " if (train_env && (train_env[0] == '1' || train_env[0] == 'y' || train_env[0] == 'Y')) {\n"
1290 " g_model.training_enabled = true;\n"
1291 " g_model.learning_rate = 1e-4f;\n"
1292 " }\n"
1293 " if (weights_path) {\n"
1294 " int dtype_rc = load_weight_dtypes(weights_path, &g_model);\n"
1295 " if (dtype_rc < 0) {\n"
1296 " fprintf(stderr, \"Failed to read weight dtype table from %%s\\n\", weights_path);\n"
1297 " return -6;\n"
1298 " }\n"
1299 " }\n"
1300 " if (g_model.training_enabled && g_model.weights_quantized) {\n"
1301 " fprintf(stderr, \"Quantized weights are inference-only; disable training for this model\\n\");\n"
1302 " return -5;\n"
1303 " }\n"
1304 " g_model.kv_cache_enabled = false;\n"
1305 " g_model.kv_cache_capacity = g_model.context_window;\n"
1306 " g_model.kv_cache_tokens = 0;\n"
1307 " if (layout_model(&g_model) != 0) return -1;\n"
1308 " if (weights_path) {\n"
1309 " if (load_model_weights(weights_path, &g_model) != 0) return -2;\n"
1310 " }\n"
1311 " g_initialized = 1;\n"
1312 " return 0;\n"
1313 "}\n\n",
1314 forward->config.num_layers,
1315 forward->config.hidden_size,
1316 forward->config.intermediate_size,
1317 forward->config.num_heads,
1318 forward->config.num_kv_heads,
1319 forward->config.vocab_size,
1320 forward->config.context_window,
1321 forward->config.rms_norm_eps,
1322 forward->config.rope_theta);
1323
1324 /* ck_model_get_info */
1325 fprintf(out,
1326 "CK_EXPORT void ck_model_get_info(CKModelInfo *info)\n"
1327 "{\n"
1328 " if (!info) return;\n"
1329 " info->num_layers = g_model.num_layers;\n"
1330 " info->hidden_size = g_model.embed_dim;\n"
1331 " info->intermediate_size = g_model.intermediate_size;\n"
1332 " info->num_heads = g_model.num_attention_heads;\n"
1333 " info->num_kv_heads = g_model.num_kv_heads;\n"
1334 " info->vocab_size = g_model.vocab_size;\n"
1335 " info->context_window = g_model.context_window;\n"
1336 " info->rms_norm_eps = g_model.rms_norm_eps;\n"
1337 " info->rope_theta = g_model.rope_theta;\n"
1338 "}\n\n");
1339
1340 /* ck_model_embed_tokens */
1341 fprintf(out,
1342 "CK_EXPORT int ck_model_embed_tokens(const int32_t *tokens, int num_tokens)\n"
1343 "{\n"
1344 " if (!g_initialized) return -1;\n"
1345 " int cap = g_model.context_window;\n"
1346 " if (g_model.kv_cache_enabled && g_model.kv_cache_capacity > 0 && g_model.kv_cache_capacity < cap) {\n"
1347 " cap = g_model.kv_cache_capacity;\n"
1348 " }\n"
1349 " if (num_tokens > cap) num_tokens = cap;\n"
1350 " if (num_tokens < 1) num_tokens = 1;\n"
1351 " g_model.active_tokens = num_tokens;\n"
1352 " if (g_model.kv_cache_enabled && !g_model.training_enabled) {\n"
1353 " g_model.kv_cache_tokens = 0;\n"
1354 " }\n"
1355 " embed_tokens(&g_model, tokens, num_tokens);\n"
1356 " return 0;\n"
1357 "}\n\n");
1358
1359 /* ck_model_forward */
1360 fprintf(out,
1361 "CK_EXPORT int ck_model_forward(float *logits_out)\n"
1362 "{\n"
1363 " if (!g_initialized) return -1;\n"
1364 " run_model_forward(&g_model);\n"
1365 " if (g_model.kv_cache_enabled && !g_model.training_enabled) {\n"
1366 " g_model.kv_cache_tokens = g_model.active_tokens;\n"
1367 " }\n"
1368 " if (logits_out && g_model.vocab_size > 0) {\n"
1369 " size_t n = (size_t)g_model.active_tokens * (size_t)g_model.vocab_size;\n"
1370 " memcpy(logits_out, ptr_f32(g_model.memory_base, g_model.logits_offset), n * sizeof(float));\n"
1371 " }\n"
1372 " return 0;\n"
1373 "}\n\n");
1374
1375 /* KV-cache helpers + decode API */
1376 fprintf(out,
1377 "CK_EXPORT int ck_model_kv_cache_enable(int capacity)\n"
1378 "{\n"
1379 " if (!g_initialized) return -1;\n"
1380 " if (g_model.training_enabled) return -4;\n"
1381 " g_model.kv_cache_enabled = true;\n"
1382 " int cap = capacity;\n"
1383 " if (cap <= 0 || cap > g_model.context_window) cap = g_model.context_window;\n"
1384 " g_model.kv_cache_capacity = cap;\n"
1385 " g_model.kv_cache_tokens = 0;\n"
1386 " g_model.active_tokens = 0;\n"
1387 " return 0;\n"
1388 "}\n\n"
1389 "CK_EXPORT void ck_model_kv_cache_reset(void)\n"
1390 "{\n"
1391 " if (!g_initialized) return;\n"
1392 " g_model.kv_cache_tokens = 0;\n"
1393 " g_model.active_tokens = 0;\n"
1394 "}\n\n"
1395 "CK_EXPORT int ck_model_kv_cache_get_tokens(void)\n"
1396 "{\n"
1397 " return g_initialized ? g_model.kv_cache_tokens : 0;\n"
1398 "}\n\n"
1399 "CK_EXPORT int ck_model_decode(int32_t token, float *logits_out)\n"
1400 "{\n"
1401 " if (!g_initialized) return -1;\n"
1402 " if (g_model.training_enabled) return -4;\n"
1403 " int ret = run_model_decode(&g_model, token);\n"
1404 " if (ret != 0) return ret;\n"
1405 " if (logits_out && g_model.vocab_size > 0) {\n"
1406 " int t = g_model.active_tokens - 1;\n"
1407 " memcpy(logits_out,\n"
1408 " ptr_f32(g_model.memory_base, g_model.logits_offset) + (size_t)t * (size_t)g_model.vocab_size,\n"
1409 " (size_t)g_model.vocab_size * sizeof(float));\n"
1410 " }\n"
1411 " return 0;\n"
1412 "}\n\n");
1413
1414 /* ck_model_get_logits - get pointer to internal logits buffer */
1415 fprintf(out,
1416 "CK_EXPORT float* ck_model_get_logits(void)\n"
1417 "{\n"
1418 " if (!g_initialized) return NULL;\n"
1419 " return ptr_f32(g_model.memory_base, g_model.logits_offset);\n"
1420 "}\n\n");
1421
1422 /* ck_model_backward */
1423 fprintf(out,
1424 "CK_EXPORT int ck_model_backward(const int32_t *tokens, const int32_t *targets, float *loss_out)\n"
1425 "{\n"
1426 " if (!g_initialized) return -1;\n"
1427 " return run_model_backward(&g_model, tokens, targets, loss_out);\n"
1428 "}\n\n");
1429
1430 /* ck_model_free */
1431 fprintf(out,
1432 "CK_EXPORT void ck_model_free(void)\n"
1433 "{\n"
1434 " if (!g_initialized) return;\n"
1435 " if (g_model.memory_base) ck_huge_free(g_model.memory_base, g_model.total_bytes);\n"
1436 " if (g_model.layers) free(g_model.layers);\n"
1437 " memset(&g_model, 0, sizeof(g_model));\n"
1438 " g_initialized = 0;\n"
1439 "}\n\n");
1440
1441 /* ck_model_get_context_window */
1442 fprintf(out,
1443 "CK_EXPORT int ck_model_get_context_window(void) { return g_initialized ? g_model.context_window : 0; }\n"
1444 "CK_EXPORT int ck_model_get_vocab_size(void) { return g_initialized ? g_model.vocab_size : 0; }\n"
1445 "CK_EXPORT int ck_model_get_hidden_size(void) { return g_initialized ? g_model.embed_dim : 0; }\n"
1446 "CK_EXPORT int ck_model_get_active_tokens(void) { return g_initialized ? g_model.active_tokens : 0; }\n"
1447 "CK_EXPORT int ck_model_is_training_enabled(void) { return g_initialized ? g_model.training_enabled : 0; }\n"
1448 "CK_EXPORT void ck_model_set_learning_rate(float lr) { if (g_initialized) g_model.learning_rate = lr; }\n"
1449 "CK_EXPORT float ck_model_get_learning_rate(void) { return g_initialized ? g_model.learning_rate : 0.0f; }\n\n"
1450 "CK_EXPORT int ck_model_enable_training(float learning_rate)\n"
1451 "{\n"
1452 " if (!g_initialized) return -1;\n"
1453 " g_model.training_enabled = true;\n"
1454 " g_model.learning_rate = learning_rate;\n"
1455 " return 0;\n"
1456 "}\n\n"
1457 "CK_EXPORT void ck_model_disable_training(void)\n"
1458 "{\n"
1459 " if (g_initialized) g_model.training_enabled = false;\n"
1460 "}\n\n"
1461 "CK_EXPORT void ck_model_optimizer_step(void)\n"
1462 "{\n"
1463 " if (!g_initialized || !g_model.training_enabled) return;\n"
1464 " sgd_update(&g_model, g_model.learning_rate);\n"
1465 "}\n\n");
1466}
1467
1468int ck_codegen_emit_runtime(const CKIRGraph *forward, const char *path, CKEmitMode mode)
1469{
1470 if (!forward || !path) {
1471 return -1;
1472 }
1473 if (ck_ir_validate_supported(forward) != 0) {
1474 return -1;
1475 }
1476
1477 FILE *out = fopen(path, "wb");
1478 if (!out) {
1479 fprintf(stderr, "ck_codegen_emit_runtime: failed to open %s: %s\n",
1480 path, strerror(errno));
1481 return -1;
1482 }
1483
1484 if (emit_runtime_preamble(out) != 0) {
1485 fclose(out);
1486 return -1;
1487 }
1488
1489 fprintf(out,
1490 "typedef enum {\n"
1491 " TASK_LM = 0,\n"
1492 " TASK_SEQ_CLS = 1\n"
1493 "} TaskType;\n\n"
1494 "typedef enum {\n"
1495 " OPTIMIZER_SGD = 0,\n"
1496 " OPTIMIZER_ADAM = 1\n"
1497 "} OptimizerType;\n\n"
1498 "typedef struct {\n"
1499 " size_t total_gradient_floats;\n"
1500 "} GradientStorage;\n\n");
1501
1502 emit_layer_offsets_struct(out);
1503 emit_model_struct(out);
1504
1505 fprintf(out,
1506 "static int ensure_layers_allocated(TransformerModel *m)\n"
1507 "{\n"
1508 " if (!m) return -1;\n"
1509 " if (!m->layers && m->num_layers > 0) {\n"
1510 " m->layers = (TrulyOptimalLayer *)calloc((size_t)m->num_layers, sizeof(TrulyOptimalLayer));\n"
1511 " if (!m->layers) return -1;\n"
1512 " }\n"
1513 " return 0;\n"
1514 "}\n\n"
1515 "static void init_weight_dtypes_uniform(TransformerModel *m, CKDataType dt)\n"
1516 "{\n"
1517 " if (!m) return;\n"
1518 " m->token_emb_dtype = dt;\n"
1519 " m->lm_head_weight_dtype = dt;\n"
1520 " m->pos_emb_dtype = CK_DT_FP32;\n"
1521 " if (ensure_layers_allocated(m) != 0) return;\n"
1522 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1523 " TrulyOptimalLayer *L = &m->layers[layer];\n"
1524 " L->wq_dtype = dt;\n"
1525 " L->wk_dtype = dt;\n"
1526 " L->wv_dtype = dt;\n"
1527 " L->wo_dtype = dt;\n"
1528 " L->w1_dtype = dt;\n"
1529 " L->w2_dtype = dt;\n"
1530 " }\n"
1531 "}\n\n"
1532 "static void refresh_weight_flags(TransformerModel *m)\n"
1533 "{\n"
1534 " if (!m) return;\n"
1535 " CKDataType base = m->token_emb_dtype;\n"
1536 " int mixed = 0;\n"
1537 " int quant = ck_dtype_is_quantized(base);\n"
1538 " if (m->lm_head_weight_dtype != base) mixed = 1;\n"
1539 " if (ck_dtype_is_quantized(m->lm_head_weight_dtype)) quant = 1;\n"
1540 " if (m->layers) {\n"
1541 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1542 " TrulyOptimalLayer *L = &m->layers[layer];\n"
1543 " if (L->wq_dtype != base || L->wk_dtype != base || L->wv_dtype != base ||\n"
1544 " L->wo_dtype != base || L->w1_dtype != base || L->w2_dtype != base) {\n"
1545 " mixed = 1;\n"
1546 " }\n"
1547 " if (ck_dtype_is_quantized(L->wq_dtype) || ck_dtype_is_quantized(L->wk_dtype) ||\n"
1548 " ck_dtype_is_quantized(L->wv_dtype) || ck_dtype_is_quantized(L->wo_dtype) ||\n"
1549 " ck_dtype_is_quantized(L->w1_dtype) || ck_dtype_is_quantized(L->w2_dtype)) {\n"
1550 " quant = 1;\n"
1551 " }\n"
1552 " }\n"
1553 " }\n"
1554 " m->weights_mixed = mixed ? true : false;\n"
1555 " m->weights_quantized = quant ? true : false;\n"
1556 " if (!mixed) {\n"
1557 " m->weight_dtype = base;\n"
1558 " }\n"
1559 "}\n\n"
1560 "static int load_weight_dtypes(const char *path, TransformerModel *m)\n"
1561 "{\n"
1562 " if (!path || !m) return -1;\n"
1563 " FILE *f = fopen(path, \"rb\");\n"
1564 " if (!f) return -1;\n"
1565 " char magic[8];\n"
1566 " if (fread(magic, 1, 8, f) != 8) {\n"
1567 " fclose(f);\n"
1568 " return -1;\n"
1569 " }\n"
1570 " if (memcmp(magic, \"BUMPWGT3\", 8) != 0) {\n"
1571 " fclose(f);\n"
1572 " return 0;\n"
1573 " }\n"
1574 " uint32_t version = 0;\n"
1575 " if (fread(&version, sizeof(uint32_t), 1, f) != 1) {\n"
1576 " fclose(f);\n"
1577 " return -1;\n"
1578 " }\n"
1579 " if (version < 3) {\n"
1580 " fclose(f);\n"
1581 " return -1;\n"
1582 " }\n"
1583 " if (fseek(f, 128, SEEK_SET) != 0) {\n"
1584 " fclose(f);\n"
1585 " return -1;\n"
1586 " }\n"
1587 " uint32_t dtype_len = 0;\n"
1588 " if (fread(&dtype_len, sizeof(uint32_t), 1, f) != 1) {\n"
1589 " fclose(f);\n"
1590 " return -1;\n"
1591 " }\n"
1592 " if (dtype_len == 0) {\n"
1593 " fclose(f);\n"
1594 " return -1;\n"
1595 " }\n"
1596 " uint8_t *dtype_buf = (uint8_t *)malloc(dtype_len);\n"
1597 " if (!dtype_buf) {\n"
1598 " fclose(f);\n"
1599 " return -1;\n"
1600 " }\n"
1601 " if (fread(dtype_buf, 1, dtype_len, f) != dtype_len) {\n"
1602 " free(dtype_buf);\n"
1603 " fclose(f);\n"
1604 " return -1;\n"
1605 " }\n"
1606 " fclose(f);\n"
1607 "\n"
1608 " size_t expected = (size_t)m->num_layers * 14u + 4u;\n"
1609 " if (dtype_len != expected) {\n"
1610 " free(dtype_buf);\n"
1611 " return -1;\n"
1612 " }\n"
1613 " if (ensure_layers_allocated(m) != 0) {\n"
1614 " free(dtype_buf);\n"
1615 " return -1;\n"
1616 " }\n"
1617 "\n"
1618 " size_t idx = 0;\n"
1619 " CKDataType token_dt = (CKDataType)dtype_buf[idx++];\n"
1620 " CKDataType pos_dt = (CKDataType)dtype_buf[idx++];\n"
1621 " if (pos_dt != CK_DT_FP32) {\n"
1622 " free(dtype_buf);\n"
1623 " return -1;\n"
1624 " }\n"
1625 " if (token_dt != CK_DT_FP32 && token_dt != CK_DT_Q4_K && token_dt != CK_DT_Q6_K) {\n"
1626 " free(dtype_buf);\n"
1627 " return -1;\n"
1628 " }\n"
1629 " m->token_emb_dtype = token_dt;\n"
1630 " m->lm_head_weight_dtype = token_dt;\n"
1631 " m->pos_emb_dtype = pos_dt;\n"
1632 "\n"
1633 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1634 " CKDataType ln1_dt = (CKDataType)dtype_buf[idx++];\n"
1635 " CKDataType ln2_dt = (CKDataType)dtype_buf[idx++];\n"
1636 " CKDataType wq_dt = (CKDataType)dtype_buf[idx++];\n"
1637 " CKDataType bq_dt = (CKDataType)dtype_buf[idx++];\n"
1638 " CKDataType wk_dt = (CKDataType)dtype_buf[idx++];\n"
1639 " CKDataType bk_dt = (CKDataType)dtype_buf[idx++];\n"
1640 " CKDataType wv_dt = (CKDataType)dtype_buf[idx++];\n"
1641 " CKDataType bv_dt = (CKDataType)dtype_buf[idx++];\n"
1642 " CKDataType wo_dt = (CKDataType)dtype_buf[idx++];\n"
1643 " CKDataType bo_dt = (CKDataType)dtype_buf[idx++];\n"
1644 " CKDataType w1_dt = (CKDataType)dtype_buf[idx++];\n"
1645 " CKDataType b1_dt = (CKDataType)dtype_buf[idx++];\n"
1646 " CKDataType w2_dt = (CKDataType)dtype_buf[idx++];\n"
1647 " CKDataType b2_dt = (CKDataType)dtype_buf[idx++];\n"
1648 "\n"
1649 " if (ln1_dt != CK_DT_FP32 || ln2_dt != CK_DT_FP32 ||\n"
1650 " bq_dt != CK_DT_FP32 || bk_dt != CK_DT_FP32 ||\n"
1651 " bv_dt != CK_DT_FP32 || bo_dt != CK_DT_FP32 ||\n"
1652 " b1_dt != CK_DT_FP32 || b2_dt != CK_DT_FP32) {\n"
1653 " free(dtype_buf);\n"
1654 " return -1;\n"
1655 " }\n"
1656 " if ((wq_dt != CK_DT_FP32 && wq_dt != CK_DT_Q4_K && wq_dt != CK_DT_Q6_K) ||\n"
1657 " (wk_dt != CK_DT_FP32 && wk_dt != CK_DT_Q4_K && wk_dt != CK_DT_Q6_K) ||\n"
1658 " (wv_dt != CK_DT_FP32 && wv_dt != CK_DT_Q4_K && wv_dt != CK_DT_Q6_K) ||\n"
1659 " (wo_dt != CK_DT_FP32 && wo_dt != CK_DT_Q4_K && wo_dt != CK_DT_Q6_K) ||\n"
1660 " (w1_dt != CK_DT_FP32 && w1_dt != CK_DT_Q4_K && w1_dt != CK_DT_Q6_K) ||\n"
1661 " (w2_dt != CK_DT_FP32 && w2_dt != CK_DT_Q4_K && w2_dt != CK_DT_Q6_K)) {\n"
1662 " free(dtype_buf);\n"
1663 " return -1;\n"
1664 " }\n"
1665 "\n"
1666 " TrulyOptimalLayer *L = &m->layers[layer];\n"
1667 " L->wq_dtype = wq_dt;\n"
1668 " L->wk_dtype = wk_dt;\n"
1669 " L->wv_dtype = wv_dt;\n"
1670 " L->wo_dtype = wo_dt;\n"
1671 " L->w1_dtype = w1_dt;\n"
1672 " L->w2_dtype = w2_dt;\n"
1673 " }\n"
1674 "\n"
1675 " CKDataType final_norm_dt = (CKDataType)dtype_buf[idx++];\n"
1676 " CKDataType final_bias_dt = (CKDataType)dtype_buf[idx++];\n"
1677 " free(dtype_buf);\n"
1678 " if (final_norm_dt != CK_DT_FP32 || final_bias_dt != CK_DT_FP32) {\n"
1679 " return -1;\n"
1680 " }\n"
1681 "\n"
1682 " refresh_weight_flags(m);\n"
1683 " return 1;\n"
1684 "}\n\n"
1685 "\n"
1686 "static int layout_model(TransformerModel *m)\n"
1687 "{\n"
1688 " if (!m) return -1;\n"
1689 " if (m->num_attention_heads <= 0 || m->embed_dim <= 0) return -1;\n"
1690 " if (m->num_kv_heads <= 0) m->num_kv_heads = m->num_attention_heads;\n"
1691 " if (m->num_attention_heads %% m->num_kv_heads != 0) return -1;\n"
1692 " if (m->context_window <= 0) m->context_window = 1;\n"
1693 " if (m->vocab_size <= 0) m->vocab_size = 1;\n"
1694 " if (m->intermediate_size <= 0) return -1;\n"
1695 " m->head_dim = m->embed_dim / m->num_attention_heads;\n"
1696 " if (m->rms_norm_eps <= 0.0f) m->rms_norm_eps = 1e-5f;\n"
1697 " if (m->rope_theta < 0.0f) m->rope_theta = 0.0f;\n"
1698 " if (m->rope_theta > 0.0f && (m->head_dim %% 2 != 0)) return -1;\n"
1699 " if (m->elem_bytes == 0) m->elem_bytes = sizeof(float);\n"
1700 " size_t elem_bytes = m->elem_bytes;\n"
1701 " m->aligned_embed_dim = align_up_elems((size_t)m->embed_dim, elem_bytes, CACHELINE_BYTES);\n"
1702 " m->aligned_head_dim = align_up_elems((size_t)m->head_dim, elem_bytes, CACHELINE_BYTES);\n"
1703 " m->aligned_attn_context_window = align_up_elems((size_t)m->context_window, elem_bytes, CACHELINE_BYTES);\n"
1704 " size_t aligned_intermediate_dim = align_up_elems((size_t)m->intermediate_size, elem_bytes, CACHELINE_BYTES);\n"
1705 " if (ensure_layers_allocated(m) != 0) return -1;\n"
1706 " if (m->weights_quantized) {\n"
1707 " /* K-quant weights require K dimension to be a multiple of 256. */\n"
1708 " if ((m->aligned_embed_dim %% 256) != 0) return -1;\n"
1709 " if ((aligned_intermediate_dim %% 256) != 0) return -1;\n"
1710 " int wo_quant = 0;\n"
1711 " if (m->layers) {\n"
1712 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1713 " if (ck_dtype_is_quantized(m->layers[layer].wo_dtype)) {\n"
1714 " wo_quant = 1;\n"
1715 " break;\n"
1716 " }\n"
1717 " }\n"
1718 " }\n"
1719 " if (wo_quant && (size_t)m->num_attention_heads * m->aligned_head_dim != m->aligned_embed_dim) return -1;\n"
1720 " }\n"
1721 "\n"
1722 " if (m->num_cores <= 0) m->num_cores = 1;\n"
1723 " m->tokens_per_core = (m->context_window + m->num_cores - 1) / m->num_cores;\n"
1724 "\n"
1725 " size_t off = 0;\n");
1726 emit_global_allocations(out);
1727 fprintf(out,
1728 " m->layers_start_offset = off;\n"
1729 "\n"
1730 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1731 " TrulyOptimalLayer *L = &m->layers[layer];\n");
1732 emit_layer_allocations(out);
1733 fprintf(out,
1734 " }\n"
1735 "\n");
1736 {
1737 const char *stride_field = ck_first_layer_buffer_name();
1738 fprintf(out,
1739 " if (m->num_layers > 1) {\n"
1740 " m->layer_stride = m->layers[1].%s_offset - m->layers[0].%s_offset;\n"
1741 " } else {\n"
1742 " m->layer_stride = 0;\n"
1743 " }\n",
1744 stride_field, stride_field);
1745 }
1746 emit_global_aliases_to_layer(out);
1747 fprintf(out,
1748 " m->total_bytes = align_up_bytes(off, CACHELINE_BYTES);\n"
1749 " m->memory_base = (uint8_t *)ck_huge_alloc(m->total_bytes);\n"
1750 " if (!m->memory_base) return -1;\n"
1751 " if (m->rope_theta > 0.0f) {\n"
1752 " rope_precompute_cache(ptr_f32(m->memory_base, m->rope_cos_cache_offset),\n"
1753 " ptr_f32(m->memory_base, m->rope_sin_cache_offset),\n"
1754 " m->context_window,\n"
1755 " m->head_dim,\n"
1756 " m->rope_theta,\n"
1757 " m->head_dim,\n"
1758 " \"none\",\n"
1759 " 1.0f);\n"
1760 " }\n"
1761 " return 0;\n"
1762 "}\n\n");
1763
1764 fprintf(out,
1765 "static void lm_head_forward(const float *hidden,\n"
1766 " const float *weights,\n"
1767 " float *logits,\n"
1768 " int T, int V, int D, int aligned_D);\n"
1769 "static void lm_head_backward(const float *hidden,\n"
1770 " const float *weights,\n"
1771 " const float *d_logits,\n"
1772 " float *d_hidden,\n"
1773 " float *d_weights,\n"
1774 " int T, int V, int D, int aligned_D);\n"
1775 "static void softmax_cross_entropy(const float *logits,\n"
1776 " const int32_t *targets,\n"
1777 " int T, int V,\n"
1778 " float *d_logits,\n"
1779 " float *loss_out);\n\n");
1780
1781 fprintf(out,
1782 "static void run_model_forward(TransformerModel *m)\n"
1783 "{\n"
1784 " uint8_t *base = m->memory_base;\n"
1785 " float *current = ptr_f32(base, m->embedded_input_offset);\n"
1786 " int aligned_intermediate_dim = (int)align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n"
1787 " int T = m->active_tokens > 0 ? m->active_tokens : m->context_window;\n"
1788 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
1789 " TrulyOptimalLayer *L = &m->layers[layer];\n"
1790 " if (!m->weights_mixed && m->weight_dtype == CK_DT_Q4_K) {\n"
1791 " CKLayerForwardParamsQ4K p = {0};\n"
1792 " p.tokens = T;\n"
1793 " p.embed_dim = m->embed_dim;\n"
1794 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
1795 " p.num_heads = m->num_attention_heads;\n"
1796 " p.num_kv_heads = m->num_kv_heads;\n"
1797 " p.head_dim = m->head_dim;\n"
1798 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
1799 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
1800 " p.intermediate_dim = m->intermediate_size;\n"
1801 " p.aligned_intermediate_dim = aligned_intermediate_dim;\n"
1802 " p.eps = m->rms_norm_eps;\n"
1803 " p.rope_pos_offset = 0;\n"
1804 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
1805 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
1806 " p.input = current;\n"
1807 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
1808 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
1809 " p.wq = cptr_void(base, L->wq_offset);\n"
1810 " p.bq = cptr_f32(base, L->bq_offset);\n"
1811 " p.wk = cptr_void(base, L->wk_offset);\n"
1812 " p.bk = cptr_f32(base, L->bk_offset);\n"
1813 " p.wv = cptr_void(base, L->wv_offset);\n"
1814 " p.bv = cptr_f32(base, L->bv_offset);\n"
1815 " p.wo = cptr_void(base, L->wo_offset);\n"
1816 " p.bo = cptr_f32(base, L->bo_offset);\n"
1817 " p.w1 = cptr_void(base, L->w1_offset);\n"
1818 " p.b1 = cptr_f32(base, L->b1_offset);\n"
1819 " p.w2 = cptr_void(base, L->w2_offset);\n"
1820 " p.b2 = cptr_f32(base, L->b2_offset);\n"
1821 " p.ln1_out = ptr_f32(base, L->ln1_out_offset);\n"
1822 " p.ln1_rstd = ptr_f32(base, L->ln1_rstd_offset);\n"
1823 " p.q = ptr_f32(base, L->q_offset);\n"
1824 " p.k = ptr_f32(base, L->k_offset);\n"
1825 " p.v = ptr_f32(base, L->v_offset);\n"
1826 " p.scores = L->scores_offset ? ptr_f32(base, L->scores_offset) : NULL;\n"
1827 " p.attn_out = ptr_f32(base, L->attn_out_offset);\n"
1828 " p.proj_tmp = ptr_f32(base, L->proj_tmp_offset);\n"
1829 " p.proj_scratch = ptr_f32(base, L->proj_scratch_offset);\n"
1830 " p.residual1 = ptr_f32(base, L->residual1_offset);\n"
1831 " p.ln2_out = ptr_f32(base, L->ln2_out_offset);\n"
1832 " p.ln2_rstd = ptr_f32(base, L->ln2_rstd_offset);\n"
1833 " p.fc1_out = ptr_f32(base, L->fc1_out_offset);\n"
1834 " p.swiglu_out = ptr_f32(base, L->swiglu_out_offset);\n"
1835 " p.mlp_out = ptr_f32(base, L->mlp_out_offset);\n"
1836 " p.output = ptr_f32(base, L->output_offset);\n"
1837 " ck_layer_forward_rmsnorm_swiglu_q4_k(&p);\n"
1838 " if (m->kv_cache_enabled && !m->training_enabled) {\n"
1839 " kv_cache_repack_head_major_inplace(p.k,\n"
1840 " p.num_kv_heads,\n"
1841 " T,\n"
1842 " m->kv_cache_capacity,\n"
1843 " p.aligned_head_dim);\n"
1844 " kv_cache_repack_head_major_inplace(p.v,\n"
1845 " p.num_kv_heads,\n"
1846 " T,\n"
1847 " m->kv_cache_capacity,\n"
1848 " p.aligned_head_dim);\n"
1849 " }\n"
1850 " current = p.output;\n"
1851 " } else if (m->weights_quantized) {\n"
1852 " CKLayerForwardParamsQ4K p = {0};\n"
1853 " p.tokens = T;\n"
1854 " p.embed_dim = m->embed_dim;\n"
1855 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
1856 " p.num_heads = m->num_attention_heads;\n"
1857 " p.num_kv_heads = m->num_kv_heads;\n"
1858 " p.head_dim = m->head_dim;\n"
1859 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
1860 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
1861 " p.intermediate_dim = m->intermediate_size;\n"
1862 " p.aligned_intermediate_dim = aligned_intermediate_dim;\n"
1863 " p.eps = m->rms_norm_eps;\n"
1864 " p.rope_pos_offset = 0;\n"
1865 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
1866 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
1867 " p.input = current;\n"
1868 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
1869 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
1870 " p.wq = cptr_void(base, L->wq_offset);\n"
1871 " p.bq = cptr_f32(base, L->bq_offset);\n"
1872 " p.wk = cptr_void(base, L->wk_offset);\n"
1873 " p.bk = cptr_f32(base, L->bk_offset);\n"
1874 " p.wv = cptr_void(base, L->wv_offset);\n"
1875 " p.bv = cptr_f32(base, L->bv_offset);\n"
1876 " p.wo = cptr_void(base, L->wo_offset);\n"
1877 " p.bo = cptr_f32(base, L->bo_offset);\n"
1878 " p.w1 = cptr_void(base, L->w1_offset);\n"
1879 " p.b1 = cptr_f32(base, L->b1_offset);\n"
1880 " p.w2 = cptr_void(base, L->w2_offset);\n"
1881 " p.b2 = cptr_f32(base, L->b2_offset);\n"
1882 " p.ln1_out = ptr_f32(base, L->ln1_out_offset);\n"
1883 " p.ln1_rstd = ptr_f32(base, L->ln1_rstd_offset);\n"
1884 " p.q = ptr_f32(base, L->q_offset);\n"
1885 " p.k = ptr_f32(base, L->k_offset);\n"
1886 " p.v = ptr_f32(base, L->v_offset);\n"
1887 " p.scores = L->scores_offset ? ptr_f32(base, L->scores_offset) : NULL;\n"
1888 " p.attn_out = ptr_f32(base, L->attn_out_offset);\n"
1889 " p.proj_tmp = ptr_f32(base, L->proj_tmp_offset);\n"
1890 " p.proj_scratch = ptr_f32(base, L->proj_scratch_offset);\n"
1891 " p.residual1 = ptr_f32(base, L->residual1_offset);\n"
1892 " p.ln2_out = ptr_f32(base, L->ln2_out_offset);\n"
1893 " p.ln2_rstd = ptr_f32(base, L->ln2_rstd_offset);\n"
1894 " p.fc1_out = ptr_f32(base, L->fc1_out_offset);\n"
1895 " p.swiglu_out = ptr_f32(base, L->swiglu_out_offset);\n"
1896 " p.mlp_out = ptr_f32(base, L->mlp_out_offset);\n"
1897 " p.output = ptr_f32(base, L->output_offset);\n"
1898 " p.wq_dtype = L->wq_dtype;\n"
1899 " p.wk_dtype = L->wk_dtype;\n"
1900 " p.wv_dtype = L->wv_dtype;\n"
1901 " p.wo_dtype = L->wo_dtype;\n"
1902 " p.w1_dtype = L->w1_dtype;\n"
1903 " p.w2_dtype = L->w2_dtype;\n"
1904 " ck_layer_forward_rmsnorm_swiglu_quant(&p);\n"
1905 " if (m->kv_cache_enabled && !m->training_enabled) {\n"
1906 " kv_cache_repack_head_major_inplace(p.k,\n"
1907 " p.num_kv_heads,\n"
1908 " T,\n"
1909 " m->kv_cache_capacity,\n"
1910 " p.aligned_head_dim);\n"
1911 " kv_cache_repack_head_major_inplace(p.v,\n"
1912 " p.num_kv_heads,\n"
1913 " T,\n"
1914 " m->kv_cache_capacity,\n"
1915 " p.aligned_head_dim);\n"
1916 " }\n"
1917 " current = p.output;\n"
1918 " } else {\n"
1919 " CKLayerForwardParams p = {0};\n"
1920 " p.tokens = T;\n"
1921 " p.embed_dim = m->embed_dim;\n"
1922 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
1923 " p.num_heads = m->num_attention_heads;\n"
1924 " p.num_kv_heads = m->num_kv_heads;\n"
1925 " p.head_dim = m->head_dim;\n"
1926 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
1927 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
1928 " p.intermediate_dim = m->intermediate_size;\n"
1929 " p.aligned_intermediate_dim = aligned_intermediate_dim;\n"
1930 " p.eps = m->rms_norm_eps;\n"
1931 " p.rope_pos_offset = 0;\n"
1932 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
1933 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
1934 " p.input = current;\n"
1935 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
1936 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
1937 " p.wq = cptr_f32(base, L->wq_offset);\n"
1938 " p.bq = cptr_f32(base, L->bq_offset);\n"
1939 " p.wk = cptr_f32(base, L->wk_offset);\n"
1940 " p.bk = cptr_f32(base, L->bk_offset);\n"
1941 " p.wv = cptr_f32(base, L->wv_offset);\n"
1942 " p.bv = cptr_f32(base, L->bv_offset);\n"
1943 " p.wo = cptr_f32(base, L->wo_offset);\n"
1944 " p.bo = cptr_f32(base, L->bo_offset);\n"
1945 " p.w1 = cptr_f32(base, L->w1_offset);\n"
1946 " p.b1 = cptr_f32(base, L->b1_offset);\n"
1947 " p.w2 = cptr_f32(base, L->w2_offset);\n"
1948 " p.b2 = cptr_f32(base, L->b2_offset);\n"
1949 " p.ln1_out = ptr_f32(base, L->ln1_out_offset);\n"
1950 " p.ln1_rstd = ptr_f32(base, L->ln1_rstd_offset);\n"
1951 " p.q = ptr_f32(base, L->q_offset);\n"
1952 " p.k = ptr_f32(base, L->k_offset);\n"
1953 " p.v = ptr_f32(base, L->v_offset);\n"
1954 " p.scores = L->scores_offset ? ptr_f32(base, L->scores_offset) : NULL;\n"
1955 " p.attn_out = ptr_f32(base, L->attn_out_offset);\n"
1956 " p.proj_tmp = ptr_f32(base, L->proj_tmp_offset);\n"
1957 " p.proj_scratch = ptr_f32(base, L->proj_scratch_offset);\n"
1958 " p.residual1 = ptr_f32(base, L->residual1_offset);\n"
1959 " p.ln2_out = ptr_f32(base, L->ln2_out_offset);\n"
1960 " p.ln2_rstd = ptr_f32(base, L->ln2_rstd_offset);\n"
1961 " p.fc1_out = ptr_f32(base, L->fc1_out_offset);\n"
1962 " p.swiglu_out = ptr_f32(base, L->swiglu_out_offset);\n"
1963 " p.mlp_out = ptr_f32(base, L->mlp_out_offset);\n"
1964 " p.output = ptr_f32(base, L->output_offset);\n"
1965 " ck_layer_forward_rmsnorm_swiglu(&p);\n"
1966 " if (m->kv_cache_enabled && !m->training_enabled) {\n"
1967 " kv_cache_repack_head_major_inplace(p.k,\n"
1968 " p.num_kv_heads,\n"
1969 " T,\n"
1970 " m->kv_cache_capacity,\n"
1971 " p.aligned_head_dim);\n"
1972 " kv_cache_repack_head_major_inplace(p.v,\n"
1973 " p.num_kv_heads,\n"
1974 " T,\n"
1975 " m->kv_cache_capacity,\n"
1976 " p.aligned_head_dim);\n"
1977 " }\n"
1978 " current = p.output;\n"
1979 " }\n"
1980 " }\n"
1981 " float *final_out = ptr_f32(base, m->final_output_offset);\n"
1982 " rmsnorm_forward(current,\n"
1983 " cptr_f32(base, m->final_ln_weight_offset),\n"
1984 " final_out,\n"
1985 " ptr_f32(base, m->final_ln_rstd_offset),\n"
1986 " T,\n"
1987 " m->embed_dim,\n"
1988 " (int)m->aligned_embed_dim,\n"
1989 " m->rms_norm_eps);\n"
1990 " if (m->vocab_size > 0) {\n"
1991 " if (m->lm_head_weight_dtype == CK_DT_Q4_K) {\n"
1992 " gemm_nt_q4_k(final_out,\n"
1993 " cptr_void(base, m->lm_head_weight_offset),\n"
1994 " NULL,\n"
1995 " ptr_f32(base, m->logits_offset),\n"
1996 " T,\n"
1997 " m->vocab_size,\n"
1998 " (int)m->aligned_embed_dim);\n"
1999 " } else if (m->lm_head_weight_dtype == CK_DT_Q6_K) {\n"
2000 " gemm_nt_q6_k(final_out,\n"
2001 " cptr_void(base, m->lm_head_weight_offset),\n"
2002 " NULL,\n"
2003 " ptr_f32(base, m->logits_offset),\n"
2004 " T,\n"
2005 " m->vocab_size,\n"
2006 " (int)m->aligned_embed_dim);\n"
2007 " } else {\n"
2008 " lm_head_forward(final_out,\n"
2009 " cptr_f32(base, m->lm_head_weight_offset),\n"
2010 " ptr_f32(base, m->logits_offset),\n"
2011 " T,\n"
2012 " m->vocab_size,\n"
2013 " m->embed_dim,\n"
2014 " (int)m->aligned_embed_dim);\n"
2015 " }\n"
2016 " }\n"
2017 "}\n\n");
2018
2019 emit_zero_grad(out);
2020 emit_sgd_update(out);
2021
2022 fprintf(out,
2023 "static int run_model_backward(TransformerModel *m,\n"
2024 " const int32_t *tokens,\n"
2025 " const int32_t *targets,\n"
2026 " float *loss_out)\n"
2027 "{\n"
2028 " if (!m || !m->training_enabled) return 0;\n"
2029 " if (!tokens || !targets) return -1;\n"
2030 " if (m->num_layers <= 0) return -1;\n"
2031 " int T = m->active_tokens > 0 ? m->active_tokens : m->context_window;\n"
2032 " int V = m->vocab_size;\n"
2033 " int D = m->embed_dim;\n"
2034 " int aligned_D = (int)m->aligned_embed_dim;\n"
2035 " uint8_t *base = m->memory_base;\n"
2036 "\n"
2037 " zero_grad(m);\n"
2038 "\n"
2039 " float *final_out = ptr_f32(base, m->final_output_offset);\n"
2040 " float *logits = ptr_f32(base, m->logits_offset);\n"
2041 " float *d_logits = ptr_f32(base, m->d_logits_offset);\n"
2042 " float *d_final_out = ptr_f32(base, m->d_final_output_offset);\n"
2043 " float *d_final_in = ptr_f32(base, m->d_final_input_offset);\n"
2044 "\n"
2045 " float loss = 0.0f;\n"
2046 " softmax_cross_entropy(logits, targets, T, V, d_logits, &loss);\n"
2047 " if (loss_out) {\n"
2048 " *loss_out = loss;\n"
2049 " }\n"
2050 " lm_head_backward(final_out,\n"
2051 " cptr_f32(base, m->lm_head_weight_offset),\n"
2052 " d_logits,\n"
2053 " d_final_out,\n"
2054 " ptr_f32(base, m->d_token_emb_offset),\n"
2055 " T, V, D, aligned_D);\n"
2056 " rmsnorm_backward(d_final_out,\n"
2057 " ptr_f32(base, m->layers[m->num_layers - 1].output_offset),\n"
2058 " cptr_f32(base, m->final_ln_weight_offset),\n"
2059 " ptr_f32(base, m->final_ln_rstd_offset),\n"
2060 " d_final_in,\n"
2061 " ptr_f32(base, m->d_final_ln_weight_offset),\n"
2062 " T, D, aligned_D);\n"
2063 "\n"
2064 " for (int layer = m->num_layers - 1; layer >= 0; --layer) {\n"
2065 " TrulyOptimalLayer *L = &m->layers[layer];\n"
2066 " CKLayerBackwardParams p = {0};\n"
2067 " p.tokens = T;\n"
2068 " p.embed_dim = m->embed_dim;\n"
2069 " p.aligned_embed_dim = (int)m->aligned_embed_dim;\n"
2070 " p.num_heads = m->num_attention_heads;\n"
2071 " p.num_kv_heads = m->num_kv_heads;\n"
2072 " p.head_dim = m->head_dim;\n"
2073 " p.aligned_head_dim = (int)m->aligned_head_dim;\n"
2074 " p.aligned_context_window = (int)m->aligned_attn_context_window;\n"
2075 " p.intermediate_dim = m->intermediate_size;\n"
2076 " p.aligned_intermediate_dim = (int)align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n"
2077 " p.eps = m->rms_norm_eps;\n"
2078 " p.rope_pos_offset = 0;\n"
2079 " p.rope_cos = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_cos_cache_offset) : NULL;\n"
2080 " p.rope_sin = (m->rope_theta > 0.0f) ? cptr_f32(base, m->rope_sin_cache_offset) : NULL;\n"
2081 " p.input = (layer == 0) ? ptr_f32(base, m->embedded_input_offset)\n"
2082 " : ptr_f32(base, m->layers[layer - 1].output_offset);\n"
2083 " p.ln1_gamma = cptr_f32(base, L->ln1_gamma_offset);\n"
2084 " p.ln2_gamma = cptr_f32(base, L->ln2_gamma_offset);\n"
2085 " p.ln1_out = cptr_f32(base, L->ln1_out_offset);\n"
2086 " p.ln1_rstd = cptr_f32(base, L->ln1_rstd_offset);\n"
2087 " p.ln2_out = cptr_f32(base, L->ln2_out_offset);\n"
2088 " p.ln2_rstd = cptr_f32(base, L->ln2_rstd_offset);\n"
2089 " p.wq = cptr_f32(base, L->wq_offset);\n"
2090 " p.bq = cptr_f32(base, L->bq_offset);\n"
2091 " p.wk = cptr_f32(base, L->wk_offset);\n"
2092 " p.bk = cptr_f32(base, L->bk_offset);\n"
2093 " p.wv = cptr_f32(base, L->wv_offset);\n"
2094 " p.bv = cptr_f32(base, L->bv_offset);\n"
2095 " p.wo = cptr_f32(base, L->wo_offset);\n"
2096 " p.bo = cptr_f32(base, L->bo_offset);\n"
2097 " p.w1 = cptr_f32(base, L->w1_offset);\n"
2098 " p.b1 = cptr_f32(base, L->b1_offset);\n"
2099 " p.w2 = cptr_f32(base, L->w2_offset);\n"
2100 " p.b2 = cptr_f32(base, L->b2_offset);\n"
2101 " p.q = cptr_f32(base, L->q_offset);\n"
2102 " p.k = cptr_f32(base, L->k_offset);\n"
2103 " p.v = cptr_f32(base, L->v_offset);\n"
2104 " p.scores = L->scores_offset ? cptr_f32(base, L->scores_offset) : NULL;\n"
2105 " p.attn_out = cptr_f32(base, L->attn_out_offset);\n"
2106 " p.residual1 = cptr_f32(base, L->residual1_offset);\n"
2107 " p.fc1_out = cptr_f32(base, L->fc1_out_offset);\n"
2108 " p.swiglu_out = cptr_f32(base, L->swiglu_out_offset);\n"
2109 " p.d_output = ptr_f32(base, L->d_output_offset);\n"
2110 " p.d_input = ptr_f32(base, L->d_input_offset);\n"
2111 " p.d_ln1_gamma = ptr_f32(base, L->d_ln1_gamma_offset);\n"
2112 " p.d_ln2_gamma = ptr_f32(base, L->d_ln2_gamma_offset);\n"
2113 " p.d_wq = ptr_f32(base, L->d_wq_offset);\n"
2114 " p.d_bq = ptr_f32(base, L->d_bq_offset);\n"
2115 " p.d_wk = ptr_f32(base, L->d_wk_offset);\n"
2116 " p.d_bk = ptr_f32(base, L->d_bk_offset);\n"
2117 " p.d_wv = ptr_f32(base, L->d_wv_offset);\n"
2118 " p.d_bv = ptr_f32(base, L->d_bv_offset);\n"
2119 " p.d_wo = ptr_f32(base, L->d_wo_offset);\n"
2120 " p.d_bo = ptr_f32(base, L->d_bo_offset);\n"
2121 " p.d_w1 = ptr_f32(base, L->d_w1_offset);\n"
2122 " p.d_b1 = ptr_f32(base, L->d_b1_offset);\n"
2123 " p.d_w2 = ptr_f32(base, L->d_w2_offset);\n"
2124 " p.d_b2 = ptr_f32(base, L->d_b2_offset);\n"
2125 " p.d_ln1_out = ptr_f32(base, L->d_ln1_out_offset);\n"
2126 " p.d_q = ptr_f32(base, L->d_q_offset);\n"
2127 " p.d_k = ptr_f32(base, L->d_k_offset);\n"
2128 " p.d_v = ptr_f32(base, L->d_v_offset);\n"
2129 " p.d_scores = ptr_f32(base, L->d_scores_offset);\n"
2130 " p.d_attn_out = ptr_f32(base, L->d_attn_out_offset);\n"
2131 " p.d_proj_tmp = ptr_f32(base, L->d_proj_tmp_offset);\n"
2132 " p.d_residual1 = ptr_f32(base, L->d_residual1_offset);\n"
2133 " p.d_ln2_out = ptr_f32(base, L->d_ln2_out_offset);\n"
2134 " p.d_fc1_out = ptr_f32(base, L->d_fc1_out_offset);\n"
2135 " p.d_swiglu_out = ptr_f32(base, L->d_swiglu_out_offset);\n"
2136 " p.d_mlp_out = ptr_f32(base, L->d_mlp_out_offset);\n"
2137 "\n"
2138 " const float *src = (layer == m->num_layers - 1)\n"
2139 " ? d_final_in\n"
2140 " : ptr_f32(base, m->layers[layer + 1].d_input_offset);\n"
2141 " memcpy(p.d_output, src, (size_t)T * (size_t)aligned_D * sizeof(float));\n"
2142 "\n"
2143 " ck_layer_backward_rmsnorm_swiglu(&p);\n"
2144 " }\n"
2145 "\n"
2146 " {\n"
2147 " TrulyOptimalLayer *L0 = &m->layers[0];\n"
2148 " embedding_backward(tokens,\n"
2149 " T,\n"
2150 " ptr_f32(base, L0->d_input_offset),\n"
2151 " ptr_f32(base, m->d_token_emb_offset),\n"
2152 " ptr_f32(base, m->d_pos_emb_offset),\n"
2153 " m->vocab_size,\n"
2154 " m->embed_dim,\n"
2155 " aligned_D,\n"
2156 " m->context_window,\n"
2157 " m->rope_theta <= 0.0f);\n"
2158 " }\n"
2159 "\n"
2160 " /* SGD update is now called separately via optimizer_step() */\n"
2161 " return 0;\n"
2162 "}\n\n");
2163
2164 fprintf(out,
2165 "static int parse_int_arg(const char *s, int *out)\n"
2166 "{\n"
2167 " if (!s || !out) return 0;\n"
2168 " char *end = NULL;\n"
2169 " long v = strtol(s, &end, 10);\n"
2170 " if (!end || *end != '\\0') return 0;\n"
2171 " *out = (int)v;\n"
2172 " return 1;\n"
2173 "}\n\n"
2174 "static int parse_float_arg(const char *s, float *out)\n"
2175 "{\n"
2176 " if (!s || !out) return 0;\n"
2177 " char *end = NULL;\n"
2178 " double v = strtod(s, &end);\n"
2179 " if (!end || *end != '\\0') return 0;\n"
2180 " *out = (float)v;\n"
2181 " return 1;\n"
2182 "}\n\n"
2183 "static void print_usage(const char *prog)\n"
2184 "{\n"
2185 " printf(\"Usage: %%s [options]\\n\", prog);\n"
2186 " printf(\" --dump Print layout summary (layer 0 only)\\n\");\n"
2187 " printf(\" --dump-all Print layout summary for all layers\\n\");\n"
2188 " printf(\" --no-forward Skip forward pass (layout + alloc only)\\n\");\n"
2189 " printf(\" --layers N Override num_layers\\n\");\n"
2190 " printf(\" --embed N Override embed_dim\\n\");\n"
2191 " printf(\" --intermediate N Override intermediate_size\\n\");\n"
2192 " printf(\" --heads N Override num_attention_heads\\n\");\n"
2193 " printf(\" --kv-heads N Override num_kv_heads\\n\");\n"
2194 " printf(\" --vocab N Override vocab_size\\n\");\n"
2195 " printf(\" --ctx N Override context_window\\n\");\n"
2196 " printf(\" --cores N Override num_cores\\n\");\n"
2197 " printf(\" --litmus Run LM head + CE + backward litmus\\n\");\n"
2198 " printf(\" --backward Run backward pass + SGD update (requires --tokens/--targets)\\n\");\n"
2199 " printf(\" --lr F SGD learning rate (default: 1e-3 when --backward)\\n\");\n"
2200 " printf(\" --steps N Training steps (default: 1)\\n\");\n"
2201 " printf(\" --log-steps Print loss per step during training\\n\");\n"
2202 " printf(\" --strict Enable strict parity mode (single-thread + double GEMM)\\n\");\n"
2203 " printf(\" --hidden PATH Load hidden activations [T x aligned_D] f32\\n\");\n"
2204 " printf(\" --weights PATH Load LM head weights [V x aligned_D] f32 (litmus)\\n\");\n"
2205 " printf(\" --targets PATH Load target tokens [T] int32\\n\");\n"
2206 " printf(\" --model-weights PATH Load full model weights (bump format)\\n\");\n"
2207 " printf(\" --tokens PATH Load token IDs [T] int32 and build embeddings\\n\");\n"
2208 " printf(\" --out-logits PATH Write logits [T x V] f32\\n\");\n"
2209 " printf(\" --out-dlogits PATH Write d_logits [T x V] f32\\n\");\n"
2210 " printf(\" --out-dhidden PATH Write d_hidden [T x aligned_D] f32\\n\");\n"
2211 " printf(\" --out-dweights PATH Write d_weights [V x aligned_D] f32\\n\");\n"
2212 " printf(\" --out-loss PATH Write loss (single f32)\\n\");\n"
2213 " printf(\" --out-weights PATH Write model weights (flat, no header)\\n\");\n"
2214 " printf(\" --help Show this help\\n\");\n"
2215 "}\n\n"
2216 "static int read_floats(const char *path, float *dst, size_t count)\n"
2217 "{\n"
2218 " if (!path || !dst) return -1;\n"
2219 " FILE *f = fopen(path, \"rb\");\n"
2220 " if (!f) {\n"
2221 " perror(\"fopen\");\n"
2222 " return -1;\n"
2223 " }\n"
2224 " size_t got = fread(dst, sizeof(float), count, f);\n"
2225 " fclose(f);\n"
2226 " return got == count ? 0 : -1;\n"
2227 "}\n\n"
2228 "static int read_ints(const char *path, int32_t *dst, size_t count)\n"
2229 "{\n"
2230 " if (!path || !dst) return -1;\n"
2231 " FILE *f = fopen(path, \"rb\");\n"
2232 " if (!f) {\n"
2233 " perror(\"fopen\");\n"
2234 " return -1;\n"
2235 " }\n"
2236 " size_t got = fread(dst, sizeof(int32_t), count, f);\n"
2237 " fclose(f);\n"
2238 " return got == count ? 0 : -1;\n"
2239 "}\n\n"
2240 "static int read_floats_file(FILE *f, float *dst, size_t count)\n"
2241 "{\n"
2242 " if (!f || !dst) return -1;\n"
2243 " size_t got = fread(dst, sizeof(float), count, f);\n"
2244 " return got == count ? 0 : -1;\n"
2245 "}\n\n"
2246 "static int read_bytes_file(FILE *f, void *dst, size_t bytes)\n"
2247 "{\n"
2248 " if (!f || !dst) return -1;\n"
2249 " size_t got = fread(dst, 1, bytes, f);\n"
2250 " return got == bytes ? 0 : -1;\n"
2251 "}\n\n"
2252 "static int write_floats_file(FILE *f, const float *src, size_t count)\n"
2253 "{\n"
2254 " if (!f || !src) return -1;\n"
2255 " size_t wrote = fwrite(src, sizeof(float), count, f);\n"
2256 " return wrote == count ? 0 : -1;\n"
2257 "}\n\n"
2258 "static int write_bytes_file(FILE *f, const void *src, size_t bytes)\n"
2259 "{\n"
2260 " if (!f || !src) return -1;\n"
2261 " size_t wrote = fwrite(src, 1, bytes, f);\n"
2262 " return wrote == bytes ? 0 : -1;\n"
2263 "}\n\n"
2264 "static int read_weight_file(FILE *f, CKDataType dtype, void *dst, size_t n_elements)\n"
2265 "{\n"
2266 " if (!f || !dst) return -1;\n"
2267 " if (dtype == CK_DT_FP32) {\n"
2268 " return read_floats_file(f, (float *)dst, n_elements);\n"
2269 " }\n"
2270 " return read_bytes_file(f, dst, ck_dtype_row_bytes(dtype, n_elements));\n"
2271 "}\n\n"
2272 "static int write_weight_file(FILE *f, CKDataType dtype, const void *src, size_t n_elements)\n"
2273 "{\n"
2274 " if (!f || !src) return -1;\n"
2275 " if (dtype == CK_DT_FP32) {\n"
2276 " return write_floats_file(f, (const float *)src, n_elements);\n"
2277 " }\n"
2278 " return write_bytes_file(f, src, ck_dtype_row_bytes(dtype, n_elements));\n"
2279 "}\n\n"
2280 "static int skip_bump_header(FILE *f)\n"
2281 "{\n"
2282 " if (!f) return -1;\n"
2283 " char magic[8];\n"
2284 " if (fread(magic, 1, 8, f) != 8) return -1;\n"
2285 " if (memcmp(magic, \"BUMPWGT3\", 8) == 0) {\n"
2286 " if (fseek(f, 128, SEEK_SET) != 0) return -1;\n"
2287 " uint32_t dtype_len = 0;\n"
2288 " if (fread(&dtype_len, sizeof(uint32_t), 1, f) != 1) return -1;\n"
2289 " if (fseek(f, (long)dtype_len, SEEK_CUR) != 0) return -1;\n"
2290 " return 1;\n"
2291 " }\n"
2292 " if (memcmp(magic, \"BUMPWGT2\", 8) == 0) {\n"
2293 " if (fseek(f, 128, SEEK_SET) != 0) return -1;\n"
2294 " return 1;\n"
2295 " }\n"
2296 " if (fseek(f, 0, SEEK_SET) != 0) return -1;\n"
2297 " return 0;\n"
2298 "}\n\n"
2299 "static int load_model_weights(const char *path, TransformerModel *m)\n"
2300 "{\n"
2301 " if (!path || !m || !m->memory_base) return -1;\n"
2302 " FILE *f = fopen(path, \"rb\");\n"
2303 " if (!f) {\n"
2304 " perror(\"fopen\");\n"
2305 " return -1;\n"
2306 " }\n"
2307 " if (skip_bump_header(f) < 0) {\n"
2308 " fclose(f);\n"
2309 " return -1;\n"
2310 " }\n"
2311 " uint8_t *base = m->memory_base;\n"
2312 " size_t aligned_intermediate = align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n"
2313 " size_t tok_elems = (size_t)m->vocab_size * m->aligned_embed_dim;\n"
2314 " if (read_weight_file(f, m->token_emb_dtype, ptr_u8(base, m->token_emb_offset), tok_elems) != 0) goto fail;\n"
2315 " if (read_floats_file(f, ptr_f32(base, m->pos_emb_offset),\n"
2316 " (size_t)m->context_window * m->aligned_embed_dim) != 0) goto fail;\n"
2317 "\n"
2318 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
2319 " TrulyOptimalLayer *L = &m->layers[layer];\n"
2320 " size_t head_w_stride = m->aligned_head_dim * m->aligned_embed_dim;\n"
2321 " size_t q_w = (size_t)m->num_attention_heads * head_w_stride;\n"
2322 " size_t kv_w = (size_t)m->num_kv_heads * head_w_stride;\n"
2323 " size_t q_b = (size_t)m->num_attention_heads * m->aligned_head_dim;\n"
2324 " size_t kv_b = (size_t)m->num_kv_heads * m->aligned_head_dim;\n"
2325 " size_t wo_w = (size_t)m->num_attention_heads * m->aligned_embed_dim * m->aligned_head_dim;\n"
2326 " size_t w1_w = (size_t)(2 * aligned_intermediate) * m->aligned_embed_dim;\n"
2327 " size_t w2_w = m->aligned_embed_dim * aligned_intermediate;\n"
2328 "\n"
2329 " if (read_floats_file(f, ptr_f32(base, L->ln1_gamma_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2330 " if (read_floats_file(f, ptr_f32(base, L->ln2_gamma_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2331 " if (read_weight_file(f, L->wq_dtype, ptr_u8(base, L->wq_offset), q_w) != 0) goto fail;\n"
2332 " if (read_floats_file(f, ptr_f32(base, L->bq_offset), q_b) != 0) goto fail;\n"
2333 " if (read_weight_file(f, L->wk_dtype, ptr_u8(base, L->wk_offset), kv_w) != 0) goto fail;\n"
2334 " if (read_floats_file(f, ptr_f32(base, L->bk_offset), kv_b) != 0) goto fail;\n"
2335 " if (read_weight_file(f, L->wv_dtype, ptr_u8(base, L->wv_offset), kv_w) != 0) goto fail;\n"
2336 " if (read_floats_file(f, ptr_f32(base, L->bv_offset), kv_b) != 0) goto fail;\n"
2337 " if (read_weight_file(f, L->wo_dtype, ptr_u8(base, L->wo_offset), wo_w) != 0) goto fail;\n"
2338 " if (read_floats_file(f, ptr_f32(base, L->bo_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2339 " if (read_weight_file(f, L->w1_dtype, ptr_u8(base, L->w1_offset), w1_w) != 0) goto fail;\n"
2340 " if (read_floats_file(f, ptr_f32(base, L->b1_offset), (size_t)(2 * aligned_intermediate)) != 0) goto fail;\n"
2341 " if (read_weight_file(f, L->w2_dtype, ptr_u8(base, L->w2_offset), w2_w) != 0) goto fail;\n"
2342 " if (read_floats_file(f, ptr_f32(base, L->b2_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2343 " }\n"
2344 "\n"
2345 " if (read_floats_file(f, ptr_f32(base, m->final_ln_weight_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2346 " if (read_floats_file(f, ptr_f32(base, m->final_ln_bias_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2347 "\n"
2348 " fclose(f);\n"
2349 " return 0;\n"
2350 "fail:\n"
2351 " fclose(f);\n"
2352 " return -1;\n"
2353 "}\n\n"
2354 "static int save_model_weights(const char *path, const TransformerModel *m)\n"
2355 "{\n"
2356 " if (!path || !m || !m->memory_base) return -1;\n"
2357 " FILE *f = fopen(path, \"wb\");\n"
2358 " if (!f) {\n"
2359 " perror(\"fopen\");\n"
2360 " return -1;\n"
2361 " }\n"
2362 " uint8_t *base = m->memory_base;\n"
2363 " size_t aligned_intermediate = align_up_elems((size_t)m->intermediate_size, m->elem_bytes, CACHELINE_BYTES);\n"
2364 " size_t tok_elems = (size_t)m->vocab_size * m->aligned_embed_dim;\n"
2365 " if (write_weight_file(f, m->token_emb_dtype, cptr_void(base, m->token_emb_offset), tok_elems) != 0) goto fail;\n"
2366 " if (write_floats_file(f, ptr_f32(base, m->pos_emb_offset),\n"
2367 " (size_t)m->context_window * m->aligned_embed_dim) != 0) goto fail;\n"
2368 "\n"
2369 " for (int layer = 0; layer < m->num_layers; ++layer) {\n"
2370 " const TrulyOptimalLayer *L = &m->layers[layer];\n"
2371 " size_t head_w_stride = m->aligned_head_dim * m->aligned_embed_dim;\n"
2372 " size_t q_w = (size_t)m->num_attention_heads * head_w_stride;\n"
2373 " size_t kv_w = (size_t)m->num_kv_heads * head_w_stride;\n"
2374 " size_t q_b = (size_t)m->num_attention_heads * m->aligned_head_dim;\n"
2375 " size_t kv_b = (size_t)m->num_kv_heads * m->aligned_head_dim;\n"
2376 " size_t wo_w = (size_t)m->num_attention_heads * m->aligned_embed_dim * m->aligned_head_dim;\n"
2377 " size_t w1_w = (size_t)(2 * aligned_intermediate) * m->aligned_embed_dim;\n"
2378 " size_t w2_w = m->aligned_embed_dim * aligned_intermediate;\n"
2379 "\n"
2380 " if (write_floats_file(f, cptr_f32(base, L->ln1_gamma_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2381 " if (write_floats_file(f, cptr_f32(base, L->ln2_gamma_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2382 " if (write_weight_file(f, L->wq_dtype, cptr_void(base, L->wq_offset), q_w) != 0) goto fail;\n"
2383 " if (write_floats_file(f, cptr_f32(base, L->bq_offset), q_b) != 0) goto fail;\n"
2384 " if (write_weight_file(f, L->wk_dtype, cptr_void(base, L->wk_offset), kv_w) != 0) goto fail;\n"
2385 " if (write_floats_file(f, cptr_f32(base, L->bk_offset), kv_b) != 0) goto fail;\n"
2386 " if (write_weight_file(f, L->wv_dtype, cptr_void(base, L->wv_offset), kv_w) != 0) goto fail;\n"
2387 " if (write_floats_file(f, cptr_f32(base, L->bv_offset), kv_b) != 0) goto fail;\n"
2388 " if (write_weight_file(f, L->wo_dtype, cptr_void(base, L->wo_offset), wo_w) != 0) goto fail;\n"
2389 " if (write_floats_file(f, cptr_f32(base, L->bo_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2390 " if (write_weight_file(f, L->w1_dtype, cptr_void(base, L->w1_offset), w1_w) != 0) goto fail;\n"
2391 " if (write_floats_file(f, cptr_f32(base, L->b1_offset), (size_t)(2 * aligned_intermediate)) != 0) goto fail;\n"
2392 " if (write_weight_file(f, L->w2_dtype, cptr_void(base, L->w2_offset), w2_w) != 0) goto fail;\n"
2393 " if (write_floats_file(f, cptr_f32(base, L->b2_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2394 " }\n"
2395 "\n"
2396 " if (write_floats_file(f, cptr_f32(base, m->final_ln_weight_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2397 " if (write_floats_file(f, cptr_f32(base, m->final_ln_bias_offset), m->aligned_embed_dim) != 0) goto fail;\n"
2398 "\n"
2399 " fclose(f);\n"
2400 " return 0;\n"
2401 "fail:\n"
2402 " fclose(f);\n"
2403 " return -1;\n"
2404 "}\n\n"
2405 "static void embed_tokens(const TransformerModel *m, const int32_t *tokens, int token_count)\n"
2406 "{\n"
2407 " if (!m || !m->memory_base || !tokens) return;\n"
2408 " const uint8_t *base = m->memory_base;\n"
2409 " float *out = ptr_f32((uint8_t *)base, m->embedded_input_offset);\n"
2410 " const float *tok_f32 = cptr_f32(base, m->token_emb_offset);\n"
2411 " const uint8_t *tok_q = (const uint8_t *)cptr_void(base, m->token_emb_offset);\n"
2412 " const float *pos = cptr_f32(base, m->pos_emb_offset);\n"
2413 " int T = m->context_window;\n"
2414 " int D = m->embed_dim;\n"
2415 " int aligned_D = (int)m->aligned_embed_dim;\n"
2416 " for (int t = 0; t < T; ++t) {\n"
2417 " float *dst = out + (size_t)t * aligned_D;\n"
2418 " if (t < token_count) {\n"
2419 " int id = tokens[t];\n"
2420 " if (id < 0 || id >= m->vocab_size) id = 0;\n"
2421 " if (m->token_emb_dtype == CK_DT_Q4_K) {\n"
2422 " size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, (size_t)aligned_D);\n"
2423 " const void *row = tok_q + (size_t)id * row_bytes;\n"
2424 " dequant_q4_k_row(row, dst, (size_t)aligned_D);\n"
2425 " } else if (m->token_emb_dtype == CK_DT_Q6_K) {\n"
2426 " size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q6_K, (size_t)aligned_D);\n"
2427 " const void *row = tok_q + (size_t)id * row_bytes;\n"
2428 " dequant_q6_k_row(row, dst, (size_t)aligned_D);\n"
2429 " } else {\n"
2430 " const float *src = tok_f32 + (size_t)id * aligned_D;\n"
2431 " memcpy(dst, src, (size_t)D * sizeof(float));\n"
2432 " }\n"
2433 " if (aligned_D > D) {\n"
2434 " memset(dst + D, 0, (size_t)(aligned_D - D) * sizeof(float));\n"
2435 " }\n"
2436 " if (m->rope_theta <= 0.0f) {\n"
2437 " const float *p = pos + (size_t)t * aligned_D;\n"
2438 " for (int d = 0; d < D; ++d) {\n"
2439 " dst[d] += p[d];\n"
2440 " }\n"
2441 " }\n"
2442 " } else {\n"
2443 " memset(dst, 0, (size_t)aligned_D * sizeof(float));\n"
2444 " }\n"
2445 " }\n"
2446 "}\n\n"
2447 "static void embed_token_at(const TransformerModel *m, int32_t token, int t)\n"
2448 "{\n"
2449 " if (!m || !m->memory_base) return;\n"
2450 " if (t < 0 || t >= m->context_window) return;\n"
2451 " const uint8_t *base = m->memory_base;\n"
2452 " float *out = ptr_f32((uint8_t *)base, m->embedded_input_offset);\n"
2453 " const float *tok_f32 = cptr_f32(base, m->token_emb_offset);\n"
2454 " const uint8_t *tok_q = (const uint8_t *)cptr_void(base, m->token_emb_offset);\n"
2455 " const float *pos = cptr_f32(base, m->pos_emb_offset);\n"
2456 " int D = m->embed_dim;\n"
2457 " int aligned_D = (int)m->aligned_embed_dim;\n"
2458 " int id = (int)token;\n"
2459 " if (id < 0 || id >= m->vocab_size) id = 0;\n"
2460 " float *dst = out + (size_t)t * aligned_D;\n"
2461 " if (m->token_emb_dtype == CK_DT_Q4_K) {\n"
2462 " size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, (size_t)aligned_D);\n"
2463 " const void *row = tok_q + (size_t)id * row_bytes;\n"
2464 " dequant_q4_k_row(row, dst, (size_t)aligned_D);\n"
2465 " } else if (m->token_emb_dtype == CK_DT_Q6_K) {\n"
2466 " size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q6_K, (size_t)aligned_D);\n"
2467 " const void *row = tok_q + (size_t)id * row_bytes;\n"
2468 " dequant_q6_k_row(row, dst, (size_t)aligned_D);\n"
2469 " } else {\n"
2470 " const float *src = tok_f32 + (size_t)id * aligned_D;\n"
2471 " memcpy(dst, src, (size_t)D * sizeof(float));\n"
2472 " }\n"
2473 " if (aligned_D > D) {\n"
2474 " memset(dst + D, 0, (size_t)(aligned_D - D) * sizeof(float));\n"
2475 " }\n"
2476 " if (m->rope_theta <= 0.0f) {\n"
2477 " const float *p = pos + (size_t)t * aligned_D;\n"
2478 " for (int d = 0; d < D; ++d) {\n"
2479 " dst[d] += p[d];\n"
2480 " }\n"
2481 " }\n"
2482 "}\n\n"
2483 "static int write_floats(const char *path, const float *src, size_t count)\n"
2484 "{\n"
2485 " if (!path || !src) return -1;\n"
2486 " FILE *f = fopen(path, \"wb\");\n"
2487 " if (!f) {\n"
2488 " perror(\"fopen\");\n"
2489 " return -1;\n"
2490 " }\n"
2491 " size_t wrote = fwrite(src, sizeof(float), count, f);\n"
2492 " fclose(f);\n"
2493 " return wrote == count ? 0 : -1;\n"
2494 "}\n\n"
2495 "static int write_float_scalar(const char *path, float v)\n"
2496 "{\n"
2497 " if (!path) return -1;\n"
2498 " FILE *f = fopen(path, \"wb\");\n"
2499 " if (!f) {\n"
2500 " perror(\"fopen\");\n"
2501 " return -1;\n"
2502 " }\n"
2503 " size_t wrote = fwrite(&v, sizeof(float), 1, f);\n"
2504 " fclose(f);\n"
2505 " return wrote == 1 ? 0 : -1;\n"
2506 "}\n\n"
2507 "static void lm_head_forward(const float *hidden,\n"
2508 " const float *weights,\n"
2509 " float *logits,\n"
2510 " int T, int V, int D, int aligned_D)\n"
2511 "{\n"
2512 " for (int t = 0; t < T; ++t) {\n"
2513 " const float *h = hidden + (size_t)t * aligned_D;\n"
2514 " float *out = logits + (size_t)t * V;\n"
2515 " for (int v = 0; v < V; ++v) {\n"
2516 " const float *w = weights + (size_t)v * aligned_D;\n"
2517 " float sum = 0.0f;\n"
2518 " for (int d = 0; d < D; ++d) {\n"
2519 " sum += h[d] * w[d];\n"
2520 " }\n"
2521 " out[v] = sum;\n"
2522 " }\n"
2523 " }\n"
2524 "}\n\n"
2525 "static void softmax_cross_entropy(const float *logits,\n"
2526 " const int32_t *targets,\n"
2527 " int T, int V,\n"
2528 " float *d_logits,\n"
2529 " float *loss_out)\n"
2530 "{\n"
2531 " double total = 0.0;\n"
2532 " for (int t = 0; t < T; ++t) {\n"
2533 " const float *row = logits + (size_t)t * V;\n"
2534 " float *drow = d_logits + (size_t)t * V;\n"
2535 " int target = targets[t];\n"
2536 " float max_logit = row[0];\n"
2537 " for (int v = 1; v < V; ++v) {\n"
2538 " if (row[v] > max_logit) max_logit = row[v];\n"
2539 " }\n"
2540 " double sum_exp = 0.0;\n"
2541 " for (int v = 0; v < V; ++v) {\n"
2542 " drow[v] = expf(row[v] - max_logit);\n"
2543 " sum_exp += drow[v];\n"
2544 " }\n"
2545 " float inv_sum = 1.0f / (float)sum_exp;\n"
2546 " for (int v = 0; v < V; ++v) {\n"
2547 " drow[v] *= inv_sum;\n"
2548 " }\n"
2549 " double logsum = (double)max_logit + log(sum_exp);\n"
2550 " total += logsum - (double)row[target];\n"
2551 " drow[target] -= 1.0f;\n"
2552 " float scale = 1.0f / (float)T;\n"
2553 " for (int v = 0; v < V; ++v) {\n"
2554 " drow[v] *= scale;\n"
2555 " }\n"
2556 " }\n"
2557 " if (loss_out) {\n"
2558 " *loss_out = (float)(total / (double)T);\n"
2559 " }\n"
2560 "}\n\n"
2561 "static void lm_head_backward(const float *hidden,\n"
2562 " const float *weights,\n"
2563 " const float *d_logits,\n"
2564 " float *d_hidden,\n"
2565 " float *d_weights,\n"
2566 " int T, int V, int D, int aligned_D)\n"
2567 "{\n"
2568 " size_t dh_count = (size_t)T * aligned_D;\n"
2569 " size_t dw_count = (size_t)V * aligned_D;\n"
2570 " for (size_t i = 0; i < dh_count; ++i) d_hidden[i] = 0.0f;\n"
2571 " for (size_t i = 0; i < dw_count; ++i) d_weights[i] = 0.0f;\n"
2572 " for (int t = 0; t < T; ++t) {\n"
2573 " const float *dlog = d_logits + (size_t)t * V;\n"
2574 " for (int d = 0; d < D; ++d) {\n"
2575 " double sum = 0.0;\n"
2576 " for (int v = 0; v < V; ++v) {\n"
2577 " sum += (double)dlog[v] * (double)weights[(size_t)v * aligned_D + d];\n"
2578 " }\n"
2579 " d_hidden[(size_t)t * aligned_D + d] = (float)sum;\n"
2580 " }\n"
2581 " }\n"
2582 " for (int v = 0; v < V; ++v) {\n"
2583 " float *dw = d_weights + (size_t)v * aligned_D;\n"
2584 " for (int d = 0; d < D; ++d) {\n"
2585 " double sum = 0.0;\n"
2586 " for (int t = 0; t < T; ++t) {\n"
2587 " sum += (double)d_logits[(size_t)t * V + v] * (double)hidden[(size_t)t * aligned_D + d];\n"
2588 " }\n"
2589 " dw[d] = (float)sum;\n"
2590 " }\n"
2591 " }\n"
2592 "}\n\n");
2593
2594 fprintf(out,
2595 "static void dump_layer_offsets(const TransformerModel *m, int layer)\n"
2596 "{\n"
2597 " const TrulyOptimalLayer *L = &m->layers[layer];\n"
2598 " printf(\"Layer %%d offsets (bytes):\\n\", layer);\n"
2599 " printf(\" ln1_gamma=%%zu ln2_gamma=%%zu wq=%%zu wk=%%zu wv=%%zu wo=%%zu w1=%%zu w2=%%zu\\n\",\n"
2600 " L->ln1_gamma_offset, L->ln2_gamma_offset, L->wq_offset, L->wk_offset,\n"
2601 " L->wv_offset, L->wo_offset, L->w1_offset, L->w2_offset);\n"
2602 " printf(\" ln1_out=%%zu q=%%zu k=%%zu v=%%zu scores=%%zu attn_out=%%zu\\n\",\n"
2603 " L->ln1_out_offset, L->q_offset, L->k_offset, L->v_offset,\n"
2604 " L->scores_offset, L->attn_out_offset);\n"
2605 " printf(\" proj_tmp=%%zu residual1=%%zu ln2_out=%%zu fc1_out=%%zu swiglu_out=%%zu mlp_out=%%zu output=%%zu\\n\",\n"
2606 " L->proj_tmp_offset, L->residual1_offset, L->ln2_out_offset,\n"
2607 " L->fc1_out_offset, L->swiglu_out_offset, L->mlp_out_offset, L->output_offset);\n"
2608 "}\n\n"
2609 "static void dump_layout(const TransformerModel *m, int dump_all)\n"
2610 "{\n"
2611 " size_t bytes = m->total_bytes;\n"
2612 " printf(\"Model config:\\n\");\n"
2613 " printf(\" layers=%%d embed=%%d intermediate=%%d heads=%%d kv_heads=%%d\\n\",\n"
2614 " m->num_layers, m->embed_dim, m->intermediate_size, m->num_attention_heads, m->num_kv_heads);\n"
2615 " printf(\" head_dim=%%d vocab=%%d ctx=%%d cores=%%d\\n\",\n"
2616 " m->head_dim, m->vocab_size, m->context_window, m->num_cores);\n"
2617 " printf(\" eps=%%.6g rope_theta=%%.6g\\n\", m->rms_norm_eps, m->rope_theta);\n"
2618 " printf(\"Aligned dims (elements): embed=%%zu head=%%zu ctx=%%zu\\n\",\n"
2619 " m->aligned_embed_dim, m->aligned_head_dim, m->aligned_attn_context_window);\n"
2620 " printf(\"Memory: total_bytes=%%zu\\n\", bytes);\n"
2621 " printf(\"Global offsets (bytes): token=%%zu pos=%%zu embedded=%%zu layers_start=%%zu\\n\",\n"
2622 " m->token_emb_offset, m->pos_emb_offset, m->embedded_input_offset, m->layers_start_offset);\n"
2623 " printf(\"Final offsets (bytes): final_ln_w=%%zu final_ln_b=%%zu final_ln_mean=%%zu final_ln_rstd=%%zu\\n\",\n"
2624 " m->final_ln_weight_offset, m->final_ln_bias_offset,\n"
2625 " m->final_ln_mean_offset, m->final_ln_rstd_offset);\n"
2626 " printf(\"LM/logits offsets (bytes): lm_head=%%zu logits=%%zu\\n\",\n"
2627 " m->lm_head_weight_offset, m->logits_offset);\n"
2628 " if (m->num_layers > 0) {\n"
2629 " dump_layer_offsets(m, 0);\n"
2630 " if (dump_all) {\n"
2631 " for (int i = 1; i < m->num_layers; ++i) {\n"
2632 " dump_layer_offsets(m, i);\n"
2633 " }\n"
2634 " }\n"
2635 " }\n"
2636 "}\n\n");
2637
2638 /* Emit either main() for standalone or API for library mode */
2639 if (mode == CK_EMIT_STANDALONE) {
2640 fprintf(out,
2641 "int main(int argc, char **argv)\n"
2642 "{\n"
2643 " int dump = 0;\n"
2644 " int dump_all = 0;\n"
2645 " int no_forward = 0;\n"
2646 " int run_litmus = 0;\n"
2647 " int run_backward = 0;\n"
2648 " const char *litmus_hidden = NULL;\n"
2649 " const char *litmus_weights = NULL;\n"
2650 " const char *litmus_targets = NULL;\n"
2651 " const char *model_weights = NULL;\n"
2652 " const char *tokens_path = NULL;\n"
2653 " const char *out_logits = NULL;\n"
2654 " const char *out_dlogits = NULL;\n"
2655 " const char *out_dhidden = NULL;\n"
2656 " const char *out_dweights = NULL;\n"
2657 " const char *out_loss = NULL;\n"
2658 " const char *out_weights = NULL;\n"
2659 " int steps = 1;\n"
2660 " int log_steps = 0;\n"
2661 " int strict = 0;\n"
2662 " int32_t *tokens = NULL;\n"
2663 " int32_t *targets = NULL;\n"
2664 " TransformerModel m = {0};\n"
2665 " memcpy(m.magic, \"BUMPWGT3\", 8);\n"
2666 " m.version = 3;\n"
2667 " m.model_type = 0;\n"
2668 " m.num_layers = %d;\n"
2669 " m.embed_dim = %d;\n"
2670 " m.intermediate_size = %d;\n"
2671 " m.num_attention_heads = %d;\n"
2672 " m.num_kv_heads = %d;\n"
2673 " m.vocab_size = %d;\n"
2674 " m.context_window = %d;\n"
2675 " m.rms_norm_eps = %.9g;\n"
2676 " m.rope_theta = %.9g;\n"
2677 " m.num_cores = 1;\n"
2678 " m.task_type = TASK_LM;\n"
2679 " m.optimizer = OPTIMIZER_SGD;\n"
2680 " m.learning_rate = 0.0f;\n"
2681 " for (int i = 1; i < argc; ++i) {\n"
2682 " if (strcmp(argv[i], \"--dump\") == 0) {\n"
2683 " dump = 1;\n"
2684 " continue;\n"
2685 " }\n"
2686 " if (strcmp(argv[i], \"--dump-all\") == 0) {\n"
2687 " dump = 1;\n"
2688 " dump_all = 1;\n"
2689 " continue;\n"
2690 " }\n"
2691 " if (strcmp(argv[i], \"--no-forward\") == 0) {\n"
2692 " no_forward = 1;\n"
2693 " continue;\n"
2694 " }\n"
2695 " if (strcmp(argv[i], \"--strict\") == 0) {\n"
2696 " strict = 1;\n"
2697 " continue;\n"
2698 " }\n"
2699 " if (strcmp(argv[i], \"--litmus\") == 0) {\n"
2700 " run_litmus = 1;\n"
2701 " continue;\n"
2702 " }\n"
2703 " if (strcmp(argv[i], \"--backward\") == 0) {\n"
2704 " run_backward = 1;\n"
2705 " continue;\n"
2706 " }\n"
2707 " if (strcmp(argv[i], \"--lr\") == 0 && i + 1 < argc) {\n"
2708 " parse_float_arg(argv[++i], &m.learning_rate);\n"
2709 " continue;\n"
2710 " }\n"
2711 " if (strcmp(argv[i], \"--help\") == 0) {\n"
2712 " print_usage(argv[0]);\n"
2713 " return 0;\n"
2714 " }\n"
2715 " if (strcmp(argv[i], \"--hidden\") == 0 && i + 1 < argc) {\n"
2716 " litmus_hidden = argv[++i];\n"
2717 " continue;\n"
2718 " }\n"
2719 " if (strcmp(argv[i], \"--weights\") == 0 && i + 1 < argc) {\n"
2720 " litmus_weights = argv[++i];\n"
2721 " continue;\n"
2722 " }\n"
2723 " if (strcmp(argv[i], \"--targets\") == 0 && i + 1 < argc) {\n"
2724 " litmus_targets = argv[++i];\n"
2725 " continue;\n"
2726 " }\n"
2727 " if (strcmp(argv[i], \"--model-weights\") == 0 && i + 1 < argc) {\n"
2728 " model_weights = argv[++i];\n"
2729 " continue;\n"
2730 " }\n"
2731 " if (strcmp(argv[i], \"--tokens\") == 0 && i + 1 < argc) {\n"
2732 " tokens_path = argv[++i];\n"
2733 " continue;\n"
2734 " }\n"
2735 " if (strcmp(argv[i], \"--out-logits\") == 0 && i + 1 < argc) {\n"
2736 " out_logits = argv[++i];\n"
2737 " continue;\n"
2738 " }\n"
2739 " if (strcmp(argv[i], \"--out-dlogits\") == 0 && i + 1 < argc) {\n"
2740 " out_dlogits = argv[++i];\n"
2741 " continue;\n"
2742 " }\n"
2743 " if (strcmp(argv[i], \"--out-dhidden\") == 0 && i + 1 < argc) {\n"
2744 " out_dhidden = argv[++i];\n"
2745 " continue;\n"
2746 " }\n"
2747 " if (strcmp(argv[i], \"--out-dweights\") == 0 && i + 1 < argc) {\n"
2748 " out_dweights = argv[++i];\n"
2749 " continue;\n"
2750 " }\n"
2751 " if (strcmp(argv[i], \"--out-loss\") == 0 && i + 1 < argc) {\n"
2752 " out_loss = argv[++i];\n"
2753 " continue;\n"
2754 " }\n"
2755 " if (strcmp(argv[i], \"--out-weights\") == 0 && i + 1 < argc) {\n"
2756 " out_weights = argv[++i];\n"
2757 " continue;\n"
2758 " }\n"
2759 " if (strcmp(argv[i], \"--steps\") == 0 && i + 1 < argc) {\n"
2760 " parse_int_arg(argv[++i], &steps);\n"
2761 " continue;\n"
2762 " }\n"
2763 " if (strcmp(argv[i], \"--log-steps\") == 0) {\n"
2764 " log_steps = 1;\n"
2765 " continue;\n"
2766 " }\n"
2767 " if (strcmp(argv[i], \"--layers\") == 0 && i + 1 < argc) {\n"
2768 " parse_int_arg(argv[++i], &m.num_layers);\n"
2769 " continue;\n"
2770 " }\n"
2771 " if (strcmp(argv[i], \"--embed\") == 0 && i + 1 < argc) {\n"
2772 " parse_int_arg(argv[++i], &m.embed_dim);\n"
2773 " continue;\n"
2774 " }\n"
2775 " if (strcmp(argv[i], \"--intermediate\") == 0 && i + 1 < argc) {\n"
2776 " parse_int_arg(argv[++i], &m.intermediate_size);\n"
2777 " continue;\n"
2778 " }\n"
2779 " if (strcmp(argv[i], \"--heads\") == 0 && i + 1 < argc) {\n"
2780 " parse_int_arg(argv[++i], &m.num_attention_heads);\n"
2781 " continue;\n"
2782 " }\n"
2783 " if (strcmp(argv[i], \"--kv-heads\") == 0 && i + 1 < argc) {\n"
2784 " parse_int_arg(argv[++i], &m.num_kv_heads);\n"
2785 " continue;\n"
2786 " }\n"
2787 " if (strcmp(argv[i], \"--vocab\") == 0 && i + 1 < argc) {\n"
2788 " parse_int_arg(argv[++i], &m.vocab_size);\n"
2789 " continue;\n"
2790 " }\n"
2791 " if (strcmp(argv[i], \"--ctx\") == 0 && i + 1 < argc) {\n"
2792 " parse_int_arg(argv[++i], &m.context_window);\n"
2793 " continue;\n"
2794 " }\n"
2795 " if (strcmp(argv[i], \"--cores\") == 0 && i + 1 < argc) {\n"
2796 " parse_int_arg(argv[++i], &m.num_cores);\n"
2797 " continue;\n"
2798 " }\n"
2799 " fprintf(stderr, \"Unknown or invalid arg: %%s\\n\", argv[i]);\n"
2800 " print_usage(argv[0]);\n"
2801 " return 1;\n"
2802 " }\n"
2803 " if (strict) {\n"
2804 " ck_set_strict_parity(1);\n"
2805 " }\n"
2806 " if (run_backward && m.learning_rate == 0.0f) {\n"
2807 " m.learning_rate = 1e-3f;\n"
2808 " }\n"
2809 " m.training_enabled = run_backward;\n"
2810 " m.weight_dtype = CK_DT_FP32;\n"
2811 " {\n"
2812 " const char *wd = getenv(\"CK_WEIGHT_DTYPE\");\n"
2813 " if (wd) {\n"
2814 " if (strcmp(wd, \"q4_k\") == 0 || strcmp(wd, \"q4_k_m\") == 0 ||\n"
2815 " strcmp(wd, \"Q4_K\") == 0 || strcmp(wd, \"Q4_K_M\") == 0) {\n"
2816 " m.weight_dtype = CK_DT_Q4_K;\n"
2817 " } else if (strcmp(wd, \"q6_k\") == 0 || strcmp(wd, \"q6_k_l\") == 0 ||\n"
2818 " strcmp(wd, \"Q6_K\") == 0 || strcmp(wd, \"Q6_K_L\") == 0) {\n"
2819 " m.weight_dtype = CK_DT_Q6_K;\n"
2820 " }\n"
2821 " }\n"
2822 " }\n"
2823 " init_weight_dtypes_uniform(&m, m.weight_dtype);\n"
2824 " refresh_weight_flags(&m);\n"
2825 " if (model_weights) {\n"
2826 " int dtype_rc = load_weight_dtypes(model_weights, &m);\n"
2827 " if (dtype_rc < 0) {\n"
2828 " fprintf(stderr, \"failed to read weight dtype table\\n\");\n"
2829 " return 1;\n"
2830 " }\n"
2831 " }\n"
2832 " if (m.training_enabled && m.weights_quantized) {\n"
2833 " fprintf(stderr, \"Quantized weights are inference-only; disable training\\n\");\n"
2834 " return 1;\n"
2835 " }\n"
2836 " if (layout_model(&m) != 0) {\n"
2837 " fprintf(stderr, \"layout_model failed\\n\");\n"
2838 " return 1;\n"
2839 " }\n"
2840 " if (model_weights) {\n"
2841 " if (load_model_weights(model_weights, &m) != 0) {\n"
2842 " fprintf(stderr, \"failed to load model weights\\n\");\n"
2843 " return 1;\n"
2844 " }\n"
2845 " }\n"
2846 " if (tokens_path) {\n"
2847 " int T = m.context_window;\n"
2848 " tokens = (int32_t *)malloc((size_t)T * sizeof(int32_t));\n"
2849 " if (!tokens) {\n"
2850 " fprintf(stderr, \"failed to alloc tokens\\n\");\n"
2851 " return 1;\n"
2852 " }\n"
2853 " if (read_ints(tokens_path, tokens, (size_t)T) != 0) {\n"
2854 " fprintf(stderr, \"failed to read tokens\\n\");\n"
2855 " free(tokens);\n"
2856 " tokens = NULL;\n"
2857 " return 1;\n"
2858 " }\n"
2859 " if (!run_backward) {\n"
2860 " embed_tokens(&m, tokens, T);\n"
2861 " free(tokens);\n"
2862 " tokens = NULL;\n"
2863 " }\n"
2864 " }\n"
2865 " if (run_backward) {\n"
2866 " if (!litmus_targets) {\n"
2867 " fprintf(stderr, \"backward requires --targets\\n\");\n"
2868 " return 1;\n"
2869 " }\n"
2870 " int T = m.context_window;\n"
2871 " targets = (int32_t *)malloc((size_t)T * sizeof(int32_t));\n"
2872 " if (!targets) {\n"
2873 " fprintf(stderr, \"failed to alloc targets\\n\");\n"
2874 " return 1;\n"
2875 " }\n"
2876 " if (read_ints(litmus_targets, targets, (size_t)T) != 0) {\n"
2877 " fprintf(stderr, \"failed to read targets\\n\");\n"
2878 " free(targets);\n"
2879 " targets = NULL;\n"
2880 " return 1;\n"
2881 " }\n"
2882 " }\n"
2883 " if (dump) {\n"
2884 " dump_layout(&m, dump_all);\n"
2885 " }\n"
2886 " if (run_litmus) {\n"
2887 " if (!litmus_hidden || !litmus_weights || !litmus_targets) {\n"
2888 " fprintf(stderr, \"litmus requires --hidden, --weights, and --targets\\n\");\n"
2889 " return 1;\n"
2890 " }\n"
2891 " int T = m.context_window;\n"
2892 " int V = m.vocab_size;\n"
2893 " int D = m.embed_dim;\n"
2894 " int aligned_D = (int)m.aligned_embed_dim;\n"
2895 " float *hidden = ptr_f32(m.memory_base, m.final_output_offset);\n"
2896 " float *weights = ptr_f32(m.memory_base, m.lm_head_weight_offset);\n"
2897 " float *logits = ptr_f32(m.memory_base, m.logits_offset);\n"
2898 " if (read_floats(litmus_hidden, hidden, (size_t)T * aligned_D) != 0) {\n"
2899 " fprintf(stderr, \"failed to read hidden\\n\");\n"
2900 " return 1;\n"
2901 " }\n"
2902 " if (read_floats(litmus_weights, weights, (size_t)V * aligned_D) != 0) {\n"
2903 " fprintf(stderr, \"failed to read weights\\n\");\n"
2904 " return 1;\n"
2905 " }\n"
2906 " int32_t *targets = (int32_t *)malloc((size_t)T * sizeof(int32_t));\n"
2907 " if (!targets) {\n"
2908 " fprintf(stderr, \"failed to alloc targets\\n\");\n"
2909 " return 1;\n"
2910 " }\n"
2911 " if (read_ints(litmus_targets, targets, (size_t)T) != 0) {\n"
2912 " fprintf(stderr, \"failed to read targets\\n\");\n"
2913 " free(targets);\n"
2914 " return 1;\n"
2915 " }\n"
2916 " float *d_logits = (float *)calloc((size_t)T * V, sizeof(float));\n"
2917 " float *d_hidden = (float *)calloc((size_t)T * aligned_D, sizeof(float));\n"
2918 " float *d_weights = (float *)calloc((size_t)V * aligned_D, sizeof(float));\n"
2919 " if (!d_logits || !d_hidden || !d_weights) {\n"
2920 " fprintf(stderr, \"failed to alloc grads\\n\");\n"
2921 " free(targets);\n"
2922 " free(d_logits);\n"
2923 " free(d_hidden);\n"
2924 " free(d_weights);\n"
2925 " return 1;\n"
2926 " }\n"
2927 " lm_head_forward(hidden, weights, logits, T, V, D, aligned_D);\n"
2928 " float loss = 0.0f;\n"
2929 " softmax_cross_entropy(logits, targets, T, V, d_logits, &loss);\n"
2930 " lm_head_backward(hidden, weights, d_logits, d_hidden, d_weights, T, V, D, aligned_D);\n"
2931 " if (out_logits) write_floats(out_logits, logits, (size_t)T * V);\n"
2932 " if (out_dlogits) write_floats(out_dlogits, d_logits, (size_t)T * V);\n"
2933 " if (out_dhidden) write_floats(out_dhidden, d_hidden, (size_t)T * aligned_D);\n"
2934 " if (out_dweights) write_floats(out_dweights, d_weights, (size_t)V * aligned_D);\n"
2935 " if (out_loss) write_float_scalar(out_loss, loss);\n"
2936 " if (!out_loss) printf(\"loss=%%.6f\\n\", loss);\n"
2937 " free(targets);\n"
2938 " free(d_logits);\n"
2939 " free(d_hidden);\n"
2940 " free(d_weights);\n"
2941 " ck_huge_free(m.memory_base, m.total_bytes);\n"
2942 " free(m.layers);\n"
2943 " return 0;\n"
2944 " }\n"
2945 " // TODO: load weights into m.memory_base using the offsets above.\n"
2946 " // TODO: write token/pos embeddings into embedded_input_offset.\n"
2947 " if (!run_backward) {\n"
2948 " if (!no_forward) {\n"
2949 " run_model_forward(&m);\n"
2950 " }\n"
2951 " } else {\n"
2952 " if (!tokens || !targets) {\n"
2953 " fprintf(stderr, \"backward requires --tokens and --targets\\n\");\n"
2954 " return 1;\n"
2955 " }\n"
2956 " if (steps < 1) steps = 1;\n"
2957 " float loss = 0.0f;\n"
2958 " for (int step = 0; step < steps; ++step) {\n"
2959 " embed_tokens(&m, tokens, m.context_window);\n"
2960 " run_model_forward(&m);\n"
2961 " if (run_model_backward(&m, tokens, targets, &loss) != 0) {\n"
2962 " fprintf(stderr, \"backward failed\\n\");\n"
2963 " return 1;\n"
2964 " }\n"
2965 " if (log_steps) {\n"
2966 " printf(\"step %%d loss=%%.6f\\n\", step, loss);\n"
2967 " }\n"
2968 " }\n"
2969 " if (out_loss) {\n"
2970 " write_float_scalar(out_loss, loss);\n"
2971 " }\n"
2972 " }\n"
2973 " if (out_logits) {\n"
2974 " write_floats(out_logits, ptr_f32(m.memory_base, m.logits_offset),\n"
2975 " (size_t)m.context_window * (size_t)m.vocab_size);\n"
2976 " }\n"
2977 " if (out_weights) {\n"
2978 " if (save_model_weights(out_weights, &m) != 0) {\n"
2979 " fprintf(stderr, \"failed to save model weights\\n\");\n"
2980 " return 1;\n"
2981 " }\n"
2982 " }\n"
2983 " ck_huge_free(m.memory_base, m.total_bytes);\n"
2984 " free(m.layers);\n"
2985 " free(tokens);\n"
2986 " free(targets);\n"
2987 " return 0;\n"
2988 "}\n",
2989 forward->config.num_layers,
2990 forward->config.hidden_size,
2991 forward->config.intermediate_size,
2992 forward->config.num_heads,
2993 forward->config.num_kv_heads,
2994 forward->config.vocab_size,
2995 forward->config.context_window,
2996 forward->config.rms_norm_eps,
2997 forward->config.rope_theta);
2998 } else {
2999 /* Library mode - emit API functions instead of main() */
3000 emit_library_api(out, forward);
3001 }
3002
3003 fclose(out);
3004 if (emit_kernel_manifest(forward, path) != 0) {
3005 return -1;
3006 }
3007 return 0;
3008}
emit_shape_expr
static void emit_shape_expr(FILE *out, const CKDimToken *shape)
Definition ckernel_codegen.c:274
op_name
static const char * op_name(CKOpType op)
Definition ckernel_codegen.c:35
ck_buffer_uses_weight_dtype
static int ck_buffer_uses_weight_dtype(const CKBufferSpec *spec)
Definition ckernel_codegen.c:87
emit_runtime_preamble
static int emit_runtime_preamble(FILE *out)
Definition ckernel_codegen.c:783
ck_codegen_emit_runtime
int ck_codegen_emit_runtime(const CKIRGraph *forward, const char *path, CKEmitMode mode)
Definition ckernel_codegen.c:1468
emit_global_offset_fields
static void emit_global_offset_fields(FILE *out)
Definition ckernel_codegen.c:159
ck_first_layer_buffer_name
static const char * ck_first_layer_buffer_name(void)
Definition ckernel_codegen.c:623
emit_offset_field
static void emit_offset_field(FILE *out, const char *name)
Definition ckernel_codegen.c:131
emit_kernel_manifest
static int emit_kernel_manifest(const CKIRGraph *forward, const char *runtime_path)
Definition ckernel_codegen.c:845
ck_plan_step_enabled
static int ck_plan_step_enabled(const CKPlanStep *step, const CKIRGraph *cfg)
Definition ckernel_codegen.c:579
emit_library_api
static void emit_library_api(FILE *out, const CKIRGraph *forward)
Definition ckernel_codegen.c:945
emit_layer_allocations
static void emit_layer_allocations(FILE *out)
Definition ckernel_codegen.c:388
emit_sgd_update
static void emit_sgd_update(FILE *out)
Definition ckernel_codegen.c:484
ck_find_kernel_spec
static const CKKernelSpec * ck_find_kernel_spec(const char *name)
Definition ckernel_codegen.c:69
emit_dim_expr
static void emit_dim_expr(FILE *out, CKDimKind dim)
Definition ckernel_codegen.c:256
emit_plan_sources
static int emit_plan_sources(FILE *f, const CKPlanStep *plan, size_t plan_count, const CKIRGraph *cfg, const char **seen, size_t *seen_count, size_t seen_cap)
Definition ckernel_codegen.c:593
ck_weight_dtype_expr
static const char * ck_weight_dtype_expr(const CKBufferSpec *spec)
Definition ckernel_codegen.c:105
emit_model_struct
static void emit_model_struct(FILE *out)
Definition ckernel_codegen.c:173
emit_global_allocations
static void emit_global_allocations(FILE *out)
Definition ckernel_codegen.c:336
ck_buffer_should_alloc
static int ck_buffer_should_alloc(const CKBufferSpec *spec)
Definition ckernel_codegen.c:82
emit_zero_grad
static void emit_zero_grad(FILE *out)
Definition ckernel_codegen.c:444
emit_unique_source
static int emit_unique_source(FILE *f, const char *path, const char **seen, size_t *seen_count, size_t seen_cap)
Definition ckernel_codegen.c:555
ck_codegen_c_skeleton
void ck_codegen_c_skeleton(const CKIRGraph *forward, const CKIRGraph *backward, FILE *out)
Definition ckernel_codegen.c:638
ck_find_buffer_spec
static const CKBufferSpec * ck_find_buffer_spec(const char *name)
Definition ckernel_codegen.c:56
emit_bump_bytes_assignment
static void emit_bump_bytes_assignment(FILE *out, const char *indent, const char *struct_prefix, const char *name, const CKDimToken *shape)
Definition ckernel_codegen.c:300
emit_layer_offsets_struct
static void emit_layer_offsets_struct(FILE *out)
Definition ckernel_codegen.c:136
emit_training_conditional_assignment
static void emit_training_conditional_assignment(FILE *out, const char *indent, const char *struct_prefix, const char *name, const CKDimToken *shape)
Definition ckernel_codegen.c:324
emit_bump_bytes_assignment_weight_dtype
static void emit_bump_bytes_assignment_weight_dtype(FILE *out, const char *indent, const char *struct_prefix, const char *name, const CKDimToken *shape, const char *dtype_expr)
Definition ckernel_codegen.c:311
emit_global_aliases_to_layer
static void emit_global_aliases_to_layer(FILE *out)
Definition ckernel_codegen.c:423
ckernel_codegen.h
CKEmitMode
CKEmitMode
Definition ckernel_codegen.h:11
CK_EMIT_STANDALONE
@ CK_EMIT_STANDALONE
Definition ckernel_codegen.h:12
CKOpType
CKOpType
Definition ckernel_ir.h:38
CK_OP_LINEAR_BWD
@ CK_OP_LINEAR_BWD
Definition ckernel_ir.h:51
CK_OP_SWIGLU
@ CK_OP_SWIGLU
Definition ckernel_ir.h:45
CK_OP_RMSNORM_BWD
@ CK_OP_RMSNORM_BWD
Definition ckernel_ir.h:47
CK_OP_SWIGLU_BWD
@ CK_OP_SWIGLU_BWD
Definition ckernel_ir.h:53
CK_OP_ADD
@ CK_OP_ADD
Definition ckernel_ir.h:42
CK_OP_SPLIT
@ CK_OP_SPLIT
Definition ckernel_ir.h:44
CK_OP_LINEAR_QKV_BWD
@ CK_OP_LINEAR_QKV_BWD
Definition ckernel_ir.h:48
CK_OP_ATTENTION_BWD
@ CK_OP_ATTENTION_BWD
Definition ckernel_ir.h:49
CK_OP_SPLIT_BWD
@ CK_OP_SPLIT_BWD
Definition ckernel_ir.h:52
CK_OP_LINEAR_QKV
@ CK_OP_LINEAR_QKV
Definition ckernel_ir.h:40
CK_OP_LINEAR
@ CK_OP_LINEAR
Definition ckernel_ir.h:43
CK_OP_RMSNORM
@ CK_OP_RMSNORM
Definition ckernel_ir.h:39
CK_OP_ADD_BWD
@ CK_OP_ADD_BWD
Definition ckernel_ir.h:50
CK_OP_ATTENTION
@ CK_OP_ATTENTION
Definition ckernel_ir.h:41
ckernel_kernel_specs.h
CKERNEL_MAX_KERNEL_SOURCES
#define CKERNEL_MAX_KERNEL_SOURCES
Definition ckernel_kernel_specs.h:23
CK_ROLE_WEIGHT
@ CK_ROLE_WEIGHT
Definition ckernel_kernel_specs.h:40
CK_ROLE_GRAD
@ CK_ROLE_GRAD
Definition ckernel_kernel_specs.h:42
CK_ROLE_INPUT
@ CK_ROLE_INPUT
Definition ckernel_kernel_specs.h:37
ck_decoder_forward_plan
const CKPlanStep ck_decoder_forward_plan[]
Definition ckernel_kernel_specs.c:99
CKDimKind
CKDimKind
Definition ckernel_kernel_specs.h:8
CK_DIM_ALIGNED_INTERMEDIATE
@ CK_DIM_ALIGNED_INTERMEDIATE
Definition ckernel_kernel_specs.h:18
CK_DIM_NUM_HEADS
@ CK_DIM_NUM_HEADS
Definition ckernel_kernel_specs.h:14
CK_DIM_ALIGNED_EMBED
@ CK_DIM_ALIGNED_EMBED
Definition ckernel_kernel_specs.h:11
CK_DIM_TOKENS
@ CK_DIM_TOKENS
Definition ckernel_kernel_specs.h:9
CK_DIM_INTERMEDIATE
@ CK_DIM_INTERMEDIATE
Definition ckernel_kernel_specs.h:17
CK_DIM_ALIGNED_CTX
@ CK_DIM_ALIGNED_CTX
Definition ckernel_kernel_specs.h:16
CK_DIM_END
@ CK_DIM_END
Definition ckernel_kernel_specs.h:20
CK_DIM_ALIGNED_HEAD
@ CK_DIM_ALIGNED_HEAD
Definition ckernel_kernel_specs.h:13
CK_DIM_HEAD_DIM
@ CK_DIM_HEAD_DIM
Definition ckernel_kernel_specs.h:12
CK_DIM_NUM_KV_HEADS
@ CK_DIM_NUM_KV_HEADS
Definition ckernel_kernel_specs.h:15
CK_DIM_VOCAB
@ CK_DIM_VOCAB
Definition ckernel_kernel_specs.h:19
CK_DIM_EMBED
@ CK_DIM_EMBED
Definition ckernel_kernel_specs.h:10
ck_decoder_backward_plan_count
const size_t ck_decoder_backward_plan_count
Definition ckernel_kernel_specs.c:129
ck_decoder_forward_plan_count
const size_t ck_decoder_forward_plan_count
Definition ckernel_kernel_specs.c:113
ck_decoder_backward_plan
const CKPlanStep ck_decoder_backward_plan[]
Definition ckernel_kernel_specs.c:115
ck_kernel_specs
const CKKernelSpec ck_kernel_specs[]
Definition ckernel_kernel_specs.c:85
ck_decoder_buffers
const CKBufferSpec ck_decoder_buffers[]
Definition ckernel_kernel_specs.c:3
ck_kernel_spec_count
const size_t ck_kernel_spec_count
Definition ckernel_kernel_specs.c:97
ck_decoder_buffer_count
const size_t ck_decoder_buffer_count
Definition ckernel_kernel_specs.c:83
CK_SCOPE_LAYER
@ CK_SCOPE_LAYER
Definition ckernel_kernel_specs.h:32
CK_SCOPE_GLOBAL
@ CK_SCOPE_GLOBAL
Definition ckernel_kernel_specs.h:33
ckernel_registry.h
ck_ir_validate_supported
int ck_ir_validate_supported(const CKIRGraph *graph)
Definition ckernel_registry.c:42
CKBufferSpec::alias_of
const char * alias_of
Definition ckernel_kernel_specs.h:51
CKBufferSpec::condition
const char * condition
Definition ckernel_kernel_specs.h:52
CKBufferSpec::scope
CKBufferScope scope
Definition ckernel_kernel_specs.h:47
CKBufferSpec::shape
CKDimToken shape[4]
Definition ckernel_kernel_specs.h:49
CKIRGraph::nodes
CKIRNode * nodes
Definition ckernel_ir.h:78
CKIRGraph::num_nodes
int num_nodes
Definition ckernel_ir.h:77
CKIRGraph::config
CKModelConfig config
Definition ckernel_ir.h:76
CKIRNode::op
CKOpType op
Definition ckernel_ir.h:68
CKIRNode::inputs
CKInputRef inputs[4]
Definition ckernel_ir.h:69
CKIRNode::id
CKKernelId id
Definition ckernel_ir.h:67
CKIRNode::n_inputs
uint8_t n_inputs
Definition ckernel_ir.h:70
CKIRNode::n_outputs
uint8_t n_outputs
Definition ckernel_ir.h:71
CKInputRef::producer
CKKernelId producer
Definition ckernel_ir.h:62
CKInputRef::out_index
uint8_t out_index
Definition ckernel_ir.h:63
CKKernelId::node
uint16_t node
Definition ckernel_ir.h:58
CKKernelId::layer
uint16_t layer
Definition ckernel_ir.h:57
CKKernelSpec::sources
const char * sources[8]
Definition ckernel_kernel_specs.h:62
CKModelConfig::rms_norm_eps
float rms_norm_eps
Definition ckernel_ir.h:31
CKModelConfig::rope_theta
float rope_theta
Definition ckernel_ir.h:32
CKPlanStep::kernel
const char * kernel
Definition ckernel_kernel_specs.h:66
CKPlanStep::condition
const char * condition
Definition ckernel_kernel_specs.h:67