Fix nvfp4 convert_and_update_tensor shape check

[skydoorkai]

Description

This is to fix #2607

For nvfp4's columnwise data , it is using enforced 2D shape. Thus, the original check would fail if rowwise_data shape is 3D shape.

To fix :
(1) expected_data should be enforced into 2D shape from rowwise_data's shape.
(2) use rowwise_data's shape as the “ground truth" shape.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

[greptile-apps]

Greptile Summary

This PR fixes a shape validation bug in nvfp4 quantization when using 3D (or N-D) tensors. Previously, the code directly compared rowwise data shapes (which preserve N-D structure) with columnwise data shapes (enforced to 2D), causing failures.

The fix introduces compressShapeTo2D() to flatten both reconstructed shapes to 2D before comparison, then uses the rowwise shape as the authoritative representation. A regression test validates that 3D tensors [32, 4, 128] now work correctly with both quantization modes.

Confidence Score: 4/5

• Safe to merge with minor testing considerations
• The fix correctly addresses the reported issue by normalizing shape comparisons to 2D. The implementation is well-documented with a clear docstring, and includes a regression test. Score is 4/5 rather than 5/5 due to limited test coverage (only one 3D shape tested) and lack of edge case validation
• No files require special attention - the changes are straightforward

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/quantizer.cpp	Added compressShapeTo2D helper and updated shape comparison logic to handle N-D rowwise tensors vs 2D columnwise tensors
tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py	Added regression test for 3D shape quantization with both rowwise and columnwise modes

_{Last reviewed commit: 36c173b}

[greptile-apps]

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (2)

transformer_engine/pytorch/csrc/quantizer.cpp
Missing include for std::accumulate

compressShapeTo2D uses std::accumulate and std::multiplies, but this file doesn’t include <numeric> (and possibly <functional>). This will fail to compile on toolchains that don’t get these transitively. Add the required headers at the top of quantizer.cpp.

---

transformer_engine/pytorch/csrc/quantizer.cpp
Shape check can reject valid tensors

compressShapeTo2D(expected_shape) only flattens the rowwise shape before comparing to shape derived from columnwise_data. If convert_shape_back_from_fp4(getTensorShape(*columnwise_data), true) ever returns a non-2D shape (or a different 2D flattening), this check will still fail even when both buffers represent the same logical tensor. Consider normalizing both sides to the same 2D convention (e.g., apply the same compression to shape as well) before comparing.

[ptrendx]

@skydoorkai I think the comments from Greptile make sense (even though they are not very high priority). Could you fix those? Other than that LGTM.

[skydoorkai]

@skydoorkai I think the comments from Greptile make sense (even though they are not very high priority). Could you fix those? Other than that LGTM.

Updated according to Greptile comments to add headers and compare 2D shapes.

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[ptrendx]

/te-ci pytorch

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{1 file reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[skydoorkai]

@ptrendx can you review again after the modification?