danielhanchen · Oct 21, 2026 · Oct 21, 2026 · Oct 21, 2026 · Oct 21, 2026 · Oct 21, 2026
@@ -48,6 +48,11 @@ def untile_scales(tiled: Any, rows: int, blocks: int) -> Any:
             f"tiled scales hold {flat.numel()} values, {rp * bp} expected for {rows}x{blocks}"
         )
     # tile (row // 128, block // 4) is 512 values: inner (row % 32, (row % 128) // 32, block % 4)
+    from .diffusion_single_file_converters import charge_host_ram
+
+    # the permuted copy, and while it is still live the trimmed copy when the block count is padded
+    trimmed = rows * blocks if blocks != bp else 0
+    charge_host_ram((rp * bp + trimmed) * flat.element_size())
     v = flat.reshape(rp // 128, bp // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rp, bp)
     return v[:rows, :blocks].contiguous()
 
@@ -59,6 +64,9 @@ def tile_scales(plain: Any) -> Any:
     rows, blocks = plain.shape
     rp, bp = _roundup(rows, 128), _roundup(blocks, 4)
     raw = plain.view(torch.uint8) if plain.element_size() == 1 else plain
+    from .diffusion_single_file_converters import charge_host_ram
+
+    charge_host_ram(2 * rp * bp * raw.element_size())  # the padded buffer and its permuted copy
     full = torch.zeros(rp, bp, dtype = raw.dtype, device = raw.device)
     full[:rows, :blocks] = raw
     out = full.reshape(rp // 128, 4, 32, bp // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1).contiguous()
@@ -69,8 +77,22 @@ def swap_nibbles(packed: Any) -> Any:
     """High-first fp4 packing <-> low-first (``0xAB`` -> ``0xBA``)."""
     import torch
 
+    from .diffusion_single_file_converters import charge_host_ram
+
     u = packed.view(torch.uint8) if packed.dtype != torch.uint8 else packed
-    return ((u << 4) | (u >> 4)).to(torch.uint8)
+    # One output (the source stays memory-mapped), filled in bounded row chunks: the expression form held
+    # the left shift, the right shift and their OR at once, three packed copies.
+    flat = u.reshape(u.shape[0], -1) if u.dim() > 1 else u.reshape(-1, 1)
+    row = max(1, int(flat.shape[1]))
+    chunk = max(1, min(int(flat.shape[0]), _DEQUANT_CHUNK_BYTES // row))
+    charge_host_ram(u.numel() + chunk * row)
+    out = torch.empty_like(flat)
+    for start in range(0, int(flat.shape[0]), chunk):
+        part = flat[start : start + chunk]
+        dst = out[start : start + chunk]
+        torch.bitwise_left_shift(part, 4, out = dst)
+        dst |= part >> 4
+    return out.reshape(u.shape)
 
 
 def decode_layer(fmt: str, weight: Any, weight_scale: Any) -> tuple[Any, Any]:
@@ -96,31 +118,35 @@ def logical_cols(fmt: str, codes: Any) -> int:
     return int(codes.shape[1]) * (2 if fmt == NVFP4 else 1)
 
 
-def dequant_block(
-    fmt: str,
-    codes: Any,
-    scale: Any,
-    *,
-    tensor_scale: Any = None,
-    pre_quant_scale: Any = None,
-    dtype: Any = None,
-) -> Any:
-    """Dense ``[N, K]`` weight; ``pre_quant_scale`` folds into the columns: ``(x * s) @ W.T == x @ (W * s).T``."""
+# Row chunks of a fallback block dequantization keep its float32 / int32 intermediates under this many bytes.
+_DEQUANT_CHUNK_BYTES = 16 << 20
+
+
+def _dequant_temp_bytes_per_row(fmt: str, cols: int, pre_quant: bool) -> int:
+    """Peak bytes of intermediates one row of ``_dequant_block_rows`` holds at once (``cols`` logical columns).
+
+    nvfp4: int32 codes (2 per packed byte: 2 * cols), two masked int32 index tensors and two float32 gathers
+    of cols / 2 each (8 * cols), the stacked float32 values (4 * cols), the scaled float32 product (4 * cols);
+    the masks and gathers are freed before the product, so the peak is codes + gathers + stack. mxfp8: float32
+    codes and their product (8 * cols). A pre-quant scale adds one float32 product."""
+    peak = 14 * cols if fmt == NVFP4 else 8 * cols
+    return peak + (4 * cols if pre_quant else 0)
+
+
+def _dequant_block_rows(fmt, codes, scale, tensor_scale, pre_quant_scale):
     import torch
 
-    dtype = dtype or torch.bfloat16
     rows = int(codes.shape[0])
     if fmt == NVFP4:
         lut = torch.tensor(
             _E2M1 + tuple(-v for v in _E2M1), dtype = torch.float32, device = codes.device
         )
         c = codes.to(torch.int32)
         values = torch.stack((lut[c & 0x0F], lut[c >> 4]), dim = -1).reshape(rows, -1, 16)
+        del c
         step = scale.to(torch.float32)
         if tensor_scale is not None:
-            step = step * torch.as_tensor(
-                tensor_scale, dtype = torch.float32, device = codes.device
-            ).reshape(-1, 1)
+            step = step * tensor_scale
         weight = (values * step.unsqueeze(-1)).reshape(rows, -1)
     elif fmt == MXFP8:
         exp = scale.view(torch.uint8).to(torch.float32) - _E8M0_BIAS
@@ -130,8 +156,50 @@ def dequant_block(
     else:
         raise ValueError(f"not a block-scaled ComfyUI format: {fmt!r}")
     if pre_quant_scale is not None:
-        weight = weight * pre_quant_scale.to(torch.float32).reshape(1, -1)
-    return weight.to(dtype)
+        weight = weight * pre_quant_scale
+    return weight
+
+
+def dequant_block(
+    fmt: str,
+    codes: Any,
+    scale: Any,
+    *,
+    tensor_scale: Any = None,
+    pre_quant_scale: Any = None,
+    dtype: Any = None,
+) -> Any:
+    """Dense ``[N, K]`` weight; ``pre_quant_scale`` folds into the columns: ``(x * s) @ W.T == x @ (W * s).T``.
+
+    Dequantized in row chunks straight into the output, so the float32 intermediates stay bounded by
+    ``_DEQUANT_CHUNK_BYTES`` instead of several times the layer; every op is per row, so the result is
+    bit-identical to one pass."""
+    import torch
+
+    from .diffusion_single_file_converters import charge_host_ram
+
+    if fmt not in BLOCK_SIZES:
+        raise ValueError(f"not a block-scaled ComfyUI format: {fmt!r}")
+    dtype = dtype or torch.bfloat16
+    rows = int(codes.shape[0])
+    cols = logical_cols(fmt, codes)
+    per_row = _dequant_temp_bytes_per_row(fmt, cols, pre_quant_scale is not None)
+    chunk = max(1, min(rows, _DEQUANT_CHUNK_BYTES // max(per_row, 1)))
+    charge_host_ram(rows * cols * torch.empty((), dtype = dtype).element_size() + chunk * per_row)
+    ts = None
+    if tensor_scale is not None:
+        ts = torch.as_tensor(tensor_scale, dtype = torch.float32, device = codes.device).reshape(-1, 1)
+    pqs = None
+    if pre_quant_scale is not None:
+        pqs = pre_quant_scale.to(torch.float32).reshape(1, -1)
+    out = torch.empty((rows, cols), dtype = dtype, device = codes.device)
+    for start in range(0, rows, chunk):
+        stop = min(rows, start + chunk)
+        part_ts = ts if ts is None or ts.shape[0] == 1 else ts[start:stop]
+        out[start:stop] = _dequant_block_rows(
+            fmt, codes[start:stop], scale[start:stop], part_ts, pqs
+        ).to(dtype)
+    return out
 
 
 def mx_quantize_activation(x: Any) -> tuple[Any, Any]:
@@ -403,6 +471,8 @@ def build_runtime_linear(
 ) -> Any:
     import torch
 
+    from .diffusion_single_file_converters import charge_host_ram
+
     rows = int(codes.shape[0])
     if fmt == NVFP4:
         from .diffusion_nvfp4_linear import nvfp4_linear_class
@@ -426,6 +496,8 @@ def build_runtime_linear(
             # FlashInfer's activation global scale is the reciprocal of ComfyUI's decode-side input_scale.
             a_gsf = (1.0 / in_scale).clone()
             cls = nvfp4_linear_class()
+        if not codes.is_contiguous():
+            charge_host_ram(int(codes.numel()) * int(codes.element_size()))
         return cls(
             cols,
             rows,
@@ -438,6 +510,8 @@ def build_runtime_linear(
             activation_scales_baked = True,
         )
     if fmt == MXFP8:
+        if not codes.is_contiguous():
+            charge_host_ram(int(codes.numel()) * int(codes.element_size()))
         return mxfp8_linear_class()(
             int(codes.shape[1]),
             rows,

@@ -29,6 +29,7 @@
 from typing import Any, Optional
 
 from .diffusion_comfy_block import BLOCK_SIZES, MXFP8, NVFP4, tiled_scale_numel
+from .diffusion_single_file_converters import host_ram_budget_scope
 
 COMFY_QUANT_SUFFIX = ".comfy_quant"
 QUANT_METADATA_KEY = "_quantization_metadata"
@@ -442,11 +443,29 @@ def refuse_comfy_quant(path: Optional[str]) -> Optional[ComfyQuantScan]:
     return scan
 
 
+def _nbytes(tensors: Any) -> int:
+    return sum(int(t.numel()) * int(t.element_size()) for t in tensors)
+
+
+def _charge_contiguous(tensor: Any) -> None:
+    """Check the copy ``tensor.contiguous()`` makes of a non-contiguous host tensor (a view costs nothing)."""
+    if not tensor.is_contiguous():
+        from .diffusion_single_file_converters import charge_host_ram
+        charge_host_ram(_nbytes([tensor]))
+
+
 def _dequant(codes: Any, scale: Any, group: int, dtype: Any) -> Any:
     """``codes * scale`` in float32, the ConvRot rotation undone, cast to ``dtype``."""
     import torch
 
+    from .diffusion_single_file_converters import charge_host_ram
+
     rows = codes.shape[0]
+    # Peak, not just the output: the float32 codes and their float32 product are live together (8 bytes per
+    # element); ConvRot's batched matmul holds the product, a float32 copy of it and its float32 result (12,
+    # measured); last the product and the output.
+    out_size = torch.empty((), dtype = dtype).element_size()
+    charge_host_ram(codes.numel() * max(12 if group else 8, 4 + out_size))
     weight = codes.to(torch.float32) * scale.to(torch.float32).reshape(-1, 1).expand(rows, 1)
     if group:
         from .diffusion_convrot import build_convrot_hadamard
@@ -482,7 +501,12 @@ def _dequant_parts(fmt: str, parts: list, dtype: Any) -> Any:
         )
         for codes, scale, extra in parts
     ]
-    return dense[0] if len(dense) == 1 else torch.cat(dense)
+    if len(dense) == 1:
+        return dense[0]
+    from .diffusion_single_file_converters import charge_host_ram
+
+    charge_host_ram(_nbytes(dense))
+    return torch.cat(dense)
 
 
 def _block_runtime_args(value: "_BlockWeight") -> Optional[tuple]:
@@ -491,6 +515,11 @@ def _block_runtime_args(value: "_BlockWeight") -> Optional[tuple]:
     import torch
 
     extras = [extra for _c, _s, extra in value.parts]
+    if len(value.parts) > 1:
+        from .diffusion_single_file_converters import charge_host_ram
+        charge_host_ram(
+            _nbytes([c for c, _s, _e in value.parts]) + _nbytes([s for _c, s, _e in value.parts])
+        )
     codes = (
         torch.cat([c for c, _s, _e in value.parts]) if len(value.parts) > 1 else value.parts[0][0]
     )
@@ -534,6 +563,7 @@ def _int8_tensor(name: str, codes: Any, scale: Any, dtype: Any) -> Optional[Any]
         "_tensor_data_names": ["zero_point", "qdata", "scale"],
     }
     stem = name.rsplit(".", 1)[0]
+    _charge_contiguous(codes)
     tensors = {
         f"{stem}._weight_qdata": codes.contiguous(),
         f"{stem}._weight_scale": scale.to(torch.float32).reshape(rows, 1).contiguous(),
@@ -591,6 +621,7 @@ def _fp8_tensor(
         "_tensor_data_names": ["qdata", "scale"],
     }
     stem = name.rsplit(".", 1)[0]
+    _charge_contiguous(codes)
     tensors = {
         f"{stem}._weight_qdata": codes.contiguous(),
         f"{stem}._weight_scale": scale.to(torch.float32)
@@ -726,7 +757,13 @@ def take(value: Any, rows: Any) -> Any:
             # float8: concatenation goes through uint8 (not implemented for float8 on every torch)
             return take(value.view(torch.uint8), rows).view(value.dtype)
         parts = [value[int(r) : int(r) + int(n)] for r, n in rows]
-        return parts[0] if len(parts) == 1 else torch.cat(parts)
+        if len(parts) == 1:
+            return parts[0]
+        from .diffusion_single_file_converters import charge_host_ram
+
+        # a reordered selection (H3's [gate; value] fc1) is a full copy of the rows it joins
+        charge_host_ram(_nbytes(parts))
+        return torch.cat(parts)
 
     converted: dict = {}
     for key in list(state):
@@ -920,6 +957,7 @@ def load_comfy_prequant(
     return model
 
 
+@host_ram_budget_scope
 def load_comfy_quant_transformer(
     transformer_cls: Any,
     path: str,
@@ -1036,6 +1074,42 @@ def _dtype_for(key: str) -> Any:
         NVFP4: nvfp4_backend,
         MXFP8: mxfp8_backend,
     }
+    # Without a runtime (fp8 below sm_89) layers widen to ``dtype`` in host RAM: refuse before an OOM.
+    from .diffusion_single_file_converters import (
+        HostRamBudget,
+        HostRamShortfallError,
+        charge_cast,
+        charge_host_ram,
+        converted_host_bytes,
+    )
+
+    # Mirrors the source loop below: a layer is kept as codes only with a runtime for its format (fp16_keep
+    # merely defers the dtype choice; the layer is still expanded) and, for a block format, no key map.
+    # nvfp4 headers describe packed uint8 (K / 2 columns); the dense weight is N x K.
+    _widened = {
+        layer.name + ".weight": (
+            (2 if layer.format == NVFP4 else 1)
+            if not backends.get(layer.format)
+            or (layer.format in BLOCK_SIZES and key_map is not None)
+            else 0
+        )
+        for layer in scan.layers.values()
+    }
+    # The header estimate refuses early; the budget it activates is charged by every allocating helper
+    # below, so a layer the estimate cannot see (runtime filters, nibble swaps) still cannot OOM the host.
+    budget = HostRamBudget(
+        path, f"expanding its quantized weights to {str(dtype).replace('torch.', '')}", logger
+    )
+    budget.check(
+        converted_host_bytes(
+            path,
+            dtype,
+            keep_fp32 = keep_fp32 if fp16_keep else (),
+            dequantized = _widened.get,
+            keep_key = keep_key,
+        )
+    )
+    budget.activate()
     sources: list = []  # (layer, codes, scale, block extras or None) kept for a runtime
     dequantized = 0
     for layer in scan.layers.values():
@@ -1078,10 +1152,12 @@ def _dtype_for(key: str) -> Any:
             continue
         wanted_dtype = keep_dtype(key) if keep_dtype is not None else None
         if wanted_dtype is not None:
+            charge_cast(value, wanted_dtype)
             state[key] = value.to(wanted_dtype)
             continue
         if fp16_keep or value.dtype == dtype:
             continue
+        charge_cast(value, dtype)
         state[key] = value.to(dtype)
     from accelerate import init_empty_weights
 
@@ -1146,14 +1222,15 @@ def _convert_tagged(width: Optional[int]) -> dict:
                 out[name] = _BlockWeight(fmt, blocks)
                 continue
             parts = [sources[i][1][r : r + n] for i, r, n in segments]
-            # one segment (a split or renamed layer) is a row slice of a contiguous tensor: no copy
+            scale_parts = [
+                sources[i][2].reshape(-1, 1).expand(sources[i][1].shape[0], 1)[r : r + n]
+                for i, r, n in segments
+            ]
+            # one segment (a split or renamed layer) is a row slice of a contiguous tensor: no copy; a
+            # reorder (a stock converter swapping gate / proj halves) is a full copy of the kept codes
+            charge_host_ram((_nbytes(parts) if len(parts) > 1 else 0) + _nbytes(scale_parts))
             codes = parts[0] if len(parts) == 1 else torch.cat(parts)
-            scale = torch.cat(
-                [
-                    sources[i][2].reshape(-1, 1).expand(sources[i][1].shape[0], 1)[r : r + n]
-                    for i, r, n in segments
-                ]
-            )
+            scale = torch.cat(scale_parts)
             if name in wanted and tuple(wanted[name].shape) != tuple(codes.shape):
                 raise ValueError(
                     f"{name}: rebuilt as {tuple(codes.shape)}, the model expects {tuple(wanted[name].shape)}"
@@ -1174,12 +1251,18 @@ def _convert_tagged(width: Optional[int]) -> dict:
         try:
             # Narrow tags first (full width is slow on a large DiT); anything unproven reruns at full width.
             converted = _convert_tagged(_NARROW_TAG_COLUMNS if sources else None)
+        except HostRamShortfallError:
+            raise  # a refusal, not an unproven tag: the full-width retry would only allocate more
         except Exception:  # noqa: BLE001 -- the full-width pass raises the real refusal
+            # Full-width float64 tags are views, but a converter that joins or splits rows materializes them
+            # (8 bytes per weight element) before any checked helper runs: check that before retrying.
+            charge_host_ram(8 * sum(int(codes.numel()) for _layer, codes, *_rest in sources))
             converted = _convert_tagged(None)
     del state
     if fp16_keep:
         for key, value in list(converted.items()):
             if torch.is_tensor(value) and value.is_floating_point():
+                charge_cast(value, _dtype_for(key))
                 converted[key] = value.to(_dtype_for(key))
 
     min_features = DEFAULT_MIN_LINEAR_FEATURES if min_features is None else int(min_features)