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

[Feature] vLLM ARM Enablement for AARCH64 CPUs #9228

Merged
merged 17 commits into from
Nov 26, 2024

Conversation

sanketkaleoss
Copy link
Contributor

@sanketkaleoss sanketkaleoss commented Oct 10, 2024

Description
This PR enables support of vLLM for AARCH64 architecture. Motivated by the requirements from (#176, #5741, etc), I implemented this PR which enables the ARM path for CPU inference.

ARM Compatibility: Modified the build scripts and configuration files to ensure compatibility with ARM processors. It currently supports float32, fp16 and bfloat16 datatypes.

Motivation
Enabling vLLM on ARM architecture broadens its usability, allowing it to run on a wider range of devices, including those with ARM processors. This enhancement is crucial for expanding the reach and applicability of vLLM in various use cases.

Checklist

  • Code changes have been tested on ARM devices (Graviton3).

Modifications

  1. Added NEON intrinsics for enabling vLLM on ARM
  2. Updated requirements-cpu.txt
  3. Added ARM path in cpu_extension.cmake

Performance Results

Model : facebook/opt-125m
Datatype : float32

image

Copy link

👋 Hi! Thank you for contributing to the vLLM project.
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 do one of these:

  • Add ready label to the PR
  • Enable auto-merge.

🚀

@sanketkaleoss sanketkaleoss changed the title vLLM ARM Enablement for AARCH64 CPUs [Feature] vLLM ARM Enablement for AARCH64 CPUs Oct 22, 2024
@sanketkaleoss
Copy link
Contributor Author

sanketkaleoss commented Oct 22, 2024

@youkaichao @ChipKerchner @mgoin @pathorn @manojnkumar I have added a new feature that allows vLLM to run on ARM CPU backend. I have tested on AWS Graviton 3E. Please have a look at this PR.

Copy link
Member

@mgoin mgoin left a comment

Choose a reason for hiding this comment

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

Nice work @sanketkaleoss! This seems reasonable to me as a base implementation. It seems only compute in fp32 is supported, is that right?
It would be good to update the cpu installation documentation with how to build and to also add a new Dockerfile.arm

message(STATUS "ARMv8 architecture detected")
list(APPEND CXX_COMPILE_FLAGS
"-mcpu=native"
"-march=armv8.6-a"
Copy link
Member

Choose a reason for hiding this comment

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

Why was this specific arch chosen?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

armv8.6-a is the first architecture that supports advance SIMD instructions, bf16 support and SVE support. That's why I chose this specific arch.

Copy link
Member

Choose a reason for hiding this comment

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

This is fine for now, but for instance I think Apple's M1 CPU uses ARMv8.4-A so we should consider supporting older versions

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I see, noted.

Copy link
Collaborator

Choose a reason for hiding this comment

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

According to wikichip, Graviton3 is also ARMv8.4A -- is the Graviton3E different? I haven't found any documentation on it

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@tlrmchlsmth You're right, graviton 3 has armv8.4. The code needs 8.6 only for bf16 dependencies. But if I just use "-mcpu=native" or "-march=armv8.4-a+bf16" the code works fine on Graviton3 instances. What do you suggest here?

Copy link
Collaborator

Choose a reason for hiding this comment

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

If -march=armv8.4-a+bf16 works, then I suggest using that for this PR.

In general, I think it's best to go with a fairly minimal ISA version, and then explicitly specify what ISA extensions we want to use, especially because ARM has many extensions that are optional for several ISA versions before they become mandatory

Copy link
Collaborator

Choose a reason for hiding this comment

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

BTW, the M1 Mac (and I think M2) does not support BF16 NEON so unfortunately this wouldn't help there.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Hi @mgoin , @tlrmchlsmth I've added the dockerfile and documentation. Changed the arm arch to native. Please have a look.

@sanketkaleoss
Copy link
Contributor Author

Nice work @sanketkaleoss! This seems reasonable to me as a base implementation. It seems only compute in fp32 is supported, is that right? It would be good to update the cpu installation documentation with how to build and to also add a new Dockerfile.arm

Thanks @mgoin for the review. It supports fp32 and bf16 as of now. Sure, I'll work on updating the documentation and adding Dockerfile too.

@youkaichao
Copy link
Member

this is great, I will hand it over to @mgoin for detailed review.

it would be better if we can support macos m chips as well, they are also arm chips.

@sanketkaleoss
Copy link
Contributor Author

this is great, I will hand it over to @mgoin for detailed review.

it would be better if we can support macos m chips as well, they are also arm chips.

Thanks @youkaichao . I'll add that to future work.

message(STATUS "ARMv8 architecture detected")
list(APPEND CXX_COMPILE_FLAGS
"-mcpu=native"
"-march=armv8.6-a"
Copy link
Collaborator

Choose a reason for hiding this comment

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

If -march=armv8.4-a+bf16 works, then I suggest using that for this PR.

In general, I think it's best to go with a fairly minimal ISA version, and then explicitly specify what ISA extensions we want to use, especially because ARM has many extensions that are optional for several ISA versions before they become mandatory

message(STATUS "ARMv8 architecture detected")
list(APPEND CXX_COMPILE_FLAGS
"-mcpu=native"
"-march=armv8.6-a"
Copy link
Collaborator

Choose a reason for hiding this comment

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

BTW, the M1 Mac (and I think M2) does not support BF16 NEON so unfortunately this wouldn't help there.


namespace vec_op {

// FIXME: FP16 is not fully supported in Torch-CPU
Copy link
Collaborator

Choose a reason for hiding this comment

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

@sanketkaleoss do you have plans to support FP16 in a future PR? I see that it's partially implemented.

Do you know what the problem with FP16 in Torch-CPU is?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I'm not sure what's the problem there. I can add the support in future PR if torch-CPU supports FP16 in future.

@caijijuhe
Copy link

Hello, I followed your Dockerfile.arm encountered during the build process。
ERROR: ld.so: object '/usr/lib/aarch64-linux-gnu/libtcmalloc_minimal.so.4' from LD_PRELOAD cannot be preloaded (cannot open shared object file): ignored.
Can you share the built image

@mergify mergify bot added ci/build documentation Improvements or additions to documentation labels Nov 4, 2024
elseif (POWER9_FOUND OR POWER10_FOUND)
message(STATUS "PowerPC detected")
# Check for PowerPC VSX support
list(APPEND CXX_COMPILE_FLAGS
"-mvsx"
"-mcpu=native"
"-mtune=native")

elseif (ASIMD_FOUND)
message(STATUS "ARMv8 architecture detected")
Copy link
Collaborator

Choose a reason for hiding this comment

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

Since it could be ARMv9:

Suggested change
message(STATUS "ARMv8 architecture detected")
message(STATUS "ARMv8 or later architecture detected")

or maybe explicitly call out NEON instead

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Okay, sure

Comment on lines 88 to 89
"-mcpu=native"
"-mtune=native"
Copy link
Collaborator

Choose a reason for hiding this comment

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

Native is OK for local development, however the binaries won't be portable at all.

I actually thought you were on the right track before this change. We really ought to explicitly specify a fairly minimal base ARM architecture (maybe ARMv8.4) and then explicitly the set of extensions that we need to build (BF16, and maybe FP16 and DotProd?)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I see, I'll try to set it as armv8.4 as a base implementation. Then, depending on the flags it would run bf16 or not.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Could you try using "-march=armv8.4-a+bf16+dotprod+fp16"?

Choose a reason for hiding this comment

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

Just for your guys reference, on arm64 -mcpu= acts as both specifying the appropriate architecture and tuning and it's generally better to use that vs -march if you're building for a specific CPU, you can find more details from here if you are running any test on Graviton: https://github.com/aws/aws-graviton-getting-started/blob/main/c-c++.md

Copy link
Contributor Author

@sanketkaleoss sanketkaleoss Nov 7, 2024

Choose a reason for hiding this comment

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

Could you try using "-march=armv8.4-a+bf16+dotprod+fp16"?

Yes, I tried it and it works. It even works with "-march=armv8.2-a+bf16+dotprod+fp16" making it compatible with Graviton2 as suggested by @ddynwzh1992 .

Copy link

mergify bot commented Nov 6, 2024

This pull request has merge conflicts that must be resolved before it can be
merged. @sanketkaleoss please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Nov 6, 2024
@sanketkaleoss
Copy link
Contributor Author

@mgoin @tlrmchlsmth I have created separate paths for FP32 and BF16 datatypes. This code now works on any arm machine from armv8.2-a onwards. Even if BF16 support is not present, it will choose FP32 path. Please have a look.

Copy link

mergify bot commented Nov 17, 2024

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @sanketkaleoss.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@sanketkaleoss
Copy link
Contributor Author

Hi @mgoin @tlrmchlsmth @ddynwzh1992 , I've implemented the suggested changes and now ARM CPU backend path supports FP32, FP16 and BF16 datatypes. Support for MAC devices has been added. I've tested on -march = armv8.2-a onwards and it runs fine even without the added extensions. Please review the changes.

@ShawnD200
Copy link
Contributor

Even if BF16 support is not present, it will choose FP32 path.

Maybe I missed something, how will it choose FP32 path if bp16 not supported? Thank you.

@sanketkaleoss
Copy link
Contributor Author

Even if BF16 support is not present, it will choose FP32 path.

Maybe I missed something, how will it choose FP32 path if bp16 not supported? Thank you.

Earlier it used to give compilation error even if we just want to use FP32 and not BF16, as there were no separate paths for FP32 and Bf16 in x86 backend case. Now, it will compile successfully even if the system doesn't have BF16 extension. User can run by selecting FP32 dtype in that case.

@sanketkaleoss
Copy link
Contributor Author

@mgoin @tlrmchlsmth @WoosukKwon @youkaichao @DarkLight1337 Your attention is required on this PR. Thanks.

@ShawnD200
Copy link
Contributor

ShawnD200 commented Nov 19, 2024

Even if BF16 support is not present, it will choose FP32 path.

Maybe I missed something, how will it choose FP32 path if bp16 not supported? Thank you.

Earlier it used to give compilation error even if we just want to use FP32 and not BF16, as there were no separate paths for FP32 and Bf16 in x86 backend case. Now, it will compile successfully even if the system doesn't have BF16 extension. User can run by selecting FP32 dtype in that case.

I thought you meant BF16 is supported without its hardware feature.

Edited: from Apple M2, BF16 has been supported.

if (remainder > 0) {
float16x8_t temp = reg.val[full_blocks];
for (int i = 0; i < remainder; ++i) {
reinterpret_cast<__fp16*>(ptr)[full_blocks * 8 + i] = vgetq_lane_f16(temp, i);
Copy link
Contributor

Choose a reason for hiding this comment

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

Hello, are you sure this line can compile? The index must be a compile-time constant, no?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thanks for the review. It compiles on my system as of now.

Copy link
Contributor

Choose a reason for hiding this comment

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

Thanks for the reply. I found that with gcc v12.3 (Ubuntu), optimization O1 and above does not require const (while O0 does), with clang v15.0 (macOS), I haven't been able to compile it.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thanks for the reply. I found that with gcc v12.3 (Ubuntu), optimization O1 and above does not require const (while O0 does), with clang v15.0 (macOS), I haven't been able to compile it.

I see, thanks for the information.

Copy link

mergify bot commented Nov 19, 2024

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @sanketkaleoss.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@sanketkaleoss
Copy link
Contributor Author

This essentially looks good to me, left a few comments. Please resolve the merge conflict and I can set the ready tag for CI to run.

I was able to test and run this on my M1 Macbook - nice work!

Server:

docker build -f Dockerfile.arm -t vllm-cpu-env --shm-size=4g . 
docker run -v ~/.cache/huggingface:/root/.cache/huggingface -p 8000:8000 --ipc=host vllm-cpu-env --model Qwen/Qwen2.5-0.5B-Instruct --dtype float16

Client:

curl http://0.0.0.0:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer DUMMY" \
  -d '{
    "model": "Qwen/Qwen2.5-0.5B-Instruct",
    "messages": [
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "What is the best pie?"
      }
    ]
  }'

Output:

{"id":"chatcmpl-210cbdc345ce48d9988a9aced2f83300","object":"chat.completion","created":1732058836,"model":"Qwen/Qwen2.5-0.5B-Instruct","choices":[{"index":0,"message":{"role":"assistant","content":"As an AI language model, I don't have personal preferences or tastes, but I can provide information about different types of pies based on popular opinion.\n\nThe best pie for everyone depends on personal taste, preferences, and dietary needs. Here are some popular types of pies:\n\n1. Pumpkin pie: It is a classic pie that uses pumpkins as a main ingredient. It is often served with crème fraîche, whipped cream, or chocolate ganache.\n\n2. Apple pie: Apple pie is a sweet and fruity pie that is popular in the autumn season. It uses apples as a main ingredient and is often served with whipped cream or chocolate ganache.\n\n3. Blueberry pie: Blueberry pie is a classic pie that uses blueberries as a main ingredient. It is often served with whipped cream or chocolate ganache.\n\n4. Chocolate pie: Chocolate pie is a sweet and rich pie that is popular in the winter season. It uses chocolate as a main ingredient and is often served with whipped cream or chocolate ganache.\n\n5. Lemon pie: Lemon pie is a fresh and zesty pie that is popular in the summer season. It uses lemons as a main ingredient and is often served with whipped cream or chocolate ganache.\n\nUltimately, the best pie for everyone depends on their personal preferences and dietary needs.","tool_calls":[]},"logprobs":null,"finish_reason":"stop","stop_reason":null}],"usage":{"prompt_tokens":25,"total_tokens":290,"completion_tokens":265,"prompt_tokens_details":null},"prompt_logprobs":null}

Thanks for the detailed review @mgoin ! I have rebased and implemented the suggested changes. Please proceed with the CI checks.

@mgoin mgoin added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 20, 2024
Copy link
Member

@mgoin mgoin left a comment

Choose a reason for hiding this comment

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

Sorry for losing track, merged with main to see if green

@youkaichao youkaichao merged commit a6760f6 into vllm-project:main Nov 26, 2024
71 of 73 checks passed
@sanketkaleoss
Copy link
Contributor Author

Sorry for losing track, merged with main to see if green

No worries, looking forward to the future work on this PR.

afeldman-nm pushed a commit to neuralmagic/vllm that referenced this pull request Dec 2, 2024
Signed-off-by: Sanket Kale <[email protected]>
Co-authored-by: Sanket Kale <[email protected]>
Co-authored-by: mgoin <[email protected]>
Signed-off-by: Andrew Feldman <[email protected]>
@animalnots
Copy link

@sanketkaleoss Hi, thank you for the PR. Does additional work needs to be done to support Qwen/Qwen2-VL-2B-Instruct ?

@sidharthrajaram
Copy link

@animalnots Did you figure that out? There's an issue related to it: #11154

sleepwalker2017 pushed a commit to sleepwalker2017/vllm that referenced this pull request Dec 13, 2024
Signed-off-by: Sanket Kale <[email protected]>
Co-authored-by: Sanket Kale <[email protected]>
Co-authored-by: mgoin <[email protected]>
@animalnots
Copy link

@animalnots Did you figure that out? There's an issue related to it: #11154

I ended up using this

satvikahuja/Easy-qwen2vlm2b-4macbook#1

Vllm never worked and instead would output "!!!!!!!!!..."

@sanketkaleoss
Copy link
Contributor Author

@sanketkaleoss Hi, thank you for the PR. Does additional work needs to be done to support Qwen/Qwen2-VL-2B-Instruct ?

Hi, sorry for the late reply. Does the model work with x86 CPU path?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ci/build documentation Improvements or additions to documentation ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Development

Successfully merging this pull request may close these issues.

9 participants