Skip to content

gqa/fuse_attention: attention sinks + sliding-window decode mask for GPT-OSS - #5260

Open
rlegithub wants to merge 1 commit into
ROCm:developfrom
rlegithub:pr/gptoss-gqa-sink
Open

rlegithub wants to merge 1 commit into
ROCm:developfrom
rlegithub:pr/gptoss-gqa-sink

Conversation

@rlegithub

Copy link
Copy Markdown

Summary

Enables the GPT-OSS attention-sink path in decode-attention fusion. GPT-OSS inserts concat(scores, sink_col) -> softmax -> slice before the V GEMM; the previous find_kv_cache_attention matcher only handled softmax_input(mask), so the sink variant didn't fuse and ran as unfused static-shape ops that truncate to T=1 at decode.

  • fuse_attention.cpp: match slice(softmax(concat(any(), sink_col))); add "concat" to valid_attn_ops
  • parse_group_query_attention.cpp: sliding-window decode-mask off-by-one + sink-slice ends=-1 (dynamic) fixes
  • rewrite_reduce.cpp: tuple guard

3 of 3 PRs enabling GPT-OSS-20B (with the fused MoE op and SLN-FP32 PRs).

Test plan

  • CI
  • Verified on gfx1151: decode attention fuses for GPT-OSS; coherent decode at ~56 tok/s. Existing (Llama-style) kv-cache attention path unaffected.

Extends find_kv_cache_attention to match the GPT-OSS attention-sink path
(slice(softmax(concat(scores, sink_col)))) so decode attention fuses (else it
runs unfused static-shape ops that truncate to T=1 at decode). Adds 'concat' to
valid_attn_ops, plus the GQA sliding-window decode-mask off-by-one + sink-slice
(ends=-1 dynamic) fixes in parse_group_query_attention, and a rewrite_reduce
tuple guard.

Co-Authored-By: Claude <noreply@anthropic.com>
@rlegithub
rlegithub marked this pull request as ready for review September 14, 2026 21:21
@rlegithub
rlegithub requested a review from causten as a code owner September 14, 2026 21:21
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