Feat: fp8 compute raw - #9478
Conversation
…sor cores InvokeAI dequantized ComfyUI 'scaled fp8' checkpoints to bf16 at load time in three near-identical implementations, discarding both the VRAM saving and the ability to use the fp8 tensor cores. Measured on an RTX 4090: the dequantize round trip makes fp8 *slower* than bf16 (0.90x on FLUX.2 Klein 9B), while torch._scaled_mm reaches 1.29x vs bf16 and 1.63x vs the dequantized path on Krea-2 Turbo, at the same VRAM and with no visible quality loss. Adds a shared invokeai/backend/quantization/fp8_scaled.py that keeps the quantization intact (weight_scale, optional calibrated input_scale, and the per-layer full_precision_matrix_mult hints), and an fp8 branch in CustomLinear._autocast_forward that falls back to the dequantized path whenever any precondition fails, rather than raising mid-generation. Wires up the Krea-2 single-file loader as the first consumer. Remaining loaders (FLUX.2, Z-Image, Qwen-Image) still dequantize eagerly. Gated behind `fp8_compute` (default off): the fp8 matmul quantizes activations too, so images change at a fixed seed. The same flag also decides whether the weights stay quantized, since keeping them fp8 without the fp8 matmul would halve VRAM but run slower.
…em on fp8 tensor cores The Krea-2 single-file loader dequantized ComfyUI 'scaled fp8' checkpoints to bf16 at load, discarding both the VRAM saving and any chance of using the fp8 tensor cores. It now keeps the quantization and hands the scales to CustomLinear, which multiplies via torch._scaled_mm where the checkpoint allows it. Measured on an RTX 4090 with krea2TurboOfficialComfy_krea2TurboFp8 at 1024x1024: 884 vs 1107 ms/step (1.25x) and 12.24 GiB resident instead of ~25 GB in bf16, so the model is fully resident rather than streamed. Images are indistinguishable from the dequantized path (PSNR 29.95 dB). The checkpoint's _quantization_metadata marks 96 of 256 layers with full_precision_matrix_mult; those are honored and stay in bf16, which is what keeps fidelity in line with ComfyUI (and costs ~23% of the speedup). Two subtleties the measurements exposed, both covered by tests now: - apply_custom_layers_to_model leaves device autocasting disabled for fully resident models, so a check living only in _autocast_forward never runs. The fp8 branch is consulted before the autocasting split. - _quantization_metadata names layers in the native scheme while the scales are extracted after the native -> diffusers rename, so the per-layer flags matched nothing. The metadata paths are pushed through the same converter. Gated behind `fp8_compute` (default off): activations are quantized too, so images change at a fixed seed.
…data read fp8 weights are force-routed to sidecar patching, and the sidecar wrapper dispatches through _autocast_forward, so the fp8 branch has to survive that route with the LoRA residual added on top. Verified on Krea-2 Turbo fp8 with a 256-layer LoRA: 1008 vs 1241 ms/step (1.231x, against 1.251x without the LoRA), all 256 fp8 modules routed to sidecar, and images equivalent between both paths (PSNR 26.83 dB). Reading the safetensors header metadata no longer fails the load. It only enriches fp8 handling with the per-layer full_precision_matrix_mult hints, so an unreadable header now warns and continues rather than raising — but it does warn, because without the hints layers the quantizer marked unsafe would silently be multiplied in fp8.
Adds a settings matrix (only `fp8_compute` is needed; `fp8_storage` is bypassed on that path) and logs when a redundant fp8_storage setting was skipped, so the case is not silent.
The Qwen-Image i2l node hardcoded vae.disable_tiling(), so a full-frame encode was the only option. At 2560x1440 that peaks at 9.26 GiB — on top of a resident multi-GB transformer, which is what makes an upscale round-trip run out of headroom exactly at this node while every other node fits. Adds `tiled` / `tile_size` input fields following the SD/SDXL i2l node, OR'd with the global force_tiled_decode setting. Off by default, so behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter. Without it the change would be inert: the cache would keep reserving the full-frame figure (10.99 GiB at 2560x1440) and evict models to honour it, no matter what the VAE actually does. Tiled, it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Measured through the node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending; real images blend far better), which is why this stays opt-in.
Both nodes reserve working memory for a full-frame operation, which at high resolutions exceeds a 24 GB card, so the model cache evicts everything else to honour it. On CUDA at 2560x1440: 19.91 GiB for the decode and 10.99 GiB for the encode. Tiling is the intended escape hatch, but it did not work on either node: - qwen_image_i2l hardcoded vae.disable_tiling(), so it could not be enabled. - qwen_image_l2i honoured the global force_tiled_decode, but computed its working-memory estimate before and independently of that flag. Tiling bounded the VAE while the cache still reserved the full-frame figure, so the memory was never freed for anything else — effectively inert. Adds `tiled` / `tile_size` input fields to both nodes following the SD/SDXL i2l/l2i nodes, OR'd with force_tiled_decode. Off by default; behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter, and both nodes resolve tile_size=0 to the VAE default (256px) before estimating. Tiled it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Without this the change would be cosmetic on i2l and remain inert on l2i. Measured through the i2l node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Verified across eight resolutions that tiled and untiled encodes produce the same latent dimensions. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending), which is why this stays opt-in. Also fixes a crash in qwen_image_i2l: `width`/`height` are `int | None`, but the workflow UI sends 0 for an unset number input, and `0 is not None` reached `image.resize((0, 0))` -> "height and width must be > 0". Non-positive values are now treated as unset, matching how tile_size uses 0.
Switching the encoder to fp8 compute left everything the checkpoint does not quantize in bf16, growing it from 4236MB to 4999MB. On a GPU already holding a ~12GB transformer that was enough to push the transformer out of a full VRAM load, which is far more expensive than the encoder change ever saved. Cast the remainder to fp8 storage, with two exclusions. Layers carrying a weight_scale keep it and go through _scaled_mm -- the cast hooks would upcast them without applying the scale. nn.Embedding is skipped because the token embedding table is the encoder's input representation: quantizing it doubles the error against bf16 (relative L2 0.0079 -> 0.0163) to save 371MiB, a bad trade for a model whose whole job is text fidelity. The old fp8_storage path did cast it, so this is strictly more accurate than what shipped before. fp8_storage 116.8 ms 4.14 GiB rel L2 0.0351 (original) fp8 compute 61.9 ms 4.88 GiB rel L2 0.0079 (previous commit) this 65.3 ms 4.50 GiB rel L2 0.0079
…llings A scaled-fp8 checkpoint may ship an input_scale of exactly 1.0, meaning the producer wrote the field without calibrating it. Taking it at face value replaces the per-forward amax scale with no scaling at all, so activations above the fp8 maximum saturate. Measured relative error against a bf16 reference, dynamic vs a 1.0 scale: |x|max 368 0.0274 0.0262 |x|max 1928 0.0277 0.4356 |x|max 30976 0.0253 0.9639 Inside +/-448 the two are equivalent -- fp8_e4m3 is a floating-point format, so a scale factor buys no relative precision the way it would for int8. Above it the unscaled path collapses. Non-finite and non-positive scales are rejected for the same reason: they cannot be a valid divisor. Also accept `.scale_input` as an alias for `.input_scale`, mirroring the `.scale_weight`/`.weight_scale` pair we already handle. Previously such a key was left in the state dict, and the Qwen3-VL loader deleted it outright, so a calibrated activation scale was silently discarded and every forward paid the amax reduction. That delete is now redundant and removed.
FP8 Compute had no user-facing documentation, and the existing FP8 Storage page actively claimed a compute path might arrive "later" — it is already here. The part worth writing down is the reproducibility constraint. torch._scaled_mm needs both operands on the same device, so a layer whose weights are still in RAM silently falls back to the dequantized BF16 path. Which layers that hits depends on how much of the model happened to fit, and that shifts between runs, so the same seed stops reproducing. Measured on a 24GB card with a ~12GB transformer at 88-95% residency: two runs with identical seed and settings differed in 98.7% of pixels; fully resident, repeated runs were bit-identical. Also corrects "FP8 + partial loading: fully supported" — true for Storage, but for Compute it costs 47% per step on top of the reproducibility loss. Regenerates settings.json, which predated both fp8 settings.
A checkpoint can ship fp8 weights with no weight_scale. The runtime already handles them — 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, discarding both the VRAM saving and the tensor cores. Krea-2 kept them by accident and said nothing about it. Only nn.Linear.weight is preserved. That restriction is not cosmetic: a Z-Image checkpoint quantized everything, 243 of its 453 fp8 tensors being 1-D biases, norm weights and a learned pad token. Keeping those fp8 saves nothing usable and breaks inference — the value reaches the activations and the next Linear gets an fp8 input, which dies in x.abs() with "abs_cuda" not implemented. A model's own _skip_layerwise_casting_patterns is honored on top, for Linears whose forward casts activations to their weight's dtype. Also stops fp8 storage from silently defeating fp8 compute: layerwise casting restores the compute dtype before every forward, so on an already-fp8 checkpoint the matmul would quietly fall back and the VRAM toggle would make the model slower with no indication why. Verified end-to-end on Z-Image unstableRevolution_V2Fp8, 1024x1024, 30 steps: transformer 11739MB -> 5881MB (both 100% resident), 1.60 -> 1.27 s/it. 1.297 -> 1.023 s/it (3 warm runs each), transformer 11740MB -> 5881MB, residency 95.5-100% -> 100%.
github-actions
Bot
added
python
lstein
left a comment
lstein
left a comment
There was a problem hiding this comment.
Adversarial review of the delta commit e93854e6ae against its stack base 49990d643e. Three blockers, one design gap, one accounting nit. The first three are reproduced below.
🔴 1. Krea-2 scaled-fp8 regression: skip_patterns silently discards weight_scale
krea2.py:432-439 routes the scaled path through the same cast_state_dict call as the raw path (keep_fp8=bool(fp8_layers) or should_keep_fp8_weights(...)) — and now hands it skip_patterns=Krea2Transformer2DModel._skip_layerwise_casting_patterns, which is ['time_embed', 'norm'].
For any layer in fp8_layers matching a skip pattern, cast_state_dict does a plain tensor.to(model_dtype) — the raw fp8 codes, scale not applied. attach_fp8_scales then skips that module because weight.dtype != FP8_DTYPE (fp8_scaled.py:259), so the scale is dropped for good. The weight ends up off by 1/weight_scale.
Repro, using the real Krea2Transformer2DModel._skip_layerwise_casting_patterns with the real extract_fp8_scaled_layers / cast_state_dict / attach_fp8_scales:
fp8_layers extracted: ['time_embed.linear_2', 'transformer_blocks.attn']
cast_state_dict kept: 1 of 2 scaled layers
attach_fp8_scales attached: 1
time_embed.linear_2 dtype=torch.bfloat16 scale=NONE rel-err vs true weight = 12405.8379
transformer_blocks.attn dtype=torch.float8_e4m3fn scale=yes rel-err vs true weight = 0.0261
This is live on real checkpoints, not a synthetic shape. Krea2Transformer2DModel has exactly two Linears matching those patterns — time_embed.linear_1 (6144×256) and time_embed.linear_2 (6144×6144, 37.7M params) — and _convert_krea2_native_to_diffusers maps tmlp.0 / tmlp.2 straight onto them (krea2.py:217-219), so they are ordinary quantized tensors in a ComfyUI export. This very likely fires on krea2TurboOfficialComfy_krea2TurboFp8, the checkpoint the base branch was benchmarked with.
This is a regression, not a pre-existing gap. The pre-PR line was if sd[k].dtype is not FP8_DTYPE: sd[k] = sd[k].to(model_dtype) — every fp8 tensor survived, so every scale got attached.
The same class of failure applies beyond skip patterns: any scaled layer that fails _is_fp8_matmul_weight (not resolvable via get_submodule, not an nn.Linear) is now dequantized without its scale.
Suggested fix: apply the skip_patterns / _is_fp8_matmul_weight filters to raw fp8 only. For scaled layers, either keep them fp8 unconditionally or fold them through dequantize_fp8_scaled first and drop them from fp8_layers. Independently, attached != len(fp8_layers) should warn rather than pass silently — that mismatch is the tell for exactly this bug, and today the log prints attached as if it were the whole story.
🔴 2. ROCm RDNA3: should_keep_fp8_weights returns True, then every forward raises
device_supports_fp8_matmul (base branch, fp8_scaled.py:341) tests torch.cuda.get_device_capability(index) >= (8, 9). On ROCm that returns the gfx arch, not an SM version. On a W7900 (gfx1100) it is (11, 0) → passes.
Verified on the card:
device_supports_fp8_matmul(W7900) -> True
should_keep_fp8_weights(W7900) -> True
_can_use_fp8_matmul -> True
FORWARD RAISED: RuntimeError torch._scaled_mm is only supported on CUDA devices
with compute capability >= 9.0 or 8.9, or ROCm MI300+
The predicate is the base branch's, but this PR is what makes it reachable for FLUX / Z-Image / raw Krea-2: previously those loaders cast the whole state dict to bf16, so _can_use_fp8_matmul short-circuited on weight.dtype != FP8_DTYPE. Now the weights stay fp8 and every Linear forward throws. _can_use_fp8_matmul's stated contract — "falling back … rather than raising, because a mid-generation failure is far worse than losing the speedup" — is defeated, because the fallback gates on the same broken predicate.
Corroborating signal: 12 of the 49 tests in tests/backend/quantization/test_fp8_scaled.py fail (not skip) on such a box, because the cuda_fp8 skipif marker uses the same predicate:
12 failed, 37 passed, 17 warnings
Fix: exclude ROCm unless MI300+, e.g. if torch.version.hip is not None: return get_device_properties(i).gcnArchName.split(':')[0] in {"gfx942", "gfx950"} — or probe once with a 16×16 _scaled_mm in a try/except and cache that result.
🔴 3. The FLUX.2 half of the PR is dead code
Flux2CheckpointModel._load_from_singlefile calls self._dequantize_fp8_weights(sd) (flux.py:919) before cast_state_dict. That method's final loop (flux.py:1094-1101) converts every tensor whose dtype contains "float8" to bf16, unconditionally — it is not gated on should_keep_fp8_weights. Nothing fp8 survives to reach cast_state_dict:
input dtypes : {'...qkv.weight': 'torch.float8_e4m3fn', '...qkv.bias': 'torch.float32'}
after _dequantize_fp8_weights: {'...qkv.weight': 'torch.bfloat16', '...qkv.bias': 'torch.float32'}
any fp8 left for cast_state_dict to keep? -> False
kept is therefore always 0, and the "FLUX.2: kept N raw fp8 weight(s)" log can never fire. The PR body lists FLUX.2 among the loaders fixed; it isn't. The keys_to_convert loop needs the same keep_fp8 gate.
🟡 4. Z-Image cannot tell raw fp8 from scaled fp8
z_image.py:275-281 deletes .scale_weight keys and never applies them, before cast_state_dict runs. So on a scaled-fp8 Z-Image checkpoint the weights are now kept quantized and multiplied with weight_scale=None, while the new log announces "Z-Image: kept N raw fp8 weight(s) quantized for the fp8 tensor cores."
Not a regression — the scale was discarded pre-PR too, and weight.to(bf16) was equally wrong — but the PR's entire premise ("raw fp8 = fp8 tensors with no weight_scale") is never actually checked here, and the new log asserts something the loader has no evidence for. Separately, the other spelling .weight_scale is not in keys_to_remove at all, so such a checkpoint dies in load_state_dict(..., strict=True) on unexpected keys.
Z-Image should go through extract_fp8_scaled_layers + attach_fp8_scales the way Krea-2 does.
🟡 5. make_room under-reserves
flux.py:657-660 and z_image.py:288-292 charge 1 byte/element for every fp8 tensor when keep_fp8, but cast_state_dict only preserves 2-D nn.Linear weights that don't match a skip pattern. Everything else lands at 2 bytes. On the Z-Image checkpoint from the PR description that's the 243 one-dimensional tensors, x_pad_token / cap_pad_token, and the t_embedder / cap_embedder Linears the loader deliberately dequantizes (cap_embedder alone is 2560×3840 ≈ 9.8M params).
The reservation is short by that amount — small relative to 6 GB, but it is a silent under-count in the one call whose job is to avoid RAM pressure. Cheapest fix is to have cast_state_dict predict/return the byte total using the same predicate.
Attacks that came back clean
Recording these so the surface that was checked is explicit:
for key in sd:with in-placesd[key] = ...— reassignment only, no resize, noRuntimeError.float8_e5m2exclusion is correct and consistent:cast_state_dictnever keeps it, and_can_use_fp8_matmulrejects it via!= FP8_DTYPE..weight_scale/.scale_weightkeys are not mistaken for.weightbykey.endswith(".weight").- Linears whose
in_features/out_featuresare not divisible by 16 fall back correctly via_can_use_fp8_matmul's alignment check, so keeping all Linears fp8 does not expose a_scaled_mmshape constraint. - The new
load_default.pyearly return does not regress a3784eb — the Qwen3-VL encoder calls_apply_fp8_to_nn_moduledirectly (krea2.py:760, 783), bypassing_apply_fp8_layerwise_casting. - Nor does it double-fire on the Krea-2 scaled path, which already returns at
krea2.py:484before reaching it. get_submoduleon a model built underinit_empty_weights()resolves fine (meta params);AttributeErroris the only exception it raises for a bad path.- FLUX's own modules do not cast activations to a weight's dtype (
timestep_embeddingkeys offt, notweight), so omittingskip_patternsin the FLUX loader is not the Z-ImageTimestepEmbedderhazard — a minor inconsistency only.
Verdict
Requesting changes. #1 is silent weight corruption on a path the base branch already shipped working; #2 turns fp8_compute into a hard mid-generation crash on RDNA3 and disables the very fallback meant to prevent that; #3 means roughly a third of the advertised scope never executes.
The direction is right and the measurements are convincing — this is about the filters being applied to the scaled path, and the device predicate being too permissive.
Conflicts in the fp8 loaders: - load_default.py: keep both guards. main added the idempotence early-return (the FP8_COMPUTE_DTYPE_ATTR marker), this branch added the "already fp8, leave it on the tensor cores" early-return. They gate different things, so both stay; the marker check runs first. - krea2.py: the transformer loader no longer calls _dequantize_scaled_fp8 — extract_fp8_scaled_layers/dequantize_fp8_scaled replaces it and honours the fp8_compute setting. main's RAM-spike fix to that helper is still relevant for its remaining callers, so the helper keeps main's dtype-aware version. target_device/model_dtype now come from main's hoisted position. - krea2.py Qwen3-VL encoder: same, plus main's .comfy_quant / scale_input key cleanup is already covered — extract_fp8_scaled_layers pops both.
…evice Review follow-ups for invoke-ai#9478. 1. A scaled fp8 layer that matched a skip pattern (or was not an nn.Linear weight) went through cast_state_dict's plain .to(dtype), which drops the weight_scale; attach_fp8_scales then skipped it for no longer being fp8, so the scale was lost and the weight ended up off by 1/weight_scale. Krea-2 hits this on ordinary ComfyUI exports — time_embed.linear_1/linear_2 are quantized like any other Linear and match the model's `time_embed` pattern. split_fp8_scaled_layers() now folds exactly those layers first, with the scale applied, and drops them from the mapping. The predicate the three callers share lives in can_stay_quantized(). attach_fp8_scales returning fewer than len(layers) now warns instead of reading as success. 2. device_supports_fp8_matmul tested compute capability >= (8, 9). On ROCm that reports the gfx arch, so RDNA3 (gfx1100 -> (11, 0)) passed and every forward then raised. It now probes once per device with a real 16x16 _scaled_mm; the capability compare is only a pre-filter. Without this the fallback that exists to avoid a mid-generation crash was gated on the same wrong answer. 3. Flux2CheckpointModel._dequantize_fp8_weights converted every float8 tensor unconditionally, before cast_state_dict ever saw the state dict — so the FLUX.2 half of the raw-fp8 path never executed and its log line could not fire. It now takes the same keep_fp8 gate. 4. Z-Image deleted .scale_weight without applying it and could not tell a scaled checkpoint from a raw one. It goes through extract_fp8_scaled_layers + attach_fp8_scales like Krea-2 now, reads the full_precision_matrix_mult hints from both the header and the .comfy_quant markers, and accepts both scale spellings. The fused-QKV split carries the quantization side-channel with it: a scale left on `...attention.qkv` keys onto a module the diffusers model does not have, so all three split weights would stay quantized but unscaled. 5. make_room charged 1 byte/element for every fp8 tensor, but only 2-D Linear weights outside the skip patterns actually stay quantized. On a checkpoint that quantized all 453 of its tensors that is most of the reservation. predict_cast_state_dict_size() answers with the same predicate the cast uses.
2 participants
Summary
A checkpoint can ship fp8 weights with no
weight_scale— plainfloat8_e4m3fntensors. The runtime already handles those:scaled_mm_lineartreatsweight_scaleas optional, andCustomLinear._can_use_fp8_matmulonly requires the fp8 dtype. The loaders never let them through, though — FLUX, FLUX.2 and Z-Image cast the whole state dict to bf16 unconditionally, so both the VRAM saving and the tensor cores are discarded 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.This adds
cast_state_dict(), which casts everything except fp8 weights when the matmul is actually available, and wires it into all four loaders with a log line so the tensor-core path is visible.Only
nn.Linear.weightis 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'. A model's own_skip_layerwise_casting_patternsis honored on top, for Linears whoseforwardcasts activations to their weight's dtype (Z-Image'sTimestepEmbedderdoes exactly that).float8_e5m2is deliberately not preserved:_scaled_mmcannot take it as the weight operand on Ada, so keeping it quantized would buy VRAM at the cost of a per-forward dequantize.Second change: 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.
Related Issues / Discussions
Stacked on
feat/fp8_scaled_compute(sharedfp8_scaled.pymodule + Krea-2 scaled support). Do not merge intomainbefore that branch.Overlaps with #9416 in
_should_use_fp8/_apply_fp8_layerwise_casting— the two guards are complementary but will conflict textually. #9416 catches GGUF/bnb payloads vianot is_floating_point() or type(...) is not torch.Tensor; an fp8 weight is floating point and a plain tensor, so it falls through that check. Whichever merges second needs a trivial rebase.QA Instructions
Requires an Ada (SM 8.9+) GPU and
fp8_compute: trueininvokeai.yaml.float8_e4m3fntensors but no.weight_scalekeys.Z-Image: kept N raw fp8 weight(s) quantized for the fp8 tensor cores.Total model sizein the[MODEL CACHE] Loaded modelline is roughly half what it was before.s/it.weight_scale) and a GGUF checkpoint must behave exactly as before.Measured on an RTX 4090, Z-Image
unstableRevolution_V2Fp8, 1024×1024, 30 steps, 1 warm-up + 3 measured runs:Note the drift on the
beforeside: 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.Merge Plan
Merge
feat/fp8_scaled_computefirst — this branch is stacked on it and its diff is meaningless standalone. Coordinate with #9416 (see above).Checklist
What's Newcopy (if doing a release after this PR)