Skip to content

[ROCm] AITER fused RoPE+KVCache#33443

Merged
vllm-bot merged 92 commits into
vllm-project:mainfrom
ROCm:fused_aiter_rope_kvcache
Feb 24, 2026
Merged

[ROCm] AITER fused RoPE+KVCache#33443
vllm-bot merged 92 commits into
vllm-project:mainfrom
ROCm:fused_aiter_rope_kvcache

Conversation

@Rohan138

@Rohan138 Rohan138 commented Jan 30, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Contributor

Purpose

Follow-up to #25954; adda a fuse_rope_kvcache Inductor pass to fuse QK RoPE and KVCache ops into the AITER fused kernel.

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

@Rohan138
Add fused qkrope+kvcache kernel in aiter registry
75e5fca 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@mergify mergify Bot added the rocm Related to AMD ROCm label Jan 30, 2026
@github-project-automation github-project-automation Bot moved this to Todo in AMD Jan 30, 2026
gemini-code-assist[bot]
gemini-code-assist Bot reviewed Jan 30, 2026
View reviewed changes

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request adds a new fused RoPE+KVCache operation for ROCm using AITER. The changes look good overall, but there are a couple of issues related to the custom op registration that could affect correctness, especially when using torch.compile. I've added comments with suggestions to address these.

