Phase 3, Section 6: StoryLM-0 head-parallel attention

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 rank-4 matmul / softmax on the existing storylm0_forward.kina graph.

1. What this section commits to

Section 5 locked 2-way column-parallel lm_head and left head-parallel attention open. H = 2 is the next StoryLM-0 axis: each logical device owns one head.

The instance is 2-way head-parallel attention:

• Logical mesh size 2.
• After the replicated QKV reshape, Q / Kᵀ / V are sliced on the head axis.
• Each shard attends with H = 1.
• concat on the head axis is the all-gather of the per-head context.

Wo and lm_head stay replicated. There is no @mesh keyword.

2. Why not a primitive

slice, rank-4 matmul, softmax, broadcast, and concat are already specified. A fused head-parallel op would invent a mesh. Splitting the head axis is the functional encoding of “shard attention on H, all-gather context.”

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]. H is divisible by the mesh size 2 (Section 2 §3.1).

Shard and gather (one head shown; head 1 uses start=1, stop=2):

q0     = slice[axis=1, start=0, stop=1] q      : [1, 1, 4, 2]
k0     = slice[axis=1, start=0, stop=1] k_t    : [1, 1, 2, 4]
v0     = slice[axis=1, start=0, stop=1] v      : [1, 1, 4, 2]
scores = q0 @ k0                               : [1, 1, 4, 4]
probs  = softmax(scores * broadcast(scale) + broadcast(mask), axis=3)
ctx0   = probs @ v0                            : [1, 1, 4, 2]
context = concat[axis=1] ctx0 ctx1             : [1, 2, 4, 2]

concat is the Section 4 all-gather (@sharded(axis=H) → @replicated) written as existing ops. The compiler still emits host insert_slice for static concat.

4. Operational meaning

QKV projections, Wo, and lm_head are the storylm0_forward algebra (replicated). Only the attention contract is sharded.

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

The result must match the unsharded batched-head attention on the same weights. Changing the last prompt token must not alter earlier logit rows.

5. Typing (derived)

Each line is existing slice / broadcast / rank-4 matmul / softmax / concat typing. H = 2 is static and divisible by 2. A single-head shard is rank 4 with a size-1 head axis, not a rank drop.

6. Surface instance

storylm0_tp_heads.kina binds tp_heads and storylm0_tp_heads. @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_heads.kina — §6 program; golden accepted by mlir-opt; storylm0_tp_heads result is tensor<1x4x8xf32>.
• scripts/validate_tp_heads.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 head’s [1, 1, 4, 2] context, 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).
• Column-parallel Wq / Wk / Wv (now Section 7).
• Combining this slice with Section 5 vocab-parallel in one program.
• The 4-device simulator or the ~100M exit demo.
• Device placement or NCCL.
• Differentiating through the shards.
• Phase 1 exit criteria.

9. Open issues

1. Mesh syntax in .kina. Still unspecified against the bootstrap grammar.
2. Projection-weight shards. Specified in Section 7.
3. Vocab + heads together. A 2-axis instance.
4. Collective lowering. concat is not kina_collective_all_gather.