[Feature] vLLM ARM Enablement for AARCH64 CPUs #9228

sanketkaleoss · 2024-10-10T06:17:01Z

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 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

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

Performance Results

Model : facebook/opt-125m
Datatype : float32

github-actions · 2024-10-10T06:17:12Z

👋 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 · 2024-10-22T05:13:24Z

@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.

mgoin

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

mgoin · 2024-10-22T10:55:57Z

cmake/cpu_extension.cmake

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


Why was this specific arch chosen?

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.

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

I see, noted.

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

@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?

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

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

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

sanketkaleoss · 2024-10-22T11:29:25Z

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 · 2024-10-24T03:24:54Z

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 · 2024-10-24T05:18:04Z

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.

tlrmchlsmth · 2024-10-25T14:18:00Z

cmake/cpu_extension.cmake

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


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

tlrmchlsmth · 2024-10-25T14:26:13Z

cmake/cpu_extension.cmake

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


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

tlrmchlsmth · 2024-10-25T14:27:23Z

csrc/cpu/cpu_types_arm.hpp

+
+namespace vec_op {
+
+// FIXME: FP16 is not fully supported in Torch-CPU


@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?

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 · 2024-10-30T11:51:37Z

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

tlrmchlsmth · 2024-11-06T00:35:21Z

cmake/cpu_extension.cmake

 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")


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

tlrmchlsmth · 2024-11-06T00:40:11Z

cmake/cpu_extension.cmake

+        "-mcpu=native"
+        "-mtune=native"


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?)

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.

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

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

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 .

mergify · 2024-11-06T09:10:21Z

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

sanketkaleoss · 2024-11-12T11:13:25Z

@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.

…ture to native

Signed-off-by: Sanket Kale <[email protected]>

sanketkaleoss changed the title ~~vLLM ARM Enablement for AARCH64 CPUs~~ [Feature] vLLM ARM Enablement for AARCH64 CPUs Oct 22, 2024

mgoin reviewed Oct 22, 2024

View reviewed changes

tlrmchlsmth reviewed Oct 25, 2024

View reviewed changes

mergify bot added ci/build documentation Improvements or additions to documentation labels Nov 4, 2024

tlrmchlsmth reviewed Nov 6, 2024

View reviewed changes

mergify bot added the needs-rebase label Nov 6, 2024

Sanket Kale and others added 10 commits November 14, 2024 17:51

Enabled path for ARM machines

3ff99cf

Enabled path for ARM machines

8f0683e

fixed typos

6683618

Added Dockerfile for ARM architecture support

cefb590

Added Documenataion for ARM architecture support and updated architec…

bdf0f30

…ture to native

Updated documentation

63a3cd8

Added compatibility with macos by creating separate fp32 and bf16 paths

903afbe

Fixed some missing conditions

908b954

fixed dco test

b8ab3bc

Signed-off-by: Sanket Kale <[email protected]>

Rebased and resolved merge conflicts

3d46d2e

sanketkaleoss force-pushed the vllm_arm_enablement branch from 9589161 to 3d46d2e Compare November 14, 2024 12:54

mergify bot removed the needs-rebase label Nov 14, 2024

Sanket Kale added 2 commits November 14, 2024 19:44

Updated documentation

c824f5e

Signed-off-by: Sanket Kale <[email protected]>

Added FP16 compatibility and updated torch version

47adfa6

Signed-off-by: Sanket Kale <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature] vLLM ARM Enablement for AARCH64 CPUs #9228

[Feature] vLLM ARM Enablement for AARCH64 CPUs #9228

sanketkaleoss commented Oct 10, 2024 •

edited

Loading

github-actions bot commented Oct 10, 2024

sanketkaleoss commented Oct 22, 2024 •

edited

Loading

mgoin left a comment

mgoin Oct 22, 2024

sanketkaleoss Oct 22, 2024

mgoin Oct 22, 2024

sanketkaleoss Oct 23, 2024

tlrmchlsmth Oct 24, 2024

sanketkaleoss Oct 25, 2024

tlrmchlsmth Oct 25, 2024

tlrmchlsmth Oct 25, 2024

sanketkaleoss Nov 5, 2024

sanketkaleoss commented Oct 22, 2024

youkaichao commented Oct 24, 2024

sanketkaleoss commented Oct 24, 2024

tlrmchlsmth Oct 25, 2024

tlrmchlsmth Oct 25, 2024

tlrmchlsmth Oct 25, 2024

sanketkaleoss Oct 26, 2024

caijijuhe commented Oct 30, 2024

tlrmchlsmth Nov 6, 2024

sanketkaleoss Nov 6, 2024

tlrmchlsmth Nov 6, 2024

sanketkaleoss Nov 6, 2024

tlrmchlsmth Nov 6, 2024

ddynwzh1992 Nov 7, 2024

sanketkaleoss Nov 7, 2024 •

edited

Loading

mergify bot commented Nov 6, 2024

sanketkaleoss commented Nov 12, 2024


		namespace vec_op {

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

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

[Feature] vLLM ARM Enablement for AARCH64 CPUs #9228

Are you sure you want to change the base?

[Feature] vLLM ARM Enablement for AARCH64 CPUs #9228

Conversation

sanketkaleoss commented Oct 10, 2024 • edited Loading

github-actions bot commented Oct 10, 2024

sanketkaleoss commented Oct 22, 2024 • edited Loading

mgoin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sanketkaleoss commented Oct 22, 2024

youkaichao commented Oct 24, 2024

sanketkaleoss commented Oct 24, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

caijijuhe commented Oct 30, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sanketkaleoss Nov 7, 2024 • edited Loading

Choose a reason for hiding this comment

mergify bot commented Nov 6, 2024

sanketkaleoss commented Nov 12, 2024

sanketkaleoss commented Oct 10, 2024 •

edited

Loading

sanketkaleoss commented Oct 22, 2024 •

edited

Loading

sanketkaleoss Nov 7, 2024 •

edited

Loading