[Task] Support SVDQuant MXFP4 kernel in Nunchaku.

[chensuyue]

The goal of this task is to successfully run the Nunchaku MXFP4 kernel, as Nunchaku currently only supports NVFP4 and INT4; the next step is to validate the accuracy of the diffusion model using the newly added MXFP4 kernel.

[changwangss]

generate mxfp4 model by https://github.com/changwangss/deepcompressor/tree/mxfp4
calibration,

python -m deepcompressor.app.diffusion.dataset.collect.calib \
 configs/model/flux.1-dev.yaml \
 configs/collect/qdiff.yaml \
 --pipeline-path /home/user2/data/xixi/FLUX.1-dev \
 --collect-root /home/user2/data/xixi/deepcompressor-datasets \
 --collect-num-samples 1 \
 --eval-num-steps 1 \
 --eval-guidance-scale 3.5 \
 --eval-batch-size 1

quantization

python -m deepcompressor.app.diffusion.ptq \
 configs/model/flux.1-dev.yaml \
 configs/svdquant/mxfp4.yaml \
 --pipeline-path /home/user2/data/xixi/FLUX.1-dev \
 --output-root /home/user2/data/xixi/deepcompressor-runs \
 --output-dirname flux-dev-mxfp4-lowrank32-nosmooth \
 --enable-smooth false \
 --wgts-enable-low-rank true \
 --wgts-low-rank-rank 32 \
 --wgts-low-rank-num-iters 1 \
 --wgts-low-rank-early-stop false \
 --calib-num-samples 1 \
 --calib-path /home/user2/data/xixi/deepcompressor-datasets/torch.bfloat16/flux.1-dev/fmeuler1-g3.5/qdiff/s1/caches \
 --calib-batch-size 1 \
 --wgts-low-rank-sample-batch-size 1 \
 --wgts-calib-range-sample-batch-size 1 \
 --wgts-calib-range-element-batch-size 16 \
 --ipts-calib-range-sample-batch-size 1 \
 --ipts-calib-range-element-batch-size 16 \
 --eval-num-samples 0 \
 --skip-eval true \
 --skip-gen true \
 --save-model /home/user2/data/xixi/deepcompressor-runs/flux-dev-mxfp4-lowrank32-nosmooth/checkpoint

convert nunchaku onefile format

python -m deepcompressor.backend.nunchaku.convert \
 --quant-path /home/user2/data/xixi/deepcompressor-runs/flux-dev-mxfp4-lowrank32-nosmooth/checkpoint \
 --output-root /home/user2/data/xixi/deepcompressor-runs/flux-dev-mxfp4-lowrank32-nosmooth/nunchaku \
 --model-name flux.1-dev
cp /home/user2/data/xixi/deepcompressor-runs/flux-dev-mxfp4-lowrank32-nosmooth/nunchaku/flux.1-dev/transformer_blocks.safetensors \
 /home/user2/data/xixi/flux.1-dev-mxfp4-lowrank32-nosmooth.safetensors

add nunchaku mxfp4 kernel support for 5090d(sm120): https://github.com/changwangss/nunchaku/tree/mxfp4

import argparse
import gc
import os
from pathlib import Path
import torch
from diffusers import FluxPipeline

BASE_MODEL = Path(os.environ.get("BASE_MODEL", "/home/user2/data/xixi/FLUX.1-dev"))
MXFP4_MODEL = Path(os.environ.get("MXFP4_MODEL", "/home/user2/data/xixi/flux.1-dev-mxfp4-lowrank-smoke.safetensors"))
NVFP4_MODEL = os.environ.get(
   "NVFP4_MODEL",
   "nunchaku-tech/nunchaku-flux.1-dev/svdq-fp4_r32-flux.1-dev.safetensors",
)
PROMPT = os.environ.get("PROMPT", "A cat holding a sign that says hello world")

def cleanup() -> None:
   gc.collect()
   if torch.cuda.is_available():
       torch.cuda.empty_cache()
       torch.cuda.ipc_collect()

def make_pipe(mode: str, offload: bool) -> FluxPipeline:
   if mode == "bf16":
       pipe = FluxPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16, local_files_only=True)
   else:
       from nunchaku import NunchakuFluxTransformer2dModel
       if mode == "mxfp4":
           transformer = NunchakuFluxTransformer2dModel.from_pretrained(
               MXFP4_MODEL,
               torch_dtype=torch.bfloat16,
               precision="mxfp4",
               device="cuda:0",
           )
       elif mode == "nvfp4":
           transformer = NunchakuFluxTransformer2dModel.from_pretrained(
               NVFP4_MODEL,
               torch_dtype=torch.bfloat16,
               precision="fp4",
               device="cuda:0",
           )
       else:
           raise ValueError(f"unknown mode: {mode}")
       pipe = FluxPipeline.from_pretrained(
           BASE_MODEL,
           transformer=transformer,
           torch_dtype=torch.bfloat16,
           local_files_only=True,
       )
   if offload:
       pipe.enable_model_cpu_offload(gpu_id=0)
   else:
[pipe.to](http://pipe.to/)("cuda:0")
   pipe.set_progress_bar_config(disable=False)
   return pipe

def main() -> None:
   parser = argparse.ArgumentParser()
   parser.add_argument("--mode", choices=["bf16", "mxfp4", "nvfp4"], required=True)
   parser.add_argument("--steps", type=int, default=int(os.environ.get("STEPS", "20")))
   parser.add_argument("--height", type=int, default=int(os.environ.get("HEIGHT", "512")))
   parser.add_argument("--width", type=int, default=int(os.environ.get("WIDTH", "512")))
   parser.add_argument("--seed", type=int, default=int(os.environ.get("SEED", "12345")))
   parser.add_argument("--guidance-scale", type=float, default=float(os.environ.get("GUIDANCE_SCALE", "3.5")))
   parser.add_argument("--out", type=Path, required=True)
   parser.add_argument("--offload", action="store_true")
   args = parser.parse_args()
   cleanup()
   pipe = make_pipe(args.mode, args.offload)
   generator = torch.Generator(device="cuda:0").manual_seed(args.seed)
   image = pipe(
       PROMPT,
       num_inference_steps=args.steps,
       guidance_scale=args.guidance_scale,
       generator=generator,
       height=args.height,
       width=args.width,
   ).images[0]
   args.out.parent.mkdir(parents=True, exist_ok=True)
   image.save(args.out)
   print(args.out)
   del pipe
   cleanup()

if __name__ == "__main__":
   main()

bf16

mxfp4

nvfp4