Architecture Decision Records

Context

Previous IR versions conflated several concerns:

• IR v2: Combined lowering and memory planning into a single phase
• IR v3: Clean layout but hardcoded architecture knowledge in Python

Decision

Split the IR pipeline into three distinct phases:

Consequences

Context

In IR v3, model architecture was embedded in build_layer_layout(). Adding new models required modifying core codegen.

Decision

Templates are the canonical definition of model architectures. The IR is a compiled artifact derived from templates + config.json.

# templates/qwen2.yaml
name: qwen2
config_mapping:
  hidden_size: embed_dim
  num_attention_heads: num_heads

layers:
  - name: decoder_layers
    repeat: "{{num_layers}}"
    ops:
      - name: pre_attn_norm
        op: RMSNorm
        input: "{{prev_output}}"
        output: normed

      - name: qkv_proj
        op: GEMM
        input: normed
        outputs: [q, k, v]

      - name: attention
        op: Attention
        inputs: [q, k, v]
        output: attn_out
        params:
          causal: true

Consequences

• New architectures: Add YAML file, not Python code
• Explicit: Architecture visible at a glance
• Validation: Can diff templates against HuggingFace

Context

Transformer inference has fundamentally different execution patterns:

• Prefill: Process entire prompt (T tokens, large buffers)
• Decode: Generate one token (T=1, KV cache reused)
• Backward: Compute gradients for training

Decision

Allocate separate activation buffers for each execution mode. Weights are shared, working memory is mode-specific.

┌─────────────────────────────────────────────────────────────┐
│                    MEMORY LAYOUT                             │
├─────────────────────────────────────────────────────────────┤
│ Weights (read-only, shared across modes)                     │
├─────────────────────────────────────────────────────────────┤
│ KV Cache (persistent across decode steps)                    │
├─────────────────────────────────────────────────────────────┤
│ Prefill Buffers [max_seq, embed_dim] ← Large                 │
├─────────────────────────────────────────────────────────────┤
│ Decode Buffers [1, embed_dim] ← 2000x smaller                │
└─────────────────────────────────────────────────────────────┘

Buffer Size Comparison (Qwen2-0.5B)

Decode buffers are ~2000x smaller, enabling better cache utilization.

Context

Buffer overflows in GEMM/attention are common and hard to diagnose. IR v2 used dynamic allocation, making bugs non-reproducible.

Decision

• All tensor offsets computed at codegen time, baked into C code
• 64-byte alignment for AVX-512 cache line optimization
• Insert 64-byte canary markers (0xDEADBEEF) between tensors in debug builds
• Export memory_layout.json for debugging tools

// Generated code - offsets are compile-time constants
#define OFFSET_EMBED_WEIGHT      64
#define OFFSET_LAYER_0_WQ        1606720
#define CANARY_0                 272629824  // Between embed and layer_0

// Verify canaries after each layer (debug mode)
bool verify_canaries(Model *model) {
    uint64_t *canary = (uint64_t*)(model->base + CANARY_0);
    for (int j = 0; j < 8; j++) {
        if (canary[j] != 0xDEADBEEFDEADBEEFULL) {
            fprintf(stderr, "CANARY VIOLATION!\n");
            return false;
        }
    }
    return true;
}

Context

Multiple kernel variants exist for each operation:

• GEMM: gemm_blocked_serial_bf16, gemm_blocked_parallel_bf16, gemm_1x1_bf16
• Attention: batch vs single-query, GQA vs MHA

Decision

Graph IR uses abstract ops. Lowering resolves to concrete kernel function names.

{"op": "GEMM",
 "inputs": ["x", "w"],
 "output": "y"}

// Prefill mode
{"kernel": "gemm_blocked_parallel_bf16"}

// Decode mode
{"kernel": "gemm_1x1_bf16"}

Result: No runtime branching for kernel selection. Mode-specific code paths are explicit.

Context

HuggingFace models use inconsistent naming conventions:

• Llama: model.layers.0.self_attn.q_proj.weight
• GPT-2: transformer.h.0.attn.c_attn.weight (combined QKV!)

Decision

Weight name mapping defined in architecture templates.

# templates/qwen2.yaml
weight_mapping:
  "model.embed_tokens.weight": "embed.weight"
  "model.layers.{i}.self_attn.q_proj.weight": "layers.{i}.attn.wq"
  "model.layers.{i}.self_attn.k_proj.weight": "layers.{i}.attn.wk"
  "model.layers.{i}.mlp.gate_proj.weight": "layers.{i}.mlp.gate"
  "model.norm.weight": "final_norm.gamma"
  "lm_head.weight": "lm_head.weight"

At codegen, validate all mapped weights exist in safetensors header. Warn about unmapped weights.