Fix overflow of padding/unpadding kernel

[adamantboy]

Description

This PR fixes a NaN issue in the fp8 padding/unpadding kernel when handling a huge shape input in Fp8Unpadding/Fp8Padding forward. Since the data type of row and row_length is int, if the value of row*row_length is bigger than max value of int, these overlow can result in an illegal memory access on cuda or forward loss/grad norm NaN
eventually.e.g:
RuntimeError: Rank 62, device 6, iteration -1: Unexpected result nan (message='found NaN in local forward loss calculation').
CUDA Error: an illegal memory access was encountered.

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:

• Modify multi_padding_kernel/multi_unpadding_kernel, convert the value of row*row_length to the size_t type

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

Fixed critical integer overflow bug in FP8 padding/unpadding CUDA kernels that caused illegal memory access and NaN values when processing large tensors.

• Changed row * row_length calculation from int to size_t by introducing row_offset variable with explicit cast
• Applied fix consistently to both multi_padding_kernel and multi_unpadding_kernel functions
• Fixed 6 instances of potential overflow: 3 in padding kernel (2 reads, 1 write, plus padding writes) and 2 in unpadding kernel (1 read, 1 write)
• The fix ensures correct memory addressing when row * row_length exceeds INT_MAX (2^31-1 ≈ 2.1B)

Confidence Score: 4/5

• This PR is safe to merge with one minor consideration about comprehensive testing
• The fix correctly addresses the overflow issue using proper type casting. The change is minimal, focused, and consistently applied to all affected code paths. However, no tests were added to verify the fix works for large tensor dimensions, which would have made this a 5/5
• No files require special attention - the single file changed has a straightforward and correct fix

Important Files Changed

Filename	Overview
transformer_engine/common/util/padding.cu	Fixed integer overflow by casting row * row_length calculation to size_t before array indexing in both padding and unpadding kernels

Sequence Diagram

sequenceDiagram
    participant CPU as CPU Host
    participant Kernel as CUDA Kernel
    participant Memory as GPU Memory
    
    CPU->>CPU: Calculate tensor dimensions
    Note over CPU: num_rows, row_length (int)
    CPU->>Kernel: Launch multi_padding_kernel/multi_unpadding_kernel
    
    Kernel->>Kernel: Calculate tile position
    Note over Kernel: tile_row, tile_col
    
    loop For each iteration
        loop For each vector element
            Kernel->>Kernel: Calculate row index (int)
            Kernel->>Kernel: Cast to size_t: row_offset = (size_t)row * row_length
            Note over Kernel: Prevents int overflow for large tensors
            Kernel->>Memory: Read input[row_offset + col + j2]
            Memory-->>Kernel: Return data
            Kernel->>Kernel: Process data
            Kernel->>Memory: Write output[row_offset + col + j2]
        end
    end
    
    Kernel-->>CPU: Complete execution

Loading

[greptile-apps]

Greptile found no issues!

From now on, if a review finishes and we haven't found any issues, we will not post anything, but you can confirm that we reviewed your changes in the status check section.

_{This feature can be toggled off in your Code Review Settings by deselecting "Create a status check for each PR".}

[BestJuly]

Please sign-off the commit to pass the DCO check. Thanks.

[BestJuly]

/te-ci pytorch

[ksivaman]

/te-ci L0 L1

[ksivaman]

LGTM