diff --git a/comfy_kitchen/tensor/base.py b/comfy_kitchen/tensor/base.py index 92e8778f..11ba712d 100644 --- a/comfy_kitchen/tensor/base.py +++ b/comfy_kitchen/tensor/base.py @@ -159,10 +159,14 @@ def __new__( qdata: torch.Tensor, layout_cls: str, params: Any, + outer_stride: Any = None, ): + # outer_stride is set by __tensor_unflatten__ under torch.compile dynamic shapes, where the + # rebuilt wrapper must report the exact (symbolic) stride PyTorch asserts. None -> contiguous. return torch.Tensor._make_wrapper_subclass( cls, params.orig_shape, + strides=outer_stride, device=qdata.device, dtype=params.orig_dtype, requires_grad=False, @@ -173,6 +177,7 @@ def __init__( qdata: torch.Tensor, layout_cls: str, params: Any, + outer_stride: Any = None, ): assert isinstance(layout_cls, str) self._qdata = qdata @@ -339,6 +344,16 @@ def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride): params_kwargs[field_name] = inner_tensors[attr_name] params = ctx["params_class"](**params_kwargs) + # torch.compile passes outer_size/outer_stride as sequences the rebuilt subclass must report + # exactly (asserted in torch._subclasses.meta_utils), and under dynamic shapes they can differ + # from the flatten-time orig_shape -- so adopt outer_size as orig_shape and forward the stride. + # Other callers (e.g. comfy.memory_management.interpret_gathered_like) pass scalar 0 + # placeholders meaning "keep stored shape", so only act on real sequences. + if isinstance(outer_size, (tuple, list, torch.Size)): + params = dataclasses.replace(params, orig_shape=tuple(outer_size)) + return QuantizedTensor( + inner_tensors["_qdata"], ctx["layout_cls"], params, outer_stride=outer_stride, + ) return QuantizedTensor(inner_tensors["_qdata"], ctx["layout_cls"], params) # ==================== Torch Dispatch ==================== diff --git a/tests/test_unflatten_dynamic_shapes.py b/tests/test_unflatten_dynamic_shapes.py new file mode 100644 index 00000000..f9c8cd50 --- /dev/null +++ b/tests/test_unflatten_dynamic_shapes.py @@ -0,0 +1,171 @@ +"""Regression tests: QuantizedTensor.__tensor_unflatten__ must honor outer_size/outer_stride. + +PyTorch's subclass fakeification (torch._subclasses.meta_utils) asserts that a tensor subclass +rebuilt via __tensor_unflatten__ reports exactly the outer_size/outer_stride it was given. Under +torch.compile with dynamic shapes these are (possibly symbolic, e.g. a fused activation row dim +s31*s81) and differ from the static orig_shape captured at quantize time. Ignoring them raises: + + AssertionError: Expected return value from ...__tensor_unflatten__() to have shape + equal to (s31, 2560), but got: torch.Size([s31*s81, 2560]) + +which crashed mxfp8 models under torch.compile. All tests except TestUnflattenRealCompileCUDA are +CPU-only and use eager-quantizable real layouts. +""" +import dataclasses + +import pytest +import torch + +from comfy_kitchen.tensor import QuantizedTensor, get_layout_class +from comfy_kitchen.tensor.base import register_layout_class + +# INT8 weight-quant and FP8 both have CPU code paths -> usable as real layouts on CPU. +CPU_LAYOUTS = [ + pytest.param("TensorWiseINT8Layout", id="int8"), + pytest.param("TensorCoreFP8Layout", id="fp8"), +] + + +def _require(layout_cls): + try: + get_layout_class(layout_cls) + except KeyError: + pytest.skip(f"{layout_cls} not registered in this build") + + +def _flatten(qt): + names, ctx = qt.__tensor_flatten__() + return {n: getattr(qt, n) for n in names}, ctx + + +class _DummyLayout: + pass + + +register_layout_class("Dummy", _DummyLayout) + + +class TestUnflattenOuterSizeCPU: + """Synthetic + real-layout coverage of the outer_size/outer_stride threading (CPU-only).""" + + @dataclasses.dataclass(frozen=True) + class _DummyParams: + scale: torch.Tensor + orig_dtype: torch.dtype + orig_shape: tuple + + def _tensor_fields(self): + return ["scale"] + + def _dummy_qt(self): + return QuantizedTensor( + torch.zeros(8, 16, dtype=torch.uint8), + "Dummy", + self._DummyParams(scale=torch.ones(8, 1), orig_dtype=torch.bfloat16, orig_shape=(8, 16)), + ) + + def test_honors_concrete_outer_size_and_stride(self): + inner, ctx = _flatten(self._dummy_qt()) + outer_size, outer_stride = (5, 16), (16, 1) + rebuilt = QuantizedTensor.__tensor_unflatten__(inner, ctx, outer_size, outer_stride) + assert tuple(rebuilt.shape) == outer_size + assert tuple(rebuilt.stride()) == outer_stride + # orig_shape must follow outer_size so layout slice-back stays consistent. + assert tuple(rebuilt._params.orig_shape) == outer_size + + def test_scalar_placeholder_keeps_orig_shape(self): + # comfy.memory_management.interpret_gathered_like calls __tensor_unflatten__(actuals, ctx, 0, 0) + # meaning "keep the stored shape"; a scalar 0 must NOT be treated as a size. + inner, ctx = _flatten(self._dummy_qt()) + rebuilt = QuantizedTensor.__tensor_unflatten__(inner, ctx, 0, 0) + assert tuple(rebuilt.shape) == (8, 16) + + @pytest.mark.parametrize("layout_cls", CPU_LAYOUTS) + def test_eager_roundtrip_placeholder(self, layout_cls): + _require(layout_cls) + torch.manual_seed(0) + w = torch.randn(48, 96, dtype=torch.float32) + qt = QuantizedTensor.from_float(w, layout_cls) + inner, ctx = _flatten(qt) + rebuilt = QuantizedTensor.__tensor_unflatten__(inner, ctx, 0, 0) + assert tuple(rebuilt.shape) == tuple(qt.shape) + assert rebuilt._layout_cls == qt._layout_cls + assert torch.equal(rebuilt.dequantize(), qt.dequantize()) + + +class TestUnflattenSymbolicShapesCPU: + """The production trigger: rebuild a REAL-layout QuantizedTensor under FakeTensorMode with a + COMPOUND symbolic outer_size (s_b*s_s), exactly what automatic-dynamic produces. Drives the real + torch._subclasses.meta_utils path. FAILS before the fix (rebuilt reports static shape), passes + after. CPU-only.""" + + @pytest.mark.parametrize("layout_cls", CPU_LAYOUTS) + def test_unflatten_compound_symbolic_outer_size(self, layout_cls): + _require(layout_cls) + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + torch.manual_seed(0) + w = torch.randn(64, 128, dtype=torch.float32) + qt = QuantizedTensor.from_float(w, layout_cls) + _, ctx = qt.__tensor_flatten__() + cols = qt.shape[1] + + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + # Harvest two backed dynamic symbols (like s31, s81) from a dynamic helper tensor. + helper = torch.randn(2, 3) + torch._dynamo.maybe_mark_dynamic(helper, 0) + torch._dynamo.maybe_mark_dynamic(helper, 1) + fh = fake_mode.from_tensor(helper, symbolic_context=None) + compound = fh.shape[0] * fh.shape[1] # e.g. s26*s49, like production s31*s81 + + inner = {n: fake_mode.from_tensor(getattr(qt, n)) for n in qt.__tensor_flatten__()[0]} + outer_size = (compound, cols) + outer_stride = (cols, 1) + + rebuilt = QuantizedTensor.__tensor_unflatten__(inner, ctx, outer_size, outer_stride) + + # This is exactly the assertion torch._subclasses.meta_utils makes internally. + assert tuple(rebuilt.shape) == tuple(outer_size) + assert tuple(rebuilt.stride()) == tuple(outer_stride) + assert tuple(rebuilt._params.orig_shape) == tuple(outer_size) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestUnflattenRealCompileCUDA: + """End-to-end production reproduction: torch.compile over a real-layout model run across TWO + different batch shapes forces automatic-dynamic, which fakeifies the QuantizedTensor with a + symbolic outer_size. Pre-fix this raises the meta_utils AssertionError. Requires a GPU whose + SM supports the layout's fast matmul (Blackwell SM12.0).""" + + @pytest.mark.parametrize("layout_cls", ["TensorWiseINT8Layout", "TensorCoreFP8Layout"]) + def test_compile_dynamic_two_shapes(self, layout_cls): + from comfy_kitchen.tensor import get_layout_class + + lc = get_layout_class(layout_cls) + if not lc.supports_fast_matmul(): + pytest.skip(f"{layout_cls} fast matmul unsupported on this GPU") + + torch.manual_seed(0) + w = torch.randn(128, 64, device="cuda", dtype=torch.bfloat16) + qw = QuantizedTensor.from_float(w, layout_cls) + + class M(torch.nn.Module): + def __init__(self, weight): + super().__init__() + self.w = weight + + def forward(self, x): + # Runtime-quantize the activation -> the subclass carries the dynamic row dim. + xq = QuantizedTensor.from_float(x, layout_cls) + return torch.nn.functional.linear(xq, self.w) + + cm = torch.compile(M(qw), dynamic=None) # automatic-dynamic on 2nd shape + with torch.no_grad(): + o1 = cm(torch.randn(8, 64, device="cuda", dtype=torch.bfloat16)) + o2 = cm(torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)) # recompile -> symbolic + assert o1.shape == (8, 128) + assert o2.shape == (16, 128)