host/host_contract.md

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

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,
`schedule_fixed`), `pipeline_engine.py --turbo` (the whole loop on the three bundles). The first three
match the diffusers-main pipeline (diffusers 85c9fa4, `Qwen/Qwen-Image-2.1-Turbo` @ `d65dbc9`) 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.

This port shares its graph code and its text-encoder and VAE graphs with RGBA-Image-2.1 (the files here
have their export-time debug locations stripped). Only the DiT weights and the sampler differ. Section 6
is the part that changed.

## 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 |
| `qi21t_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 |
| `qi21t_dit_full_int8lin_dyn_iofp32.aimodel` | the same as the bf16 DiT | the same | the same |
| `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, fp32 boundary. Two bundles with one interface; use either. `bf16`: bf16 weights
  and compute. `int8lin`: the 232 Linear weights in int8 (symmetric, one bf16 scale per block of
  32), dequantized inside the graph; norms, scales and compute in bf16.
- 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), the Turbo checkpoint's
`processor/tokenizer.json`. 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).

The Turbo `tokenizer.json` differs from the base checkpoint's only in its offset and decoder
settings (`trim_offsets`, `add_prefix_space`, `use_regex`). Vocabulary, merges and added tokens are
the same, and both files give the processor's ids on every prompt the port was gated on.

## 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 six files are byte-identical to
RGBA-Image-2.1's `host/`. 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, a fixed 8-step grid, no CFG)

The checkpoint carries its own sigmas: `sample_sigmas` in its `model_index.json`, copied to
`host/scheduler.json`. The pipeline passes them to the scheduler in place of a linspace and ignores
`num_inference_steps`. `use_dynamic_shifting` is false, so there is no `mu` and the grid is the same
at every image size. All arrays are fp32.

```
sigmas    = fp32(sample_sigmas)        # 1.0, 0.978453, 0.95418, 0.926626, 0.89508, 0.845148, 0.704534, 0.414568
sigmas    = shift·sigmas / (1 + (shift − 1)·sigmas)                  # shift 1.0: unchanged, bit for bit
                                                                     # shift_terminal null: no terminal stretch
timesteps = sigmas · num_train_timesteps                             # fp32
sigmas    = sigmas ++ [0]                                            # 9 values
for i in 0 ..< 8:
    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
```

On this grid `t` equals `sigmas[i]` bit for bit at all 8 steps; keep `timesteps[i] / 1000` anyway,
the form the pipeline uses. `qi21_sched.schedule_fixed` reproduces the reference sigmas, timesteps
and all 8 Euler steps bit-exactly at 256², 512² and 1024².

RGBA-Image-2.1 runs 40 steps on a shifted linspace instead (its `host/host_contract.md`, §6). The
Turbo DiT is a distilled checkpoint, not the base DiT run for fewer steps: on the same noise and
prompt, the fp32 Turbo pipeline's 8-step image and the fp32 base pipeline's 40-step image are
21.75 dB (256²) and 23.35 dB (512²) apart (white-composited RGB PSNR).

## 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.

The VAE bundles hold the base checkpoint's fp32 VAE weights. The Turbo checkpoint's `vae/` is a bf16
copy of those weights (all 238 tensors are the round-to-nearest-even of the fp32 ones). Decoding the
same latent, the two sets of weights differ by up to 0.074–0.106 on the [-1, 1] scale: uint8 PSNR
56.95 / 58.95 / 58.15 dB at 256² / 512² / 1024², at most 10–14 levels on one channel, 1 level on
alpha. A host that compares against the diffusers Turbo pipeline sees that band; the bundles match a
fp32 decode with their own weights to max|Δ| 3.0e-5 – 4.2e-5.

## 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, with a dynamic or a static image axis. What runs is an ahead-of-time compile with
`--expect-frequent-reshapes`, loaded with `SpecializationOptions.default()`:

```
xcrun coreai-build compile qi21t_dit_full_bf16_dyn_iofp32.aimodel \
    --output aot/qi21t_dit_full_bf16_dyn_iofp32 \
    --platform macOS --architecture h16c --preferred-compute gpu --expect-frequent-reshapes
# -> aot/qi21t_dit_full_bf16_dyn_iofp32/qi21t_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). On that Mac the bf16 DiT compiles in 107.6 s to 26.52 GiB, the int8 DiT
in 301 s to 20.30 GiB; the encoder's compiled copy is 14.10 GiB. With
`--preferred-compute neural-engine` the DiT compiles to 0 Neural Engine regions and runs on the GPU.
Whether a Swift host using `GraphModel(computeUnits: .gpu)` hits the same Neural Engine region has
not been tested.