nvfp4: static input-scale calibration (capture once, inject forever)

Justified by live data: sampler phase timing shows sampling at 89% of
the chain node, all 792 quantized layers on the fast path, and the probe
puts dynamic input quantization at ~0.30 ms/layer-call vs ~0.11 ms with
a provided scale (amax dominates) - a top-line per-call cost.

New input_scales widget on the loader:
- "auto-calibrate", first load without a sidecar: forward-pre-hooks
  record each quantized layer's running input amax (one generation is
  enough; sidecar under refinements/nvfp4_scales/<checkpoint>.json,
  per-machine and rebuilt in one run whenever absent).
- every later load: injects input_scale = amax * 1.25 / (448*6) into
  the state dict next to each quantized weight - core's loader
  registers it natively and skips the per-call amax from then on.
  Margin leaves clip headroom (NVFP4 clips above the scale).

Tests: suffix-matched injection across checkpoint prefixes, value math,
core Linear registering input_scale and forwarding, sidecar read.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

nvfp4_patch.py

@@ -38,10 +38,123 @@ from __future__ import annotations
 
 import json
 import logging
+import os
 import re
 
 import torch
 
+# ---------------------------------------------------------------------------
+# Static input-scale calibration
+# ---------------------------------------------------------------------------
+# Core quantizes each layer's INPUT dynamically (amax reduction + quantize) on every
+# call unless the layer has a static `input_scale` parameter. Live probe on sm_120:
+# the dynamic path costs ~0.30 ms/layer-call and a provided scale cuts it to ~0.11 ms
+# (amax dominates) - across ~800 quantized layers that is a top-line cost of every
+# model call. Calibration is capture-then-freeze: a "calibrate" run records each
+# quantized layer's input amax via forward-pre-hooks into a per-checkpoint sidecar;
+# later loads inject `input_scale = amax * MARGIN / (448 * 6)` into the state dict,
+# where core's loader picks it up natively. MARGIN leaves headroom because NVFP4
+# clips activations above the scale; per the disposable-keys philosophy the sidecar
+# is per-machine and rebuilt in one run whenever absent.
+CALIBRATION_MARGIN = 1.25
+_F8F4_MAX = 448.0 * 6.0  # F8_E4M3_MAX * F4_E2M1_MAX
+INPUT_SCALE_MODES = ["dynamic", "auto-calibrate"]
+
+
+def _scales_path(checkpoint_path):
+    base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "nvfp4_scales")
+    return os.path.join(base, os.path.basename(checkpoint_path) + ".json")
+
+
+def load_calibration(checkpoint_path):
+    """{module_path: input_amax} recorded by a previous calibrate run, or None."""
+    try:
+        p = _scales_path(checkpoint_path)
+        if not os.path.exists(p):
+            return None
+        with open(p, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        amax = data.get("amax") or {}
+        return {str(k): float(v) for k, v in amax.items()} or None
+    except Exception as e:  # noqa: BLE001
+        logging.warning(f"[FunPack NVFP4] could not read calibration sidecar: {e}")
+        return None
+
+
+def inject_input_scales(sd, target_weight_keys, calibration, margin=CALIBRATION_MARGIN):
+    """Add `<layer>.input_scale` next to each quantized weight whose module path has a
+    calibrated amax. Keys are matched by suffix so any checkpoint prefix
+    ('model.diffusion_model.' or none) works. Returns the number injected."""
+    injected = 0
+    for k in target_weight_keys:
+        prefix = k[: -len("weight")]           # '...transformer_blocks.3.attn1.to_q.'
+        module_path = prefix.rstrip(".")
+        for cal_key, amax in calibration.items():
+            if module_path == cal_key or module_path.endswith("." + cal_key):
+                sd[prefix + "input_scale"] = torch.tensor(
+                    float(amax) * margin / _F8F4_MAX, dtype=torch.float32)
+                injected += 1
+                break
+    return injected
+
+
+def install_calibration_capture(model, checkpoint_path):
+    """Forward-pre-hooks on every quantized Linear recording the running input amax.
+    The sidecar is written on the first full sweep and refreshed periodically; hooks
+    live for the model instance's lifetime (a calibration run is expected to be
+    slightly SLOWER - it pays the amax it is trying to eliminate, once)."""
+    store = {"amax": {}, "writes": 0, "calls": 0}
+    path = _scales_path(checkpoint_path)
+
+    def _flush():
+        try:
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            tmp = path + ".tmp"
+            with open(tmp, "w", encoding="utf-8") as f:
+                json.dump({"amax": store["amax"], "margin": CALIBRATION_MARGIN}, f)
+            os.replace(tmp, path)
+            store["writes"] += 1
+        except Exception as e:  # noqa: BLE001
+            logging.warning(f"[FunPack NVFP4] calibration flush failed: {e}")
+
+    hooked = 0
+    try:
+        for name, m in model.model.diffusion_model.named_modules():
+            if getattr(m, "quant_format", None) != "nvfp4":
+                continue
+
+            def _hook(_module, args, _name=name):
+                try:
+                    x = args[0] if args else None
+                    if isinstance(x, torch.Tensor):
+                        a = float(x.detach().abs().amax())
+                        if a > store["amax"].get(_name, 0.0):
+                            store["amax"][_name] = a
+                    store["calls"] += 1
+                    if store["calls"] % 4096 == 0:
+                        _flush()
+                except Exception:
+                    pass
+
+            m.register_forward_pre_hook(_hook)
+            hooked += 1
+    except Exception as e:  # noqa: BLE001
+        logging.warning(f"[FunPack NVFP4] calibration capture install failed: {e}")
+        return 0
+
+    if hooked:
+        # First flush after one full model call so a single generation is enough.
+        def _first_flush_hook(_module, _args):
+            if store["writes"] == 0 and len(store["amax"]) >= hooked:
+                _flush()
+        try:
+            last = [m for _, m in model.model.diffusion_model.named_modules()
+                    if getattr(m, "quant_format", None) == "nvfp4"][-1]
+            last.register_forward_pre_hook(_first_flush_hook)
+        except Exception:
+            pass
+    return hooked
+
 QUANTIZE_SCOPES = [
     "video blocks",         # video-branch Linears in middle blocks (default, safest)
     "video + cross-modal",  # + audio_to_video_attn / video_to_audio_attn bridges
@@ -330,10 +443,19 @@ class FunPackNVFP4ModelLoader:
                                               "tooltip": "First N transformer blocks stay full precision."}),
                 "keep_last_blocks": ("INT", {"default": 2, "min": 0, "max": 16,
                                              "tooltip": "Last N transformer blocks stay full precision."}),
+                "input_scales": (INPUT_SCALE_MODES, {
+                    "default": "dynamic",
+                    "tooltip": "auto-calibrate: first run records each layer's input range "
+                               "(slightly slower) into a per-checkpoint sidecar; every later "
+                               "load injects static input scales, removing the per-call amax "
+                               "cost (~0.2 ms x every quantized layer x every step). Reload "
+                               "the model after the calibration run to activate.",
+                }),
             }
         }
 
-    def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks):
+    def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks,
+             input_scales="dynamic"):
         import comfy.model_management as mm
         import comfy.sd
         import comfy.utils
@@ -377,9 +499,27 @@ class FunPackNVFP4ModelLoader:
             f"(scope='{quantize_scope}', full-precision blocks: {sorted(keep) or 'none'})."
         )
 
+        calibration = load_calibration(path) if input_scales == "auto-calibrate" else None
+        if calibration:
+            injected = inject_input_scales(sd, targets, calibration)
+            logging.info(
+                f"[FunPack NVFP4] static input scales ACTIVE on {injected}/{len(targets)} layers "
+                f"(calibrated {len(calibration)}, margin {CALIBRATION_MARGIN}x) — per-call input "
+                "amax eliminated. Delete the sidecar under refinements/nvfp4_scales/ to recalibrate."
+            )
+
         model = comfy.sd.load_diffusion_model_state_dict(sd, metadata=metadata)
         if model is None:
             raise RuntimeError("ComfyUI could not detect the model type after NVFP4 patching.")
+
+        if input_scales == "auto-calibrate" and not calibration:
+            hooked = install_calibration_capture(model, path)
+            logging.info(
+                f"[FunPack NVFP4] CALIBRATING input scales on {hooked} layers this session "
+                "(slightly slower). Run one full generation, then reload the model (change any "
+                "widget or restart) — static scales activate on the next load."
+            )
+
         try:
             log_nvfp4_diagnostics(model, device)
         except Exception as e:  # noqa: BLE001

tests/test_nvfp4_patch.py

@@ -187,3 +187,55 @@ def test_diagnostics_reports_layer_census(capsys):
     assert "nvfp4 layers in loaded model: 3" in text
     assert "EMULATION" in text
     assert "GEMM probe skipped" in text or "GEMM probe" in text
+
+
+def test_input_scale_injection_and_core_load(tmp_path, monkeypatch):
+    """Calibration sidecar -> inject_input_scales matches by suffix across checkpoint
+    prefixes, and core's MixedPrecisionOps Linear registers the injected input_scale
+    and still forwards correctly."""
+    from nvfp4_patch import inject_input_scales, _scales_path
+    import nvfp4_patch as np_mod
+
+    ckpt = str(tmp_path / "model.safetensors")
+    sidecar = _scales_path(ckpt)
+    assert sidecar.endswith("model.safetensors.json")
+
+    cal = {"transformer_blocks.1.attn1.to_q": 12.5}
+    w = _w()
+    sd = {"model.diffusion_model.transformer_blocks.1.attn1.to_q.weight": w,
+          "model.diffusion_model.transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)}
+    out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0)
+    n = inject_input_scales(out, targets, cal)
+    assert n == 1
+    key = "model.diffusion_model.transformer_blocks.1.attn1.to_q.input_scale"
+    assert key in out and out[key].dtype == torch.float32
+    assert float(out[key]) == pytest.approx(12.5 * np_mod.CALIBRATION_MARGIN / (448.0 * 6.0))
+
+    # core load consumes it: input_scale becomes a registered parameter
+    import comfy.ops
+    ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=[])
+    lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16)
+    layer_sd = {k.replace("model.diffusion_model.transformer_blocks.1.attn1.to_q.", ""): v
+                for k, v in out.items()}
+    missing, unexpected = lin.load_state_dict(layer_sd, strict=False)
+    assert not missing and not unexpected
+    assert isinstance(getattr(lin, "input_scale", None), torch.nn.Parameter)
+    x = torch.randn(2, 4, D, dtype=torch.bfloat16) * 0.5
+    with torch.no_grad():
+        y = lin.forward(x)
+    y_ref = torch.nn.functional.linear(x, w)
+    rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8)
+    assert y.shape == y_ref.shape and rel < 0.3
+
+    # no calibration entry for a layer -> nothing injected for it
+    sd2 = {"transformer_blocks.2.attn1.to_q.weight": _w()}
+    out2, targets2, _ = quantize_state_dict_nvfp4(sd2, "video blocks", 0, 0)
+    assert inject_input_scales(out2, targets2, cal) == 0
+
+
+def test_load_calibration_reads_sidecar(tmp_path, monkeypatch):
+    import nvfp4_patch as np_mod
+    monkeypatch.setattr(np_mod, "_scales_path", lambda p: str(tmp_path / "s.json"))
+    assert np_mod.load_calibration("whatever") is None
+    (tmp_path / "s.json").write_text('{"amax": {"a.b": 3.0}, "margin": 1.25}')
+    assert np_mod.load_calibration("whatever") == {"a.b": 3.0}