FLUX.2 [klein] 4B/9B + Z-Image-Turbo on a shared DiT scaffold

[merceod]

What does this PR do?

Stacked on #286 (engine/dit-scaffold); this PR adds the two model packages, their configs, tests and benchmarks.

Numerics: the eager path is bit-exact with diffusers for klein-4B and 9B (text-to-image, single- and two-reference edits) and Z-Image-Turbo (including captions beyond the captured buckets). With the shipped defaults the klein transformer is bit-exact in the served path too (compile_exact_ops: [norms] keeps the norms on the eager kernels inside the compiled forward; a plain compile landed at a median 35-38 dB from the eager path over the 100 prompts); the remaining deviation is the autotuned VAE, at least 53 dB on every prompt, and vae_compile: false makes the served output bit-exact. Z-Image ships the plain compile (median 34 dB) because keeping its norms eager costs 45%; compile_exact_ops: [norms] is the documented alternative. Also verified served: rows are independent of their batch neighbours bit for bit, edits batch and match their solo images, the API probe (formats, n, error paths, uploads, disconnects) passes, and a 60-minute mixed-size soak ran 15983 images with 0 errors and flat memory.

Benchmarks (protocol table; 1x H100 80GB HBM3 on one node, driver 580.173.02, torch 2.11.0+cu128, FlashInfer 0.6.18.post1; 100 fixed prompts, seed = prompt index; latency n=20 after 3 warmups, throughput median of 3 x 32 images per concurrency level; M* rows re-measured 2026-09-21 on the shipped defaults, baselines from 2026-09-18)

System	model	size / steps	output	B=1 latency median	p95	images/s @4	images/s @8	images/s @16	peak VRAM	notes
mstar	flux2_klein	1024x1024 / 4	png (default)	0.364 s	0.371 s	3.18 (3.07-3.19)	3.23 (2.95-3.23)	3.22 (2.92-3.25)	35.4 GiB	n=20; 32 images x 3 repeats per level
mstar	flux2_klein_9b	1024x1024 / 4	png (default)	0.644 s	0.648 s	1.71 (1.66-1.71)	1.72 (1.67-1.72)	1.72 (1.71-1.72)	40.4 GiB	n=20; 32 images x 3 repeats per level
mstar	z_image_turbo	1024x1024 / 8	png (default)	0.862 s	0.885 s	1.19 (1.18-1.19)	1.17 (1.17-1.19)	1.18 (1.17-1.19)	43.5 GiB	n=20; 32 images x 3 repeats per level
sglang	Tongyi-MAI/Z-Image-Turbo	1024x1024 / 8	server default	0.863 s	0.897 s	1.15 (1.14-1.15)	1.16 (1.15-1.17)	1.18 (1.17-1.18)	67.7 GiB	n=20; 32 images x 3 repeats per level
sglang	black-forest-labs/FLUX.2-klein-4B	1024x1024 / 4	server default	0.387 s	0.435 s	2.82 (2.77-2.83)	2.73 (2.72-2.75)	2.86 (2.85-2.97)	64.7 GiB	n=20; 32 images x 3 repeats per level
sglang	black-forest-labs/FLUX.2-klein-9B	1024x1024 / 4	jpeg (default)	0.638 s	0.682 s	1.63 (1.62-1.63)	1.63 (1.63-1.65)	1.66 (1.65-1.67)	78.6 GiB	n=20; 32 images x 3 repeats per level
sglang_png	black-forest-labs/FLUX.2-klein-4B	1024x1024 / 4	png	0.502 s	0.613 s	2.22 (2.20-2.24)	2.28 (2.27-2.28)	2.30 (2.26-2.30)	64.6 GiB	n=20; 32 images x 3 repeats per level
vllm_omni	Tongyi-MAI/Z-Image-Turbo	1024x1024 / 8	server default	1.670 s	1.678 s	0.62 (0.62-0.62)	0.62 (0.62-0.62)	0.62 (0.62-0.62)	24.4 GiB	n=20; 32 images x 3 repeats per level
vllm_omni	black-forest-labs/FLUX.2-klein-4B	1024x1024 / 4	server default	0.438 s	0.459 s	2.72 (2.71-2.72)	2.71 (2.70-2.71)	2.66 (2.65-2.67)	19.9 GiB	n=20; 32 images x 3 repeats per level
vllm_omni	black-forest-labs/FLUX.2-klein-9B	1024x1024 / 4	png (default)	0.688 s	0.695 s	1.60 (1.60-1.60)	1.60 (1.60-1.60)	1.60 (1.60-1.60)	37.7 GiB	n=20; 32 images x 3 repeats per level
vllm_omni_jpeg	black-forest-labs/FLUX.2-klein-4B	1024x1024 / 4	jpeg	0.385 s	0.393 s	2.65 (2.64-2.66)	2.64 (2.64-2.65)	2.65 (2.64-2.65)	19.9 GiB	n=20; 32 images x 3 repeats per level

