Skip to content

arm: NEON + i8mm vec_dot for PTQ1_0 - #290

Open
karusrus wants to merge 1 commit into
PrismML-Eng:prismfrom
karusrus:arm-neon-ptq1_0-pr
Open

karusrus wants to merge 1 commit into
PrismML-Eng:prismfrom
karusrus:arm-neon-ptq1_0-pr

Conversation

@karusrus

Copy link
Copy Markdown

Overview

Follow-up to #265, which dropped its PTQ1_0 NEON kernel and asked for a validated one. This adds a NEON vec_dot for PTQ1_0 x Q8_0 on ARM, with an i8mm path for prompt processing.

  • Trits are decoded in 8-bit lanes: ((q + (q >> 1)) >> 1) >> 6 equals (3q) >> 8 for every byte value (checked for all 256), so there is no widening to u16. One block becomes eight int8x16 vectors in element order.
  • Each 32-wide Q8_0 sub-block is reduced with sdot and accumulated in float32x4; one horizontal add per row.
  • With __ARM_FEATURE_MATMUL_INT8, PTQ1_0 uses nrows = 2 and computes a 2x2 tile with SMMLA, so every decoded trit vector is used for two activation columns.

Results

Snapdragon 7 Gen 4 (Cortex-A720/A520, dotprod + i8mm), Ternary Bonsai 2 27B PTQ1_0, CPU only, -t 5, runs alternated:

build pp64 t/s tg16 t/s
generic 1.22 1.09
this PR 2.33 1.31

test-backend-ops perf, MUL_MAT m=4096 k=14336: n=1 65.9 GFLOPS, n=512 153.4 GFLOPS.

Tests: with #265 merged on top, test-quantize-fns passes, including its exact packed ternary vec_dot check ("ternary packed dot products: ok (0 failures)"). That check uses nrc == 1. The nrc == 2 (i8mm) path was checked end to end: greedy 24-token output on the 27B is byte-identical to the generic path. The file also builds for armv7a (NEON, no i8mm).

Additional information

Requirements

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: YES. The kernel was written with Claude Code. I built, tested and measured it on my device.

@github-actions github-actions Bot added the ggml label Sep 28, 2026
@bri-prism

Copy link
Copy Markdown
Collaborator

@karusrus @kiljoy001, thanks for already working this out on #265. Both kernels test bit-exact on an M5 Pro, including the i8mm two-row path. Our plan is to merge #265 first and then #290, which only needs its arch-fallback.h alias lines dropped when rebasing.

@karusrus

karusrus commented Oct 5, 2026

Copy link
Copy Markdown
Author

Thanks @bri-prism, sounds good. Once #265 is in, I'll rebase this on master and drop the arch-fallback.h alias lines: the ARM PTQ1_0 alias from #265 goes away, since this PR adds the NEON kernel, and the PQ2_0 x Q8_K alias stays. If it's easier for you to do that while merging, that works too.

PTQ1_0 had only the generic vec_dot on ARM. This adds a NEON kernel:

- trits are decoded in 8-bit lanes: ((q + (q >> 1)) >> 1) >> 6 equals
  (3q) >> 8 for every byte value, so no widening is needed; a 128-trit
  block becomes eight int8x16 vectors in element order
- each 32-wide Q8_0 sub-block is reduced with sdot and accumulated into
  a float32x4 with vmlaq_n_f32, one horizontal add per row
- with __ARM_FEATURE_MATMUL_INT8, PTQ1_0 uses nrows = 2 and the nrc == 2
  path computes a 2x2 tile with vmmlaq_s32, reusing decoded trits for
  two activation columns

Snapdragon 7 Gen 4, Ternary Bonsai 2 27B, -t 5: pp64 1.22 -> 2.33 t/s,
tg16 1.09 -> 1.31 t/s against the generic path. test-quantize-fns passes;
greedy output matches the generic path on the same device.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@karusrus
karusrus force-pushed the arm-neon-ptq1_0-pr branch from f77a146 to 9098196 Compare October 6, 2026 11:15
@karusrus

karusrus commented Oct 6, 2026 •

Copy link
Copy Markdown
Author

@bri-prism rebased on prism now that #265 is in. Both ARM aliases in arch-fallback.h are gone: the PTQ1_0 one because this PR adds the NEON kernel, and the PQ2_0 x Q8_K one because #265 brought a NEON kernel for it. I said earlier that the second alias would stay; that was before #265 landed in its final form. Nothing else changed.

test-quantize-fns passes on arm64 (NEON + dotprod). The i8mm path is the same code you tested on the M5 Pro.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants