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 mismatch error in NVFP4 quantization when using N-dimensional tensors with both rowwise and columnwise quantization modes.

Key changes:

• Added compressShapeTo2D utility to flatten N-D shapes to 2D for comparison
• Modified shape validation in convert_and_update_tensor to compare compressed 2D shapes instead of raw shapes
• Uses rowwise shape as the ground truth since it preserves the original N-D structure
• Added comprehensive test case for 3D tensor quantization with both modes

Technical context:
NVFP4 columnwise quantization enforces 2D storage by flattening all dimensions except the last, while rowwise preserves the original N-D shape. The previous direct shape comparison failed for N-D tensors (N>2) because it compared a 3D rowwise shape like [32, 4, 128] against a 2D columnwise shape like [128, 128]. The fix compresses both to equivalent 2D representations ([128, 128]) before comparison.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The fix correctly addresses the shape mismatch issue with a well-documented helper function, includes proper test coverage for the regression, and follows the existing code patterns. The logic is sound: compressing both shapes to 2D for comparison is the right approach given that columnwise storage is always 2D while rowwise preserves N-D structure.
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/quantizer.cpp	Adds compressShapeTo2D helper to properly compare rowwise (N-D) and columnwise (2D) tensor shapes for NVFP4 quantization
tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py	Adds regression test for 3D tensor quantization with both rowwise and columnwise modes, verifying shape consistency

Sequence Diagram

sequenceDiagram
    participant User
    participant NVFP4Quantizer
    participant convert_and_update_tensor
    participant compressShapeTo2D

    User->>NVFP4Quantizer: quantize 3D tensor [32,4,128]
    NVFP4Quantizer->>NVFP4Quantizer: Create rowwise data [32,4,64]
    NVFP4Quantizer->>NVFP4Quantizer: Create columnwise data [128,64]
    
    User->>NVFP4Quantizer: update_quantized()
    NVFP4Quantizer->>convert_and_update_tensor: Validate shapes
    
    convert_and_update_tensor->>convert_and_update_tensor: convert_shape_back_from_fp4(columnwise, true)<br/>Returns: [128, 128]
    convert_and_update_tensor->>convert_and_update_tensor: convert_shape_back_from_fp4(rowwise, false)<br/>Returns: [32, 4, 128]
    
    convert_and_update_tensor->>compressShapeTo2D: Compress rowwise [32,4,128]
    compressShapeTo2D-->>convert_and_update_tensor: [128, 128]
    
    convert_and_update_tensor->>compressShapeTo2D: Compress columnwise [128,128]
    compressShapeTo2D-->>convert_and_update_tensor: [128, 128]
    
    convert_and_update_tensor->>convert_and_update_tensor: Compare [128,128] == [128,128] ✓
    convert_and_update_tensor->>convert_and_update_tensor: Use rowwise shape [32,4,128] as ground truth
    
    convert_and_update_tensor-->>User: Success

Loading

_{Last reviewed commit: 075ebf6}

[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}

[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}