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 Overview

Greptile Summary

Fixes shape validation for NVFP4 tensors when both rowwise and columnwise data are present. The issue occurred because columnwise data enforces 2D shapes while rowwise data can be 3D or higher, causing the direct shape comparison to fail.

The fix:

• Adds compressShapeTo2D() helper to flatten all dimensions except the last into a single dimension
• Compares compressed 2D versions of both shapes for validation
• Uses rowwise shape as the ground truth for the final tensor shape

This properly handles cases where rowwise data has 3D+ shapes while maintaining shape consistency validation.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The fix is well-targeted and addresses a specific bug where shape validation failed for 3D+ rowwise tensors. The logic is sound: compressing both shapes to 2D for comparison ensures they're compatible while preserving the original rowwise shape for actual use. The implementation is clean and the helper function is simple and correct.
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/quantizer.cpp	Adds helper function to compress shapes to 2D and fixes shape validation for NVFP4 tensors when rowwise data is 3D+

Sequence Diagram

sequenceDiagram
    participant Caller
    participant convert_and_update_tensor
    participant convert_shape_back_from_fp4
    participant compressShapeTo2D

    Caller->>convert_and_update_tensor: NVFP4Tensor with rowwise & columnwise data
    
    alt columnwise_data exists
        convert_and_update_tensor->>convert_shape_back_from_fp4: getTensorShape(*columnwise_data), transpose=true
        convert_shape_back_from_fp4-->>convert_and_update_tensor: shape (2D enforced)
        
        alt rowwise_data also exists
            convert_and_update_tensor->>convert_shape_back_from_fp4: getTensorShape(*rowwise_data), transpose=false
            convert_shape_back_from_fp4-->>convert_and_update_tensor: expected_shape (could be 3D+)
            
            convert_and_update_tensor->>compressShapeTo2D: expected_shape
            compressShapeTo2D-->>convert_and_update_tensor: expected_shape_2d
            
            convert_and_update_tensor->>compressShapeTo2D: shape
            compressShapeTo2D-->>convert_and_update_tensor: shape_2d
            
            convert_and_update_tensor->>convert_and_update_tensor: Validate shape_2d == expected_shape_2d
            convert_and_update_tensor->>convert_and_update_tensor: Use expected_shape as ground truth
        end
    else only rowwise_data exists
        convert_and_update_tensor->>convert_shape_back_from_fp4: getTensorShape(*rowwise_data), transpose=false
        convert_shape_back_from_fp4-->>convert_and_update_tensor: shape
    end
    
    convert_and_update_tensor-->>Caller: TensorWrapper with correct shape

Loading

[greptile-apps]

_{1 file 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}