Closing the gap: Wan 2.2 14B VAE decoder - tt-xla vs tt-metal

[vkovinicTT]

Here we discuss the gap between our compiler path for the Wan 2.2 VAE decoder and
tt-metal's hand-written one, and the changes needed to close it, ordered by performance impact. Suggestions welcome :)

[vkovinicTT]

Code links point at tt-metal/tt-mlir main. A few tt-mlir conv3d-config pieces aren't upstreamed yet, so those link to the vkovinic/wan_galaxy branch instead (marked (branch)).

Headline. The two biggest levers are running full-T instead of the frame-by-frame loop and spatial sharding (which we can't do yet - no halo CCL). A symptom of how broken the sharded path is today: the unsharded decoder on 1 chip is currently faster than our channel-sharded version on 4 chips, more detailed explanation later.

tt-xla vs tt-metal - gaps by priority

Table of gaps ordered by impact. For more details read section below.

#	Area	tt-xla (ours)	tt-metal	Owner	Effort	Impact
1	Temporal	decode one latent frame at a time (sequential loop); dynamo unrolls all 21 iterations into one graph → every op runs ×21	full-T: all frames in one forward pass - Galaxy's large DRAM holds the whole-clip activations (a few large chunks on smaller meshes)	model	small	highest
2	Parallelism / CCL	channel sharding (Megatron col→row), resblocks only; upsamplers / attention / boundary convs replicated → two all-reduce sources (Megatron output sum + small per-norm channel-sum)	spatial H+W fracture on everything via halo-exchange CCL; no channel all-reduce	compiler	large	highest
3	conv3d blocking	blocking from a static full-channel lookup table; sharded per-device channels so we miss the table values → most convs use the worst-case (1,1,1,32,32) default	hand-swept _BLOCKINGS per (mesh, res, layer)	compiler	medium	high
4	RMSNorm	decomposed F.normalize chain, no fused op (+ a channel-sum all-reduce when sharded)	one fused rms_norm + SiLU kernel	compiler	medium	medium
5	Attention	per-frame T=1, replicated, serial (×21)	full-T → independent B·T frames fractured across mesh	model (+compiler)	medium	medium
6	Misc	1×1 convs kept as conv3d kernel-1; NCDHW↔NDHWC permutes around convs	1×1 convs as Linear; NDHWC end-to-end	compiler	small	low–med

#1 is small-effort / highest-impact → the natural first move; #2 is large-effort but unlocks full-T at 720p and resolves #3's root cause. #5 (attention) stays medium on purpose: it's one block at the cheapest (latent) resolution, and its B·T parallelism is a free byproduct of #1 + #2 (you need T>1 and a distribution mechanism, which those provide) - so there's little standalone win
to chase here.

---

More detailed overview

The gaps reinforce each other, so they read best as one story. The decoder is slow for two compounding reasons: we run it one frame at a time (duplicating the whole graph ~21×), and we shard across channels - the only axis available without a halo CCL - which is also the axis that breaks conv3d.

1. Full-T temporal pass - kill the ×21 unroll · model · small · highest

The latents flowing through the decoder are large, so diffusers decodes the video as a stream - one latent frame at a time (AutoencoderKLWan._decode loops for i in range(num_frame), carrying a feat_cache between frames) to keep peak memory low. Under torch.compile dynamo unrolls that loop, so the graph ends up with all ~21 frame iterations laid out back-to-back and run sequentially - each on a tiny T=1 slice, so the hardware is launch-bound rather than doing real work.

On Galaxy memory isn't the constraint, so we don't need to stream: we can run all frames through the decoder in one forward pass (full-T), exactly what metal does there (vae_t_chunk_size=None). For us it's a model-side patch - drop the per-frame loop and feat_cache and call self.decoder(x) on the whole clip. It's #1 because it's a multiplier: collapsing the ×21 shrinks the conv3d (#3) and CCL costs by the same factor for almost no effort.

Where to look:

• diffusers [AutoencoderKLWan._decode (https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_wan.py#L1187-L1205) (the
  for i in range(num_frame) loop)
• metal's preset that selects full-T,
  pipeline_wan.py#L82-L100.

2. Spatial (H/W) sharding via halo-exchange CCL - the keystone · compiler · large · highest

We shard across channels (Megatron col→row on the resblocks) only because the compiler has no halo-exchange CCL
(discussion, #8216 issue). It works but it's the wrong axis, and it causes three problems at once: the expensive high-res upsamplers / conv_out aren't in the pattern so they stay replicated; the row-parallel half of each block needs an all-reduce; and splitting channels changes per-device channel counts, which is what makes conv3d miss its
blocking table (#3).

Metal shards the spatial dims instead: neighbor_pad_persistent_buffer is a halo exchange that gives each device its neighbours' border rows/columns (folding the padding-mask and causal-T-pad into the same CCL), so convs run locally on their H/W tile. Channels stay whole (so blocking works) and the activation splits across up to 32 devices (so full-T fits at 720p). The fix is heavy - add a halo CCL - but it's the keystone: unblocks full-T at high res, removes the channel all-reduces, and fixes #3's root cause.

Where to look:

• metal's halo CCL inside the conv forward, vae_wan2_1.py#L464-L474;
• our channel-sharding specs, shard_vae_decoder_specs / _megatron_pair_specs here.

3. conv3d blocking config · compiler · medium · high

Every ttnn.conv3d carries a blocking config that decides how its work tiles across cores - and that's most of conv3d performance. In our graph 650 of 692 convs get the worst-case default (1,1,1,32,32); only the replicated conv_in/conv_out are tuned. Why: blocking comes from one static table mirroring tt-metal's full-channel _DEFAULT_BLOCKINGS, but channel sharding feeds it reduced per-device channels that aren't in the table (and 3×1×1/1×1×1 kernels have no rows), and unlike conv2d, conv3d does no cost-model search to recover, even though an OpModel<Conv3dOp>
exists.

This is the reason of why sharded loses to unsharded (which keeps channels full, so it does hit the table). Fix: drive conv3d blocking from the OpModel like conv2d - shape-driven, so it covers the sharded shapes and odd kernels automatically. Or by sharding the model spatially we will have channels replicated, so we will hit the table for conv3d config.

Where to look:

The lookup table / heuristic in tt-xla:

• (branch) Conv3dConfigHeuristic.cpp#L36-L73,
  where it's applied
• (branch) LegalOpConfigAnalysis.cpp#L143-L188,
  the no-search rule book
• (branch) OpRules/ConvRules.cpp#L140-L149,
• TTNNOpModel.cpp#L5190,
• where the trivial default is stamped TTIRToTTNN.cpp#L2070

What tt-metal folks do (it seems they have even bigger static table hahah)

• metal's real per-(mesh,res) table conv3d.py#L150
  (with _DEFAULT_BLOCKINGS #L456).

4. RMSNorm not fused · compiler · medium · medium

Our WanRMS_norm (diffusers) is written with F.normalize, so it lowers to a chain of primitives rather than one norm op:

abs(x) → pow(·,2) → sum(channel) → sqrt → clamp(ε) → div(x, ·) → ·(√D·γ)

That is RMSNorm (the √D scale turns sum-of-squares into mean-of-squares), just written with div/sum/clamp over the channel dim instead of the mul·rsqrt/mean-over-last-dim shape the compiler's RMSNormFusionPattern looks for - so it never matches, and on the sharded channel axis each norm also fires an all-reduce.

We can fuse it by extending RMSNormFusionPattern to accept this form too, collapsing the chain into one ttnn.rms_norm (the sharded reduce then rides the existing DistributedRMSNorm path). tt-metal goes one step further: it doesn't just use rms_norm, it fuses the SiLU into the same kernel (dit_rms_norm_unary_fused(..., activation=SILU)) - which we could match as well
once the norm fuses.

Where to look:

• our matcher that won't match: TTIRFusing.cpp#L2623-L2630
• metal's fused norm+SiLU: normalization.py#L54-L61 (dit_rms_norm_unary_fused(..., activation=...)), used with SiLU at vae_wan2_1.py#L520

5. Replicated high-res layers + attention · model (+compiler) · medium · medium

The decoder has one attention block (mid-block): a spatial self-attention where, within a frame, every position attends to every other - so the sequence is the H·W grid and the batch is B·T, each (batch, frame) an independent problem.

That independence is the whole story: frame-by-frame leaves B·T = 1, so we run the attention 21× in sequence and, since the mid-block isn't sharded, replicated on every device (3 of 4 chips repeating identical work). Metal, on full-T, has all frames at once and just fractures the B·T batch across the mesh, running them in parallel.

This is pending since we first need to run full all frames at once which we mentioned in #1 in order for us to benefit this.

6. Smaller compiler wins · compiler · small · low-med

Two independent cleanups. The 1×1×1 convs (post_quant_conv, conv_shortcut) are really matmuls but we keep them as conv3d, so they ride the same blocking machinery and bad default; metal makes them Linear. And since TTNN conv3d is channels-last (NDHWC) while our torch graph is channels-first (NCDHW), the compiler wraps every conv in a permute-in/permute-out pair - pure data shuffling - and back-to-back convs leave inverse permutes an EraseInverseOps-style pass could cancel.

Neither is a bottleneck; the layout churn mostly vanishes once convs stop bouncing layouts.