[PyTorch] Fix shape and size() for columnwise-only quantized tensors

[pggPL]

Description

Float8BlockwiseQTensor.shape, Float8TensorStorage.size() and Float8BlockwiseQTensorStorage.size() all report the logical shape of a tensor whose data is only stored columnwise/transposed. Each of the three gets it wrong in a different way.

1. Float8BlockwiseQTensor.shape returns the transposed buffer shape. Blockwise stores columnwise data transposed, and the fast path added in 9dac78e ("CPU Overhead Optimizations", #2559) does not undo that, so the property disagrees with both size() and the wrapper metadata:

q.set_usage(rowwise=False, columnwise=True)
t = q.make_empty((128, 256), dtype=torch.bfloat16, device="cuda")

t.shape                                  # (256, 128)  <- wrong
t.size()                                 # (128, 256)
torch._C.TensorBase.shape.__get__(t)     # (128, 256)

That commit added the equivalent property to four classes at once; Float8Tensor and NVFP4Tensor undo the reordering. MXFP8Tensor reads the same but is already correct, since it stores columnwise data untransposed.

2. size(dim) raises TypeError on the columnwise-only path, for FP8 and blockwise alike. Both storages forward *args to the underlying buffer and then reorder the result. size(dim) returns an int, which the reorder then indexes:

t.size(0)   # TypeError: 'int' object is not subscriptable   (FP8)
t.size(0)   # TypeError: 'int' object is not iterable        (blockwise)

Forwarding dim is also wrong in principle here, because buffer dim i is not logical dim i under a transposed layout. This predates #2559 — it came in with the classes themselves (#1505, #1513).

3. Float8TensorStorage.size() flattens to 2D, which contradicts dim() and Float8Tensor.shape for rank ≥ 3:

t = q.quantize(x)                                        # x is (4, 128, 256)
t.update_usage(rowwise_usage=False, columnwise_usage=True)

t.dim()      # 3
t.shape      # (4, 128, 256)
t.size()     # (128, 1024)   <- 2 entries

All of this is reachable from quantize() followed by update_usage(rowwise_usage=False), which is the normal way to release rowwise data once it is no longer needed.

NVFP4 is deliberately left alone: it reports columnwise-only tensors flattened and warns about it, and its shape and size() agree with each other.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Float8BlockwiseQTensor.shape undoes the columnwise transpose, with a 2D fast path.
• Float8BlockwiseQTensorStorage.size() and Float8TensorStorage.size() rebuild the full logical shape before applying a dim argument, instead of forwarding it to the buffer.
• Float8TensorStorage.size() applies the same rotation as Float8Tensor.shape rather than flattening to 2D.
• test_shape_matches_size in tests/pytorch/test_quantized_tensor.py covers every entry of _quantization_list, three usage combinations and 2D/3D shapes, asserting that shape and size() describe the same tensor and that size(dim) agrees with it for positive and negative dims.

Note on CPU overhead

#2559 introduced the shape property specifically to cut CPU overhead, so the fix keeps the rowwise path untouched and adds a 2D fast path on the columnwise-only branch. Measured on an RTX Ada, make_empty((128, 256)), columnwise-only, ns/call:

variant	ns
before (returns the wrong shape)	75
this PR: 2D fast path, index into torch.Size	186
torch.Tensor.size(self) (wrapper metadata)	197
intermediate list() + slices	311
Float8BlockwiseQTensorStorage.size() before this PR	456

.shape therefore stays cheaper than the size() that already computed the same rotation. Reading the wrapper metadata benchmarks similarly but was avoided: buffers are reassigned after construction in a number of places, and the metadata would not follow.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

---

Note on local verification: my workstation is an RTX Ada, so fp8_blockwise, mxfp8 and nvfp4 are not in _quantization_list there and only the fp8 variants of the new test ran locally. The other paths were verified by direct measurement across all seven quantizers, three usage combinations and both ranks, using make_empty, which allocates regardless of arch support. test_numerics shows 28 failures both with and without these changes, so they are pre-existing on this GPU.

One thing worth watching in CI: making Float8TensorStorage.size() rank-preserving changes the rank it returns for transpose-only tensors of rank ≥ 3. If anything in the stack relied on the 2D flattening, it would only surface on Hopper+.

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

Corrects logical shape reporting for columnwise-only FP8 tensors.

• Restores logical dimension ordering for transposed Float8 and blockwise storage buffers.
• Makes size() and size(dim) consistent with tensor shape for positive and negative dimensions.
• Adds coverage across quantization formats, usage modes, and 2D/3D shapes, while skipping unsupported row-scaled NVFP4 columnwise-only allocation.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tests/pytorch/test_quantized_tensor.py	Adds shape and size consistency coverage and correctly skips the unsupported row-scaled NVFP4 columnwise-only case before allocation.
transformer_engine/pytorch/tensor/float8_blockwise_tensor.py	Converts transposed columnwise buffer dimensions back into the logical blockwise tensor shape.
transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py	Reconstructs the complete logical shape before handling an optional dimension argument.
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Preserves rank when reconstructing transpose-only FP8 tensor sizes and applies dimension indexing afterward.

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into quantized_tenso..." | Re-trigger Greptile}

[pggPL]

@vthumbe1503

[vthumbe1503]

LGTM. Thanks for catching and fixing this bug. We probably need a wider test suite for FP8 BS. I had found similar code coverage missing for FP8 BS missing w.r.t dequantize not being autograd aware in this PR

Also, a minor note based on comments in the description, CPU overhead issues are more pronounced on Grace Systems, but indeed can be a problem in other systems as well as demonstrated in the description

[vthumbe1503]

We might want to dig deep into why we dont need to apply the same shape preservation for NVFP4 tensors. But given it already exists in the code base, makes sense to defer it to a different PR.

[vthumbe1503]

/te-ci pytorch