Phase 1, Section 19: Functional KV-cache append

StoryLM-0 decode slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses Phase 0 concat, Section 5 rank-4 matmul, and existing softmax / broadcast.

1. What this section commits to

Autoregressive decode reuses past keys and values. The update is append: concatenate the new position onto a host-held cache, then attend with a single new query.

There is no kv_cache keyword, IR constructor, or FFI. The append is Phase 0 concat on the sequence axis. The StoryLM-0 layouts match storylm0_forward.kina:

k_t : Tensor[T, [B, H, Dh, S], @host]    (K transposed for Q @ Kᵀ)
v   : Tensor[T, [B, H, S, Dh], @host]

llm_generate.kina stays a loop-carried matrix sketch. A while body must return the same carry types it bound; a growing S cannot be a loop-carried tensor. Multi-step generation that grows the cache stays a host loop around this one-step function.

2. Why not a primitive

concat, rank-4 matmul, softmax, and broadcast are already specified. A fused cache op would invent mutability, a T_max slot, and a write index that Phase 0 does not define. Scatter into a preallocated buffer is IR-internal, not a user form (Section 14).

3. Operands

All values are host tensors. Bootstrap tensor-literal payloads are discarded.

B = 1, H = 2, Dh = 2, S_past = 3, S_new = 1. The concat axis is static (the checker already requires that). Placements match. T is a floating dtype from the implemented frem set. The instance is finite F32.

A single new query attending to every cached position (including itself) needs no additive causal mask. Prefill masking stays the host-provided mask in storylm0_forward.kina.

4. Operational meaning

k_next  = concat[axis=3] k_past k_new     : [B, H, Dh, S_past+S_new]
v_next  = concat[axis=2] v_past v_new     : [B, H, S_past+S_new, Dh]
scores  = q_new @ k_next                  : [B, H, 1, S_past+S_new]
scaled  = scores * broadcast(scale)
probs   = softmax(scaled, axis=3)
context = probs @ v_next                  : [B, H, 1, Dh]

concat is Phase 0 §13: non-concat dims match; the concat axis sums. Host static concat lowers to tensor.insert_slice (LLVM 18 cannot bufferize tensor.concat). matmul is Section 5 variadic-batch (rank 4). softmax is Section 5.

No in-place write. Each call produces new tensors. The host stores k_next / v_next and passes them as k_past / v_past on the next step, with a new q_new / k_new / v_new.

5. Typing (derived)

Each line is existing concat / matmul / broadcast / softmax typing. After the two concats:

k_next  : Tensor[T, [1, 2, 2, 4], @host]
v_next  : Tensor[T, [1, 2, 4, 2], @host]
context : Tensor[T, [1, 2, 1, 2], @host]

The wrapper result is a constructed flat tuple (Section 12 flatten):

(k_next, v_next, context)

6. Surface instance

let kv_step = fn(
  q_new: tensor<F32, [1, 2, 1, 2], @host>,
  k_past: tensor<F32, [1, 2, 2, 3], @host>,
  k_new: tensor<F32, [1, 2, 2, 1], @host>,
  v_past: tensor<F32, [1, 2, 3, 2], @host>,
  v_new: tensor<F32, [1, 2, 1, 2], @host>,
  scale: tensor<F32, [1, 1, 1, 1], @host>
) -> (
  tensor<F32, [1, 2, 2, 4], @host>,
  tensor<F32, [1, 2, 4, 2], @host>,
  tensor<F32, [1, 2, 1, 2], @host>
) { … §4 lets …; (k_next, v_next, context) }

No new tokens. The step sits outside any reverseMode function. Projection (x @ W_k, reshape, transpose) is the existing decoder graph; this section’s example takes already laid-out Kᵀ / V.

7. Tests

• storylm0_kv_step.kina — §6 program; golden accepted by mlir-opt; the step fn result is (tensor<1x2x2x4xf32>, tensor<1x2x4x2xf32>, tensor<1x2x1x2xf32>).
• scripts/validate_kv_cache.py compiles the example to a .so, injects patterned finite F32 tensors, and matches §4 (NumPy concatenate + batched matmul + softmax; the same algebra in PyTorch when present).

8. Out of scope

• Growing S inside while (carry types are fixed).
• Preallocated T_max buffers and a write index.
• A user-facing scatter / slice-assign.
• Multi-step generation, tokenization, or sampling in the same program.
• Differentiating through the append.
• Device concat / decode.
• Changing storylm0_forward.kina prefill.

9. Open issues

1. Host generation loop. Specified as static-S decode plus host argmax in Section 20.
2. Prefill → decode handoff. Specified in Section 21.
3. T_max scratch. Specified as pad-with-zeros + slice-prefix in Section 22.
4. Concat VJP. Unspecified. Do not guess a payload.