SIMD Dot-Product Kernels for Quantized Embeddings

Computation Graphs & ML Infrastructure advanced Building with AI ~35 min read

Overview

Embedding retrieval often begins with an apparently simple operation: compute a similarity score between a query vector and many stored vectors, then keep the highest-scoring candidates. When vectors are normalized, dot product ranking is equivalent to cosine similarity ranking, so the hot path becomes a tight multiply-accumulate loop. If the corpus is small or freshness matters more than approximate indexing, exhaustive scoring can be simpler and more explainable than building a nearest-neighbor index for every update.

The problem is that exhaustive scoring is brutally sensitive to memory bandwidth and instruction throughput. A 1536-dimensional float32 embedding uses 6 KiB. Millions of vectors quickly turn into many GiB of resident memory, and every query must stream through that memory before the language model can receive context. Quantizing embeddings to int8 cuts storage by 4x, but it also changes the arithmetic: the kernel now needs to widen bytes, multiply, accumulate into wider lanes, reduce horizontally, and preserve enough numerical fidelity for ranking.

This topic teaches the implementation pattern behind a production-style SIMD dot-product kernel for quantized embeddings. The progression starts with a scalar Go loop, removes loop-carried dependencies and bounds checks, switches from float32 to int8, then moves the hot loop into x86 assembly using AVX2 and AVX-512 VNNI instructions. The same shape applies to vector search systems, rerank prefilters, compressed feature stores, recommendation retrieval, and any service that scores many fixed-length byte vectors on CPU.

Core Concepts

Dot Product as Retrieval Primitive

A dot product multiplies matching vector elements and sums the products. For normalized embeddings, it gives the same ordering as cosine similarity, which makes it a practical scoring primitive for exact vector retrieval. The retrieval loop scores every candidate vector against the query, maintains a small top-k heap or selection buffer, and returns the highest-scoring chunks as context.

Quantized Embeddings

Integer quantization stores each component in fewer bits, commonly one signed byte per component. This lowers memory footprint and improves cache residency, but the kernel must accumulate in a larger integer type because byte products overflow quickly. In practice, an int8 dot product commonly multiplies bytes, widens intermediate products, and accumulates into int32 lanes before applying any scale correction outside the innermost loop.

SIMD Widen-Multiply-Accumulate

SIMD turns repeated scalar arithmetic into vector arithmetic over many elements per instruction. For quantized dot products, the useful instructions are not just byte loads. The key is combining sign extension, pairwise multiply-add, vector accumulation, and horizontal reduction. AVX2 can process a block by widening bytes to words and using packed multiply-add instructions. AVX-512 VNNI goes further by fusing byte dot products directly into 32-bit accumulators.

Key Concept: Keep the Hot Loop Mechanically Small

The fastest dot-product kernels do very little per block: load contiguous bytes, widen or compensate signedness, multiply-add into vector accumulators, advance pointers, and loop. Length checks, feature detection, scaling, tail handling, and top-k maintenance should sit outside the inner block loop whenever possible.

How It Works

From Scalar Loop to Bounds-Check-Eliminated Loop

The baseline Go implementation is easy to read but has loop-carried dependencies and runtime bounds checks. Unrolling the loop creates independent multiply operations that the CPU can pipeline. Proving equal lengths and slicing to a fixed block can help the compiler eliminate repeated bounds checks in memory-safe languages.

func DotInt8Scalar(a, b []int8) int32 {
    if len(a) != len(b) {
        panic("vectors must have equal length")
    }

    var sum int32
    for i := 0; i < len(a); i += 4 {
        // Full-slice expressions restrict capacity and help the compiler reason
        // about all four indexed accesses as one checked block.
        aa := a[i : i+4 : i+4]
        bb := b[i : i+4 : i+4]
        sum += int32(aa[0]) * int32(bb[0])
        sum += int32(aa[1]) * int32(bb[1])
        sum += int32(aa[2]) * int32(bb[2])
        sum += int32(aa[3]) * int32(bb[3])
    }
    return sum
}

AVX2 Kernel Shape

An AVX2 byte dot-product loop can process a fixed block per iteration. One practical shape sign-extends 16 bytes from each vector into 16-bit lanes, uses a multiply-add instruction to produce 32-bit partial sums, and accumulates those partial sums in a vector register. After the block loop, the vector accumulator is folded horizontally into one scalar result.

