Skip to content

[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 14 commits into from
Aug 11, 2025

Conversation

vllmellm
Copy link
Contributor

@vllmellm vllmellm commented Aug 8, 2025

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.

Purpose

Integrate aiter rope ops in the RotaryEmbedding module, which boosts model performance.

Benchmark Results

meta-llama/Meta-Llama-3-8B-Instruct

Metric With Aiter Rope Without Aiter Rope
Request Throughput (req/s) 2.52 2.53
Output Token Thpt (tok/s) 2,089.46 1,995.06
Total Token Thpt (tok/s) 4,608.15 4,525.02
Mean TTFT (ms) 431.57 499.96
Median TTFT (ms) 131.37 179.53
P99 TTFT (ms) 1,423.43 1,557.27
Mean TPOT (ms) 23.32 25.15
Median TPOT (ms) 16.01 16.12
P99 TPOT (ms) 184.30 206.32
Mean ITL (ms) 15.95 16.01
Median ITL (ms) 15.24 15.07
P99 ITL (ms) 49.05 50.28

deepseek-ai/DeepSeek-V2-Lite-Chat

Metric With Aiter Rope Without Aiter Rope
Request Throughput (req/s) 3.12 3.07
Output Token Thpt (tok/s) 2,993.27 2,986.46
Total Token Thpt (tok/s) 6,112.77 6,055.55
Mean TTFT (ms) 127.90 204.32
Median TTFT (ms) 81.33 109.66
P99 TTFT (ms) 453.34 852.63
Mean TPOT (ms) 11.70 11.66
Median TPOT (ms) 11.42 11.52
P99 TPOT (ms) 14.14 13.11
Mean ITL (ms) 11.68 11.70
Median ITL (ms) 10.82 10.90
P99 ITL (ms) 39.48 39.37

IMPORTANT NOTE: You must use --compilation-config '{ "custom_ops": ["+rotary_embedding"]}' to enable this custom ops.

benchmark setting

python vllm/benchmarks/benchmark_serving.py --backend vllm --model "$model_name" --dataset-name random --num-prompts 50 --request-rate 10 --random-input-len 1000 --random-output-len 1000

Test Plan

Test models that are afftected by this change, using lm_eval on gsm8k dataset.

environment setting

Step 1: run vllm serve

VLLM_USE_V1=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_LINEAR=0 SAFETENSORS_FAST_GPU=1

vllm serve $MODEL_NAME --compilation-config '{ "custom_ops": ["+rotary_embedding"]}' --trust-remote-code --swap-space 16 --distributed-executor-backend mp

Step 2: run lm_eval

lm_eval --model local-completions --tasks gsm8k --model_args model=$MODEL_NAME,base_url=http://localhost:8000/v1/completions --trust_remote_code --num_fewshot 5 --batch_size 256

Tested models:

  • meta-llama/Meta-Llama-3-8B-Instruct (tests
    Llama3RotaryEmbedding)
  • deepseek-ai/DeepSeek-V2-Lite-Chat (tests DeepseekScalingRotaryEmbedding)
  • tencent/Hunyuan-A13B-Pretrain (tests DynamicNTKAlphaRotaryEmbedding)
  • NousResearch/Yarn-Mistral-7b-128k (tests YaRNScalingRotaryEmbedding)
  • Qwen/Qwen3-235B-A22B-FP8
  • mistralai/Mixtral-8x7B-Instruct-v0.1
  • mistralai/Mistral-7B-Instruct-v0.1

Test Result

meta-llama/Meta-Llama-3-8B-Instruct

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.7566 0.7528 0.0118 0.0119
strict-match 5 exact_match 0.7589 0.7551 0.0118 0.0118

deepseek-ai/DeepSeek-V2-Lite-Chat (-tp 1)

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.6626 0.6611 0.0130 0.0130
strict-match 5 exact_match 0.6566 0.6520 0.0131 0.0131

tencent/Hunyuan-A13B-Pretrain (-tp 2)

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.6422 0.6346 0.0132 0.0133
strict-match 5 exact_match 0.3548 0.3450 0.0132 0.0131

NousResearch/Yarn-Mistral-7b-128k (-tp 2)

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.2790 0.2813 0.0124 0.0124
strict-match 5 exact_match 0.2767 0.2790 0.0123 0.0124

Qwen/Qwen3-235B-A22B-FP8

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.8802 0.8741 0.0089 0.0091
strict-match 5 exact_match 0.8605 0.8544 0.0095 0.0097

mistralai/Mixtral-8x7B-Instruct-v0.1

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.6497 0.6475 0.0131 0.0132
strict-match 5 exact_match 0.6452 0.6429 0.0132 0.0132

mistralai/Mistral-7B-Instruct-v0.1 (-tp 2)

Tasks Version Filter n-shot Metric Value (Before) Value (After) Stderr (Before) Stderr (After)
gsm8k 3 flexible-extract 5 exact_match 0.3404 0.3381 0.0131 0.013
strict-match 5 exact_match 0.3336 0.3328 0.0130 0.013

vllmellm added 12 commits August 5, 2025 09:39
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Copy link

github-actions bot commented Aug 8, 2025

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

@mergify mergify bot added the rocm Related to AMD ROCm label Aug 8, 2025
Copy link
Contributor

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

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 integrates AITensaR (AITER) Rotary Position Embedding (RoPE) operations for ROCm to improve performance. The changes look promising and the benchmark results are positive. I've found a critical bug that could lead to a runtime error and a typo in a function name that should be corrected for maintainability. Please see my detailed comments.

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
@vllmellm vllmellm marked this pull request as ready for review August 8, 2025 12:49
@DarkLight1337 DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 11, 2025
Copy link
Member

@DarkLight1337 DarkLight1337 left a comment

Choose a reason for hiding this comment

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

LGTM, can you merge from main to fix CI?

@vllm-bot vllm-bot merged commit 9c97a1c into vllm-project:main Aug 11, 2025
37 of 41 checks passed
jingyu-ml pushed a commit to jingyu-ml/vllm that referenced this pull request Aug 11, 2025
…-project#22521)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Signed-off-by: jingyu <jingyu@omniml.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants