Repository navigation
[BugFix] int64-widen MHC kernel offsets to prevent int32 overflow on large N - #1600
Open
selimfedakar wants to merge 1 commit into
Open
selimfedakar wants to merge 1 commit into
selimfedakar wants to merge 1 commit into
Conversation
The MHC Triton kernels (_mhc_mm_norm_fwd/bwd, _mhc_pre_fwd/bwd, _mhc_post_res_fwd/bwd) indexed the [N, K]/[N, M]/[N, C] activation tensors with an int32 row program id (pid_n = tl.program_id(0)). For N * K (etc.) > 2**31 the offset n_offs * stride overflows INT_MAX and wraps negative, causing an illegal memory access on large-N inputs (reproduced: mhc_mm_norm_fwd at N=600000, K=4096 -> CUDA illegal memory access). Widen pid_n to int64 in all six kernels, matching the existing int64 fixes in RMSNorm (linkedin#804/linkedin#1482), rope (linkedin#1342) and softmax (linkedin#1497). Numerics unchanged; existing MHC tests pass. Adds a regression test (N*K > 2**31, 20 GiB guard). Signed-off-by: Ahmet Selim Fedakar <a.selimfedakar@gmail.com>
This branch has not been deployed
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
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.
Summary
The MHC Triton kernels indexed the
[N, K]/[N, M]/[N, C]activationtensors with an int32 row program id (
pid_n = tl.program_id(0)), so offsetsn_offs * stride_*are int32. ForN * K(resp.N * M,N * C)> 2**31theoffset overflows
INT_MAX, wraps negative, and the kernel accesses memory out ofbounds — an illegal memory access on large-N inputs.
Reproduced on an NVIDIA RTX PRO 6000 Blackwell (sm_120):
mhc_mm_norm_fwdwithN=600000, K=4096(N*K = 2,457,600,000 > 2**31) →CUDA error: an illegal memory access was encountered.This widens the row program id to int64 in all six N-indexing kernels
(
_mhc_mm_norm_fwd,_mhc_mm_norm_bwd_fused,_mhc_pre_fwd,_mhc_pre_bwd,_mhc_post_res_fwd,_mhc_post_res_bwd), matching the existing int64 fixes inRMSNorm (#804, #1482), rope (#1342) and softmax (#1497). House idiom
tl.program_id(0).to(tl.int64). Numerics unchanged for in-range shapes.Testing Done
Hardware Type: NVIDIA RTX PRO 6000 Blackwell (sm_120), torch 2.8.0+cu128,
triton 3.4.0.
Before:
mhc_mm_norm_fwd(N=600000, K=4096)→ CUDA illegal memory access.After: forward matches a token-chunked reference (max abs diff 2.5e-3).
Added regression test
test_mhc_mm_norm_int32_offset_overflow(
N*K > 2**31, guarded by a 20 GiB free-memory skip).pytest test/transformers/test_mhc.py: 33 passed before and after the fix(no regression) + the new overflow test passes.
ruff check+ruff formaton the changed files: clean.run
make testto ensure correctnessrun
make checkstyleto ensure code stylerun
make test-convergence(N/A — kernel addressing fix, no numerics change)