Fix CI failure related to bug in MXFP8 copy implementation

[vthumbe1503]

Description

This PR introduced a bug during the last few commits where in, implementation of copy_ method in torch dispatch accessed mxfp8 attributes without checking if the both source and estination tensors are indeed mxfp8.

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:

• Fix MXFP8 tensor's copy method.

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

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

Greptile Overview

Greptile Summary

Fixed an AttributeError introduced in the FSDP2 support implementation by reordering conditional checks in the copy_ operation handler within MXFP8Tensor.__torch_dispatch__.

Key Changes:

• Moved rowwise_matches and columnwise_matches evaluations inside the type check block
• Previously, attribute access on _rowwise_data and _columnwise_data occurred before verifying both src and dst were MXFP8Tensor instances
• This caused AttributeError when the copy_ operation was invoked with non-MXFP8Tensor objects
• The fix ensures attributes are only accessed after type validation, preventing runtime errors

The change maintains the same logical flow but correctly guards attribute access with type checks, allowing the code to fall back to base class behavior when tensors don't match expected types.

Confidence Score: 5/5

• This PR is safe to merge with no risk
• The fix addresses a clear bug (AttributeError) by properly ordering type checks before attribute access. The logic is identical to the original implementation, just with corrected guard placement. Pre-commit hooks have validated formatting.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
transformer_engine/pytorch/tensor/mxfp8_tensor.py	5/5	Fixed AttributeError in copy_ operation by moving attribute access checks inside type validation block

Sequence Diagram

sequenceDiagram
participant Caller
participant MXFP8Tensor
participant torch_dispatch
Caller->>torch_dispatch: copy_(dst, src)
torch_dispatch->>MXFP8Tensor: __torch_dispatch__(copy_.default, ...)
alt Both src and dst are MXFP8Tensor
MXFP8Tensor->>MXFP8Tensor: Check isinstance(src, MXFP8Tensor)
MXFP8Tensor->>MXFP8Tensor: Check isinstance(dst, MXFP8Tensor)
Note over MXFP8Tensor: Type checks pass, safe to access attributes
MXFP8Tensor->>MXFP8Tensor: Check rowwise_matches
MXFP8Tensor->>MXFP8Tensor: Check columnwise_matches
alt Layout matches
MXFP8Tensor->>MXFP8Tensor: Copy rowwise data if present
MXFP8Tensor->>MXFP8Tensor: Copy columnwise data if present
MXFP8Tensor-->>Caller: Return dst
else Layout mismatch
MXFP8Tensor->>torch_dispatch: super().__torch_dispatch__()
torch_dispatch-->>Caller: Default behavior
end
else Type check fails
MXFP8Tensor->>torch_dispatch: super().__torch_dispatch__()
torch_dispatch-->>Caller: Default behavior
end

Loading

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[ksivaman]

LGTM