15 min read15 min left

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:

y=αx.WT+βCy = \alpha x . W^T + \beta C x:(1,K)x: (1, K), W:(M,K)W: (M, K), α=1\alpha=1 and β=0\beta=0

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 NN weights represented using MM bytes, for the case of the Q4KQ4K quant we will cover, N=256N = 256 and M=144M = 144, meaning M/N=0.56M/N = 0.56 bytes per weight (~4bit quant).

diagram 2026 07 21 23 58
Figure 1: Weight block

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:

diagram 2026 07 26 01 22

The weight blocks span out across the KK dimension of the weight tensor, meaning that we have a total of N∗K/256N * K/256 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 (1,1024)(1, 1024) tile of the activation vector xx and a (2,1024)(2, 1024) tile of the weights vector (transposed), producing a partial (1,2)(1,2) result of the output vector.

diagram 2026 07 22 00 13
Figure 2: Weight block

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 WTW^T).

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 WTW^T 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.

diagram 2026 07 22 13 14
Figure 3: Groups per team

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

diagram 2026 07 23 12 53
Figure 4: Team A scales & minimums
diagram 2026 07 22 19 48
Figure 5: Team B scales & minimums

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.

diagram 2026 07 22 00 28
Figure 6: Weight ownership by each lane
Each lane also owns its corresponding xix_i so that it can perform its partial dot product (i.e: L0 owns q0..q..7q_0 .. q..7 & x0..x7x_0..x_7).

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 xtile(1,256)@wtile(256,1)x_{tile}(1,256) @ w_{tile}(256,1). The following is the expression of the compute needed to get the partial output (corresponding to a single tile) sgs_g and mgm_g can be factored out for each group expression, the same happens with dd and dmd_m but as a “global” factor

yblock=∑i=0255xi(d s⌊i/32⌋ qi−dm m⌊i/32⌋)=∑g=07∑i=32g32g+31xi(d sg qi−dm mg)pg=∑i=32g32g+31xi(d sg qi−dm mg)=d sg∑i=32g32g+31xiqi⏟Qg−dm mg∑i=32g32g+31xi⏟Xg,yblock=∑g=07pg=∑g=07(d sg Qg−dm mg Xg)=d∑g=07sgQg−dm∑g=07mgXg.\begin{aligned} y_{block} &= \sum_{i=0}^{255} x_i\left(d\,s_{\lfloor i/32\rfloor}\,q_i -d_m\,m_{\lfloor i/32\rfloor}\right) \\[12pt] &= \sum_{g=0}^{7} \sum_{i=32g}^{32g+31} x_i\left(d\,s_g\,q_i-d_m\,m_g\right) \\[16pt] p_g &= \sum_{i=32g}^{32g+31} x_i\left(d\,s_g\,q_i-d_m\,m_g\right) \\[12pt] &= d\,s_g \underbrace{\sum_{i=32g}^{32g+31}x_iq_i}_{Q_g} - d_m\,m_g \underbrace{\sum_{i=32g}^{32g+31}x_i}_{X_g}, \\[16pt] y_{block} &= \sum_{g=0}^{7}p_g \\[12pt] &= \sum_{g=0}^{7} \left(d\,s_g\,Q_g-d_m\,m_g\,X_g\right) \\[12pt] &= d\sum_{g=0}^{7}s_gQ_g - d_m\sum_{g=0}^{7}m_gX_g. \end{aligned}

A dequantized weight is wi=d sg qi−dm mgw_i = d\,s_g\,q_i-d_m\,m_g and the full dot product + dequant is ∑i=0256xi wi\sum_{i=0}^{256}x_i \ w_i The main work for each lane is calculating XgX_g & QgQ_g for the corresponding 8 weight portion of each group it owns. Because activations are all stored in FP16, calculating XgX_g is very straightforward, you load the corresponding portion and you sum reduce them. However, when computing QgQ_g you have to account for the shifts each qiq_i has because of how they are stored.

diagram 2026 07 22 15 26
Figure 7: Loaded ushort

When loading a 2 byte portion of the qiq_i 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 q1<<8q_1 << 8 (equivalent to q1∗28q_1 * 2^8), in the next iteration you load q4∗28q_4 * 2^8 (see figure). Instead of correcting each shift you do the shifted acumulation and correct in the end (save redundant shifts) ((q1∗28)x1+(q4∗28)x4)∗1/28\\((q_1 * 2^8) x_1 + (q_4 * 2^8) x_4) * 1/2^8 You only perform the correcting shift at the end of the ∑i=07(xi wi)\sum_{i=0}^{7}(x_i \ w_i) accumulation.

