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/: 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.