Fix int32 overflow in permute_bwd_kernel row offsets - #172
Open
truong-v wants to merge 1 commit into
Open
Conversation
perm_row is loaded as int32 and multiplied by num_cols in 32-bit arithmetic; the product is widened to 64-bit only afterwards, so once perm_row * num_cols reaches 2**31 the offset wraps negative and the kernel reads far below the gradient buffer. The store address (pid_m * num_cols) wraps the same way at larger row counts. At hidden=7168, topK=8 the read wrap is reached at ~37.4k tokens. Widen the program id and the loaded row index to int64 before any scaling.
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
permute_bwd_kernelloadsperm_rowfrom the int32inv_idxtable and multiplies it bynum_colsin 32-bit arithmetic; the product is widened to 64-bit only afterwards, so onceperm_row * num_colsreaches 2^31 the offset wraps negative and the kernel reads far below the gradient buffer. The store address (pid_m * num_cols) wraps the same way atnum_topK-times-larger row counts. Withhidden = 7168andtopK = 8the read wrap is reached at ~37.4k tokens:moe_permute_topk_bwd_opcrashes with an illegal memory access (or silently returns garbage if the wrapped address is mapped).Changes:
pid_mand the loadedperm_rowtotl.int64before any scaling, so every row offset in the kernel is computed in 64-bit. This is the same spelling TransformerEngine's_unpermute_kerneluses, and the same fix class as [Bugfix][Mamba] IMA in causal_conv1d kernel for long sequences vllm-project/vllm#41617.Measured on an L40S: at
tokens=38000, topK=8, hidden=7168(bf16, 304000-row buffer, 4.4 GB) the op crashes with an illegal memory access before this PR and completes after it, with sampled output rows bit-identical to a reference computed in the kernel's accumulation order. Overhead of the 64-bit addressing, full-opdo_benchattokens=8192, topK=8, hidden=7168(no wrap, interleaved runs x3): 1.5663–1.5672 ms before, 1.5702–1.5706 ms after (~+0.2%).Tests
tests/kernels/permute/has no pytest coverage (its scripts are a standalone harness with a hard-coded device). Addedtests/kernels/test_moe_permute_bwd.py: two small exact-correctness tests formoe_permute_topk_bwd_opagainst a reference with the kernel's accumulation order, plustest_moe_permute_topk_bwd_large_offsets, which drives a >2^31-element gradient buffer (skipped on GPUs with <10 GB memory). Reverting the kernel change alone makes the large-offsets test fail with the illegal memory access while the small tests still pass; all three pass with the fix.Fixes #171
Environment
main@ 9b5b362