feat(fp8): run scaled and raw fp8 checkpoints on the fp8 tensor cores

[Pfannkuchensack]

Summary

Adds FP8 Compute: running checkpoints that already ship FP8 weights directly on the FP8 tensor cores via torch._scaled_mm, instead of unpacking them to BF16 at load. This is distinct from the existing FP8 Storage toggle, which quantizes a full-precision model yourself and still does the math in BF16. New global setting fp8_compute in invokeai.yaml, default false.

Two kinds of checkpoint are covered, because both exist in the wild:

Scaled FP8 (FP8 weight + a weight_scale, the ComfyUI convention). Previously these were dequantized at load. They now stay quantized, with the scale attached to the layer. Wired into FLUX.1, FLUX.2, Anima, Krea-2, Z-Image and the Qwen3-VL and Mistral encoders. Handling that turned out to need more than reading one key:

• full_precision_matrix_mult hints arrive via two transports — the safetensors header _quantization_metadata and per-layer .comfy_quant uint8-JSON tensors. The latter marks 96 of 256 layers in PM_Krea2_turboV2FP8 and was previously ignored entirely. Honored by default; fp8_compute_full_precision_hints: false overrides the producer.
• Both spellings (weight_scale/scale_weight, input_scale/scale_input) are read everywhere. Reading only one meant a checkpoint using the other spelling had its scales deleted without ever being applied, leaving every quantized weight off by 1/weight_scale with nothing logged.
• Uncalibrated input_scale placeholders (exactly 1.0, non-finite, <= 0) are discarded rather than applied.

FLUX.1 is the cheap case and the most rewarding one. The checkpoint is already in the BFL layout Flux implements and qkv stays one fused Linear, so no key conversion, no qkv split and no scale copying are needed -- the per-tensor scale attaches to exactly the module it was computed for. The model is 100% nn.Linear with no shape outside the _scaled_mm 16-divisibility rule, so every quantized layer reaches the tensor cores. Before this, a scaled-FP8 FLUX.1 checkpoint did not load at all: its 629 scale and marker keys reached load_state_dict as unexpected keys.

One fix falls out of that. convert_bundle_to_flux_transformer_checkpoint folded anything ending in scale to bf16 for the model's own RMSNorm parameters -- and .weight_scale ends in scale too, so a quantization scale lost 8 mantissa bits before the loader ever saw it.

Raw FP8 (FP8 weight, no scale at all). The runtime already handled these — scaled_mm_linear treats weight_scale as optional — but the loaders never let them through: FLUX, FLUX.2 and Z-Image cast the whole state dict to BF16 unconditionally, discarding both the VRAM saving and the tensor cores before the model is built. Krea-2 kept them by accident, via a dtype check written for the scaled path, and logged nothing about it. A shared cast_state_dict() now casts everything except qualifying FP8 weights when the matmul is actually available, wired into all four loaders with a log line.

Only nn.Linear.weight is preserved. That restriction is not cosmetic. A Z-Image checkpoint in the wild quantized everything: 243 of its 453 FP8 tensors are 1-D biases, norm weights and a learned pad token. Keeping those quantized saves nothing usable and breaks inference — the FP8 value reaches the activations, the next Linear receives an FP8 input, and the forward dies with "abs_cuda" not implemented for 'Float8_e4m3fn'. Skip patterns alone are not enough here; you end up chasing one layer after the next. The predicate is _is_fp8_matmul_weight(): suffix .weight and dim() >= 2 and the module really is a torch.nn.Linear. A model's own _skip_layerwise_casting_patterns is honored on top, for Linears whose forward casts activations to their weight's dtype (Z-Image's TimestepEmbedder does exactly that).

FLUX.2 is the hard one, and the only place the side-channel needs real work. Its keys are renamed BFL -> diffusers and its fused qkv is split into three projections; both steps drop a scale silently, leaving the weights quantized but unscaled. Worse, the block renames are substring tests, so img_attn.proj.weight_scale satisfies "img_attn.proj.weight" in rest and the scale would be written to the weight's destination key, replacing the weight itself. Side-channel keys are therefore no longer routed through the weight converter at all: each one is placed by pushing a probe <base>.weight through the real converter and reading back where it landed, so placement cannot drift from the rename rules. A fused qkv yields three destinations — a scalar scale is copied to each, a per-channel scale split like the weight. Layer hints get the same one-to-many rename, since they name layers BFL-side while the scales are read diffusers-side.

The FLUX.2 [dev] Mistral text encoder turned out to be the largest single item in this series. It ships as scaled FP8 — 210 quantized Linears — and the loader folded every scale into BF16 at load. That doubles it: 16.8 GiB on disk becomes 32.3 GiB in memory, which does not fit on a 24 GB card even on its own, and the dequantization ran for 13 minutes before OOMing. It now loads in 5.8 s at 17,180 MB, fully resident. That is 15.5 GiB saved — more than any transformer it feeds.

Anima did not load at all either. Its _filter_non_model_keys drops derived buffers and exporter metadata but not the quantization side-channel, and the loader raises on unexpected keys — 500 of them on a plain scaled export, 749 on one that also ships comfy_quant markers. Anima itself is the easy shape (separate q_proj/k_proj/v_proj, one prefix strip), but it exposed a quieter bug shared with the FLUX.2 work: _quantization_metadata names layers in the checkpoint's own scheme, net.-prefixed, while the scales are read after the prefix is stripped. Read without renaming, the header matches nothing and every full_precision_matrix_mult flag is dropped silently. In both loaders the names are now pushed through the real rename rather than restated.

float8_e5m2 is deliberately not preserved: _scaled_mm cannot take it as the weight operand on Ada, so keeping it quantized would buy VRAM at the cost of a per-forward dequantize.

FP8 storage no longer silently defeats FP8 compute. Layerwise casting installs hooks that restore the compute dtype before every forward, so on an already-FP8 checkpoint the matmul would quietly fall back to the dequantized path — the VRAM toggle would make the model slower with no indication why. It is now skipped, with a log line, when the weights are already quantized and the matmul is in use.

Device support is probed by actually attempting a small FP8 matmul, not inferred from the device name, so an unsupported build falls back cleanly instead of failing partway through a generation.

Not in scope

• Qwen-Image is untouched by FP8 compute.
• SDXL was measured and deliberately left out. Its UNet is 87% nn.Linear and every shape clears the _scaled_mm 16-divisibility rule, so it would suit this well — but it loads through diffusers' from_single_file/from_pretrained, which owns dtype conversion and leaves no state-dict seam, and FP8 SDXL checkpoints barely exist. Not worth the loader surgery today.

Related Issues / Discussions

Follow-up to #8945 (FP8 storage), #9231 (hook-based casting) and #9412 (compute-dtype resolution).

Stacked on #9416, which is stacked on #9415 -> #9414. All four touch _should_use_fp8 and _apply_fp8_to_nn_module; the overlap is already resolved on this branch, so merging in that order is conflict-free. _apply_fp8_to_nn_module now carries both extension points: extra_skip_patterns (declarative, from #9414) and skip (programmatic, used here to leave scaled-FP8 layers alone). They are complementary — a scaled layer holds an FP8 weight plus a weight_scale, and the casting hooks would upcast it without applying that scale, i.e. a silently wrong weight.

QA Instructions

Requires an FP8-matmul-capable GPU (SM 8.9+ — Ada/Hopper/Blackwell) and fp8_compute: true in invokeai.yaml.

Set enable_partial_loading: false before measuring anything. FP8 compute only applies to layers actually resident on the GPU; anything still in RAM takes the normal path. Under partial loading the same seed does not reproduce (measured: mean pixel difference 0 -> 16-22 across 98.7% of pixels) and run-to-run noise is as large as any settings difference. This is documented with a :::danger callout in fp8-storage.mdx.

Raw FP8 — e.g. Z-Image unstableRevolution_V2Fp8 (FP8 tensors, no .weight_scale keys):

1. Generate. The log must show kept N raw fp8 weight(s) and no FP8 layerwise casting enabled line. If you see the casting line, storage won and the measurement is meaningless.
2. Total model size should roughly halve.
3. Generate at least 4 images and compare warm s/it. One image is not enough — cold load and residency drift both hide in it.

Scaled FP8, FLUX.1 — comfyanonymous/flux_dev_scaled_fp8_test (needs a separate CLIP-L, T5 and VAE; the bundled ones in an all-in-one checkpoint are discarded by the loader and cannot be selected):

4. Log shows FLUX: kept 314 layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled) and the transformer at ~11,350 MB. On main this checkpoint fails to load with 629 unexpected keys.
5. Enable FP8 Storage on the same model as well. Expect FP8 storage skipped ... 314 weight(s) are already fp8 and unchanged timings — not a slowdown.

