feat(model): load nvfp4 comfy_quant checkpoints via sdnq

nvfp4 layers keep their packed 4-bit codes and land on SDNQ grouped
quantization as float4_e2m1fn: the nibble order is swapped once at load,
the e4m3 block scales are unswizzled from the cuBLAS tile layout, and
the fp32 global scale folds into them as per-group scales. Marker
orig_shape acts as a cross-check and alignment padding is sliced
against the model dimensions.

- reject nvfp4 markers with unknown group sizes or convrot flags
- accept uint8 storage in the krea2 real-file check

pipelines/native_transformer.py

@@ -74,7 +74,9 @@ COMFY_QUANT_MARKER = ".comfy_quant"
 COMFY_QUANT_FORMATS: dict[str, str] = {  # comfy_quant format string -> SDNQ weights_dtype
     "int8_tensorwise": "int8",
     "float8_e4m3fn": "float8_e4m3fn",
+    "nvfp4": "float4_e2m1fn",
 }
+NVFP4_GROUP_SIZE = 16  # fixed by the format definition; markers restate it
 
 
 class OverrideArchMismatch(Exception):
@@ -521,6 +523,11 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict
                 f"for {name!r} is malformed ({type(e).__name__}: {e})"
             ) from e
         if meta.get("convrot"):
+            if fmt == "nvfp4":
+                raise OverrideArchMismatch(
+                    f"Load model: type={type_name} native_transformer comfy_quant layer "
+                    f"{name!r} combines nvfp4 with convrot, which the format does not define"
+                )
             group_size = int(meta.get("convrot_groupsize", 256))
             is_pow4 = group_size >= 4 and (group_size & (group_size - 1)) == 0 and (group_size.bit_length() & 1) == 1
             if not is_pow4:
@@ -528,6 +535,11 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict
                     f"Load model: type={type_name} native_transformer comfy_quant layer "
                     f"{name!r} convrot_groupsize={group_size} is not a power of 4"
                 )
+        if fmt == "nvfp4" and int(meta.get("group_size", NVFP4_GROUP_SIZE)) != NVFP4_GROUP_SIZE:
+            raise OverrideArchMismatch(
+                f"Load model: type={type_name} native_transformer comfy_quant layer "
+                f"{name!r} nvfp4 group_size={meta.get('group_size')} is not {NVFP4_GROUP_SIZE}"
+            )
         marked[name] = meta
         formats.add(fmt)
     unsupported = sorted(formats - set(COMFY_QUANT_FORMATS))
@@ -548,15 +560,18 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict
     return marked, next(iter(formats))
 
 
-def remap_comfy_quant(state_dict: dict, marked_names: set[str]) -> dict:
+def remap_comfy_quant(state_dict: dict, marked_names: set[str], defer_scales: bool = False) -> dict:
     """Translate comfy_quant tensor naming to SDNQ naming.
 
     Renames ``<name>.weight_scale`` to ``<name>.scale`` (reshaped to 2-D,
     ``[]`` becomes ``[1, 1]``, since SDNQ transposes scales in place) and
     drops the ``<name>.comfy_quant`` markers plus optional
     ``<name>.input_scale`` activation-calibration sidecars (SDNQ re-derives
-    activation scales dynamically). Returns a new dict; the input (which may
-    be the cached state dict) is not mutated.
+    activation scales dynamically). With ``defer_scales`` the scale rename is
+    skipped: block-scaled formats (nvfp4) carry swizzled scale tensors whose
+    transform needs the layer dimensions, so the prequantized builder handles
+    them per layer. Returns a new dict; the input (which may be the cached
+    state dict) is not mutated.
     """
     scale_suffix = ".weight_scale"
     input_scale_suffix = ".input_scale"
@@ -566,13 +581,78 @@ def remap_comfy_quant(state_dict: dict, marked_names: set[str]) -> dict:
             continue
         if key.endswith(input_scale_suffix) and key[: -len(input_scale_suffix)] in marked_names:
             continue
-        if key.endswith(scale_suffix) and key[: -len(scale_suffix)] in marked_names:
+        if not defer_scales and key.endswith(scale_suffix) and key[: -len(scale_suffix)] in marked_names:
             remapped[f"{key[: -len(scale_suffix)]}.scale"] = value.reshape(-1, 1)
             continue
         remapped[key] = value
     return remapped
 
 
