Kina.ai

Learn

Phase 0 Companion: Real-World ML / LLM Examples

Companion to the Phase 0 specification. Programs from real LLM and ML workloads, shown alongside their PyTorch equivalents to make the cost and shape of the design tangible.

1. Framing

The previous walkthrough showed primitives. This one shows production-relevant code: things that appear in actual LLM forward passes (LLaMA, Mistral, GPT family) and things that hit Phase 0's deliberate limits.

Each example has three parts:

The examples are organized by what they reveal:

2. Things that work cleanly

These programs type-check naturally in Phase 0 with no escape hatches and only modest verbosity.

2.1 RMSNorm (LLaMA-style normalization)

The normalization layer used in LLaMA, Mistral, and most modern LLMs. Replaces LayerNorm; cheaper because there's no mean subtraction.

PyTorch

def rmsnorm(x, weight, eps=1e-6):
    # x: [B, S, D]  weight: [D]
    var = x.pow(2).mean(-1, keepdim=True)
    rstd = torch.rsqrt(var + eps)
    return x * rstd * weight

Three lines of math. Looks elegant, and is — but most of the elegance comes from PyTorch's implicit broadcasting and keepdim.

Phase 0 surface

let rmsnorm [B, S, D]
    (x : Tensor[F32, [B, S, D]])
    (weight : Tensor[F32, [D]])
    (eps : Tensor[F32, []])           // scalar
    : Tensor[F32, [B, S, D]] =
  let x_sq         = x * x in                                  // [B, S, D]
  let mean_sq      = reduce_keep[mean, axis=2] x_sq in         // [B, S, 1]
  let eps_bcast    = broadcast eps [B, S, 1] in
  let denom_inner  = mean_sq + eps_bcast in                    // [B, S, 1]
  let rstd         = rsqrt denom_inner in                      // [B, S, 1]
  let rstd_bcast   = broadcast rstd [B, S, D] in
  let normed       = x * rstd_bcast in                         // [B, S, D]
  let weight_bcast = broadcast weight [B, S, D] in
  normed * weight_bcast

Notes

Nine lines vs PyTorch's three. But every shape transition is visible: each broadcast tells you exactly what's being expanded and to what. There are no SMT obligations to discharge here — every shape relationship is syntactic. The type-checker runs in microseconds.

This is the good case for the design. The verbosity is real but bounded; the lowering will fuse most of these into one or two kernels (the broadcast-then-multiply patterns are exactly what fusion is for); the explicitness pays off when you need to debug or optimize.

2.2 SwiGLU FFN block

The feed-forward block from LLaMA. Three matmuls (gate, up, down) and an elementwise gate-times-up.

PyTorch

def swiglu_ffn(x, w_gate, w_up, w_down):
    # x: [B, S, D]  w_gate, w_up: [D, H]  w_down: [H, D]
    gate = F.silu(x @ w_gate)
    up   = x @ w_up
    return (gate * up) @ w_down

Phase 0 surface

let swiglu_ffn [B, S, D, H]
    (x      : Tensor[F32, [B, S, D]])
    (w_gate : Tensor[F32, [D, H]])
    (w_up   : Tensor[F32, [D, H]])
    (w_down : Tensor[F32, [H, D]])
    : Tensor[F32, [B, S, D]] =
  let w_gate_b = broadcast w_gate [B, D, H] in
  let gate_pre = matmul x w_gate_b in                          // [B, S, H]
  let gate     = silu gate_pre in
  let w_up_b   = broadcast w_up [B, D, H] in
  let up       = matmul x w_up_b in                            // [B, S, H]
  let gated    = gate * up in
  let w_down_b = broadcast w_down [B, H, D] in
  matmul gated w_down_b

Notes

Eight lines for what PyTorch does in three. Each weight matrix needs an explicit broadcast to add the batch dimension, which is the same cost we saw in the linear layer example. No SMT calls; pure shape unification.

This is also a good case. SwiGLU is the bulk of the compute in modern LLMs (FFN is typically 2/3 of transformer parameters), and the type system handles it without complaint.

3. Things that need verbose workarounds

