Fix int32 overflows in W4A4 kernels at high resolution - #944
Open
danielwoz wants to merge 1 commit into
Open
Conversation
danielwoz
force-pushed
the
fix-int32-overflow-w4a4-large-m
branch
from
July 10, 2026 02:50
40fd085 to
552f0b8
Compare
danielwoz
changed the title
Several W4A4 code paths compute row-major offsets or size checks as int * int, which overflows once a single call's padded token count times feature width exceeds 2^31 (e.g. a video DiT FFN at 3840x3840 resolution). In debug builds the oscales size asserts abort spuriously; in release builds the device-side pointer offsets wrap and the kernels read or write out of bounds. Promote the affected products to int64: - gemm_w4a4_launch_impl.cuh: the two oscales asserts in quantize_w4a4_act_fuse_lora and the matching assert in quantize_w4a4_act - gemm_w4a4.cuh: the input offset in quantize_w4a4_fuse_lora_kernel - gemm_base.cuh: the fp16 output store offset in the GEMM epilogue - epilogues.cuh: the QKV projection epilogue output offset
danielwoz
force-pushed
the
fix-int32-overflow-w4a4-large-m
branch
from
July 10, 2026 03:02
4c1b404 to
20e7bf7
Compare
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
1 participant
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.
Several W4A4 code paths compute row-major offsets or size checks in 32-bit int. At high enough resolution (for example a video DiT running at 3840x3840, where a single FFN call exceeds 2^31 elements), the oscales size asserts overflow and abort in debug builds. In release builds the asserts are compiled out and the device-side pointer offsets wrap instead, so the kernels silently read or write out of bounds.
This change promotes the affected products to int64: the two oscales asserts in quantize_w4a4_act_fuse_lora and the matching assert in quantize_w4a4_act, the input pointer offset in quantize_w4a4_fuse_lora_kernel, the fp16 output store offset in the GEMM epilogue (gemm_base.cuh), and the QKV projection epilogue output offset (epilogues.cuh). Remaining per-tile and packed-store offsets stay far below 2^31 at these sizes and are left unchanged.