JPEG output (output_format: jpeg) saves about 30 ms per image on the M* rows. Exactness costs: klein with compile: false 0.461 s / 2.5 images/s (compiled VAE) or 0.523 s bit-exact; Z-Image with compile_exact_ops: [norms] 1.25 s / 0.80 images/s.

Reproduction (same GPU, same node, back to back; the client is benchmark/flux2_klein/bench_images.py)

# M*
mstar serve flux2_klein --port 8000
python benchmark/flux2_klein/bench_images.py --port 8000 --model flux2_klein --mode latency --n 20 --warmup 3 \
    --prompts commons/bench/data/image/prompts_100.txt --size 1024x1024 --steps 4 --tag mstar --out mstar_latency.json
python benchmark/flux2_klein/bench_images.py --port 8000 --model flux2_klein --mode throughput --concurrency 4 8 16 --n 32 ...
# SGLang 0.5.19 (its default output is JPEG; add --output-format png to the client for like-for-like)
sglang serve --model-path black-forest-labs/FLUX.2-klein-4B --port 30000 --performance-mode speed \
    --dit-cpu-offload false --text-encoder-cpu-offload false --attention-backend fa --warmup-mode server \
    --output-path '' --batching-mode dynamic --batching-max-size 16 --batching-delay-ms 5
# vLLM-Omni 0.28
vllm serve black-forest-labs/FLUX.2-klein-4B --omni --port 8002 \
    --default-sampling-params '{"0": {"num_inference_steps":4,"guidance_scale":1.0}}'
# served-vs-eager PSNR over the prompt set
python benchmark/flux2_klein/direct_pipeline.py --prompts-file prompts_100.txt --count 100 --out-dir sdpa_ref
python benchmark/flux2_klein/psnr.py --dirs sdpa_ref mstar_png --pattern 'mstar_{:03d}.png' --count 100

How was it tested?

• CPU: python -m pytest test/modular at the head, 1002 passed / 142 skipped (tiny-config bit-exactness of both ports, graph/walk structure, batching, output formats, knobs, the PSNR and summary tools).
• GPU parity vs diffusers oracles (test/flux2_klein/record_oracle.py --refs 2, test/z_image/record_oracle.py): test_flux2_klein_reference_equivalence.py for 4B and 9B (FLUX2_KLEIN_REPO), test_z_image_reference_equivalence.py, all steps 0.0 and images PSNR inf.
• GPU batch invariance: test_flux2_klein_batch_invariance.py, test_z_image_batch_invariance.py; served probes benchmark/flux2_klein/{batch_invariance_client,edit_probe,api_probe,soak}.py.

Known gaps

• No CUDA-graph buckets for edits or non-1024^2 sizes yet: they run the compiled eager path and each new shape pays one 5-7 s compile on first use.
• The autotuned VAE can pick different conv kernels in different server processes (about 64 dB apart on the same seeds); vae_compile: false avoids it. Z-Image's transformer output also depends on the batch size it ran in (cuBLAS/cuDNN algorithm choice, 31-39 dB vs the single-row image; rows are independent of each other).
• Not exercised on a GPU: LoRA with a real adapter. CFG for the non-distilled -base checkpoints is not implemented.
• klein-9B weights are under the FLUX Non-Commercial License (flagged in the registry, docs and README).

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant