Phase 2

toolkit/util/comfy_quant_export.py

@@ -0,0 +1,104 @@
+"""Export toolkit-quantized modules into ComfyUI ``comfy_quant`` checkpoints —
+the inverse of toolkit/util/comfy_quant_import.py.
+
+Every quantized OstrisLinear whose backend has a comfy storage format emits
+its quantized tensors plus the ``<prefix>.comfy_quant`` uint8 JSON marker:
+
+  - convrot8 (int8_tensorwise + convrot): weight int8 [out, in], fp32
+    weight_scale [out, 1] (comfy_kitchen's per-channel convention)
+  - nvfp4: high-nibble-first packed fp4 pairs, e4m3 block scales re-swizzled
+    to the cuBLAS 128x4 tile layout, fp32 weight_scale_2 per-tensor scale and
+    optional AWQ pre_quant_scale
+  - float8_e4m3fn: fp8_e4m3 weight + one fp32 per-tensor weight_scale
+  - convrotcomfyw4a4: via convrot_quant.export_comfy_convrot_w4a4
+
+Biases are NOT emitted here — they are ordinary parameters and flow through
+the regular state_dict path.
+"""
+
+import json
+from typing import Dict, List, Tuple
+
+import torch
+
+from toolkit.util.ostris_quant import OstrisLinear
+
+
+def comfy_quant_marker(conf: dict) -> torch.Tensor:
+    return torch.tensor(list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8)
+
+
+@torch.no_grad()
+def export_comfy_quantized_layers(
+    root: torch.nn.Module,
+) -> Tuple[Dict[str, torch.Tensor], List[str], List[str]]:
+    """Comfy-format state-dict entries for every quantized OstrisLinear in
+    ``root``. Returns ``(entries, exported_names, unexportable_names)`` —
+    entries are keyed by module path in root's layout (run
+    convert_state_dict_on_save afterwards for the checkpoint layout);
+    unexportable_names lists quantized modules whose backend has no comfy
+    storage format (the caller decides whether to dequantize instead)."""
+    from toolkit.util.convrot_quant import (
+        ConvRotComfyW4A4Quantizer,
+        export_comfy_convrot_w4a4,
+    )
+    from toolkit.util.nvfp4_quant import swap_nvfp4_nibbles, swizzle_nvfp4_scales
+
+    entries: Dict[str, torch.Tensor] = {}
+    exported: List[str] = []
+    unexportable: List[str] = []
+
+    for name, module in root.named_modules():
+        if not isinstance(module, OstrisLinear):
+            continue
+
+        if hasattr(module, "cr8_qdata"):
+            rot = int(getattr(module, "cr8_rot_size", 1) or 1)
+            conf = {"format": "int8_tensorwise"}
+            if rot > 1:
+                conf.update({"convrot": True, "convrot_groupsize": rot})
+            entries[f"{name}.weight"] = module.cr8_qdata.detach().cpu().contiguous()
+            entries[f"{name}.weight_scale"] = (
+                module.cr8_scales.view(torch.float32)
+                .detach()
+                .cpu()
+                .reshape(module.out_features, 1)
+                .contiguous()
+            )
+            entries[f"{name}.comfy_quant"] = comfy_quant_marker(conf)
+        elif hasattr(module, "nv4_qdata"):
+            entries[f"{name}.weight"] = swap_nvfp4_nibbles(
+                module.nv4_qdata.detach().cpu()
+            )
+            scales = module.nv4_scales.view(torch.float8_e4m3fn).detach().cpu()
+            entries[f"{name}.weight_scale"] = swizzle_nvfp4_scales(
+                scales.reshape(module.out_features, module.in_features // 16)
+            ).view(torch.float8_e4m3fn)
+            entries[f"{name}.weight_scale_2"] = (
+                module.nv4_pts.view(torch.float32).detach().cpu().reshape(())
+            )
+            if hasattr(module, "nv4_pre_scale"):
+                entries[f"{name}.pre_quant_scale"] = (
+                    module.nv4_pre_scale.view(torch.float32).detach().cpu().contiguous()
+                )
+            entries[f"{name}.comfy_quant"] = comfy_quant_marker({"format": "nvfp4"})
+        elif hasattr(module, "f8_qdata"):
+            entries[f"{name}.weight"] = module.f8_qdata.detach().cpu().contiguous()
+            entries[f"{name}.weight_scale"] = (
+                module.f8_scale.view(torch.float32).detach().cpu().reshape(())
+            )
+            entries[f"{name}.comfy_quant"] = comfy_quant_marker(
+                {"format": "float8_e4m3fn", "full_precision_matrix_mult": True}
+            )
+        elif isinstance(module.ostris_quantizer, ConvRotComfyW4A4Quantizer):
+            layer_entries = export_comfy_convrot_w4a4(module, f"{name}.")
+            layer_entries.pop(f"{name}.bias", None)
+            entries.update(
+                {k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in layer_entries.items()}
+            )
+        else:
+            unexportable.append(name)
+            continue
+        exported.append(name)
+
+    return entries, exported, unexportable

toolkit/util/comfy_quant_import.py

@@ -71,6 +71,135 @@ class Int8Embedding(torch.nn.Module):
         return out.to(input_ids.device).reshape(*input_ids.shape, self.embedding_dim)
 
 
+@torch.no_grad()
+def split_fused_quantized_keys(
+    state_dict: Dict[str, torch.Tensor],
+    prefix: str,
+    dst_prefixes,
+) -> Dict[str, torch.Tensor]:
+    """Split one fused quantized comfy entry (``<prefix>.weight`` /
+    ``.weight_scale`` / ``.comfy_quant`` / ...) into equal row ranges under
+    ``dst_prefixes`` (out-dim concat order). Exact for every supported format:
+    int8 rows and their per-row scales slice; fp8's per-tensor scale and
+    nvfp4's weight_scale_2 / pre_quant_scale are shared by every split; nvfp4
+    block scales are unswizzled, row-split, and re-swizzled. Mutates and
+    returns state_dict. Used by classes whose module layout splits a fused
+    checkpoint projection (e.g. qkv -> to_q/to_k/to_v)."""
+    from toolkit.util.nvfp4_quant import swizzle_nvfp4_scales
+
+    marker = state_dict.pop(f"{prefix}.comfy_quant")
+    conf = parse_comfy_quant_blob(marker)
+    fmt = conf.get("format")
+
+    weight = state_dict.pop(f"{prefix}.weight")
+    scale = state_dict.pop(f"{prefix}.weight_scale", None)
+    pts = state_dict.pop(f"{prefix}.weight_scale_2", None)
+    pre = state_dict.pop(f"{prefix}.pre_quant_scale", None)
+    bias = state_dict.pop(f"{prefix}.bias", None)
+    state_dict.pop(f"{prefix}.input_scale", None)
+
+    n = len(dst_prefixes)
+    if weight.shape[0] % n != 0:
+        raise ValueError(
+            f"{prefix}: fused out dim {weight.shape[0]} does not split into {n}"
+        )
+    rows = weight.shape[0] // n
+
+    scale_parts = None
+    if scale is not None:
+        if fmt == "nvfp4":
+            in_features = weight.shape[1] * 2  # packed fp4 pairs
+            full = unswizzle_nvfp4_scales(
+                scale.view(torch.float8_e4m3fn), weight.shape[0], in_features // 16
+            )
+            scale_parts = [
+                swizzle_nvfp4_scales(p).view(torch.float8_e4m3fn)
+                for p in full.split(rows, dim=0)
+            ]
+        elif scale.ndim == 0 or scale.numel() == 1:
+            scale_parts = [scale.clone() for _ in range(n)]
+        else:
+            scale_parts = list(scale.reshape(weight.shape[0], -1).split(rows, dim=0))
+
+    for i, dst in enumerate(dst_prefixes):
+        state_dict[f"{dst}.comfy_quant"] = marker.clone()
+        state_dict[f"{dst}.weight"] = weight[i * rows : (i + 1) * rows].contiguous()
+        if scale_parts is not None:
+            state_dict[f"{dst}.weight_scale"] = scale_parts[i].contiguous()
+        if pts is not None:
+            state_dict[f"{dst}.weight_scale_2"] = pts.clone()
+        if pre is not None:
+            state_dict[f"{dst}.pre_quant_scale"] = pre.clone()
+        if bias is not None:
+            state_dict[f"{dst}.bias"] = bias[i * rows : (i + 1) * rows].contiguous()
+    return state_dict
+
+
+@torch.no_grad()
+def fuse_split_quantized_keys(
+    state_dict: Dict[str, torch.Tensor],
+    src_prefixes,
+    prefix: str,
+) -> Dict[str, torch.Tensor]:
+    """Inverse of split_fused_quantized_keys: concatenate N split quantized
+    comfy entries back into one fused entry (out-dim concat in src order).
+    All parts must share the same format config; fp8 parts must share the same
+    per-tensor scale (true for entries produced by the splitter). Mutates and
+    returns state_dict."""
+    from toolkit.util.nvfp4_quant import swizzle_nvfp4_scales
+
+    markers = [state_dict.pop(f"{p}.comfy_quant") for p in src_prefixes]
+    confs = [parse_comfy_quant_blob(m) for m in markers]
+    if any(c != confs[0] for c in confs[1:]):
+        raise ValueError(f"{prefix}: split parts carry different quant configs")
+    fmt = confs[0].get("format")
+
+    weights = [state_dict.pop(f"{p}.weight") for p in src_prefixes]
+    scales = [state_dict.pop(f"{p}.weight_scale", None) for p in src_prefixes]
+    ptss = [state_dict.pop(f"{p}.weight_scale_2", None) for p in src_prefixes]
+    pres = [state_dict.pop(f"{p}.pre_quant_scale", None) for p in src_prefixes]
+    biases = [state_dict.pop(f"{p}.bias", None) for p in src_prefixes]
+
+    state_dict[f"{prefix}.comfy_quant"] = markers[0]
+    weight = torch.cat(weights, dim=0).contiguous()
+    state_dict[f"{prefix}.weight"] = weight
+    if scales[0] is not None:
+        if fmt == "nvfp4":
+            in_features = weight.shape[1] * 2
+            rows = [w.shape[0] for w in weights]
+            full = torch.cat(
+                [
+                    unswizzle_nvfp4_scales(
+                        s.view(torch.float8_e4m3fn), r, in_features // 16
+                    )
+                    for s, r in zip(scales, rows)
+                ],
+                dim=0,
+            )
+            state_dict[f"{prefix}.weight_scale"] = swizzle_nvfp4_scales(full).view(
+                torch.float8_e4m3fn
+            )
+        elif scales[0].ndim == 0 or scales[0].numel() == 1:
+            if any(
+                not torch.equal(s.reshape(-1), scales[0].reshape(-1)) for s in scales[1:]
+            ):
+                raise ValueError(
+                    f"{prefix}: per-tensor scales differ across split parts"
+                )
+            state_dict[f"{prefix}.weight_scale"] = scales[0]
+        else:
+            state_dict[f"{prefix}.weight_scale"] = torch.cat(
+                [s.reshape(w.shape[0], -1) for s, w in zip(scales, weights)], dim=0
+            ).contiguous()
+    if ptss[0] is not None:
+        state_dict[f"{prefix}.weight_scale_2"] = ptss[0]
+    if pres[0] is not None:
+        state_dict[f"{prefix}.pre_quant_scale"] = pres[0]
+    if biases[0] is not None:
+        state_dict[f"{prefix}.bias"] = torch.cat(biases, dim=0).contiguous()
+    return state_dict
+
+
 def _to_ostris(module: torch.nn.Linear, quantizer, orig_dtype: torch.dtype) -> OstrisLinear:
     if "weight" in module._parameters:
         del module._parameters["weight"]
@@ -128,7 +257,19 @@ def import_comfy_quantized_layers(
                 "expected nn.Linear or nn.Embedding"
             )
 
-        if fmt == "int8_tensorwise":
+        if fmt == "float8_e4m3fn":
+            # fp8_e4m3 weight + fp32 per-tensor scale, dequantized matmul
+            from toolkit.util.float8_quant import Float8Quantizer
+
+            quantizer = get_ostris_quantizer("float8_e4m3fn")
+            Float8Quantizer.attach_(
+                module,
+                weight.view(torch.float8_e4m3fn)
+                if weight.dtype != torch.float8_e4m3fn
+                else weight,
+                weight_scale,
+            )
+        elif fmt == "int8_tensorwise":
             rot = int(conf.get("convrot_groupsize", 256)) if conf.get("convrot") else 1
             quantizer = get_ostris_quantizer("convrot8")
             module.register_buffer("cr8_qdata", weight.contiguous(), persistent=False)
@@ -158,7 +299,7 @@ def import_comfy_quantized_layers(
         else:
             raise ValueError(
                 f"Unsupported comfy quant format {fmt!r} on {prefix} "
-                "(supported: int8_tensorwise, nvfp4)"
+                "(supported: int8_tensorwise, nvfp4, float8_e4m3fn)"
             )
 
         # drop unused calibration extras if present
@@ -174,4 +315,40 @@ def import_comfy_quantized_layers(
             )
         converted += 1
 
+    # legacy ComfyUI scaled-fp8 checkpoints (e.g. the wan *_fp8_scaled files):
+    # a top-level ``scaled_fp8`` marker tensor plus per-layer fp8 ``weight``
+    # and scalar fp32 ``scale_weight`` — the float8 backend's exact storage.
+    # ``scale_input`` (activation quant) is dropped; matmuls run dequantized.
+    if "scaled_fp8" in state_dict:
+        from toolkit.util.float8_quant import Float8Quantizer
+
+        state_dict.pop("scaled_fp8")
+        for scale_key in [k for k in state_dict if k.endswith(".scale_weight")]:
+            prefix = scale_key[: -len(".scale_weight")]
+            module_path = key_map(prefix) if key_map is not None else prefix
+            module = root.get_submodule(module_path)
+            if not isinstance(module, torch.nn.Linear):
+                raise ValueError(
+                    f"scaled_fp8 entry {prefix} points at {type(module).__name__}, "
+                    "expected nn.Linear"
+                )
+            weight = state_dict.pop(f"{prefix}.weight")
+            scale = state_dict.pop(scale_key)
+            state_dict.pop(f"{prefix}.scale_input", None)
+            quantizer = get_ostris_quantizer("float8_e4m3fn")
+            Float8Quantizer.attach_(
+                module,
+                weight
+                if weight.dtype == torch.float8_e4m3fn
+                else weight.view(torch.float8_e4m3fn),
+                scale,
+            )
+            _to_ostris(module, quantizer, orig_dtype)
+            bias = state_dict.pop(f"{prefix}.bias", None)
+            if bias is not None and module.bias is not None:
+                module._parameters["bias"] = torch.nn.Parameter(
+                    bias.detach().clone(), requires_grad=False
+                )
+            converted += 1
+
     return state_dict, converted

toolkit/util/float8_quant.py

@@ -0,0 +1,53 @@
+"""ComfyUI-style float8 weight storage as an Ostris backend.
+
+Matches the comfy_quant ``{"format": "float8_e4m3fn",
+"full_precision_matrix_mult": true}`` layout: the weight stored as
+torch.float8_e4m3fn plus one fp32 per-tensor scale, matmuls running on the
+dequantized weight (W8A16 numerics). Used both to import comfy fp8/fp8-mixed
+checkpoints and to quantize/export in that format.
+"""
+
+from typing import Optional
+
+import torch
+
+from toolkit.util.ostris_quant import OstrisLinear, OstrisQuantizer
+
+FLOAT8_QTYPES = ["float8_e4m3fn"]
+
+F8_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+class Float8Quantizer(OstrisQuantizer):
+    """fp8_e4m3 weight + fp32 per-tensor scale, dequantized matmul."""
+
+    def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
+        scale = (weight_fp32.abs().max() / F8_MAX).clamp(min=1e-12)
+        q = (weight_fp32 / scale).clamp(-F8_MAX, F8_MAX).to(torch.float8_e4m3fn)
+        self.attach_(module, q, scale)
+
+    @staticmethod
+    def attach_(
+        module: torch.nn.Module,
+        qweight: torch.Tensor,  # float8_e4m3fn (out, in)
+        scale: torch.Tensor,  # fp32 scalar
+    ) -> None:
+        """Register the quantized representation on the module. Used both by
+        quantize_ and by importers of pre-quantized checkpoints."""
+        module.register_buffer("f8_qdata", qweight.contiguous(), persistent=False)
+        module.register_buffer(
+            "f8_scale",
+            scale.detach().float().clone().reshape(1).view(torch.uint8),
+            persistent=False,
+        )
+
+    def dequantize(self, module: "OstrisLinear") -> torch.Tensor:
+        scale = module.f8_scale.view(torch.float32)[0]
+        return module.f8_qdata.to(torch.float32) * scale
+
+    @torch.no_grad()
+    def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None:
+        w = fp_weight.to(torch.float32)
+        scale = (w.abs().max() / F8_MAX).clamp(min=1e-12)
+        module.f8_qdata.copy_((w / scale).clamp(-F8_MAX, F8_MAX).to(torch.float8_e4m3fn))
+        module.f8_scale.copy_(scale.reshape(1).view(torch.uint8))

toolkit/util/nvfp4_quant.py

@@ -59,6 +59,25 @@ def swap_nvfp4_nibbles(packed: torch.Tensor) -> torch.Tensor:
     return ((packed << 4) | (packed >> 4)).contiguous()
 
 
+def swizzle_nvfp4_scales(scales: torch.Tensor) -> torch.Tensor:
+    """Inverse of unswizzle_nvfp4_scales: row-major (rows, cols) block scales
+    into the cuBLAS 128x4-tile layout ComfyUI checkpoints store (comfy_kitchen's
+    ``to_blocked``). Pads to tile boundaries when needed."""
+    rows, cols = scales.shape
+    n_row_blocks = (rows + 127) // 128
+    n_col_blocks = (cols + 3) // 4
+    padded_rows = n_row_blocks * 128
+    padded_cols = n_col_blocks * 4
+    if (padded_rows, padded_cols) != (rows, cols):
+        padded = scales.new_zeros(padded_rows, padded_cols)
+        padded[:rows, :cols] = scales
+        scales = padded
+    x = scales.reshape(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
+    x = x.reshape(n_row_blocks, n_col_blocks, 4, 32, 4)
+    x = x.transpose(2, 3).reshape(-1, 32, 16)
+    return x.reshape(-1).contiguous()
+
+
 class Nvfp4Quantizer(OstrisQuantizer):
     """Block-16 nvfp4 weights, full-precision activations. One instance is
     shareable across modules."""

toolkit/util/ostris_quant.py

@@ -186,6 +186,7 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
     """Resolve a qtype string to a quantizer backend instance, or None if the qtype
     does not belong to a custom backend. Add new backends here."""
     from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer
+    from toolkit.util.float8_quant import FLOAT8_QTYPES, Float8Quantizer
     from toolkit.util.nvfp4_quant import NVFP4_QTYPES, Nvfp4Quantizer
     from toolkit.util.orbit_quant import ORBIT_QTYPES, OrbitQuantizer
     from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
@@ -200,6 +201,8 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
         quantizer = get_convrot_quantizer(qtype)
     elif qtype in NVFP4_QTYPES:
         quantizer = Nvfp4Quantizer()
+    elif qtype in FLOAT8_QTYPES:
+        quantizer = Float8Quantizer()
     elif qtype in UINTX_QTYPES:
         quantizer = UIntXQuantizer(UINTX_QTYPES[qtype])
     if quantizer is not None: