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

examples/models/qwen_image/.gitignore

@@ -0,0 +1,17 @@
+# prepared weights (downloaded/converted via tensor_preparation.py)
+tmp_qwen_image/
+*.safetensors
+
+# per-rank sharded weight cache (qwen_image.py --weight-cache-dir)
+weight_cache/
+
+# compiled kernels
+build/
+build_*/
+
+# generated images
+*.png
+
+# python caches
+__pycache__/
+*.pyc

examples/models/qwen_image/README.md

@@ -0,0 +1,131 @@
+# Qwen-Image (20B MMDiT) on Trainium
+
+NKIPy port of [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image)
+(`QwenImagePipeline`, Apache-2.0) — a 20B dual-stream **MMDiT** text-to-image
+model with 3D RoPE, QK-RMSNorm, per-stream modulation, a Qwen2.5-VL text
+encoder, a 16-channel video-style VAE, and rectified-flow sampling.
+
+The port is **device-only** end-to-end: the denoiser, text encoder, and VAE all
+run on trn2 (TP=4, fused on-device sampling), producing correct 512px images.
+
+## Sample output
+
+![Qwen-Image sample: a coffee shop entrance with a chalkboard sign](assets/sample.png)
+
+*512px, TP=4, 50 steps, guidance 4.0 — prompt: "a coffee shop entrance with a
+chalkboard sign".*
+
+## Run
+
+```bash
+bash demo.sh "a coffee shop storefront with a chalkboard sign reading 'Qwen Cafe'"
+```
+
+`demo.sh` reuses the compilation cache, runs the CPU correctness tests, then
+generates one image via `torchrun`. Env knobs: `TP` `HEIGHT` `WIDTH` `STEPS`
+`OUTPUT`; `CLEAN=1` forces a recompile; `SKIP_TESTS=1` skips the CPU tests.
+
+Or drive it directly:
+
+```bash
+torchrun --nproc-per-node 4 qwen_image.py "your prompt" \
+    --height 512 --width 512 --steps 50
+```
+
+**Tensor parallelism is required** — 20B bf16 (~40 GB) doesn't fit on one 24 GB
+core. TP=4 fits (20.4 GB/core) and is the max: the text encoder caps TP at
+`num_kv_heads = 4`. The model is downloaded once by the diffusers pipeline (HF
+cache).
+
+### Startup caching & resident mode
+
+A cold launch is dominated by one-time host-side weight prep, not compile or
+generation, so two mechanisms amortize it:
+
+- **Shard cache** (on by default, `--weight-cache-dir`, default `./weight_cache`):
+  the per-rank sharded weights are extracted + written once, then reloaded on
+  every later launch — turning the ~350 s extract/shard into a memory-mapped
+  read (and dropping peak RAM from ~200 GB to ~45 GB). `--no-weight-cache`
+  disables it.
+- **Resident mode** (`--prompts-file FILE` or `--interactive`): keep the process
+  alive and generate an image per prompt. The weight upload (~25 s) and NEFF
+  load (~17 s) — which no on-disk cache can remove — are then paid once; each
+  additional image costs only `generate()` (e.g. ~2.8 s at 8 steps, ~17 s at 50).
+- **Text-length bucketing** (`--text-bucket N`, default 64): the text-encoder
+  sequence and the denoiser text length are rounded up to a multiple of `N`, so
+  varying prompts reuse one compiled kernel per bucket instead of recompiling
+  per exact length (a different-length prompt in a resident session then costs
+  ~2.8 s instead of a ~90 s recompile). Padding is exact — the encoder is causal
+  (real tokens never attend to the pad tail) and the denoiser masks the pad text
+  tokens — but because the kernel *shapes* change, the bf16 GEMM reduction order
+  differs, so a given prompt's pixels differ slightly (same scene) from the
+  `--text-bucket 0` (per-exact-length, bitwise-stable) result.
+
+```bash
+# one image per line, process stays resident
+torchrun --nproc-per-node 4 qwen_image.py --prompts-file prompts.txt --steps 50
+```
+
+CPU correctness tests: `uv run pytest tests/`. The on-device TP check needs
+hardware and is opt-in: `QWEN_IMAGE_TP_DEVICE_TEST=1 uv run pytest
+tests/test_tp_device.py`.
+
+## What runs on device vs host
+
+Everything with real FLOPs runs on device; the host keeps only trivial glue.
+
+- **Denoiser** — 60-block MMDiT, TP=4. The fused sampling step
+  (`QwenImageDenoiser.sample` / `kernels/transformer.py:denoise_step`) runs CFG +
+  FlowMatchEuler on device with the packed latent resident across all steps; the
+  host feeds only the per-step scalars `[true_cfg_scale, dt]` and once-uploaded
+  text embeds/masks.
+- **Text encoder** (`kernels/text_encoder.py`) — a prefill Qwen2.5 decoder LM (28
+  layers, hidden 3584, GQA 28/4, SwiGLU 18944, RoPE θ=1e6, qkv **with** bias, no
+  QK-norm), returning the last hidden state. Reuses the qwen3 block kernels;
+  sharded Megatron-style (q/o by query heads, k/v by KV heads, MLP by
+  intermediate). Device rel_l2 6.7e-3 bf16.
+- **VAE decoder** (`kernels/vae.py`, fp32, one-shot) — the T=1 2D-collapsed
+  decoder (see below). Device rel_l2 4.4e-6 at 64×64 latent → 512×512 image.
+  Replicated per rank.
+- **Host** keeps: tokenizer + chat template + embedding-table lookup, latent
+  pack / `_unpack_latents` + denorm, and the scalar `mu`/sigma flow-match
+  schedule — data-dependent / string-I/O / scalar work that buys nothing on
+  device.
+
+Weight footprint: the 7B encoder (~14 GB bf16) at TP=4 adds ~3.5 GB/core on top
+of the sharded denoiser (→ ~24 GB/core, tight; TP=8 gives headroom).
+
+The confirmed model structure lives in `config.py` (defaults match
+`Qwen/Qwen-Image`); `get_config`/`get_vae_config` read the diffusers config JSON
+and fall back to those defaults when diffusers is unavailable.
+
+## Lessons / limits (durable)
+
+Non-obvious things this port surfaced — worth knowing before touching the
+kernels or scaling the config.
+
+- **Interleaved RoPE convention.** Qwen RoPE uses the *interleaved* `(2i, 2i+1)`
+  pair convention, **not** qwen3's `(i, i+half)` split. Getting this wrong is a
+  subtle, hard-to-spot numerical mismatch. (`kernels/rope3d.py`, complex path
+  reformulated as real interleaved for the tracer.)
+- **VAE latent denormalization.** The pipeline stores `latents_std` as its
+  *reciprocal* and does `latents / std_recip` (== ×std). Dividing by the raw std
+  instead is off by up to std² (~10×) and produces dark/green-tinted output.
+- **1024px exceeds the 2 GB HLO-proto limit.** The whole 60-block denoiser is
+  unrolled into one HLO graph. At grid 64×64 (1024px, 4096 image + 34 text
+  tokens) the serialized graph exceeds protobuf's 2 GB limit and compile fails in
+  `to_proto().SerializeToString()`. This hits **both** the fused `denoise_step`
+  and the non-fused `qwenimage_forward`. **Works today:** ≤512px (grid ≤32×32).
+  The durable fix (a device-side scan/loop over one compiled block, weights
+  indexed per iteration) is independent of TP — TP shrinks per-core *weights*,
+  not *graph* size.
+- **Tracer dead-weight pruning.** The tracer prunes weights that don't reach the
+  output (e.g. the last block's text-stream MLP / out-proj — only the image
+  stream feeds `final_layer`), so the driver must filter runtime inputs to
+  `kernel.input_tensors_info`.
+- **VAE at T=1 collapses 3D→2D.** For text-to-image we decode a single frame, and
+  at T=1 the WAN-style 3D causal video VAE reduces *exactly* to a 2D conv decoder
+  (`w[:, :, -1]`, `feat_cache=None`, `time_conv` skipped). This is what made the
+  VAE port a day instead of a week.
+- **Dynamic timestep shifting.** Qwen-Image's scheduler passes
+  `mu = calculate_shift(image_seq_len)` to `set_timesteps` — match the pipeline.

