TurboQuant Attention in CUDA: Transform the Query, Reuse the Work

I recently solved a CUDA challenge about attention scores from a TurboQuant-compressed key cache. The useful shortcut came before the kernel: move the expensive work onto the queries and reuse it.

An attention score is a dot product between a query, the current token's question, and a cached key representing an earlier token. This challenge asks for raw scores only: no softmax, value lookup, or $1/\sqrt{D}$ scaling.

The keys arrived in pieces

The keys arrive as codebook indices. Each points to a representative number, called a centroid. Before quantization, the keys were rotated by a matrix $\Pi$.

There is also a correction for quantization loss: signs $\sigma$, a magnitude $\gamma$, and a projection matrix S_mat. These come from the residual-correction idea in the TurboQuant paper.

All of this is prepared before the scoring kernels run. They read compressed keys; they do not compress them. Here, B counts queries, S counts cached tokens, D is the vector length, and C is the codebook size. The output has shape [B, S], one score for every query and cached token.

What does gamma store? Subtract the centroid-based key approximation from the original key to get the residual, then take its length. For example, a residual [0.3, 0.15] has $\gamma=\sqrt{0.3^2+0.15^2}\approx0.3354$. The compressor stores this magnitude and the signs of the projected residual. gamma measures the leftover error, not the length of the query or original key.

The obvious approach is to reconstruct every key, then take the dot product. But I only need the score. Do I need to build the key at all?

Move the transforms to the query

Let $c_s$ be the centroid vector for cached token $s$. Writing vectors as columns, the reconstructed key is:

The transpose, $T$, exchanges a matrix's rows and columns. Moving these matrices across the dot product gives:

Now I compute $A_b=\Pi q_b$ and $T_b=S_{\mathrm{mat}}q_b$ once per query. Every cached key reuses them. The score is unchanged apart from floating-point rounding; the original quantization loss remains.

A tiny example before CUDA

Take one query, $q=[1,2]$, two cached tokens, and this codebook. That makes B=1, S=2, D=2, and C=4.

index:      0      1      2      3
centroid: -0.75  -0.25   0.25   0.75

Use an identity rotation and, just for easy hand arithmetic, the toy projection below. This matrix is an illustration, not a sampled Gaussian projection.

Pi    = [[1, 0], [0,  1]]
S_mat = [[1, 1], [1, -1]]
A = Pi × q    = [1,  2]
T = S_mat × q = [3, -1]

Suppose the compressed inputs contain:

For token 0, the centroid score is 1 × (-0.75) + 2 × (-0.25) = -1.25. The residual dot product is 3 × 1 + (-1) × 1 = 2. Scale that by $\sqrt{\pi/2}/2$ and gamma=0.2, giving a correction of about 0.250663.

Add the two parts: -1.25 + 0.250663 = -0.999337. For token 1, the centroid score is 1 × 0.25 + 2 × 0.75 = 1.75, and the residual dot product is 3 × 1 + (-1) × (-1) = 4. With gamma=0.4, we get:

The output is [[-0.999337, 2.752651]], rounded to six decimals. Token 1 has the higher score for this query. These are scores, not probabilities.

Two kernels, two jobs

The first kernel assigns one thread to each (query, coordinate) pair. It uses b = idx / D and d = idx % D, then loops over the query to calculate one entry of A and one of T.

The second uses b = idx / S and s = idx % S to assign one thread to each (query, cached token) pair. Each thread loops over all D coordinates and writes one complete score. In our tiny example, each kernel launches a block of 256 threads, but only the first two do useful work; the rest exit at the bounds check.

The scoring loop is the heart of the solution:

unsigned char code = K_idx[s * D + d];
mse_sum += A[b * D + d] * codebook[code];
residual_sum += T[b * D + d]
              * (float)qjl_signs[s * D + d];

After the loop, it writes:

scores[b * S + s] = mse_sum
    + (1.2533141373155001f / D) * gamma[s] * residual_sum;

Here, 1.253314... is $\sqrt{\pi/2}$. The codebook size C is unused because the challenge guarantees valid indices.

