Phase 3, Section 5: StoryLM-0 as the first sharding workload

Phase 3: Sharding (the distribution axis). Not a Phase 3 exit-criteria replacement. Not a new primitive. Uses Phase 0 slice / concat and Section 5 matmul on the existing storylm0_forward.kina graph.

1. What this section commits to

AGENTS.md step 5 says to use the same tiny causal model as the Phase 3 sharding workload before scaling model size. Sections 2–4 sketch @mesh, @sharded, @replicated, and collective coercions. The compiler does not parse those refinements yet. This section locks the first numeric instance so later mesh work has a known algebra.

The instance is 2-way column-parallel lm_head:

Logical mesh size 2 (not the Section 1 4-device simulator).
lm_head [D, V] = [4, 8] is split on axis 1 into two [4, 4] shards.
Each shard computes [B, T, 4] logits.
concat on the vocab axis is the all-gather of those shards.

Everything else in the forward graph is replicated (the same host tensors, computed once). There is no @mesh keyword in this example.

2. Why not a primitive

slice, matmul, broadcast, and concat are already specified. A fused tensor-parallel op would invent a mesh and a collective ABI that Phase 3 Milestone 1 has not implemented. Column-parallel vocab is the functional encoding of “shard lm_head on V, all-gather logits.”

3. Operands (numeric instance)

Widths match storylm0_forward.kina: B = 1, T = 4, H = 2, Dh = 2, D = 4, V = 8. Host tensors. Token ids [1, 4, 2, 6] (same as the forward numeric check). V is divisible by the mesh size 2 (Section 2 §3.1).

Full graph inputs are the Section 5 / storylm0_forward table. The only changed line is the vocabulary projection:

Name	Type	Role
`residual`	`Tensor[F32, [1, 4, 4], @host]`	`x + attn @ Wo`
`lm_head`	`Tensor[F32, [4, 8], @host]`	full vocabulary

Shard and gather:

lm0    = slice[axis=1, start=0, stop=4] lm_head   : [4, 4]
lm1    = slice[axis=1, start=4, stop=8] lm_head   : [4, 4]
z0     = residual @ broadcast(lm0)                : [1, 4, 4]
z1     = residual @ broadcast(lm1)                : [1, 4, 4]
logits = concat[axis=2] z0 z1                     : [1, 4, 8]

concat here is the Section 4 all-gather (@sharded(axis=V) → @replicated) written as existing ops. The compiler still emits host insert_slice for static concat; it does not emit kina_collective_all_gather.

4. Operational meaning

Attention, residual, and the QKV / Wo projections are the storylm0_forward algebra (replicated). Only the last matmul is sharded.

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

The result must match the unsharded residual @ broadcast(lm_head) on the same weights. Changing the last prompt token must not alter earlier logit rows (same causal non-leak as the forward check).

5. Typing (derived)

Each line is existing slice / broadcast / matmul / concat typing. V = 8 is static and divisible by 2.

6. Surface instance

storylm0_tp_vocab.kina binds tp_logits and storylm0_tp_vocab. @main calls the full graph with host placeholders. The numeric check loads both fns from the same .so.

No new tokens. The steps sit outside any reverseMode function. @mesh / @sharded / @replicated do not appear.

7. Tests

storylm0_tp_vocab.kina — §6 program; golden accepted by mlir-opt; storylm0_tp_vocab result is tensor<1x4x8xf32>.
scripts/validate_tp_vocab.py compiles the example to a .so, injects the storylm0_forward weights and prompt [1, 4, 2, 6], and matches (a) unsharded forward logits, (b) each shard’s [1, 4, 4] columns, and (c) the causal non-leak on a changed last token (NumPy; the same algebra in PyTorch when present).

8. Out of scope

Parsing or type-checking @mesh, @sharded, @replicated.
GSPMD propagation and collective insertion (Sections 3–4).
The Section 1 4-device threaded simulator.
The ~100M-parameter Phase 3 exit demo.
Head-parallel attention (now Section 6), embedding-row parallel, or contracting-axis (all-reduce) parallelism.
Device placement or NCCL.
Differentiating through the shards.
Phase 1 exit criteria.

9. Open issues

Mesh syntax in .kina. Section 2’s @mesh([4]) sketch is not in the bootstrap grammar. Do not guess a token.
4-way vocab split. V = 8 also divides by 4; this instance stays 2-way.
Head-parallel attention. Specified in Section 6.
Collective lowering. concat is not kina_collective_all_gather.