Phase 3, Section 7: StoryLM-0 QKV projection shards

Phase 3: Sharding (the distribution axis). Not a Phase 3 exit-criteria replacement. Not a new primitive. Uses Phase 0 slice / concat / reshape / transpose and Section 6 per-head attention on the existing storylm0_forward.kina graph.

1. What this section commits to

Section 6 locked 2-way head-parallel attention after a replicated QKV projection, and left column-parallel Wq / Wk / Wv open. H · Dh = 4 is the next StoryLM-0 axis: each logical device owns one head’s projection columns.

The instance is 2-way column-parallel QKV:

• Logical mesh size 2.
• Wq / Wk / Wv [D, H·Dh] = [4, 4] are sliced on axis 1 into two [4, 2] shards.
• Each shard projects x to that head’s flat Q / K / V, then reshapes to H = 1.
• Per-head attention is the Section 6 contract.
• 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-3 matmul, reshape, transpose, rank-4 matmul, softmax, broadcast, and concat are already specified. A fused projection-shard op would invent a mesh. Splitting the H·Dh columns is the functional encoding of “shard QKV on the output axis, attend locally, 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·Dh is divisible by the mesh size 2 (Section 2 §3.1).

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

wq0     = slice[axis=1, start=0, stop=2] wq     : [4, 2]
q_flat0 = x @ broadcast(wq0)                    : [1, 4, 2]
q0      = reshape/transpose q_flat0             : [1, 1, 4, 2]
k0      = reshape/transpose (x @ broadcast(wk0)): [1, 1, 2, 4]
v0      = reshape/transpose (x @ broadcast(wv0)): [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]

q_flat0.reshape([1, 4, 1, 2]).transpose([0, 2, 1, 3]) is the Section 5 Q layout with H = 1. Keys use transpose([0, 2, 3, 1]) so Kᵀ is [1, 1, 2, 4].

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

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

x        = gather embed token_ids
context  = tp_proj(x, Wq, Wk, Wv, 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. proj_q0(x, Wq) must match the unsharded first-head query after reshape.

Row-parallel Wo would need an all-reduce after the output projection. That collective is unspecified; Wo stays replicated.

5. Typing (derived)

Each line is existing slice / broadcast / rank-3 matmul / reshape / transpose / rank-4 matmul / softmax / concat typing. H·Dh = 4 is static and divisible by 2. A single-head shard is rank 4 with a size-1 head axis, not a rank drop. Rank-3 [1, 4, 4] @ [1, 4, 2] is the existing batched matmul.

6. Surface instance

storylm0_tp_proj.kina binds proj_q0, tp_proj, and storylm0_tp_proj. @main calls the full graph with host placeholders. The numeric check loads those 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_proj.kina — §6 program; golden accepted by mlir-opt; storylm0_tp_proj result is tensor<1x4x8xf32>.
• scripts/validate_tp_proj.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) proj_q0 against the first-head query, (c) tp_proj against the concatenated context, and (d) 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).
• Row-parallel Wo (would need an all-reduce).
• 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. Vocab + heads together. A 2-axis instance.
3. Collective lowering. concat is not kina_collective_all_gather.
4. Row-parallel Wo. Needs a specified all-reduce.