Skip to content
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

Fused ROPE and reshape cache kernel #229

Open
wants to merge 14 commits into
base: main
Choose a base branch
from

Conversation

maleksan85
Copy link

@maleksan85 maleksan85 commented Oct 11, 2024

Added support to ROPE kernel to store values right into cache.

LLama3.1 8B fp16 test without using fused ROPE:

root@banff-cyxtera-s82-5:~/workspace/vllm# HIP_VISIBLE_DEVICES=5 VLLM_FUSED_ROPE_W_KV_CACHE=0 python benchmarks/benchmark_latency.py --model /data/models/Meta-Llama-3.1-8B --input_len 1024 --output_len 1024 --batch-size=16

Avg latency: 9.835484731818239 seconds

with fused kernel:

root@banff-cyxtera-s82-5:~/workspace/vllm# HIP_VISIBLE_DEVICES=5 VLLM_FUSED_ROPE_W_KV_CACHE=1 python benchmarks/benchmark_latency.py --model /data/models/Meta-Llama-3.1-8B --input_len 1024 --output_len 1024 --batch-size=16

Avg latency: 9.86210185677434 seconds

Also perf in test (although it is not very fair as kernels were run only once). Taking some test in the end of run of test suite:

tests/kernels/test_fused_rope_and_reshape_cache.py::test_fused_rotary_embedding_with_reshape_cache[5-0-False-dtype1-None-128-8-32-16-8-auto-1024-16] Non fused call 0.22136466577649117 ms
Fused run 0.0467388890683651 ms

Shapes are pretty small, so kernel start overhead might be sufficient.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant