fix(dsa): align packed-CP indexer causal masks - #14
Closed
JackRao123 wants to merge 3 commits into
Closed
Conversation
…-k paths
The cuDNN indexer_forward_wrapper applies TOP-LEFT-aligned causal masking by
default (row i keeps keys j <= i), but both packed-CP indexer top-k paths hand
it query chunks whose rows sit at ABSOLUTE causal positions inside a key
prefix cropped to key_end:
- _indexer_topk_from_score_chunks (single packed THD sequence, CP front/back
segments): each row chunk's q[0] sits at global key index
bottom_right_key_start + row_start;
- _indexer_topk_multi_packed_cp_thd (multi-document packed THD): each THD
segment's q[0] sits at doc-local key index
segment_k_lengths - segment_q_lengths.
Without the offsets, every zigzag chunk except cp_rank 0's front chunk is
masked to a chunk-local window (~seq/2cp keys) instead of its true causal
prefix, so the downstream top-k silently selects from the wrong keys
(measured 2-25% overlap vs an exact fp32 torch reference at cp_size=32,
GLM-5.2 indexer dims), and at 131k tokens the mismatched -inf pattern
surfaces as cudaErrorIllegalAddress inside indexer_top_k. Training does NOT
crash at short sequence lengths - it just learns on a wrong sparse-attention
pattern.
Passing the kernel's q_causal_offsets argument ("global uncompressed token
index for each batch/THD segment's local q[0]", cudnn 1.25.0) at both call
sites makes index parity exact (overlap 1.0000, zero causally-out-of-bounds
indices) vs the torch reference for single-doc rows at 8k/32k/131k across
cp ranks 0/15/31 and multi-doc packs [8192,4096]/[65536,65536]/[131008,64],
and eliminates the 131k IMA (reproduced standalone at docs=[65536,65536],
cp_rank=31 before the fix).
Note test_cudnn_indexer_topk_single_packed_cp_real_kernel_uses_bottom_right_alignment
(the one test that runs the real kernel on this path) is currently disabled as
flaky (cutlass ThrMma build issue); the remaining tests mock the kernel and
mask this defect.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 9, 2026
pstefa1707
approved these changes
Jul 10, 2026
Keep the per-path invariants close to the code without duplicating the failure-mode explanation at both call sites. Signed-off-by: Jack Rao <jack.rao@baseten.co>
Keep the packed-CP rationale in the LM#14 description instead of duplicating it beside both call sites. Signed-off-by: Jack Rao <jack.rao@baseten.co>
JackRao123
changed the title
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.
Summary
Pass
q_causal_offsetsto the cuDNN DSA indexer when it selects top-k keys for packed context-parallel queries:These query chunks begin at an absolute causal position within their cropped key prefix. Without the offsets, cuDNN applies a chunk-local causal mask, producing incorrect top-k indices for most CP chunks and potentially surfacing as an illegal memory access at long context.
Validation
Integration target and merge order
The deleted
trainers-main-20260907target has been replaced bytrainers-mainat the identical base commit (038760cd), so this branch needs no rebase.This PR and Megatron-Bridge#17 can be reviewed in parallel, but this PR must merge first. If it is squash- or rebase-merged, Bridge#17 will repin its Megatron-LM submodule to the landed
trainers-mainSHA before merging. Trainers will then advance its Bridge gitlink.Follow-up
The real-kernel packed-CP test remains disabled because of an unrelated flaky CUTLASS
ThrMmabuild issue. Re-enable it, and add a multi-document analogue, once that build issue is resolved.