-
-
Notifications
You must be signed in to change notification settings - Fork 9.4k
[ROCm][AITER] Support AITER Rope ops in RotaryEmbedding Module. #22521
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.
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
Merged
Merged
Changes from all commits
Commits
Show all changes
14 commits
Select commit
Hold shift + click to select a range
29d44d1
integrate aiter rope module
vllmellm 8af1c71
fix device type
vllmellm 2ec0a28
add aiter rope hip forward func
vllmellm ab16de6
fix bugs in function registeration and invocation
vllmellm f778d2a
fix shape mismatch
vllmellm 607f62b
fix pre-commit issues
vllmellm 2e22977
clean up forward_hip
vllmellm 7f016e7
bugfix
vllmellm 04458da
improve code
vllmellm 54c56af
Merge remote-tracking branch 'origin/main' into aiter-rope
vllmellm 3234686
bugfix
vllmellm 9708eb2
remove unnecessary command
vllmellm 6abeb5a
fix miss-spelling function name
vllmellm 5c8101a
Merge remote-tracking branch 'origin/main' into aiter-rope
vllmellm File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
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
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
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
127 changes: 127 additions & 0 deletions
127
vllm/model_executor/layers/rotary_embedding/rocm_aiter_rope_ops.py
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,127 @@ | ||
# SPDX-License-Identifier: Apache-2.0 | ||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project | ||
|
||
from typing import Optional | ||
|
||
import torch | ||
|
||
import vllm.envs as envs | ||
from vllm.platforms import current_platform | ||
from vllm.utils import direct_register_custom_op | ||
|
||
|
||
def is_rocm_rotary_embedding_enabled() -> bool: | ||
return (current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER) | ||
|
||
|
||
def rocm_aiter_rotary_emb_without_key_forward_hip_impl( | ||
positions: torch.Tensor, | ||
sin: torch.Tensor, | ||
cos: torch.Tensor, | ||
query: torch.Tensor, | ||
offsets: Optional[torch.Tensor] = None, | ||
rotate_style: int = 0, | ||
is_nope_first: bool = False, | ||
) -> None: | ||
import aiter as ops | ||
if offsets is None: | ||
ops.rope_cached_positions_fwd_inplace( | ||
query, | ||
cos, | ||
sin, | ||
positions, | ||
rotate_style, | ||
reuse_freqs_front_part=True, | ||
nope_first=is_nope_first, | ||
) | ||
else: | ||
ops.rope_cached_positions_offsets_fwd_inplace( | ||
query, | ||
cos, | ||
sin, | ||
positions, | ||
offsets, | ||
rotate_style, | ||
reuse_freqs_front_part=True, | ||
nope_first=is_nope_first, | ||
) | ||
|
||
|
||
def rocm_aiter_rotary_emb_with_key_forward_hip_impl( | ||
positions: torch.Tensor, | ||
sin: torch.Tensor, | ||
cos: torch.Tensor, | ||
query: torch.Tensor, | ||
key: torch.Tensor, | ||
offsets: Optional[torch.Tensor] = None, | ||
rotate_style: int = 0, | ||
is_nope_first: bool = False, | ||
) -> None: | ||
import aiter as ops | ||
if offsets is None: | ||
ops.rope_cached_positions_2c_fwd_inplace( | ||
query, | ||
key, | ||
cos, | ||
sin, | ||
positions, | ||
rotate_style, | ||
reuse_freqs_front_part=True, | ||
nope_first=is_nope_first, | ||
) | ||
else: | ||
ops.rope_cached_positions_offsets_2c_fwd_inplace( | ||
query, | ||
key, | ||
cos, | ||
sin, | ||
positions, | ||
offsets, | ||
rotate_style, | ||
reuse_freqs_front_part=True, | ||
nope_first=is_nope_first, | ||
) | ||
|
||
|
||
def rocm_aiter_rotary_emb_with_key_forward_hip_fake( | ||
positions: torch.Tensor, | ||
sin: torch.Tensor, | ||
cos: torch.Tensor, | ||
query: torch.Tensor, | ||
key: torch.Tensor, | ||
offsets: Optional[torch.Tensor] = None, | ||
rotate_style: int = 0, | ||
is_nope_first: bool = False, | ||
) -> None: | ||
pass | ||
|
||
|
||
def rocm_aiter_rotary_emb_without_key_forward_hip_fake( | ||
positions: torch.Tensor, | ||
sin: torch.Tensor, | ||
cos: torch.Tensor, | ||
query: torch.Tensor, | ||
offsets: Optional[torch.Tensor] = None, | ||
rotate_style: int = 0, | ||
is_nope_first: bool = False, | ||
) -> None: | ||
pass | ||
|
||
|
||
if is_rocm_rotary_embedding_enabled(): | ||
|
||
direct_register_custom_op( | ||
op_name="rocm_aiter_rotary_emb_with_key_forward_hip", | ||
op_func=rocm_aiter_rotary_emb_with_key_forward_hip_impl, | ||
mutates_args=["key", "query"], | ||
fake_impl=rocm_aiter_rotary_emb_with_key_forward_hip_fake, | ||
dispatch_key=current_platform.dispatch_key, | ||
) | ||
|
||
direct_register_custom_op( | ||
op_name="rocm_aiter_rotary_emb_without_key_forward_hip", | ||
op_func=rocm_aiter_rotary_emb_without_key_forward_hip_impl, | ||
mutates_args=["query"], | ||
fake_impl=rocm_aiter_rotary_emb_without_key_forward_hip_fake, | ||
dispatch_key=current_platform.dispatch_key, | ||
) |
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.
Uh oh!
There was an error while loading. Please reload this page.