Extensions/ComfyUI-fp8_cublas
ComfyUI Extension

ComfyUI-fp8_cublas

ComfyUI custom node that patches loaded models (Flux, SDXL, etc.) to run Linear layers in FP8 or MXFP8, unlocking tensor-core throughput on Ada Lovelace (RTX 4090) and…

By marduk191·Created 4 months ago·Updated 4 months ago· 0
marduk191/ComfyUI-fp8_cublas
Nodes1
On cloudLocal install
Categoryoptimization
Stars0
Updated4 months ago

Nodes (1)

FP8 / MXFP8 Model Patch (SM89+)
optimization
Readme

fp8_cublas — FP8 / MXFP8 Model Patch for ComfyUI

ComfyUI custom node that patches loaded models (Flux, SDXL, etc.) to run Linear layers in FP8 or MXFP8, unlocking tensor-core throughput on Ada Lovelace (RTX 4090) and Blackwell (RTX 5090) GPUs.

Inspired by: — the technique of using FP16 accumulators for FP8 GEMM on SM89, now accessible via torch._scaled_mm(use_fast_accum=True) and the IST-DASLab gemm-fp8 CUTLASS kernel.

Hardware Support

| GPU | SM | FP8 backend | MXFP8 | Notes | |-----|----|------------|-------|-------| | RTX 4090 | 89 | gemm_fp8 (CUTLASS, SM89-tuned) | — | Requires optional install | | RTX 4090 | 89 | torch._scaled_mm (fallback) | — | Used if gemm-fp8 not installed | | RTX 5090 | 120 | torch._scaled_mm (native Blackwell) | ✅ | No extra install needed | | RTX 3090 / older | <89 | ❌ not supported | ❌ | FP8 tensor cores require SM89+ |

Performance

RTX 4090 (SM89) — FP8 GEMM throughput

xychart-beta
    title "RTX 4090 FP8 GEMM Throughput (TFLOPS)"
    x-axis ["BF16 baseline", "cuBLASLt FP32 accum", "torch._scaled_mm", "gemm_fp8"]
    y-axis "TFLOPS" 0 --> 500
    bar [82, 330, 380, 473]

RTX 5090 (SM120) — FP8 / MXFP8 throughput

xychart-beta
    title "RTX 5090 FP8 / MXFP8 Throughput (TFLOPS)"
    x-axis ["BF16 baseline", "FP8 (torch._scaled_mm)", "MXFP8 (torchao)"]
    y-axis "TFLOPS" 0 --> 900
    bar [210, 838, 838]

MXFP8 matches FP8 throughput on SM120 while offering better quantization accuracy (per-block scaling: 1 scale per 32 elements vs 1 scale per tensor).

VRAM — weight storage per 1B parameters

xychart-beta
    title "VRAM for Weights (GB per 1B parameters)"
    x-axis ["BF16", "FP8", "MXFP8", "FP4 (future)"]
    y-axis "GB" 0 --> 2.5
    bar [2.0, 1.0, 1.0, 0.5]

FP8 and MXFP8 both halve weight VRAM vs BF16.

Installation

1. Install the node

Via ComfyUI Manager (search for fp8_cublas) or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/marduk191/ComfyUI-fp8_cublas

2. Install torchao (MXFP8 / RTX 5090)

torchao is bundled with recent ComfyUI. If missing:

pip install torchao>=0.10.0

3. Install gemm-fp8 (RTX 4090 only — optional but recommended)

Enables the IST-DASLab CUTLASS kernel that matches the 473 TFLOPS technique. Without it, the node falls back to torch._scaled_mm (still ~4.6× faster than BF16).

git clone --recurse-submodules https://github.com/IST-DASLab/gemm-fp8
cd gemm-fp8
pip install -e .

You'll see a console warning at ComfyUI startup if you have an RTX 4090 but gemm-fp8 is not installed.

Usage

Find the node in ComfyUI under optimization → FP8 / MXFP8 Model Patch (SM89+).

[Load Checkpoint] ──► [FP8 / MXFP8 Model Patch] ──► [KSampler]
                              ▲
                        mode: fp8 | mxfp8
                  accumulator: fp16 | fp32

Node Inputs

| Input | Type | Options | Description | |-------|------|---------|-------------| | model | MODEL | — | Any ComfyUI model (Flux, SDXL, SD3, etc.) | | mode | combo | fp8, mxfp8 | Quantization format | | accumulator | combo | fp16, fp32 | FP8 mode only — FP16 = max speed, FP32 = safer numerics |

Mode Reference

| Mode | Scaling | Hardware | Accuracy | Speed | |------|---------|----------|----------|-------| | fp8 | Per-tensor (1 scale/tensor) | SM89+ | Good | ★★★★☆ | | mxfp8 | Per-block (1 scale/32 elems) | SM100+ (RTX 5090) | Better | ★★★★★ |

Recommended Settings by GPU

| GPU | Mode | Accumulator | |-----|------|-------------| | RTX 4090 | fp8 | fp16 | | RTX 5090 | mxfp8 | — (N/A) | | RTX 5090 (speed priority) | fp8 | fp16 |

SageAttention Compatibility

Safe to combine — they patch different operations:

  • FP8ModelPatch patches nn.Linear layers (projections, feed-forward)
  • Patch Sage Attention patches the attention SDPA kernel

Recommended combo for RTX 5090:

[Model] → [FP8ModelPatch  mode=mxfp8] → [Patch Sage Attention  sageattn_qk_int8_pv_fp8_cuda++] → [KSampler]

Apply FP8ModelPatch first, then SageAttention on top.

Console Output

On startup you'll see which backend was selected:

# RTX 5090 — optimal path, no action needed
[FP8ModelPatch] Device: NVIDIA GeForce RTX 5090 (SM120) | mode=mxfp8

# RTX 4090 + gemm-fp8 installed —  technique active
[FP8ModelPatch] Patched 256 layers → FP8 | backend=gemm_fp8 (SM89) | accum=torch.float16

# RTX 4090 without gemm-fp8 — warns and falls back
[FP8ModelPatch] WARNING: RTX 4090 (SM89) detected but gemm-fp8 is not installed.
  For the optimized SM89 FP8 kernel, install:
    git clone --recurse-submodules https://github.com/IST-DASLab/gemm-fp8
    cd gemm-fp8 && pip install -e .
  Falling back to torch._scaled_mm — still faster than BF16, but not peak SM89.

Credits