Where Phase 0's design decisions (matmul-with-one-batch-dim, no implicit broadcasting) show their cost.

3.1 Multi-head self-attention with softmax

The canonical LLM kernel. Four leading dimensions ([B, H, S, D_h]), softmax in the middle, the whole thing typically fused in production via FlashAttention.

PyTorch (idiomatic)

def mha(q, k, v, scale):
    # q, k, v: [B, H, S, D_h]
    scores  = (q @ k.transpose(-2, -1)) * scale          # [B, H, S, S]
    weights = F.softmax(scores, dim=-1)                  # [B, H, S, S]
    return weights @ v                                   # [B, H, S, D_h]

Three lines. PyTorch hides:

Phase 0 surface

let mha [B, H, S, D_h]
    (q     : Tensor[F32, [B, H, S, D_h]])
    (k     : Tensor[F32, [B, H, S, D_h]])
    (v     : Tensor[F32, [B, H, S, D_h]])
    (scale : Tensor[F32, []])
    : Tensor[F32, [B, H, S, D_h]] =

  // ---- Step 1: K^T ----
  let k_t = transpose k [0, 1, 3, 2] in                  // [B, H, D_h, S]

  // ---- Step 2: Flatten [B, H] so matmul has one batch dim ----
  let dim BH = B * H in                                  // axiomatic; outside Presburger
  let q_flat   = reshape q   [BH, S, D_h] in
  let k_t_flat = reshape k_t [BH, D_h, S] in
  let v_flat   = reshape v   [BH, S, D_h] in

  // ---- Step 3: Scaled dot-product ----
  let scores      = matmul q_flat k_t_flat in            // [BH, S, S]
  let scale_b     = broadcast scale [BH, S, S] in
  let scaled      = scores * scale_b in

  // ---- Step 4: Softmax (composite of 6 ops) ----
  let max_row     = reduce_keep[max, axis=2] scaled in   // [BH, S, 1]
  let max_b       = broadcast max_row [BH, S, S] in
  let shifted     = scaled - max_b in                    // numerical stability
  let exp_shifted = exp shifted in
  let sum_row     = reduce_keep[sum, axis=2] exp_shifted in  // [BH, S, 1]
  let sum_b       = broadcast sum_row [BH, S, S] in
  let weights     = exp_shifted / sum_b in               // [BH, S, S]

  // ---- Step 5: weights @ V ----
  let out_flat    = matmul weights v_flat in             // [BH, S, D_h]

  // ---- Step 6: Unflatten ----
  reshape out_flat [B, H, S, D_h]

Notes

About 23 lines of body, vs PyTorch's 3. Three sources of verbosity:

  1. The flatten dance (Step 2 + Step 6). Three reshapes plus a fresh BH variable, all because matmul takes one batch dim. As the previous walkthrough flagged, this is the strongest argument for adding [..., M, K] matmul to the operation set in Phase 1.

  2. Softmax is six operations (Step 4). The numerical-stability max-shift, an exp, a reduce-sum, two broadcasts, and a divide. PyTorch hides this in F.softmax, but underneath, those six ops are what runs. Phase 7 surface syntax can sugar softmax(scaled, dim=-1) to expand to the six-op composite.

  3. Three explicit broadcasts for scalar scaling and softmax intermediates.

This is the bad case for current Phase 0. About a third of the verbosity is fundamental (softmax just is six ops underneath); the other two thirds is recoverable via better primitives or surface sugar. Two concrete fixes:

Together: 23 lines → ~10 lines, much closer to PyTorch's 3. Both are Phase 1 candidates.