Scaled FP8, Krea-2 — e.g. PM_Krea2_turboV2FP8:

6. Log shows kept N layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled) plus, on this checkpoint, the note that 96 of 256 layers are marked full_precision_matrix_mult.
7. Toggle fp8_compute_full_precision_hints: false — those layers then also run on the tensor cores, which is worth ~1.5x on a checkpoint with no marked layers (see below).

Scaled FP8, FLUX.2 — black-forest-labs/FLUX.2-klein-9b-fp8 or -4b-fp8 (4B needs a Qwen3-4B encoder; the loader rejects an 8B one):

8. Log shows FLUX.2: kept N layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled) and no FP8 layerwise casting enabled line. On main the scales are folded and the model runs as FP8 storage.
9. FLUX.2 [dev] additionally logs Mistral encoder: kept 210 layer(s) in fp8. On main that encoder OOMs on a 24 GB card after a ~13 minute dequantization; here it loads in seconds at ~17,180 MB.

Scaled FP8, Anima — e.g. Bedovyy/Anima-FP8 or pachiiahri/anima-fp8-comfyui:

10. Log shows Anima: kept 250 layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled) and the transformer at ~2,488 MB instead of ~3,988 MB. On main the model fails to load with several hundred unexpected keys.

Regressions:

11. fp8_compute: false — behavior unchanged, weights dequantize to BF16 as before.
12. A GGUF and a bnb-NF4 checkpoint must be unaffected.
13. A raw fp8 checkpoint (fp8 weights, no weight_scale anywhere) must be unaffected by the scaled path: it still logs kept N raw fp8 weight(s) and no kept N layer(s) in fp8 line. Verified on the FLUX.1 raw bundle — 0 scaled layers detected, 314 weights kept with fp8_compute on, 0 with it off.
14. A bundled FLUX.1 checkpoint (model.diffusion_model. prefixed, text encoders and VAE inside) must still load its transformer. Note the bundled encoders and VAE are discarded by the loader and cannot be selected — that is pre-existing behaviour, not a regression.

Measured — RTX 4090, torch 2.7.1+cu128

Z-Image unstableRevolution_V2Fp8, 1024x1024, 30 steps, 1 warm-up + 3 measured:

	before	after
transformer	11,740 MB	5,881 MB
residency	95.5 - 100%	100%
s/it (warm)	1.19 / 1.34 / 1.36 -> 1.297	1.01 / 1.06 / 1.00 -> 1.023
graph total (warm)	40.0 s	30.9 s

Note the drift on the before side: at 11.7 GB the transformer no longer stays fully resident, so per-step time degrades run over run. The FP8 version is not just faster on average, it is stable.

That 1.27x is bought by residency, not by the tensor cores — at 11.7 GB the BF16 transformer no longer stayed fully resident. How much the matmul itself buys depends on the checkpoint, and the spread is large.

FLUX.1 dev scaled FP8 — the same question with residency held constant

1024x1024, 20 steps, 1 warm-up + 3 measured. FP8 storage is the honest baseline here: at BF16 the transformer is 23.8 GB and does not fit on a 24 GB card alongside T5, so "turn FP8 off" is not an option a user has.

arm	fp8_compute	fp8_storage	warm (s)	mean	VRAM	resident
A	true	false	16.4 / 16.5 / 16.4	16.4 s	11,353 MB	100%
B	false	true	26.1 / 26.5 / 24.5	25.7 s	11,350 MB	100%
C	true	true	18.1 / 16.4 / 16.4	17.0 s	11,353 MB	100%

1.57x, on the pinned torch 2.7.1. The two arms differ by 3 MB of VRAM and are both fully resident, so this is the matmul and nothing else. Arm C is the guard above doing its job — without it, storage would have taken the matmul away and landed on arm B's numbers.

Anima — no speedup, but the checkpoint loads at all

Anima base v1.0, 1024x1024, 35 steps, CFG 4.5, same seed both arms. The baseline here is BF16, since Anima fits either way.

arm	warm (s)	mean	VRAM
FP8 scaled	14.5 / 14.4 / 14.1	14.3 s	2,488 MB
BF16	14.4 / 14.4 / 14.5	14.4 s	3,988 MB

38% less VRAM at identical speed. Same composition, same lighting, same framing at the same seed; the differences are in fine structure (spoke detail, stone texture) — not a degraded or dithered image. For Anima the point is not the matmul anyway: before this the checkpoint raised on 500 unexpected keys.

