feat(examples): Qwen-Image 20B MMDiT text-to-image on Trainium

[ymwangg]

Summary

Adds a device-only, TP-mandatory Qwen-Image 20B MMDiT text-to-image pipeline to examples/models/qwen_image/. The full pipeline runs on trn2 (TP=4):

• Denoiser — 60-block dual-stream MMDiT (3D RoPE, QK-RMSNorm, per-stream modulation) with fused on-device CFG + FlowMatchEuler sampling; the packed latent stays resident across all steps and the host feeds only per-step scalars.
• Text encoder — prefill Qwen2.5-VL decoder LM (28 layers, GQA 28/4, SwiGLU), reusing the qwen3 block kernels, Megatron-sharded. Device rel_l2 6.7e-3 bf16.
• VAE decoder — fp32 one-shot; at T=1 the WAN-style 3D causal video VAE collapses exactly to a 2D conv decoder. Device rel_l2 4.4e-6 at 64×64 → 512×512.
• Host keeps only trivial glue (tokenizer/chat template, embedding lookup, latent pack/denorm, scalar flow-match schedule).

Produces correct 512px images. End-to-end demo.sh reuses the compile cache, runs CPU correctness tests, then generates via torchrun.

Notable details

• RoPE perf fix: interleave-scatter rewrite drops the denoiser from 14.1s → 0.30s/step. Qwen uses the interleaved (2i, 2i+1) pair convention, not qwen3's (i, i+half) split.
• VAE denorm: the pipeline stores latents_std as its reciprocal (latents / std_recip); dividing by raw std instead is off by ~std².
• 1024px limit: the 60-block denoiser unrolls into one HLO graph; at grid 64×64 the serialized proto exceeds protobuf's 2 GB limit. Works today at ≤512px; durable fix is a device-side scan over one compiled block.

Test plan

• CPU correctness tests: uv run pytest examples/models/qwen_image/tests/
• On-device TP check (opt-in, needs hardware): QWEN_IMAGE_TP_DEVICE_TEST=1 uv run pytest tests/test_tp_device.py
• Validated end-to-end on trn2 at TP=4, 512px.

See examples/models/qwen_image/README.md for full run instructions and the durable lessons/limits.