From 199e43751bb4417e481fd86f18da519da5bfca5f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 27 Jul 2026 15:21:18 +0200 Subject: [PATCH 1/3] [PyTorch] Fix Float8BlockwiseQTensor.shape for columnwise-only tensors Float8BlockwiseQTensor.shape is a fast path added to avoid PyObject lookups. For a columnwise-only tensor it returned the raw columnwise buffer shape, but blockwise stores columnwise data transposed, so the property disagreed with both Float8BlockwiseQTensorStorage.size() and the wrapper's own metadata: make_empty((128, 256)) yielded .shape == (256, 128) while .size() == (128, 256). Apply the same reorder size() does. Float8Tensor and NVFP4Tensor already de-transpose in their equivalent fast paths; MXFP8Tensor needs no change because it stores columnwise data untransposed. Cover the invariant for every quantization scheme and usage combination with test_shape_matches_size, rather than only for blockwise. Introduced in 9dac78e76 ("CPU Overhead Optimizations", #2559). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_quantized_tensor.py | 32 +++++++++++++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 4 ++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2f77ecd66..75e4f694ed 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -757,6 +757,38 @@ def test_shape_with_none_data( f"after setting data to None on {type(x_test).__name__}" ) + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize( + "rowwise, columnwise", + [(True, True), (True, False), (False, True)], + ids=["rowwise_columnwise", "rowwise_only", "columnwise_only"], + ) + def test_shape_matches_size( + self, + *, + quantization: str, + rowwise: bool, + columnwise: bool, + shape: Iterable[int] = (128, 256), + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """shape and size() agree with the requested shape for every usage combination. + + Both are fast paths derived from whichever data buffer is present, and + classes that store columnwise data transposed have to undo that. A + columnwise-only tensor is where the two can drift apart. + """ + quantizer = make_quantizer(quantization, device=device) + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + if (quantizer.rowwise_usage, quantizer.columnwise_usage) != (rowwise, columnwise): + pytest.skip(f"{quantization} does not support this usage combination") + + x = quantizer.make_empty(shape, dtype=dtype, device=device) + + assert tuple(x.shape) == tuple(shape), f"{type(x).__name__}.shape is {tuple(x.shape)}" + assert tuple(x.size()) == tuple(shape), f"{type(x).__name__}.size() is {tuple(x.size())}" + @pytest.mark.parametrize( "quantization", _quantization_list + (["nvfp4_2d"] if nvfp4_available else []), diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 09fde86f17..ce5f57e73c 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -551,7 +551,9 @@ def shape(self): if self._rowwise_data is not None: return self._rowwise_data.shape if self._columnwise_data is not None: - return self._columnwise_data.shape + # Columnwise data is stored transposed, matching size() in the storage. + dims = list(self._columnwise_data.shape) + return torch.Size(dims[1:] + dims[:1]) return torch.Tensor.size(self) @property From 0001d1bb7debbfb0968458ee1108710fede50bf1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 27 Jul 2026 16:06:00 +0200 Subject: [PATCH 2/3] Skip row-scaled NVFP4 columnwise-only case and speed up the 2D shape path Row-scaled NVFP4 accepts set_usage(rowwise=False) but its allocator asserts on rowwise usage, so test_shape_matches_size hit an NVTE_CHECK failure instead of skipping. Filter that combination out before set_usage. Also give Float8BlockwiseQTensor.shape a 2D fast path: indexing torch.Size directly measures ~186 ns/call against ~311 ns for building an intermediate list, on the columnwise-only branch. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_quantized_tensor.py | 4 ++++ .../pytorch/tensor/float8_blockwise_tensor.py | 6 ++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 75e4f694ed..4784710ac4 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -780,6 +780,10 @@ def test_shape_matches_size( columnwise-only tensor is where the two can drift apart. """ quantizer = make_quantizer(quantization, device=device) + # Row-scaled NVFP4 accepts set_usage(rowwise=False) but rejects the + # allocation itself, so it has to be filtered out up front. + if getattr(quantizer, "row_scaled_nvfp4", False) and not rowwise: + pytest.skip(f"{quantization} requires rowwise usage") quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) if (quantizer.rowwise_usage, quantizer.columnwise_usage) != (rowwise, columnwise): pytest.skip(f"{quantization} does not support this usage combination") diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ce5f57e73c..6699cb1ea1 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -552,8 +552,10 @@ def shape(self): return self._rowwise_data.shape if self._columnwise_data is not None: # Columnwise data is stored transposed, matching size() in the storage. - dims = list(self._columnwise_data.shape) - return torch.Size(dims[1:] + dims[:1]) + dims = self._columnwise_data.shape + if len(dims) == 2: + return torch.Size((dims[1], dims[0])) + return torch.Size(tuple(dims[1:]) + (dims[0],)) return torch.Tensor.size(self) @property From 0efc6c5257d0fe564da94f87261031640500b804 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 27 Jul 2026 16:33:04 +0200 Subject: [PATCH 3/3] [PyTorch] Fix size() on the columnwise-only path for FP8 and blockwise Both storages forwarded *args straight to the underlying buffer, then reordered the result. That works for the rowwise branch, where the buffer matches the logical shape, but not for the transposed one: size(dim) returned an int from the buffer, which the reorder then tried to index, so every size(dim) call raised TypeError. Forwarding dim is also wrong in principle there, since buffer dim i is not logical dim i. Rebuild the logical shape in full first, then index into it. Float8TensorStorage also flattened the transpose-only shape to 2D, which disagreed with both dim() and Float8Tensor.shape for rank >= 3; it now applies the same rotation the shape property does. Reachable from quantize() followed by update_usage(rowwise_usage=False). NVFP4 is left alone: it reports columnwise-only tensors flattened on purpose and warns about it, and shape and size() agree there. Extend test_shape_matches_size over 2D and 3D shapes, asserting that shape and size() describe the same tensor and that size(dim) agrees with it for positive and negative dims. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_quantized_tensor.py | 32 ++++++++++++++----- .../float8_blockwise_tensor_storage.py | 15 +++++---- .../tensor/storage/float8_tensor_storage.py | 13 ++++++-- 3 files changed, 43 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 4784710ac4..fd8e220205 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -763,21 +763,23 @@ def test_shape_with_none_data( [(True, True), (True, False), (False, True)], ids=["rowwise_columnwise", "rowwise_only", "columnwise_only"], ) + @pytest.mark.parametrize("shape", [(128, 256), (4, 128, 256)], ids=["2d", "3d"]) def test_shape_matches_size( self, *, quantization: str, rowwise: bool, columnwise: bool, - shape: Iterable[int] = (128, 256), + shape: Iterable[int], dtype: torch.dtype = torch.bfloat16, device: torch.device = "cuda", ) -> None: - """shape and size() agree with the requested shape for every usage combination. + """shape, size() and size(dim) stay consistent for every usage combination. - Both are fast paths derived from whichever data buffer is present, and - classes that store columnwise data transposed have to undo that. A - columnwise-only tensor is where the two can drift apart. + Both shape and size() are derived from whichever data buffer is present, + and classes that store columnwise data transposed have to undo that. A + columnwise-only tensor is where they can drift apart -- from each other, + and from the shape the tensor was allocated with. """ quantizer = make_quantizer(quantization, device=device) # Row-scaled NVFP4 accepts set_usage(rowwise=False) but rejects the @@ -789,9 +791,23 @@ def test_shape_matches_size( pytest.skip(f"{quantization} does not support this usage combination") x = quantizer.make_empty(shape, dtype=dtype, device=device) - - assert tuple(x.shape) == tuple(shape), f"{type(x).__name__}.shape is {tuple(x.shape)}" - assert tuple(x.size()) == tuple(shape), f"{type(x).__name__}.size() is {tuple(x.size())}" + name = type(x).__name__ + + # shape and size() must describe the same tensor, whichever buffer they + # end up reading. + assert tuple(x.shape) == tuple(x.size()), f"{name}: {tuple(x.shape)} vs {tuple(x.size())}" + + # size(dim) must agree with the full shape, including negative indices. + # It cannot be served by forwarding dim to a transposed buffer. + for dim in range(len(x.shape)): + assert x.size(dim) == x.shape[dim], f"{name}.size({dim}) is {x.size(dim)}" + neg = dim - len(x.shape) + assert x.size(neg) == x.shape[neg], f"{name}.size({neg}) is {x.size(neg)}" + + # NVFP4 deliberately reports columnwise-only tensors flattened to 2D and + # warns about it, so only the ranks it preserves are checked here. + if not (isinstance(quantizer, NVFP4Quantizer) and not rowwise and len(shape) > 2): + assert tuple(x.shape) == tuple(shape), f"{name}.shape is {tuple(x.shape)}" @pytest.mark.parametrize( "quantization", diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 464a7c3b23..06516d0b9e 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -317,12 +317,15 @@ def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._rowwise_data is not None: return self._rowwise_data.size(*args, **kwargs) - dims = list(self._columnwise_data.size(*args, **kwargs)) - reordered = [] - for i in range(1, len(dims)): - reordered.append(dims[i]) - reordered.append(dims[0]) - return torch.Size(reordered) + # Columnwise data is stored transposed, so a dim argument cannot be + # forwarded to it: rebuild the logical shape first, then index into it. + dims = self._columnwise_data.shape + if len(dims) == 2: + shape = torch.Size((dims[1], dims[0])) + else: + shape = torch.Size(tuple(dims[1:]) + (dims[0],)) + dim = args[0] if args else kwargs.get("dim") + return shape if dim is None else shape[dim] def view(self, shape): """Reshape the leading (token) dims without dequantizing. diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 374d0e1e72..429ddde97f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -5,7 +5,6 @@ """Mixin class holding data specific for Float8Tensor""" from __future__ import annotations -import math from typing import Any, Dict, Optional, Tuple, Union import torch @@ -179,8 +178,16 @@ def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._data is not None: return self._data.size(*args, **kwargs) - size = self._transpose.size(*args, **kwargs) - return torch.Size([size[-1], math.prod(size[:-1])]) + # The transpose is stored as [last, *leading], so a dim argument cannot + # be forwarded to it: rebuild the logical shape first, then index into + # it. This matches the shape property on Float8Tensor. + dims = self._transpose.shape + if len(dims) == 2: + shape = torch.Size((dims[1], dims[0])) + else: + shape = torch.Size(tuple(dims[1:]) + (dims[0],)) + dim = args[0] if args else kwargs.get("dim") + return shape if dim is None else shape[dim] @property def device(self):