Phase 1, Section 20: Host generation loop

StoryLM-0 decode slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses Section 19 concat/attention, Section 5 gather / rank-4 matmul / softmax, and Section 17 argmax.

1. What this section commits to

Autoregressive generation is a host driver around static-shape Kina steps. Each step consumes one new token and the host-held cache, returns the appended cache and [B, 1, V] logits, then the host applies Section 17 argmax and feeds the class id back.

There is no generate keyword, in-graph PRNG, or growing S inside while. A while carry must keep the same types; cache length is a static shape. The host therefore calls a family of static-S decode functions (one compiled signature per past length) and stores k_next / v_next between calls.

llm_generate.kina stays a loop-carried matrix sketch.

2. Why not a primitive

gather, matmul, reshape, transpose, concat, softmax, broadcast, and argmax are already specified. A fused generate op would invent a token clock, a T_max slot, and a write index that Phase 0 does not define.

3. Operands (numeric instance)

Layouts and widths match storylm0_forward.kina: B = 1, H = 2, Dh = 2, D = 4, V = 8. Host tensors. Bootstrap literal payloads are discarded.

The first step has no past cache (S = 0 is not a default shape). That step attends the new token to itself and returns S = 1 Kᵀ / V. Later steps use S ∈ {1, 2} in this instance and return S+1.

No additive causal mask: one new query attends to every cached position, including itself.

4. Operational meaning

One decode step, existing ops only:

x       = gather embed token_ids          : [B, 1, D]
q_flat  = x @ broadcast(wq)               : [B, 1, D]
k_flat  = x @ broadcast(wk)
v_flat  = x @ broadcast(wv)
q       = reshape/transpose q_flat        : [B, H, 1, Dh]
k_new   = reshape/transpose k_flat        : [B, H, Dh, 1]
v_new   = reshape/transpose v_flat        : [B, H, 1, Dh]
k_next  = concat[axis=3] k_past k_new     : [B, H, Dh, S+1]
v_next  = concat[axis=2] v_past v_new     : [B, H, S+1, Dh]
scores  = q @ k_next
probs   = softmax(scores * broadcast(scale), axis=3)
context = probs @ v_next                  : [B, H, 1, Dh]
merged  = transpose/reshape context       : [B, 1, D]
logits  = (x + merged @ broadcast(wo)) @ broadcast(lm_head)
                                          : [B, 1, V]

The self-attend first step omits the two concats and uses k_new / v_new in place of k_next / v_next.

Host loop for this instance (seed token t0, three new tokens):

(k1, v1, z1) = decode_self(t0, …)
t1           = argmax(z1, axis=2)          : [B, 1] I64
(k2, v2, z2) = decode_s1(t1, k1, v1, …)
t2           = argmax(z2, axis=2)
(k3, v3, z3) = decode_s2(t2, k2, v2, …)
t3           = argmax(z3, axis=2)
sequence     = [t0, t1, t2, t3]

argmax is Section 17 (lowest index on ties). The host may substitute Section 18 Gumbel-max for any step; this instance is greedy.

Weights are left-padded with a unit batch ([4, 4] → [1, 4, 4], [4, 8] → [1, 4, 8]) so [B, 1, D] @ W type-checks. Same broadcast rule as storylm0_forward.kina.

5. Typing (derived)

Each line is existing gather / matmul / reshape / transpose / concat / softmax / broadcast / argmax typing. After three host steps the sequence is four I64 class ids. Cache ranks stay 4; only the static sequence axis changes, and each change is a different compiled fn.

6. Surface instance

storylm0_generate.kina binds decode_self, decode_s1, decode_s2, and greedy_one. @main calls decode_s2 with host placeholders so the module is compile-only. The numeric check loads every fn from the same .so.

No new tokens. The steps sit outside any reverseMode function.

7. Tests

• storylm0_generate.kina — §6 program; golden accepted by mlir-opt; decode_s2 result is (tensor<1x2x2x3xf32>, tensor<1x2x3x2xf32>, tensor<1x1x8xf32>).
• scripts/validate_generate.py compiles the example to a .so, injects patterned finite F32 weights and seed token 3, runs the §4 host loop, and matches (a) the incremental KV algebra and (b) a full causal forward on the generated prefix (NumPy; the same algebra in PyTorch when present).

8. Out of scope

• Growing S inside while.
• Preallocated T_max buffers and a write index.
• Prefill → decode handoff that slices storylm0_forward.kina Kᵀ / V (now Section 21).
• The companion’s 32-token checkpoint test (now Section 23).
• Tokenization, corpus I/O, or trained weights.
• Sampled generate in the numeric instance (Section 18 remains available to the host).
• Differentiating through the loop.
• Device decode.

9. Open issues

1. Prefill handoff. Specified as a cache-returning forward plus last-logit slice in Section 21.
2. T_max scratch. Specified as pad-with-zeros + slice-prefix in Section 22.
3. Sampled rollout. Specified with the 32-token host loop in Section 23.
4. Polymorphic S. One decode fn for every past length needs a specified shape-polymorphic calling convention.