+def unswizzle_block_scales(scales: torch.Tensor, rows: int, groups: int) -> torch.Tensor:
+    """Invert the cuBLAS 2D block-scaling factor layout back to row-major
+    ``[rows, groups]``, dropping the tile alignment padding.
+
+    Block-scaled comfy_quant checkpoints store per-group scales pre-tiled for
+    the hardware kernels: the ``[roundup(rows, 128), roundup(groups, 4)]``
+    grid is rearranged into 32x16 tiles as described in
+    https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+    """
+    row_blocks = -(rows // -128)
+    col_blocks = -(groups // -4)
+    tiles = scales.reshape(-1, 32, 4, 4).transpose(1, 2)
+    tiles = tiles.reshape(row_blocks, col_blocks, 4, 32, 4).reshape(row_blocks, col_blocks, 128, 4)
+    return tiles.permute(0, 2, 1, 3).reshape(row_blocks * 128, col_blocks * 4)[:rows, :groups]
+
+
+def adopt_nvfp4_layer(sd: dict, name: str, linear: torch.nn.Linear, component_name: str, meta: dict) -> None:
+    """Rewrite one nvfp4 layer's tensors in place into SDNQ layout.
+
+    Slices tile-alignment padding off the packed weight, swaps the nibble
+    order (the container packs the even element into the high nibble, SDNQ
+    unpacks low-first), unswizzles the e4m3 block scales, and folds the fp32
+    global scale into them as ``[out, groups, 1]`` fp32 grouped scales.
+    """
+    out_features, in_features = linear.out_features, linear.in_features
+    if in_features % NVFP4_GROUP_SIZE != 0:
+        raise OverrideArchMismatch(
+            f"Load model: transformer=native {component_name} comfy_quant marked module "
+            f"{name!r} nvfp4 group size {NVFP4_GROUP_SIZE} does not divide in_features={in_features}"
+        )
+    orig_shape = meta.get("orig_shape")
+    if orig_shape is not None and tuple(orig_shape) != (out_features, in_features):
+        raise OverrideArchMismatch(
+            f"Load model: transformer=native {component_name} comfy_quant marked module "
+            f"{name!r} stores orig_shape={tuple(orig_shape)} but the model expects {(out_features, in_features)}"
+        )
+    weight = sd.get(f"{name}.weight")
+    scales = sd.pop(f"{name}.weight_scale", None)
+    global_scale = sd.pop(f"{name}.weight_scale_2", None)
+    if scales is None or global_scale is None:
+        raise OverrideArchMismatch(
+            f"Load model: transformer=native {component_name} comfy_quant marked module "
+            f"{name!r} is missing nvfp4 weight_scale/weight_scale_2 tensors"
+        )
+    packed_columns = in_features // 2
+    if weight is None or weight.ndim != 2 or weight.shape[0] < out_features or weight.shape[1] < packed_columns:
+        found = tuple(weight.shape) if weight is not None else "missing"
+        raise OverrideArchMismatch(
+            f"Load model: transformer=native {component_name} comfy_quant marked module "
+            f"{name!r} packed weight {found} cannot hold {(out_features, packed_columns)}"
+        )
+    groups = in_features // NVFP4_GROUP_SIZE
+    expected_scales = -(out_features // -128) * 128 * -(groups // -4) * 4
+    if scales.numel() != expected_scales:
+        raise OverrideArchMismatch(
+            f"Load model: transformer=native {component_name} comfy_quant marked module "
+            f"{name!r} block scales hold {scales.numel()} values, expected {expected_scales}"
+        )
+    weight = weight[:out_features, :packed_columns]
+    weight = torch.bitwise_or(torch.bitwise_left_shift(torch.bitwise_and(weight, 15), 4), torch.bitwise_right_shift(weight, 4))
+    scales = unswizzle_block_scales(scales, out_features, groups)
+    sd[f"{name}.weight"] = weight
+    sd[f"{name}.scale"] = scales.to(torch.float32).mul_(global_scale.to(torch.float32)).unsqueeze(-1)
+
+
 def partition_siblings(
     state_dict: dict,
     siblings: dict[str, SiblingSpec],
@@ -727,10 +807,13 @@ def build_component_prequantized(
     the marked layers onto SDNQ layers without dequantizing.
 
     The supported formats are subsets of SDNQ's symmetric quantization: same
-    storage layout (unpacked 8-bit ``[out, in]``), same dequant math
-    (``weight * scale``, no zero point), so tensors are adopted bit-exact;
-    convrot layers map onto SDNQ's Hadamard support, whose identical regular
-    Hadamard construction lets the dequantizer undo the stored rotation.
+    dequant math (``weight * scale``, no zero point), so tensors are adopted
+    bit-exact. The 8-bit formats share the storage layout directly (unpacked
+    ``[out, in]``); convrot layers map onto SDNQ's Hadamard support, whose
+    identical regular Hadamard construction lets the dequantizer undo the
+    stored rotation; nvfp4 layers keep their packed 4-bit codes (nibble order
+    swapped once) and land on SDNQ's grouped quantization with the block and
+    global scales folded into fp32 per-group scales.
     The file dictates which layers are quantized, independent of the user's
     quantization settings, and floating-point SDNQ params are not cast to the
     target dtype (the fp32 scales must survive). Layers are assembled in
@@ -753,8 +836,9 @@ def build_component_prequantized(
     matmul_dtype = "int8" if dtype_dict[weights_dtype]["is_integer"] else "float8_e4m3fn"
     storage_dtype = dtype_dict[weights_dtype]["storage_dtype"]
     target_dtype = dtype if dtype is not None else devices.dtype
+    is_nvfp4 = comfy_format == "nvfp4"
     marked_names = set(marker_meta)
-    sd = remap_comfy_quant(state_dict, marked_names)
+    sd = remap_comfy_quant(state_dict, marked_names, defer_scales=is_nvfp4)
 
     # Civitai relabels these containers freely; trust the marker only as far
     # as the stored tensors actually match it.
@@ -773,7 +857,7 @@ def build_component_prequantized(
     quantization_config = SDNQConfig(
         weights_dtype=weights_dtype,
         quantized_matmul_dtype=matmul_dtype,
-        group_size=-1,
+        group_size=NVFP4_GROUP_SIZE if is_nvfp4 else -1,
         use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
         dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
         add_skip_keys=False,
@@ -809,20 +893,23 @@ def build_component_prequantized(
                 f"Load model: transformer=native {component_name} comfy_quant marked "
                 f"module {name!r} convrot_groupsize={hadamard_group_size} does not divide in_features={linear.in_features}"
             )
+        if is_nvfp4:
+            adopt_nvfp4_layer(sd, name, linear, component_name, marker_meta[name])
+        layer_shape = torch.Size((linear.out_features, linear.in_features))
         linear.sdnq_dequantizer = SDNQDequantizer(
             result_dtype=target_dtype,
-            result_shape=None,
-            original_shape=torch.Size((linear.out_features, linear.in_features)),
+            result_shape=layer_shape if is_nvfp4 else None,
+            original_shape=layer_shape,
             original_stride=(linear.in_features, 1),
-            quantized_weight_shape=torch.Size((linear.out_features, linear.in_features)),
+            quantized_weight_shape=torch.Size((linear.out_features, linear.in_features // NVFP4_GROUP_SIZE, NVFP4_GROUP_SIZE)) if is_nvfp4 else layer_shape,
             weights_dtype=weights_dtype,
             quantized_matmul_dtype=matmul_dtype,
             hadamard_group_size=hadamard_group_size,
-            group_size=-1,
+            group_size=NVFP4_GROUP_SIZE if is_nvfp4 else -1,
             svd_rank=32,
             svd_steps=8,
             use_quantized_matmul=False,
-            re_quantize_for_matmul=False,
+            re_quantize_for_matmul=is_nvfp4,
             use_stochastic_rounding=False,
             use_hadamard=use_hadamard,
             layer_class_name="Linear",

test/test-krea2-transformer.py

@@ -190,7 +190,7 @@ def run_comfy_quant_real_file():
     storage_dtypes = {m.weight.dtype for m in sdnq_layers}
     print(f"comfy_quant real file: {len(sdnq_layers)} SDNQ layers, storage {storage_dtypes}")
     assert len(sdnq_layers) == expected_layers, f"expected {expected_layers} SDNQ layers, got {len(sdnq_layers)}"
-    assert storage_dtypes <= {torch.int8, torch.float8_e4m3fn}, f"unexpected storage dtypes: {storage_dtypes}"
+    assert storage_dtypes <= {torch.int8, torch.float8_e4m3fn, torch.uint8}, f"unexpected storage dtypes: {storage_dtypes}"
     assert len(storage_dtypes) == 1, "adopted weights must share one storage dtype"
     assert transformer.blocks[0].attn.wq.__class__.__name__ == "SDNQLinear"
     assert getattr(transformer, "quantization_config", None) is not None

test/test-native-transformer.py

@@ -550,6 +550,93 @@ def test_transcode_quant_metadata_no_matches_is_noop():
     assert out is sd
 
 
+# ============================================================
+# nvfp4: sdnq codec, scale unswizzle, detection
+# ============================================================
+
+E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0)
+
+
+def to_blocked_reference(matrix: torch.Tensor) -> torch.Tensor:
+    """cuBLAS block-scaling tile layout as written into nvfp4 containers;
+    the loader's unswizzle must invert this exactly."""
+    rows, cols = matrix.shape
+    row_blocks = -(rows // -128)
+    col_blocks = -(cols // -4)
+    padded = torch.zeros((row_blocks * 128, col_blocks * 4), dtype=matrix.dtype)
+    padded[:rows, :cols] = matrix
+    blocks = padded.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3)
+    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
+    return rearranged.reshape(row_blocks * 128, col_blocks * 4)
+
+
+def test_unswizzle_block_scales_roundtrip():
+    for rows, groups in ((128, 4), (8, 3), (200, 10)):
+        mat = torch.arange(rows * groups, dtype=torch.float32).reshape(rows, groups)
+        back = nt.unswizzle_block_scales(to_blocked_reference(mat), rows, groups)
+        assert torch.equal(back, mat), f'unswizzle mismatch for {(rows, groups)}'
+
+
+def test_nvfp4_codec_ocp_table():
+    """SDNQ's float4_e2m1fn decodes OCP FP4 E2M1 exactly, including the
+    subnormal codes 1/9 as +/-0.5; nvfp4 containers adopt it directly."""
+    from modules.sdnq.packed_float import unpack_float
+    packed = torch.tensor([(2 * j) | (((2 * j) + 1) << 4) for j in range(8)], dtype=torch.uint8)
+    dec = unpack_float(packed, 'float4_e2m1fn', torch.Size([16]))
+    for code in range(16):
+        assert float(dec[code]) == E2M1_VALUES[code], f'code {code}: {float(dec[code])} != {E2M1_VALUES[code]}'
+
+
+def test_nvfp4_pack_ocp_grid_roundtrip():
+    """Every OCP grid value survives a pack/unpack round trip exactly
+    (subnormals included), and off-grid values land inside the value set.
+    Exact nearest-rounding near the grid midpoints is not asserted: the
+    packer's staged rounding may resolve boundary values to either side."""
+    from modules.sdnq.packed_float import pack_float, unpack_float
+    grid = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, 0.0]
+    vals = torch.tensor(grid, dtype=torch.float32)
+    out = unpack_float(pack_float(vals, 'float4_e2m1fn'), 'float4_e2m1fn', vals.shape)
+    assert torch.equal(out, vals), f'grid values must roundtrip exactly: {out.tolist()}'
+    off_grid = torch.tensor([0.2, 0.6, 0.9, 1.2, 2.4, 5.5, -0.4, -0.9, -1.7, -3.4, -5.9, 0.05, 0.99, -0.99], dtype=torch.float32)
+    out = unpack_float(pack_float(off_grid, 'float4_e2m1fn'), 'float4_e2m1fn', off_grid.shape)
+    for v, o in zip(off_grid.tolist(), out.tolist()):
+        assert abs(o) in {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0}, f'{v} -> {o} outside the OCP set'
+
+
+def test_detect_comfy_nvfp4_marker_accepted():
+    import json
+    payload = json.dumps({'format': 'nvfp4', 'group_size': 16, 'orig_dtype': 'torch.float16', 'orig_shape': [32, 32]})
+    sd = {
+        'in_proj.weight': torch.zeros((32, 16), dtype=torch.uint8),
+        'in_proj.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8),
+    }
+    marked, fmt = nt.detect_comfy_quant(sd, 'Test')
+    assert fmt == 'nvfp4'
+    assert marked['in_proj']['orig_shape'] == [32, 32]
+
+
+def test_detect_comfy_nvfp4_bad_group_size_raises():
+    import json
+    payload = json.dumps({'format': 'nvfp4', 'group_size': 32})
+    sd = {'a.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8)}
+    try:
+        nt.detect_comfy_quant(sd, 'Test')
+        raise AssertionError('expected OverrideArchMismatch')
+    except nt.OverrideArchMismatch as e:
+        assert 'group_size' in str(e)
+
+
+def test_detect_comfy_nvfp4_convrot_raises():
+    import json
+    payload = json.dumps({'format': 'nvfp4', 'convrot': True, 'convrot_groupsize': 16})
+    sd = {'a.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8)}
+    try:
+        nt.detect_comfy_quant(sd, 'Test')
+        raise AssertionError('expected OverrideArchMismatch')
+    except nt.OverrideArchMismatch as e:
+        assert 'convrot' in str(e)
+
+
 # ============================================================
 # is_noop_converter
 # ============================================================
@@ -1436,6 +1523,135 @@ def test_load_comfy_full_precision_mm_excluded_from_matmul():
             os.unlink(path)
 
 
+def nvfp4_fixture(dim: int, with_marker: bool = True) -> tuple[dict, torch.Tensor]:
+    """A valid nvfp4 container for one marked linear, plus the reference
+    dequantized weight per the format definition: E2M1 decode * unswizzled
+    e4m3 block scale * fp32 global scale."""
+    import json
+    groups = dim // 16
+    gen = torch.Generator().manual_seed(7)
+    codes = torch.randint(0, 16, (dim, dim), generator=gen, dtype=torch.uint8)
+    block_scales = (torch.rand(dim, groups, generator=gen) * 2 + 0.5).to(torch.float8_e4m3fn)
+    global_scale = torch.tensor(0.0125, dtype=torch.float32)
+    lut = torch.tensor(E2M1_VALUES, dtype=torch.float32)
+    reference = lut[codes.long()].reshape(dim, groups, 16) * (block_scales.float() * global_scale).unsqueeze(-1)
+    reference = reference.reshape(dim, dim)
+    packed = torch.bitwise_or(torch.bitwise_left_shift(codes[:, 0::2], 4), codes[:, 1::2])  # even element in the high nibble
+    sd = {
+        'model.diffusion_model.in_proj.weight': packed,
+        'model.diffusion_model.in_proj.weight_scale': to_blocked_reference(block_scales),
+        'model.diffusion_model.in_proj.weight_scale_2': global_scale,
+        'model.diffusion_model.in_proj.bias': torch.randn(dim, generator=gen).to(torch.float16),
+        'model.diffusion_model.out_proj.weight': torch.randn(dim, dim, generator=gen).to(torch.float16),
+        'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16),
+    }
+    if with_marker:
+        marker = json.dumps({'format': 'nvfp4', 'group_size': 16, 'orig_dtype': 'torch.float16', 'orig_shape': [dim, dim]})
+        sd['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(marker.encode()), dtype=torch.uint8)
+    return sd, reference
+
+
+def test_load_comfy_nvfp4_end_to_end():
+    """Full pipeline on an nvfp4 container: packed uint8 weights survive with
+    swapped nibble order, block scales unswizzle and fold with the global
+    scale into fp32 grouped scales, and the dequantized weight matches the
+    format reference exactly."""
+    fd, path = tempfile.mkstemp(suffix='.safetensors')
+    try:
+        dim = 32
+        raw, reference = nvfp4_fixture(dim)
+        write_fixture(raw, fd, path)
+
+        with ComfyTestEnv(dim):
+            spec = nt.TransformerSpec(cls=MockMiniTransformer)
+            transformer, _ = nt.load(
+                local_file=path,
+                repo_id='fake/repo',
+                spec=spec,
+                diffusers_cfg={},
+                dtype=torch.float32,
+            )
+
+        in_proj = transformer.in_proj
+        assert in_proj.__class__.__name__ == 'SDNQLinear'
+        dq = in_proj.sdnq_dequantizer
+        assert dq.weights_dtype == 'float4_e2m1fn'
+        assert dq.group_size == 16
+        assert tuple(dq.quantized_weight_shape) == (dim, dim // 16, 16)
+        assert dq.re_quantize_for_matmul is True
+        assert in_proj.weight.dtype == torch.uint8
+        assert tuple(in_proj.weight.shape) == (dim, dim // 2)
+        assert in_proj.scale.dtype == torch.float32
+        assert tuple(in_proj.scale.shape) == (dim, dim // 16, 1)
+
+        dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
+        assert torch.equal(dequantized.detach().cpu(), reference)
+
+        x = torch.randn(2, dim)
+        out = in_proj(x)
+        ref = torch.nn.functional.linear(x, reference, in_proj.bias.detach().cpu())
+        assert torch.allclose(out.detach().cpu(), ref, atol=1e-5)
+
+        assert transformer.out_proj.__class__ is torch.nn.Linear
+    finally:
+        if os.path.exists(path):
+            os.unlink(path)
+
+
+def test_load_comfy_nvfp4_metadata_form():
+    """Marker-less nvfp4 (header metadata only, no orig_shape/group_size):
+    the format constant and module dimensions fill the gaps."""
+    fd, path = tempfile.mkstemp(suffix='.safetensors')
+    try:
+        dim = 32
+        raw, reference = nvfp4_fixture(dim, with_marker=False)
+        layers = {'model.diffusion_model.in_proj': {'format': 'nvfp4'}}
+        write_fixture(raw, fd, path, metadata=quant_metadata(layers))
+
+        with ComfyTestEnv(dim):
+            spec = nt.TransformerSpec(cls=MockMiniTransformer)
+            transformer, _ = nt.load(
+                local_file=path,
+                repo_id='fake/repo',
+                spec=spec,
+                diffusers_cfg={},
+                dtype=torch.float32,
+            )
+
+        in_proj = transformer.in_proj
+        assert in_proj.__class__.__name__ == 'SDNQLinear'
+        assert in_proj.sdnq_dequantizer.weights_dtype == 'float4_e2m1fn'
+        dequantized = in_proj.sdnq_dequantizer(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
+        assert torch.equal(dequantized.detach().cpu(), reference)
+    finally:
+        if os.path.exists(path):
+            os.unlink(path)
+
+
+def test_load_comfy_nvfp4_orig_shape_mismatch_raises():
+    """A marker whose orig_shape disagrees with the model config is a wrong
+    file for the class; reject so the base-repo fallback engages."""
+    import json
+    fd, path = tempfile.mkstemp(suffix='.safetensors')
+    try:
+        dim = 32
+        raw, _ = nvfp4_fixture(dim)
+        marker = json.dumps({'format': 'nvfp4', 'group_size': 16, 'orig_shape': [dim, dim * 2]})
+        raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(marker.encode()), dtype=torch.uint8)
+        write_fixture(raw, fd, path)
+
+        with ComfyTestEnv(dim):
+            spec = nt.TransformerSpec(cls=MockMiniTransformer)
+            try:
+                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={}, dtype=torch.float32)
+                raise AssertionError('expected OverrideArchMismatch')
+            except nt.OverrideArchMismatch as e:
+                assert 'orig_shape' in str(e)
+    finally:
+        if os.path.exists(path):
+            os.unlink(path)
+
+
 def test_detect_comfy_mixed_formats_raises():
     """One file mixing int8 and fp8 layers has no single SDNQ mapping; reject."""
     sd = {
@@ -1953,6 +2169,21 @@ def run_all():
     ]:
         run_test(cat, fn)
 
+    log.warning('=== nvfp4 codec / unswizzle / load ===')
+    cat = category('nvfp4')
+    for fn in [
+        test_unswizzle_block_scales_roundtrip,
+        test_nvfp4_codec_ocp_table,
+        test_nvfp4_pack_ocp_grid_roundtrip,
+        test_detect_comfy_nvfp4_marker_accepted,
+        test_detect_comfy_nvfp4_bad_group_size_raises,
+        test_detect_comfy_nvfp4_convrot_raises,
+        test_load_comfy_nvfp4_end_to_end,
+        test_load_comfy_nvfp4_metadata_form,
+        test_load_comfy_nvfp4_orig_shape_mismatch_raises,
+    ]:
+        run_test(cat, fn)
+
     log.warning('=== noop_converter detection ===')
     cat = category('noop')
     for fn in [