perf(sse41): vectorize grouped-int4 GEMV - #1286
Conversation
3f817f7 to
d007179
Compare
|
@OPS-NeoRetro i'll try to stay on top of keeping this up to date with main but would love to get it merged! LMK if there's anything else i can do |
9dcbb36 to
e5ee024
Compare
Register the extracted GEMV and SSE4.1 headers for their engine, test, and adapter consumers. This satisfies the current dev dependency gate and makes header edits rebuild affected binaries.
e5ee024 to
a3ec030
Compare
OPS-NeoRetro
left a comment
OPS-NeoRetro
left a comment
There was a problem hiding this comment.
Good job, but you have to organize the code more logically. Seem like I'll have to or someone must open a PR to make the Makefiles smaller in size. @cameron, please re-test this rebase.
| uint32_t packed=(uint32_t)q4[(int64_t)(o+0)*rb+byte] | ||
| |(uint32_t)q4[(int64_t)(o+1)*rb+byte]<<8 | ||
| |(uint32_t)q4[(int64_t)(o+2)*rb+byte]<<16 | ||
| |(uint32_t)q4[(int64_t)(o+3)*rb+byte]<<24; |
There was a problem hiding this comment.
I think it's possible to compose this only out of SSE4.1 intrinsics
There was a problem hiding this comment.
Addressed in 3b24221. The row loader now uses _mm_cvtsi32_si128 followed by three _mm_insert_epi8 calls, reading exactly one byte per row. Added exhaustive packed-byte checks across all four lanes and one-byte-row coverage; the forced-SSE4.1 exactness and ASan/UBSan runs pass.
|
|
||
| #if defined(__SSE4_1__) | ||
| #include <immintrin.h> | ||
| #include "sse41_kernels.h" | ||
| #endif | ||
|
|
||
| #if defined(__SSE4_1__) && !defined(__AVX2__) | ||
| /* Load one packed byte from each of four output rows, then unpack their low and | ||
| * high offset nibbles into four f32 lanes. Keeping independent output rows in | ||
| * the lanes preserves the scalar operation order within every row. */ | ||
| static inline void colibri_i4_rows4(const uint8_t *q4,int rb,int o,int byte, | ||
| __m128 *lo,__m128 *hi){ | ||
| const __m128i m4=_mm_set1_epi8(0x0F), b8=_mm_set1_epi8(8); | ||
| uint32_t packed=(uint32_t)q4[(int64_t)(o+0)*rb+byte] | ||
| |(uint32_t)q4[(int64_t)(o+1)*rb+byte]<<8 | ||
| |(uint32_t)q4[(int64_t)(o+2)*rb+byte]<<16 | ||
| |(uint32_t)q4[(int64_t)(o+3)*rb+byte]<<24; | ||
| __m128i by=_mm_cvtsi32_si128((int)packed); | ||
| __m128i qlo=_mm_sub_epi8(_mm_and_si128(by,m4),b8); | ||
| __m128i qhi=_mm_sub_epi8(_mm_and_si128(_mm_srli_epi16(by,4),m4),b8); | ||
| *lo=_mm_cvtepi32_ps(_mm_cvtepi8_epi32(qlo)); | ||
| *hi=_mm_cvtepi32_ps(_mm_cvtepi8_epi32(qhi)); | ||
| } | ||
|
|
||
| static inline __m128 colibri_f32_rows4(const float *p,int stride,int o,int i){ | ||
| return _mm_set_ps(p[(int64_t)(o+3)*stride+i],p[(int64_t)(o+2)*stride+i], | ||
| p[(int64_t)(o+1)*stride+i],p[(int64_t)(o+0)*stride+i]); | ||
| } | ||
|
|
||
| /* Process four output rows at once without a horizontal reduction. Each lane | ||
| * uses the scalar kernel's pair sum, scale multiply, and accumulator add in | ||
| * the same order, so the result can remain byte-identical on pre-FMA CPUs. */ | ||
| static void matmul_i4_grouped_sse41_rows4(float *y,const float *x, | ||
| const uint8_t *q4,const float *scale, | ||
| int S,int I,int O,int gs,int rb,int ng, | ||
| int o4){ | ||
| #pragma omp parallel for schedule(static) | ||
| for(int o=0;o<o4;o+=4){ | ||
| for(int s=0;s<S;s++){ | ||
| const float *xs=x+(int64_t)s*I; __m128 a=_mm_setzero_ps(); | ||
| for(int g=0;g*gs<I;g++){ | ||
| int base=g*gs,end=base+gs; if(end>I) end=I; | ||
| __m128 sc=colibri_f32_rows4(scale,ng,o,g); int i=base; | ||
| for(;i+1<end;i+=2){ | ||
| __m128 lo,hi; colibri_i4_rows4(q4,rb,o,i>>1,&lo,&hi); | ||
| __m128 pair=_mm_add_ps(_mm_mul_ps(_mm_set1_ps(xs[i]),lo), | ||
| _mm_mul_ps(_mm_set1_ps(xs[i+1]),hi)); | ||
| a=_mm_add_ps(a,_mm_mul_ps(pair,sc)); | ||
| } | ||
| if(i<end){ | ||
| __m128 lo,hi; colibri_i4_rows4(q4,rb,o,i>>1,&lo,&hi); (void)hi; | ||
| a=_mm_add_ps(a,_mm_mul_ps(_mm_mul_ps(_mm_set1_ps(xs[i]),lo),sc)); | ||
| } | ||
| } | ||
| colibri_sse41_storeu_ps(y+(int64_t)s*O+o,a); | ||
| } | ||
| } | ||
| } | ||
| #endif |
There was a problem hiding this comment.
Why you don't put this in sse41_kernels.h because sse41_kernels.h contain SSE4.1 codepaths?
There was a problem hiding this comment.
Moved the row helpers and both the single-matrix and fused gate/up kernels into c/sse41_kernels.h in 3b24221. quant.h and colibri.c keep their dispatch and scalar fallbacks. The shared header includes its own intrinsic and integer-type headers, and the arithmetic order is preserved.
…get.h (JustVugg#1712) next to this branch's header
|
Merged current dev into your branch (Makefile prerequisite lines only: serve_budget.h from #1712, the build-flags stamp from #1707, sse41_kernels.h kept). On the merge test_i4_grouped passes both in the default build and compiled with -msse4.1 and no AVX2, which is the path this PR adds, and the kimi/olmoe framing tests pass. Merging once CI is green, into 1.12.1. |
…ot.h and the new CPU-only test rules next to this branch's sse41_kernels.h and SSE4.1 gates
Use SSE4.1 byte inserts instead of scalar shift/OR packing. Move the single and fused grouped-int4 kernels beside their shared row helpers, with the existing dispatch and arithmetic order retained. Exercise every packed byte per lane, allocation-boundary loads, and fused output byte identity.
|
Addressed both inline comments: intrinsics for byte packing, and both grouped-int4 SSE4.1 kernels moved into the shared header. |
3 participants
Note
Authorship: I started this from performance work in my agent-driven local inference benchmarking project, which I've been using for months (https://sloptimized.ccb.gg/) to integrate new models into my workflow. The implementation, tests, benchmarks, and most of this write-up are ~95% agent-authored.
Summary
Builds on #1239, now merged into
dev.This adds an SSE4.1 path for grouped-int4 W4A32 GEMV and fused gate/up GEMV on pre-AVX2 CPUs. It vectorizes four independent output rows at a time, shares activation broadcasts in the fused path, and retains the scalar path for odd group sizes and output-row tails.
The key property is exactness. An earlier prototype vectorized along K and needed a horizontal reduction, which changed the floating-point addition tree. This version vectorizes across output rows. Each SIMD lane performs the same pair-sum, group-scale multiply, and accumulator-add sequence as one scalar row, with no horizontal reduction.
The branch also adds forced-SSE exactness tests at
-O1,-O3, and-O3 -ffp-contract=off. The SSE4.1 row helpers and both grouped-int4 kernels live together insse41_kernels.h; byte packing uses SSE4.1 insert intrinsics.Scope and dependency
__SSE4_1__ && !__AVX2__.Rebased onto
devafter #1239 merged. This PR now contains only the grouped-int4 optimization, its exactness tests, and header dependency updates. The SSE4.1 foundation and Olmoe architecture guard are already upstream. The rebase preserves the current sharedidot.hkernel layout and upstream oracle header dependencies.Performance
All measurements use this branch at
2d04e3cee4abf103edca801371f8f6aa38e4412on a Dell R720 with 2x Xeon E5-2680 v2 (Ivy Bridge), 20 physical/40 logical cores, two NUMA nodes, and 354.9 GiB RAM.Isolated fused kernel
Production shape
S=1, I=6144, O=2048, gs=64, fixed seed0x902a371d, 100 fused gate/up calls per sample, five samples per path, interleaved NUMA allocation. Scalar and SSE4.1 order was alternated.Every scalar and SSE4.1 sample produced checksum
b5253d3717f0ae12. The benchmark reference is an independent copy of the scalar f32 operation tree, and the check usesmemcmprather than a numeric tolerance.Full GLM-5.2 check
As a secondary end-to-end check, I ran the grouped-int4 GLM-5.2 artifact on the same host with 2x V100 32 GB, 40 OpenMP threads, one serial client, one KV slot, a full 17-token prefix-cache hit, and five 32-token samples per build.
This is a 1.797x median decode gain and 41.6% lower median TTFT. All ten responses had identical text and canonical-token hashes, and all profiles reported zero expert disk service and zero felt expert-read wait.
The SSE4.1 run had a wide 81.7% sample spread, so I consider the isolated kernel benchmark the primary performance evidence. A separate matched non-streaming run also favored SSE4.1 by 1.612x (1.172634 versus 0.727504 tok/s), but it is retained as provisional because it did not record TTFT.
Build identities and benchmark commands
Builds used GCC 15.2.0, CUDA 12.9,
ARCH=x86-64,CUDA_ARCH=sm_70,-mno-avx2, and-mno-fma:/nix/store/iyfsa6rk1qlh9fximhy56i7zyr26r7c2-colibri-pr1286-sse41-1.9.0-2d04e3c/nix/store/h0cllrcpc52isdf530c6qrzr8m0arhx7-colibri-pr1286-scalar-1.9.0-2d04e3cKernel benchmark:
COLIBRI_KERNEL_BENCH_PACKAGE=/nix/store/iyfsa6rk1qlh9fximhy56i7zyr26r7c2-colibri-pr1286-sse41-1.9.0-2d04e3c \ COLIBRI_KERNEL_BENCH_THREADS=20 ./colibri/benchmark-grouped-int4-cpu.sh # Repeat with COLIBRI_KERNEL_BENCH_THREADS=40.Full-model controls were identical except for
COLIBRI_BENCH_PACKAGE:COLIBRI_BENCH_PREDICT=32 COLIBRI_BENCH_REPETITIONS=5 \ COLIBRI_BENCH_OMP_THREADS=40 COLIBRI_BENCH_PIPE=0 \ COLIBRI_BENCH_CUDA_DENSE=0 COLIBRI_BENCH_GPUS=0,1 \ COLIBRI_BENCH_CUDA_EXPERT_GB=58 COLIBRI_BENCH_DIRECT=0 \ COLIBRI_BENCH_CACHE_ROUTE=0 COLIBRI_BENCH_MTP=0 \ COLIBRI_BENCH_DRAFT=0 COLIBRI_BENCH_DROP_CACHES=1 \ COLIBRI_BENCH_PACKAGE="$package" ./colibri/benchmark-glm52-v100.shComplete retained records are in the public benchmark catalog under
r4v20,r4v40,r4fs0, andr4fv0. Rejected harness attempts are retained as invalid rather than omitted.Validation
Reviewer follow-up
3b242215, 2026-09-23, on top of maintainer merge621ae67b:nix shell nixpkgs#gcc14 --command make -C c -j8 checkpassed; complete C suite passed; Python suite ran 1,312 tests with 145 skippedForced SSE4.1 tests at
-O1,-O3, and-O3 -ffp-contract=off: scalar-exact output for aligned, partial-group, odd-I, single-row, output-tail, and one-byte-row shapesThe row unpacker checks all 256 packed-byte values in every lane, with tight and strided rows; fused gate/up output is compared with
memcmpin forced-SSE4.1 buildsTargeted ASan/UBSan: pass for single and fused grouped-int4 kernels, forced SSE4.1 with AVX/AVX2/FMA disabled, including allocation-boundary byte loads
Standalone shared-header builds pass for baseline x86-64 and SSE4.1; GCC 14 assembly uses byte-memory
pinsrbinstructions for row insertionPrior platform and dependency checks passed: AArch64 Linux and arm64 macOS exclude forced-x86 tests; x86-64 Linux and Windows include all six; header edits rebuild affected engine/test targets
Fresh merge check immediately before the follow-up push:
git merge-tree --write-tree HEAD upstream/devpassed againstdevat7a5bdee8make -C c checkCUDA changes were tested with
make -C c cuda-test(not applicable: no CUDA changes)Performance claims include hardware, commands, and repeatable measurements
Compatibility