feat(model): load convrot comfy_quant checkpoints via sdnq hadamard

ConvRot is the regular Hadamard rotation SDNQ implements: identical
construction, normalization, axis, and dequant order. Per-layer markers
map onto the dequantizer's use_hadamard and hadamard_group_size; group
sizes must be powers of 4 and divide in_features, else base-repo
fallback. Detection carries per-layer metadata since files mix plain
and rotated layers.

pipelines/native_transformer.py

@@ -427,21 +427,25 @@ def check_forbidden_markers(
             )
 
 
-def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] | None:
+def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict], str] | None:
     """Detect ``comfy_quant`` pre-quantized layers in a state dict.
 
     ComfyUI-format quantized checkpoints mark each quantized layer with a
     ``<name>.comfy_quant`` uint8 tensor whose bytes are a JSON object naming
     the storage format, alongside a ``<name>.weight_scale`` tensor. Returns
-    the set of marked module names and the format string, or ``None`` when no
-    markers are present. Raises :class:`OverrideArchMismatch` for malformed
-    markers, unsupported formats, or a file mixing formats across layers, so
-    the caller's base-repo fallback engages.
+    a mapping of marked module names to their parsed marker JSON plus the
+    file's format string, or ``None`` when no markers are present. ConvRot
+    markers (``convrot``/``convrot_groupsize``) are accepted per layer: the
+    rotation is the same regular Hadamard SDNQ implements, so flagged layers
+    map onto ``use_hadamard`` (regular Hadamards only exist for power-of-4
+    group sizes). Raises :class:`OverrideArchMismatch` for
+    malformed markers, unsupported formats or group sizes, or a file mixing
+    formats across layers, so the caller's base-repo fallback engages.
     """
     marker_keys = [key for key in state_dict if key.endswith(COMFY_QUANT_MARKER)]
     if not marker_keys:
         return None
-    marked: set[str] = set()
+    marked: dict[str, dict] = {}
     formats: set[str] = set()
     for key in marker_keys:
         name = key[: -len(COMFY_QUANT_MARKER)]
@@ -454,14 +458,14 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str]
                 f"for {name!r} is malformed ({type(e).__name__}: {e})"
             ) from e
         if meta.get("convrot"):
-            # ConvRot-flagged layers may require an inverse rotation at
-            # runtime; adopting them as plain tensors would corrupt outputs.
-            log.error(f'Load model: type={type_name} quant=comfy format={fmt} convrot=true not supported')
-            raise OverrideArchMismatch(
-                f"Load model: type={type_name} native_transformer comfy_quant layer "
-                f"{name!r} carries a convrot rotation flag, not supported"
-            )
-        marked.add(name)
+            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:
+                raise OverrideArchMismatch(
+                    f"Load model: type={type_name} native_transformer comfy_quant layer "
+                    f"{name!r} convrot_groupsize={group_size} is not a power of 4"
+                )
+        marked[name] = meta
         formats.add(fmt)
     unsupported = sorted(formats - set(COMFY_QUANT_FORMATS))
     if unsupported:
@@ -644,7 +648,7 @@ def build_component_prequantized(
     state_dict: dict,
     config: dict,
     cls: type,
-    marked_names: set[str],
+    marker_meta: dict[str, dict],
     comfy_format: str,
     dtype,
     acceptable_missing: tuple[str, ...],
@@ -656,7 +660,9 @@ def build_component_prequantized(
 
     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.
+    (``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.
     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
@@ -679,6 +685,7 @@ 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
+    marked_names = set(marker_meta)
     sd = remap_comfy_quant(state_dict, marked_names)
 
     # Civitai relabels these containers freely; trust the marker only as far
@@ -722,6 +729,14 @@ def build_component_prequantized(
                 f"Load model: transformer=native {component_name} comfy_quant marked "
                 f"module {name!r} is {linear.__class__.__name__}, expected Linear"
             )
+        # convrot weights are stored rotated by SDNQ's regular Hadamard; the dequantizer undoes it
+        use_hadamard = bool(marker_meta[name].get("convrot"))
+        hadamard_group_size = int(marker_meta[name].get("convrot_groupsize", 256)) if use_hadamard else 256
+        if use_hadamard and linear.in_features % hadamard_group_size != 0:
+            raise OverrideArchMismatch(
+                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}"
+            )
         linear.sdnq_dequantizer = SDNQDequantizer(
             result_dtype=target_dtype,
             result_shape=None,
@@ -730,14 +745,14 @@ def build_component_prequantized(
             quantized_weight_shape=torch.Size((linear.out_features, linear.in_features)),
             weights_dtype=weights_dtype,
             quantized_matmul_dtype=matmul_dtype,
-            hadamard_group_size=256,
+            hadamard_group_size=hadamard_group_size,
             group_size=-1,
             svd_rank=32,
             svd_steps=8,
             use_quantized_matmul=False,
             re_quantize_for_matmul=False,
             use_stochastic_rounding=False,
-            use_hadamard=False,
+            use_hadamard=use_hadamard,
             layer_class_name="Linear",
         )
         wrapped = get_sdnq_wrapper_class(linear, dequant_forward)
@@ -857,17 +872,18 @@ def build_component(
 
         comfy_quant = detect_comfy_quant(state_dict, cls.__name__) if component_name == "transformer" else None
         if comfy_quant is not None:
-            marked_names, comfy_format = comfy_quant
+            marker_meta, comfy_format = comfy_quant
+            convrot_count = sum(1 for meta in marker_meta.values() if meta.get("convrot"))
             log.info(
                 f'Load model: transformer=native {component_name} quant=comfy '
-                f'format={comfy_format} layers={len(marked_names)} keys={len(state_dict)} cls={cls.__name__}'
+                f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} keys={len(state_dict)} cls={cls.__name__}'
             )
             component = build_component_prequantized(
                 component_name=component_name,
                 state_dict=state_dict,
                 config=config,
                 cls=cls,
-                marked_names=marked_names,
+                marker_meta=marker_meta,
                 comfy_format=comfy_format,
                 dtype=dtype,
                 acceptable_missing=acceptable_missing,

test/test-native-transformer.py

@@ -323,7 +323,8 @@ def test_detect_comfy_valid_markers():
     detected = nt.detect_comfy_quant(sd, 'Test')
     assert detected is not None
     marked, fmt = detected
-    assert marked == {'blocks.0.attn.wq', 'blocks.0.mlp.up'}
+    assert set(marked) == {'blocks.0.attn.wq', 'blocks.0.mlp.up'}
+    assert marked['blocks.0.attn.wq'] == {'format': 'int8_tensorwise'}
     assert fmt == 'int8_tensorwise'
 
 
@@ -351,20 +352,37 @@ def test_detect_comfy_malformed_marker_raises():
         assert 'malformed' in str(e)
 
 
-def test_detect_comfy_convrot_flag_raises():
-    """Markers carrying convrot=true may need runtime inverse rotation; adopting
-    them as plain tensors would corrupt outputs, so they must be rejected."""
+def test_detect_comfy_convrot_accepted():
+    """ConvRot markers map onto SDNQ's Hadamard support; detection carries the
+    per-layer fields through instead of rejecting."""
     import json
-    payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256})
+    payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256, 'per_row': True})
     sd = {
         'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
         'blocks.0.attn.wq.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8),
+        'blocks.0.mlp.up.weight': torch.zeros((4, 4), dtype=torch.int8),
+        'blocks.0.mlp.up.comfy_quant': comfy_marker('int8_tensorwise'),
+    }
+    marked, fmt = nt.detect_comfy_quant(sd, 'Test')
+    assert fmt == 'int8_tensorwise'
+    assert marked['blocks.0.attn.wq'].get('convrot') is True
+    assert marked['blocks.0.attn.wq'].get('convrot_groupsize') == 256
+    assert not marked['blocks.0.mlp.up'].get('convrot')
+
+
+def test_detect_comfy_convrot_bad_groupsize_raises():
+    """Regular Hadamards only exist for power-of-4 sizes; any other convrot
+    group size cannot be a compatible rotation."""
+    import json
+    payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 8})
+    sd = {
+        'blocks.0.attn.wq.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)
+        assert 'power of 4' in str(e)
 
 
 def test_detect_comfy_marker_missing_format_field_raises():
@@ -1446,25 +1464,26 @@ def test_load_fused_bf16_end_to_end():
             os.unlink(path)
 
 
-def test_load_fused_convrot_falls_back():
-    """The one real Ideogram 4 civitai file is convrot-flagged: conversion
-    succeeds, then detection rejects it into the base-repo fallback."""
+def test_load_fused_convrot_end_to_end():
+    """Fused qkv with a convrot marker: the converter duplicates the marker to
+    the three split layers, each of which loads with hadamard geometry. Row
+    slicing commutes with the in-features rotation, so slices stay lossless."""
     import json as json_mod
     fd, path = tempfile.mkstemp(suffix='.safetensors')
     try:
-        dim = 8
+        dim, group_size = 8, 4
         raw = fused_fixture(dim, quantized=True)
-        payload = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256, 'per_row': True})
+        payload = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': group_size, 'per_row': True})
         raw['model.diffusion_model.layers.0.attention.qkv.comfy_quant'] = torch.tensor(list(payload.encode()), dtype=torch.uint8)
