Sharding of WAN 2.2 DiT

[vkovinicTT]

Here we discuss sharding of DiT through compiler, differences between compiler and tt-metal's sharding and gaps we need to fill in order to catch up to tt-metal folks.

[vkovinicTT]

Here we walk through how we shard the Wan 2.2 14B DiT (WanTransformer3DModel) - the dual-expert transformer that runs every denoising step in the Wan 2.2 A14B pipeline. This is the DiT counterpart to the VAE sharding discussion. Implementation lives in shared.py::shard_dit_specs and apply_dit_sp_activation_sharding.

---

DiT structure - what we're sharding

WanTransformer3DModel  (~14B params per expert, 2 experts in A14B)
├── rope                  WanRotaryPosEmbed  (buffers, no params)
├── patch_embedding       Conv3d(C_in → D, k=(1,2,2), s=(1,2,2))      ← entry, replicated
├── condition_embedder    WanTimeTextImageEmbedding
│   ├── time_embedder     Linear + SiLU + Linear  (→ D)               ← Megatron pair
│   ├── time_proj         Linear(D → 6·D)                             ← replicated
│   └── text_embedder     Linear + GELU + Linear  (→ D)               ← Megatron pair
├── blocks[0..N-1]        WanTransformerBlock × 40                    ← the bulk
│   each block:
│     ├── norm1, norm3    FP32LayerNorm  (non-affine)
│     ├── attn1           WanAttention (self, with RoPE)              ← Megatron pair
│     ├── attn2           WanAttention (cross, K/V from prompt)       ← Megatron pair
│     ├── norm2           FP32LayerNorm  (affine)
│     ├── ffn             FeedForward (D → ffn_dim → D, gelu-approximate) ← Megatron pair
│     └── scale_shift_table  Parameter (1, 6, D)                      ← replicated
├── norm_out              FP32LayerNorm  (non-affine)
├── proj_out              Linear(D → out_channels · prod(patch_size)) ← exit, replicated
└── scale_shift_table     Parameter (1, 2, D)                         ← replicated

The model has num_heads = 40 attention heads and num_layers = 40 blocks. Per forward, the activation enters the block stack flat as (B, L, D) where L = F'·H'·W' is the number of patchified video tokens.

---

Two sharding axes

Following tt-metal's terminology:

Role	What it shards	Size on (2, 4) mesh	Size on (8, 4) mesh
TP (tensor-parallel)	hidden dim D and FFN inner dim - head-aligned	4	4
SP (sequence-parallel)	flat token dim L of activations	2	8

We fix TP=4 on every production mesh because num_heads = 40 = 10 × 4 divides cleanly and each device ends up with a whole number of attention heads. SP scales with whatever's left over after pinning TP.

shard_dit_specs handles TP via weight annotations. apply_dit_sp_activation_sharding handles SP via sharding_constraint ops on intermediate activations - since SP shards a dim (L) that has no analog in any parameter.

[vkovinicTT]

Why Megatron is right for the DiT

A DiT block is textbook Megatron-LM: two column→row matmul pairs (attention QKV+O, FFN up+down) with a non-linearity between. Same template the VAE residual block uses, just at transformer scale:

• Column-parallel (to_q/k/v, ff1): split the output dim. The output is head-aligned (40 heads split 4 ways → 10 heads per device), so each device naturally owns a contiguous head group. No CCL.
• Row-parallel (to_out, ff2): split the input dim. Consumes the partitioned activation directly; the matmul output is a partial sum that needs one all-reduce to materialize.

Three such pairs per block give us 3 all-reduces per block at minimum (attn1.to_out, attn2.to_out, ffn.ff2). No all_gather between the column and row sides - the activation stays sharded across the non-linearity. The norms inside attention (norm_q/k) reduce over the full hidden dim, which is sharded, so each one needs a small stats AR.

---

Block-level sharding strategy

Inside the block, the activation flips between two sharding states depending on whether we're inside or outside a Megatron pair:

STATE A - between Megatron pairs   :  TP-replicated D, SP-sharded L
   per-device shape: (B, L/sp, D)

STATE B - inside a Megatron pair   :  TP-sharded D, SP-sharded L
   per-device shape: (B, L/sp, D/tp)

The transitions are baked into the matmul shapes:

