Quantization

How ComfyUI represents and runs quantized weights, primarily FP8. The full design and rationale is in QUANTIZATION.md; this page summarizes what that document says and where each piece lives in code.

Purpose

Run large diffusion models on smaller GPUs by storing weights in lower precision (FP8 today, more layouts in the future). The system needs to:

• Carry quantization metadata (scales, layouts) alongside tensors.
• Allow per-layer precision so sensitive layers can stay full-precision.
• Dispatch optimized kernels for common ops (linear, matmul) and fall back to dequantize-then-default for everything else.

Layout

comfy/
├── quant_ops.py         # QuantizedTensor, QuantizedLayout, two-level op registry
├── ops.py               # MixedPrecisionOps + the integration with model loading
└── float.py             # FP8 cast helpers
QUANTIZATION.md          # The design doc

Key abstractions

Type / function	File	What it is
QuantizedTensor	comfy/quant_ops.py	A torch.Tensor subclass that uses __torch_dispatch__ to route ops through registries
QuantizedLayout	comfy/quant_ops.py	The base class for "how is this format represented" — defines quantize + dequantize
register_layout_op	comfy/quant_ops.py	Decorator to register a fast-path op handler for a (op, layout) pair
Generic op registry	comfy/quant_ops.py	Routes layout-agnostic ops (.to, .clone, .reshape, …)
QUANT_ALGOS	comfy/quant_ops.py	Map from format string (e.g., "float8_e4m3fn") to QuantizedLayout
MixedPrecisionOps	comfy/ops.py	nn.Module mode that loads selected layers as QuantizedTensors
Linear._load_from_state_dict	MixedPrecisionOps	The hook that turns weights into QuantizedTensor at load time
pick_operations	comfy/ops.py	Picks MixedPrecisionOps when a model config has layer_quant_config

How a quantized checkpoint loads

sequenceDiagram
    participant Loader as comfy.sd.load_checkpoint
    participant Detect as model_detection
    participant Cfg as ModelXXX (supported_models)
    participant Ops as pick_operations (comfy.ops)
    participant MP as MixedPrecisionOps.Linear
    participant Tensor as QuantizedTensor
    Loader->>Detect: detect by state-dict shape
    Detect-->>Loader: ModelConfig (with layer_quant_config from metadata)
    Loader->>Ops: pick based on config
    Ops-->>Loader: MixedPrecisionOps
    Loader->>MP: build Linear(s) using MixedPrecisionOps
    Loader->>MP: load_state_dict
    MP->>MP: for each layer, check _layer_quant_config
    alt layer is quantized
        MP->>Tensor: build QuantizedTensor with weight + scale + …
    else layer is full precision
        MP->>MP: load weight as plain tensor in compute dtype
    end

The detection step reads _quantization_metadata from the safetensors metadata. If present, the model config is annotated with a layer_quant_config mapping layer_name → {format, layout-specific params}. pick_operations sees that and selects MixedPrecisionOps.

At inference time

When a node calls torch.nn.functional.linear(x, weight) and weight is a QuantizedTensor, __torch_dispatch__ routes the call through the registered linear handler for that layout. The handler can:

• Call an FP8 matmul kernel (when the device supports it and the activations meet alignment requirements).
• Or dequantize and fall back to the default linear.

For ops without a registered fast path, the default behavior is dequantize(QuantizedTensor) → high_precision_tensor → run_op. This keeps every op working even when no kernel is available.

Scaling parameters

The four scaling parameters defined in QUANTIZATION.md:

Parameter	What it is
weight_scale	Per-tensor or per-channel scale for the weight
weight_scale_2	Global scale for double-scaling recipes
pre_quant_scale	Smoothing scale for salient weights
input_scale	Quantization scale for activations

Activation scales (input_scale) require post-training calibration because they depend on actual inputs. Weight scales come for free from absmax(weight) / dynamic_range.

Checkpoint format

Quantized checkpoints are still ordinary safetensors files:

• Quantized weights are stored in their storage dtype (e.g., uint8 for packed FP8).
• Each quantized tensor has companion *_scale tensors stored alongside.
• The metadata dict carries _quantization_metadata JSON: format_version, layers (per-layer format).

comfy.utils.load_torch_file parses the metadata; the model loader pulls out the _quantization_metadata block and passes it through model detection.

Why per-layer

Not every layer survives quantization equally. The first/last projections, normalization layers, and small attention output projections often need full precision; the heavy MLP and attention matmuls are the wins. MixedPrecisionOps lets a checkpoint encode exactly that decision.

This is also a forward-compatibility hook: future layouts (NF4, INT8, double-scaled FP8) plug into the same QuantizedLayout abstraction without changes to the executor or the rest of the model code.

Integration points

• Activated in pick_operations (comfy/ops.py) when a model config carries layer_quant_config.
• Quantized tensors flow through the rest of the engine via __torch_dispatch__, so Sampling pipeline and Model management work unchanged.
• Tested in tests-unit/comfy_quant/.

Where to start a change

• Adding a layout: subclass QuantizedLayout in comfy/quant_ops.py; add an entry to QUANT_ALGOS; register fast-path ops via @register_layout_op for the most-used ops (linear, matmul, to).
• Adding compute support for new hardware: extend the FP8 cast helpers in comfy/float.py and the dispatch in MixedPrecisionOps Linear.forward.
• Producing quantized checkpoints: external — see the calibration workflow described in QUANTIZATION.md.