Conversation
…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>
Collaborator
|
Unit tests need to be added for this. |
rlegithub
marked this pull request as ready for review
September 14, 2026 19:26
pfultz2
requested changes
Sep 14, 2026
pfultz2
left a comment
Collaborator
There was a problem hiding this comment.
Unit test needs to be added and the comments cleaned up.
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
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
find_kv_cache_attention::get_attn_instructionswalks attention inputs backward and pulls valid side-input ops into the attention submodule. Becausereduce_*ops are valid attention ops (softmax), it also pulled a preceding RMSNorm'sreduce_suminto 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 result→get_tuning_configthrows → "No valid tuned compilation".Fix: skip
reduce_*ops in the backward side-inputexpand. The attention's own softmax reduces are on the directstart→endpath (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