What the three profiles suggest

model	gain	quantized Linears
FLUX.1 dev	1.57x	314 of 314
FLUX.2 Klein 9B	1.10x	144
Anima base	1.00x	250

The trend is consistent with coverage driving the speedup, but it does not establish it — Anima is a 2 B model where the matmul is unlikely to be the bottleneck at this resolution regardless. Treated as an observation, not an explanation.

FLUX.2 Klein — a smaller win, and a VRAM trade-off worth knowing

Same conditions, 8 steps. FP8 storage is again the honest baseline, since these checkpoints ship quantized.

arm	fp8_compute	fp8_storage	model	warm (s)	mean	VRAM	layers in FP8
A	true	false	Klein 9B	18.5 / 18.5 / 18.4	18.5 s	9,030 MB	144
B	false	true	Klein 9B	20.1 / 20.5 / 20.5	20.4 s	8,708 MB	–
C	true	false	Klein 4B	8.4 / 10.0 / 8.1	8.9 s	3,902 MB	100

1.10x on Klein 9B — well short of FLUX.1's 1.57x. The likely reason is coverage: the FLUX.1 checkpoint quantizes 314 of 314 Linears, Klein 9B only 144. The matmul can only act where FP8 actually is; the rest runs BF16 either way. Not separately measured.

Note the direction of the VRAM column. FP8 compute is larger than FP8 storage here — 9,030 MB vs 8,708 MB — because storage quantizes every eligible Linear while compute keeps only the ones the checkpoint shipped. Storage is smaller and slower, compute is bigger and faster. Anyone expecting "FP8 compute always shrinks the model further" will be surprised, so it is worth stating.

Cold load follows the usual pattern: 48.6 s with compute against 119.0 s with storage, the BF16 round trip again.

FLUX.2 [dev] renders end to end with an FP8 encoder and an FP8 transformer, but is not a valid timing subject on 24 GB: at 32.2 B parameters the transformer alone is 33.8 GB resident (60.0 GiB without FP8 compute), so it needs partial loading and the measurement rules above rule it out.

Why Krea-2 shows nothing and FLUX.1 shows 1.57x

Measured on Krea-2 at 100% residency in both arms (12,228 MB vs 12,533 MB):

torch / CUDA	FP8 storage	FP8 compute	gain
2.7.1 / 12.8 (shipped)	1.431	1.433	1.00x — none
2.11.0 / 12.8	1.459	1.266	1.15x
2.11.0 / 13.0	1.479	1.265	1.17x

The leading suspect is full_precision_matrix_mult coverage. PM_Krea2_turboV2FP8 marks 96 of its 256 layers, and those dequantize on every forward instead of using the tensor cores; the scaled FLUX.1 checkpoint marks none of its 314. That is not the whole story — turning the hints off on Krea-2 was separately measured at only ~10% (57% with LoRAs), which does not close a 1.00x-to-1.57x gap on its own, so architecture and layer shapes plausibly contribute. What the FLUX.1 arms do establish is the part that matters for reviewing this PR: the flat Krea-2 result is not a property of torch 2.7.1. On the pinned version, the same code path is worth 1.57x on a checkpoint that lets it run.

cu130 changes nothing on Ada either way (0.08%).

What FP8 compute buys regardless of the matmul:

• Cold load. FP8 storage must dequantize a scaled-FP8 checkpoint to BF16 and then re-quantize it back to FP8 via layerwise casting — measured 139-189 s cold versus 12-40 s on Krea-2, and 82.9 s versus 30.7 s on FLUX.1. FP8 compute just leaves the scales in place.
• RAM. That same roundtrip drove the server process to a 58.3 GB working set for a 12 GB model (cf. fix(model-loaders): stop materializing scaled-fp8 checkpoints in float32 #9429).
• Fit. Where a model does not fit at BF16 — the Z-Image case above — halving it is what makes full residency possible in the first place.

fp8_compute_full_precision_hints: true costs ~10% without LoRAs (57% with) and is still the right default: it is the fidelity the quantizer asked for.

Merge Plan

Merge after #9414 -> #9415 -> #9416. No DB migration, no redux slice, no new dependency. Adds two invokeai.yaml settings (fp8_compute, fp8_compute_full_precision_hints), both opt-in via fp8_compute defaulting to false, so existing installs are unaffected.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable) — tests/backend/quantization/test_fp8_scaled.py plus captured key-layout fixtures for scaled-FP8 Z-Image and mixed-FP8 FLUX.2
• ❗Changes to a redux slice have a corresponding migration — n/a, no redux changes
• Documentation added / updated (if applicable) — docs/src/content/docs/configuration/fp8-storage.mdx, renamed to "FP8 Storage & Compute"
• Updated What's New copy (if doing a release after this PR)

[lstein]

Adversarial review — round 2 (d8f16182dc)

All three round-1 blockers are fixed, and I verified each:

• Krea-2 skip-pattern regression — split_fp8_scaled_layers() now folds the scale into any layer the cast would dequantize anyway, before cast_state_dict() runs. time_embed.linear_1/2 come out correct.
• ROCm RDNA3 crash — _probe_fp8_matmul() replaces the capability guess, and torch.version.hip is not None forces the probe rather than short-circuiting it. gfx1100 fails the probe and falls back.
• FLUX.2 dead branch — _dequantize_fp8_weights(..., keep_fp8=...) now spares raw fp8, so the FLUX.2 half of the raw path is reachable.

The shared fp8_scaled.py is a real improvement over three divergent copies, and routing cast_state_dict / predict_cast_state_dict_size / split_fp8_scaled_layers through one can_stay_quantized() predicate closes the drift that caused round 1's blocker.

New round, new findings. Everything below is reproduced against d8f16182dc.

---

Blockers

1. Krea-2 native format: last.linear's weight_scale is dropped — regression vs main

The loader now converts native → diffusers keys before extracting the scales, on the stated grounds that "the renames are substring-based on .weight, so a sibling .weight_scale is carried to the same destination key automatically."

That holds for the str.replace() rules and the startswith prefix rules, but five branches in _convert_krea2_native_to_diffusers use exact key equality — k == "last.linear.weight", "last.linear.bias", "last.norm.scale", "last.modulation.lin", "txtmlp.0.scale". A sibling last.linear.weight_scale matches none of them and is left behind while its weight moves to final_layer.linear.weight.

Two further keys the substring rules also miss, for the same reason:

• .scale_weight spelling — k.replace(".attn.wk.weight", ".attn.to_k.weight") does not match .attn.wk.scale_weight. On a native checkpoint using the old spelling, every within-block scale is orphaned.
• .input_scale — likewise never renamed, so calibrated activation scales are silently discarded and every forward pays the amax reduction.

In all three cases extract_fp8_scaled_layers() pops the orphan key (so nothing trips load_state_dict) and drops it, because sd["<native path>.weight"] no longer exists. The weight stays quantized with no scale attached, and since the layer never enters fp8_layers, warn_on_unattached_scales() cannot see it either. With fp8_compute on it runs through _scaled_mm unscaled; with it off it is cast to bf16 unscaled. Off by 1/weight_scale, silently, both ways.

Reproduced:

sd = {
    "blocks.0.attn.wq.weight": fp8, "blocks.0.attn.wq.weight_scale": t(0.5),
    "blocks.0.attn.wq.input_scale": t(0.25),
    "blocks.0.attn.wk.weight": fp8, "blocks.0.attn.wk.scale_weight": t(0.5),
    "last.linear.weight":      fp8, "last.linear.weight_scale":      t(0.5),
}
conv   = _convert_krea2_native_to_diffusers(sd)
layers = extract_fp8_scaled_layers(conv)
# recovered: ['transformer_blocks.0.attn.to_q']
# UNSCALED : final_layer.linear.weight
# UNSCALED : transformer_blocks.0.attn.to_k.weight
# input_scale recovered for to_q: None

main gets last.linear right, because _dequantize_scaled_fp8() ran before the conversion:

final_layer.linear.weight: dtype=torch.bfloat16 val=0.5   # main
final_layer.linear.weight: dtype=float8, no scale         # this PR

Related: _dequantize_scaled_fp8() is now dead code — nothing outside tests/ calls it. Its eight tests, including the new test_the_scale_weight_spelling_is_folded_too, assert a behaviour the shipping path no longer has.

2. float8_e5m2 scaled weights lose their scale entirely

extract_fp8_scaled_layers() pops every scale key up front, then keeps only the layers whose weight is float8_e4m3fn:

if weight is None or getattr(weight, "dtype", None) != FP8_DTYPE:
    continue

For an e5m2 weight the scale has already been removed from sd and is now discarded. Nothing folds it — dequantize_fp8_scaled() never sees the layer, and cast_state_dict() does a plain .to(bf16):

sd = {"lin.weight": (ones(32,32)*2).to(torch.float8_e5m2), "lin.weight_scale": t(0.25)}
extract_fp8_scaled_layers(sd)          # -> {}   (and the scale key is gone)
# resulting weight: 2.0    expected: 0.5

Excluding e5m2 from staying quantized is the right call and is well argued in the docstring; excluding it from scale recovery is not. On main, FLUX.2's and Krea-2's dequantizers folded the scale regardless of dtype, so this is a regression on both.

Suggested shape: extract for any float8* dtype, and let can_stay_quantized() keep the e4m3fn-only restriction — split_fp8_scaled_layers() then folds the e5m2 layers with the scale properly applied.

3. Block-wise (2-D) weight_scale now raises instead of being expanded

The FLUX.2 loader keeps _dequantize_fp8_weights() as a "Safety net for scale layouts the shared extractor does not model (block-wise scales whose shape has to be expanded to the weight's). It is a no-op … because extraction has already taken the scales it understood."

Extraction takes all scales, not only the ones it understands — the weight_scales[path] = sd.pop(key) loop is unconditional. So the repeat_interleave expansion is unreachable for exactly the layouts it was written for, and dequantize_fp8_scaled() has no equivalent:

sd = {"b.weight": ones(64,128).to(fp8), "b.weight_scale": full((64,2), 0.5)}
dequantize_fp8_scaled(sd, extract_fp8_scaled_layers(sd), torch.bfloat16)
# RuntimeError: The size of tensor a (64) must match the size of tensor b (128) at
#               non-singleton dimension 0

main expands this correctly. With fp8_compute off — the default — the model now fails to load; with it on, the mismatch surfaces later as a shape error in scaled_mm_linear's out * weight_scale.reshape(1, -1), mid-generation.

---

Major

4. _quantization_metadata header hints are dropped for any prefixed checkpoint

The PR fixes the "hints are named in the checkpoint's scheme" bug for the rename step (remap_flux2_layer_paths, _remap_native_layer_paths, _strip_anima_prefix_from_layer_paths) but not for the prefix strip that runs before it. The header is read from the file, so its layer names still carry model.diffusion_model. / diffusion_model. / net., while the state dict has had the prefix removed.

meta = {"_quantization_metadata": '{"layers": {"model.diffusion_model.blocks.0.attn.wq":
                                   {"full_precision_matrix_mult": true}}}'}