+        raw['model.diffusion_model.layers.0.attention.qkv.weight_scale'] = torch.rand(3 * dim, 1, dtype=torch.float32)
         write_fixture(raw, fd, path)
         with ComfyTestEnv(dim):
-            raised = False
-            try:
-                nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={})
-            except nt.OverrideArchMismatch as e:
-                raised = True
-                assert 'convrot' in str(e)
-            assert raised, 'expected OverrideArchMismatch'
+            transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32)
+        attn = transformer.layers[0].attention
+        for name in ('to_q', 'to_k', 'to_v'):
+            dq = getattr(attn, name).sdnq_dequantizer
+            assert dq.use_hadamard is True, f'{name} lost the convrot flag'
+            assert dq.hadamard_group_size == group_size
     finally:
         if os.path.exists(path):
             os.unlink(path)
@@ -1534,6 +1553,103 @@ def test_load_transformer_secondary_slot_syncs_tracker():
             os.unlink(path)
 
 
+def regular_hadamard(n: int) -> torch.Tensor:
+    """Regular Hadamard construction the ConvRot format and SDNQ share:
+    4x4 base, Kronecker recursion, 1/sqrt(n) normalization. Symmetric and
+    involutory, so rotation and inverse are the same matrix."""
+    h4 = torch.tensor([[1., 1., 1., -1.], [1., 1., -1., 1.], [1., -1., 1., 1.], [-1., 1., 1., 1.]])
+    h = h4
+    size = 4
+    while size < n:
+        h = torch.kron(h, h4)
+        size *= 4
+    return h / (n ** 0.5)
+
+
+def comfy_convrot_quantize(weight: torch.Tensor, group_size: int):
+    """Reference ConvRot int8 quantizer (mirrors the ComfyUI runtime): rotate
+    grouped in-features by the regular Hadamard, then row-wise symmetric int8."""
+    h = regular_hadamard(group_size)
+    out_f, in_f = weight.shape
+    rotated = (weight.reshape(out_f, -1, group_size) @ h.T).reshape(out_f, in_f)
+    scale = rotated.abs().amax(dim=-1, keepdim=True) / 127
+    q = rotated.div(scale).round().clamp(-128, 127).to(torch.int8)
+    return q, scale
+
+
+def test_load_comfy_convrot_end_to_end():
+    """ConvRot parity: a layer quantized with the reference ConvRot math
+    loads as an SDNQ hadamard layer whose dequant matches the reference dequant
+    and recovers the original weight within int8 quantization error. The plain
+    second layer proves per-layer mixing within one file."""
+    import json as json_mod
+    fd, path = tempfile.mkstemp(suffix='.safetensors')
+    try:
+        dim, group_size = 8, 4
+        original = torch.randn(dim, dim)
+        q, scale = comfy_convrot_quantize(original, group_size)
+        marker = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': group_size, 'per_row': True})
+        raw = {
+            'model.diffusion_model.in_proj.weight': q,
+            'model.diffusion_model.in_proj.weight_scale': scale,
+            'model.diffusion_model.in_proj.comfy_quant': torch.tensor(list(marker.encode()), dtype=torch.uint8),
+            'model.diffusion_model.in_proj.bias': torch.zeros(dim, dtype=torch.float16),
+            'model.diffusion_model.out_proj.weight': torch.randint(-128, 127, (dim, dim), dtype=torch.int8),
+            'model.diffusion_model.out_proj.weight_scale': torch.tensor(0.03125, dtype=torch.float32),
+            'model.diffusion_model.out_proj.comfy_quant': comfy_marker('int8_tensorwise'),
+            'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16),
+        }
+        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
+        dq = in_proj.sdnq_dequantizer
+        assert dq.use_hadamard is True
+        assert dq.hadamard_group_size == group_size
+        assert tuple(in_proj.scale.shape) == (dim, 1)
+
+        dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None).detach().cpu()
+        h = regular_hadamard(group_size)
+        reference = ((q.float() * scale).reshape(dim, -1, group_size) @ h.T).reshape(dim, dim)
+        assert torch.allclose(dequantized, reference, atol=1e-5), 'SDNQ hadamard dequant diverges from the ConvRot reference'
+        assert torch.allclose(dequantized, original, atol=0.15), 'dequant does not recover the original weight'
+
+        # plain layer in the same file stays rotation-free
+        assert transformer.out_proj.sdnq_dequantizer.use_hadamard is False
+    finally:
+        if os.path.exists(path):
+            os.unlink(path)
+
+
+def test_load_comfy_convrot_nondivisible_falls_back():
+    """A convrot group size that does not divide in_features cannot be undone;
+    the file must fall back to the base repo."""
+    import json as json_mod
+    fd, path = tempfile.mkstemp(suffix='.safetensors')
+    try:
+        dim = 8
+        raw = comfy_fixture(dim)
+        marker = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 16})
+        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)
+            raised = False
+            try:
+                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
+            except nt.OverrideArchMismatch as e:
+                raised = True
+                assert 'does not divide' in str(e)
+            assert raised, 'expected OverrideArchMismatch'
+    finally:
+        if os.path.exists(path):
+            os.unlink(path)
+
+
 def test_build_component_comfy_preempts_sdnq_fresh_quant():
     """When SDNQ on-load quant settings are active (quant_type=SDNQConfig), a
     comfy_quant file must still take the pre-quantized path: fresh quant of
@@ -1613,7 +1729,8 @@ def run_all():
         test_detect_comfy_valid_markers,
         test_detect_comfy_unsupported_format_raises,
         test_detect_comfy_malformed_marker_raises,
-        test_detect_comfy_convrot_flag_raises,
+        test_detect_comfy_convrot_accepted,
+        test_detect_comfy_convrot_bad_groupsize_raises,
         test_detect_comfy_marker_missing_format_field_raises,
         test_remap_comfy_renames_and_reshapes_scale,
         test_remap_comfy_passes_unmarked_keys_verbatim,
@@ -1690,6 +1807,8 @@ def run_all():
         test_load_comfy_marker_dtype_mismatch_raises,
         test_load_comfy_unsupported_format_raises_mismatch,
         test_load_comfy_marker_for_unknown_module_raises_mismatch,
+        test_load_comfy_convrot_end_to_end,
+        test_load_comfy_convrot_nondivisible_falls_back,
         test_load_transformer_syncs_loaded_unet,
         test_load_transformer_secondary_slot_syncs_tracker,
         test_build_component_comfy_preempts_sdnq_fresh_quant,
@@ -1710,7 +1829,7 @@ def run_all():
         test_ideogram4_converter_does_not_mutate_input,
         test_load_fused_comfy_end_to_end,
         test_load_fused_bf16_end_to_end,
-        test_load_fused_convrot_falls_back,
+        test_load_fused_convrot_end_to_end,
     ]:
         run_test(cat, fn)