What the rearrangement saves

At the challenge's performance shape, there are 32 queries and 32,768 cached keys, each with 128 coordinates. Transforming 32 queries is much less matrix work than reconstructing 32,768 keys.

With 256 threads per block, the first kernel launches 16 blocks for 4,096 query coordinates. The second launches 4,096 blocks for 1,048,576 scores. The GPU schedules those blocks across its available streaming multiprocessors.

The two temporary query arrays occupy 32 KiB, compared with 16 MiB for a materialized FP32 key matrix. This comparison excludes the inputs and output scores. The challenge stores each index and sign in a byte; my kernel does not add bit packing.

This is an algebraic saving, not a measured GPU speedup. Each scoring thread still loops over all coordinates, and solve allocates temporary arrays on every call.

What I liked about this problem was that a small change in the equation removed an entire intermediate array. Before making a reconstruction faster, it helped to ask whether I needed to reconstruct it at all.

Full CUDA solution

Here are both kernels and the solve entry point together. Use Copy code to copy the complete solution.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

#include <cuda_runtime.h>
#include <math.h>

// Kernel 1: Precompute transformed queries, reused for every cached token.
// A[b, d] = sum_j Pi[d, j]    * Q[b, j]
// T[b, d] = sum_j S_mat[d, j] * Q[b, j]
__global__ void preprocess_queries(
    const float* Q,
    const float* Pi,
    const float* S_mat,
    float* A,
    float* T,
    int B,
    int D
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * D;
    if (idx >= total) return;

    int b = idx / D;
    int d = idx % D;
    float a_sum = 0.0f;
    float t_sum = 0.0f;
    for (int j = 0; j < D; ++j) {
        float q = Q[b * D + j];
        a_sum += Pi[d * D + j] * q;
        t_sum += S_mat[d * D + j] * q;
    }
    A[b * D + d] = a_sum;
    T[b * D + d] = t_sum;
}

// Kernel 2: One raw attention score per thread.
// No softmax and no 1/sqrt(D) scaling.
__global__ void attention_scores_kernel(
    const float* A,
    const float* T,
    const unsigned char* K_idx,
    const signed char* qjl_signs,
    const float* gamma,
    const float* codebook,
    float* scores,
    int B,
    int S,
    int D
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * S;
    if (idx >= total) return;

    int b = idx / S;
    int s = idx % S;
    float mse_sum = 0.0f;
    float residual_sum = 0.0f;
    for (int d = 0; d < D; ++d) {
        unsigned char code = K_idx[s * D + d];
        float centroid = codebook[code];
        mse_sum += A[b * D + d] * centroid;
        float sign = (float)qjl_signs[s * D + d];
        residual_sum += T[b * D + d] * sign;
    }
    const float correction_scale = 1.2533141373155001f / (float)D;
    float residual = correction_scale * gamma[s] * residual_sum;
    scores[idx] = mse_sum + residual;
}

// Required entry point. All pointers refer to device memory.
extern "C" void solve(
    const float* Q,
    const unsigned char* K_idx,
    const signed char* qjl_signs,
    const float* gamma,
    const float* Pi,
    const float* S_mat,
    const float* codebook,
    float* scores,
    int B,
    int S,
    int D,
    int C
) {
    (void)C;
    float* A = nullptr;
    float* T = nullptr;
    size_t bytes = (size_t)B * D * sizeof(float);
    cudaMalloc(&A, bytes);
    cudaMalloc(&T, bytes);
    constexpr int THREADS = 256;
    int preprocess_blocks = (B * D + THREADS - 1) / THREADS;
    preprocess_queries<<<preprocess_blocks, THREADS>>>(Q, Pi, S_mat, A, T, B, D);
    int total_scores = B * S;
    int score_blocks = (total_scores + THREADS - 1) / THREADS;
    attention_scores_kernel<<<score_blocks, THREADS>>>(
        A, T, K_idx, qjl_signs, gamma, codebook, scores, B, S, D
    );
    cudaFree(A);
    cudaFree(T);
}

Download the .cu file · TurboQuant paper