# after _strip_comfyui_prefix + _remap_native_layer_paths:
#   hint key : model.diffusion_model.blocks.0.attn.to_q
#   layer    : transformer_blocks.0.attn.to_q   full_precision_matmul=False

The producer said "do not multiply this layer in fp8" and it is multiplied in fp8, silently — the precise failure the hint plumbing exists to prevent. Same structure in FluxCheckpointModel (bundled checkpoints, after convert_bundle_to_flux_transformer_checkpoint strips the prefix), Flux2CheckpointModel, and ZImageCheckpointModel. Invisible on the checkpoints you measured because PM_Krea2_turboV2FP8 ships the flags as .comfy_quant tensors, which ride through the strip inside the state dict.

Pushing the hint names through the prefix strip as well — the same "run it through the real function" trick already used for the renames — covers all four loaders.

5. FLUX.2: the adaLN scale/shift swap is not applied to a per-output-channel weight_scale

final_layer.adaLN_modulation.1.weight has its two halves swapped by _flux2_swap_scale_shift(). The side-channel is placed by probing the converter for the destination key, which is correct, but the value is copied verbatim:

weight rows after convert : [12, 16, 20, 0, 4, 8]     # halves swapped
weight_scale after convert: [ 1,  2,  3, 4, 5, 6]     # unchanged

Rows 0–2 now hold the original rows 3–5 but are scaled by the original rows 0–2's factors. Per-tensor scalars are unaffected, and split_qkv_sidechannel handles the fused-qkv case correctly — this is the one converter transform that reorders rows and is not mirrored onto the scale.

---

Minor / robustness

6. _strip_anima_prefix_from_layer_paths can abort the load. zip(layer_names, stripped.keys(), strict=True) raises ValueError if _strip_anima_bundle_prefix drops any name — which it does for any hint name outside the detected prefix. Uncaught, so the whole Anima load fails:

   _strip_anima_prefix_from_layer_paths(["net.blocks.0.attn.q_proj", "final_layer.linear"])
   # ValueError: zip() argument 2 is shorter than argument 1

   _remap_native_layer_paths in krea2.py wraps the equivalent call in try/except; this one should too.

