Fix exclusive upper bounds in block mask iterators - #20
Open
ACEEE-1222 wants to merge 1 commit into
Open
Conversation
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
Treat
max_block_idxas an exclusive upper bound in all block-mask iteratormask_valmethods.The iterators store the number of valid entries in
max_block_idx, so validindices are in
[0, max_block_idx). The previousindex > max_block_idxchecks allowed
index == max_block_idxto proceed.Why this matters
The forward kernel advances the iterator before testing whether it reached the
end:
mask_block_idx++; mask_val = blockmask.mask_val(mask_block_idx); is_last_block = mask_block_idx >= max_block_idx || mask_val == -1;Consequently,
mask_val(max_block_idx)must be safe and return-1.Depending on optimization and the surrounding allocation, the old condition
can read the next mask row/head or cross the allocation boundary.
Reproduction
Environment:
A frozen LongBench
multi_newswindow beginning at dataset offset 90deterministically reproduces the issue with
CUDA_LAUNCH_BLOCKING=1.Before this change:
block_sparse_attn_cuda.fwd_block;After this change, using the same 16-sample window:
Running offset 94 by itself succeeds with both builds. The fault therefore
depends on the preceding call sequence or allocator layout; it is not an
input-only crash. Replaying offsets 90-94 is the minimal verified reproducer.
Additional validation
0.003906259.677126945462078e-05both builds with identical checksums.
the real model path is needed to expose the faulty boundary.
Scope
The same exclusive-bound correction is applied to:
This is an iterator correctness fix rather than a Blackwell-only code path,
although the deterministic reproduction was observed on an SM120 GPU.