// Pseudo-assembly shape for an AVX2 int8 dot-product block loop.
// Names are illustrative; exact Go assembler spelling is architecture-specific.

zero ymm_acc
loop:
    if remaining < 16: goto reduce
    ymm_a16 = sign_extend_16_i8_to_i16(load_16_bytes(a_ptr))
    ymm_b16 = sign_extend_16_i8_to_i16(load_16_bytes(b_ptr))
    ymm_prod32 = multiply_adjacent_i16_and_add_to_i32(ymm_a16, ymm_b16)
    ymm_acc = add_i32(ymm_acc, ymm_prod32)
    a_ptr += 16
    b_ptr += 16
    remaining -= 16
    goto loop
reduce:
    return horizontal_sum_i32_lanes(ymm_acc)

AVX-512 VNNI Signedness Compensation

AVX-512 VNNI provides VPDPBUSD, which multiplies unsigned bytes from one operand by signed bytes from the other operand, groups four byte-products into each 32-bit lane, and accumulates. If both embedding vectors are stored as signed bytes, one operand must be shifted into unsigned range. The compensation identity is: compute (a + 128) * b, separately compute 128 * b, then subtract the overshoot from the final accumulator.

// Algebra used when a VNNI instruction expects one unsigned byte vector.
// For signed byte vectors a and b:
//
//     sum(a * b) == sum((a + 128) * b) - 128 * sum(b)
//
// A VNNI kernel can vectorize both terms:
//     acc_pos += vpdpbusd(a + 128, b)
//     acc_bias += vpdpbusd(128, b)
//     acc = acc_pos - acc_bias

Practical Applications

Exact Embedding Search for Fresh Corpora

Exact scoring is attractive when new vectors arrive continuously and index rebuild latency is unacceptable. A quantized SIMD kernel lets a service keep a simple appendable vector store, scan all active vectors, and return deterministic top-k results. This is useful for code search, per-tenant document stores, and small-to-medium corpora where operational simplicity beats approximate indexing complexity.

Candidate Prefiltering Before Reranking

Many retrieval stacks use a cheap first-stage scorer followed by a heavier reranker. A byte-vector dot-product kernel can serve as that first stage, generating a few hundred candidates quickly before invoking a cross-encoder, LLM judge, or graph-aware reranker. The kernel does not need perfect semantic fidelity; it needs stable recall at the candidate cutoff.

CPU-Only Retrieval in Edge or On-Prem Systems

Not every deployment can depend on a GPU or managed vector database. SIMD kernels let CPU-only systems support local retrieval with predictable binary distribution and no external service dependency. Go assembly with a scalar fallback is especially useful for single-binary services that need fast paths on modern x86 servers while still running on unsupported machines.

Implementation Patterns

Dispatch Once, Then Call a Stable Kernel

Do not branch on CPU capabilities inside the scoring loop. Bind the fastest safe implementation at initialization time, then expose one public function. Keep the scalar version as the oracle for correctness tests and as the fallback for unsupported architectures.

package dot

import "golang.org/x/sys/cpu"

var DotInt8 func(a, b []int8) int32 = dotInt8Scalar

func init() {
    switch {
    case cpu.X86.HasAVX512VNNI:
        DotInt8 = dotInt8VNNI
    case cpu.X86.HasAVX2:
        DotInt8 = dotInt8AVX2
    }
}

func ScoreAll(query []int8, vectors [][]int8, scores []int32) {
    for i, v := range vectors {
        scores[i] = DotInt8(query, v)
    }
}

Separate Kernel Contracts from Retrieval Policy

The assembly kernel should have a narrow contract: equal-length vectors, expected alignment or safe unaligned loads, fixed block multiple if required, and an integer dot-product result. Top-k selection, dequantization scales, filtering, authorization, and tracing belong in surrounding Go code. This separation keeps the assembly verifiable and limits the blast radius of architecture-specific code.

Test Against a Scalar Oracle

Every optimized kernel should be property-tested against the scalar implementation across vector lengths, signs, random data, all-zero inputs, maximum positive and negative bytes, and tails. For retrieval, also test rank stability after quantization by comparing top-k overlap against a float baseline on representative embeddings.

Performance Considerations

Common Pitfalls

Further Reading