Skip to content

feat(backend-native): add NEON kernels for AArch64 - #192

Open
dv333 wants to merge 1 commit into
integrallis:mainfrom
dv333:apple-silicon-neon-kernels
Open

dv333 wants to merge 1 commit into
integrallis:mainfrom
dv333:apple-silicon-neon-kernels

Conversation

@dv333

@dv333 dv333 commented Sep 18, 2026

Copy link
Copy Markdown

Summary

On AArch64 every quantized matrix kernel runs the scalar fallback, because model-kernels only has AVX2 paths. This adds NEON versions of the Q4_0, Q5_0, Q8_0, Q4_K and Q6_K kernels. They are selected at runtime when the CPU has the dot-product extension: all Apple Silicon, Graviton 2 and later, and Ampere Altra.

Results are bit-identical to the scalar kernels, so outputs and qualification evidence don't move.

Results

Apple M5 Pro (15 cores), macOS 27, OpenJDK 25.0.2. Gemma 3 1B Q4_K_M (the qualified bartowski.google-gemma-3-1b-it-gguf.q4_k_m marker, rust-ffm backend), models.native.quantizedDecode=true, default thread counts:

Kernel library Decode Prefill (378-token prompt)
Released 0.3.42 macos-aarch64 library 17.8 tokens/s 35 tokens/s
This branch with JMODELS_KERNELS_DISABLE_NEON=1 18.2 tokens/s
This branch 48.5 tokens/s 165 tokens/s
  • Decode: steady-state runs 2 and 3 of 3, 182 generated tokens each.
  • Prefill: runs 2 to 4, each prompt starting with a different number so the prompt cache can't help.
  • Setup: all rows ran back to back in one session, loading the library through -Dmodels.native.kernels.library.
  • Output: the greedy completion is byte-identical (same SHA-256) with the released library and with this branch.

Design

  • Only the exact integer arithmetic is vectorized. That covers the Q8 block sums for Q4_0, Q5_0 and Q8_0, the Q4_K group dot products, and the Q6_K per-lane sums. The Q6_K sums keep the scalar kernel's index & 7 lane partition.
  • Floating-point steps are unchanged and run in the scalar order. That's what makes the results bit-identical. The AVX2 Q5_0 path, for comparison, reduces its lanes in a different order.
  • sdot goes through inline assembly because vdotq_s32 is still unstable (Tracking Issue for NEON dot product intrinsics rust-lang/rust#117224). It can switch to the intrinsic once that stabilizes.
  • Selection uses is_aarch64_feature_detected!("dotprod") and falls back to scalar otherwise. JMODELS_KERNELS_DISABLE_NEON=1 forces scalar, which allows A/B comparisons with a single library.
  • Scope of the change: the new code lives in neon.rs. lib.rs only gains the DotKernel variants, kernel selection and the two dispatch points. There's no ABI change.
  • Not covered, still scalar on AArch64: Q5_K, activation quantization, attention and the Gated DeltaNet kernels.

Testing

  • cargo test --release: 23 passed. That includes 5 new tests that compare each NEON kernel with its scalar counterpart bit for bit. They cover single rows and batched row ranges, batch sizes 1 and 3, activations that include -128 and 127, and random scales, bit planes and packed K-quant scales.
  • cargo clippy --release --all-targets -- -D warnings and cargo fmt --check: clean.
  • gradle :backend-native:check -PmodelsNativePlatform=macos-aarch64 (the task the macOS AArch64 CI job runs): build successful, 42 Java tests passed.
  • Not run locally: linux-aarch64 and windows-aarch64. The code is plain core::arch::aarch64 plus one sdot instruction, so CI should cover both.

A related finding, not changed here

With faster kernels, one decode worker per core starts to cost throughput. The Java thread that runs between native calls competes with the spinning workers. On the same M5 Pro with this branch:

  • models.native.kernels.decodeThreads=12 (out of 15) gave 59 tokens/s, against 48.5 at the default.
  • In an earlier run while other work was going on, 15 threads fell to 26 tokens/s while 12 stayed near 60.

The Linux EPYC profile runs 8 kernel threads on 8 vCPUs, so it may be affected too. I'm happy to open a separate issue with the data.

@dv333

dv333 commented Sep 18, 2026

Copy link
Copy Markdown
Author

@bsbodden The CI workflows are waiting for approval since this is my first PR here. I can also share the small Java program behind the benchmark numbers, or move the decode-thread note into its own issue, whichever helps.

@dv333 dv333 left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Pls review

AArch64 CPUs ran the scalar fallback for every quantized matrix kernel,
because the crate only had AVX2 paths. This adds NEON versions of the
Q4_0, Q5_0, Q8_0, Q4_K and Q6_K kernels, selected at runtime when the CPU
has the dot-product extension (all Apple Silicon, Graviton 2 and later,
Ampere Altra).

Only the exact integer arithmetic is vectorized: the Q8 block sums, the
Q4_K group dot products and the Q6_K per-lane sums. Every floating-point
step keeps the scalar kernels' operations and order, so results are
bit-identical to the scalar path. New tests assert that for single rows
and batched row ranges of every format.

sdot is emitted through inline assembly because vdotq_s32 is not yet
stable (rust-lang/rust#117224). JMODELS_KERNELS_DISABLE_NEON=1 selects
the scalar kernels, which allows A/B comparisons with one library.

Gemma 3 1B Q4_K_M on an Apple M5 Pro, models.native.quantizedDecode=true,
default thread counts: decode 17.8 -> 48.5 tokens/s, prefill of a
378-token prompt 35 -> 165 tokens/s, identical output.
@dv333
dv333 force-pushed the apple-silicon-neon-kernels branch from be717b8 to 3f602ef Compare September 18, 2026 21:50
@dv333

dv333 commented Sep 18, 2026

Copy link
Copy Markdown
Author

@bsbodden can you pls review this PR?

@bsbodden

Copy link
Copy Markdown
Member

I will... thanks for the contribution!

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.

2 participants