feat(fp8): enable FP8 storage for Z-Image

[Pfannkuchensack]

Summary

Z-Image was excluded from FP8 storage in #8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete.

Removing it alone is not enough. Our hook-based cast (#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its input to it:

weight_dtype = self.mlp[0].weight.dtype   # float8_e4m3fn, our pre-hook has not fired yet
if weight_dtype.is_floating_point:        # float8 IS floating point -> True
    t_freq = t_freq.to(weight_dtype)      # the INPUT becomes float8

The compute_dtype fallback branch right below is never reached, because float8.is_floating_point short-circuits first. F.linear then dies with:

RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn'

which is exactly why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's own list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more, never fewer.

Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel.

Frontend: drop the !isZImage guard that hid the FP8 Storage switch in the model's Default Settings.

Related Issues / Discussions

Base of a four-PR series: this one, then #9415 (Anima), #9416 (quantized guard), #9478 (FP8 compute). All four touch _should_use_fp8 and _apply_fp8_to_nn_module, so merging in that order is conflict-free.

Follow-up to #8945 (FP8 storage) and #9231 (hook-based casting).

QA Instructions

Needs a CUDA GPU — _should_use_fp8 returns False on any other device, so the whole path is inert elsewhere.

Setup: Model Manager → a Z-Image model → Default Settings. The FP8 Storage switch is now visible for Z-Image (that is the frontend change). Enable it → Save. The cache invalidates on save, so the next generation reloads with FP8 storage.

1. Z-Image-Turbo (diffusers format) — generate. Look for this line in the log:

   FP8 layerwise casting enabled for Z-Image-Turbo (storage=float8_e4m3fn, compute=torch.bfloat16, param_size=5880MB)
   

   and the transformer resident at ~5880MB instead of ~11.5GB. On main (with the exclusion removed but without the skip-pattern change) this run fails with "addmm_cuda" not implemented for 'Float8_e4m3fn'.

2. A single-file Z-Image checkpoint (e.g. a ComfyUI-style .safetensors) — same, this is the ZImageCheckpointModel wiring. Verified locally with a 14.37GB checkpoint, also resident at 5880MB.

3. Compare output quality against the same seed/steps/CFG with FP8 off. Expect the same composition with slightly less micro-detail, not a different or broken image.

4. Toggle FP8 back off — unchanged behaviour.

5. Regression check on another base (FLUX or SDXL) with FP8 enabled, since the skip-pattern change affects every diffusers model: it can only skip more modules than before, so expect unchanged or marginally better output at marginally higher VRAM.

Unit tests:

uv run --extra cuda --extra test pytest tests/backend/model_manager/load -q --no-cov

Merge Plan

Plain merge — no DB schema, no redux slice, no API schema change. Note for the frontend: the FP8 switch now appears for Z-Image models that may already have a persisted fp8_storage value from before the UI hid it; those take effect on next load, which is the intended behaviour.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, no redux changes
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR) — n/a

[lstein]

Adversarial review at bfd4942870, in a worktree against main. No blockers — the central claim reproduces under direct attack. Seven non-blocking notes below.

The core claim reproduces

I built a real ZImageTransformer2DModel (diffusers 0.39.0) and ran a full CPU forward three ways:

	Linears cast	forward
no fp8	0	OK (baseline)
fp8, no declared patterns	37	RuntimeError: mat1 and mat2 must have the same dtype, but got Float8_e4m3fn and Float
fp8, with ['t_embedder', 'cap_embedder']	34	OK

Skipped: t_embedder.mlp.0, t_embedder.mlp.2, cap_embedder.1. t_embedder is the model's only TimestepEmbedder, so there is no second instance of this hazard. Note that diffusers 0.39 does have the compute_dtype fallback branch — the PR body is right that it is unreachable, because float8_e4m3fn.is_floating_point short-circuits first.

Attacks that failed

