[ARM] [SDPA] SVE implementation of MHASingleToken for FP32 #27273
Conversation
NishantPrabhuFujitsu
requested review from
ilya-lavrenov
and removed request for
a team
github-actions
bot
added
category: CPU
.gitmodules
Outdated
| @@ -87,3 +87,6 @@ | |||
| [submodule "src/plugins/intel_cpu/thirdparty/shl"] | |||
| path = src/plugins/intel_cpu/thirdparty/shl | |||
| url = https://github.com/openvinotoolkit/shl.git | |||
| [submodule "thirdparty/src/plugins/intel_npu/thirdparty/level-zero"] | |||
There was a problem hiding this comment.
Please remove all submodule changes from the PR.
There was a problem hiding this comment.
I have (tried my best to) clean this up. Please let me know if it's not reverted yet.
If the changes haven't reverted, please guide me on how I can fix it. Sorry and thanks!
There was a problem hiding this comment.
Submodule changes are still there. I just pushed the commit that reverts all unnecessery changes e2d5f11. Please apply it on top of your branch.
Please don't include any submodule changes in the commits. I would recommend to call in your working folder to have actual state.
git submodule init
git submodule update
There was a problem hiding this comment.
@dmitry-gorokhov I have applied your commit as a patch and now the changes seem to have reverted. Please take a look and let me know.
Meanwhile, I was trying to rebase my branch with master but I receive merge conflicts like so (conflict on the whole folder?):
Could you please help?
There was a problem hiding this comment.
You can use "git merge ... " instead of "git rebase ..." to pick up latest master changes.
Or squash all 10 commits from your branch into single one and call "git rebase ..." after that
| @@ -246,6 +249,79 @@ static constexpr size_t vec_len_f16_neon = vec_len_neon / sizeof(ov::float16); | |||
| #endif | |||
|
|
|||
| #ifdef OPENVINO_ARCH_ARM64 | |||
| #if defined(__ARM_FEATURE_SVE) && !defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) | |||
There was a problem hiding this comment.
- Could you please clarify why we need
!defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)check here? - Lets use
HAVE_SVEinstead of__ARM_FEATURE_SVE
There was a problem hiding this comment.
- It was a hotfix I had added to silence some errors when testing out my changes initially. They are no longer needed, so I have removed them in the latest commit.
- Updated in the latest commit.
src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/mha_single_token.cpp
Outdated
Show resolved
Hide resolved
github-actions
bot
added
category: inference
src/inference/src/system_conf.cpp
Outdated
| # elif defined(__aarch64__) && defined(__APPLE__) | ||
| int64_t result(0); | ||
| size_t size = sizeof(result); | ||
| const std::string& cap = "hw.optional.sve"; |
There was a problem hiding this comment.
I suppose it's just a placeholder? I don't see such HW capability on macOS
There was a problem hiding this comment.
I checked for current or upcoming SVE support on Apple Metal, but couldn't find any. To ensure I don't miss it, I checked on Perplexity and this is what it suggested. I'd like your suggestion on whether it should be kept or removed.
There was a problem hiding this comment.
I think it is ok to return false for Apple so far. We can consider SVE/SME support on M4 as separate activity.
There was a problem hiding this comment.
Fixed to return "false" in latest commit.
github-actions
bot
removed
category: GPU
NishantPrabhuFujitsu
force-pushed
the
mha-single-token-arm-sve-f32
branch
from
14ca0b6
to
76d0305
Compare
NishantPrabhuFujitsu
force-pushed
the
mha-single-token-arm-sve-f32
branch
from
76d0305
to
43ad2e1
Compare
Details:
MHASingleTokenfor SVE-128, SVE-256 and SVE-512 platforms.exp_ps_<isa>using fewer FMA operations. Executes ~18% faster and has better output precision.Note: I am aware of the Neon FP16 implementation of SDPA added recently. To accommodate for this, the current SVE changes will be used only if the hardware does not have ARM FP16 support. I will follow up with SVE FP16 implementations soon.
[SVE] Benchmarking results
Below are the benchmarking results of execution time of each ported function. Measurements were performed by running each function individually on dummy inputs (128 fp32 elements) for 1,000,000 iterations and computing average time (in micro-seconds).
Execution time of
MHASingleTokenas a whole was also measured for two LLMs, the results of which are shown below. For LlaMA-3-8B, the SVE-128 and SVE-512 systems at my disposal did not have enough memory, so only SVE-256 results are shown. While there is an improvement overall, these results could be contaminated with run-to-run variation due to the small execution time of the kernel.Benchmarking details: Prompt length of 108 tokens was used; total time for generating 50 tokens was measured and average execution time was computed.
New exponential implementation
It is based on the discussion in these slides (this is based on a past talk in Fujitsu hence the document is in Japanese, sorry!). The algorithm followed is slightly different from the current implementation, in that it uses
fexpainstruction available on ARM and requires only 3 Taylor expansion terms (2 FMA operations) to be precise until the 8th decimal place.Our benchmarking results showed this implementation to be 44%-58% faster than the existing Neon implementation. It is ~18% faster than the SVE implementation of the current algorithm in Neon.
In this PR, the new implementation is called by default. The SVE port of the existing Neon implementation has also been retained, if needed.