From ee3d998a5d90dc275703f0539770a52173de2718 Mon Sep 17 00:00:00 2001
From: silveroxides <ishimarukaito@gmail.com>
Date: Tue, 10 Nov 2026 17:29:49 +0100
Subject: [PATCH] Support w6a8_int8 quantization decoding and shape unpacking

---
 nodes/quantization_guard.py           |  8 ++++-
 tests/test_lora_quantization_guard.py | 50 +++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/nodes/quantization_guard.py b/nodes/quantization_guard.py
index eb76259..eaba3be 100644
--- a/nodes/quantization_guard.py
+++ b/nodes/quantization_guard.py
@@ -280,6 +280,12 @@ class DiffusionQuantization:
             if len(shape) != 2:
                 raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}")
             return shape[0], shape[1] * 2
+        if fmt == "w6a8_int8":
+            if len(shape) != 2:
+                raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}")
+            if (shape[1] * 8) % 6 != 0:
+                raise ValueError(f"[{self.label}] Invalid 6-bit packed width {shape[1]} for {key}")
+            return shape[0], shape[1] * 8 // 6
         return shape
 
     def logical_dtype(self, key: str, compute_dtype: torch.dtype) -> torch.dtype:
@@ -342,7 +348,7 @@ class DiffusionQuantization:
                 "quant_group_size": 64,
                 "linear_dtype": config.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
             }
-        elif fmt == "asym_w4a8_int8":
+        elif fmt in ("asym_w4a8_int8", "w6a8_int8"):
             if f"{stem}.weight_correction" in tensors:
                 raise ValueError(
                     f"[{self.label}] Core's W4A8 loader does not apply weight_correction for {key}"
diff --git a/tests/test_lora_quantization_guard.py b/tests/test_lora_quantization_guard.py
index 2e89d36..5d4efef 100644
--- a/tests/test_lora_quantization_guard.py
+++ b/tests/test_lora_quantization_guard.py
@@ -593,3 +593,53 @@ def test_nvfp4_stream_decode_matches_core_full_dequantization(monkeypatch, tmp_p
         assert actual.shape == original.shape
     finally:
         handler.__exit__(None, None, None)
+
+
+def test_w6a8_int8_stream_decode_matches_core_dequantization(tmp_path, modules):
+    from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+
+    layout_name = QUANT_ALGOS["w6a8_int8"]["comfy_tensor_layout"]
+    layout = get_layout_class(layout_name)
+    original = torch.randn(4, 256, dtype=torch.float32)
+    qdata, params = layout.quantize(original, bits=6, codebook=False)
+    path = tmp_path / "w6a8.safetensors"
+    key = "model.layer.weight"
+    metadata = {
+        "_quantization_metadata": json.dumps({
+            "layers": {
+                "model.layer": {
+                    "format": "w6a8_int8",
+                    "group_size": params.group_size,
+                    "convrot_groupsize": params.convrot_groupsize,
+                }
+            }
+        })
+    }
+    save_file({
+        key: qdata,
+        "model.layer.weight_s_rel": params.scale,
+        "model.layer.weight_s_channel": params.s_channel,
+    }, str(path), metadata=metadata)
+
+    handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
+    try:
+        decoder = modules["guard"].DiffusionQuantization(handler, str(path), "W6A8 test")
+        assert decoder.logical_shape(key) == tuple(original.shape)
+        keys = decoder.required_keys(key)
+        stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for source_key in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == source_key
+                tensors[source_key] = batch[0][1]
+            actual = decoder.decode(key, tensors, torch.float32)
+        finally:
+            for source_key in tensors:
+                handler.mark_processed(source_key)
+            stream.close()
+        expected = QuantizedTensor(qdata, layout_name, params).dequantize()
+        torch.testing.assert_close(actual, expected)
+        assert actual.shape == original.shape
+    finally:
+        handler.__exit__(None, None, None)