Skip to content

v3.4.3 - Z Image Hybrid ConvRot NVFP4: Tensor Core (TC / W4A4) opt-in

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 24 Aug 20:34
· 59 commits to main since this release
v3.4.3
9fee210
EN 中文

1. Overview

v3.4.3 adds an opt-in Tensor Core (TC / W4A4) execution path for the Z Image / ZIT Hybrid ConvRot NVFP4 models. Until now these models ran only on the Comfy parity path — W4A16 (NVFP4 weights × fp16 activations, stock GEMM + online act rotate) — which keeps the 4-bit weight compression but never uses the FP4 Tensor Core. With a calibrated per-layer input_scale, the Linear hot path now switches to W4A4 TC (NVFP4 weights × 4-bit rotated activations on the raw cublas_gemm_blockwise_fp4 GEMM), unlocking the real Blackwell FP4 Tensor Core throughput.

The path is gated by trajectory-fidelity validation (a deterministic per-step latent-divergence comparison): TC vs parity final-cosine is 0.951 vs 0.952 (noise-level difference, 0 / 20 bifurcations), so TC adds no systematic quality loss while delivering the Tensor Core speedup.

2. Parity vs TC

parity (previous default) TC (new opt-in)
Weights NVFP4 (E2M1 + block scale) NVFP4 (E2M1 + block scale)
Activations fp16 4-bit NVFP4 (rotated, E2M1)
GEMM stock a @ w.T (fp16) raw cublas_gemm_blockwise_fp4 (FP4)
Tensor Core FP4 ✗ ✓
Prerequisite — calibrated input_scale

Parity was the default because TC needs an activation quantization scale (input_scale) that the earlier published checkpoints did not carry. Forcing TC on an uncalibrated checkpoint collapses quality (decoded SSIM ≈ 0.18) — this is why the path is opt-in and gated.

3. input_scale calibration

  • input_scale = amax / 2688, measured in the rotated domain (the same domain the ConvRot weights are stored in).
  • Produced by a standalone step: Z_Image/calib_input_scale_nvfp4.py — it samples the model's per-layer activation amax and writes *.input_scale keys into a *_calib.safetensors checkpoint.
  • Measured, not searched: input_scale is a running absmax ÷ 2688, not a histogram / quality knob — there is no per-layer search.
  • The loader reads *.input_scale from the checkpoint to decide TC eligibility.

4. Loader opt-in

  • Priority: HSWQ_ZI_FORCE_PARITY=1 > HSWQ_ZI_FORCE_TC=1 > auto-detect *.input_scale.
  • checkpoint_has_input_scale() detects the calibrated keys; zi_use_tensorcore() resolves the effective mode.
  • HSWQ_ZI_FORCE_PARITY=1 keeps the old W4A16 behaviour; HSWQ_ZI_FORCE_TC=1 forces W4A4 even without calibration (intentionally allowed, but collapses quality on uncalibrated files).
  • A warning-filter patch (_patch_load_model_weights_warnings()) silences the expected input_scale / comfy_quant load warnings that the parity path previously spammed.

5. Quality gate — trajectory divergence

  • Tool: benchmark/zi_convrot_nvfp4_traj_compare.py — a deterministic per-step latent trajectory comparator (torch.backends.cudnn.deterministic = True, fixed seeds), measuring the final-latent cosine vs the reference.
  • Bifurcation = a single-step cosine drop > 0.05 (divergence from the reference trajectory).
  • Same-image threshold = final-cos ≥ 0.98.
  • Results (hybrid nv89, 20 seeds):
    • TC: final-cos mean 0.95137, min 0.86664, 0 / 20 bifurcations
    • parity: final-cos mean 0.95233, min 0.86509, 0 / 20 bifurcations
    • → TC ≈ parity (mean difference 0.001, noise-level).
  • Contrast — native full-NVFP4 (180 layers, no INT8 protection): mean 0.90945, min 0.60344, 1 / 20 bifurcations (seed 12 diverged). HSWQ hybrid (INT8-protected high-impact layers) removes the bifurcation.

6. Benchmark (test5 = moodyProMix nv90 config)

  • Decoded SSIM 0.9772
  • VRAM 12368 → 5935 MB (−52%) vs FP16
  • Wall time 3.39 s vs 7.18 s (2.12×)
  • TC GEMM hits 2160, dequant fallbacks 0
  • 160.5 TFLOPS (~18% of Blackwell FP4 peak)

7. Files added / modified

Node repo (ComfyUI-HSWQ-Loader-and-Tools):

Kind File Change
Modified nodes/zimage_nvfp4/load_unet.py TC (W4A4) opt-in gating (checkpoint_has_input_scale() / zi_use_tensorcore()), HSWQ_ZI_FORCE_TC=1 override, input_scale/comfy_quant load-warning filter
Modified nodes/zimage_nvfp4/zi_nvfp4_forward.py accumulate TC GEMM FLOPs (_TC_FLOPS); nvfp4_forward_stats() returns tc_flops

Upstream HSWQ repo (Hybrid-Sensitivity-Weighted-Quantization) — tools / validation:

Kind File Purpose
New Z_Image/calib_input_scale_nvfp4.py measure input_scale = amax / 2688 (rotated domain)
New benchmark/zi_convrot_nvfp4_bench_v3.py --tc force-TensorCore bench + TFLOPS
New benchmark/zi_convrot_nvfp4_traj_compare.py deterministic trajectory divergence comparator
Modified nvfp4_addmm_patch.py / nvfp4_comfy_parity.py (vendored) GEMM-mode counters (scaled_mm hits vs dequant fallbacks; parity NVFP4/INT8 forward)

8. Links

Assets 2
Loading