Skip to content

fix(attention): map grouped decode heads to KV groups - #177

Open
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/triton-decode-gqa-head-mapping
Open

fix(attention): map grouped decode heads to KV groups#177
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/triton-decode-gqa-head-mapping

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Aug 25, 2026

Copy link
Copy Markdown

Summary

  • map each grouped-decode head tile within its KV group before computing query-head indices
  • mask partial tail tiles at the KV-group boundary so they cannot read the next group with the previous KV head
  • launch one set of head tiles per KV head
  • cover GQA configurations (48, 2), (40, 2), and (80, 4), whose groups span multiple 16-head tiles with a partial tail

Problem

The previous stage-1 mapping partitioned the global query-head axis into consecutive tiles, then inferred the KV head from the tile index. When GROUP > 16 and GROUP % 16 != 0, a tile could cross a KV-group boundary. For example, with 48 query heads and 2 KV heads, query heads 24-31 were assigned to KV head 0 instead of KV head 1.

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.

1 participant