[ROCm] AITER fused RoPE+KVCache by Rohan138 · Pull Request #33443 · vllm-project/vllm

Rohan138 · 2026-01-30T19:27:35Z

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.

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

gemini-code-assist

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.

vllm/_aiter_ops.py

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

vllm/config/compilation.py

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>

… into fused_aiter_rope_kvcache

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

ProExpertProg

Great work, thanks for addressing all of the comments!!

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>

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> Signed-off-by: Andrii Skliar <askliar@nvidia.com>

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>

Add fused qkrope+kvcache kernel in aiter registry

75e5fca

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

mergify bot added the rocm Related to AMD ROCm label Jan 30, 2026

github-project-automation bot added this to AMD Jan 30, 2026

github-project-automation bot moved this to Todo in AMD Jan 30, 2026

gemini-code-assist bot reviewed Jan 30, 2026

View reviewed changes

vllm/_aiter_ops.py Outdated Show resolved Hide resolved

vllm/_aiter_ops.py Outdated Show resolved Hide resolved

Rohan138 added 4 commits January 30, 2026 14:31

Add compilation/config necessities

e9a9d90

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Add initial gpt-oss unit test

f711495

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

wip

ee525ce

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

wip

1abaa17

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Rohan138 force-pushed the fused_aiter_rope_kvcache branch from 7a8ecc2 to 1abaa17 Compare January 31, 2026 00:22

remove split from pass

74eba5d

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

tjtanaa reviewed Feb 2, 2026

View reviewed changes

vllm/config/compilation.py Outdated Show resolved Hide resolved

Rohan138 mentioned this pull request Feb 3, 2026

[Bug]: RotaryEmbedding CustomOp does not work with gpt-oss #33663

Closed

1 task

Rohan138 force-pushed the fused_aiter_rope_kvcache branch from 5a202fa to 74eba5d Compare February 3, 2026 07:33

Rohan138 mentioned this pull request Feb 3, 2026

[Bug]: RotaryEmbedding+KVCache ops unable to pattern match for ROCmAiterTritonRopeReshapeKVCacheFusionPass #33666

Closed

1 task

Rohan138 added 7 commits February 3, 2026 16:38

matching works!

6125fad

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

pass works, need to add kv cache context init to the unit test

b4e39b3

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

take conditional out of customop

e6b9793

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Merge branch 'main' into fused_aiter_rope_kvcache

74384e8

wip

68ad95d

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

rename to head_size

821cc93

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

UNIT TEST PASSES

12a4e14

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

mergify bot added the v1 label Feb 5, 2026

Rohan138 added 7 commits February 5, 2026 12:03

fix spelling

6bafc79

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Merge branch 'main' into fused_aiter_rope_kvcache

7d49ab2

fetch layers and match by layer name manually

8fe5b3b

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

fetch layers and match by layer name manually

149096c

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

check slot_mapping in customop

1f6f9ac

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

move k v and kvcache sharing checks out of custom op

d4f05e8

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

drop check in selfattention, add check in cross and whisper

aab9a41

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Rohan138 and others added 6 commits February 20, 2026 16:40

Add GH issue to compare func outputs

8bc61c7

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

drop tolerance from scatter_split_replace UT

3c40e19

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Update vllm/compilation/passes/fusion/rope_kvcache_fusion.py

1b09335

Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>

Merge branch 'fused_aiter_rope_kvcache' of https://github.com/ROCm/vllm…

75283e4

… into fused_aiter_rope_kvcache

do fp8 kv cache in layer

9fe52cc

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Add example graph to ScatterSplitReplacementPass

ce6fde3

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

ProExpertProg approved these changes Feb 21, 2026

View reviewed changes

github-project-automation bot moved this from To Triage to Ready in gpt-oss Issues & Enhancements Feb 21, 2026

ProExpertProg enabled auto-merge (squash) February 21, 2026 01:35

ProExpertProg added the torch.compile label Feb 21, 2026

github-project-automation bot added this to torch.compile integration Feb 21, 2026

github-project-automation bot moved this to To triage in torch.compile integration Feb 21, 2026

ProExpertProg moved this from To triage to In review in torch.compile integration Feb 21, 2026

Merge branch 'main' into fused_aiter_rope_kvcache

2d2cc34

gshtras added a commit to ROCm/vllm that referenced this pull request Feb 23, 2026

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

Merge ROCm:fused_aiter_rope_kvcache (vllm-project#33443)

ac2e745

vllm-bot merged commit 2ff4e51 into vllm-project:main Feb 24, 2026
66 of 68 checks passed

github-project-automation bot moved this from Todo to Done in AMD Feb 24, 2026

github-project-automation bot moved this from Ready to Done in gpt-oss Issues & Enhancements Feb 24, 2026

github-project-automation bot moved this from In review to Done in torch.compile integration Feb 24, 2026

This was referenced Feb 24, 2026

[ROCm]: Enable customop and rope+kvcache fusion for AITER RoPE #35180

Merged

[ROCm][WIP]: Fused aiter rope kvcache mla #35245

Draft

Rohan138 mentioned this pull request Mar 5, 2026

vllm 0.17.0 gpt-oss updates SemiAnalysisAI/InferenceX#867

Closed

ProExpertProg mentioned this pull request Mar 6, 2026

[Hardware][AMD] Add fused QK RoPE and reshape & cache flash support for ROCm #28850

Closed

5 tasks

cquil11 mentioned this pull request Mar 7, 2026

[AMD] GPT-OSS vLLM 0.17.0 AMD update SemiAnalysisAI/InferenceX#889

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[ROCm] AITER fused RoPE+KVCache#33443

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

Rohan138 commented Jan 30, 2026 •

edited by github-actions bot

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ProExpertProg left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Uh oh!

Conversation

Rohan138 commented Jan 30, 2026 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ProExpertProg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Rohan138 commented Jan 30, 2026 •

edited by github-actions bot

Loading