Skip to content

Fix exclusive upper bounds in block mask iterators - #20

Open
ACEEE-1222 wants to merge 1 commit into
mit-han-lab:mainfrom
ACEEE-1222:fix/blockmask-exclusive-upper-bound
Open

Fix exclusive upper bounds in block mask iterators#20
ACEEE-1222 wants to merge 1 commit into
mit-han-lab:mainfrom
ACEEE-1222:fix/blockmask-exclusive-upper-bound

Conversation

@ACEEE-1222

Copy link
Copy Markdown

Summary

Treat max_block_idx as an exclusive upper bound in all block-mask iterator
mask_val methods.

The iterators store the number of valid entries in max_block_idx, so valid
indices are in [0, max_block_idx). The previous index > max_block_idx
checks allowed index == max_block_idx to 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:

  • NVIDIA RTX 6000D, compute capability 12.0
  • CUDA 13.1 compiler
  • PyTorch 2.11.0+cu130
  • Qwen3-4B geometry: 32 query heads, 8 KV heads, head dimension 128
  • block size 128, causal block-sparse prefill

A frozen LongBench multi_news window beginning at dataset offset 90
deterministically reproduces the issue with CUDA_LAUNCH_BLOCKING=1.

Before this change:

  • offsets 90-93 complete;
  • offset 94 fails directly in block_sparse_attn_cuda.fwd_block;
  • the process exits with a CUDA illegal-memory-access error.

After this change, using the same 16-sample window:

  • all offsets 90-105 complete;
  • the process exits successfully after 16/16 rows.

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

  • BF16 output parity against PyTorch SDPA at sequence length 512:
    • maximum absolute error: 0.00390625
    • mean absolute error: 9.677126945462078e-05
  • Forward/backward outputs and Q/K/V gradients are finite.
  • Old and patched kernels produce identical parity metrics and gradient norms.
  • 32 repeated launches spanning 8K, 16K, 32K, and 40,384 tokens complete for
    both builds with identical checksums.
  • CUDA memcheck reports no errors for the small parity case in either build;
    the real model path is needed to expose the faulty boundary.

Scope

The same exclusive-bound correction is applied to:

  • forward streaming;
  • forward exact streaming;
  • forward block-sparse;
  • backward streaming;
  • backward block-sparse iterators.

This is an iterator correctness fix rather than a Blackwell-only code path,
although the deterministic reproduction was observed on an SM120 GPU.

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.

2 participants