Skip to content

Fuse GQA local window as kv-cache attention #4617

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
turneram wants to merge 5 commits into
base: develop
Choose a base branch
from
+14 −6
Open

Fuse GQA local window as kv-cache attention #4617

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f5c4491
Update matcher
turneram Feb 12, 2026
c50b1a3
Merge remote-tracking branch 'origin/develop' into gqa-local-window
turneram Feb 12, 2026
6e581b4
Rocmlir commit
turneram Feb 12, 2026
f5d1fb7
Formatting
turneram Feb 17, 2026
ee45340
Changelog
turneram Feb 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ Full documentation for MIGraphX is available at
* Fixed an issue with `reshape_lazy`'s shape computation that was leading to invalid reshapes (#4594).

### Optimized
* Added optimized fusion for local_window mode of GQA operator
Copy link

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

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

The phrasing "Added optimized fusion" is inconsistent with other entries in the Optimized section. Based on the pattern used in previous releases (e.g., lines 77-81), entries in the Optimized section should start with a past tense verb describing what was optimized rather than "Added". Consider rephrasing to something like "Optimized fusion for local_window mode of GQA operator" or "Improved fusion to support local_window mode of GQA operator".

Suggested change
* Added optimized fusion for local_window mode of GQA operator
* Optimized fusion for local_window mode of GQA operator.

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

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

The changelog entry is missing a period at the end. All other entries in the Optimized section end with a period for consistency.

Suggested change
* Added optimized fusion for local_window mode of GQA operator
* Added optimized fusion for local_window mode of GQA operator.

Copilot uses AI. Check for mistakes.

### Removed

Expand Down
2 changes: 1 addition & 1 deletion requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,4 +29,4 @@ pybind/pybind11@3e9dfa2866941655c56877882565e7577de6fc7b --build
msgpack/msgpack-c@cpp-3.3.0 -DMSGPACK_BUILD_TESTS=Off -DMSGPACK_BUILD_EXAMPLES=Off -DCMAKE_POLICY_VERSION_MINIMUM=3.5
sqlite3@3.50.4 -DCMAKE_POSITION_INDEPENDENT_CODE=On
ROCm/composable_kernel@ad0db05b040bacda751c65c705261b8a0a7ed25d --cmake subdir -DCMAKE_DIR=codegen -DCMAKE_POSITION_INDEPENDENT_CODE=On -DBUILD_TESTING=Off
ROCm/rocMLIR@9f843291bab85785796985b60f9a2840e7aad302 -DBUILD_FAT_LIBROCKCOMPILER=On -DLLVM_INCLUDE_TESTS=Off
ROCm/rocMLIR@7b4aa8860e1002f98a2fbde411bf1a6289763d6a -DBUILD_FAT_LIBROCKCOMPILER=On -DLLVM_INCLUDE_TESTS=Off
Copy link
Member

@bdevorem bdevorem Feb 18, 2026

Choose a reason for hiding this comment

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

I think Justin's PR with the local window fix needs to rebase; the MLIR debug failures you're seeing are solved in rocMLIR#2237, which was merged yesterday

17 changes: 12 additions & 5 deletions src/fuse_attention.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -762,7 +762,7 @@
auto matcher() const
{
static const std::unordered_set<std::string> skip_set = {
"multibroadcast", "reshape", "unsqueeze"};
"multibroadcast", "broadcast", "reshape", "unsqueeze", "squeeze"};

auto keys =
match::skip(match::name(skip_set))(match::name("concat_past_present")).bind("pres_k");
Expand All @@ -775,12 +775,18 @@
auto attn_scores = match::any_of(scale, gemm1);
auto causal_mask =
match::name("where")(match::arg(0)(broadcasted_const), match::arg(2)(attn_scores));
auto local_window_comp = match::skip(match::name(skip_set))(
match::name("convert")(match::arg(0)(match::name("greater"))));

Check warning on line 779 in src/fuse_attention.cpp

View workflow job for this annotation

GitHub Actions / cppcheck 

style: Too many nested parentheses can affect readability; consider using variables instead. [migraphx-MatcherNestedParentheses]
auto local_window_mask =
match::name("where")(match::arg(0)(match::any_of(local_window_comp, broadcasted_const)),
match::arg(2)(match::any_of(causal_mask, scale, gemm1)));
auto greater = match::name("greater")(match::arg(1)(match::any().bind("total_sl")));
auto conv_greater =
match::skip(match::name("unsqueeze"))(match::name("convert")(match::arg(0)(greater)));
auto bc_greater = match::name("multibroadcast")(match::arg(0)(conv_greater));
auto mask = match::name("where")(match::arg(0)(bc_greater),
match::arg(2)(match::any_of(causal_mask, scale, gemm1)));
auto bc_greater = match::name("multibroadcast")(match::arg(0)(conv_greater));
auto mask = match::name("where")(
match::arg(0)(bc_greater),
match::arg(2)(match::any_of(local_window_mask, causal_mask, scale, gemm1)));
auto attn_probabilities = match::skip(match::name("convert"))(
match::softmax_input(match::skip(match::name("convert"))(mask)));
auto values =
Expand Down Expand Up @@ -825,7 +831,8 @@
"broadcast",
"multibroadcast",
"@literal",
"unsqueeze"};
"unsqueeze",
"squeeze"};

auto is_valid_attn_op = [&](auto i) {
return i->get_operator().attributes().get("pointwise", false) or
Expand Down
Loading