How llama.cpp GEMV Kernel Works on CUDA
During the decoding stage of an LLM, many of its linear layers multiply a weight matrix by one hidden-state vector. This matrix-vector multiplication, commonly known as GEMV, is used in cases where inputs are not batched (batch-size-one decoding, generally used in on-device/local inference). While prompt processing (prefill) or batched serving can use matrix-matrix multiplication (GEMM), GEMV kernels are used for on-device inference, and optimizing them can help bring down the latency for generating every single token. In this article, we study the FP16 CUDA GEMV kernel used in llama.cpp (we are using a simplified implementation derived from this pinned version’s FP16 weight path).
Let’s take a brief moment to go through the GEMV kernel equation. For batch-size-one decoding, we only process a single token at a time, so this single token gets converted into an embedding of shape \([\mathrm{d}]\) and gets passed through the attention and MLP layers. Given an intermediate input / hidden state \(x\) of shape \([\mathrm{d}]\), we have to multiply it with a weight \(W\) (let’s assume its shape is \([\mathrm{M}, \mathrm{d}]\)) to get a result \(y\) of shape \([\mathrm{M}]\). So for each element in \(y\), we need to calculate:
\[y_i = \sum_j W_{i,j} x_j\]For the single-vector path we are considering, the GEMV kernel launches a single block for processing each element of y. So if we are multiplying an x vector of shape 4096 with a W of shape 14336 x 4096, the kernel would launch 14336 blocks, with each block doing the dot product. The weights are stored in row-major order, and each block works on a single row.
We are using a simplified version of the llama.cpp GEMV FP16 kernel here. For simplicity, this version uses a fixed block size of 256 threads, uses FP16 for both the weights and the hidden state, and uses FP32 for accumulation and output.
Implementation note: The original kernel can dynamically select block size, stores the hidden state in FP32, and can choose between FP16 or FP32 accumulation. We omit those additional paths to focus more on the kernel’s core computation and reduction pattern.
Now, let’s go through our kernel: (you can find the code in this Github repo)
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
constexpr int kBlockSize = 256;
constexpr int kWarpSize = 32;
__inline__ __device__ float warp_reduce_sum(float value) {
for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
value += __shfl_xor_sync(0xffffffff, value, offset, kWarpSize);
}
return value;
}
template <int BlockSize>
__global__ void dense_half2_kernel(
const half * __restrict__ W, const half * __restrict__ x,
float * __restrict__ y,
int out_features, int in_features) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
if (row >= out_features) {
return;
}
const half *row_ptr = W + static_cast<size_t>(row) * in_features;
const half2 *row_ptr2 = reinterpret_cast<const half2 *>(row_ptr);
const half2 *x2 = reinterpret_cast<const half2 *>(x);
float sum = 0.0f;
for (int col2 = tid; col2 < in_features / 2; col2 += BlockSize) {
const float2 w = __half22float2(row_ptr2[col2]);
const float2 v = __half22float2(x2[col2]);
sum += w.x * v.x + w.y * v.y;
}
sum = warp_reduce_sum(sum);
constexpr int kNumWarps = BlockSize / kWarpSize;
__shared__ float warp_sums[kNumWarps];
const int lane = tid & (kWarpSize - 1);
const int warp_id = tid / kWarpSize;
if (lane == 0) {
warp_sums[warp_id] = sum;
}
__syncthreads();
if (warp_id == 0) {
float block_sum = lane < kNumWarps ? warp_sums[lane] : 0.0f;
block_sum = warp_reduce_sum(block_sum);
if (lane == 0) {
y[row] = block_sum;
}
}
}
Essentially, each thread loads elements and accumulates a partial dot product. Within each warp, these partial values are reduced, and then the block reduces the warp-reduced sums into the final row result. Now we are going to have a closer look at the operations.
A quick explanation of warps: a warp is a group of adjacent threads (threads with consecutive thread IDs) that the GPU schedules together. We can call the threads in the same warp lanes. Lanes in a warp can directly read a value held by another lane using registers, without using shared or global memory. This makes some operations like reduction very efficient.
Let’s focus on the loop internal operations first:
26
27
28
29
30
31
float sum = 0.0f;
for (int col2 = tid; col2 < in_features / 2; col2 += BlockSize) {
const float2 w = __half22float2(row_ptr2[col2]);
const float2 v = __half22float2(x2[col2]);
sum += w.x * v.x + w.y * v.y;
}
Instead of processing a single weight * input per iteration in the loop, we process two weight * input values. After getting the row and tid on lines 16-17, we convert the half pointer (which is FP16) into half2 pointers (lines 23-24) so that we can use the __half22float2 operation that would convert a half2 into a float2. Using half2 reduces the number of loop iterations on line 27 (compared to converting each half value into float individually) and allows four-byte vectorized loads.
Since we are processing 2 values per iteration, for a 4,096-wide row, there are 2,048 half2 values. Since we have 256 threads in a block, the loop assigns the first 256 positions to the 256 threads:
27
for (int col2 = tid; col2 < in_features / 2; col2 += BlockSize)
So, on the first iteration, thread 0 reads position 0, thread 1 reads position 1, and so on through thread 255. Within one warp, 32 neighboring threads read 32 neighboring half2 values. This is the coalesced access pattern the kernel wants for both the weight row and input vector. Coalesced reads let the hardware combine a warp’s weight loads into as few memory transactions as possible, which is what matters most here since streaming the weight matrix tend to dominate the kernel’s memory traffic.
On the next iteration, every thread advances by 256 positions. The block then covers the next 256-half2 tile, without overlap. Repeating that pattern covers all 2,048 positions.
Implementation note: For simplicity, we kept hidden state size even (4096) and assumed four-byte-aligned FP16 buffers.
After the load loop, every thread has only part of the dot product. The kernel reduces those 256 values in two stages. First, we use the warp_reduce_sum function here, which repeatedly exchanges and adds values between threads in the same warp using __shfl_xor_sync.
This operation reduces within each 32-thread warp. __shfl_xor_sync moves values between threads in a warp through registers, which means this operation needs neither global memory nor shared memory (for more details, check this guide). After doing the warp-wise reduction using warp_reduce_sum, we have 256 / 32 or 8 warp-wise sums, the result of which is stored at the threads of each warp. So we need to collect these results and sum them to get the final result. For that, we allocate a shared-memory array and collect the warp-wise sums in the shared memory. A __syncthreads() becomes required here so that the collection is done across all warps before we proceed to the next operations.
32
33
34
35
36
37
38
39
40
41
sum = warp_reduce_sum(sum);
constexpr int kNumWarps = BlockSize / kWarpSize;
__shared__ float warp_sums[kNumWarps];
const int lane = tid & (kWarpSize - 1);
const int warp_id = tid / kWarpSize;
if (lane == 0) {
warp_sums[warp_id] = sum;
}
__syncthreads();
Finally, we load these shared memory values into the registers for warp 0, and perform another warp-wise reduction to get the final sum. We want to use the efficient warp_reduce_sum here, but for that we want to ensure that all the other threads (or lanes) in the warp do not contain values that might affect our reduction. So for all those lanes in the 32-lane warp, we set the value of block_sum to 0.0f and finally perform another warp reduction. (For this specific example, lanes 0 - 7 contain the actual values, the remaining lanes contain 0.0f)
43
44
45
46
47
48
49
if (warp_id == 0) {
float block_sum = lane < kNumWarps ? warp_sums[lane] : 0.0f;
block_sum = warp_reduce_sum(block_sum);
if (lane == 0) {
y[row] = block_sum;
}
}
One important thing to note here: the GPU needs to stream the large weights from global memory and multiply them by the hidden state vector. For example: the weights can be as large as 14336x4096 for MLP layers, while the hidden state would contain 4096 elements in this case. Streaming the weights from memory can introduce considerable latency, which is why GEMV is generally memory-bound. In large matrix-matrix multiplication or GEMM, you can load a weight tile into on-chip memory to multiply with multiple hidden states and amortize this memory traffic, but GEMV needs to pay this heavy tax for each hidden state.