• STATE A → STATE B at every column-parallel projection (to_q/k/v, ff1) - pure local matmul, no CCL.
• STATE B → STATE A at every row-parallel projection (to_out, ff2) - emits a partial sum, AR materializes the full D.

The block's entry and exit both sit in STATE A, so the 40 blocks chain transparently.

---

One block, end to end

                  spatial : STATE A
                      │
                      ▼
              ┌── norm1 ─────────────────────────────────────────────┐
              │  non-affine LN + adaLN (scale_msa, shift_msa)        │
              │  no params, fully local                              │ 0 CCL
              └──────────────────────────────────────────────────────┘
                      │
              ┌── attn1 (self-attention) ────────────────────────────┐
              │  to_q/k/v   col-parallel  STATE A → STATE B          │
              │  norm_q/k   RMSNorm across full D, γ TP-sharded      │ 2 small stats AR
              │  head split + RoPE apply (local)                     │
              │  AG of K, V along SP axis on L                       │ 2 large SP-AG
              │  SDPA       per-device on local Q × full K, V        │
              │  to_out     row-parallel  STATE B → STATE A          │ 1 large TP-AR
              └──────────────────────────────────────────────────────┘
                      │
                      + gate_msa * attn1_out  (STATE A + STATE A)
                      ▼
              ┌── norm2 ─────────────────────────────────────────────┐
              │  affine LN, γ/β replicated, fully local              │ 0 CCL
              └──────────────────────────────────────────────────────┘
                      │
              ┌── attn2 (cross-attention to prompt) ─────────────────┐
              │  to_q       col-parallel on spatial                  │
              │  to_k/v     col-parallel on prompt (already SP-rep)  │ no K/V AG needed
              │  norm_q/k   stats AR                                 │ 2 small stats AR
              │  SDPA       per-device, no ring                      │
              │  to_out     row-parallel + AR  STATE B → STATE A     │ 1 large TP-AR
              └──────────────────────────────────────────────────────┘
                      │
                      + attn2_out  (STATE A + STATE A)
                      ▼
              ┌── norm3 + adaLN ─────────────────────────────────────┐ 0 CCL
              └──────────────────────────────────────────────────────┘
                      │
              ┌── ffn ───────────────────────────────────────────────┐
              │  net[0].proj  col-parallel  STATE A → inner-sharded  │
              │  GELU         local on inner-sharded                 │
              │  net[2]       row-parallel + AR                      │ 1 large TP-AR
              └──────────────────────────────────────────────────────┘
                      │
                      + c_gate_msa * ffn_out  (STATE A + STATE A)
                      ▼
                  spatial : STATE A

Per-block CCL bill at the largest mesh we care about (Galaxy 8×4, sp=8, tp=4):

Op	Count	Size
attn1 K and V SP-AGs	2	big
attn1.to_out AR	1	big
attn2.to_out AR	1	big
ffn.net[2] AR	1	big
norm_q/k stats ARs (× 2 attns)	4	small
Total	5 big + 4 small

On the 1×4 mesh, SP collapses to a no-op so the two K/V SP-AGs disappear - that block costs just 3 big + 4 small CCLs.

---

Why cross-attention is "free" on the SP axis

In attn1 (self-attention), K and V are derived from spatial so they share its SP-sharded L axis. Each device only sees its own slice of L tokens, but SDPA needs every Q token to attend against the full K, V - so we all-gather K and V along SP before the dot-product.

attn2 (cross-attention) is different. K and V come from the prompt, which has its own L_prompt axis that's unrelated to the video's L. The prompt is replicated on the SP axis to start with, so attn2's K and V are already SP-replicated and no SP-AG is needed. That's why cross-attention costs half the CCL of self-attention.

[vkovinicTT]

Sequence parallelism - sharding the L axis of activations

SP shards the activation's L axis across sp_factor devices. Since no weight has an L axis, we can't express SP via shard_dit_specs - we annotate intermediate activations directly with sharding_constraint ops.

The dataflow at runtime:

