Skip to content

Preserve active GQA mailbox contributions during initialization - #3601

Open
1sgtpepper wants to merge 3 commits into
NVIDIA:mainfrom
1sgtpepper:fix/gqa-mailbox-initialization
Open

1sgtpepper wants to merge 3 commits into
NVIDIA:mainfrom
1sgtpepper:fix/gqa-mailbox-initialization

Conversation

@1sgtpepper

@1sgtpepper 1sgtpepper commented Sep 8, 2026 •

Copy link
Copy Markdown

Fixes #3598.

Initialize only inactive split mailboxes in EPILOG_warp. Active producers may already have written their contributions when a receiving CTA reaches initialization; clearing those entries can erase valid data. Existing transaction waits order producer completion before reduction. Inactive entries retain their required −infinity/zero/zero values.

The shared epilog serves both contiguous and paged GQA.

Tests

Build 93_blackwell_low_latency_gqa_mailboxes with CUTLASS_NVCC_ARCHS=100a, then run:

compute-sanitizer --tool racecheck --error-exitcode 3 ./build/examples/93_blackwell_low_latency_gqa/93_blackwell_low_latency_gqa_mailboxes

On B200/CUDA 13.1.1, all-active, one-active and partial-tail cases pass numerically with zero racecheck hazards. Restoring the original header produces 54 wrong values; counts vary with scheduling. CMake build.

Each active split processes one tile, isolating this regression from the scratch-reuse race in #3600. With both fixes, the contiguous and paged example checks pass. Another 16 contiguous regression cases pass under racecheck and memcheck.

Signed-off-by: 1sgtpepper <cynejarviszarceno@gmail.com>
Signed-off-by: 1sgtpepper <cynejarviszarceno@gmail.com>
Signed-off-by: 1sgtpepper <cynejarviszarceno@gmail.com>
@github-actions

Copy link
Copy Markdown

This PR has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this PR if it is no longer required. Otherwise, please respond with a comment indicating any updates. This PR will be labeled inactive-90d if there is no activity in the next 60 days.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Example 93 GQA initialization overwrites active remote contributions

1 participant