Phase 1, Section 21: Prefill → decode handoff

StoryLM-0 decode slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses the storylm0_forward.kina graph, Phase 0 slice, Section 19 concat/attention, and Section 17 argmax.

1. What this section commits to

Section 20 starts generation from a single seed that self-attends. Real decode starts after a prefill: a T>1 causal forward that produces logits and a cache, then one static-S decode step that appends a generated token.

There is no prefill keyword. Prefill is the existing forward graph with the Kᵀ / V tensors returned beside logits. The host slices the last time step of those logits (Phase 0 slice), applies argmax, and calls a Section 20 decode fn whose past length equals the prefill T.

storylm0_forward.kina stays logits-only. This section’s example does not grow S inside while. T_max scratch is pad-with-zeros + slice-prefix in Section 22.

2. Why not a primitive

gather, rank-4 matmul, softmax, broadcast, slice, concat, and argmax are already specified. A fused prefill op would invent a cache ABI that the forward graph already computes.

3. Operands (numeric instance)

Widths match storylm0_forward.kina with T = 3: B = 1, H = 2, Dh = 2, D = 4, V = 8. Host tensors. Bootstrap literal payloads are discarded.

Prefill:

Prefill result:

k_t    : Tensor[F32, [1, 2, 2, 3], @host]
v      : Tensor[F32, [1, 2, 3, 2], @host]
logits : Tensor[F32, [1, 3, 8], @host]

The last-logit slice is slice[axis=1, start=2, stop=3]: [1, 3, 8] → [1, 1, 8]. Host .so lowering runs --lower-affine after --expand-strided-metadata so extract_slice of a function-boundary memref can convert. That result is a view of the logits buffer (not a fresh allocation). Decode then uses Section 20’s S_past = 3 step (decode_s3).

4. Operational meaning

Prefill is the Section 5 / storylm0_forward algebra at T = 3, except the function returns the cache:

x      = gather embed token_ids
q, k_t, v = reshape/transpose (x @ Wq/Wk/Wv)
logits = residual + (attn(q, k_t, v, mask, scale) @ Wo) @ lm_head
result = (k_t, v, logits)

Host handoff:

(k, v, z) = prefill(prompt, mask, …)
z_last    = slice[axis=1, start=2, stop=3] z     : [1, 1, 8]
t3        = argmax(z_last, axis=2)               : [1, 1] I64
(k4, v4, z4) = decode_s3(t3, k, v, …)
t4        = argmax(z4, axis=2)
sequence  = prompt ++ [t3, t4]

decode_s3 is Section 20 §4 with S = 3 (concat on the sequence axis, no causal mask: one new query attends to every cached position). argmax is Section 17.

The cache handed off is the full prefill Kᵀ / V, not a prefix slice of those tensors. slice is only for the last logit row.

5. Typing (derived)

Each line is existing gather / matmul / reshape / transpose / softmax / broadcast / slice / concat / argmax typing.

6. Surface instance

storylm0_prefill.kina binds prefill, last_logits, decode_s3, and greedy_one. @main calls prefill with host placeholders. The numeric check loads every fn from the same .so.

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

7. Tests

• storylm0_prefill.kina — §6 program; golden accepted by mlir-opt; prefill result is (tensor<1x2x2x3xf32>, tensor<1x2x3x2xf32>, tensor<1x3x8xf32>).
• scripts/validate_prefill.py compiles the example to a .so, injects patterned finite F32 weights and prompt [1, 4, 2], runs the §4 handoff, and matches (a) a full causal forward on the prompt (logits and Kᵀ / V) and (b) a causal forward last-position on prompt ++ [t3] after the decode step (NumPy; the same algebra in PyTorch when present).

8. Out of scope

• Growing S inside while.
• Preallocated T_max buffers and a write index.
• Changing storylm0_forward.kina to return a cache.
• The companion’s 32-token checkpoint test (now Section 23).
• Sampled rollout.
• Differentiating through prefill or the handoff.
• Device prefill / decode.

9. Open issues

1. T_max scratch. Specified as pad-with-zeros + slice-prefix in Section 22.
2. Prefill T family. One prefill signature per prompt length, same static-S stance as Section 20.
3. Sampled rollout after prefill. The 32-token sampled loop is Section 23 (re-forward, not a cache-returning prefill family).
4. Packed cache layout. A single tensor holding K and V is unspecified.