Skip to content

Conv1d optimised kernel for aarch64 - #54093

Open
almayne wants to merge 5 commits into
vllm-project:mainfrom
almayne:conv1d
Open

Conv1d optimised kernel for aarch64#54093
almayne wants to merge 5 commits into
vllm-project:mainfrom
almayne:conv1d

Conversation

@almayne

@almayne almayne commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Purpose

Advanced SIMD implementation for aarch64 of the causal conv1d kernel used in Gated DeltaNets. Includes faster silu implementation with fexp.

Test Plan

Run related tests:

python -m pytest \
  tests/kernels/mamba/test_causal_conv1d.py \
  tests/kernels/mamba/cpu/test_cpu_gdn_ops.py \
  tests/kernels/mamba/test_cpu_short_conv.py \
  tests/v1/attention/test_gdn_metadata_builder.py -v

Run e2e inference on Qwen 3.5 MoE model and demonstrate improved performance.

OMP_NUM_THREADS=64  vllm bench throughput \
  --num-prompts 256 \
  --seed 0 \
  --dataset-name sharegpt \
 --dataset-path sharegpt_qwen35_256.json \
  --output-len 256 \
  --max-model-len 4096 \
  --max-num-batched-tokens 4096 \
  --model Qwen/Qwen3.5-35B-A3B \
  --tensor-parallel-size 1 \
  --no-enable-prefix-caching \
  --num-warmups 5

Test Result

Unit tests all pass.

Perf uplift of ~9% on 64 Neoverse V2 cores:

  • reqs/s 1.0882x
  • total tokens/s 1.0904x
  • output tokens/s 1.0904

almayne added 2 commits August 27, 2026 14:55
@almayne
Added aarch64 specific causal conv1d kernel.
21fcff6 
Co-authored-by: GPT 5.6 Sol

Signed-off-by: Anna Mayne <anna.mayne@arm.com>
@almayne
Make sure aarch64 has the correct fallbacks for conv state dim first …
5ddde48 
…storage.

Signed-off-by: Anna Mayne <anna.mayne@arm.com>
claude[bot]
claude Bot reviewed Aug 27, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added ci/build cpu Related to CPU backends labels Aug 27, 2026
@nicholaskh-ai nicholaskh-ai mentioned this pull request Sep 8, 2026
Open
fadara01
fadara01 reviewed Sep 8, 2026
View reviewed changes

@fadara01 fadara01 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thank you for your PR!

Added some comments

Comment thread csrc/cpu/sgl-kernels/conv_arm.h
Comment thread csrc/cpu/sgl-kernels/conv_arm.h
Comment on lines +76 to +87
bfloat16x8_t va_pair =
zip1q_bf16(va[0], va[1]); // channels 0-3, inputs for taps 0-1
vc0 = vbfdotq_f32(vc0, va_pair,
vb0); // 4 independent length-two dot products
va_pair = zip2q_bf16(va[0], va[1]); // channels 4-7, inputs for taps 0-1
vc1 = vbfdotq_f32(vc1, va_pair, vb1);
va_pair = zip1q_bf16(va[2], va[3]); // channels 0-3, inputs for taps 2-3
vc0 =
vbfdotq_f32(vc0, va_pair,
vb2); // Add to the previous result for the same channels
va_pair = zip2q_bf16(va[2], va[3]); // channels 4-7, inputs for taps 2-3
vc1 = vbfdotq_f32(vc1, va_pair, vb3);

@fadara01 fadara01 Sep 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

maybe we can we load more bf16 pairs from B and have more independent bfdot streams here to better exploit the CPUs SIMD pipelines?

does this improve perf in a micro benchmark for this op?

@almayne almayne Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I tried a 16-channel tile instead, doubling the number of bfdot instructions. This showed small improvements (~9%) in my microbenchmark for prefill with 1024 tokens and no change with 128 tokens. There was no improvement for decode across a range of batch sizes. And e2e tests using vllm bench throughput showed no perf uplift.

Comment thread cmake/cpu_extension.cmake Outdated
Comment thread vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
conv_state_view = conv_buf[:, :, : width - 1]
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
decode_conv_state = conv_state_view[decode_state_indices].contiguous()
decode_mixed_qkv = causal_conv1d_update_torch(

@fadara01 fadara01 Sep 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

are there any un-used functions we can trim from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py

@almayne almayne Sep 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I don't think so. See:
https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/ops/causal_conv1d.py#L1282-L1285

https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/short_conv.py#L102-L108

Comment thread csrc/cpu/torch_bindings.cpp
Comment thread vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py Outdated
Comment thread vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py Outdated
Comment thread vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py Outdated
Comment thread vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py Outdated
Comment thread vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py Outdated
fadara01
fadara01 reviewed Sep 9, 2026
View reviewed changes
Comment thread csrc/cpu/sgl-kernels/conv.cpp Outdated
@almayne
Moved weight packing to model-load time. Improved test coverage. Mino…
8e471ac 
…r changes from PR feedback.

Signed-off-by: Anna Mayne <anna.mayne@arm.com>
fadara01
fadara01 approved these changes Sep 15, 2026
View reviewed changes

@fadara01 fadara01 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM, thank you

@fadara01

fadara01 commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

@bigPYJ1151 - could you please take a look?

@fadara01

fadara01 commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

@bigPYJ1151 could you please have a look?
this enables fused conv1d for Arm and unifies the our path with x86

@almayne almayne changed the title Conv1d Conv1d optimised kernel for aarch64 Sep 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@claude claude[bot] claude[bot] left review comments
@bigPYJ1151 bigPYJ1151 Awaiting requested review from bigPYJ1151 bigPYJ1151 is a code owner
@tdoublep tdoublep Awaiting requested review from tdoublep tdoublep is a code owner
@tomeras91 tomeras91 Awaiting requested review from tomeras91 tomeras91 is a code owner
@mgoin mgoin Awaiting requested review from mgoin mgoin is a code owner
@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner
@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner
@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner
@zyongye zyongye Awaiting requested review from zyongye zyongye is a code owner
@AndreasKaratzas AndreasKaratzas Awaiting requested review from AndreasKaratzas AndreasKaratzas is a code owner
+1 more reviewer
@fadara01 fadara01 fadara01 approved these changes
Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

@fadara01 fadara01

Labels

ci/build cpu Related to CPU backends

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@almayne @fadara01