Skip to content

Fix int32 overflow in qdq_mxfp4 causing GPU memory access fault above INT32_MAX elements - #43

Open
zjin-lcf wants to merge 2 commits into
amd:release/0.12from
zjin-lcf:fix-mxfp4-int32-overflow
Open

zjin-lcf wants to merge 2 commits into
amd:release/0.12from
zjin-lcf:fix-mxfp4-int32-overflow

Conversation

@zjin-lcf

@zjin-lcf zjin-lcf commented Sep 19, 2026 •

Copy link
Copy Markdown

Problem

qdq_mxfp4 / qdq_mxfp4_ truncate numel() into an int and index elements with an int. On any tensor with more than INT32_MAX (2,147,483,647) elements the count wraps negative, still passes the divisibility check, and the kernel launches with a garbage grid dimension. The process dies with:

`launch_qdq_mxfp4` truncated `a.numel()` into an `int`, and
`qdq_mxfp4_kernel` indexed elements with an `int`. For a tensor with more
than INT32_MAX elements this wraps negative, still passes the
divisibility check, and launches with a garbage grid dimension, so the
process dies with "Memory access fault by GPU node-N ... Reason: Unknown".

Use 64-bit types for the element count and the element index, and check
the grid dimension explicitly. This mirrors what the sibling
mxfp4/dequantize.cu already does.

Minimal reproducer (bf16, hidden 14336) on MI300A / ROCm 7.2:

  2,132,795,392 elements (under INT32_MAX) -> OK
  2,162,155,520 elements (over  INT32_MAX) -> memory access fault

With this change both cases pass.
@zjin-lcf
zjin-lcf marked this pull request as draft September 19, 2026 21:02
Addresses review feedback: rather than erroring out when the block count
exceeds the maximum grid dimension, make the kernel a grid-stride loop and
clamp the grid. This removes the ceiling entirely instead of moving it.

The kernel derives each MXFP4 group's shared scale from a 32-lane butterfly
shuffle, so every active warp must be fully populated. That holds here
because the caller already requires `numel % block_size == 0`, and the
stride is a multiple of block_size, so the loop bound always lands on a
block boundary -- a block either runs an iteration with all of its threads
or not at all.

Verified on MI300A / ROCm 7.2 against the stock kernel:
  - bit-identical output on 8,388,608 elements, mixed magnitudes
  - bit-identical with the grid clamped to 1024 blocks, forcing 64
    grid-stride iterations per thread (exercises the multi-iteration path)
  - 2,162,155,520 elements (over INT32_MAX): faulted before, passes now
@zjin-lcf

zjin-lcf commented Sep 19, 2026 •

Copy link
Copy Markdown
Author

Updated per review.

1. Is numel a multiple of block_size?

Yes — guaranteed by the precondition immediately above the grid computation, which selects block_size from {128, 64} precisely so the division is exact and rejects anything else:

2. Grid-stride loop

Agreed, and switched to it. The previous revision converted one failure mode into another — it raised an error when grid_size > INT_MAX instead of removing the ceiling. qdq_mxfp4_kernel is now a grid-stride loop and the launcher clamps the grid, so the limit is gone rather than relocated. For any tensor that fits in memory today the clamp is inactive and each thread still handles exactly one element, so there is no change to the common path.

Verification (MI300A / gfx942, ROCm 7.2), differential against the stock kernel rather than a hand-written reference:

test result
8,388,608 elements, mixed magnitudes (1e-3 to 1e3) bit-identical to stock
same input, grid clamped to 1024 blocks → 64 grid-stride iterations per thread bit-identical to stock
2,132,795,392 elements (under INT32_MAX) OK before and after
2,162,155,520 elements (over INT32_MAX) faulted before, passes now

@zjin-lcf
zjin-lcf marked this pull request as ready for review September 19, 2026 21:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant