Fuse the rotate-half concat into attention and let the K rotary write the KV cache directly

[pfultz2]

Motivation

In RoPE models the rotary embedding needs rotate_half(x) = concat(x[h:], x[:h]). On Phi-3.5 decode this concat ran as a standalone concat_kernel for every layer (32 launches per token): fuse_attention captures the pointwise rotary math on Q into the rocMLIR attention prologue, but it stops at concat, so the half-swap stayed a copy kernel in front of the fused attention. On the K side the rotary is a pointwise kernel that fuse_concat_past_present retargets to write straight into the cache, but only when the producer's shape was identical to the concat input, which excluded both the reshaped rotary producer and the seq-major prefill layout.

Removing those copies is pure launch-count savings on a launch-bound decode path.

Technical Details

• simplify_algebra: find_concat_swapped_halves. A concat of the two equal halves of one tensor in swapped order (identical slicing on every other axis) is rewritten as where(mask, hi, lo) over the concat axis split into {2, half}, with the halves as unsqueezed broadcast views and a {true, false} mask literal, followed by a reshape. find_op_shape_transform_op then sinks the reshape past the rotary mul/add, so the attention prologue becomes slice → where → mul → add → reshape → transpose → dot and rocMLIR compiles it with no concat kernel. The K rotary becomes a single where_mul_mul_add pointwise.
• fuse_concat_past_present: invert the view chain. The producer no longer has to match the concat input shape. The view ops between the producer and concat_past_present (reshape_lazy, squeeze, unsqueeze, transpose) are inverted with shape_transform_descriptor::invert and applied to the cache slot view, which becomes the producer's output buffer and reported output_shape. This fuses the 5-D rotary producer at decode and the seq-major (non-standard) producer at prefill, so the K append copy disappears in both cases. The inversion helper is invert_view_chain, now exported from eliminate_concat.hpp and shared with eliminate_concat.
• Tests. simplify_algebra: positive single-axis and multi-axis cases plus unequal halves, different sources and mismatched other-axis slices. fuse_concat_past_present: reshaped, transposed and prefill seq-major producers, and an unmergeable reshape that must not fuse. eliminate_concat: invert_view_chain reshape, transpose+reshape, non-lazy and broadcast cases.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

• • Added: New functionality.
• • Changed: Changes to existing functionality.
• • Removed: Functionality or support that has been removed. (Compared to a previous release)
• • Optimized: Component performance that has been optimized or improved.
• • Resolved Issues: Known issues from a previous version that have been resolved.
• • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

[Copilot]

🟡 Changes recommended

The new cache-layout paths lack runtime verification, and the required optimization changelog entry is missing.

2 open findings

• Add GPU verification for stride-retargeting cache writes · New
• Add changelog entry for eliminated KV-cache copies · New

What changed in this PR

Optimizes RoPE attention by eliminating rotate-half concat copies and writing rotary K results directly into the KV cache.

Changes:

• Rewrites swapped-half concats into fusible pointwise operations.
• Inverts view chains for direct cache writes.
• Adds structural and helper tests.

Review was single-pass without agent fan-out.

File	Description
src/​simplify_algebra.cpp	Adds swapped-half concat rewrite.
src/​eliminate_concat.cpp	Extracts view-chain inversion logic.
src/​include/​migraphx/​eliminate_concat.hpp	Exports the inversion helper.
src/​targets/​gpu/​fuse_concat_past_present.cpp	Retargets viewed producers into cache slots.
test/​simplify_algebra_test.cpp	Tests concat rewrite eligibility.
test/​eliminate_concat_test.cpp	Tests view-chain inversion.
test/​gpu/​fuse_concat_past_present.cpp	Tests reshaped and transposed producer fusion.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

❌ Patch coverage is 95.71429% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/eliminate_concat.cpp	88.89%	3 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #5379      +/-   ##
===========================================
- Coverage    92.71%   92.71%   -0.01%     
===========================================
  Files          635      635              
  Lines        36203    36851     +648     
===========================================
+ Hits         33565    34163     +598     
- Misses        2638     2688      +50     

Files with missing lines	Coverage Δ
src/include/migraphx/eliminate_concat.hpp	100.00% <ø> (ø)
src/simplify_algebra.cpp	97.68% <100.00%> (+0.07%)	⬆️
src/eliminate_concat.cpp	94.00% <88.89%> (+1.95%)	⬆️

... and 9 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

🟡 Changes recommended

Cache retargeting can create invalid instruction ordering and trigger shape-propagation failures.

4 open findings

• Insert inverse views before their producer · New
• Insert separate retargeted producer for cache dependency · New
• Add GPU verification for stride-retargeting cache writes
• Add changelog entry for eliminated KV-cache copies

🧠 Review effort: Balanced

[TedThemistokleous]

shouldn't this be inputs.size() % 2 != 0 since we want to check two equal halves? Right now this assumes only two inputs. This should ensure every concat checked then has inputs that would be divisible evenly