diff --git a/examples/models/qwen_image/.gitignore b/examples/models/qwen_image/.gitignore
new file mode 100644
index 0000000..6021efc
--- /dev/null
+++ b/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
diff --git a/examples/models/qwen_image/README.md b/examples/models/qwen_image/README.md
new file mode 100644
index 0000000..72d889c
--- /dev/null
+++ b/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.
diff --git a/examples/models/qwen_image/assets/sample.png b/examples/models/qwen_image/assets/sample.png
new file mode 100644
index 0000000..69f9e4f
Binary files /dev/null and b/examples/models/qwen_image/assets/sample.png differ
diff --git a/examples/models/qwen_image/config.py b/examples/models/qwen_image/config.py
new file mode 100644
index 0000000..86cdddd
--- /dev/null
+++ b/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
diff --git a/examples/models/qwen_image/demo.sh b/examples/models/qwen_image/demo.sh
new file mode 100755
index 0000000..cbc4744
--- /dev/null
+++ b/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 "=========================================="
diff --git a/examples/models/qwen_image/kernels/__init__.py b/examples/models/qwen_image/kernels/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/models/qwen_image/kernels/attention.py b/examples/models/qwen_image/kernels/attention.py
new file mode 100644
index 0000000..b6ed427
--- /dev/null
+++ b/examples/models/qwen_image/kernels/attention.py
@@ -0,0 +1,130 @@
+"""Joint (double-stream) attention for the Qwen-Image MMDiT.
+
+Mirrors diffusers ``QwenDoubleStreamAttnProcessor2_0``. Qwen-Image runs **one**
+attention over the concatenation of the text and image tokens (rather than
+separate self- and cross-attention):
+
+    1. project image tokens with to_q/to_k/to_v, text tokens with
+       add_q_proj/add_k_proj/add_v_proj;
+    2. reshape to heads, apply per-head QK-RMSNorm (norm_q/norm_k for image,
+       norm_added_q/norm_added_k for text);
+    3. apply 3D RoPE to q/k of both streams;
+    4. concat in order [text, image] along the sequence axis;
+    5. scaled-dot-product attention (bidirectional, no causal mask); an optional
+       key-padding mask masks text padding tokens;
+    6. split back, project image out via to_out, text out via to_add_out.
+
+The score matrix is O((Ltext+Limg)^2); the image tokens dominate. A tuned
+device backend can be swapped in later — this hand-rolled softmax is the
+portable, CPU-testable reference.
+"""
+
+import numpy as np
+
+from .rmsnorm import rmsnorm_kernel
+from .rope3d import apply_rotary_emb
+from .softmax import softmax_kernel
+
+
+def _to_heads(x, n_heads, head_dim):
+    """(B, L, n_heads*head_dim) -> (B, n_heads, L, head_dim)."""
+    B, L, _ = x.shape
+    return x.reshape(B, L, n_heads, head_dim).transpose(0, 2, 1, 3)
+
+
+def _from_heads(x):
+    """(B, n_heads, L, head_dim) -> (B, L, n_heads*head_dim)."""
+    B, H, L, d = x.shape
+    return x.transpose(0, 2, 1, 3).reshape(B, L, H * d)
+
+
+def joint_attention_kernel(
+    img, txt,
+    # image-stream projections
+    iq_w, iq_b, ik_w, ik_b, iv_w, iv_b, io_w, io_b, iq_g, ik_g,
+    # text-stream projections
+    tq_w, tq_b, tk_w, tk_b, tv_w, tv_b, to_w, to_b, tq_g, tk_g,
+    n_heads, head_dim, eps,
+    img_cos, img_sin, txt_cos, txt_sin,
+    txt_mask_bias=None, local_heads=None, all_reduce_fn=None,
+):
+    """Joint attention over concat([text, image]).
+
+    Args:
+        img: (B, Limg, hidden) modulated image tokens.
+        txt: (B, Ltext, hidden) modulated text tokens.
+        *_w / *_b: projection weights / biases. Under tensor parallelism these
+            are the *sharded* slices: q/k/v are (hidden, local_heads*head_dim),
+            o is (local_heads*head_dim, hidden), and the o/to biases are the full
+            replicated (hidden,).
+        iq_g/ik_g/tq_g/tk_g: QK-RMSNorm gains (head_dim,).
+        img_cos/img_sin: (Limg, head_dim) RoPE tables for image tokens.
+        txt_cos/txt_sin: (Ltext, head_dim) RoPE tables for text tokens.
+        txt_mask_bias: optional (B, 1, 1, Ltext) additive bias for text padding
+            (applied only over the text key positions).
+        local_heads: heads owned by this rank under tensor parallelism.
+        all_reduce_fn: callable summing the output projections across ranks
+            (row-parallel reduction); always applied (TP is required).
+
+    Returns:
+        (img_out, txt_out): (B, Limg, hidden), (B, Ltext, hidden).
+    """
+    Ltext = txt.shape[1]
+    h = n_heads if local_heads is None else local_heads
+
+    # projections (head axis may be sharded -> ``h`` local heads)
+    iq = _to_heads(np.matmul(img, iq_w) + iq_b, h, head_dim)
+    ik = _to_heads(np.matmul(img, ik_w) + ik_b, h, head_dim)
+    iv = _to_heads(np.matmul(img, iv_w) + iv_b, h, head_dim)
+    tq = _to_heads(np.matmul(txt, tq_w) + tq_b, h, head_dim)
+    tk = _to_heads(np.matmul(txt, tk_w) + tk_b, h, head_dim)
+    tv = _to_heads(np.matmul(txt, tv_w) + tv_b, h, head_dim)
+
+    # QK-RMSNorm over head_dim (per-head)
+    iq = rmsnorm_kernel(iq, iq_g, eps=eps)
+    ik = rmsnorm_kernel(ik, ik_g, eps=eps)
+    tq = rmsnorm_kernel(tq, tq_g, eps=eps)
+    tk = rmsnorm_kernel(tk, tk_g, eps=eps)
+
+    # 3D RoPE
+    iq = apply_rotary_emb(iq, img_cos, img_sin)
+    ik = apply_rotary_emb(ik, img_cos, img_sin)
+    tq = apply_rotary_emb(tq, txt_cos, txt_sin)
+    tk = apply_rotary_emb(tk, txt_cos, txt_sin)
+
+    # concat [text, image] along sequence
+    q = np.concatenate([tq, iq], axis=2)
+    k = np.concatenate([tk, ik], axis=2)
+    v = np.concatenate([tv, iv], axis=2)
+
+    # scaled dot-product attention (bidirectional)
+    scores = np.matmul(q, k.transpose(0, 1, 3, 2)).astype(np.float32)
+    scores = scores / np.float32(np.sqrt(head_dim))
+    if txt_mask_bias is not None:
+        # bias applies over the text key positions only; image keys unmasked.
+        # Build the joint-key bias by concatenating the text bias with a zero
+        # image-key block (trace-friendly: no in-place slice assignment).
+        B = scores.shape[0]
+        Limg = scores.shape[-1] - Ltext
+        txt_b = np.broadcast_to(txt_mask_bias.astype(np.float32), (B, 1, 1, Ltext))
+        img_b = np.zeros((B, 1, 1, Limg), dtype=np.float32)
+        bias = np.concatenate([txt_b, img_b], axis=-1)
+        scores = scores + bias
+    weights = softmax_kernel(scores).astype(v.dtype)
+    out = np.matmul(weights, v)  # (B, H, Ljoint, head_dim)
+
+    joint = _from_heads(out)  # (B, Ljoint, local_heads*head_dim)
+    txt_out = joint[:, :Ltext, :]
+    img_out = joint[:, Ltext:, :]
+
+    # row-parallel output projection: matmul the local slice, all-reduce the
+    # partial sums, then add the (replicated) bias exactly once.
+    img_out = np.matmul(img_out, io_w)
+    txt_out = np.matmul(txt_out, to_w)
+    img_out = all_reduce_fn(img_out)
+    txt_out = all_reduce_fn(txt_out)
+    if io_b is not None:
+        img_out = img_out + io_b
+    if to_b is not None:
+        txt_out = txt_out + to_b
+    return img_out, txt_out
diff --git a/examples/models/qwen_image/kernels/embeddings.py b/examples/models/qwen_image/kernels/embeddings.py
new file mode 100644
index 0000000..d5196f9
--- /dev/null
+++ b/examples/models/qwen_image/kernels/embeddings.py
@@ -0,0 +1,84 @@
+"""Input embeddings for the Qwen-Image MMDiT.
+
+Three pieces run before the transformer blocks:
+
+1. ``img_in`` — a Linear projecting the patchified latent (in_channels=64) to
+   ``inner_dim`` (3072). The latent is already patchified by the pipeline, so
+   this is a plain matmul (no strided patch-embed conv).
+2. ``txt_norm`` (RMSNorm over joint_attention_dim) + ``txt_in`` (Linear to
+   inner_dim) — normalize and project the Qwen2.5-VL text embeddings.
+3. ``time_text_embed`` (``QwenTimestepProjEmbeddings``) — sinusoidal timestep
+   embedding (256 ch, flip_sin_to_cos, scale 1000) → SiLU-MLP to ``inner_dim``.
+   This ``temb`` conditions every block's modulation and the final layer.
+
+The sinusoidal frequency table is a comptime numpy constant; only the timestep
+values are runtime tensors.
+"""
+
+import numpy as np
+
+from .rmsnorm import rmsnorm_kernel
+
+
+def _silu(x):
+    xf = x.astype(np.float32)
+    return (xf * (1.0 / (1.0 + np.exp(-xf)))).astype(x.dtype)
+
+
+def img_in_kernel(latent, weight, bias):
+    """Project patchified latent (B, Limg, in_channels) -> (B, Limg, inner_dim)."""
+    return np.matmul(latent, weight) + bias
+
+
+def txt_in_kernel(text, norm_weight, in_weight, in_bias, eps=1e-6):
+    """RMSNorm(text) then Linear -> (B, Ltext, inner_dim).
+
+    Args:
+        text: (B, Ltext, joint_attention_dim).
+        norm_weight: (joint_attention_dim,) RMSNorm gain.
+        in_weight: (joint_attention_dim, inner_dim), in_bias: (inner_dim,).
+    """
+    x = rmsnorm_kernel(text, norm_weight, eps=eps)
+    return np.matmul(x, in_weight) + in_bias
+
+
+def timestep_embedding(timesteps, dim=256, max_period=10000, scale=1000.0,
+                       flip_sin_to_cos=True, downscale_freq_shift=0.0):
+    """Sinusoidal timestep embedding matching diffusers ``get_timestep_embedding``.
+
+    Qwen-Image uses ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0,
+    scale=1000)``: the frequency argument is scaled by 1000, embeddings are
+    ``[sin, cos]`` then flipped to ``[cos, sin]``.
+
+    ``timesteps`` is a runtime (B,) tensor; the frequency table is comptime.
+    Returns (B, dim).
+    """
+    half = dim // 2
+    exponent = -np.log(max_period) * np.arange(half, dtype=np.float32)
+    exponent = exponent / (half - downscale_freq_shift)
+    freqs = np.exp(exponent)  # (half,)
+
+    args = np.expand_dims(timesteps.astype(np.float32), -1) * np.expand_dims(freqs, 0)
+    args = scale * args
+    emb = np.concatenate([np.sin(args), np.cos(args)], axis=-1)  # (B, dim)
+    if flip_sin_to_cos:
+        emb = np.concatenate([emb[:, half:], emb[:, :half]], axis=-1)
+    return emb
+
+
+def time_text_embed_kernel(timesteps, proj_weight, proj_bias, emb_weight, emb_bias,
+                           dtype):
+    """QwenTimestepProjEmbeddings: sinusoid -> Linear -> SiLU -> Linear.
+
+    Args:
+        timesteps: (B,) diffusion timestep.
+        proj_weight: (256, inner_dim), proj_bias: (inner_dim,)  [linear_1]
+        emb_weight: (inner_dim, inner_dim), emb_bias: (inner_dim,)  [linear_2]
+    Returns:
+        temb: (B, inner_dim) conditioning embedding.
+    """
+    sin_emb = timestep_embedding(timesteps).astype(dtype)  # (B, 256)
+    h = np.matmul(sin_emb, proj_weight) + proj_bias
+    h = _silu(h)
+    temb = np.matmul(h, emb_weight) + emb_bias
+    return temb
diff --git a/examples/models/qwen_image/kernels/feedforward.py b/examples/models/qwen_image/kernels/feedforward.py
new file mode 100644
index 0000000..52c653d
--- /dev/null
+++ b/examples/models/qwen_image/kernels/feedforward.py
@@ -0,0 +1,37 @@
+"""Per-stream feed-forward for the Qwen-Image MMDiT.
+
+Matches diffusers ``FeedForward(activation_fn="gelu-approximate")``: a GELU
+(tanh approximation) MLP, no gating. Both ``img_mlp`` and ``txt_mlp`` use it with
+``dim_out == dim`` and inner dim ``4 * dim``.
+"""
+
+import numpy as np
+
+
+def _gelu_tanh(x):
+    xf = x.astype(np.float32)
+    inner = np.sqrt(2.0 / np.pi) * (xf + 0.044715 * xf * xf * xf)
+    out = 0.5 * xf * (1.0 + np.tanh(inner))
+    return out.astype(x.dtype)
+
+
+def feedforward_kernel(x, up_weight, up_bias, down_weight, down_bias,
+                       all_reduce_fn):
+    """Args:
+        x: (B, L, dim)
+        up_weight: (dim, inner), up_bias: (inner,)     [column-parallel under TP]
+        down_weight: (inner, dim), down_bias: (dim,)   [row-parallel under TP]
+        all_reduce_fn: callable summing the row-parallel output across ranks;
+            always applied (TP is required). ``inner`` is the local (sharded)
+            intermediate size and ``down_bias`` is the full replicated (dim,),
+            added once after the reduction.
+    """
+    h = np.matmul(x, up_weight)
+    if up_bias is not None:
+        h = h + up_bias
+    h = _gelu_tanh(h)
+    out = np.matmul(h, down_weight)
+    out = all_reduce_fn(out)
+    if down_bias is not None:
+        out = out + down_bias
+    return out
diff --git a/examples/models/qwen_image/kernels/final_layer.py b/examples/models/qwen_image/kernels/final_layer.py
new file mode 100644
index 0000000..0d4e19b
--- /dev/null
+++ b/examples/models/qwen_image/kernels/final_layer.py
@@ -0,0 +1,50 @@
+"""Final output head for the Qwen-Image MMDiT.
+
+``norm_out`` (``AdaLayerNormContinuous``): non-affine LayerNorm modulated by a
+(scale, shift) derived from ``temb`` via SiLU-Linear, applied to the image
+stream only. ``proj_out``: Linear(inner_dim -> patch_size**2 * out_channels).
+
+diffusers' ``QwenImageTransformer2DModel.forward`` returns the proj_out result
+directly (shape (B, Limg, patch**2 * out_ch)); the pipeline unpatchifies it into
+(B, out_ch, H, W). ``unpatchify`` here mirrors that pipeline step for M4.
+"""
+
+import numpy as np
+
+from .layernorm import layernorm_kernel
+
+
+def _silu(x):
+    xf = x.astype(np.float32)
+    return (xf * (1.0 / (1.0 + np.exp(-xf)))).astype(x.dtype)
+
+
+def final_layer(img, temb, norm_lin_w, norm_lin_b, proj_w, proj_b, eps=1e-6):
+    """AdaLayerNormContinuous(norm_out) + proj_out on the image stream.
+
+    Args:
+        img: (B, Limg, inner_dim) image tokens from the last block.
+        temb: (B, inner_dim) timestep conditioning.
+        norm_lin_w: (inner_dim, 2*inner_dim), norm_lin_b: (2*inner_dim,).
+        proj_w: (inner_dim, patch**2 * out_ch), proj_b: (patch**2 * out_ch,).
+    Returns:
+        (B, Limg, patch**2 * out_ch).
+    """
+    emb = np.matmul(_silu(temb), norm_lin_w) + norm_lin_b  # (B, 2*inner_dim)
+    scale, shift = np.split(emb, 2, axis=-1)  # each (B, inner_dim)
+    x = layernorm_kernel(img, eps=eps)
+    x = x * (1 + np.expand_dims(scale, 1)) + np.expand_dims(shift, 1)
+    return np.matmul(x, proj_w) + proj_b
+
+
+def unpatchify(x, patch_size, out_channels, gh, gw):
+    """(B, gh*gw, patch**2 * out_ch) -> (B, out_ch, gh*patch, gw*patch).
+
+    Mirrors the QwenImagePipeline unpatchify: tokens are row-major over the
+    (gh, gw) patch grid, each carrying a (out_ch, patch, patch) block.
+    """
+    B = x.shape[0]
+    p = patch_size
+    x = x.reshape(B, gh, gw, out_channels, p, p)
+    x = x.transpose(0, 3, 1, 4, 2, 5)  # (B, out_ch, gh, p, gw, p)
+    return x.reshape(B, out_channels, gh * p, gw * p)
diff --git a/examples/models/qwen_image/kernels/layernorm.py b/examples/models/qwen_image/kernels/layernorm.py
new file mode 100644
index 0000000..9517265
--- /dev/null
+++ b/examples/models/qwen_image/kernels/layernorm.py
@@ -0,0 +1,19 @@
+"""Non-affine LayerNorm for the Qwen-Image MMDiT blocks.
+
+Both streams use ``LayerNorm(elementwise_affine=False)`` before adaLN-style
+modulation, so there are no learnable gain/bias. Computed in fp32 to limit
+numerical error, then cast back to the input dtype.
+"""
+
+import numpy as np
+
+
+def layernorm_kernel(x, eps: float = 1e-6, compute_dtype=np.float32):
+    original_dtype = x.dtype
+    x = x.astype(compute_dtype)
+
+    mean = np.mean(x, axis=-1, keepdims=True)
+    var = np.mean(np.square(x - mean), axis=-1, keepdims=True)
+    z = (x - mean) / np.sqrt(var + eps)
+
+    return z.astype(original_dtype)
diff --git a/examples/models/qwen_image/kernels/mmdit_block.py b/examples/models/qwen_image/kernels/mmdit_block.py
new file mode 100644
index 0000000..362a90a
--- /dev/null
+++ b/examples/models/qwen_image/kernels/mmdit_block.py
@@ -0,0 +1,114 @@
+"""A single Qwen-Image MMDiT dual-stream block.
+
+Mirrors diffusers ``QwenImageTransformerBlock`` (still-image path, no
+``zero_cond_t`` / ``modulate_index``). Both the image stream (``img``) and the
+text stream (``txt``) carry their own norms, modulation, and MLP; they meet only
+in the joint attention.
+
+Order:
+    1. modulation: img_mod/txt_mod = Linear(SiLU(temb)) -> 6*dim, split into
+       (mod1, mod2), each (shift, scale, gate).
+    2. norm1 + modulate both streams -> joint attention -> gated residual.
+    3. norm2 + modulate both streams -> per-stream MLP -> gated residual.
+
+norm1/norm2 are non-affine LayerNorm; all conditioning enters through the
+modulation. Weights arrive as a flat dict ``w`` keyed by the canonical names in
+``weight_layout.py``.
+"""
+
+import numpy as np
+
+from .attention import joint_attention_kernel
+from .feedforward import feedforward_kernel
+from .layernorm import layernorm_kernel
+
+
+def _silu(x):
+    xf = x.astype(np.float32)
+    return (xf * (1.0 / (1.0 + np.exp(-xf)))).astype(x.dtype)
+
+
+def _modulation(temb, mod_w, mod_b, hidden):
+    """SiLU -> Linear(dim, 6*dim), returning the two (shift, scale, gate) triples.
+
+    Args:
+        temb: (B, hidden) timestep conditioning.
+        mod_w: (hidden, 6*hidden), mod_b: (6*hidden,).
+    Returns:
+        (mod1, mod2), each a tuple (shift, scale, gate) of (B, 1, hidden).
+    """
+    params = np.matmul(_silu(temb), mod_w) + mod_b  # (B, 6*hidden)
+    mod1, mod2 = np.split(params, 2, axis=-1)  # each (B, 3*hidden)
+
+    def _split3(m):
+        shift, scale, gate = np.split(m, 3, axis=-1)  # each (B, hidden)
+        return (
+            np.expand_dims(shift, 1),
+            np.expand_dims(scale, 1),
+            np.expand_dims(gate, 1),
+        )
+
+    return _split3(mod1), _split3(mod2)
+
+
+def _apply_mod(x, shift, scale):
+    return x * (1 + scale) + shift
+
+
+def mmdit_block(img, txt, w, temb, n_heads, head_dim, hidden, eps,
+                img_cos, img_sin, txt_cos, txt_sin, txt_mask_bias=None,
+                local_heads=None, all_reduce_fn=None):
+    """Run one MMDiT block.
+
+    Args:
+        img: (B, Limg, hidden) image tokens.
+        txt: (B, Ltext, hidden) text tokens.
+        w: dict of this block's weights (canonical keys). Under tensor
+            parallelism the attention/MLP weights are the sharded slices; the
+            replicated modulation/norms keep full hidden.
+        temb: (B, hidden) timestep conditioning.
+        img_cos/img_sin/txt_cos/txt_sin: RoPE tables (see rope3d).
+        txt_mask_bias: optional (B, 1, 1, Ltext) additive text-padding bias.
+        local_heads: attention heads owned by this rank (None -> n_heads / TP=1).
+        all_reduce_fn: optional collective summing the row-parallel outputs
+            (attention out-proj and MLP down-proj) across ranks.
+    Returns:
+        (img, txt) updated streams (full hidden; replicated across ranks).
+    """
+    (img_shift1, img_scale1, img_gate1), (img_shift2, img_scale2, img_gate2) = _modulation(
+        temb, w["img_mod_w"], w["img_mod_b"], hidden
+    )
+    (txt_shift1, txt_scale1, txt_gate1), (txt_shift2, txt_scale2, txt_gate2) = _modulation(
+        temb, w["txt_mod_w"], w["txt_mod_b"], hidden
+    )
+
+    # 1. norm1 + modulate -> joint attention -> gated residual
+    img_mod = _apply_mod(layernorm_kernel(img, eps=eps), img_shift1, img_scale1)
+    txt_mod = _apply_mod(layernorm_kernel(txt, eps=eps), txt_shift1, txt_scale1)
+
+    img_attn, txt_attn = joint_attention_kernel(
+        img_mod, txt_mod,
+        w["iq_w"], w["iq_b"], w["ik_w"], w["ik_b"], w["iv_w"], w["iv_b"],
+        w["io_w"], w["io_b"], w["iq_g"], w["ik_g"],
+        w["tq_w"], w["tq_b"], w["tk_w"], w["tk_b"], w["tv_w"], w["tv_b"],
+        w["to_w"], w["to_b"], w["tq_g"], w["tk_g"],
+        n_heads, head_dim, eps,
+        img_cos, img_sin, txt_cos, txt_sin,
+        txt_mask_bias=txt_mask_bias,
+        local_heads=local_heads, all_reduce_fn=all_reduce_fn,
+    )
+    img = img + img_gate1 * img_attn
+    txt = txt + txt_gate1 * txt_attn
+
+    # 2. norm2 + modulate -> per-stream MLP -> gated residual
+    img_mod2 = _apply_mod(layernorm_kernel(img, eps=eps), img_shift2, img_scale2)
+    img_ff = feedforward_kernel(img_mod2, w["iff0_w"], w["iff0_b"], w["iff2_w"], w["iff2_b"],
+                                all_reduce_fn=all_reduce_fn)
+    img = img + img_gate2 * img_ff
+
+    txt_mod2 = _apply_mod(layernorm_kernel(txt, eps=eps), txt_shift2, txt_scale2)
+    txt_ff = feedforward_kernel(txt_mod2, w["tff0_w"], w["tff0_b"], w["tff2_w"], w["tff2_b"],
+                                all_reduce_fn=all_reduce_fn)
+    txt = txt + txt_gate2 * txt_ff
+
+    return img, txt
diff --git a/examples/models/qwen_image/kernels/rmsnorm.py b/examples/models/qwen_image/kernels/rmsnorm.py
new file mode 100644
index 0000000..cce8231
--- /dev/null
+++ b/examples/models/qwen_image/kernels/rmsnorm.py
@@ -0,0 +1,24 @@
+"""RMSNorm for Qwen-Image.
+
+Used in two places: the text-stream input norm (``txt_norm``, over
+``joint_attention_dim``) and the per-head QK-RMSNorm inside joint attention
+(``norm_q/norm_k`` and ``norm_added_q/norm_added_k``, over ``head_dim``).
+
+Matches the qwen3 example's ``rmsnorm_kernel``; computed in fp32 to limit
+numerical error, then cast back to the input dtype.
+"""
+
+import numpy as np
+
+
+def rmsnorm_kernel(x, weight, eps: float, compute_dtype=np.float32):
+    original_dtype = x.dtype
+    x = x.astype(compute_dtype)
+    weight = weight.astype(compute_dtype)
+
+    z = np.mean(np.square(x), axis=-1, keepdims=True)
+    z = (z + eps).astype(x.dtype)
+    z = x / np.sqrt(z)
+
+    res = z * weight
+    return res.astype(original_dtype)
diff --git a/examples/models/qwen_image/kernels/rope3d.py b/examples/models/qwen_image/kernels/rope3d.py
new file mode 100644
index 0000000..4fe9be6
--- /dev/null
+++ b/examples/models/qwen_image/kernels/rope3d.py
@@ -0,0 +1,148 @@
+"""3D RoPE for the Qwen-Image MMDiT (``QwenEmbedRope``).
+
+Qwen-Image applies rotary position embedding over a 3D (frame, height, width)
+grid for the image tokens and a 1D range for the text tokens, then rotates the
+attention queries/keys of *both* streams before the joint attention.
+
+Two things differ from the qwen3 example's RoPE:
+
+* **Interleaved complex convention.** diffusers builds per-position unit-magnitude
+  complex frequencies and multiplies ``view_as_complex(x reshaped to (...,-1,2))``
+  by them (``apply_rotary_emb_qwen`` with ``use_real=False``). So dimension pair
+  ``(2i, 2i+1)`` of the head is one complex number rotated by angle
+  ``theta_i``, rather than qwen3's ``(i, i+half)`` half-split.
+* **3D axes with ``scale_rope`` centering.** ``head_dim`` is partitioned into
+  ``axes_dims_rope`` (frame, H, W); the H/W position indices are centered around
+  zero (negative half then positive half), the frame axis is not.
+
+Positions are comptime constants (grid size is known at trace time), so the
+angle table is built in numpy and baked into the graph; only the runtime rotate
+(``apply_rotary_emb``) touches tensors. We keep everything real-valued
+(cos/sin), which matches ``x_out = x*cos + rotate(x)*sin`` — the real form of the
+complex multiply — and needs no complex dtype on device.
+"""
+
+import numpy as np
+
+
+def _axis_freqs(index, dim, theta):
+    """Per-position angles for one RoPE axis.
+
+    Args:
+        index: (P,) integer positions.
+        dim: this axis' slice of head_dim (must be even); yields dim//2 angles.
+        theta: RoPE base.
+    Returns:
+        (P, dim//2) angles = outer(index, 1/theta**(arange(0,dim,2)/dim)).
+    """
+    assert dim % 2 == 0
+    inv_freq = 1.0 / (theta ** (np.arange(0, dim, 2, dtype=np.float64) / dim))
+    return np.outer(index.astype(np.float64), inv_freq)  # (P, dim//2)
+
+
+def compute_rope_freqs(frame, height, width, txt_len, axes_dims, theta=10000.0,
+                       scale_rope=True, max_pos=4096):
+    """Build the (image, text) RoPE angle tables (comptime constant).
+
+    Mirrors ``QwenEmbedRope._compute_video_freqs`` + the text slice in
+    ``forward``. Returns per-position angles (not yet cos/sin) so callers can
+    materialise cos/sin at the desired dtype.
+
+    Args:
+        frame, height, width: image grid (frame is 1 for a still image). These
+            are the RoPE grid extents; the pipeline passes H_latent//patch,
+            W_latent//patch here.
+        txt_len: number of text tokens.
+        axes_dims: (d_frame, d_h, d_w) partition of head_dim, each even.
+        theta: RoPE base.
+        scale_rope: center the H/W position indices around zero (Qwen default).
+        max_pos: size of the cached position range (diffusers uses 4096).
+
+    Returns:
+        vid_angles: (frame*height*width, head_dim//2)
+        txt_angles: (txt_len, head_dim//2)
+    """
+    d_f, d_h, d_w = axes_dims
+
+    # Positive index range [0..max_pos) and the negative (centering) range,
+    # matching diffusers: neg index = arange(max_pos).flip(0) * -1 - 1.
+    pos_index = np.arange(max_pos)
+    neg_index = np.arange(max_pos)[::-1] * -1 - 1  # [..., -3, -2, -1]
+
+    pos_f = _axis_freqs(pos_index, d_f, theta)  # (max_pos, d_f//2)
+    pos_h = _axis_freqs(pos_index, d_h, theta)
+    pos_w = _axis_freqs(pos_index, d_w, theta)
+    neg_h = _axis_freqs(neg_index, d_h, theta)
+    neg_w = _axis_freqs(neg_index, d_w, theta)
+
+    # frame axis: positions [0..frame), broadcast over (frame, H, W)
+    fr = pos_f[:frame].reshape(frame, 1, 1, -1)
+    fr = np.broadcast_to(fr, (frame, height, width, d_f // 2))
+
+    if scale_rope:
+        h_ang = np.concatenate([neg_h[-(height - height // 2):], pos_h[:height // 2]], axis=0)
+        w_ang = np.concatenate([neg_w[-(width - width // 2):], pos_w[:width // 2]], axis=0)
+    else:
+        h_ang = pos_h[:height]
+        w_ang = pos_w[:width]
+    hh = np.broadcast_to(h_ang.reshape(1, height, 1, -1), (frame, height, width, d_h // 2))
+    ww = np.broadcast_to(w_ang.reshape(1, 1, width, -1), (frame, height, width, d_w // 2))
+
+    vid_angles = np.concatenate([fr, hh, ww], axis=-1).reshape(frame * height * width, -1)
+
+    # text tokens occupy positions after the image grid center
+    if scale_rope:
+        max_vid_index = max(height // 2, width // 2)
+    else:
+        max_vid_index = max(height, width)
+    txt_f = pos_f[max_vid_index:max_vid_index + txt_len]
+    txt_h = pos_h[max_vid_index:max_vid_index + txt_len]
+    txt_w = pos_w[max_vid_index:max_vid_index + txt_len]
+    txt_angles = np.concatenate([txt_f, txt_h, txt_w], axis=-1)  # (txt_len, head_dim//2)
+
+    return vid_angles, txt_angles
+
+
+def cos_sin_from_angles(angles, dtype=np.float32):
+    """Turn (S, head_dim//2) angles into interleaved (S, head_dim) cos/sin.
+
+    Each angle governs a dimension *pair*; interleaving to (cos0,cos0,cos1,...)
+    lets ``apply_rotary_emb`` operate on the (2i, 2i+1) real/imag pairs without
+    reshaping. Returned as comptime constants.
+    """
+    cos = np.cos(angles)
+    sin = np.sin(angles)
+    cos = np.repeat(cos, 2, axis=-1).astype(dtype)  # (S, head_dim)
+    sin = np.repeat(sin, 2, axis=-1).astype(dtype)
+    return cos, sin
+
+
+def apply_rotary_emb(x, cos, sin):
+    """Rotate q/k with interleaved RoPE (real form of the complex multiply).
+
+    Args:
+        x: (B, H, S, head_dim) queries or keys.
+        cos, sin: (S, head_dim) interleaved tables from ``cos_sin_from_angles``.
+
+    For complex z = (x_even + i*x_odd) rotated by (cos + i*sin):
+        out_even = x_even*cos - x_odd*sin
+        out_odd  = x_even*sin + x_odd*cos
+    which equals ``x*cos + rotate_half_interleaved(x)*sin`` where
+    ``rotate_half_interleaved([a,b]) = [-b, a]`` per pair.
+    """
+    orig_dtype = x.dtype
+    xf = x.astype(np.float32)
+    cos = cos.astype(np.float32).reshape((1, 1) + cos.shape)
+    sin = sin.astype(np.float32).reshape((1, 1) + sin.shape)
+
+    x_even = xf[..., 0::2]
+    x_odd = xf[..., 1::2]
+    # interleave [-x_odd, x_even] back to full width. Build it with a stack +
+    # reshape rather than assigning into strided slices of an empty array: the
+    # strided-slice assignment lowers to a stride-2 scatter that dominated the
+    # whole block (~75 ms/block, ~95% of the cte-backend cost); stack+reshape
+    # lowers to a cheap concat + contiguous reshape.
+    rot = np.stack([-x_odd, x_even], axis=-1).reshape(xf.shape)
+
+    out = xf * cos + rot * sin
+    return out.astype(orig_dtype)
diff --git a/examples/models/qwen_image/kernels/softmax.py b/examples/models/qwen_image/kernels/softmax.py
new file mode 100644
index 0000000..2f9a948
--- /dev/null
+++ b/examples/models/qwen_image/kernels/softmax.py
@@ -0,0 +1,6 @@
+import numpy as np
+
+
+def softmax_kernel(x):
+    exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
+    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
diff --git a/examples/models/qwen_image/kernels/text_encoder.py b/examples/models/qwen_image/kernels/text_encoder.py
new file mode 100644
index 0000000..58036ee
--- /dev/null
+++ b/examples/models/qwen_image/kernels/text_encoder.py
@@ -0,0 +1,139 @@
+"""Qwen2.5 text-encoder forward for Qwen-Image, on device.
+
+Qwen-Image conditions on the last hidden state of a Qwen2.5-VL model run
+**text-only** (no image tokens), so the vision tower is unused and the encoder
+is a standard Qwen2.5 decoder LM: 28 layers, hidden 3584, 28 query heads / 4 KV
+heads (GQA), SwiGLU MLP (intermediate 18944), RMSNorm, RoPE (theta 1e6),
+head_dim 128. Differences vs the qwen3 example's kernels:
+
+* q/k/v are **separate** projections **with bias** (not a fused qkv, no QK-norm);
+* we run a single **prefill** pass over the whole prompt and return the *last
+  hidden state* (before the LM head) — this is an encoder, not a generator, so
+  there is no KV cache / decode loop and no sampling.
+
+Attention is causal (it is a decoder LM; diffusers runs it with the default
+causal mask). Weights arrive as a flat ``**weights`` dict keyed by
+``text_weight_layout.py``. RoPE cos/sin and the causal mask are comptime numpy
+constants baked into the graph (same pattern as the MMDiT / qwen3 kernels).
+
+TP: q/k/v/o and the MLP shard exactly like the MMDiT denoiser (heads /
+intermediate), with an all-reduce after o_proj and down_proj. ``local_heads`` /
+``local_kv_heads`` / ``all_reduce_fn`` come from the config.
+"""
+
+import numpy as np
+
+from .rmsnorm import rmsnorm_kernel
+from .softmax import softmax_kernel
+
+
+def _rope_tables(seq_len, head_dim, theta, dtype):
+    """Comptime cos/sin for standard (half-split) RoPE. (seq_len, head_dim)."""
+    inv_freq = 1.0 / (theta ** (np.arange(0, head_dim, 2, dtype=np.float64) / head_dim))
+    t = np.arange(seq_len, dtype=np.float64)
+    freqs = np.outer(t, inv_freq)  # (seq, head_dim/2)
+    emb = np.concatenate([freqs, freqs], axis=-1)  # (seq, head_dim)
+    return np.cos(emb).astype(dtype), np.sin(emb).astype(dtype)
+
+
+def _apply_rope(x, cos, sin):
+    """x: (B, H, S, d); cos/sin: (S, d). Half-split rotation (Qwen2 convention)."""
+    d = x.shape[-1]
+    half = d // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    rot = np.concatenate([-x2, x1], axis=-1)
+    cos = cos.reshape((1, 1) + cos.shape)
+    sin = sin.reshape((1, 1) + sin.shape)
+    return x * cos + rot * sin
+
+
+def _swiglu(x, gate_w, up_w, down_w, all_reduce_fn):
+    """SwiGLU MLP: down(silu(gate(x)) * up(x)). gate/up column-parallel, down row."""
+    g = np.matmul(x, gate_w)
+    u = np.matmul(x, up_w)
+    gf = g.astype(np.float32)
+    act = (gf * (1.0 / (1.0 + np.exp(-gf)))).astype(g.dtype)
+    h = act * u
+    out = np.matmul(h, down_w)
+    return all_reduce_fn(out)
+
+
+def _layer(x, w, i, cos, sin, causal_bias, n_heads, n_kv_heads, head_dim, eps,
+           local_heads, local_kv_heads, all_reduce_fn):
+    """One Qwen2.5 decoder layer (pre-norm attention + pre-norm SwiGLU)."""
+    B, S, _ = x.shape
+    hq = local_heads if local_heads is not None else n_heads
+    hkv = local_kv_heads if local_kv_heads is not None else n_kv_heads
+
+    # ── attention ──
+    h = rmsnorm_kernel(x, w[f"l{i}.attn_norm"], eps=eps)
+    q = np.matmul(h, w[f"l{i}.q_w"]) + w[f"l{i}.q_b"]
+    k = np.matmul(h, w[f"l{i}.k_w"]) + w[f"l{i}.k_b"]
+    v = np.matmul(h, w[f"l{i}.v_w"]) + w[f"l{i}.v_b"]
+
+    q = q.reshape(B, S, hq, head_dim).transpose(0, 2, 1, 3)
+    k = k.reshape(B, S, hkv, head_dim).transpose(0, 2, 1, 3)
+    v = v.reshape(B, S, hkv, head_dim).transpose(0, 2, 1, 3)
+
+    q = _apply_rope(q, cos, sin)
+    k = _apply_rope(k, cos, sin)
+
+    # GQA: repeat KV heads
+    n_rep = hq // hkv
+    if n_rep > 1:
+        k = np.repeat(k, n_rep, axis=1)
+        v = np.repeat(v, n_rep, axis=1)
+
+    scores = np.matmul(q, k.transpose(0, 1, 3, 2)).astype(np.float32)
+    scores = scores / np.float32(np.sqrt(head_dim))
+    scores = scores + causal_bias  # (1,1,S,S) comptime
+    attn = softmax_kernel(scores).astype(v.dtype)
+    o = np.matmul(attn, v)  # (B, hq, S, d)
+    o = o.transpose(0, 2, 1, 3).reshape(B, S, hq * head_dim)
+
+    o = np.matmul(o, w[f"l{i}.o_w"])
+    o = all_reduce_fn(o)
+    x = x + o
+
+    # ── SwiGLU MLP ──
+    h = rmsnorm_kernel(x, w[f"l{i}.mlp_norm"], eps=eps)
+    ff = _swiglu(h, w[f"l{i}.gate_w"], w[f"l{i}.up_w"], w[f"l{i}.down_w"],
+                 all_reduce_fn=all_reduce_fn)
+    x = x + ff
+    return x
+
+
+def text_encoder_forward(hidden, configs, **weights):
+    """Prefill forward returning the last hidden state (B, S, hidden).
+
+    Args:
+        hidden: (B, S, hidden) token embeddings (host does the embedding lookup
+            from the input ids; the embedding table is huge and the lookup is
+            data-dependent, so it stays on host).
+        configs: text-encoder config (see TextEncoderConfig).
+    Returns:
+        (B, S, hidden) last hidden state (post final norm), matching
+        ``hidden_states[-1]`` of the diffusers encoder.
+    """
+    n_heads = configs.num_heads
+    n_kv_heads = configs.num_kv_heads
+    head_dim = configs.head_dim
+    eps = configs.rms_norm_eps
+    tp = getattr(configs, "tp_size", 1) or 1
+    local_heads = n_heads // tp if tp > 1 else None
+    local_kv_heads = max(1, n_kv_heads // tp) if tp > 1 else None
+    all_reduce_fn = configs.all_reduce_fn
+
+    B, S, _ = hidden.shape
+    cos, sin = _rope_tables(S, head_dim, configs.rope_theta, configs.dtype)
+    # causal mask (comptime): upper triangle -> large negative
+    causal = np.triu(np.ones((S, S), dtype=np.float32) * -1e9, k=1)
+    causal_bias = causal.reshape(1, 1, S, S)
+
+    x = hidden
+    for i in range(configs.num_layers):
+        x = _layer(x, weights, i, cos, sin, causal_bias, n_heads, n_kv_heads,
+                   head_dim, eps, local_heads, local_kv_heads, all_reduce_fn)
+
+    x = rmsnorm_kernel(x, weights["final_norm"], eps=eps)
+    return x
diff --git a/examples/models/qwen_image/kernels/text_weight_layout.py b/examples/models/qwen_image/kernels/text_weight_layout.py
new file mode 100644
index 0000000..43c8cdd
--- /dev/null
+++ b/examples/models/qwen_image/kernels/text_weight_layout.py
@@ -0,0 +1,30 @@
+"""Flat weight-key scheme for the Qwen2.5 text encoder.
+
+Per-layer keys are prefixed ``l{i}.``; shared (embedding-side) weights use bare
+names. As elsewhere, ``nn.Linear`` weights are transposed (out,in)->(in,out) for
+the ``x @ W`` convention; RMSNorm gains and biases pass through.
+
+The token-embedding table and LM head are **not** here — the host does the
+embedding lookup (huge table, data-dependent gather) and there is no LM head in
+encoder mode.
+"""
+
+# per-layer short keys
+LAYER_KEYS = [
+    "attn_norm", "mlp_norm",
+    "q_w", "q_b", "k_w", "k_b", "v_w", "v_b", "o_w",
+    "gate_w", "up_w", "down_w",
+]
+
+SHARED_KEYS = ["final_norm"]
+
+
+def layer_key(i, short):
+    return f"l{i}.{short}"
+
+
+def present_keys(flat, num_layers):
+    keys = [k for k in SHARED_KEYS if k in flat]
+    for i in range(num_layers):
+        keys += [layer_key(i, s) for s in LAYER_KEYS if layer_key(i, s) in flat]
+    return keys
diff --git a/examples/models/qwen_image/kernels/tp.py b/examples/models/qwen_image/kernels/tp.py
new file mode 100644
index 0000000..53f9a19
--- /dev/null
+++ b/examples/models/qwen_image/kernels/tp.py
@@ -0,0 +1,84 @@
+"""Tensor-parallel helpers for the Qwen-Image MMDiT.
+
+The 20B model (~40 GB bf16) does not fit on one trn2 core (24 GB), so the
+denoiser is sharded Megatron-style across ``tp_size`` cores:
+
+* **Attention** shards by heads: each core owns ``num_heads // tp_size`` heads
+  and holds the corresponding slices of q/k/v (column-parallel, output dim) and
+  the output projection (row-parallel, input dim). Heads are independent over the
+  full sequence, so joint attention is unaffected; each core computes the partial
+  contribution of its heads and the output projections are summed with an
+  **all-reduce**.
+* **MLP** shards by intermediate: ff0 column-parallel (output dim), ff2
+  row-parallel (input dim), summed with an all-reduce after ff2.
+* Modulation, LayerNorms, input/output projections, and residuals stay
+  **replicated** — every core holds full-hidden activations after each
+  all-reduce, so the elementwise/residual math is identical across ranks.
+
+``all_reduce_fn`` abstracts the collective: the kernels always apply it after
+each row-parallel projection. The device driver passes a real
+``nkipy.distributed.collectives.all_reduce`` wrapper (tensor parallelism is
+required — the 20B model does not fit on one core).
+"""
+
+import numpy as np
+
+
+def make_all_reduce(tp_size):
+    """Return an all-reduce callable summing across the full replica group.
+
+    Requires ``tp_size > 1`` (tensor parallelism is mandatory). Imported lazily
+    so importing this module doesn't require the distributed runtime.
+    """
+    if tp_size is None or tp_size <= 1:
+        raise ValueError("tensor parallelism is required (tp_size > 1)")
+
+    import nkipy.distributed.collectives as cc
+    import torch.distributed as dist
+
+    def _all_reduce(x):
+        return cc.all_reduce(
+            x, replica_groups=[list(range(dist.get_world_size()))], reduce_op=np.add
+        )
+
+    return _all_reduce
+
+
+def shard_heads(weight, bias, rank, tp_size, n_heads, head_dim, axis):
+    """Slice a head-partitioned projection weight for ``rank``.
+
+    ``axis`` is the head-carrying axis: the output dim for q/k/v (column-parallel,
+    weight (hidden_in, hidden_out)) or the input dim for the output projection
+    (row-parallel, weight (hidden_in, hidden_out)). Biases are sharded only for
+    column-parallel projections (axis == 1); row-parallel biases are added once
+    on the full output and stay replicated.
+    """
+    local_heads = n_heads // tp_size
+    start = rank * local_heads * head_dim
+    end = start + local_heads * head_dim
+
+    if axis == 1:  # column-parallel: shard output dim (and bias)
+        w = weight[:, start:end]
+        b = None if bias is None else bias[start:end]
+    else:  # row-parallel: shard input dim; bias stays full (replicated)
+        w = weight[start:end, :]
+        b = bias
+    return np.ascontiguousarray(w), (None if b is None else np.ascontiguousarray(b))
+
+
+def shard_intermediate(weight, bias, rank, tp_size, axis):
+    """Slice an MLP weight along the intermediate dim for ``rank``.
+
+    ``axis == 1`` (ff0, column-parallel): shard output/intermediate dim + bias.
+    ``axis == 0`` (ff2, row-parallel): shard input/intermediate dim; bias full.
+    """
+    inter = weight.shape[axis]
+    local = inter // tp_size
+    start, end = rank * local, (rank + 1) * local
+    if axis == 1:
+        w = weight[:, start:end]
+        b = None if bias is None else bias[start:end]
+    else:
+        w = weight[start:end, :]
+        b = bias
+    return np.ascontiguousarray(w), (None if b is None else np.ascontiguousarray(b))
diff --git a/examples/models/qwen_image/kernels/transformer.py b/examples/models/qwen_image/kernels/transformer.py
new file mode 100644
index 0000000..3fc8e73
--- /dev/null
+++ b/examples/models/qwen_image/kernels/transformer.py
@@ -0,0 +1,166 @@
+"""Full Qwen-Image MMDiT forward pass.
+
+Given a patchified latent (B, Limg, in_channels), text embeddings (B, Ltext,
+joint_attention_dim), a (B,) timestep, and the image grid ``img_shape`` =
+(frame, gh, gw), produce ``proj_out`` (B, Limg, patch**2 * out_channels) —
+matching diffusers ``QwenImageTransformer2DModel.forward`` (which returns before
+unpatchify; the pipeline unpatchifies).
+
+All weights arrive as flat ``**weights`` kwargs (the tracer turns each top-level
+ndarray into an HLO parameter but does not recurse into dicts); ``regroup_weights``
+rebuilds the shared/per-block structure. The RoPE tables are comptime constants
+built from the (fixed) image grid + text length.
+"""
+
+import numpy as np
+from config import Config
+
+from .embeddings import img_in_kernel, time_text_embed_kernel, txt_in_kernel
+from .final_layer import final_layer, unpatchify
+from .mmdit_block import mmdit_block
+from .rope3d import compute_rope_freqs, cos_sin_from_angles
+from .weight_layout import regroup_weights
+
+
+def _core(latent, text, timestep, img_shape, shared, blocks, configs,
+          txt_mask_bias=None):
+    """Shared MMDiT body. Returns proj_out (B, Limg, patch**2 * out_channels).
+
+    Tensor parallelism (``configs.tp_size`` > 1) shards the attention heads and
+    MLP intermediate across ranks; ``configs.all_reduce_fn`` sums the
+    row-parallel outputs. The blocks receive already-sharded weights (from
+    ``tp.shard_*`` in the driver) and ``local_heads`` = n_heads // tp_size.
+    """
+    hidden = configs.hidden_size
+    n_heads = configs.num_heads
+    head_dim = configs.head_dim
+    eps = configs.norm_eps
+    frame, gh, gw = img_shape
+
+    tp_size = getattr(configs, "tp_size", 1) or 1
+    local_heads = n_heads // tp_size if tp_size > 1 else None
+    all_reduce_fn = configs.all_reduce_fn
+    Ltext = text.shape[1]
+
+    # ── input projections ──
+    img = img_in_kernel(latent, shared["img_in.weight"], shared["img_in.bias"])
+    txt = txt_in_kernel(
+        text, shared["txt_norm.weight"],
+        shared["txt_in.weight"], shared["txt_in.bias"], eps=1e-6,
+    )
+
+    # ── timestep conditioning ──
+    temb = time_text_embed_kernel(
+        timestep,
+        shared["time.proj.weight"], shared["time.proj.bias"],
+        shared["time.emb.weight"], shared["time.emb.bias"],
+        configs.dtype,
+    )
+
+    # ── RoPE tables (comptime) ──
+    vid_ang, txt_ang = compute_rope_freqs(
+        frame, gh, gw, Ltext, configs.axes_dims_rope,
+        theta=configs.rope_theta, scale_rope=True,
+    )
+    img_cos, img_sin = cos_sin_from_angles(vid_ang, dtype=configs.dtype)
+    txt_cos, txt_sin = cos_sin_from_angles(txt_ang, dtype=configs.dtype)
+
+    # ── transformer blocks ──
+    for w in blocks:
+        img, txt = mmdit_block(
+            img, txt, w, temb, n_heads, head_dim, hidden, eps,
+            img_cos, img_sin, txt_cos, txt_sin, txt_mask_bias=txt_mask_bias,
+            local_heads=local_heads, all_reduce_fn=all_reduce_fn,
+        )
+
+    # ── final layer (image stream only) ──
+    out = final_layer(
+        img, temb,
+        shared["norm_out.linear.weight"], shared["norm_out.linear.bias"],
+        shared["proj_out.weight"], shared["proj_out.bias"], eps=1e-6,
+    )
+    return out
+
+
+def qwenimage_forward(latent, text, timestep, img_shape, configs: Config,
+                      text_mask=None, unpatch=False, **weights):
+    """MMDiT denoiser.
+
+    Args:
+        latent: (B, Limg, in_channels) patchified latent.
+        text: (B, Ltext, joint_attention_dim) text embeddings.
+        timestep: (B,) diffusion timestep.
+        img_shape: (frame, gh, gw) patch grid for RoPE / unpatchify.
+        text_mask: optional (B, Ltext) {0,1} mask for text padding.
+        unpatch: if True, unpatchify to (B, out_channels, gh*p, gw*p); otherwise
+            return the raw proj_out (B, Limg, patch**2 * out_channels), matching
+            the diffusers transformer output.
+    """
+    shared, blocks = regroup_weights(weights, configs.num_layers)
+
+    txt_mask_bias = None
+    if text_mask is not None:
+        txt_mask_bias = (1.0 - text_mask.astype(configs.dtype)) * np.float32(-1e9)
+        txt_mask_bias = np.expand_dims(txt_mask_bias, axis=[1, 2])  # (B,1,1,Ltext)
+
+    out = _core(latent, text, timestep, img_shape, shared, blocks, configs,
+                txt_mask_bias=txt_mask_bias)
+
+    if unpatch:
+        frame, gh, gw = img_shape
+        out = unpatchify(out, configs.patch_size, configs.out_channels, gh, gw)
+    return out
+
+
+def denoise_step(
+    latents,       # (B, Limg, in_channels)  current packed latent (resident)
+    cond_text,     # (B, Ltext, joint_dim)   conditional text embeddings
+    neg_text,      # (B, Ltext, joint_dim)   unconditional text embeddings
+    timestep,      # (B,) current flow-match timestep (already /1000)
+    coeffs,        # (2,) [true_cfg_scale, dt]  host-precomputed scalars
+    img_shape, configs: Config,
+    cond_mask=None, neg_mask=None, **weights,
+):
+    """Fused one flow-match sampling step, keeping the latent resident on device.
+
+    Runs the MMDiT denoiser on the CFG batch (cond + uncond stacked), applies
+    Qwen-Image "true CFG" (norm-rescaled guidance), and performs the
+    FlowMatchEuler update ``prev = sample + dt * model_output`` — all on device.
+    The host supplies only the per-step scalars (``true_cfg_scale``, ``dt``);
+    both are functions of the fixed sigma schedule.
+
+    Model output and packed latents share the same shape (in_channels ==
+    patch**2 * out_channels == 64), so the step stays in packed space (unpatchify
+    happens once, on host, after the loop).
+
+    Returns prev_sample: (B, Limg, in_channels) latent for the next step.
+    """
+    shared, blocks = regroup_weights(weights, configs.num_layers)
+
+    # stack cond + uncond so the denoiser runs once over a 2B batch
+    latent_in = np.concatenate([latents, latents], axis=0)
+    text_in = np.concatenate([cond_text, neg_text], axis=0)
+    ts_in = np.concatenate([timestep, timestep], axis=0)
+
+    txt_mask_bias = None
+    if cond_mask is not None:
+        mask = np.concatenate([cond_mask, neg_mask], axis=0)
+        txt_mask_bias = (1.0 - mask.astype(configs.dtype)) * np.float32(-1e9)
+        txt_mask_bias = np.expand_dims(txt_mask_bias, axis=[1, 2])
+
+    noise = _core(latent_in, text_in, ts_in, img_shape, shared, blocks, configs,
+                  txt_mask_bias=txt_mask_bias)
+    noise = noise.astype(np.float32)
+    noise_cond, noise_uncond = np.split(noise, 2, axis=0)
+
+    # true CFG: combine then rescale to the conditional norm (per-token, over ch)
+    cfg_scale = coeffs[0]
+    comb = noise_uncond + cfg_scale * (noise_cond - noise_uncond)
+    cond_norm = np.sqrt(np.sum(noise_cond * noise_cond, axis=-1, keepdims=True))
+    comb_norm = np.sqrt(np.sum(comb * comb, axis=-1, keepdims=True))
+    model_output = comb * (cond_norm / comb_norm)
+
+    # FlowMatchEuler (non-stochastic): prev = sample + dt * model_output
+    sample = latents.astype(np.float32)
+    prev_sample = sample + coeffs[1] * model_output
+    return prev_sample.astype(configs.dtype)
diff --git a/examples/models/qwen_image/kernels/vae.py b/examples/models/qwen_image/kernels/vae.py
new file mode 100644
index 0000000..64b69a2
--- /dev/null
+++ b/examples/models/qwen_image/kernels/vae.py
@@ -0,0 +1,154 @@
+"""Qwen-Image VAE decoder (latents -> pixels) on Trainium.
+
+``AutoencoderKLQwenImage`` is a WAN-style 3D **causal video** VAE. For
+text-to-image we decode a **single frame** (T=1), and at T=1 the entire temporal
+machinery collapses to a plain 2D conv decoder (verified numerically to 0.0 diff
+against diffusers ``vae.decode``):
+
+* each ``QwenImageCausalConv3d`` (3x3x3, causal left-pad of 2 in time) sees a
+  T=1 input, so the two leading temporal taps multiply zero-padding and only the
+  **last** temporal tap contributes -> a 2D conv with weight ``w[:, :, -1]``. The
+  weight extraction bakes this collapse in, so the kernel just calls ``conv2d``.
+* ``feat_cache`` stays ``None`` for the single (first) frame, and the
+  ``upsample3d`` resamplers' temporal ``time_conv`` is on the "Rep" first-frame
+  path and is **skipped** entirely. Spatial upsampling is nearest-2x + conv3x3.
+
+Decoder structure (diffusers ``QwenImageDecoder3d``, Qwen-Image config
+base_dim=96, z_dim=16, dim_mult=[1,2,4,4], num_res_blocks=2):
+    conv_in (3x3)
+    mid_block: resnet -> attention (1x1 conv, single head) -> resnet
+    4x up_block: (num_res_blocks+1) residual blocks, then (all but last)
+        nearest-2x upsample + conv3x3
+    norm_out (RMS over channels) -> SiLU -> conv_out (3x3)
+
+Nonlinearity is SiLU. Norm is ``QwenImageRMS_norm``: L2-normalize over the
+channel dim, then ``* sqrt(dim) * gamma`` (no bias in this checkpoint). Weights
+arrive flat via ``**weights`` (see ``vae_weight_layout``). Runs in fp32 (the
+decoder is numerically sensitive and one-shot at the end of sampling); conv is
+the weak path on the PE array, so this is correctness-first, not a perf win.
+"""
+
+import numpy as np
+from nkipy.core import tensor_apis
+
+
+def _silu(x):
+    # sigmoid via clip to avoid overflow on large-magnitude activations
+    return x * (1 / (1 + np.exp(-np.clip(x, -30, 30))))
+
+
+def _softmax(x):
+    e = np.exp(x - np.max(x, axis=-1, keepdims=True))
+    return e / np.sum(e, axis=-1, keepdims=True)
+
+
+def rms_norm(x, gamma, eps=1e-12):
+    """QwenImageRMS_norm over (B, C, H, W): L2-normalize channels * sqrt(C) * gamma.
+
+    ``F.normalize(x, dim=1)`` == x / sqrt(sum_c x^2 + eps); the layer then scales
+    by ``sqrt(dim)`` and the per-channel ``gamma`` (no bias in this checkpoint).
+    Computed in fp32.
+    """
+    dtype = x.dtype
+    B, C, H, W = x.shape
+    xf = x.astype(np.float32)
+    denom = np.sqrt(np.sum(np.square(xf), axis=1, keepdims=True) + eps)
+    normalized = xf / denom
+    scale = np.float32(np.sqrt(C))
+    g = gamma.astype(np.float32).reshape(1, C, 1, 1)
+    return (normalized * scale * g).astype(dtype)
+
+
+def conv2d(x, w, b, stride=1, padding=1):
+    out = tensor_apis.conv2d(x, w, stride=stride, padding=padding)
+    return out + b.reshape(1, -1, 1, 1)
+
+
+def resnet_block(x, w, prefix):
+    """QwenImageResidualBlock: (RMS-SiLU-conv3x3) x2 + shortcut.
+
+    The shortcut is a 1x1 conv when in/out channels differ, else identity
+    (diffusers uses ``nn.Identity`` -> no shortcut weight in the flat dict).
+    """
+    if prefix + "conv_shortcut.weight" in w:
+        shortcut = conv2d(x, w[prefix + "conv_shortcut.weight"],
+                          w[prefix + "conv_shortcut.bias"], padding=0)
+    else:
+        shortcut = x
+
+    h = rms_norm(x, w[prefix + "norm1.gamma"])
+    h = _silu(h)
+    h = conv2d(h, w[prefix + "conv1.weight"], w[prefix + "conv1.bias"], padding=1)
+    h = rms_norm(h, w[prefix + "norm2.gamma"])
+    h = _silu(h)
+    h = conv2d(h, w[prefix + "conv2.weight"], w[prefix + "conv2.bias"], padding=1)
+    return shortcut + h
+
+
+def attention(x, w, prefix):
+    """QwenImageAttentionBlock: single-head spatial self-attention.
+
+    RMS-norm, then 1x1-conv qkv (a per-pixel linear), attention over H*W, 1x1-conv
+    proj, residual. The 1x1 convs are applied as matmuls over the (HW, C) layout.
+    """
+    B, C, H, W = x.shape
+    h = rms_norm(x, w[prefix + "norm.gamma"])
+    # (B, C, H, W) -> (B, HW, C)
+    seq = h.reshape(B, C, H * W).transpose(0, 2, 1)
+    # to_qkv: 1x1 conv (out 3C) == linear on channels. Weight stored (in, 3C).
+    qkv = np.matmul(seq, w[prefix + "to_qkv.weight"]) + w[prefix + "to_qkv.bias"]
+    q, k, v = np.split(qkv, 3, axis=-1)
+
+    scores = (q @ k.transpose(0, 2, 1)).astype(np.float32) / np.float32(np.sqrt(C))
+    attn = _softmax(scores).astype(v.dtype)
+    out = attn @ v  # (B, HW, C)
+    out = np.matmul(out, w[prefix + "proj.weight"]) + w[prefix + "proj.bias"]
+
+    out = out.transpose(0, 2, 1).reshape(B, C, H, W)
+    return x + out
+
+
+def upsample_nearest(x):
+    """Nearest-exact 2x spatial upsampling via repeat (QwenImageUpsample)."""
+    x = np.repeat(x, 2, axis=2)
+    x = np.repeat(x, 2, axis=3)
+    return x
+
+
+def vae_decode(latents, configs, **weights):
+    """Decode (1, z_dim, h, w) latents to (1, 3, H, W) pixels in [-1, 1].
+
+    ``latents`` must already be denormalized (``z * std + mean``) on the host,
+    matching ``QwenImagePipeline`` (the driver does this before calling). The
+    single-frame temporal dim is squeezed out on the host so the kernel is pure
+    2D; output is the T=1 frame (1, 3, H, W).
+    """
+    from .vae_weight_layout import regroup_vae_weights
+
+    w = regroup_vae_weights(weights)
+    nrb = configs.vae_num_res_blocks
+    n_up = len(configs.vae_dim_mult)
+
+    # post_quant_conv (1x1) then conv_in (3x3)
+    x = conv2d(latents, w["post_quant.weight"], w["post_quant.bias"], padding=0)
+    x = conv2d(x, w["conv_in.weight"], w["conv_in.bias"], padding=1)
+
+    # mid block: resnet, attention, resnet
+    x = resnet_block(x, w, "mid.resnets.0.")
+    x = attention(x, w, "mid.attentions.0.")
+    x = resnet_block(x, w, "mid.resnets.1.")
+
+    # up blocks: (nrb + 1) residual blocks, then spatial upsample (all but last)
+    for i in range(n_up):
+        for j in range(nrb + 1):
+            x = resnet_block(x, w, f"up.{i}.resnets.{j}.")
+        if f"up.{i}.upsample.weight" in w:
+            x = upsample_nearest(x)
+            x = conv2d(x, w[f"up.{i}.upsample.weight"],
+                       w[f"up.{i}.upsample.bias"], padding=1)
+
+    # output head: RMS-norm -> SiLU -> conv_out (3x3)
+    x = rms_norm(x, w["norm_out.gamma"])
+    x = _silu(x)
+    x = conv2d(x, w["conv_out.weight"], w["conv_out.bias"], padding=1)
+    return x
diff --git a/examples/models/qwen_image/kernels/vae_weight_layout.py b/examples/models/qwen_image/kernels/vae_weight_layout.py
new file mode 100644
index 0000000..8a9b273
--- /dev/null
+++ b/examples/models/qwen_image/kernels/vae_weight_layout.py
@@ -0,0 +1,19 @@
+"""Flat weight-key scheme for the Qwen-Image VAE decoder kernel.
+
+Keys mirror the ``QwenImageDecoder3d`` module tree with short prefixes:
+    conv_in.*, conv_out.*, norm_out.gamma
+    mid.resnets.{0,1}.*, mid.attentions.0.*
+    up.{i}.resnets.{j}.*, up.{i}.upsample.{weight,bias}
+
+Conv weights are collapsed from 3D (out, in, kt, kh, kw) to 2D (out, in, kh, kw)
+during extraction (the T=1 last-temporal-tap collapse; see ``kernels/vae.py``),
+so the kernel consumes plain (out, in, kh, kw) directly via nkipy conv2d. The
+attention ``to_qkv``/``proj`` are 1x1 convs stored as (in, out) matmul weights.
+``regroup_vae_weights`` is an identity pass-through (keys are already flat) so
+the kernel and prep share one naming source of truth.
+"""
+
+
+def regroup_vae_weights(flat):
+    # keys are already the canonical flat names; nothing to regroup
+    return dict(flat)
diff --git a/examples/models/qwen_image/kernels/weight_layout.py b/examples/models/qwen_image/kernels/weight_layout.py
new file mode 100644
index 0000000..af8e4ef
--- /dev/null
+++ b/examples/models/qwen_image/kernels/weight_layout.py
@@ -0,0 +1,70 @@
+"""Canonical flat weight-key scheme shared by weight prep and the kernel.
+
+The compiled ``qwenimage_forward`` kernel takes all weights as flat ``**weights``
+kwargs (the tracer only turns top-level ndarrays into HLO parameters — it does
+not recurse into dicts). Per-block weights are prefixed ``b{layer}_``; shared
+weights use bare names. ``regroup_weights`` rebuilds the nested structure.
+
+Each MMDiT block is dual-stream: an image stream (``img_*``) and a text stream
+(``txt_*``). Both feed one joint attention. Naming maps to the diffusers
+``QwenImageTransformerBlock`` submodules:
+
+    img_mod / txt_mod  -> SiLU-Linear(dim, 6*dim) modulation  (weight+bias)
+    img: to_q/to_k/to_v/to_out  + norm_q/norm_k (QK-RMSNorm)
+    txt: add_q_proj/add_k_proj/add_v_proj/to_add_out + norm_added_q/norm_added_k
+    img_mlp / txt_mlp  -> FeedForward gelu-approximate (fc1/fc2 weight+bias)
+
+The layout is grouped by stream so a later tensor-parallel split (shard heads +
+MLP intermediate) can slice per-stream keys uniformly.
+"""
+
+# short per-block keys -> filled by tensor_prep from the diffusers state dict.
+# ``*_w`` weights, ``*_b`` biases, ``*_g`` RMSNorm gains.
+BLOCK_KEYS = [
+    # modulation (SiLU -> Linear to 6*hidden), one per stream
+    "img_mod_w", "img_mod_b", "txt_mod_w", "txt_mod_b",
+    # image-stream attention projections
+    "iq_w", "iq_b", "ik_w", "ik_b", "iv_w", "iv_b", "io_w", "io_b",
+    "iq_g", "ik_g",  # QK-RMSNorm gains (norm_q / norm_k)
+    # text-stream attention projections
+    "tq_w", "tq_b", "tk_w", "tk_b", "tv_w", "tv_b", "to_w", "to_b",
+    "tq_g", "tk_g",  # QK-RMSNorm gains (norm_added_q / norm_added_k)
+    # image-stream MLP (gelu-approximate)
+    "iff0_w", "iff0_b", "iff2_w", "iff2_b",
+    # text-stream MLP
+    "tff0_w", "tff0_b", "tff2_w", "tff2_b",
+]
+
+# shared (non-block) weights
+SHARED_KEYS = [
+    "img_in.weight", "img_in.bias",
+    "txt_norm.weight",
+    "txt_in.weight", "txt_in.bias",
+    "time.proj.weight", "time.proj.bias",     # TimestepEmbedding fc1
+    "time.emb.weight", "time.emb.bias",        # TimestepEmbedding fc2
+    "norm_out.linear.weight", "norm_out.linear.bias",  # AdaLayerNormContinuous
+    "proj_out.weight", "proj_out.bias",
+]
+
+
+def block_key(layer_id, short_name):
+    return f"b{layer_id}_{short_name}"
+
+
+def regroup_weights(flat, num_layers):
+    """Split a flat ``{key: tensor}`` dict into (shared, blocks).
+
+    Returns:
+        shared: dict of SHARED_KEYS -> tensor
+        blocks: list (len num_layers) of dict short_name -> tensor
+    """
+    shared = {k: flat[k] for k in SHARED_KEYS if k in flat}
+    blocks = []
+    for layer_id in range(num_layers):
+        blk = {}
+        for short in BLOCK_KEYS:
+            key = block_key(layer_id, short)
+            if key in flat:
+                blk[short] = flat[key]
+        blocks.append(blk)
+    return shared, blocks
diff --git a/examples/models/qwen_image/qwen_image.py b/examples/models/qwen_image/qwen_image.py
new file mode 100644
index 0000000..3473912
--- /dev/null
+++ b/examples/models/qwen_image/qwen_image.py
@@ -0,0 +1,772 @@
+"""Qwen-Image text-to-image on Trainium — driver / entry point.
+
+Runs the full pipeline on device (denoiser + Qwen2.5 text encoder + VAE); the
+host keeps only the tokenizer + embedding lookup, latent pack/denorm, and the
+scalar flow-match schedule. The host-side glue here (packing, img_shapes,
+true-CFG, flow-match step) mirrors ``diffusers.QwenImagePipeline`` exactly. See
+README.md for the architecture and the device/host split.
+
+Tensor parallelism is required: 20B bf16 (~40 GB) does not fit on one 24 GB trn2
+core, so launch with torchrun. TP=4 is the validated layout; TP is capped at
+num_kv_heads=4 by the text encoder.
+"""
+
+import argparse
+import os
+import time
+
+import numpy as np
+import torch
+from config import Config, get_config
+from kernels.text_encoder import text_encoder_forward
+from kernels.transformer import denoise_step
+from kernels.vae import vae_decode
+from kernels.weight_layout import BLOCK_KEYS, SHARED_KEYS, block_key
+
+
+def _unpack_latents(latents, height, width, vae_scale_factor):
+    """Packed tokens -> (B, C, 1, H, W) latent. Matches the pipeline (5-D for the
+    3D Qwen VAE)."""
+    batch_size, num_patches, channels = latents.shape
+    height = 2 * (int(height) // (vae_scale_factor * 2))
+    width = 2 * (int(width) // (vae_scale_factor * 2))
+    latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+    return latents.reshape(batch_size, channels // 4, 1, height, width)
+
+
+class _DeviceModule:
+    """Shared plumbing for the on-device components (denoiser / text encoder /
+    VAE): upload the pre-sharded weights once, compile a kernel per input shape
+    (cached), and run it.
+
+    The tracer prunes weights that don't reach the output, so every compiled
+    kernel is invoked with only the subset in ``kernel.input_tensors_info``
+    (``_compile`` returns that subset as ``wkeys``).
+    """
+
+    def __init__(self, weights, config, fp32=False):
+        from nkipy.runtime import DeviceTensor
+        from nkipy.runtime.device_tensor import bfloat16
+
+        self._DeviceTensor = DeviceTensor
+        self._bf16 = bfloat16
+        self.config = config
+        self.device_weights = self._upload(weights, fp32)
+        self._weight_keys = self._select_weight_keys()
+        self._kernels = {}
+
+    def _upload(self, weights, fp32):
+        """torch/numpy weights -> ``{key: DeviceTensor}`` (bf16, or fp32 for the
+        numerically-sensitive VAE)."""
+        DT = self._DeviceTensor
+        out = {}
+        for key, tensor in weights.items():
+            if fp32:
+                arr = tensor.detach().cpu().numpy() if isinstance(tensor, torch.Tensor) else tensor
+                out[key] = DT.from_numpy(np.ascontiguousarray(arr.astype(np.float32)), name=key)
+            else:
+                src = tensor.to(torch.bfloat16) if tensor.is_floating_point() else tensor
+                out[key] = DT.from_torch(src, name=key)
+        return out
+
+    def _select_weight_keys(self):
+        """Weights to feed the kernel, before tracer pruning. Default: all of
+        them in a deterministic (sorted) order; the denoiser / text encoder
+        narrow this to their own canonical key scheme.
+
+        The order matters: it becomes the kernel's HLO parameter order, which
+        feeds the compile cache's content hash. Sorting keeps that order
+        independent of the weight-dict's source order (extraction preserves
+        insertion order; the shard cache round-trips through safetensors, which
+        re-sorts keys) so a cached and a freshly-extracted run hit the same
+        compiled NEFF."""
+        return sorted(self.device_weights)
+
+    def _compile(self, key, kernel_fn, name, placeholders, build_dir):
+        """Compile once (cached by ``key``) and return ``(kernel, wkeys)``, where
+        ``wkeys`` is the weight subset the traced graph actually reads."""
+        if key in self._kernels:
+            return self._kernels[key]
+        from nkipy.runtime import DeviceKernel
+
+        weight_kwargs = {k: self.device_weights[k] for k in self._weight_keys}
+        kernel = DeviceKernel.compile_and_load(
+            kernel_fn, name=name, configs=self.config, build_dir=build_dir,
+            additional_compiler_args=self.config.additional_compiler_args_nkipy,
+            **placeholders, **weight_kwargs,
+        )
+        wkeys = [k for k in self._weight_keys if k in set(kernel.input_tensors_info)]
+        self._kernels[key] = (kernel, wkeys)
+        return kernel, wkeys
+
+    def _run(self, kernel, wkeys, host_inputs, out):
+        """Assemble host inputs + the live weight subset, run into ``out``."""
+        inputs = dict(host_inputs)
+        inputs.update({k: self.device_weights[k] for k in wkeys})
+        kernel(inputs=inputs, outputs={"output0": out})
+        return out
+
+
+class QwenImageDenoiser(_DeviceModule):
+    """The Trainium MMDiT denoiser (host encoder / VAE live in the driver).
+
+    Compiles the fused ``denoise_step`` once per text length; ``sample`` runs the
+    full denoising loop with the packed latent resident on device (CFG +
+    FlowMatchEuler on device, no per-step host round-trip).
+    """
+
+    def __init__(self, weights, config: Config, img_shape, batch_size=1):
+        self.img_shape = img_shape  # (frame, gh, gw)
+        self.batch_size = batch_size
+        t = time.time()
+        super().__init__(weights, config)
+        print(f"[qwen-image] denoiser weights ready in {time.time() - t:.2f}s "
+              f"(TP={getattr(config, 'tp_size', 1)})")
+
+    def _select_weight_keys(self):
+        keys = [k for k in SHARED_KEYS if k in self.device_weights]
+        for i in range(self.config.num_layers):
+            keys += [block_key(i, s) for s in BLOCK_KEYS
+                     if block_key(i, s) in self.device_weights]
+        return keys
+
+    def _step_kernel_for(self, txt_len):
+        """Compile (once) the fused ``denoise_step`` kernel for a text length.
+
+        cond and uncond prompts can have different token counts, so we compile
+        per text length on demand (usually 1-2 distinct lengths).
+        """
+        if txt_len in self._kernels:
+            return self._kernels[txt_len]
+        cfg = self.config
+        frame, gh, gw = self.img_shape
+        Limg = frame * gh * gw
+        B = self.batch_size
+        DT, bf16 = self._DeviceTensor, self._bf16
+
+        placeholders = dict(
+            latents=DT.from_numpy(np.empty((B, Limg, cfg.in_channels), dtype=bf16), "latents"),
+            cond_text=DT.from_numpy(np.empty((B, txt_len, cfg.joint_attention_dim), dtype=bf16), "cond_text"),
+            neg_text=DT.from_numpy(np.empty((B, txt_len, cfg.joint_attention_dim), dtype=bf16), "neg_text"),
+            timestep=DT.from_numpy(np.empty((B,), dtype=np.float32), "timestep"),
+            coeffs=DT.from_numpy(np.empty((2,), dtype=np.float32), "coeffs"),
+            img_shape=self.img_shape,
+            cond_mask=DT.from_numpy(np.empty((B, txt_len), dtype=bf16), "cond_mask"),
+            neg_mask=DT.from_numpy(np.empty((B, txt_len), dtype=bf16), "neg_mask"),
+        )
+        print(f"[qwen-image] compiling denoise_step "
+              f"({cfg.num_layers} blocks, grid {gh}x{gw}, txt {txt_len})")
+        t = time.time()
+        kernel, wkeys = self._compile(txt_len, denoise_step,
+                                      f"denoise_step_t{txt_len}", placeholders, "./build")
+        print(f"[qwen-image] --> step kernel (txt {txt_len}) ready in {time.time() - t:.2f}s")
+        return kernel, wkeys
+
+    def sample(self, init_latents, cond_text, neg_text, cond_mask, neg_mask,
+               step_coeffs):
+        """Full denoising loop with the latent resident on device (fused step).
+
+        Runs ``denoise_step`` once per step; the packed latent stays on device
+        across all steps (no per-step host round-trip). Text embeddings/masks are
+        uploaded once. ``cond_text``/``neg_text`` must share a text length (pad +
+        mask on the host). ``step_coeffs`` is a list of per-step ``[cfg_scale, dt]``.
+
+        All tensors are torch; returns the final packed latent (torch, B×Limg×C).
+        """
+        DT, bf16 = self._DeviceTensor, self._bf16
+
+        def to_np(t, dt):
+            return t.detach().to(torch.float32).cpu().numpy().astype(dt)
+
+        txt_len = cond_text.shape[1]
+        kernel, wkeys = self._step_kernel_for(txt_len)
+
+        B, Limg, C = init_latents.shape
+        latents = DT.from_numpy(to_np(init_latents, bf16), "latents")
+        nxt = DT.from_numpy(np.empty((B, Limg, C), dtype=bf16), "latents_next")
+        cond_d = DT.from_numpy(to_np(cond_text, bf16), "cond_text")
+        neg_d = DT.from_numpy(to_np(neg_text, bf16), "neg_text")
+        cmask_d = DT.from_numpy(to_np(cond_mask, bf16), "cond_mask")
+        nmask_d = DT.from_numpy(to_np(neg_mask, bf16), "neg_mask")
+
+        for cfg_scale, dt, ts_val in step_coeffs:
+            ts_d = DT.from_numpy(np.full((B,), ts_val, dtype=np.float32), "timestep")
+            coeffs_d = DT.from_numpy(np.array([cfg_scale, dt], dtype=np.float32), "coeffs")
+            host_inputs = {
+                "latents": latents, "cond_text": cond_d, "neg_text": neg_d,
+                "timestep": ts_d, "coeffs": coeffs_d,
+                "cond_mask": cmask_d, "neg_mask": nmask_d,
+            }
+            self._run(kernel, wkeys, host_inputs, nxt)
+            latents, nxt = nxt, latents  # swap; no host round-trip of latent data
+
+        return latents.torch().to(torch.float32)
+
+
+class DeviceTextEncoder(_DeviceModule):
+    """The Trainium Qwen2.5 text encoder (``kernels/text_encoder.py``).
+
+    Runs the text-only Qwen2.5-VL decoder LM on device and returns the last
+    hidden state. The host keeps the tokenizer, chat template, and embedding-
+    table lookup (huge, data-dependent gather); the device runs the 28-layer
+    transformer, TP-sharded like the denoiser (the 7B encoder doesn't fit one
+    core). Compiles ``text_encoder_forward`` once per sequence length.
+    """
+
+    def __init__(self, weights, te_config):
+        t = time.time()
+        super().__init__(weights, te_config)
+        print(f"[qwen-image] text-encoder weights ready in {time.time() - t:.2f}s "
+              f"(TP={getattr(te_config, 'tp_size', 1)})")
+
+    def _select_weight_keys(self):
+        from kernels.text_weight_layout import present_keys
+        return present_keys(self.device_weights, self.config.num_layers)
+
+    def _kernel_for(self, seq_len, hidden):
+        if seq_len in self._kernels:
+            return self._kernels[seq_len]
+        DT, bf16, cfg = self._DeviceTensor, self._bf16, self.config
+        placeholders = {"hidden": DT.from_numpy(np.empty((1, seq_len, hidden), dtype=bf16), "hidden")}
+        print(f"[qwen-image] compiling text_encoder_forward "
+              f"({cfg.num_layers} layers, seq {seq_len})")
+        t = time.time()
+        kernel, wkeys = self._compile(seq_len, text_encoder_forward,
+                                      f"text_encoder_forward_s{seq_len}", placeholders, "./build_te")
+        print(f"[qwen-image] --> text encoder (seq {seq_len}) ready in {time.time() - t:.2f}s")
+        return kernel, wkeys
+
+    def encode(self, embeds):
+        """Run the encoder. ``embeds``: (1, S, hidden) torch token embeddings.
+
+        Returns the last hidden state as a torch tensor (1, S, hidden).
+        """
+        B, S, H = embeds.shape
+        assert B == 1, "device text encoder processes one prompt at a time"
+        kernel, wkeys = self._kernel_for(S, H)
+        DT = self._DeviceTensor
+        hidden_np = embeds.detach().to(torch.float32).cpu().numpy().astype(self._bf16)
+        host_inputs = {"hidden": DT.from_numpy(hidden_np, "hidden")}
+        out = DT.from_numpy(np.empty((1, S, H), dtype=self._bf16), "out")
+        self._run(kernel, wkeys, host_inputs, out)
+        return out.torch().to(torch.float32)
+
+
+def _round_up(n, bucket):
+    """Round ``n`` up to the next multiple of ``bucket`` (``bucket <= 0`` -> no
+    change). Used to snap prompt/text lengths onto a small set of buckets so
+    varying prompts reuse a few compiled kernels instead of recompiling per
+    exact length."""
+    if bucket <= 0:
+        return int(n)
+    return ((int(n) + bucket - 1) // bucket) * bucket
+
+
+def encode_prompt_device(pipe, device_encoder, prompt, text_bucket=0):
+    """Replicate ``QwenImagePipeline._get_qwen_prompt_embeds`` with the encoder
+    forward on device.
+
+    The host does everything data-dependent (chat-template wrapping, tokenize,
+    embedding lookup, masked-hidden extraction, ``drop_idx`` prefix slice); the
+    device runs only the transformer. Batch=1 (cond and neg are encoded
+    separately). Returns ``(prompt_embeds, prompt_embeds_mask)`` in the same
+    layout as ``encode_prompt``.
+
+    With ``text_bucket > 0`` the encoder input is right-padded to the next
+    bucket multiple so different-length prompts hit one compiled kernel per
+    bucket. This is exact: the encoder is causal, so a real token at position
+    ``i < L`` never attends to a later pad position — its hidden state is
+    identical with or without the padding — and we slice the pad tail back off
+    before the masked-hidden extraction.
+    """
+    template = pipe.prompt_template_encode
+    drop_idx = pipe.prompt_template_encode_start_idx
+    txt = [template.format(prompt)]
+    txt_tokens = pipe.tokenizer(
+        txt, max_length=pipe.tokenizer_max_length + drop_idx,
+        padding=True, truncation=True, return_tensors="pt")
+    input_ids = txt_tokens.input_ids
+    attn_mask = txt_tokens.attention_mask
+
+    # host embedding lookup (the table stays on host)
+    with torch.no_grad():
+        embeds = pipe.text_encoder.get_input_embeddings()(input_ids)
+    # right-pad the encoder input to a bucket length (causal-safe, see docstring)
+    L = embeds.shape[1]
+    s_fixed = _round_up(L, text_bucket)
+    if s_fixed > L:
+        embeds = torch.cat(
+            [embeds, embeds.new_zeros(embeds.shape[0], s_fixed - L, embeds.shape[2])], dim=1)
+    # device transformer -> last hidden state, then drop the padded tail
+    hidden = device_encoder.encode(embeds.to(torch.float32))[:, :L]
+
+    # masked-hidden extraction + drop template prefix (matches diffusers)
+    bool_mask = attn_mask.bool()
+    valid = bool_mask.sum(dim=1).tolist()
+    selected = hidden[bool_mask]
+    split = torch.split(selected, valid, dim=0)
+    split = [e[drop_idx:] for e in split]
+    max_len = max(e.size(0) for e in split)
+    prompt_embeds = torch.stack(
+        [torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) for u in split])
+    mask = torch.stack(
+        [torch.cat([torch.ones(u.size(0)), torch.zeros(max_len - u.size(0))]) for u in split])
+    if mask.all():
+        mask = None
+    return prompt_embeds, mask
+
+
+class DeviceVAEDecoder(_DeviceModule):
+    """The Trainium VAE decoder (``kernels/vae.py``).
+
+    Decodes a single image (T=1), where ``AutoencoderKLQwenImage``'s 3D causal
+    video VAE collapses exactly to a 2D conv decoder. The host squeezes the
+    temporal dim and does the latent denormalization; the device runs the conv
+    stack. Runs in fp32 (numerically sensitive, one-shot). Compiles per latent
+    spatial size. Replicated per rank (small, not sharded).
+    """
+
+    def __init__(self, weights, vae_config):
+        super().__init__(weights, vae_config, fp32=True)
+
+    def _kernel_for(self, h, w):
+        if (h, w) in self._kernels:
+            return self._kernels[(h, w)]
+        DT, cfg = self._DeviceTensor, self.config
+        placeholders = {"latents": DT.from_numpy(np.empty((1, cfg.z_dim, h, w), dtype=np.float32), "latents")}
+        print(f"[qwen-image] compiling vae_decode (latent {h}x{w})")
+        t = time.time()
+        kernel, wkeys = self._compile((h, w), vae_decode,
+                                      f"vae_decode_{h}x{w}", placeholders, "./build_vae")
+        print(f"[qwen-image] --> vae kernel ({h}x{w}) ready in {time.time() - t:.2f}s")
+        return kernel, wkeys
+
+    def decode(self, latents):
+        """Args: denormalized latents torch (1, z_dim, H, W). Returns pixels
+        torch (1, 3, H*8, W*8) in [-1, 1] (unclamped; caller clamps)."""
+        _, _, h, w = latents.shape
+        kernel, wkeys = self._kernel_for(h, w)
+        DT, cfg = self._DeviceTensor, self.config
+        lat_np = latents.detach().to(torch.float32).cpu().numpy().astype(np.float32)
+        host_inputs = {"latents": DT.from_numpy(lat_np, "latents")}
+        scale = 2 ** (len(cfg.dim_mult) - 1)
+        out = DT.from_numpy(
+            np.empty((1, cfg.out_channels, h * scale, w * scale), dtype=np.float32), "pixels")
+        self._run(kernel, wkeys, host_inputs, out)
+        return out.torch().to(torch.float32)
+
+
+def load_host_pipeline(model_name, dtype=torch.bfloat16):
+    """Load the diffusers QwenImagePipeline (host encoder + VAE + scheduler)."""
+    from diffusers import QwenImagePipeline
+
+    return QwenImagePipeline.from_pretrained(model_name, torch_dtype=dtype)
+
+
+def _pad_text(embeds, mask, target_len):
+    """Right-pad (B, L, D) embeds to ``target_len`` and return a matching mask.
+
+    ``encode_prompt`` returns ``mask=None`` when every token is valid; we
+    materialise an all-ones mask so the padded tail can be masked to 0 (the
+    fused denoise_step always takes explicit masks).
+    """
+    B, L, D = embeds.shape
+    if mask is None:
+        mask = embeds.new_ones(B, L)
+    if L == target_len:
+        return embeds, mask
+    pad = target_len - L
+    embeds = torch.cat([embeds, embeds.new_zeros(B, pad, D)], dim=1)
+    mask = torch.cat([mask, mask.new_zeros(B, pad)], dim=1)
+    return embeds, mask
+
+
+def generate(pipe, denoiser, prompt, negative_prompt, config, height, width,
+             guidance_scale, num_steps, text_encoder, vae_decoder, seed=0):
+    """Full pipeline: text-encode, prepare latents + flow-match schedule, run the
+    device-resident sampling loop, then unpack + VAE-decode to pixels.
+
+    Requires true-CFG (``guidance_scale > 1.0``) — Qwen-Image always uses it, and
+    the fused device step is built around the cond/neg pair.
+    """
+    if guidance_scale <= 1.0:
+        raise ValueError("qwen-image requires true-CFG (guidance_scale > 1.0)")
+    device = "cpu"
+    vae_scale = pipe.vae_scale_factor
+    text_bucket = getattr(config, "text_bucket", 0)
+
+    # text encode on device
+    prompt_embeds, prompt_mask = encode_prompt_device(pipe, text_encoder, prompt, text_bucket)
+    neg_embeds, neg_mask = encode_prompt_device(pipe, text_encoder, negative_prompt, text_bucket)
+
+    # init latents (packed) + flow-match schedule
+    num_ch = pipe.transformer.config.in_channels // 4
+    gen = torch.Generator().manual_seed(seed)
+    latents = pipe.prepare_latents(
+        1, num_ch, height, width, prompt_embeds.dtype, device, gen)
+
+    # Qwen-Image's scheduler uses dynamic shifting: mu is derived from the image
+    # token count (matches QwenImagePipeline.__call__).
+    from diffusers.pipelines.qwenimage.pipeline_qwenimage import calculate_shift
+
+    sigmas = np.linspace(1.0, 1 / num_steps, num_steps)
+    image_seq_len = latents.shape[1]
+    mu = calculate_shift(
+        image_seq_len,
+        pipe.scheduler.config.get("base_image_seq_len", 256),
+        pipe.scheduler.config.get("max_image_seq_len", 4096),
+        pipe.scheduler.config.get("base_shift", 0.5),
+        pipe.scheduler.config.get("max_shift", 1.15),
+    )
+    pipe.scheduler.set_timesteps(sigmas=sigmas, mu=mu, device=device)
+    timesteps = pipe.scheduler.timesteps
+
+    # device-resident loop: CFG + FlowMatchEuler run on device in denoise_step.
+    # cond/neg share a text length (pad + mask); dt comes from the scheduler
+    # sigmas (prev = sample + dt * model_output). Bucketing the shared length
+    # lets varying prompts reuse one compiled denoise_step (the padded text
+    # tokens are masked out via cond_mask/neg_mask, so the result is exact).
+    tlen = _round_up(max(prompt_embeds.shape[1], neg_embeds.shape[1]), text_bucket)
+    c_emb, c_mask = _pad_text(prompt_embeds, prompt_mask, tlen)
+    n_emb, n_mask = _pad_text(neg_embeds, neg_mask, tlen)
+    sched_sigmas = pipe.scheduler.sigmas  # (num_steps+1,)
+    step_coeffs = []
+    for i, t in enumerate(timesteps):
+        dt = float(sched_sigmas[i + 1] - sched_sigmas[i])
+        step_coeffs.append((float(guidance_scale), dt, float(t) / 1000.0))
+    latents = denoiser.sample(latents, c_emb, n_emb, c_mask, n_mask, step_coeffs)
+
+    # unpack + VAE decode on device. Denormalize exactly like QwenImagePipeline:
+    # the pipeline stores latents_std as its RECIPROCAL and then does
+    # ``latents / latents_std_recip`` (== latents * std) + mean. Matching that.
+    latents = _unpack_latents(latents, height, width, vae_scale)  # (1, z, 1, H, W)
+    latents = latents.to(pipe.vae.dtype)
+    z_dim = pipe.vae.config.z_dim
+    latents_mean = torch.tensor(pipe.vae.config.latents_mean).view(1, z_dim, 1, 1, 1).to(latents)
+    latents_std_recip = (1.0 / torch.tensor(pipe.vae.config.latents_std)).view(1, z_dim, 1, 1, 1).to(latents)
+    latents = latents / latents_std_recip + latents_mean
+    # device VAE: squeeze the single frame -> (1, z, H, W), decode on device
+    image = vae_decoder.decode(latents[:, :, 0])
+    image = (image / 2 + 0.5).clamp(0, 1)
+    return (image.permute(0, 2, 3, 1).cpu().float().numpy() * 255).round().astype(np.uint8)
+
+
+# Bump when the extraction/sharding layout changes so stale shard caches are
+# not silently reused (the content the cache stores must match what the kernels
+# expect). Part of the cache path.
+_SHARD_CACHE_VERSION = 1
+
+
+def _shard_cache_path(cache_dir, model_name, tag, tp_size, rank):
+    """Per-(model, component, tp, rank) shard-cache file. Model name is slugged so
+    it is safe as a directory component."""
+    slug = model_name.replace("/", "__")
+    return os.path.join(
+        cache_dir, f"v{_SHARD_CACHE_VERSION}", slug,
+        f"{tag}_tp{tp_size}_rank{rank}.safetensors")
+
+
+def _cached_weights(cache_path, build_fn, store_dtype, log, label, enabled):
+    """Load a rank's sharded weights from ``cache_path`` if present, else build
+    them with ``build_fn`` (extract + shard from the host pipeline) and write the
+    cache. Floating tensors are stored as ``store_dtype`` (bf16 for the
+    denoiser/text-encoder — exactly what gets uploaded — fp32 for the VAE), which
+    both halves the cache/load size and skips the fp32 extraction blow-up on
+    cached runs. Returns a ``{key: torch.Tensor}`` dict ready to upload.
+
+    Extraction is the whole cost of weight prep (~330 s at TP=4: transpose +
+    copy + fp32 cast over 20B params, ranks contending for memory bandwidth); a
+    cache hit replaces it with a memory-mapped read of a few GB.
+    """
+    import safetensors.torch as st
+
+    if enabled and os.path.exists(cache_path):
+        t = time.time()
+        weights = st.load_file(cache_path)
+        log(f"[qwen-image] {label}: loaded shard cache "
+            f"({len(weights)} tensors) in {time.time() - t:.2f}s")
+        return weights
+
+    t = time.time()
+    weights = build_fn()
+    log(f"[qwen-image] {label}: extracted+sharded in {time.time() - t:.2f}s")
+    if store_dtype is not None:
+        weights = {
+            k: (v.to(store_dtype) if v.is_floating_point() else v).contiguous()
+            for k, v in weights.items()
+        }
+    if enabled:
+        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+        # Per-rank filename (each rank writes its own shard), so writes never
+        # collide; the temp+rename keeps a concurrent reader from seeing a
+        # half-written file.
+        tmp = f"{cache_path}.tmp.{os.getpid()}"
+        st.save_file(weights, tmp)
+        os.replace(tmp, cache_path)
+        log(f"[qwen-image] {label}: wrote shard cache {cache_path}")
+    return weights
+
+
+def _build_device_denoiser_weights(pipe, config, tp_size, rank):
+    """Extract the denoiser's flat weights from the host pipeline's transformer
+    and slice this rank's tensor-parallel shard."""
+    from weight_extract import extract_flat_weights, shard_flat_weights
+
+    flat = extract_flat_weights(pipe.transformer, config.num_layers, dtype=np.float32)
+    shard = shard_flat_weights(flat, rank, tp_size, config.num_layers,
+                               config.num_heads, config.head_dim)
+    return {k: torch.from_numpy(np.ascontiguousarray(v)) for k, v in shard.items()}
+
+
+def _make_text_encoder_config(pipe, tp_size):
+    """Build the ``TextEncoderConfig`` from the host pipeline (cheap; needed on
+    every run to construct ``DeviceTextEncoder``, cached or not)."""
+    from config import TextEncoderConfig
+    from kernels.tp import make_all_reduce
+
+    te_hf = pipe.text_encoder.config.text_config
+    # rope_theta is top-level on older transformers, nested under
+    # ``rope_parameters``/``rope_scaling`` on newer ones (e.g. 5.5.x).
+    rope_theta = getattr(te_hf, "rope_theta", None)
+    if rope_theta is None:
+        params = getattr(te_hf, "rope_parameters", None) or getattr(te_hf, "rope_scaling", None) or {}
+        rope_theta = params.get("rope_theta", 1000000.0)
+    te_cfg = TextEncoderConfig(
+        num_layers=te_hf.num_hidden_layers,
+        hidden_size=te_hf.hidden_size,
+        num_heads=te_hf.num_attention_heads,
+        num_kv_heads=te_hf.num_key_value_heads,
+        head_dim=getattr(te_hf, "head_dim", te_hf.hidden_size // te_hf.num_attention_heads),
+        intermediate_size=te_hf.intermediate_size,
+        rms_norm_eps=te_hf.rms_norm_eps,
+        rope_theta=rope_theta,
+    )
+    te_cfg.tp_size = tp_size
+    te_cfg.all_reduce_fn = make_all_reduce(tp_size)
+    return te_cfg
+
+
+def _build_text_encoder_weights(pipe, te_cfg, tp_size, rank):
+    """Extract the Qwen2.5 encoder weights from the host pipeline and slice this
+    rank's TP shard. The host pipeline's encoder stays resident for the
+    embedding-table lookup (only the transformer layers run on device)."""
+    from weight_extract import extract_text_encoder_weights, shard_text_encoder_weights
+
+    lm = pipe.text_encoder.model.language_model
+    flat = extract_text_encoder_weights(lm, te_cfg.num_layers, dtype=np.float32)
+    flat = shard_text_encoder_weights(
+        flat, rank, tp_size, te_cfg.num_layers,
+        te_cfg.num_heads, te_cfg.num_kv_heads, te_cfg.head_dim)
+    return {k: torch.from_numpy(np.ascontiguousarray(v)) for k, v in flat.items()}
+
+
+def _build_vae_weights(pipe):
+    """Extract the VAE decoder weights (T=1 2D-collapsed decoder, replicated per
+    rank) from the host pipeline."""
+    from weight_extract import extract_vae_decoder_weights
+
+    flat = extract_vae_decoder_weights(pipe.vae, dtype=np.float32)
+    return {k: torch.from_numpy(np.ascontiguousarray(v)) for k, v in flat.items()}
+
+
+def _indexed_output(template, idx):
+    """``output.png`` -> ``output_000.png`` for the idx-th resident image."""
+    base, ext = os.path.splitext(template)
+    return f"{base}_{idx:03d}{ext}"
+
+
+def _generate_and_save(pipe, denoiser, text_encoder, vae_decoder, config, args,
+                       prompt, seed, out_path, rank):
+    """Run one image end-to-end and (rank 0) save it. Returns wall time (s).
+
+    Within a resident process the kernels are already loaded, so repeated calls
+    reuse them (the denoiser / text-encoder / VAE cache compiled kernels in
+    ``self._kernels`` and the runtime caches the loaded NEFF), i.e. no re-upload
+    and no NEFF reload — only a new text length triggers a one-time recompile.
+    """
+    t = time.time()
+    images = generate(pipe, denoiser, prompt, args.negative_prompt, config,
+                      args.height, args.width, args.guidance_scale, args.steps,
+                      text_encoder, vae_decoder, seed=seed)
+    dt = time.time() - t
+    if rank == 0:
+        from PIL import Image
+        Image.fromarray(images[0]).save(out_path)
+    return dt
+
+
+def _resident_loop(pipe, denoiser, text_encoder, vae_decoder, config, args,
+                   rank, log):
+    """Keep the process (and all TP ranks) resident, generating one image per
+    prompt so the one-time setup — weight upload (~25s) + NEFF load (~17s) that
+    no on-disk cache can remove — is paid once and amortized across every image.
+
+    Rank 0 is the source of prompts (``--prompts-file``, else interactive
+    stdin); each ``(prompt, seed)`` is broadcast to all ranks so they run
+    ``generate`` together — the kernels use collectives, so every rank must
+    participate on every image. A ``None`` prompt ends the loop on all ranks.
+    """
+    import torch.distributed as dist
+
+    src = None
+    if rank == 0 and args.prompts_file:
+        with open(args.prompts_file) as f:
+            prompts = [ln.strip() for ln in f if ln.strip()]
+        log(f"[qwen-image] resident mode: {len(prompts)} prompt(s) from {args.prompts_file}")
+        src = iter(prompts)
+    elif rank == 0:
+        log("[qwen-image] resident mode: interactive (blank line / Ctrl-D to quit)")
+
+    idx = 0
+    while True:
+        # Rank 0 picks the next prompt; all ranks agree via broadcast so they
+        # step through generate() in lockstep.
+        if rank == 0:
+            if src is not None:
+                prompt = next(src, None)
+            else:
+                try:
+                    prompt = input("prompt> ").strip() or None
+                except EOFError:
+                    prompt = None
+            payload = [prompt, int(args.seed) + idx]
+        else:
+            payload = [None, None]
+        dist.broadcast_object_list(payload, src=0)
+        prompt, seed = payload
+        if prompt is None:
+            break
+        out_path = _indexed_output(args.output, idx)
+        dt = _generate_and_save(pipe, denoiser, text_encoder, vae_decoder,
+                                config, args, prompt, seed, out_path, rank)
+        log(f"[qwen-image] image {idx} ({args.steps} steps) in {dt:.2f}s "
+            f"-> {out_path}  |  {prompt[:50]!r}")
+        idx += 1
+    log(f"[qwen-image] resident mode: generated {idx} image(s)")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prompt", nargs="?", default="a coffee shop entrance with a chalkboard sign")
+    parser.add_argument("--negative-prompt", default=" ")
+    parser.add_argument("--model", default="Qwen/Qwen-Image")
+    # 512px (grid 32x32) is the validated working resolution; native 1328px
+    # exceeds the compiler's 2 GB HLO-proto limit (see README.md).
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=4.0)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output", default="output.png")
+    parser.add_argument(
+        "--weight-cache-dir", default="./weight_cache",
+        help="directory for the per-rank sharded weight cache; a cache hit skips "
+             "the ~330s host-side extract/shard on every launch")
+    parser.add_argument(
+        "--no-weight-cache", action="store_true",
+        help="disable the shard cache (always extract+shard from the pipeline)")
+    parser.add_argument(
+        "--prompts-file", default=None,
+        help="resident batch mode: generate one image per non-empty line of this "
+             "file, keeping the process alive so the one-time weight upload + NEFF "
+             "load are paid once and amortized across every image")
+    parser.add_argument(
+        "--interactive", action="store_true",
+        help="resident REPL mode: read prompts from stdin (rank 0) until a blank "
+             "line or EOF; same amortization as --prompts-file")
+    parser.add_argument(
+        "--text-bucket", type=int, default=64,
+        help="round prompt text length up to this multiple so varying prompts "
+             "reuse one compiled text-encoder + denoise kernel per bucket instead "
+             "of recompiling per exact length (padded tokens are masked; exact). "
+             "0 disables (compile per exact length)")
+    args = parser.parse_args()
+
+    # ── tensor-parallel setup (torchrun) ──────────────────────────────────────
+    # NEURON_RT_DISABLE_EXECUTION_BARRIER must be 0 (set before nkipy import) for
+    # correct multi-rank collectives.
+    import torch.distributed as dist
+
+    tp_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if tp_size <= 1:
+        raise SystemExit(
+            "qwen-image requires tensor parallelism: the 20B model (~40 GB bf16) "
+            "does not fit on one 24 GB core. Launch with e.g. "
+            "`torchrun --nproc-per-node 4 qwen_image.py ...` (TP=4 fits; TP is "
+            "capped at num_kv_heads=4 by the text encoder).")
+    os.environ["NEURON_RT_DISABLE_EXECUTION_BARRIER"] = "0"
+    os.environ.setdefault("NEURON_RT_ROOT_COMM_ID", "localhost:61455")
+    dist.init_process_group()
+    rank = dist.get_rank()
+    os.environ["NEURON_RT_VISIBLE_CORES"] = str(rank)
+    torch.set_num_threads(1)
+
+    def log(msg):
+        if rank == 0:
+            print(msg)
+
+    from kernels.tp import make_all_reduce
+    config = get_config(args.model, args.steps)
+    config.model_name = args.model
+    config.tp_size = tp_size
+    config.all_reduce_fn = make_all_reduce(tp_size)
+    config.text_bucket = args.text_bucket
+
+    from config import get_vae_config
+
+    cache_enabled = not args.no_weight_cache
+    cache_dir = args.weight_cache_dir
+
+    def cpath(tag):
+        return _shard_cache_path(cache_dir, args.model, tag, tp_size, rank)
+
+    setup_start = time.time()
+    log(f"[qwen-image] loading host pipeline {args.model} (TP={tp_size})")
+    pipe = load_host_pipeline(args.model)
+    vae_scale = pipe.vae_scale_factor
+    gh, gw = args.height // vae_scale // 2, args.width // vae_scale // 2
+
+    # Weight prep goes through a per-rank shard cache (bf16 for the
+    # denoiser/text-encoder, fp32 for the VAE). A cache hit skips the dominant
+    # cost of a cold launch — the host-side extract/shard (~330s at TP=4).
+    denoiser_weights = _cached_weights(
+        cpath("denoiser"),
+        lambda: _build_device_denoiser_weights(pipe, config, tp_size, rank),
+        torch.bfloat16, log, "denoiser weights", cache_enabled)
+    denoiser = QwenImageDenoiser(denoiser_weights, config, (1, gh, gw), batch_size=1)
+
+    te_cfg = _make_text_encoder_config(pipe, tp_size)
+    te_weights = _cached_weights(
+        cpath("text_encoder"),
+        lambda: _build_text_encoder_weights(pipe, te_cfg, tp_size, rank),
+        torch.bfloat16, log, "text-encoder weights", cache_enabled)
+    text_encoder = DeviceTextEncoder(te_weights, te_cfg)
+
+    vae_cfg = get_vae_config(args.model)
+    vae_weights = _cached_weights(
+        cpath("vae"), lambda: _build_vae_weights(pipe),
+        torch.float32, log, "VAE weights", cache_enabled)
+    vae_decoder = DeviceVAEDecoder(vae_weights, vae_cfg)
+
+    dist.barrier()
+    setup_done = time.time()
+
+    if args.prompts_file or args.interactive:
+        # Resident mode: setup is paid once; each image costs only generate().
+        _resident_loop(pipe, denoiser, text_encoder, vae_decoder, config, args,
+                       rank, log)
+    else:
+        log("[qwen-image] generating")
+        dt = _generate_and_save(pipe, denoiser, text_encoder, vae_decoder, config,
+                                args, args.prompt, args.seed, args.output, rank)
+        log(f"[qwen-image] --> {args.steps} steps in {dt:.2f}s")
+        if rank == 0:
+            print(f"[qwen-image] saved {args.output}")
+
+    log(f"[qwen-image] (setup {setup_done - setup_start:.1f}s done before generation)")
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/models/qwen_image/tests/test_pipeline_plumbing.py b/examples/models/qwen_image/tests/test_pipeline_plumbing.py
new file mode 100644
index 0000000..43c535e
--- /dev/null
+++ b/examples/models/qwen_image/tests/test_pipeline_plumbing.py
@@ -0,0 +1,39 @@
+"""Validate the host-driver latent plumbing without the 20B weights.
+
+``_unpack_latents`` in ``qwen_image.py`` must match the diffusers
+``QwenImagePipeline`` static method exactly (it is a copy; this guards against
+drift) and round-trip the pipeline's packed layout back to the original latent.
+
+    cd examples/models/qwen_image
+    python -m pytest tests/test_pipeline_plumbing.py -v
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+torch = pytest.importorskip("torch")
+_diff = pytest.importorskip("diffusers", reason="diffusers with Qwen-Image required")
+from diffusers.pipelines.qwenimage.pipeline_qwenimage import QwenImagePipeline  # noqa: E402
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from qwen_image import _unpack_latents  # noqa: E402
+
+
+def test_unpack_matches_diffusers():
+    B, C, H, W = 1, 4, 8, 8
+    vae_scale = 2  # so unpack recovers H,W from packed
+    latent = torch.randn(B, C, 1, H, W)  # Qwen VAE latents are 5-D (B,C,frame,H,W)
+
+    # the runtime gets already-packed latents from ``pipe.prepare_latents``; source
+    # the packed input the same way (pipeline packs the frame-squeezed latent).
+    packed = QwenImagePipeline._pack_latents(latent[:, :, 0], B, C, H, W)
+
+    # unpack round-trip (channels = C*4, recover to 5-D)
+    height, width = H * vae_scale, W * vae_scale
+    ours_un = _unpack_latents(packed, height, width, vae_scale)
+    ref_un = QwenImagePipeline._unpack_latents(packed, height, width, vae_scale)
+    assert torch.allclose(ours_un, ref_un)
+    # and the packed->unpacked recovers the original latent
+    assert torch.allclose(ours_un, latent, atol=1e-5)
diff --git a/examples/models/qwen_image/tests/test_rope_rmsnorm.py b/examples/models/qwen_image/tests/test_rope_rmsnorm.py
new file mode 100644
index 0000000..73fe355
--- /dev/null
+++ b/examples/models/qwen_image/tests/test_rope_rmsnorm.py
@@ -0,0 +1,101 @@
+"""CPU validation of the M1 primitives (3D RoPE + RMSNorm) vs diffusers.
+
+Run from the example dir with the project venv:
+
+    cd examples/models/qwen_image
+    python -m pytest tests/test_rope_rmsnorm.py -v
+
+Requires ``diffusers`` (>=0.39) and ``torch`` for the reference.
+"""
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+_qwen = pytest.importorskip(
+    "diffusers.models.transformers.transformer_qwenimage",
+    reason="diffusers with Qwen-Image support required",
+)
+from diffusers.models.normalization import RMSNorm  # noqa: E402
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from kernels import rope3d, rmsnorm  # noqa: E402
+
+
+def _rel_l2(a, b):
+    a = a.astype(np.float64)
+    b = b.astype(np.float64)
+    return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12))
+
+
+AXES = (16, 56, 56)
+HEAD_DIM = sum(AXES)  # 128
+THETA = 10000.0
+
+
+@pytest.mark.parametrize("frame,h,w,txt_len", [(1, 8, 8, 16), (1, 16, 12, 32)])
+def test_rope_freqs_match_diffusers(frame, h, w, txt_len):
+    """Our comptime angle tables must reproduce QwenEmbedRope's cos/sin."""
+    rope = _qwen.QwenEmbedRope(theta=int(THETA), axes_dim=list(AXES), scale_rope=True)
+    with torch.no_grad():
+        vid_freqs, txt_freqs = rope(
+            (frame, h, w), device=torch.device("cpu"), max_txt_seq_len=txt_len
+        )
+    # diffusers returns complex freqs (S, head_dim//2); compare cos/sin.
+    ref_vid = vid_freqs.resolve_conj().cpu().numpy()  # complex
+    ref_txt = txt_freqs.resolve_conj().cpu().numpy()
+
+    vid_ang, txt_ang = rope3d.compute_rope_freqs(
+        frame, h, w, txt_len, AXES, theta=THETA, scale_rope=True
+    )
+    assert vid_ang.shape == (frame * h * w, HEAD_DIM // 2)
+    assert txt_ang.shape == (txt_len, HEAD_DIM // 2)
+
+    # angle -> unit complex; compare cos & sin parts
+    assert _rel_l2(np.cos(vid_ang), ref_vid.real) < 1e-5
+    assert _rel_l2(np.sin(vid_ang), ref_vid.imag) < 1e-5
+    assert _rel_l2(np.cos(txt_ang), ref_txt.real) < 1e-5
+    assert _rel_l2(np.sin(txt_ang), ref_txt.imag) < 1e-5
+
+
+@pytest.mark.parametrize("frame,h,w", [(1, 8, 8), (1, 16, 12)])
+def test_apply_rotary_matches_diffusers(frame, h, w):
+    """Rotating q with our real-form apply must match apply_rotary_emb_qwen."""
+    B, H = 2, 3
+    S = frame * h * w
+    rng = np.random.default_rng(0)
+    # our kernel uses (B, H, S, D); diffusers' apply_rotary_emb_qwen expects
+    # (B, S, H, D) here (it unsqueezes freqs at the head axis), so transpose.
+    x = rng.standard_normal((B, H, S, HEAD_DIM)).astype(np.float32)
+
+    rope = _qwen.QwenEmbedRope(theta=int(THETA), axes_dim=list(AXES), scale_rope=True)
+    with torch.no_grad():
+        vid_freqs, _ = rope((frame, h, w), device=torch.device("cpu"), max_txt_seq_len=8)
+        x_bshd = torch.from_numpy(x).permute(0, 2, 1, 3).contiguous()  # (B,S,H,D)
+        ref = _qwen.apply_rotary_emb_qwen(
+            x_bshd, vid_freqs, use_real=False
+        ).permute(0, 2, 1, 3).contiguous().cpu().numpy()  # back to (B,H,S,D)
+
+    vid_ang, _ = rope3d.compute_rope_freqs(frame, h, w, 8, AXES, theta=THETA, scale_rope=True)
+    cos, sin = rope3d.cos_sin_from_angles(vid_ang, dtype=np.float32)
+    out = rope3d.apply_rotary_emb(x, cos, sin)
+
+    assert _rel_l2(out, ref) < 1e-5
+
+
+@pytest.mark.parametrize("dim", [128, 3584])
+def test_rmsnorm_matches_diffusers(dim):
+    rng = np.random.default_rng(1)
+    x = rng.standard_normal((2, 7, dim)).astype(np.float32)
+    g = rng.standard_normal((dim,)).astype(np.float32)
+
+    ref_mod = RMSNorm(dim, eps=1e-6, elementwise_affine=True)
+    with torch.no_grad():
+        ref_mod.weight.copy_(torch.from_numpy(g))
+        ref = ref_mod(torch.from_numpy(x)).cpu().numpy()
+
+    out = rmsnorm.rmsnorm_kernel(x, g, eps=1e-6)
+    assert _rel_l2(out, ref) < 1e-5
diff --git a/examples/models/qwen_image/tests/test_tp.py b/examples/models/qwen_image/tests/test_tp.py
new file mode 100644
index 0000000..0b0ca10
--- /dev/null
+++ b/examples/models/qwen_image/tests/test_tp.py
@@ -0,0 +1,46 @@
+"""TP weight-sharding math (no kernel trace, no hardware).
+
+Validates that the Megatron-style shard slicers partition the weights cleanly:
+column-parallel q shards must concatenate back to the full weight with no
+overlap or gap. End-to-end TP equivalence is validated on device by
+``tests/test_tp_device.py`` (torchrun; skips without hardware).
+
+    cd examples/models/qwen_image
+    python -m pytest tests/test_tp.py -v
+"""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+_diff = pytest.importorskip("diffusers", reason="diffusers with Qwen-Image required")
+from diffusers import QwenImageTransformer2DModel  # noqa: E402
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from weight_extract import extract_flat_weights, shard_flat_weights  # noqa: E402
+
+
+HF = dict(
+    num_layers=2, num_attention_heads=8, attention_head_dim=24,
+    joint_attention_dim=40, in_channels=16, out_channels=4, patch_size=2,
+    axes_dims_rope=(8, 8, 8),
+)
+
+
+def test_column_parallel_shards_reconstruct():
+    """Column-parallel q shards must concat back to the full weight (no overlap)."""
+    torch.manual_seed(0)
+    model = QwenImageTransformer2DModel(**HF).eval()
+    flat = extract_flat_weights(model, HF["num_layers"], dtype=np.float32)
+    for tp_size in (2, 4):
+        shards = [
+            shard_flat_weights(flat, r, tp_size, HF["num_layers"],
+                               HF["num_attention_heads"], HF["attention_head_dim"])
+            for r in range(tp_size)
+        ]
+        recon = np.concatenate([shards[r]["b0_iq_w"] for r in range(tp_size)], axis=1)
+        assert recon.shape == flat["b0_iq_w"].shape
+        assert np.allclose(recon, flat["b0_iq_w"])
diff --git a/examples/models/qwen_image/tests/test_tp_device.py b/examples/models/qwen_image/tests/test_tp_device.py
new file mode 100644
index 0000000..99db961
--- /dev/null
+++ b/examples/models/qwen_image/tests/test_tp_device.py
@@ -0,0 +1,177 @@
+"""Device tensor-parallel validation for the Qwen-Image MMDiT.
+
+Shards a reduced-config denoiser across ``TP`` cores, compiles and runs
+``qwenimage_forward`` with the real Neuron all-reduce collective, and (on rank 0)
+compares against the single-core diffusers baseline. This exercises the actual
+multi-core path — sharded weights + collectives — without the 20B download.
+
+    cd examples/models/qwen_image
+    torchrun --nproc-per-node 4 tests/test_tp_device.py --num-layers 2 --num-heads 8
+
+The reduced config keeps head_dim=128 (real value) but shrinks heads / text dim /
+blocks so it fits and compiles fast; num_heads must be divisible by TP.
+
+Needs Neuron hardware + multiple cores, so the collected ``pytest`` entry
+(`test_tp_device`) **skips by default** on a normal CPU test run. Opt in with
+``QWEN_IMAGE_TP_DEVICE_TEST=1`` (optionally ``QWEN_IMAGE_TP_SIZE=<n>``) and it
+shells out to the ``torchrun`` invocation above.
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+# runnable from tests/: put the example root (config, kernels, weight_extract) on the path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+
+def _rel_l2(a, b):
+    a = a.astype(np.float64)
+    b = b.astype(np.float64)
+    return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-layers", type=int, default=2)
+    parser.add_argument("--num-heads", type=int, default=8)
+    parser.add_argument("--head-dim", type=int, default=128)
+    parser.add_argument("--joint-dim", type=int, default=64)
+    parser.add_argument("--grid", type=int, default=8)
+    parser.add_argument("--txt-len", type=int, default=16)
+    parser.add_argument("--tol", type=float, default=0.05)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+
+    # must be set before importing nkipy: keep the NRT per-execution barrier on
+    # for correct multi-rank collectives (SPMD-with-collectives).
+    os.environ["NEURON_RT_DISABLE_EXECUTION_BARRIER"] = "0"
+    os.environ.setdefault("NEURON_RT_ROOT_COMM_ID", "localhost:61375")
+    dist.init_process_group()
+    rank = dist.get_rank()
+    tp_size = dist.get_world_size()
+    os.environ["NEURON_RT_VISIBLE_CORES"] = str(rank)
+    torch.set_num_threads(1)
+
+    assert args.num_heads % tp_size == 0, "num_heads must be divisible by TP size"
+
+    from config import Config
+    from diffusers import QwenImageTransformer2DModel
+    from kernels.tp import make_all_reduce
+    from kernels.transformer import qwenimage_forward
+    from kernels.weight_layout import BLOCK_KEYS, SHARED_KEYS, block_key
+    from nkipy.runtime import DeviceKernel, DeviceTensor
+    from nkipy.runtime.device_tensor import bfloat16
+    from weight_extract import extract_flat_weights, shard_flat_weights
+
+    d = args.head_dim
+    a0 = (d // 3) & ~1
+    a1 = (d // 3) & ~1
+    a2 = d - a0 - a1
+    hf = dict(
+        num_layers=args.num_layers, num_attention_heads=args.num_heads,
+        attention_head_dim=args.head_dim, joint_attention_dim=args.joint_dim,
+        in_channels=16, out_channels=4, patch_size=2, axes_dims_rope=(a0, a1, a2),
+    )
+
+    torch.manual_seed(args.seed)
+    model = QwenImageTransformer2DModel(**hf).eval()
+
+    frame, gh, gw = 1, args.grid, args.grid
+    B, Limg = 2, frame * gh * gw
+    rng = np.random.default_rng(args.seed)
+    latent = rng.standard_normal((B, Limg, hf["in_channels"])).astype(np.float32)
+    text = rng.standard_normal((B, args.txt_len, hf["joint_attention_dim"])).astype(np.float32)
+    timestep = rng.uniform(0, 1000, size=(B,)).astype(np.float32)
+
+    # single-core reference (rank 0 only)
+    ref = None
+    if rank == 0:
+        with torch.no_grad():
+            ref = model(
+                hidden_states=torch.from_numpy(latent),
+                encoder_hidden_states=torch.from_numpy(text),
+                timestep=torch.from_numpy(timestep),
+                img_shapes=[(frame, gh, gw)] * B, return_dict=True,
+            ).sample.cpu().numpy()
+
+    # this rank's shard
+    full = extract_flat_weights(model, hf["num_layers"], dtype=bfloat16)
+    shard = shard_flat_weights(full, rank, tp_size, hf["num_layers"],
+                               hf["num_attention_heads"], hf["attention_head_dim"])
+
+    cfg = Config(
+        num_layers=hf["num_layers"], num_heads=hf["num_attention_heads"],
+        head_dim=hf["attention_head_dim"], joint_attention_dim=hf["joint_attention_dim"],
+        in_channels=hf["in_channels"], out_channels=hf["out_channels"],
+        patch_size=hf["patch_size"], axes_dims_rope=hf["axes_dims_rope"],
+        dtype=bfloat16, tp_size=tp_size, all_reduce_fn=make_all_reduce(tp_size),
+    )
+
+    present = [k for k in SHARED_KEYS if k in shard]
+    for i in range(cfg.num_layers):
+        present += [block_key(i, s) for s in BLOCK_KEYS if block_key(i, s) in shard]
+    device_weights = {k: DeviceTensor.from_numpy(shard[k], name=k) for k in present}
+
+    latent_d = DeviceTensor.from_numpy(latent.astype(bfloat16), "latent")
+    text_d = DeviceTensor.from_numpy(text.astype(bfloat16), "text")
+    timestep_d = DeviceTensor.from_numpy(timestep.astype(np.float32), "timestep")
+
+    if rank == 0:
+        print(f"[qwen-image-tp] compiling (TP={tp_size}, {cfg.num_layers} blocks, "
+              f"{cfg.num_heads}x{cfg.head_dim} heads, grid {gh}x{gw})")
+    kernel = DeviceKernel.compile_and_load(
+        qwenimage_forward, name="qwenimage_forward_tp",
+        latent=latent_d, text=text_d, timestep=timestep_d,
+        img_shape=(frame, gh, gw), configs=cfg, build_dir="./build_tp",
+        additional_compiler_args=cfg.additional_compiler_args_nkipy,
+        **device_weights,
+    )
+    out_dim = cfg.patch_size * cfg.patch_size * cfg.out_channels
+    out_d = DeviceTensor.from_numpy(np.empty((B, Limg, out_dim), dtype=bfloat16), "out")
+    expected = set(kernel.input_tensors_info)
+    inputs = {"latent": latent_d, "text": text_d, "timestep": timestep_d}
+    inputs.update({k: v for k, v in device_weights.items() if k in expected})
+    kernel(inputs=inputs, outputs={"output0": out_d})
+    dev = out_d.torch().to(torch.float32).numpy()
+
+    dist.barrier()
+    if rank == 0:
+        rel = _rel_l2(dev, ref)
+        print(f"[qwen-image-tp] TP={tp_size} device vs diffusers: rel_l2={rel:.4e}")
+        print("[qwen-image-tp] PASS" if rel < args.tol else "[qwen-image-tp] FAIL")
+    dist.destroy_process_group()
+
+
+def test_tp_device():
+    """Collected pytest entry: skips unless opted in (needs Neuron hardware +
+    multiple cores), then drives the torchrun harness above and checks it PASSes.
+
+    A single pytest process can't be the N ranks, so we shell out to torchrun on
+    this same file — its ``__main__`` block runs ``main()`` per rank.
+    """
+    import subprocess
+
+    import pytest
+
+    if os.environ.get("QWEN_IMAGE_TP_DEVICE_TEST") != "1":
+        pytest.skip("device TP test needs Neuron hardware; opt in with "
+                    "QWEN_IMAGE_TP_DEVICE_TEST=1")
+    root = Path(__file__).resolve().parents[1]
+    nproc = os.environ.get("QWEN_IMAGE_TP_SIZE", "4")
+    proc = subprocess.run(
+        ["torchrun", "--nproc-per-node", nproc, str(Path(__file__)),
+         "--num-layers", "2", "--num-heads", "8"],
+        cwd=str(root), capture_output=True, text=True)
+    print(proc.stdout)
+    print(proc.stderr)
+    assert "[qwen-image-tp] PASS" in proc.stdout, proc.stderr or proc.stdout
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/models/qwen_image/tests/test_vae.py b/examples/models/qwen_image/tests/test_vae.py
new file mode 100644
index 0000000..7500db9
--- /dev/null
+++ b/examples/models/qwen_image/tests/test_vae.py
@@ -0,0 +1,73 @@
+"""CPU validation of the Qwen-Image VAE decoder (T=1 2D collapse) vs diffusers.
+
+The device VAE (``kernels/vae.py``) decodes a single image (T=1), where the 3D
+causal *video* VAE collapses exactly to a 2D conv decoder. This test builds a
+small ``AutoencoderKLQwenImage``, decodes a fixed latent through both diffusers
+and our numpy kernel, and checks rel_l2 < 1e-3.
+
+    cd examples/models/qwen_image
+    python -m pytest tests/test_vae.py -v
+"""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+_diff = pytest.importorskip("diffusers", reason="diffusers with Qwen-Image required")
+from diffusers.models.autoencoders.autoencoder_kl_qwenimage import (  # noqa: E402
+    AutoencoderKLQwenImage,
+)
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from config import VAEConfig  # noqa: E402
+from kernels.vae import vae_decode  # noqa: E402
+from weight_extract import extract_vae_decoder_weights  # noqa: E402
+
+
+def _rel_l2(a, b):
+    a = a.astype(np.float64)
+    b = b.astype(np.float64)
+    return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12))
+
+
+def _build_small_vae(base_dim, z_dim, dim_mult, num_res_blocks):
+    torch.manual_seed(0)
+    return AutoencoderKLQwenImage(
+        base_dim=base_dim, z_dim=z_dim, dim_mult=list(dim_mult),
+        num_res_blocks=num_res_blocks, temperal_downsample=[False, True, True],
+    ).eval()
+
+
+@pytest.mark.parametrize("base_dim,dim_mult", [(8, (1, 2, 4, 4)), (12, (1, 2, 2, 4))])
+def test_vae_decode_matches_diffusers(base_dim, dim_mult):
+    z_dim, nrb = 16, 2
+    vae = _build_small_vae(base_dim, z_dim, dim_mult, nrb)
+
+    rng = np.random.default_rng(0)
+    lat = rng.standard_normal((1, z_dim, 1, 4, 4)).astype(np.float32)  # (B,C,T,H,W)
+
+    with torch.no_grad():
+        ref = vae.decode(torch.from_numpy(lat), return_dict=False)[0].cpu().numpy()
+    ref = ref[:, :, 0]  # squeeze the single frame -> (1, 3, H, W)
+
+    cfg = VAEConfig(z_dim=z_dim, base_dim=base_dim, dim_mult=dim_mult,
+                    num_res_blocks=nrb)
+    flat = extract_vae_decoder_weights(vae, dtype=np.float32)
+    out = vae_decode(lat[:, :, 0], cfg, **flat)  # host squeezes T; kernel is 2D
+
+    assert out.shape == ref.shape, (out.shape, ref.shape)
+    rel = _rel_l2(out, ref)
+    assert rel < 1e-3, f"rel_l2={rel:.3e}"
+
+
+def test_vae_8x_upscale_shape():
+    """Full Qwen VAE spatial ratio is 8x (dim_mult of length 4 -> 3 upsamples)."""
+    vae = _build_small_vae(8, 16, (1, 2, 4, 4), 2)
+    cfg = VAEConfig(z_dim=16, base_dim=8, dim_mult=(1, 2, 4, 4), num_res_blocks=2)
+    flat = extract_vae_decoder_weights(vae, dtype=np.float32)
+    lat = np.zeros((1, 16, 5, 5), dtype=np.float32)
+    out = vae_decode(lat, cfg, **flat)
+    assert out.shape == (1, 3, 40, 40)  # 5 * 8
diff --git a/examples/models/qwen_image/weight_extract.py b/examples/models/qwen_image/weight_extract.py
new file mode 100644
index 0000000..3a7c0a8
--- /dev/null
+++ b/examples/models/qwen_image/weight_extract.py
@@ -0,0 +1,330 @@
+"""Map a diffusers ``QwenImageTransformer2DModel`` state dict to the flat
+canonical key scheme (``kernels/weight_layout.py``).
+
+Torch ``nn.Linear`` weights are stored as (out, in); our numpy kernels use the
+``x @ W`` convention, so linear weights are transposed to (in, out). RMSNorm /
+LayerNorm gains and all biases pass through unchanged.
+
+Kept dependency-light (numpy/torch only) so it can run against an in-memory
+model — the driver, the reduced-config tests, and ``tests/test_tp_device.py``
+all call these functions directly on a loaded diffusers module.
+"""
+
+import numpy as np
+import torch
+
+from kernels.tp import shard_heads, shard_intermediate
+from kernels.weight_layout import BLOCK_KEYS, SHARED_KEYS, block_key
+
+
+def _np(p):
+    """torch tensor -> numpy, routing bf16 through fp32 (torch can't .numpy() bf16)."""
+    p = p.detach().cpu()
+    if p.dtype == torch.bfloat16:
+        p = p.to(torch.float32)
+    return p.numpy()
+
+
+def _t(p):
+    """torch Linear weight (out,in) -> numpy (in,out)."""
+    return _np(p).T.copy()
+
+
+def _v(p):
+    return _np(p).copy()
+
+
+def _block_weights(sd, i):
+    """Extract one block's weights (canonical short keys) from a state dict.
+
+    ``sd`` keys are the ``transformer_blocks.{i}.`` submodule names.
+    """
+    def g(name):
+        return sd[f"transformer_blocks.{i}.{name}"]
+
+    return {
+        "img_mod_w": _t(g("img_mod.1.weight")), "img_mod_b": _v(g("img_mod.1.bias")),
+        "txt_mod_w": _t(g("txt_mod.1.weight")), "txt_mod_b": _v(g("txt_mod.1.bias")),
+        "iq_w": _t(g("attn.to_q.weight")), "iq_b": _v(g("attn.to_q.bias")),
+        "ik_w": _t(g("attn.to_k.weight")), "ik_b": _v(g("attn.to_k.bias")),
+        "iv_w": _t(g("attn.to_v.weight")), "iv_b": _v(g("attn.to_v.bias")),
+        "io_w": _t(g("attn.to_out.0.weight")), "io_b": _v(g("attn.to_out.0.bias")),
+        "iq_g": _v(g("attn.norm_q.weight")), "ik_g": _v(g("attn.norm_k.weight")),
+        "tq_w": _t(g("attn.add_q_proj.weight")), "tq_b": _v(g("attn.add_q_proj.bias")),
+        "tk_w": _t(g("attn.add_k_proj.weight")), "tk_b": _v(g("attn.add_k_proj.bias")),
+        "tv_w": _t(g("attn.add_v_proj.weight")), "tv_b": _v(g("attn.add_v_proj.bias")),
+        "to_w": _t(g("attn.to_add_out.weight")), "to_b": _v(g("attn.to_add_out.bias")),
+        "tq_g": _v(g("attn.norm_added_q.weight")), "tk_g": _v(g("attn.norm_added_k.weight")),
+        "iff0_w": _t(g("img_mlp.net.0.proj.weight")), "iff0_b": _v(g("img_mlp.net.0.proj.bias")),
+        "iff2_w": _t(g("img_mlp.net.2.weight")), "iff2_b": _v(g("img_mlp.net.2.bias")),
+        "tff0_w": _t(g("txt_mlp.net.0.proj.weight")), "tff0_b": _v(g("txt_mlp.net.0.proj.bias")),
+        "tff2_w": _t(g("txt_mlp.net.2.weight")), "tff2_b": _v(g("txt_mlp.net.2.bias")),
+    }
+
+
+def extract_flat_weights(model, num_layers, dtype=np.float32):
+    """Return a flat ``{key: ndarray}`` dict for ``qwenimage_forward``.
+
+    Args:
+        model: a diffusers ``QwenImageTransformer2DModel`` (or any module whose
+            ``named_parameters`` follow the same naming).
+        num_layers: number of transformer blocks to extract.
+    """
+    sd = dict(model.named_parameters())
+
+    flat = {
+        "img_in.weight": _t(sd["img_in.weight"]), "img_in.bias": _v(sd["img_in.bias"]),
+        "txt_norm.weight": _v(sd["txt_norm.weight"]),
+        "txt_in.weight": _t(sd["txt_in.weight"]), "txt_in.bias": _v(sd["txt_in.bias"]),
+        "time.proj.weight": _t(sd["time_text_embed.timestep_embedder.linear_1.weight"]),
+        "time.proj.bias": _v(sd["time_text_embed.timestep_embedder.linear_1.bias"]),
+        "time.emb.weight": _t(sd["time_text_embed.timestep_embedder.linear_2.weight"]),
+        "time.emb.bias": _v(sd["time_text_embed.timestep_embedder.linear_2.bias"]),
+        "norm_out.linear.weight": _t(sd["norm_out.linear.weight"]),
+        "norm_out.linear.bias": _v(sd["norm_out.linear.bias"]),
+        "proj_out.weight": _t(sd["proj_out.weight"]), "proj_out.bias": _v(sd["proj_out.bias"]),
+    }
+
+    for i in range(num_layers):
+        for short, val in _block_weights(sd, i).items():
+            flat[block_key(i, short)] = val
+
+    if dtype is not None:
+        flat = {k: v.astype(dtype) for k, v in flat.items()}
+    return flat
+
+
+# per-block key -> (kind, axis) sharding rule. Column-parallel (axis 1) shards
+# the output dim; row-parallel (axis 0) shards the input dim. Keys not listed
+# (modulation, QK-norm gains) are replicated on every rank.
+_HEAD_COL = ("head", 1)   # q/k/v: shard output (head) dim + bias
+_HEAD_ROW = ("head", 0)   # o: shard input (head) dim; bias replicated
+_INT_COL = ("inter", 1)   # ff0: shard intermediate dim + bias
+_INT_ROW = ("inter", 0)   # ff2: shard intermediate (input) dim; bias replicated
+
+_SHARD_RULES = {
+    "iq_w": _HEAD_COL, "ik_w": _HEAD_COL, "iv_w": _HEAD_COL, "io_w": _HEAD_ROW,
+    "tq_w": _HEAD_COL, "tk_w": _HEAD_COL, "tv_w": _HEAD_COL, "to_w": _HEAD_ROW,
+    "iff0_w": _INT_COL, "iff2_w": _INT_ROW,
+    "tff0_w": _INT_COL, "tff2_w": _INT_ROW,
+}
+# bias that rides along with a column-parallel weight (sharded the same way)
+_BIAS_OF = {
+    "iq_w": "iq_b", "ik_w": "ik_b", "iv_w": "iv_b", "io_w": "io_b",
+    "tq_w": "tq_b", "tk_w": "tk_b", "tv_w": "tv_b", "to_w": "to_b",
+    "iff0_w": "iff0_b", "iff2_w": "iff2_b", "tff0_w": "tff0_b", "tff2_w": "tff2_b",
+}
+
+
+def shard_flat_weights(flat, rank, tp_size, num_layers, n_heads, head_dim):
+    """Slice a full flat weight dict into ``rank``'s tensor-parallel shard.
+
+    Attention q/k/v/o shard by heads, MLP ff0/ff2 by intermediate; everything
+    else (modulation, norms, QK-norm gains, shared input/output projections) is
+    replicated. Row-parallel biases (o, ff2) stay full — they are added once
+    after the all-reduce (see the kernels). Returns a new flat dict.
+    """
+    if tp_size <= 1:
+        return dict(flat)
+
+    out = {}
+    # shared weights are replicated
+    for k in SHARED_KEYS:
+        if k in flat:
+            out[k] = flat[k]
+
+    for i in range(num_layers):
+        for short in BLOCK_KEYS:
+            key = block_key(i, short)
+            if key not in flat:
+                continue
+            rule = _SHARD_RULES.get(short)
+            if rule is None:
+                # replicated (modulation weights/biases, QK-norm gains, and the
+                # biases already handled alongside their weight)
+                if short not in _BIAS_OF.values():
+                    out[key] = flat[key]
+                continue
+            kind, axis = rule
+            w = flat[key]
+            b_short = _BIAS_OF[short]
+            b = flat.get(block_key(i, b_short))
+            if kind == "head":
+                sw, sb = shard_heads(w, b, rank, tp_size, n_heads, head_dim, axis)
+            else:
+                sw, sb = shard_intermediate(w, b, rank, tp_size, axis)
+            out[key] = sw
+            if axis == 1:
+                # column-parallel: bias sharded with the weight
+                out[block_key(i, b_short)] = sb
+            elif b is not None:
+                # row-parallel: keep the full replicated bias
+                out[block_key(i, b_short)] = b
+    return out
+
+
+# ── text encoder (Qwen2.5 LM) ────────────────────────────────────────────────
+
+def extract_text_encoder_weights(lm, num_layers, dtype=np.float32):
+    """Flatten a Qwen2.5 text-model (``te.model.language_model``) to flat keys.
+
+    ``lm`` is the ``Qwen2_5_VLTextModel``; we take its decoder layers + final
+    norm. The embedding table / LM head are left on host.
+    """
+    from kernels.text_weight_layout import layer_key
+
+    sd = dict(lm.named_parameters())
+    out = {"final_norm": _v(sd["norm.weight"])}
+    for i in range(num_layers):
+        p = f"layers.{i}."
+        out[layer_key(i, "attn_norm")] = _v(sd[p + "input_layernorm.weight"])
+        out[layer_key(i, "mlp_norm")] = _v(sd[p + "post_attention_layernorm.weight"])
+        out[layer_key(i, "q_w")] = _t(sd[p + "self_attn.q_proj.weight"])
+        out[layer_key(i, "q_b")] = _v(sd[p + "self_attn.q_proj.bias"])
+        out[layer_key(i, "k_w")] = _t(sd[p + "self_attn.k_proj.weight"])
+        out[layer_key(i, "k_b")] = _v(sd[p + "self_attn.k_proj.bias"])
+        out[layer_key(i, "v_w")] = _t(sd[p + "self_attn.v_proj.weight"])
+        out[layer_key(i, "v_b")] = _v(sd[p + "self_attn.v_proj.bias"])
+        out[layer_key(i, "o_w")] = _t(sd[p + "self_attn.o_proj.weight"])
+        out[layer_key(i, "gate_w")] = _t(sd[p + "mlp.gate_proj.weight"])
+        out[layer_key(i, "up_w")] = _t(sd[p + "mlp.up_proj.weight"])
+        out[layer_key(i, "down_w")] = _t(sd[p + "mlp.down_proj.weight"])
+    if dtype is not None:
+        out = {k: v.astype(dtype) for k, v in out.items()}
+    return out
+
+
+# ── VAE decoder (AutoencoderKLQwenImage, T=1 2D collapse) ────────────────────
+
+def _conv3d_to_2d(p):
+    """Collapse a causal-conv3d weight (out, in, kt, kh, kw) to 2D (out, in, kh,
+    kw) by taking the last temporal tap. At T=1 the causal left-padding zeros the
+    earlier taps, so only ``w[:, :, -1]`` contributes (verified 0.0 diff)."""
+    return _np(p)[:, :, -1, :, :].copy()
+
+
+def _conv1x1_to_linear(p):
+    """A 1x1 ``nn.Conv2d`` weight (out, in, 1, 1) -> (in, out) matmul weight."""
+    return _np(p)[:, :, 0, 0].T.copy()
+
+
+def extract_vae_decoder_weights(vae, dtype=np.float32):
+    """Flatten an ``AutoencoderKLQwenImage`` decoder to the flat 2D key scheme.
+
+    Collapses every 3D causal conv to a 2D conv (last temporal tap); the
+    attention 1x1 convs become (in, out) matmul weights. ``post_quant_conv`` (a
+    1x1 conv applied to the latent before the decoder proper) is included so the
+    kernel input is the raw (denormalized) latent. The encoder / quant_conv are
+    unused for decoding.
+    """
+    sd = dict(vae.named_parameters())
+    dec = "decoder."
+    cfg_mult = vae.config.dim_mult
+    nrb = vae.config.num_res_blocks
+
+    out = {}
+
+    def conv3(dst, src):
+        out[dst + ".weight"] = _conv3d_to_2d(sd[src + ".weight"])
+        out[dst + ".bias"] = _v(sd[src + ".bias"])
+
+    # post_quant_conv (1x1x1 causal-conv3d) applied to the latent before the
+    # decoder proper -> 2D 1x1 conv
+    out["post_quant.weight"] = _conv3d_to_2d(sd["post_quant_conv.weight"])
+    out["post_quant.bias"] = _v(sd["post_quant_conv.bias"])
+
+    conv3("conv_in", dec + "conv_in")
+    conv3("conv_out", dec + "conv_out")
+    out["norm_out.gamma"] = _v(sd[dec + "norm_out.gamma"])
+
+    # mid block: resnets 0/1 + attention 0
+    for j in (0, 1):
+        _extract_resnet(sd, out, f"mid.resnets.{j}.", dec + f"mid_block.resnets.{j}.")
+    ap = dec + "mid_block.attentions.0."
+    out["mid.attentions.0.norm.gamma"] = _v(sd[ap + "norm.gamma"])
+    out["mid.attentions.0.to_qkv.weight"] = _conv1x1_to_linear(sd[ap + "to_qkv.weight"])
+    out["mid.attentions.0.to_qkv.bias"] = _v(sd[ap + "to_qkv.bias"])
+    out["mid.attentions.0.proj.weight"] = _conv1x1_to_linear(sd[ap + "proj.weight"])
+    out["mid.attentions.0.proj.bias"] = _v(sd[ap + "proj.bias"])
+
+    # up blocks
+    n_up = len(cfg_mult)
+    for i in range(n_up):
+        up = dec + f"up_blocks.{i}."
+        for j in range(nrb + 1):
+            _extract_resnet(sd, out, f"up.{i}.resnets.{j}.", up + f"resnets.{j}.")
+        # spatial upsampler conv (resample.1); time_conv is skipped at T=1
+        usw = up + "upsamplers.0.resample.1.weight"
+        if usw in sd:
+            out[f"up.{i}.upsample.weight"] = _np(sd[usw]).copy()  # already 2D conv
+            out[f"up.{i}.upsample.bias"] = _v(sd[up + "upsamplers.0.resample.1.bias"])
+
+    if dtype is not None:
+        out = {k: v.astype(dtype) for k, v in out.items()}
+    return out
+
+
+def _extract_resnet(sd, out, dst, src):
+    """One QwenImageResidualBlock -> flat 2D keys (norm gammas, conv1/conv2, and
+    an optional 1x1 conv_shortcut)."""
+    out[dst + "norm1.gamma"] = _v(sd[src + "norm1.gamma"])
+    out[dst + "norm2.gamma"] = _v(sd[src + "norm2.gamma"])
+    out[dst + "conv1.weight"] = _conv3d_to_2d(sd[src + "conv1.weight"])
+    out[dst + "conv1.bias"] = _v(sd[src + "conv1.bias"])
+    out[dst + "conv2.weight"] = _conv3d_to_2d(sd[src + "conv2.weight"])
+    out[dst + "conv2.bias"] = _v(sd[src + "conv2.bias"])
+    if src + "conv_shortcut.weight" in sd:
+        # shortcut is a 1x1x1 causal-conv3d -> 2D 1x1 conv (out, in, 1, 1)
+        out[dst + "conv_shortcut.weight"] = _conv3d_to_2d(sd[src + "conv_shortcut.weight"])
+        out[dst + "conv_shortcut.bias"] = _v(sd[src + "conv_shortcut.bias"])
+
+
+def shard_text_encoder_weights(flat, rank, tp_size, num_layers, n_heads,
+                               n_kv_heads, head_dim):
+    """Slice the flat text-encoder weights into ``rank``'s TP shard.
+
+    Megatron-style, mirroring the denoiser (``shard_flat_weights``): attention
+    q/o shard by query heads, k/v shard by KV heads (GQA), and the SwiGLU MLP
+    shards by intermediate (gate/up column-parallel, down row-parallel). q/k/v
+    biases ride with their column-parallel weight; the row-parallel o/down have
+    no bias in Qwen2.5. RMSNorm gains (attn_norm/mlp_norm/final_norm) replicate.
+
+    Requires ``tp_size`` to divide both ``n_heads`` and ``n_kv_heads`` (so
+    ``tp_size <= n_kv_heads``; the 4-KV-head encoder shards up to TP=4).
+    """
+    from kernels.text_weight_layout import SHARED_KEYS, layer_key
+    from kernels.tp import shard_heads, shard_intermediate
+
+    if tp_size <= 1:
+        return dict(flat)
+    if n_heads % tp_size or n_kv_heads % tp_size:
+        raise ValueError(
+            f"text-encoder TP={tp_size} must divide n_heads={n_heads} and "
+            f"n_kv_heads={n_kv_heads} (KV heads cap TP at {n_kv_heads})")
+
+    out = {k: flat[k] for k in SHARED_KEYS if k in flat}
+    for i in range(num_layers):
+        # replicated norms
+        for s in ("attn_norm", "mlp_norm"):
+            out[layer_key(i, s)] = flat[layer_key(i, s)]
+        # attention: q/o by query heads, k/v by KV heads
+        for w_short, b_short, nh in (("q_w", "q_b", n_heads),
+                                     ("k_w", "k_b", n_kv_heads),
+                                     ("v_w", "v_b", n_kv_heads)):
+            sw, sb = shard_heads(flat[layer_key(i, w_short)],
+                                 flat[layer_key(i, b_short)],
+                                 rank, tp_size, nh, head_dim, axis=1)
+            out[layer_key(i, w_short)] = sw
+            out[layer_key(i, b_short)] = sb
+        ow, _ = shard_heads(flat[layer_key(i, "o_w")], None,
+                            rank, tp_size, n_heads, head_dim, axis=0)
+        out[layer_key(i, "o_w")] = ow
+        # SwiGLU MLP: gate/up column-parallel, down row-parallel
+        for w_short in ("gate_w", "up_w"):
+            sw, _ = shard_intermediate(flat[layer_key(i, w_short)], None,
+                                       rank, tp_size, axis=1)
+            out[layer_key(i, w_short)] = sw
+        dw, _ = shard_intermediate(flat[layer_key(i, "down_w")], None,
+                                   rank, tp_size, axis=0)
+        out[layer_key(i, "down_w")] = dw
+    return out