Comment thread vllm/_aiter_ops.py Outdated
Comment thread vllm/_aiter_ops.py Outdated
@Rohan138
Add compilation/config necessities
e9a9d90 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
Add initial gpt-oss unit test
f711495 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
wip
ee525ce 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
wip
1abaa17 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138 Rohan138 force-pushed the fused_aiter_rope_kvcache branch from 7a8ecc2 to 1abaa17 Compare January 31, 2026 00:22
@Rohan138
remove split from pass
74eba5d 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
tjtanaa
tjtanaa reviewed Feb 2, 2026
View reviewed changes
Comment thread vllm/config/compilation.py Outdated
@Rohan138 Rohan138 mentioned this pull request Feb 3, 2026
Closed
1 task
@Rohan138 Rohan138 force-pushed the fused_aiter_rope_kvcache branch from 5a202fa to 74eba5d Compare February 3, 2026 07:33
@Rohan138 Rohan138 mentioned this pull request Feb 3, 2026
Closed
1 task
@Rohan138
matching works!
6125fad 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
pass works, need to add kv cache context init to the unit test
b4e39b3 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
take conditional out of customop
e6b9793 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
Merge branch 'main' into fused_aiter_rope_kvcache
74384e8
@Rohan138
wip
68ad95d 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
rename to head_size
821cc93 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
UNIT TEST PASSES
12a4e14 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@mergify mergify Bot added the v1 label Feb 5, 2026
@Rohan138
fix spelling
6bafc79 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
Merge branch 'main' into fused_aiter_rope_kvcache
7d49ab2
@Rohan138
fetch layers and match by layer name manually
8fe5b3b 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
fetch layers and match by layer name manually
149096c 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
check slot_mapping in customop
1f6f9ac 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
move k v and kvcache sharing checks out of custom op
d4f05e8 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@Rohan138
drop check in selfattention, add check in cross and whisper
aab9a41 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
@ProExpertProg ProExpertProg moved this from To triage to In review in torch.compile integration Feb 21, 2026
gshtras added a commit to ROCm/vllm that referenced this pull request Feb 23, 2026
@gshtras
Merge ROCm:fused_aiter_rope_kvcache (vllm-project#33443)
dbab37f
gshtras added a commit to ROCm/vllm that referenced this pull request Feb 23, 2026
@gshtras
Merge ROCm:fused_aiter_rope_kvcache (vllm-project#33443)
ac2e745
@vllm-bot vllm-bot merged commit 2ff4e51 into vllm-project:main Feb 24, 2026
66 of 68 checks passed
@github-project-automation github-project-automation Bot moved this from Todo to Done in AMD Feb 24, 2026
@github-project-automation github-project-automation Bot moved this from In review to Done in torch.compile integration Feb 24, 2026
This was referenced Feb 24, 2026
Merged
Closed
@Rohan138 Rohan138 mentioned this pull request Mar 5, 2026
Closed
@ProExpertProg ProExpertProg mentioned this pull request Mar 6, 2026
Closed
5 tasks
@cquil11 cquil11 mentioned this pull request Mar 7, 2026
Merged
Copilot AI pushed a commit to machov/vllm that referenced this pull request Mar 10, 2026
@Rohan138 @charlifu
@ProExpertProg @dllehr-amd
[ROCm] AITER fused RoPE+KVCache (vllm-project#33443)
d129bc6 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: charlifu <charlifu@amd.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: charlifu <charlifu@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
jiangkuaixue123 pushed a commit to jiangkuaixue123/vllm that referenced this pull request Apr 28, 2026
@Rohan138 @charlifu
@ProExpertProg @dllehr-amd @jiangkuaixue123
[ROCm] AITER fused RoPE+KVCache (vllm-project#33443)
19cd2f9 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: charlifu <charlifu@amd.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: charlifu <charlifu@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
mystous pushed a commit to mystous/vllm_hybrid that referenced this pull request May 10, 2026
@Rohan138 @charlifu
@ProExpertProg @dllehr-amd
[ROCm] AITER fused RoPE+KVCache (vllm-project#33443)
d9bae0c 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: charlifu <charlifu@amd.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: charlifu <charlifu@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026
@gshtras
Merge ROCm:fused_aiter_rope_kvcache (vllm-project#33443)
746bdc0
my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026
@Rohan138 @charlifu
@ProExpertProg @dllehr-amd
[ROCm] AITER fused RoPE+KVCache (vllm-project#33443)
e8fc270 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: charlifu <charlifu@amd.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: charlifu <charlifu@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026
@gshtras
Merge ROCm:fused_aiter_rope_kvcache (vllm-project#33443)
614cba0
my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026
@Rohan138 @charlifu
@ProExpertProg @dllehr-amd
[ROCm] AITER fused RoPE+KVCache (vllm-project#33443)
e7c0263 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: charlifu <charlifu@amd.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: charlifu <charlifu@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
0826joyce pushed a commit to 0826joyce/vllm-serving-optimization that referenced this pull request May 19, 2026
@Rohan138 @charlifu
@ProExpertProg @dllehr-amd
[ROCm] AITER fused RoPE+KVCache (vllm-project#33443)
67c59bb 
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: charlifu <charlifu@amd.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: charlifu <charlifu@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
@Rohan138 Rohan138 deleted the fused_aiter_rope_kvcache branch May 31, 2026 16:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ProExpertProg ProExpertProg ProExpertProg approved these changes

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@youkaichao youkaichao Awaiting requested review from youkaichao youkaichao is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@mgoin mgoin Awaiting requested review from mgoin mgoin is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

@houseroad houseroad Awaiting requested review from houseroad houseroad is a code owner

@hmellor hmellor Awaiting requested review from hmellor

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

@zou3519 zou3519 Awaiting requested review from zou3519 zou3519 is a code owner

@LucasWilkinson LucasWilkinson Awaiting requested review from LucasWilkinson LucasWilkinson is a code owner

@tjtanaa tjtanaa Awaiting requested review from tjtanaa tjtanaa is a code owner

@gshtras gshtras Awaiting requested review from gshtras

@tdoublep tdoublep Awaiting requested review from tdoublep tdoublep is a code owner

@zhuohan123 zhuohan123 Awaiting requested review from zhuohan123 zhuohan123 is a code owner

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat alexm-redhat is a code owner

@njhill njhill Awaiting requested review from njhill njhill is a code owner

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

gpt-oss Related to GPT-OSS models ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm torch.compile v1

Projects

AMD
Status: Done
gpt-oss Issues & Enhancements
Status: Done
torch.compile integration
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

7 participants

@Rohan138 @ProExpertProg @tjtanaa @gshtras @vllm-bot @charlifu @dllehr-amd