• Parity with diffusers. enable_layerwise_casting unions DEFAULT_SKIP_MODULES_PATTERN + _keep_in_fp32_modules + _skip_layerwise_casting_patterns, so the union here matches. The re.search-on-leaf-name vs diffusers' prefix-prune divergence is equivalent for every substring pattern any supported model declares.
• transformer.dtype after casting. z_image_denoise.py:589 and :703 do latents.to(transformer.dtype) — exactly the hazard backend/util/fp8.py exists to prevent, and Z-Image was previously immune only by being excluded. It does not break: get_parameter_dtype returns the first floating-point param in named_parameters() order, which is the root-level x_pad_token Parameter — not a Linear/Conv, so never cast. Verified empirically (model.dtype -> torch.bfloat16, forward OK). See note 7.
• LoRA on fp8 weights. LayerPatcher._is_any_part_of_layer_fp8 already forces sidecar patching.
• Cache accounting. ModelCache.put sizes with calc_model_size_by_data after the loader returns, so the fp8 shrink is reflected.
• Meta params. load_state_dict(sd, assign=True) is strict, so nothing meta-resident reaches the cast.
• Test sensitivity. Both new tests are load-bearing: ignoring extra_skip_patterns inside _apply_fp8_to_nn_module fails both; dropping only the caller argument fails one.
• tests/backend/model_manager: 1192 passed, 142 skipped. ruff check + ruff format --check clean.

Non-blocking notes

1. Stale docs, directly contradicted by this PR. docs/src/content/docs/configuration/fp8-storage.mdx:72 still has the row | Z-Image (any variant) | No — dtype mismatch with skipped layers |, and :105 lists Z-Image in the exclusion list a user is told to check when VRAM does not drop. The checklist marks docs n/a, but this page documents precisely the exclusion being deleted. The sentence "Within a supported model, norm layers, position/patch embeddings, and proj_in/proj_out are skipped" also now needs "plus the model's own declared list".

2. Unreserved peak-RAM overshoot in the new single-file wiring. assign=True makes each param alias its sd tensor, and sd is still live when the cast runs at z_image.py:489, so param.data.to(fp8) allocates a second copy while the bf16 originals stay reachable. Verified: after the cast, the sd entries are still bf16 and full size. make_room() reserved only the bf16 size, so peak is ~1.5x the reservation — roughly 11.5 GB reserved vs ~17.4 GB actual for Z-Image. One line fixes it (sd.clear() before the cast). The same shape pre-exists at krea2.py:374, so if you touch it, worth doing both.

3. The new wiring is untested. Deleting z_image.py:489 leaves the entire tests/backend/model_manager suite green. tests/backend/model_manager/load/test_krea2_loader_boundaries.py already drives _load_from_singlefile directly, so the pattern exists.

4. _keep_in_fp32_modules is not honored, although diffusers unions it too. I checked every currently-loaded model and it is inert: Wan's time_embedder sits under condition_embedder, scale_shift_table is a bare Parameter, and Krea-2's entries are all norm*. Worth a comment so the next model that declares one does not lose it silently.

5. The cross-model impact is understated. "Strict superset of our defaults (FLUX/SD3, UNet, CogView4)" is true, but omits the two supported models with an actual delta. Measured on meta-device builds with real configs:

model	params leaving fp8	savings lost
Wan 14B (condition_embedder, patch_embedding)	232 M	~221 MiB
Krea-2 (time_embed)	39 M	~38 MiB
FLUX.1 / Qwen-Image	0	0

Right direction — it is what diffusers intends — but Wan users on tight VRAM will see ~220 MiB more usage, and QA step 5 currently names only FLUX/SDXL, the two bases where nothing changes. Worth adding Wan to the regression check.

6. The new comment blesses a pre-existing bug. "including ComfyUI fp8 checkpoints, whose scale metadata was filtered out above" — dropping .scale_weight / scaled_fp8 and casting the raw fp8 codes to bf16 is not discarding metadata; for a ComfyUI scaled-fp8 Z-Image checkpoint it loads unscaled weights. Pre-existing (same issue I raised on #9478), but the comment now reads as though it is a safe no-op. Suggest rewording so it does not imply correctness for scaled-fp8 inputs.

7. transformer.dtype is dormant, not fixed. Now that Z-Image is fp8-eligible, z_image_denoise.py:589 / :703 should use get_model_compute_dtype(transformer). They are correct today only because x_pad_token happens to be parameter zero; move the pad tokens into a submodule, or add a Linear ahead of them, and the denoise loop starts feeding float8 into F.linear. Cheap hardening, and it removes the coupling to diffusers' parameter ordering.

Approving — note 1 is the only one I would call close to required, since it is user-facing and now says the opposite of what the code does.