7. A transient probe failure disables fp8 compute for the process. _probe_fp8_matmul catches every exception and the result is cached permanently (reset_fp8_matmul_support_cache is documented "for tests"). The probe runs during model loading, i.e. under real VRAM pressure — a CUDA OOM on the 16×16 allocation or on cuBLAS workspace init sticks as "this GPU cannot do fp8". Worth distinguishing an allocation failure from a genuine unsupported-op error, or not caching the former.

8. Mistral: _drop_quantization_metadata is skipped on the fp8 path, so keys ending .scale and starting scaled_fp8 — which that function deliberately dropped and is_scale_metadata_key does not cover — now reach load_state_dict. Harmless with strict=False, but they show up in the "ignored N unexpected keys" debug line.

9. Anima lost its non-float guard. The old loop was if sd[k].is_floating_point(): sd[k] = sd[k].to(model_dtype); cast_state_dict() casts unconditionally. No Anima checkpoint key appears to be non-float today (the integer-ish buffers are non-persistent and filtered by _NON_MODEL_KEY_SUFFIXES), so this is latent rather than live — but the guard was there on purpose.

10. Partial-load accounting. weight_scale/input_scale are non-persistent buffers, so they are invisible to CachedModelWithPartialLoad._state_dict_bytes, and _move_non_persistent_buffers_to_device re-copies all ~314 of them (buffer.to(device, copy=True), unconditionally) on every partial_load_to_vram. Both are small; noting so the accounting gap is a known one.

11. Doc/description drift. The fp8_compute text in the InvokeAIAppConfig class docstring omits the "reproducibility requires full residency" paragraph that the Field(description=...) carries. _QUANTIZED_MODEL_FORMATS and the matching list in MainModelDefaultSettings.tsx both omit sdnq_quantized (the _is_quantized_param backstop covers it, so this is cosmetic).

12. CI. python-checks and py3.11: windows-cpu are both cancelled, not failed — worth a re-run so the rollup is clean.

---

Attacked and held

• In-place sd[key] = ... while iterating for key in sd — no insert/delete, safe.
• The three-way agreement between cast_state_dict, predict_cast_state_dict_size and split_fp8_scaled_layers — all funnel through can_stay_quantized(); the round-1 drift is gone. predict_* no longer under-counts biases/norms.
• scaled_mm_linear operand layout: weight is (out, in) contiguous, so .t() is genuinely column-major; the unpad out[: x.shape[0] - pad] is right because x was rebound to the padded tensor; the per-channel weight scale is applied after the unpad and before the reshape.
• fp8 storage hooks vs. the matmul: the forward pre-hook restores the compute dtype before CustomLinear.forward runs, so _can_use_fp8_matmul never sees an fp8 weight on a storage-cast layer. The Qwen3-VL skip= lambda keys off weight_scale, which attach_fp8_scales has already set at that point.
• LoRA over fp8: _is_any_part_of_layer_fp8 forces sidecar patching, the sidecar's base call routes to _autocast_forward → fp8 matmul, and the non-LoRA residual branch goes through _cast_weight_bias_for_input, which now applies the scale rather than a bare .to().
• weight_scale buffers survive wrap_custom_layer (shared __dict__), are moved by _move_non_persistent_buffers_to_device under partial loading and by model.to() on full load.
• should_keep_fp8_weights short-circuits on is_fp8_matmul_enabled(), so a default install never runs the CUDA probe during loading. get_config() is lru_cached, so the per-forward is_fp8_matmul_enabled() is cheap.
• transformer.dtype leaking float8 into latents.to(...): Z-Image's first parameter is x_pad_token (1-D → never kept fp8), Mistral's is embed_tokens (an nn.Embedding), Anima has no .dtype. No float8 escapes. get_model_compute_dtype's "first non-fp8 float param" fallback also covers the loaders that now return early without setting the marker.
• FLUX.1's load_state_dict(..., strict=True): both scale spellings, both input-scale spellings, .comfy_quant and the stray scaled_fp8 are all popped before the load.
• iter_weight_scale_pairs skipping pairs whose .weight is absent, and extract_fp8_scaled_layers dropping a scale with no fp8 weight — correct, and the reason findings 1 and 2 are silent rather than loud.