input (5D, replicated on all axes)
       ↓ patch_embedding (replicated weight) → 5D (B, D, F', H', W') replicated
       ↓ flatten(2).transpose(1, 2)          → 3D (B, L, D) replicated
       ↓  ★ sharding_constraint "replicated"   ★ sharding_constraint "L on SP"   ← introduce SP here
       ↓
       (B, L/sp, D)  ← STATE A
       ↓ block 0
       ↓ ...
       ↓ block 39 → (B, L/sp, D) ← STATE A
       ↓ proj_out (replicated weight)        → (B, L/sp, 64)
       ↓  ★ sharding_constraint "replicated"  ← AG L back to full before unpatchify
       ↓ reshape → permute → flatten ×3
       output (replicated)

The interesting part is the double constraint at the block-stack entry (and an analogous one inside rope.forward). It's there to work around a tt-mlir lowering issue with reshape:

What we want:                            What happens with a single constraint:

    5D replicated                            5D replicated
        ↓                                       ↓
    reshape (5D → 3D)                       reshape (5D → 3D)
        ↓                                       ↓
    sharding_constraint("L on SP")          sharding_constraint("L on SP")
        ↓                                       ↑ ← Shardy back-propagates
    block stack ...                              │   the shard INTO the
                                                 │   reshape, partitions one
                                                 │   of (F', H', W'), and
                                                 │   tt-mlir fails to update
                                                 │   the reshape's static
                                                 │   result type to per-device
                                                 │   shape → "number of elements
                                                 │   doesn't match" error

The fix is two constraints back to back, both downstream of the reshape:

    reshape (5D → 3D)
        ↓
    sharding_constraint(replicated)   ← terminates back-propagation here
        ↓
    sharding_constraint(L on SP)      ← shard requested between two non-reshape ops
        ↓                                  → Shardy inserts a clean scatter
    block stack ...                         (reshape is never partitioned)

We apply this pattern in two places: at the block-stack entry (after flatten+transpose) and on the rope output (rope's body ends in a similar reshape). The exit at proj_out only needs a single constraint because the direction is reversed there - we're going sharded → replicated upstream of the unpatchify reshape, which doesn't trigger the bug; the all-gather lands cleanly before the reshape.

[vkovinicTT]

Full DiT layout - what's sharded where

WanTransformer3DModel
├── rope                  buffers, replicated; output marked SP-on-L via forward hook
├── patch_embedding       weight replicated, output stays replicated
├── condition_embedder
│   ├── time_embedder     Megatron col→row pair (TP on D)
│   ├── time_proj         replicated  (lone col-par would cost an AG every block)
│   └── text_embedder     Megatron col→row pair (TP on D)
├── blocks[0..39]
│   ├── norm1, norm3      no params (non-affine LN)
│   ├── norm2.weight/bias replicated  (matches STATE A)
│   ├── attn1, attn2
│   │   ├── to_q/k/v      ("tp", None)        col-parallel on D
│   │   ├── norm_q/k      ("tp",)             gamma matches sharded act
│   │   └── to_out[0]     (None, "tp")        row-parallel + AR → STATE A
│   ├── ffn
│   │   ├── net[0].proj   ("tp", None)        col-parallel on inner
│   │   └── net[2]        (None, "tp")        row-parallel + AR
│   └── scale_shift_table replicated
├── norm_out              no params
├── proj_out              replicated  (one explicit AG hook to force replicated output)
└── scale_shift_table     replicated  (final-layer adaLN)

Per forward at sp=2 we end up with ~5 large CCLs and ~4 small stats ARs per block (× 40 blocks), plus a handful of one-time CCLs in condition_embedder and at the block-stack boundaries. At 480p, warm time drops from ~77 s (TP-only on 1×4) to ~41 s (TP+SP on 2×4) - roughly the 2× speedup we expected from sp=2. PCC dips modestly (from 0.9950 to 0.9247), which is a separate compiler-precision wrinkle we're still chasing - not a structural problem with the sharding design.

---

How we differ from tt-metal - and what we'd need to close the gap

tt-metal's Wan 2.2 DiT (see metal implementation) is structurally very close to ours - Megatron col→row pairs, TP on heads, SP on L. The gap comes from two things they have and we don't: fused matmul-with-CCL kernels (plus ring SDPA) at the kernel level, and a hand-written model they're free to restructure however they like. Our diffusers-based path leaves the model untouched, so any equivalent of their model-side rewrites has to come from compiler pattern-matching - or from us monkey-patching diffusers.

Side-by-side per-block

	tt-xla (this design)	tt-metal	What metal uses	Kind of gap
Activation contract between Megatron pairs	STATE A: replicated D on TP (Plan A)	STATE B: sharded D on TP (Plan B)	enabled by fused AG+MM	kernel fusion
Self-attn K/V across SP	explicit AG of K, V before plain SDPA	overlapped via ring SDPA	ring_joint_scaled_dot_product_attention	compiler primitive
Col-parallel matmul (e.g. to_qkv, ff1)	AG of input → matmul (two ops)	fused AG+MM in one kernel	all_gather_minimal_matmul_async	kernel fusion
Row-parallel matmul (e.g. to_out, ff2)	matmul + AR (two ops)	fused MM+RS in one kernel	minimal_matmul_strided_reduce_scatter_async	kernel fusion
Gated residual add after attn/ffn	separate elementwise add+mul ops	fused into the MM+RS kernel	addcmul fused into RS write	kernel fusion
QKV projection	three separate Linears (diffusers default)	one fused matmul + chunk	ColParallelLinear(D, 3D, chunks=3)	model rewrite
RMSNorm + RoPE inside attention	stats-AR LN + separate RoPE apply	one fused kernel	wan_fused_rmsnorm_post_allgather	kernel fusion or model rewrite
Patch embedding	Conv3d (replicated weight); its downstream flatten+transpose triggers a reshape partitioner bug we work around with back-to-back constraints	Conv3d replaced by a hand-written reshape+permute+Linear - mathematically the same, but no reshape lands in the compiled graph	model rewrite

What we'd need to close the gap

Pure kernel / compiler work - model stays as-is:

1. Ring SDPA primitive. Eliminates the two large SP-AGs of K and V per self-attention block and overlaps K/V circulation with SDPA compute.
2. Fused AG+matmul (all_gather_minimal_matmul_async-equivalent). Turns the explicit AG of input into a phase of the matmul. Roughly halves the bandwidth on every col-parallel projection.
3. Fused matmul+reduce-scatter with optional addcmul write. Lets row-parallel projections emit per-device-sharded outputs instead of AR-materialized ones, and folds the residual+gate into the same kernel.
4. Fix the reshape partitioner bug. Lets us drop the back-to-back-constraint workaround at the block-stack entry and inside rope.forward. Alternative: monkey-patch the model the way metal did (replace Conv3d patchify with reshape+permute+Linear, precompute rope tables outside the traced graph) - but a compiler-side fix is the cleaner answer.

Either compiler pattern-match or model rewrite:

5. Fused RMSNorm + RoPE + head-split. As a compiler pass it's a pure kernel win; as a model rewrite it's a custom op call replacing norm + apply_rotary_emb. Either way it removes 2–3 dispatches per attention per block.

Model rewrite only:

6. Fused QKV projection. Replace to_q / to_k / to_v (three Linears) with one Linear(D → 3D) + chunk(3). diffusers ships a fuse_qkv_projections() for some pipelines; for Wan we'd need to monkey-patch the attention module. In principle a compiler pass could spot three concurrent matmuls reading the same input, but no such pass exists today.

The split matters because items 1–4 keep shard_dit_specs + apply_dit_sp_activation_sharding as the only thing we own - diffusers stays vanilla. Items 5–6 push us into model-surgery territory, where we'd be carrying patches against upstream diffusers. The performance picture above (warm time, PCC) is what falls out of items 1–4 alone; the kernel-level fusions are where the next big chunk lives.

[vkovinicTT]

Notes: 480p DiT slower after the SP tile-alignment padding

Adding the SP sequence-length padding to the DiT (pad L to a multiple of SP*32 so the
mesh_partition scatter is tile-aligned) makes 480p warm inference go from ~5.6s to ~8.9s/iter.
720p needs that padding to run at all on galaxy, but 480p ran fine without it, so on 480p this is pure
overhead.

The cost isn't the padding tokens (8 of them). To keep attention exact, the patch swaps the
attention processor to slice K and V back to the unpadded length before SDPA in every block.
Diffing the before/after runtime IR, everything is essentially the same (same matmuls, SDPA,
all-gathers, reduce-scatters, memory) except ttnn.slice_static goes from 776 to 857 - the
extra 81 are the per-block K/V slice-backs (40 blocks × K,V, plus one final hidden slice). That's where the ~3.4s goes.

Anyone picking this up: the lever is the K/V slice-back, not the padding.