Support Online FP8 & NF4 Weight Quantization

[Uxito-Ada]

Description

Add online FP8 & NF4 weight-only quantization for DiT pipelines with torchao and bitsandbytes backend:

• introduce QuantType.TORCHAO_FP8 and QuantType.BNB_NF4
• Qwen-Image FP8 and NF4 examples on H100

Motivation

Online weight-only quantization is necessary especially for models not natively providing quantized checkpoints:

• FP8 can reduce memory size to some extend and keep high quality at the same.
• NF4 can sharply reduce memory size further, which is useful for low-memory GPUs, and also keep high quality compared with lossy solution e.g. torchao INT4.
  Detailed data and comparison can be found in the following.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Performance improvement
• Code refactoring
• Documentation update
• Other (please describe):

Changes Made

• Add QuantType.TORCHAO_FP8, QuantType.BNB_NF4, QuantKernelBackend.TORCHAO and QuantKernelBackend.BITSANDBYTES
• Add telefuser.ops.torchao_fp8_linear and telefuser.ops.bnb_nf4_linear
• Added qwen_image_t2i_torchao_fp8_h100.py and qwen_image_t2i_bnb_nf4_h100.py , as well as according documents

Benchmark Setup

• GPU: NVIDIA H100 80GB HBM3
• Model: Qwen-Image-2512
• Prompt: A cat playing piano
• Resolution: 1328x1328, 928x1664
• Steps: 8, 16, 32
• Loop: 2 warmup, 5 repeats
• Dtype:
  • BF16 (baseline)
  • Qwen-Image Lightning FP8
  • BNB NF4
  • TorchAO INT4
  • TorchAO FP8
  • llmcompressor FP8
• Metrics:
  • Latency/Throughput
  • Memory
  • Quality

Performance Results

• Though the performance of torchao FP8 is slightly weaker than Lightning FP8, it is still necessary for models without FP8 checkpoint.
• llmcompressor is not chosen is because that it does not save any memory: it is found that llmcompressor does not convert much Qwen-Image's DiT BF16 linears to FP8. Although llmcompressor declares that it has W8A8, it shows much slower throughput, even worse than the other low bits.
• BNB NF4 is chosen rather than Tochao NF4 is because that it shows much better image quality.

Image Quality Comparison

Testing

• Unit tests pass (pytest tests/)
• Manual testing performed
• Benchmarks added/updated (if applicable)

Test commands:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python examples/qwen_image/qwen_image_t2i_torchao_fp8_h100.py \
  --model_root Qwen/Qwen-Image-2512 \
  --prompt "A cat playing piano" \
  --aspect_ratio 1:1 \
  --num_inference_steps 16 \
  --seed 42

cd /root/heyang/TeleFuser
source .venv/bin/activate
uv pip install -U bitsandbytes

export LD_LIBRARY_PATH=$VIRTUAL_ENV/lib/python3.10/site-packages/nvidia/nvjitlink/lib:$VIRTUAL_ENV/lib/python3.10/site-packages/nvidia/cuda_runtime/lib:$VIRTUAL_ENV/lib/python3.10/site-packages/nvidia/cublas/lib:$LD_LIBRARY_PATH

CUDA_VISIBLE_DEVICES=3 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python examples/qwen_image/qwen_image_t2i_bnb_nf4_h100.py   --model_root /hhb-data/aigc/model_zoo/Qwen-Image-2512   --prompt "A cat playing piano"   --aspect_ratio 1:1   --num-inference-steps 16   --seed 42   --output qwen_image_bnb_nf4.png

Checklist

• Code follows the project's coding standards (ruff)
• Pre-commit hooks pass (pre-commit run --all-files)
• All tests pass (pytest tests/)
• New tests added for new functionality
• Documentation updated (example and PR notes)
• Commit messages are clear and descriptive
• PR title follows the convention: [TYPE] Brief description

GPU Architecture Support

• SM80 (Ampere, Ada Lovelace)
• SM90 (Hopper H100)
• SM100+ (Blackwell)

[Uxito-Ada]

WIP on documents