Engineering note
Writing a high performant GEMV + dequant for Metal
Notes on writing a high performant quant linear kernel for Metal (hitting 90% of Metal's MBU).
What is a GEMV
A GEMV is one of the most common operations in LLMs, during the decode phase, each forward pass executes this operation hundreds of times. On a modern architecture like Qwen 3.5, each Decoder layer (regardless of GDN or FullAttention) has many projections (Q, K & V as an example) as well as MLPs processing big shapes. Mathematically the GEMV we will be covering is:
, , and
Because this ends up taking almost most of the execution time of each forward pass on decode, it is important to have high performant GEMV kernels to achieve the highest possible decode/generation tps (tokens per second).
Dequantization
In order to run bigger models on computers with low amounts of memory, models are quantized, so that each weight takes less space. There are metrics such as perplexity, which measure the overall degradation in a generative prediction confidence across a text corpus, or KL divergence, which directly quantifies token-by-token probability distribution drift between the original and quantized outputs. These and other benchmarks are important to choose how much of a degradation we tolerate, in order to reduce the total used memory.
UD-Q4_K_XL
There are many ways to quantize a model, some methods use the same quant type for all the layers in the model (i.e: FP8 or NVFP4), while other quantization models introduce a mixed quantization schema, where it uses a specific quant type for each different layer of the model. The UD-Q4_K_XL is a method introduced by Unsloth, specific for the Qwen3.5 series of models. It contains Q4_K, Q5_K, Q6_K, Q8_0 & IQ4_XS all across the model. Quantization methods
Fusing GEMV and Dequantization
The specific case I will be talking about is the fused Q4_K + GEMV, Q5_K and Q6_k are very similar (they only changes some scale and minimum storage) and Q8_0 doesn’t require dequantization, so it is similar to a performant GEMV.
What is a weight block?
A weight block is a group of weights represented using bytes, for the case of the quant we will cover, and , meaning bytes per weight (~4bit quant).
These weights are divided into 8 groups of 32, which share the same group scales and minimums. These scales and minimums (globals and per group) are used to dequantize each of the 256 weights in the following way:
The weight blocks span out across the dimension of the weight tensor, meaning that we have a total of weight blocks, for each layer’s tensor, these weight blocks can be processed in parallel and later be reduced if we tile across this K dimension.
SIMD & Threadgroups map to data
Inside metal’s backend a SIMD (NVIDIA warp) is a group of 32 lanes (NVIDIA threads) executing instructions with a SIMD control mechanism. For this strategy each SIMD will own a tile of the activation vector and a tile of the weights vector (transposed), producing a partial result of the output vector.
There will also be a further logical division which is dividing the 32 lanes in 4 groups of 8 called “cohorts”, each of these cohorts processes two full quant weight blocks (one from each column of ).
After tuning to get the best results, I found 64 lanes (2 SIMDs) per threadgroup, to be the best configuration possible. These SIMDs are scheduled to the same GPU core (NVIDIA SMs) so that they can share the threadgroup memory and synchronize using barriers, so tuning the amount of SIMDs per GPU core will determine the occupancy: many SIMDs inside a threadgroup may not fit in the register file, reducing occupancy, while low SIMDs may not be enough to hide latency when one SIMD is executing a memory operation (context switching to hide latency).
Lanes map to data
The analysis will be simplified to a single block per cohort, because processing the second block (from the next column) is the same processing with different indexing (done in a for loop)Each of the cohorts is logically subdivided into “teams” of 4 lanes, where each of them processes 4 of the 8 groups that make a weight block.
The way these teams are organized have its purpose in the way the scales and minimums for each group are stored. By loading two contiguous bytes of local scales & minimums, each team ends up owning the groups showed above. All 4 lanes in the team LOAD these local scales and minimums from GMEM to REG
Because we process each group (32 weights) in teams of 4 lanes, each lane will “own” a contiguous portion of 8 weights from each group that the team owns.
Compute
Since this is a fused dequant + GEMV kernel, we have to dequantize the weights owned by each lane using global and group scale & minimums and also compute . The following is the expression of the compute needed to get the partial output (corresponding to a single tile) and can be factored out for each group expression, the same happens with and but as a “global” factor
A dequantized weight is and the full dot product + dequant is The main work for each lane is calculating & for the corresponding 8 weight portion of each group it owns. Because activations are all stored in FP16, calculating is very straightforward, you load the corresponding portion and you sum reduce them. However, when computing you have to account for the shifts each has because of how they are stored.
When loading a 2 byte portion of the array from the weight block, you get four 4 weights, where three of them are shifted (4, 8 and 12 slots). Because these shifts are the same for all the quants loaded by a lane. i.e: In the first iteration of the loading you load (equivalent to ), in the next iteration you load (see figure). Instead of correcting each shift you do the shifted acumulation and correct in the end (save redundant shifts) You only perform the correcting shift at the end of the accumulation.
The following code shows that each is responsible for a pair of groups (i.e: G1 and G2 for Team A) and you can see how the accumulation is done without the shifts (only masking the corresponding weights).
for (ushort i = 0; i < 4; ++i) {
// q1 for the first 64 bytes and q2 for the second 64 bytes
acc1[0] += yl[2 * i] * float(q1[i] & 0x000f);
acc1[1] += yl[2 * i + 1] * float(q1[i] & 0x0f00);
acc1[2] += yl[2 * i + 8] * float(q1[i] & 0x00f0);
acc1[3] += yl[2 * i + 9] * float(q1[i] & 0xf000);
acc2[0] += yh[2 * i] * float(q2[i] & 0x000f);
acc2[1] += yh[2 * i + 1] * float(q2[i] & 0x0f00);
acc2[2] += yh[2 * i + 8] * float(q2[i] & 0x00f0);
acc2[3] += yh[2 * i + 9] * float(q2[i] & 0xf000);
}
The for loop ends up being unrolled producing loads of 8 contiguous bytes per lane, if we go back to figure 6, we can see that the loads are coalesced L0 loads , L1 loads , … Because they are coalesced they happen on a single read to memory. Across the whole SIMD we have 4 cohorts so the loads from memory should happen in 4 transactions (per column of ). NOTE: The real amount of transactions depend on undocumented cache/sector boundaries and alignment (in theory they should be 4). between cohorts of 8 lanes. The following is the optimized air code (similar to LLVM IR), which is an intermediate representation before producing the final binary.
# ASM instruction that maps to the q1 or q2 loads
# i16, xyzw means four 16-bit components = 8 byte LOAD
0x230: load r29l_r29h_r30l_r30h, dr7, 2, lsl 2, i16, xyzw, a
When results are reduced we apply the corresponding shifts and multiply each term of the expresion by its corresponding group minimums and scales, as well as the global ones (dh).
The dot(sumy, float4(sc8[2], sc8[3], sc8[6], sc8[7])) performs for the four groups each lane contributes to.
sumf[row] += float(dh[0]) * ((acc1[0] + acc1[1] / 256.0f) * sc8[0] +
(acc1[2] + acc1[3] / 256.0f) * sc8[1] / 16.0f +
(acc2[0] + acc2[1] / 256.0f) * sc8[4] +
(acc2[2] + acc2[3] / 256.0f) * sc8[5] / 16.0f) -
float(dh[1]) * dot(sumy, float4(sc8[2], sc8[3], sc8[6], sc8[7]));
Below is a diagram that visually shows the contributions made by a single team with regards to a weight block. The other team does the same with groups (G3, G4, G7 and G8).
Noting that the group and team separation is just made for dequantization purposes. In the end when we have the partials of each team, each group and each weight block, those are all sum reduced because they belong to the same output tile.
float sum = simd_sum(sumf[row]);
Kernel code
template<uint K, uint N, ushort ROWS = 2>
[[max_total_threads_per_threadgroup(64)]]
kernel void decode_q4k(device half* dst [[buffer(0)]], device const half* src [[buffer(1)]],
device const uchar* weights [[buffer(2)]],
ushort lane [[thread_index_in_simdgroup]],
ushort simd_group [[simdgroup_index_in_threadgroup]],
uint3 group [[threadgroup_position_in_grid]]) {
constexpr ushort kmask1 = 0x3f3f, kmask2 = 0x0f0f, kmask3 = 0xc0c0;
ushort ix = lane / 8, it = lane % 8, iq = it / 4, ir = it % 4;
uint nb = K / 256, first_row = (group.x * 2 + simd_group) * ROWS;
float yl[16], yh[16], sumf[ROWS] = {0.0f};
device const half* src4 = src + ix * 256 + 64 * iq + 8 * ir;
for (uint ib = ix; ib < nb; ib += 4) {
float4 sumy = 0.0f;
for (ushort i = 0; i < 8; ++i) {
// load(x_i) -> sum(x_i)
yl[i] = float(src4[i]); sumy[0] += yl[i];
yl[i + 8] = float(src4[i + 32]); sumy[1] += yl[i + 8];
yh[i] = float(src4[i + 128]); sumy[2] += yh[i];
yh[i + 8] = float(src4[i + 160]); sumy[3] += yh[i + 8];
}
for (ushort row = 0; row < ROWS; ++row) { // one iter per column of w^t owned by SIMD
long o = long(first_row + row) * nb * 144 + ib * 144;
device const ushort* sc = reinterpret_cast<device const ushort*>(weights + o + 4) + iq;
device const ushort* q1 = reinterpret_cast<device const ushort*>(weights + o + 16) + 16 * iq + 4 * ir;
device const ushort* q2 = q1 + 32;
device const half* dh = reinterpret_cast<device const half*>(weights + o);
ushort sc16[4];
thread const uchar* sc8 = reinterpret_cast<thread const uchar*>(sc16);
// load 6 bytes that the team owns (mi, si) in i = {1, 2, 5, 6} or {3, 4, 7, 8}
sc16[0] = sc[0] & kmask1;
sc16[1] = sc[2] & kmask1;
sc16[2] = (sc[4] & kmask2) | ((sc[0] & kmask3) >> 2);
sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
float4 acc1 = 0.0f, acc2 = 0.0f;
// 4(loop) * 2(q1,q2) * 2(2byte LOADS) = 16 bytes loaded per lane * 8 (threads_per_cohort) = 128 bytes (256 weight block)
for (ushort i = 0; i < 4; ++i) {
// sum(xi . qi) unshifted (only mask to isolate the rest 12b)
acc1[0] += yl[2 * i] * float(q1[i] & 0x000f);
acc1[1] += yl[2 * i + 1] * float(q1[i] & 0x0f00);
acc1[2] += yl[2 * i + 8] * float(q1[i] & 0x00f0);
acc1[3] += yl[2 * i + 9] * float(q1[i] & 0xf000);
acc2[0] += yh[2 * i] * float(q2[i] & 0x000f);
acc2[1] += yh[2 * i + 1] * float(q2[i] & 0x0f00);
acc2[2] += yh[2 * i + 8] * float(q2[i] & 0x00f0);
acc2[3] += yh[2 * i + 9] * float(q2[i] & 0xf000);
}
// remove the shifts + apply group scales and mins + global scales and mins
sumf[row] += float(dh[0]) * ((acc1[0] + acc1[1] / 256.0f) * sc8[0] +
(acc1[2] + acc1[3] / 256.0f) * sc8[1] / 16.0f +
(acc2[0] + acc2[1] / 256.0f) * sc8[4] +
(acc2[2] + acc2[3] / 256.0f) * sc8[5] / 16.0f) -
float(dh[1]) * dot(sumy, float4(sc8[2], sc8[3], sc8[6], sc8[7]));
}
src4 += 4 * 256;
}
for (ushort row = 0; row < ROWS; ++row) {
float sum = simd_sum(sumf[row]);
if (lane == 0 && first_row + row < N) dst[first_row + row] = half(sum);
}
}
Roofline Analysis
Measuring performance with metrics such as execution time or tokens per second, is not a fair analysis because these are highly dependent on the hardware the kernels are running on. A better analysis comes by checking if we are using the hardware at our disposal, to the maximum of its capabilities, which is what the Roofline model provides.
Roughly explained the roofline model says that a given kernels’s maximum performance is limited by three factors: 1. The arithmetic intensity, how many FLOP it does per byte (loaded and stored), 2. The maximum theoretical bandwidth and 3. The maximum theoretical of the hardware we are running the kernel on.There are some caveats to this, since there are different levels in the memory hierarchy each with different bandwidths (i.e: GMEM, SMEM, …), and also different compute units (i.e: SP, Tensor cores, …)each with different .
The arithmetic intensity will place the kernel on the -axis of the roofline plot. If it falls on the left of the ridge point (red font) it will be limited by the max bandwidth (memory bound kernel) and if it falls to the right it will be limited by the max theoretical of the device (compute bound kernel).
For simplicity since this is clearly a memory bound kernel, the analysis will be made with the higher level of the hierarchy. The will be from the GPU cores (no special hardware unit) and the bandwidth will be the DRAM one.
In the case of this Q4_K decode gemv where in FP16, in Q4_K (144 bytes per 256 weight block), and we have tiles of processed by each kernel:
- : bytes
- : bytes
- : bytes Activations are read once per SIMD (reuse across SIMDs from same threadgroup doesn’t help, not much reusing) and each SIMD covers two output rows. Therefore, we have SIMDs, each reading bytes of activations ( total bytes). Some of these end up being cached which moves the AI to the right
- FLOPs. Without counting dequantization FLOPs (implementation overhead)
All of the experiments were run on an M2 pro macbook which has a theoretical peak and a LPDDR5 DRAM with bandwidth. Giving an AI ridge point (point of memory bound into compute bound transition):
so we know that our maximum atainable performance, will be limited by the bandwidth of our device. The caveats mentioned earlier about different AI’s are skipped since cache reuse isn’t really possible (each row of weights is read and only used once), so they won’t really move the needle and turn this into a compute bound kernel.
Performance & Conclusion
Taking the kernel to put as an example, the measured time was , and a measured bandwidth of which is a of the memory bound limit, which is related to the overhead introduced by dequantization. This suggests that further gains will likely come from reducing dequantization overhead and improving its overlap with memory accesses (to hide latency).
If you have read up to this point, thanks for reading. I am open for suggestions on improving my writing style, knowledge about these kernels or other optimizations that may come to your mind.
- , , and
- The analysis will be simplified to a single block per cohort, because processing the second block (from the next column) is the same processing with different indexing (done in a for loop)
- All 4 lanes in the team LOAD these local scales and minimums from GMEM to REG
- Each lane also owns its corresponding so that it can perform its partial dot product (i.e: L0 owns & ).
- and can be factored out for each group expression, the same happens with and but as a “global” factor
- A dequantized weight is and the full dot product + dequant is
- i.e: In the first iteration of the loading you load (equivalent to ), in the next iteration you load (see figure). Instead of correcting each shift you do the shifted acumulation and correct in the end (save redundant shifts)
- L0 loads , L1 loads , … Because they are coalesced they happen on a single read to memory. Across the whole SIMD we have 4 cohorts so the loads from memory should happen in 4 transactions (per column of ). NOTE: The real amount of transactions depend on undocumented cache/sector boundaries and alignment (in theory they should be 4).
- The
dot(sumy, float4(sc8[2], sc8[3], sc8[6], sc8[7]))performs for the four groups each lane contributes to. - There are some caveats to this, since there are different levels in the memory hierarchy each with different bandwidths (i.e: GMEM, SMEM, …), and also different compute units (i.e: SP, Tensor cores, …)each with different .
- Activations are read once per SIMD (reuse across SIMDs from same threadgroup doesn’t help, not much reusing) and each SIMD covers two output rows. Therefore, we have SIMDs, each reading bytes of activations ( total bytes). Some of these end up being cached which moves the AI to the right
- Without counting dequantization FLOPs (implementation overhead)