3.2 KV-cache append (single step works, loops don't)

LLM inference appends the new token's K/V to a cache after each step. The cache grows by 1 each step.

PyTorch

def kv_append(k_cache, v_cache, k_new, v_new):
    # k_cache, v_cache: [B, H, S_cache, D_h]
    # k_new, v_new:     [B, H, 1, D_h]
    k_updated = torch.cat([k_cache, k_new], dim=2)
    v_updated = torch.cat([v_cache, v_new], dim=2)
    return k_updated, v_updated

Phase 0 surface

let kv_append [B, H, S_cache, D_h]
    (k_cache : Tensor[F32, [B, H, S_cache, D_h]])
    (v_cache : Tensor[F32, [B, H, S_cache, D_h]])
    (k_new   : Tensor[F32, [B, H, 1, D_h]])
    (v_new   : Tensor[F32, [B, H, 1, D_h]])
    : (Tensor[F32, [B, H, S_cache + 1, D_h]],
       Tensor[F32, [B, H, S_cache + 1, D_h]]) =
  let k_updated = concat[axis=2] k_cache k_new in
  let v_updated = concat[axis=2] v_cache v_new in
  (k_updated, v_updated)

Notes

This actually works in Phase 0. Concat at axis 2 produces S_cache + 1 along that axis — pure Presburger arithmetic, the type-checker handles it without escape hatches. The signature even encodes the fact that the output cache is one larger.

What doesn't work in Phase 0: a loop of these calls. Generating a sequence of tokens means calling kv_append repeatedly, and after N calls the cache shape is [B, H, S_cache + N, D_h] where N is dynamic. Loops require:

So: a single attention step types beautifully. Building an inference loop on top requires Phase 2.

This is informative. Phase 0 + Phase 1 gets you a typed forward pass for one transformer block. Useful for benchmarks, papers, validation. Real LLM inference needs the loop, which means Phase 2 is on the critical path for "actually running models."

4. Things that hit Phase 0 limits

Programs that show where the deliberate restrictions become genuine blockers.

4.1 Grouped-Query Attention head expansion

Modern LLMs use GQA (LLaMA 2/3, Mistral): several Q heads share each K/V head, saving KV-cache memory. Before attention, the K/V tensors are repeated to match Q's head count.

PyTorch

def expand_kv_for_gqa(k, v, n_q_heads):
    # k, v: [B, n_kv_heads, S, D_h]
    # output: [B, n_q_heads, S, D_h]
    n_kv_heads = k.shape[1]
    repeat = n_q_heads // n_kv_heads
    k = k.repeat_interleave(repeat, dim=1)
    v = v.repeat_interleave(repeat, dim=1)
    return k, v

Phase 0 surface (attempted)

let expand_kv_for_gqa [B, KV_H, Q_H, S, D_h]
    requires (Q_H mod KV_H = 0)
    (k : Tensor[F32, [B, KV_H, S, D_h]])
    (v : Tensor[F32, [B, KV_H, S, D_h]])
    : (Tensor[F32, [B, Q_H, S, D_h]],
       Tensor[F32, [B, Q_H, S, D_h]]) =

  let dim REP = Q_H / KV_H in                             // OK: Q_H mod KV_H = 0

  // Standard repeat-via-broadcast pattern:
  // [B, KV_H, S, D_h] -> [B, KV_H, 1, S, D_h]
  //                  -> [B, KV_H, REP, S, D_h]      (broadcast)
  //                  -> [B, Q_H, S, D_h]            (reshape, merging KV_H * REP)

  let k_5d = reshape k [B, KV_H, 1, S, D_h] in            // axiomatic volume eq
  let k_expanded = broadcast k_5d [B, KV_H, REP, S, D_h] in

  // Now we want to merge KV_H and REP into Q_H. We need: KV_H * REP = Q_H.
  // The type system has REP = Q_H / KV_H and Q_H mod KV_H = 0, but cannot
  // automatically derive KV_H * REP = Q_H because variable*variable is not
  // in the refinement language.

  let dim merged = KV_H * REP in                          // axiomatic; outside Presburger
  let assert_eq merged Q_H in                             // additional axiom
  let k_merged = reshape k_expanded [B, merged, S, D_h] in

  // ... but the return type wants [B, Q_H, S, D_h], and we have [B, merged, S, D_h].
  // We need yet another assertion or coercion to bridge merged ≡ Q_H.
  ...

Notes

This doesn't quite work without multiple escape-hatch assertions. The mathematical fact (Q_H mod KV_H = 0) ⇒ KV_H · (Q_H / KV_H) = Q_H is true and should be derivable, but Presburger doesn't reason about variable-times-variable products. Phase 0 forces the user to assert it axiomatically — which works but reads as ugly boilerplate that isn't even logically sound (the assertion isn't proved; the compiler trusts the user).

