Phase 1, Section 18: Seeded categorical sampling

StoryLM-0 decode slice. Not a Phase 1 exit-criteria replacement. Not a new primitive. Uses Phase 0 log / neg / div / add / broadcast and Section 17 argmax.

1. What this section commits to

Section 17 locks greedy class selection. Stochastic decoding still needs a seeded categorical sample from softmax(logits / T).

Randomness is a Phase 2 effect (Phase 0 README; design memo 02). This section does not add a sample keyword, an RNG op, or a Random handler. The host supplies a Uniform(0, 1) tensor u (seeded outside Kina) and a positive temperature T. Kina applies the Gumbel-max algebra and Section 17 argmax.

Given u, the program is deterministic. Reproducible seeds live on the host.

2. Why not a primitive

log, neg, div, add, broadcast, and argmax are already specified. A fused sample op would invent an in-graph PRNG, a seed clock, and a NaN policy that Phase 0 does not define.

Top-k and top-p need a specified sort / prefix-sum selection. Those ops are not in the implemented set. They stay out of scope.

3. Operands

All values are host tensors. Bootstrap tensor-literal payloads are discarded, so the numeric instance does not write uniforms or T into .kina source.

u has the same shape, dtype, and placement as logits. Every element of u satisfies 0 < u_i < 1. T > 0. T = 0, u_i ∈ {0, 1}, and NaN are unspecified.

[1, 1] broadcasts to [1, 4, 8] under the existing Phase 0 §11 rule (left-pad with 1s). Host lowering maps each size-1 source axis to affine 0 after that pad (same rule as same-rank [1, 1] → [4, 8]). Rank-0 Tensor[T, []] is allowed by the type rule; it is not the numeric instance.

T is a floating dtype from the implemented frem set. The instance is finite F32.

4. Operational meaning

After broadcasting T to the logits shape, for every site and class i:

g_i      = -log(-log(u_i))
scores_i = logits_i / T + g_i
token    = argmax(scores)     (Section 17; lowest index on ties)

This is the Gumbel-max trick: token is distributed as Categorical(softmax(logits / T)) when u_i are i.i.d. Uniform(0, 1).

No top-k, top-p, nucleus, or typical decoding. No in-graph seed update. Each call consumes the u the host passed.

5. Typing (derived)

Each line is existing unary / elementwise / broadcast / reduce typing. log requires a floating operand. After broadcast, div and add require equal shapes. argmax on the class axis yields I64 (Section 17).

tokens : Tensor[I64, [1, 4], @host]

6. Surface instance

let sample = fn(
  logits: tensor<F32, [1, 4, 8], @host>,
  u: tensor<F32, [1, 4, 8], @host>,
  temperature: tensor<F32, [1, 1], @host>
) -> tensor<I64, [1, 4], @host> {
  let log_u = log(u) in
  let neg_log_u = -log_u in
  let log_neg_log_u = log(neg_log_u) in
  let gumbel = -log_neg_log_u in
  let t_b = temperature.broadcast([1, 4, 8]) in
  let scaled = logits / t_b in
  let scores = scaled + gumbel in
  let tokens = scores.reduce(op=argmax, axis=2) in
  tokens
}

No new tokens. The composition sits outside any reverseMode function.

7. Tests

• storylm0_sample.kina — §6 program; golden accepted by mlir-opt; the sample fn result is tensor<1x4xi64>.
• scripts/validate_sample.py compiles the example to a .so, injects patterned finite F32 logits, host Uniform(ε, 1-ε) from a fixed RandomState seed, and T = 0.8, and matches §4 (NumPy; the same algebra in PyTorch when present).

8. Out of scope

• An in-graph PRNG, seed tensor, or Random effect.
• Top-k, top-p, typical, or beam search.
• Differentiating through the sample (Section 17: no argmax VJP; no Gumbel-softmax relaxation).
• Device argmax / device sampling.
• Composing decoder forward + this sample as a generation loop (greedy host loop is Section 20).
• KV-cache.

9. Open issues

1. Top-k / top-p. Need a specified sort or prefix selection before this algebra can mask a subset of classes.
2. KV-cache / generate. KV-cache append is Section 19. The host greedy loop is Section 20. Sampled 32-token rollout is Section 23. llm_generate.kina remains a matrix sketch.
3. Phase 2 Random. A later effect can own the host Uniform draw; this section's algebra stays the same.
4. u ∈ {0, 1} / T ≤ 0 / NaN. Unspecified. Do not guess a payload.