host/host_contract.md

# Host contract — RGBA-Image-2.1 (Core AI port of Qwen-Image-2.1)

Everything a host computes around the three graphs, for text-to-image, batch 1, no CFG.
The Python reference of every step below is in the zoo's
[`conversion/qwenimage21/`](https://github.com/john-rocky/coreai-model-zoo/tree/main/conversion/qwenimage21):
`qi21_tokenize.py` (tokenizer), `qi21_host.py` (RoPE, pack/unpack), `qi21_sched.py` (sampler),
`pipeline_engine.py` (the whole loop on the three bundles). The first three match the diffusers-main
pipeline (`Qwen/Qwen-Image-2.1` @ `790c926`) exactly: the same token ids, bit-identical RoPE tables,
bit-identical sigmas and Euler steps. `pipeline_engine.py` is scored against the pipeline's fp32 images.

## 1. The graphs

Every graph has one function, `main`. All tensors crossing a graph boundary are fp32, except
`input_ids` (int32).

| bundle | inputs | output | axes |
| --- | --- | --- | --- |
| `qi21_encoder_dynL_w16a32_ids_iofp32.aimodel` | `input_ids [1,Lfull]` int32 | `hidden [1,Lfull,4096]` | `Lfull` 16..512 |
| `qi21_dit_full_bf16_dyn_iofp32.aimodel` | `img_tokens [1,N,64]`, `txt_feats [1,L,4096]`, `timestep [1]`, `txt_cos [1,L,64]`, `txt_sin [1,L,64]`, `img_cos [1,N,64]`, `img_sin [1,N,64]` | `vel [1,N,64]` | `L` 8..512, `N` 64..4096 |
| `qi21_vae_{256,512,1024}_fp32.aimodel` | `latents_packed [1,N,64]` | `image [1,4,S,S]` | fixed: `N = (S/16)²` |

- Encoder: the Qwen3-VL-8B text stack (36 layers). bf16 weights, fp32 compute. `embed_tokens` is
  inside the graph. The output is the residual stream after the last layer, **before** the final
  RMSNorm. Do not apply a norm on the host.
- DiT: 32 blocks, bf16 weights and compute, fp32 boundary.
- VAE: decoder only, fp32. The graph unpacks the tokens and applies `latents * std + mean` itself.

## 2. Tokenize

Tokenizer: `tokenizer/tokenizer.json` (Qwen2 BPE, byte-level). No BOS token. An empty prompt is
replaced by a single space `" "`.

Template (text-to-image):

```
<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n
```

`\n` is a newline character. Encode the whole string once, with no padding and no truncation, to get
`ids` (length `Lfull`).

**`drop_idx`** is the number of tokens of the system part alone,
`<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n`. Compute it by encoding
that string. With this tokenizer it is 14:
`[151644, 8948, 198, 1092, 30782, 408, 323, 23643, 279, 3897, 9934, 13, 151645, 198]`.
Check that `ids[:drop_idx]` equals those tokens.

`Lfull` must be 16..512 (the encoder's axis), and the DiT needs `L = Lfull − drop_idx` in 8..512.
The template around the empty-prompt substitute `" "` is 23 tokens (`L` = 9).

## 3. Encode

```
hidden        = encoder(input_ids = ids as int32 [1, Lfull])      # [1, Lfull, 4096]
prompt_embeds = hidden[:, drop_idx : Lfull]                       # [1, L, 4096], L = Lfull - drop_idx
```

Pass exactly `Lfull` tokens; the axis is dynamic, so no padding is needed. (Attention is causal, so
pad tokens never reach the first `Lfull` outputs, but a longer input moves the GPU result by about
4e-6 relative.) Run the encoder once per image.

## 4. Sequence layout and RoPE

The DiT's sequence is `[text L | image N]`: the `L` prompt tokens, then the image tokens in raster
order. One image token covers one 16×16 px tile of the output: for an `S×S` image,
`h = w = S/16` and `N = h·w` (256² → 256 tokens, 512² → 1024, 1024² → 4096). There is no 2×2 packing.

Each token has a 3-axis position `(frame, height, width)`:

- text token `i` (0-based) → `(i, i, i)`;
- image token at row `y`, column `x` (0-based, token index `y·w + x`) →
  `(L, y − (h − h//2), x − (w − w//2))`. The grid is centred on zero. For `h = 16`, rows go −8..7.

`host/rope_axis{a}_{cos,sin}.f32` hold one row per position −1024..8191: fp32, little-endian,
row-major `[9216, P_a]` with `P = (8, 28, 28)`. The row for position `p` is `p + 1024`. A token's
64 values are the three rows concatenated in axis order:

```
cos[token] = axis0_cos[f + 1024] ++ axis1_cos[hpos + 1024] ++ axis2_cos[wpos + 1024]   # 8 + 28 + 28
sin[token] = the same with the _sin tables
txt_cos, txt_sin = rows of the L text tokens     # [1, L, 64]
img_cos, img_sin = rows of the N image tokens    # [1, N, 64]
```

The tables are `torch.polar(1, outer(pos, 1 / 10000^(arange(0, d, 2)/d)))` for axis widths
`d = 16, 56, 56`, the same numbers as `QwenImage21Rope.freqs` of diffusers main (bit-exact). Axes 1
and 2 have the same width, so their files are byte-identical. The four tensors depend only on `L`,
`h` and `w`: build them once per image, not per step.

## 5. Noise and pack

The initial latent is standard normal noise: 64 channels × `h` rows × `w` columns, channel-major.
The reference pipeline draws it as `randn((1, 1, 64, h, w))` and packs it into DiT tokens:

```
x[0, y·w + x_, c] = z[c, y, x_]          # = z.view(1, 64, h·w).transpose(1, 2)
```

Any N(0, 1) noise generates an image. Matching the Python engine for a given seed needs its exact
draw: `torch.randn((1, 1, 64, h, w), generator=torch.Generator("cpu").manual_seed(seed))`.

## 6. Sampler (FlowMatch Euler, 40 steps, no CFG)

Constants in `host/scheduler.json` (from the checkpoint's `scheduler_config.json`). All arrays fp32.

```
sigmas   = linspace(1, 1/steps, steps)                           # steps = 40
mu       = N · m + b,  m = (max_shift − base_shift) / (max_image_seq_len − base_image_seq_len),
                       b = base_shift − m · base_image_seq_len    # N = image tokens
sigmas   = e^mu / (e^mu + (1/sigmas − 1))                        # time_shift_type "exponential"
sigmas   = 1 − (1 − sigmas) / ((1 − sigmas[-1]) / (1 − shift_terminal))
timesteps = sigmas · num_train_timesteps                          # fp32
sigmas   = sigmas ++ [0]                                         # 41 values
for i in 0 ..< steps:
    t   = timesteps[i] / 1000                                    # fp32; this is the DiT `timestep` input
    vel = dit(img_tokens = x, txt_feats = prompt_embeds, timestep = [t], txt_cos, txt_sin, img_cos, img_sin)
    x   = x + (sigmas[i+1] − sigmas[i]) · vel                    # fp32
```

`mu` is 0.5 at 256² (N = 256), 0.5387… at 512², 0.6935… at 1024². Keep `t = timesteps[i] / 1000` as
written (multiply by 1000, then divide) to match the reference bit for bit. `qi21_sched.py`
reproduces the reference sigmas, timesteps and every Euler step bit-exactly at 256² and 512².

## 7. Decode

```
image = vae_S(latents_packed = x)          # x after the last step, [1, N, 64], exactly as the sampler holds it
rgba8 = round(clip(image · 0.5 + 0.5, 0, 1) · 255)   # [1, 4, S, S] -> channels R, G, B, A
```

Feed the sampler's latent unchanged: no unpacking and no `* std + mean` on the host (the graph does
both). The four channels are what the reference pipeline saves as an RGBA PNG.

## 8. Compile before running on the GPU

On macOS 27.0 (26A428) the Python runtime crashes on the 32-block DiT when the `.aimodel` is
compiled just-in-time: the MPSGraph delegate forms a Neural Engine region inside the graph and the
ANE inference fails (`ANERegion.mm:414 failed assertion … Code=-19`). A plain ahead-of-time compile
fails the same way. What runs is an ahead-of-time compile with `--expect-frequent-reshapes`, loaded
with `SpecializationOptions.default()`:

```
xcrun coreai-build compile qi21_dit_full_bf16_dyn_iofp32.aimodel \
    --output aot/qi21_dit_full_bf16_dyn_iofp32 \
    --platform macOS --architecture h16c --preferred-compute gpu --expect-frequent-reshapes
# -> aot/qi21_dit_full_bf16_dyn_iofp32/qi21_dit_full_bf16_dyn_iofp32.h16c.aimodelc
```

Do the same for the encoder and the VAE bundle you use; every gate of this port ran all three this
way, on an M4 Max with `--architecture h16c` (without `--architecture`, `coreai-build` compiles for
every supported architecture). The compiled DiT is about 1.9× its `.aimodel` (27 GB); the encoder's
stays at 14.10 GiB. Whether a Swift host using `GraphModel(computeUnits: .gpu)` hits the same Neural
Engine region has not been tested.