Three responses are possible:

  1. Accept the verbosity and add the axioms. Sound only because a runtime reshape-volume check catches mistakes. Workable; ugly.
  2. Add a repeat_interleave primitive to the operation set (Phase 1). Its typing rule encodes the divisibility requirement and produces the merged dimension symbolically without going through reshape twice. This is the right answer for production code; GQA is everywhere.
  3. Strengthen the refinement language to handle restricted variable products (specifically: products where one factor is dim_a / dim_b and the divisibility refinement is in scope). Research-shaped; not a Phase 1 deliverable.

(2) is concrete and small. Adding repeat_interleave (and its dual flatten_pair) as primitives in Phase 1 would handle GQA, multi-head attention's flatten dance, and several other patterns at once.

4.2 Token embedding lookup

The first operation in any LLM: turn token IDs into vectors via table lookup.

PyTorch

def embed(token_ids, embedding_table):
    # token_ids:        [B, S] of int
    # embedding_table:  [V, D]
    # returns:          [B, S, D]
    return embedding_table[token_ids]    # gather along dim 0

Phase 0 surface

// Cannot be expressed in Phase 0.
//
// The operation `embedding_table[token_ids]` is a gather, which Phase 0
// (Section 2 §3) explicitly defers. The output shape [B, S, D] depends on
// the values of token_ids (each entry indexes into V), which is data-
// dependent indexing.

Notes

This is a hard limit. Embedding lookup is gather, and gather is on the deferred list in Section 2 §3 (alongside scatter, sort, scan).

You can't write a workable LLM forward pass without embedding. Phase 1 must include gather as a primitive. This is non-negotiable for any practical use of the language.

What gather requires that Phase 0 doesn't have:

The type rule is reasonable: gather : Tensor[T, [V, D]] → Tensor[I64, [B, S]] → Tensor[T, [B, S, D]] with a runtime-checked precondition that all indices are in range [0, V). The runtime check connects to memo 1's symbolic-shapes work (the output shape is [B, S, D] which is statically known; the bounds check is a runtime predicate).

Adding gather is a Phase 1 work item.

4.3 Cross-entropy loss (forward-only)

The loss function for next-token prediction. Even just the forward pass hits Phase 0 limits.

PyTorch

def cross_entropy_forward(logits, targets):
    # logits:  [B, V]  (un-normalized scores)
    # targets: [B] of int
    log_probs = F.log_softmax(logits, dim=-1)
    nll = -log_probs.gather(1, targets.unsqueeze(1)).squeeze(1)
    return nll.mean()

Phase 0 surface

Cannot be expressed cleanly. The reasons stack:

Notes

Cross-entropy loss forward is the simplest possible training-related operation, and Phase 0 can't express it. This is fine — Phase 0 is an explicitly forward-pass-only foundation — but it's worth being clear about.

The smallest viable Phase 1 extension to enable training-shaped code: add gather. With gather, cross-entropy forward is straightforward. Backward still needs Phase 2 (effects, AD).

5. Composition: a real transformer block fragment

What a complete piece of a transformer looks like with the workable parts. RMSNorm + SwiGLU FFN with residual connection. (Skipping attention because §3.1 already showed how verbose that is.)

PyTorch

def transformer_block_partial(x, norm_w, w_gate, w_up, w_down, eps):
    # x: [B, S, D]
    # All weights as appropriate shapes.
    h = rmsnorm(x, norm_w, eps)
    h = swiglu_ffn(h, w_gate, w_up, w_down)
    return x + h     # residual

Phase 0 surface

let transformer_block_partial [B, S, D, H]
    (x      : Tensor[F32, [B, S, D]])
    (norm_w : Tensor[F32, [D]])
    (w_gate : Tensor[F32, [D, H]])
    (w_up   : Tensor[F32, [D, H]])
    (w_down : Tensor[F32, [H, D]])
    (eps    : Tensor[F32, []])
    : Tensor[F32, [B, S, D]] =

  let h_norm   = rmsnorm[B, S, D] x norm_w eps in
  let h_ffn    = swiglu_ffn[B, S, D, H] h_norm w_gate w_up w_down in
  x + h_ffn

