Skip to content

ggml-cpu : build TQ with SIMD and use the standard vec op - #4

Open
mverrilli wants to merge 1 commit into
elusznik:turboquant-cpu-tbq-prfrom
mverrilli:pr21089-fix-simd
Open

ggml-cpu : build TQ with SIMD and use the standard vec op#4
mverrilli wants to merge 1 commit into
elusznik:turboquant-cpu-tbq-prfrom
mverrilli:pr21089-fix-simd

Conversation

@mverrilli

Copy link
Copy Markdown

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

  • ~7x faster TBQ dequant on the KV-cache read (i5-8600)
  • Quantize side also vectorized
  • Output unchanged.
  • Checked non-TBQ types, unaffected (same to_float pointer, ppl the same)

Requirements

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: Yes. Diagnose build-target/arch-flags root cause, simple refactor.

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).

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread ggml/src/ggml-turboq-impl.h
Comment thread ggml/src/ggml-turboq-impl.h
Comment on lines +165 to +181
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;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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;

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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.

Comment thread ggml/src/ggml-turboq-impl.h
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant