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

This PR updates NVFP4Quantizer::convert_and_update_tensor to handle cases where rowwise NVFP4 data retains an N-D “original” shape while columnwise NVFP4 data is stored with an enforced 2D shape. It introduces compressShapeTo2D to flatten the rowwise-derived expected shape for comparison, and then uses the rowwise-derived shape as the ground-truth shape for subsequent buffer allocations/coercions.

This fits into the PyTorch NVFP4 quantization flow by ensuring shape consistency checks and downstream buffer creation use a consistent logical shape when both rowwise and columnwise representations are present.

Confidence Score: 3/5

• Moderate merge risk due to compile/shape-check correctness issues in the new logic.
• Change is localized and addresses a real mismatch scenario, but the new helper introduces a likely missing standard header include and the shape equality check normalizes only the rowwise side, which can still trigger false mismatches depending on what convert_shape_back_from_fp4(..., true) returns for columnwise buffers.
• transformer_engine/pytorch/csrc/quantizer.cpp

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/quantizer.cpp	Adds compressShapeTo2D and changes NVFP4 row/col shape consistency check to compare columnwise shape against a 2D-flattened rowwise shape, then uses rowwise shape as ground truth; needs missing standard header includes and the one-sided normalization can still cause mismatches.

Sequence Diagram

sequenceDiagram
    participant Py as Python NVFP4Tensor
    participant Q as NVFP4Quantizer::convert_and_update_tensor
    participant R as rowwise_data buffer
    participant C as columnwise_data buffer

    Py->>Q: convert_and_update_tensor(tensor)
    Q->>Py: read attrs (_rowwise_data/_columnwise_data/...)

    alt columnwise_data exists
        Q->>C: getTensorShape(columnwise_data)
        C-->>Q: shape(fp4-col)
        Q-->>Q: convert_shape_back_from_fp4(..., true)
        alt rowwise_data also exists
            Q->>R: getTensorShape(rowwise_data)
            R-->>Q: shape(fp4-row)
            Q-->>Q: convert_shape_back_from_fp4(..., false)
            Q-->>Q: compressShapeTo2D(expected_shape)
            Q-->>Q: NVTE_CHECK(col_shape == expected_shape_2d)
            Q-->>Q: shape = expected_shape (ground truth)
        end
    else only rowwise_data exists
        Q->>R: getTensorShape(rowwise_data)
        R-->>Q: shape(fp4-row)
        Q-->>Q: convert_shape_back_from_fp4(..., false)
    end

    Q-->>Py: allocate missing buffers using inferred shape
    Q-->>Py: return updated TensorWrapper + py::object

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.