kernel tutorials · tt-metal · multi-chip

Megatron and CCLs.

The two Megatron matmul cuts, then every collective in ttnn — drawn in one visual language: cell color = which device the data came from, grey Σ = summed, dashed = empty. Everything below is vector — zoom as far as you like.

Megatron tensor parallelism: Z = (X·A)·B

Two ways to cut a weight across devices:

column-split — shards the output
A X Y0 Y1 Y2 Y3
  • needs the full X on every device (replicated input)
  • yields exact values — each device holds a different part of Y
  • comm none (GeLU applies locally)
row-split — shards the input
B Y0 Y1 Y2 Y3 Z0 Z1 Z2 Z3
  • needs a 1/N input slice — exactly what column-split left behind
  • yields full-size results, partial values
  • comm one all-reduce, at the very end

Lock and key: column-split leaves Y sharded exactly the way row-split wants its input — zero communication between the matmuls, one all-reduce at the end. Reversed, the partials would need an all-reduce before the GeLU.

The same pattern runs both halves of a transformer layer — GeLU in the MLP, softmax·V in attention; both sit safely in the middle because column-split values are exact:

MLP block X replicated W_up/gate col-split sharded · exact W_down GeLU (local) row-split full-size · partial all-reduce Z replicated
Attention block X replicated W_qkv packed q k v q k v q k v q k v col-split by heads q k v q k v q k v q k v Q·K·V heads · exact attention (local) sharded · exact W_o row-split full-size · partial all-reduce Z replicated

The full block — left operand beside the result, top operand above it: Y appears where X meets A, the partial Z's where Y meets B.

Device 0 Device 1 Device 2 Device 3 replicated A — 4×8, column-split A0 A1 A2 A3 B — 8×4, row-split B0 B1 B2 B3 X X — 2×4, replicated Y0 Y1 Y2 Y3 Y = X·A — 2×8, column-split Z0 Z1 Z2 Z3 Z0..Z3 = Yi·Bi — four full-size 2×4 partials All-Reduce (sum) Z Z — 2×4, replicated Yi is produced on the device that already owns Bi (same color) — no communication between the two matmuls, one All-Reduce at the very end.

Every ttnn CCL, on 4 devices

Each panel is before → after, one row per device. Move ops shuffle colored chunks around unchanged; reduce ops blend everyone's colors into a grey Σ. (reduce_to_root is currently SDPA-specialized in ttnn, and all_to_all ships as the MoE dispatch/combine pair — the panels show the underlying primitive.)

cell color = which device the data came from Σ = summed data (all colors mixed) empty / no data MOVE ops — data travels, nothing is computed (colors survive) broadcast D0 D1 D2 D3 The root's tensor is copied to every device (root = D0 here). all_broadcast D0 D1 D2 D3 broadcast from every root at once: each device ends with a LIST of all 4 tensors. send_recv (point-to-point) D0 D1 D2 D3 One sender copies to one receiver (here D0 → D2). The only non-collective op. all_gather D0 D1 D2 D3 Device i brings shard i → everyone gets the full concatenation. Memory grows ×4. mesh_partition D0 D1 D2 D3 Inverse of all_gather: each device keeps only slice i of the replicated input. all_to_all (dispatch / combine) D0 D1 D2 D3 Chunk j of device i moves to device j — a transpose of the grid (MoE token routing). REDUCE ops — tensors are summed elementwise (colors blend into grey Σ) reduce_to_root D0 D1 D2 D3 Sum of everyone's tensors lands only on the root (SDPA tree-reduction, root = D0). Σ all_reduce D0 D1 D2 D3 Everyone gets the full sum. = reduce_scatter + all_gather. Σ Σ Σ Σ reduce_scatter D0 D1 D2 D3 Everyone's tensors are summed; device i keeps only slice i of the sum. Shrinks ÷4. Σ Σ Σ Σ

How to remember them

  1. Two families, that's it. Data can only move (gather / broadcast / scatter / to-all) or be summed (anything with reduce in the name).
  2. all_ prefix = no root. Every device ends with the result. No all_ = one special root device: broadcast from root, reduce to root.
  3. scatter / partition = sliced result. Each device keeps only its 1/N piece — look for the staircase in the panels.
  4. Inverses. mesh_partition undoes all_gather: their panels are mirror images (staircase ↔ full rows).
  5. Compositions. all_reduce = reduce_scatter + all_gather — literally how it's implemented on a ring. all_broadcast = broadcast run from every root at once.
  6. all_to_all is a transpose. Horizontal color stripes become vertical ones. Picture the (device × chunk) grid rotating and you can't forget it.
  7. Memory check by counting cells. all_gather grows ×N, mesh_partition & reduce_scatter shrink ÷N, all_reduce & all_to_all keep size.

On the wire — the T3K ring, animated

T3K wires its 8 chips in a ring — each one talks only to its two neighbours. All eight drawn below, one row per device. Watch a hop at a time: move packets are copied on unchanged, reduce packets pick up a stripe at every stop and turn grey Σ when the last contribution lands.

✏️ Want to edit or reuse these diagrams? Grab the editable vector source: megatron-and-ccls.excalidraw — open it at excalidraw.com or with the VS Code Excalidraw extension. The SVGs above are rendered straight from that file, so nothing was rasterized.