diagram 2026 07 22 15 32
Figure 8: Shift correcting

The following code shows that each acciacc_i 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 b0..b7b_0..b7, L1 loads b8..b15b_8..b15, … 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 WTW^T). 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 ∑gmgXg\sum_{g}m_gX_g 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).

diagram 2026 07 22 19 33
Figure 9: Visual diagram of contributions

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.

diagram 2026 07 25 15 53
Figure 10: Roofline Plot

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 FLOP/sFLOP/s 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 FLOP/sFLOP/s.

The arithmetic intensity will place the kernel on the xx-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 FLOP/sFLOP/s 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 FLOP/sFLOP/s 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 y=x@WTy = x @ W^T where x:(1,K)x: (1, K) in FP16, WT:(K,N)W^T: (K, N) in Q4_K (144 bytes per 256 weight block), and we have tiles of (1,K)@(K,2)=(2,)(1, K) @ (K, 2) = (2,) processed by each kernel:

  • WTW^T: 144256NK\frac{144}{256} N K bytes
  • yy: 2N2N bytes
  • xx: KNKN 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 N/2N/2 SIMDs, each reading 2K2K bytes of activations (KNKN total bytes). Some of these end up being cached which moves the AI to the right
  • 2NK2NK FLOPs. Without counting dequantization FLOPs (implementation overhead)
AI≈2NK1.5625NK+2N→1.28 FLOP/byteAI \approx \frac{2NK}{1.5625NK + 2N} \rightarrow 1.28\ \text{FLOP/byte}

All of the experiments were run on an M2 pro macbook which has a theoretical peak 6.8TFLOP/s6.8\text{TFLOP}/s and a LPDDR5 DRAM with 200GBs200 GBs bandwidth. Giving an AI ridge point (point of memory bound into compute bound transition):

AIridge=6.8 TFLOP/s200 GB/s≈34 FLOP/byteAI_{\text{ridge}} = \frac{6.8\ \text{TFLOP/s}}{200\ \text{GB/s}} \approx 34\ \text{FLOP/byte}

AIkernel<<AIridgeAI_{kernel} << AI_{ridge} 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 K=N=4096K=N=4096 kernel to put as an example, the measured time was 63μs63 \mu s, and a measured bandwidth of B=152.3GBsB = 152.3GBs which is a 76%76\% 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.


  1. x:(1,K)x: (1, K), W:(M,K)W: (M, K), α=1\alpha=1 and β=0\beta=0
  2. The analysis will be simplified to a single block per cohort, because processing the second block (from the next WTW^T column) is the same processing with different indexing (done in a for loop)
  3. All 4 lanes in the team LOAD these local scales and minimums from GMEM to REG
  4. Each lane also owns its corresponding xix_i so that it can perform its partial dot product (i.e: L0 owns q0..q..7q_0 .. q..7 & x0..x7x_0..x_7).
  5. sgs_g and mgm_g can be factored out for each group expression, the same happens with dd and dmd_m but as a “global” factor
  6. A dequantized weight is wi=d sg qi−dm mgw_i = d\,s_g\,q_i-d_m\,m_g and the full dot product + dequant is ∑i=0256xi wi\sum_{i=0}^{256}x_i \ w_i
  7. i.e: In the first iteration of the loading you load q1<<8q_1 << 8 (equivalent to q1∗28q_1 * 2^8), in the next iteration you load q4∗28q_4 * 2^8 (see figure). Instead of correcting each shift you do the shifted acumulation and correct in the end (save redundant shifts) ((q1∗28)x1+(q4∗28)x4)∗1/28\\((q_1 * 2^8) x_1 + (q_4 * 2^8) x_4) * 1/2^8
  8. L0 loads b0..b7b_0..b7, L1 loads b8..b15b_8..b15, … 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 WTW^T). NOTE: The real amount of transactions depend on undocumented cache/sector boundaries and alignment (in theory they should be 4).
  9. The dot(sumy, float4(sc8[2], sc8[3], sc8[6], sc8[7])) performs ∑gmgXg\sum_{g}m_gX_g for the four groups each lane contributes to.
  10. 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 FLOP/sFLOP/s.
  11. 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 N/2N/2 SIMDs, each reading 2K2K bytes of activations (KNKN total bytes). Some of these end up being cached which moves the AI to the right
  12. Without counting dequantization FLOPs (implementation overhead)