Conversation
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. |
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
force-pushed
the
apple-silicon-neon-kernels
branch
from
September 18, 2026 21:50
be717b8 to
3f602ef
Compare
Author
|
@bsbodden can you pls review this PR? |
Member
|
I will... thanks for the contribution! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On AArch64 every quantized matrix kernel runs the scalar fallback, because
model-kernelsonly 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_mmarker, rust-ffm backend),models.native.quantizedDecode=true, default thread counts:macos-aarch64libraryJMODELS_KERNELS_DISABLE_NEON=1-Dmodels.native.kernels.library.Design
index & 7lane partition.sdotgoes through inline assembly becausevdotq_s32is still unstable (Tracking Issue for NEON dot product intrinsics rust-lang/rust#117224). It can switch to the intrinsic once that stabilizes.is_aarch64_feature_detected!("dotprod")and falls back to scalar otherwise.JMODELS_KERNELS_DISABLE_NEON=1forces scalar, which allows A/B comparisons with a single library.neon.rs.lib.rsonly gains theDotKernelvariants, kernel selection and the two dispatch points. There's no ABI change.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 warningsandcargo fmt --check: clean.gradle :backend-native:check -PmodelsNativePlatform=macos-aarch64(the task the macOS AArch64 CI job runs): build successful, 42 Java tests passed.linux-aarch64andwindows-aarch64. The code is plaincore::arch::aarch64plus onesdotinstruction, 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.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.