ggml-cpu : build TQ with SIMD and use the standard vec op - #4
Conversation
Was compiled in ggml-base (no arch flags) so it ran scalar. Move to a ggml-cpu TU and use the standard vec op (ggml_vec_dot_f32).
There was a problem hiding this comment.
Code Review
This pull request refactors the TurboQuant implementation to support arch-built SIMD dequantization on the CPU backend. It introduces a shared header ggml-turboq-impl.h for rotation and codebook routines, adds a SIMD-specific turboq.c for the CPU backend, and updates the CPU operators to prefer SIMD dequantization when available. The code review feedback highlights opportunities to eliminate dynamic memory allocations (malloc/free) and prevent potential thread-local memory leaks by replacing dynamically allocated thread-local pointers and arrays with stack-allocated or static thread-local arrays of fixed size.
| static TURBOQ_TLS float * tl_Q = NULL; | ||
| static TURBOQ_TLS float * tl_Q_row = NULL; | ||
| static TURBOQ_TLS int64_t tl_Q_dim = 0; | ||
| static TURBOQ_TLS uint64_t tl_Q_seed = 0; | ||
|
|
||
| static const float * turboq_get_rotation(int64_t d, uint64_t seed) { | ||
| if (tl_Q != NULL && tl_Q_dim == d && tl_Q_seed == seed) { | ||
| return tl_Q; | ||
| } | ||
| // Regenerate | ||
| free(tl_Q); | ||
| free(tl_Q_row); | ||
| tl_Q = (float *)malloc(d * d * sizeof(float)); | ||
| tl_Q_row = (float *)malloc(d * d * sizeof(float)); | ||
| tl_Q_dim = d; | ||
| tl_Q_seed = seed; | ||
|
|
There was a problem hiding this comment.
Thread-local pointers allocated via malloc are not automatically freed when a thread exits in C, leading to memory leaks if threads are dynamically created and destroyed. Since d is always TURBOQ_KV_DIM (128), we can use static thread-local arrays of fixed size TURBOQ_KV_DIM * TURBOQ_KV_DIM to completely avoid dynamic allocation and prevent memory leaks on thread exit.
static TURBOQ_TLS float tl_Q[TURBOQ_KV_DIM * TURBOQ_KV_DIM];
static TURBOQ_TLS float tl_Q_row[TURBOQ_KV_DIM * TURBOQ_KV_DIM];
static TURBOQ_TLS int64_t tl_Q_dim = 0;
static TURBOQ_TLS uint64_t tl_Q_seed = 0;
static TURBOQ_TLS bool tl_Q_initialized = false;
static const float * turboq_get_rotation(int64_t d, uint64_t seed) {
assert(d <= TURBOQ_KV_DIM);
if (tl_Q_initialized && tl_Q_dim == d && tl_Q_seed == seed) {
return tl_Q;
}
// Regenerate
tl_Q_dim = d;
tl_Q_seed = seed;
tl_Q_initialized = true;There was a problem hiding this comment.
@elusznik I think perhaps build-once, never-free, shared-global to address. Basically one shared Q built once standard mutex-guarded init, then read lock-free. Leaving this out since it was pre-existing and I don't want to introduce a threading change in the mix.
Was compiled in ggml-base (no arch flags) so it ran scalar. Move to a ggml-cpu TU and use the standard vec op (ggml_vec_dot_f32).
Overview
This moves the rotation + codec into a ggml-cpu translation unit (built with ARCH_FLAGS) and replaces the hand-rolled AVX2 matvec with ggml_vec_dot_f32, so it vectorizes on every supported CPU arch.
Additional information
Requirements