feat: Support mxfp8#12907
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds MXFP8 support across the codebase: MixedPrecisionOps._load_from_state_dict gains a 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/ops.py`:
- Around line 804-818: pick_operations() must treat "mxfp8" like the other FP8
formats when fp8_compute is False: add "mxfp8" to the disabled set so
unsupported backends will trigger the fallback that sets _full_precision_mm =
True instead of continuing the FP8 load path; update the logic that builds the
disabled set (the function pick_operations and any variable named disabled or
similar) to include "mxfp8" whenever not fp8_compute, and ensure downstream
selection uses that disabled set so the MXFP8 branch in the load path (the code
constructing layout_cls.Params with torch.float8_e8m0fnu) is bypassed for
non-FP8-capable backends.
In `@comfy/quant_ops.py`:
- Around line 11-12: The module currently references TensorCoreMXFP8Layout
without a local definition so add a fallback stub for TensorCoreMXFP8Layout in
the existing import-fallback block (same place other layout stubs live) so the
name always exists when comfy_kitchen is absent; then guard its use by wrapping
the register_layout_op(...) call that mentions TensorCoreMXFP8Layout (and any
other registrations using it) with a check that the real comfy_kitchen-provided
class is present (or use a boolean flag set during import fallback) and remove
or conditionally include TensorCoreMXFP8Layout from the module export list
(__all__) so it is only exported when the real implementation is available.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 2fad1690-a310-4a61-a995-c7388ed86646
📒 Files selected for processing (2)
comfy/ops.pycomfy/quant_ops.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/ops.py`:
- Around line 804-818: The MXFP8 branch handling when self.quant_format ==
"mxfp8" uses torch.float8_e8m0fnu directly
(block_scale.view(torch.float8_e8m0fnu)), which will fail on PyTorch < 2.8 even
if supports_fp8_compute() returned true; add an explicit PyTorch version check
(e.g., compute torch_version_numeric >= (2,8) or similar) before using or
returning torch.float8_e8m0fnu in this code path and in any related places (see
get_supported_float8_types(), supports_fp8_compute()); if the PyTorch version is
too old, raise a clear error or fall back to a safe path rather than calling
block_scale.view with the missing dtype, and ensure layout_cls.Params is only
constructed with a valid dtype.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 40be2d6c-a9a7-4b46-93ec-e385006b5312
📒 Files selected for processing (2)
comfy/ops.pycomfy/quant_ops.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/quant_ops.py`:
- Around line 144-145: QUANT_ALGOS is adding the "mxfp8" entry unconditionally
while TensorCoreMXFP8Layout is only registered when _CK_AVAILABLE is true,
causing get_layout_class() to return None and later dereference
layout_cls.Params; update the code so that the "mxfp8" entry is only appended to
QUANT_ALGOS when _CK_AVAILABLE is true (same conditional that calls
register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)), ensuring
the availability check (_CK_AVAILABLE), the registration call
(register_layout_class/TensorCoreMXFP8Layout), and the QUANT_ALGOS update are
guarded together to prevent layout_cls from being None in comfy/ops.py where
layout_cls.Params is used.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: b5649e78-76c7-450b-9961-add3116d725d
📒 Files selected for processing (1)
comfy/quant_ops.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/float.py`:
- Line 247: The return uses the dtype symbol torch.float8_e8m0fnu directly which
breaks on older PyTorch versions; update the code around the return in
comfy/float.py (the to_blocked(...).view(...) call that uses block_scales_e8m0,
flatten=False and returns output_fp8) to safely resolve the dtype via a
try/except or getattr fallback (e.g., attempt to use torch.float8_e8m0fnu and
fall back to a compatible dtype or raise a clear error) so it matches the
defensive pattern used elsewhere (like comfy/model_management.py).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 54d0b618-c271-4eb1-9f92-d20bead18547
📒 Files selected for processing (2)
comfy/float.pycomfy/quant_ops.py
comfy/ops.py
Outdated
| if not fp8_compute: | ||
| disabled.add("float8_e4m3fn") | ||
| disabled.add("float8_e5m2") | ||
| disabled.add("mxfp8") |
There was a problem hiding this comment.
Did you check if this works on ada?
There was a problem hiding this comment.
Right, no it doesn't work... comfy-kitchen does handle it so outputs work with error spam, so better to disable for non-blackwell, adding that.
|
Successfully installed comfy-kitchen-0.2.8 and the error: |
Are you using a LoRA? If so, which one? Is it maybe saved in fp8? |
no any lora, even without distilled-lora, same error this is workflow |
Okay, it's the pytorch version, confirmed same error in 2.9.1 while in 2.10.0 it works. |
yes, it work now. thx kijai. Have a nice day! |
As this is already supported in comfy-kitchen. Blackwell only in practice.
Tested with
https://huggingface.co/Kijai/LTX2.3_comfy/blob/main/diffusion_models/experimental/ltx-2.3-22b-distilled_transformer_only_mxfp8_block32.safetensors
Example on 5090
BF16:
8/8 [02:09<00:00, 16.14s/it]bf16_00001-audio.mp4
mxfp8:
8/8 [01:17<00:00, 9.68s/it]ltx23_mxfp8_00001-audio.mp4