Skip to content

fuse_attention: exclude side-input reductions from the decode (kv_cache) attention group - #5255

Open
rlegithub wants to merge 1 commit into
ROCm:developfrom
rlegithub:pr/fuse-attention-exclude-side-input-reductions
Open

rlegithub wants to merge 1 commit into
ROCm:developfrom
rlegithub:pr/fuse-attention-exclude-side-input-reductions

Conversation

@rlegithub

Copy link
Copy Markdown

Summary

find_kv_cache_attention::get_attn_instructions walks attention inputs backward and pulls valid side-input ops into the attention submodule. Because reduce_* ops are valid attention ops (softmax), it also pulled a preceding RMSNorm's reduce_sum into the group. That reduction is an intermediate feeding the attention GEMM, and rocMLIR's TosaToRock reduction lowering can only lower a reduce whose output traces to a kernel result → tosa.reduce_sum can't trace the reduction output to a kernel resultget_tuning_config throws → "No valid tuned compilation".

Fix: skip reduce_* ops in the backward side-input expand. The attention's own softmax reduces are on the direct start→end path (find_instructions_between) and are unaffected, so the attention stays fused; only the RMSNorm reduce is left as its own kernel. A reduce feeding a GEMM can never trace to a kernel result, so excluding it cannot regress a working fusion.

This restores fused (flash) decode attention for models with a pre-attention RMSNorm on gfx1151.

Note: this is a pragmatic fusion-boundary fix. The more complete solution — teaching rocMLIR to fuse an intermediate reduction into the GEMM prologue (norm + attention as one kernel) — is a larger rocMLIR change tracked separately.

Test plan

  • CI
  • Verified on gfx1151: int4 LLM (Llama-4-family) decodes coherently with fused attention; Llama-3.2-1B int4 regression passes.

…roup (middle-way; option-a future)

MIDDLE-WAY SOLUTION. Enables FUSED (flash) decode attention for the Muse
Glimmer int4 path on gfx1151 by excluding reduce_* ops from the backward
side-input `expand` in find_kv_cache_attention::get_attn_instructions.

Root cause: the grouping pulled the preceding RMSNorm's reduce_sum into the
attention submodule. rocMLIR (TosaToRock traceToRes) can only lower a tosa
reduce whose output traces to a kernel result; a reduce feeding the attention
gemm is an intermediate and cannot trace -> get_tuning_config throws ->
"No valid tuned compilation". The attention's OWN softmax reduces are on the
direct start->end path (find_instructions_between) and are unaffected, so the
attention stays fused. A reduce feeding a gemm can never trace to a kernel
result, so excluding it cannot regress any working fusion.

This is a pragmatic, blast-radius-minimized fix at the MIGraphX fusion
boundary. The ULTIMATE SOLUTION (option A) -- teaching rocMLIR to fuse an
intermediate reduction into the gemm prologue so RMSNorm + attention become a
single kernel -- is planned as future work (rocMLIR TosaToRock reduction
lowering).

Verified on the clean build (gfx1151): Muse Glimmer int4 decodes coherently
with fused attention ("The capital of France is Paris.."), first 8 token IDs
exact-match the oracle, ~10 tok/s (fused) vs ~9.7 (decomposed), 0 fatal, no
reduce-trace throw. Llama-3.2-1B-int4 regression passes (coherent, 68.9 tok/s).

Co-Authored-By: Claude <noreply@anthropic.com>
@pfultz2

pfultz2 commented Sep 13, 2026

Copy link
Copy Markdown
Collaborator

Unit tests need to be added for this.

@rlegithub
rlegithub marked this pull request as ready for review September 14, 2026 19:26
@rlegithub
rlegithub requested a review from causten as a code owner September 14, 2026 19:26

@pfultz2 pfultz2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unit test needs to be added and the comments cleaned up.

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