← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
Enumerations | Functions
ckernel_codegen.h File Reference
#include "ckernel_ir.h"
#include <stdio.h>

Go to the source code of this file.

Enumerations

enum  CKEmitMode { CK_EMIT_STANDALONE = 0 , CK_EMIT_LIBRARY = 1 }
 

Functions

void ck_codegen_c_skeleton (const CKIRGraph *forward, const CKIRGraph *backward, FILE *out)
 
int ck_codegen_emit_runtime (const CKIRGraph *forward, const char *path, CKEmitMode mode)
 

Enumeration Type Documentation

◆ CKEmitMode

enum CKEmitMode

Code generation output mode.

Enumerator
CK_EMIT_STANDALONE 
CK_EMIT_LIBRARY 

Definition at line 11 of file ckernel_codegen.h.

11 {
12 CK_EMIT_STANDALONE = 0, /* Emit with main() for standalone executable */
13 CK_EMIT_LIBRARY = 1, /* Emit as library with API functions, no main() */
14} CKEmitMode;
CKEmitMode
CKEmitMode
Definition ckernel_codegen.h:11
CK_EMIT_STANDALONE
@ CK_EMIT_STANDALONE
Definition ckernel_codegen.h:12
CK_EMIT_LIBRARY
@ CK_EMIT_LIBRARY
Definition ckernel_codegen.h:13

Function Documentation

◆ ck_codegen_c_skeleton()

void ck_codegen_c_skeleton ( const CKIRGraph *  forward,
const CKIRGraph *  backward,
FILE *  out 
)

Emit a C skeleton for forward + backward execution based on the IR.

This does not yet generate full pointer arithmetic or memory planning. It is intended as a starting point that:

  • Defines a model config / runtime context
  • Shows a per-layer forward loop over IR nodes
  • Sketches a backward loop over the backward IR

Definition at line 638 of file ckernel_codegen.c.

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}
op_name
static const char * op_name(CKOpType op)
Definition ckernel_codegen.c:35
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
CKModelConfig::rms_norm_eps
float rms_norm_eps
Definition ckernel_ir.h:31
CKModelConfig::rope_theta
float rope_theta
Definition ckernel_ir.h:32

References CKIRGraph::config, CKModelConfig::context_window, CKModelConfig::hidden_size, CKIRNode::id, CKIRNode::inputs, CKModelConfig::intermediate_size, CKKernelId::layer, CKIRNode::n_inputs, CKIRNode::n_outputs, CKKernelId::node, CKIRGraph::nodes, CKModelConfig::num_heads, CKModelConfig::num_kv_heads, CKModelConfig::num_layers, CKIRGraph::num_nodes, CKIRNode::op, op_name(), op_name(), op_name(), op_name(), CKInputRef::out_index, CKInputRef::producer, CKModelConfig::rms_norm_eps, CKModelConfig::rope_theta, and CKModelConfig::vocab_size.

Referenced by main().

◆ ck_codegen_emit_runtime()

int ck_codegen_emit_runtime ( const CKIRGraph *  forward,
const char *  path,
CKEmitMode  mode 
)

Emit a C runtime file that stitches kernels for the given forward IR.

Parameters
forwardThe forward IR graph
pathOutput file path
modeCK_EMIT_STANDALONE for executable with main(), CK_EMIT_LIBRARY for shared object with API functions

Returns 0 on success, non-zero on failure.

Definition at line 1468 of file ckernel_codegen.c.

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_runtime_preamble
static int emit_runtime_preamble(FILE *out)
Definition ckernel_codegen.c:783
ck_first_layer_buffer_name
static const char * ck_first_layer_buffer_name(void)
Definition ckernel_codegen.c:623
emit_kernel_manifest
static int emit_kernel_manifest(const CKIRGraph *forward, const char *runtime_path)
Definition ckernel_codegen.c:845
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
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
emit_zero_grad
static void emit_zero_grad(FILE *out)
Definition ckernel_codegen.c:444
emit_layer_offsets_struct
static void emit_layer_offsets_struct(FILE *out)
Definition ckernel_codegen.c:136
emit_global_aliases_to_layer
static void emit_global_aliases_to_layer(FILE *out)
Definition ckernel_codegen.c:423
ck_ir_validate_supported
int ck_ir_validate_supported(const CKIRGraph *graph)
Definition ckernel_registry.c:42

References CK_EMIT_STANDALONE, ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_first_layer_buffer_name(), ck_ir_validate_supported(), CKIRGraph::config, CKModelConfig::context_window, emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_aliases_to_layer(), emit_global_allocations(), emit_global_allocations(), emit_global_allocations(), emit_global_allocations(), emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), emit_kernel_manifest(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_allocations(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_layer_offsets_struct(), emit_library_api(), emit_library_api(), emit_library_api(), emit_library_api(), emit_model_struct(), emit_model_struct(), emit_model_struct(), emit_model_struct(), emit_runtime_preamble(), emit_runtime_preamble(), emit_runtime_preamble(), emit_runtime_preamble(), emit_sgd_update(), emit_sgd_update(), emit_sgd_update(), emit_sgd_update(), emit_zero_grad(), emit_zero_grad(), emit_zero_grad(), emit_zero_grad(), CKModelConfig::hidden_size, CKModelConfig::intermediate_size, CKModelConfig::num_heads, CKModelConfig::num_kv_heads, CKModelConfig::num_layers, CKModelConfig::rms_norm_eps, CKModelConfig::rope_theta, and CKModelConfig::vocab_size.

Referenced by main().