examples/models/qwen_image/config.py

@@ -0,0 +1,171 @@
+from dataclasses import dataclass, field
+
+import numpy as np
+from neuronxcc.nki.language import bfloat16
+
+# to control compiler_args
+DTYPE = bfloat16
+
+
+@dataclass
+class Config:
+    """Qwen-Image (QwenImageTransformer2DModel) configuration.
+
+    Defaults match Qwen/Qwen-Image. The denoiser is a dual-stream **MMDiT**: 60
+    blocks, each running an image stream and a text stream through per-stream
+    LayerNorm + adaLN-style modulation, QKV with QK-RMSNorm + 3D RoPE, a single
+    joint attention over concat([text, image]), then per-stream gated MLPs.
+
+    The text tokens participate in the same attention as the image tokens
+    (joint attention over concat([text, image])), so there is no separate
+    cross-attention.
+    """
+
+    # MMDiT transformer
+    num_layers: int = 60
+    num_heads: int = 24
+    head_dim: int = 128
+    # hidden/inner dim derived: num_heads * head_dim == 3072
+    patch_size: int = 2
+
+    # latent space (VAE-encoded image); Qwen-Image uses a 16-channel VAE, and
+    # the DiT sees patchified latents so in_channels = 16 * patch_size**2 = 64.
+    in_channels: int = 64
+    out_channels: int = 16
+
+    # text conditioning (Qwen2.5-VL encoder, frozen)
+    joint_attention_dim: int = 3584  # Qwen2.5-VL hidden size (text stream input)
+    pooled_projection_dim: int = 768
+
+    # 3D RoPE axis dims (frame, height, width); sum == head_dim (16+56+56 == 128)
+    axes_dims_rope: tuple = (16, 56, 56)
+    rope_theta: float = 10000.0
+    guidance_embeds: bool = False
+
+    # sampling (FlowMatchEulerDiscreteScheduler, rectified flow)
+    num_inference_steps: int = 50
+    true_cfg_scale: float = 4.0
+
+    norm_eps: float = 1e-6
+    max_batch_size: int = 1  # x2 internally for classifier-free guidance
+    dtype: np.dtype = DTYPE
+    additional_compiler_args_nkipy: str = "--lnc 1"
+
+    # tensor parallelism (required): shard attention heads + MLP intermediate
+    # across tp_size cores (>=4 needed to fit the 20B weights on trn2's 24 GB/core).
+    # ``all_reduce_fn`` must be set by the driver to a real collective — the
+    # kernels always apply it after each row-parallel projection.
+    tp_size: int = 1
+    all_reduce_fn: object = None
+
+    @property
+    def hidden_size(self) -> int:
+        return self.num_heads * self.head_dim
+
+
+@dataclass
+class TextEncoderConfig:
+    """Qwen2.5 text-encoder config (the text-only LM inside Qwen2.5-VL).
+
+    Defaults match Qwen-Image's text_encoder. Used by the on-device encoder
+    (``kernels/text_encoder.py``); the host keeps the tokenizer, chat template,
+    and embedding-table lookup.
+    """
+
+    num_layers: int = 28
+    hidden_size: int = 3584
+    num_heads: int = 28
+    num_kv_heads: int = 4
+    head_dim: int = 128  # 3584 / 28
+    intermediate_size: int = 18944
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1000000.0
+    vocab_size: int = 152064
+
+    # chat-template wrapping (matches QwenImagePipeline)
+    prompt_template: str = (
+        "<|im_start|>system\nDescribe the image by detailing the color, shape, "
+        "size, texture, quantity, text, spatial relationships of the objects and "
+        "background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+    )
+    template_drop_idx: int = 34
+    tokenizer_max_length: int = 1024
+
+    dtype: np.dtype = DTYPE
+    additional_compiler_args_nkipy: str = "--lnc 1"
+    tp_size: int = 1
+    all_reduce_fn: object = None
+
+
+@dataclass
+class VAEConfig:
+    """AutoencoderKLQwenImage decoder config (the T=1 2D-collapsed decoder).
+
+    Defaults match Qwen-Image's VAE. Used by the on-device VAE decoder
+    (``kernels/vae.py``); the host keeps the latent denormalization and the
+    single-frame squeeze/unsqueeze.
+    """
+
+    z_dim: int = 16
+    base_dim: int = 96
+    dim_mult: tuple = (1, 2, 4, 4)
+    num_res_blocks: int = 2
+    out_channels: int = 3
+
+    dtype: np.dtype = np.float32  # VAE runs in fp32 (numerically sensitive)
+    additional_compiler_args_nkipy: str = "--lnc 1"
+
+    # kernel reads these short aliases
+    @property
+    def vae_dim_mult(self):
+        return self.dim_mult
+
+    @property
+    def vae_num_res_blocks(self):
+        return self.num_res_blocks
+
+
+def get_vae_config(model_name: str) -> "VAEConfig":
+    """Build a VAEConfig from a diffusers Qwen-Image checkpoint (falls back to
+    Qwen-Image defaults if diffusers is unavailable)."""
+    try:
+        from diffusers import AutoencoderKLQwenImage
+
+        hf = AutoencoderKLQwenImage.load_config(model_name, subfolder="vae")
+        return VAEConfig(
+            z_dim=hf["z_dim"],
+            base_dim=hf["base_dim"],
+            dim_mult=tuple(hf["dim_mult"]),
+            num_res_blocks=hf["num_res_blocks"],
+        )
+    except Exception:
+        return VAEConfig()
+
+
+def get_config(model_name: str, num_inference_steps: int) -> Config:
+    """Build a Config from a diffusers Qwen-Image checkpoint.
+
+    Reads the DiT (QwenImageTransformer2DModel) config JSON. Falls back to the
+    Qwen-Image defaults if diffusers is unavailable so the module stays
+    importable on machines without the dependency.
+    """
+    try:
+        from diffusers import QwenImageTransformer2DModel
+
+        hf = QwenImageTransformer2DModel.load_config(model_name, subfolder="transformer")
+        config = Config(
+            num_layers=hf["num_layers"],
+            num_heads=hf["num_attention_heads"],
+            head_dim=hf["attention_head_dim"],
+            patch_size=hf["patch_size"],
+            in_channels=hf["in_channels"],
+            out_channels=hf["out_channels"],
+            joint_attention_dim=hf["joint_attention_dim"],
+            pooled_projection_dim=hf.get("pooled_projection_dim", 768),
+            axes_dims_rope=tuple(hf.get("axes_dims_rope", (16, 56, 56))),
+            guidance_embeds=hf.get("guidance_embeds", False),
+            num_inference_steps=num_inference_steps,
+        )
+    except Exception:
+        config = Config(num_inference_steps=num_inference_steps)
+    return config

examples/models/qwen_image/demo.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+# End-to-end demo for Qwen-Image (20B MMDiT) on Trainium.
+# Usage: bash demo.sh ["your prompt"]
+#
+# Runs the CPU correctness tests and generates one image at 512px end-to-end on
+# device: the denoiser (fused CFG + FlowMatchEuler sampling loop), the Qwen2.5
+# text encoder, and the VAE decoder all run on the TP cores.
+#
+# The expensive one-time work is cached, so reruns are fast:
+#   * The ~40 GB model is downloaded once by the diffusers pipeline into the HF
+#     cache; the driver extracts + shards all weights in-memory (no repack step).
+#   * Compiled kernels are kept in build/ build_te/ build_vae/ and reused across
+#     runs (content-hash keyed, so kernel edits recompile automatically).
+#
+# Env knobs: TP HEIGHT WIDTH STEPS OUTPUT (see below); CLEAN=1 forces a kernel
+# recompile; SKIP_TESTS=1 skips the CPU tests.
+#
+# Notes:
+#   * The 20B model does not fit on one core (~40 GB bf16 > 24 GB/core), so it is
+#     tensor-parallel over $TP cores via torchrun. TP is required; TP=4 fits (and
+#     is the max — the text encoder caps TP at num_kv_heads=4).
+#   * 512px (grid 32x32) is the working resolution; native 1024px currently
+#     exceeds the compiler's 2 GB HLO-proto limit (see README.md).
+set -e
+
+MODEL="Qwen/Qwen-Image"
+TP=${TP:-4}                       # tensor-parallel degree (override: TP=8 bash demo.sh)
+HEIGHT=${HEIGHT:-512}
+WIDTH=${WIDTH:-512}
+STEPS=${STEPS:-20}
+PROMPT="${1:-a photorealistic coffee shop storefront at golden hour, a wooden chalkboard sign reading 'Qwen Cafe', warm cinematic lighting, highly detailed}"
+OUTPUT="${OUTPUT:-output.png}"
+
+echo "=========================================="
+echo "Qwen-Image (20B MMDiT) on Trainium — demo"
+echo "  TP=$TP  ${WIDTH}x${HEIGHT}  steps=$STEPS"
+echo "=========================================="
+
+# Step 1: compilation cache. Kept across runs so reruns skip the neuronx-cc
+# compiles (denoiser ~33s, text encoder, VAE). The caches are content-hash
+# keyed, so editing a kernel recompiles automatically — no stale-cache risk.
+# Force a clean rebuild with CLEAN=1.
+echo ""
+if [ "${CLEAN:-0}" = "1" ]; then
+    echo "[1/3] CLEAN=1 -> clearing compilation caches (build/ build_te/ build_vae/)..."
+    rm -rf build/ build_te/ build_vae/ 2>/dev/null || true
+    echo "OK caches cleared"
+else
+    echo "[1/3] Reusing compilation cache (set CLEAN=1 to force a recompile)."
+fi
+
+# Step 2: CPU correctness tests (kernels validated vs diffusers). Skip on
+# reruns with SKIP_TESTS=1.
+echo ""
+if [ "${SKIP_TESTS:-0}" = "1" ]; then
+    echo "[2/3] SKIP_TESTS=1 -> skipping CPU correctness tests."
+else
+    echo "[2/3] Running CPU correctness tests..."
+    python -m pytest tests/ -q
+fi
+
+# Step 3: generate one image end-to-end (TP; denoiser + text encoder + VAE on
+# device). The driver downloads the model (first run) and extracts + shards all
+# weights in-memory from the host pipeline.
+echo ""
+echo "[3/3] Generating an image (TP=$TP)..."
+echo "=========================================="
+torchrun --nproc-per-node "$TP" qwen_image.py "$PROMPT" \
+    --model "$MODEL" \
+    --height "$HEIGHT" --width "$WIDTH" --steps "$STEPS" \
+    --output "$OUTPUT"
+
+echo ""
+echo "=========================================="
+echo "OK Demo complete! Image written to $OUTPUT"
+echo "=========================================="