Skip to content

Accelerate ternary v3.3 mask decoding with AVX2 and Arm NEON - #1

Merged
MohsenDehghankar merged 1 commit into
UIC-InDeXLab:mainfrom
dmangla3:optimize-ternary-coefficient-fma
Aug 10, 2026
Merged

Accelerate ternary v3.3 mask decoding with AVX2 and Arm NEON#1
MohsenDehghankar merged 1 commit into
UIC-InDeXLab:mainfrom
dmangla3:optimize-ternary-coefficient-fma

Conversation

@dmangla3

Copy link
Copy Markdown
Contributor

Summary

  • decode v3.3 positive/negative scatter masks through an aligned coefficient table
  • accumulate all active output rows with AVX2 or Arm NEON while preserving the scalar fallback and ABI
  • make the ternary CPU Makefile portable across x86, Apple Silicon, and non-Darwin OpenMP builds
  • add deterministic correctness coverage for random, alternating, and sparse masks at every SIMD group size (k=5..8)
  • add a same-source scalar-vs-SIMD benchmark and record the Apple M4 results

Apple M4 benchmark

Single-threaded v3.3 kernel latency for a 3200 x 8640 ternary matrix. Preprocessing is excluded. Each result is the median of 41 runs after 10 warmups.

k SIMD Scalar Speedup
6 5.653 ms 8.653 ms 1.53x
8 14.466 ms 31.164 ms 2.15x

The benchmark compiles the scalar comparison from the same source with RSR_DISABLE_SIMD and verifies the two outputs before timing. max_abs_difference was 0 for both recorded runs.

These are shape- and hardware-specific kernel measurements, not end-to-end model throughput. Arm k=4 deliberately keeps the scalar mask walk because the SIMD table path did not improve latency at this production-sized shape.

Verification

  • make -C kernels/bit_1_58/cpu clean all
  • python -m pytest -q tests/test_ternary_rsr.py (188 passed)
  • python -m py_compile benchmarking/bit_1_58/bench_v3_3_simd.py
  • git diff --check

The repository Linux workflow will additionally rebuild both CPU kernel families and run the complete test suite on AVX2.

@MohsenDehghankar
MohsenDehghankar merged commit 99bcccd into UIC-InDeXLab:main Aug 10, 2026
1 check passed
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