Notes

Once rmsnorm and swiglu_ffn are defined (Examples 2.1 and 2.2), the block composition is clean. Four lines, almost identical to PyTorch's structure.

This is the case the design optimizes for: abstraction at the function level recovers ergonomics. The verbosity lives inside rmsnorm and swiglu_ffn; the composition reads naturally. Library-level functions hide the explicit broadcasting and reshaping; user-level code reads more like PyTorch.

The argument for Phase 7 surface syntax shifts: it isn't just about making core IR programs less verbose, it's about giving users fluent composition. Most ML code is composition (block calls block calls block calls primitives), and composition is already fine in Phase 0. Where Phase 7 would help most:

Surface ergonomics need to make primitive-level code readable. They don't need to do anything special for composition — the type system already does that.

6. Reflections — what these examples teach us

Stepping back from the individual programs, the patterns that emerge:

6.1 The Phase 1 primitive additions, prioritized

In order of impact:

  1. Variadic-batch matmul ([..., M, K] @ [..., K, N]). Eliminates the flatten dance. Clean typing rule. High leverage.
  2. Gather (and probably scatter as its dual). Required for embedding, cross-entropy, top-k routing. Non-negotiable for practical use.
  3. softmax as a primitive operation with internal decomposition. Lowering decomposes; users see one op. Big readability win.
  4. repeat_interleave and flatten_pair (or just a more general repeat primitive). Handles GQA cleanly. Avoids needing the variable-product escape hatch in common cases.
  5. A scalar-broadcast shorthand. x * scalar shouldn't require an explicit broadcast at the IR level (the dimension-1 broadcast is mechanical). This is borderline between "core IR" and "Phase 7 surface" — possibly acceptable to put it in the core if we want.

These are the additions that take the operation set from "can express forward passes with effort" to "can express forward passes idiomatically."

6.2 The Phase 2 dependencies

Three cases needed Phase 2:

Phase 2 is on the critical path for "real" use. Phase 0 + Phase 1 alone produces a research artifact (typed forward passes for ML benchmarks); Phase 2 produces something that can train models.

6.3 The verbosity is mostly composable

The biggest reaction to §3.1 (multi-head attention) is "this is unusable at 23 lines." But the response should be: users almost never write attention from scratch. They write it once, put it in a library, and call attention(q, k, v) thereafter. The 23 lines live in one place; the rest of the codebase calls a 1-line invocation.

This is roughly how every framework works. PyTorch's nn.MultiheadAttention is hundreds of lines internally; users see one constructor. The Phase 0 verbosity of §3.1 means the first implementer of attention has a hard time; everyone else has it easy.

6.4 The escape hatches will accumulate

The variable-product axioms in §3.1 (matmul flatten) and §4.1 (GQA expansion) and the implicit reshape volume axioms throughout — these are individually sound under "trust the user, runtime catches mistakes" but they accumulate. A complex attention layer might have 4–6 axiomatic assertions, each one a small soundness gap.

Two responses:

The escape hatches are an honest feature of Phase 0's design but not a feature to be proud of. Phase 1 should chip away at them.

7. Conclusion

Five real-world ML/LLM patterns, two PyTorch comparisons per. Phase 0 handles the simpler ones (RMSNorm, SwiGLU) with mild verbosity, struggles with the complex ones (full multi-head attention) due to the matmul-batched-with-one-dim restriction, and cannot express several critical patterns (embedding, GQA cleanly, cross-entropy, anything with a loop or training).

The walkthrough surfaces a clear Phase 1 prioritization (variadic matmul, gather, softmax, repeat primitives) and confirms that Phase 0 is genuinely forward-pass-only — Phase 2 (effects, AD, control flow) is on the critical path for any practical use.

This is the right problem to be in. Phase 0 was supposed to be the foundation, not the language. The walkthrough is doing what walkthroughs are supposed to do: showing the foundation's edges so Phase 1 can shape them deliberately.

End of real-world examples walkthrough.