⌂ Guide home · 8 · The compressed GEMV← 7 · The 12-bit format | 9 · Router and Triton →

8 · The compressed GEMV, line by line (`gemvz_kernel`)

gemvz_kernel (lines 275–472) is the workhorse of the submission. It runs 12 of the 30 scored shapes and produces most of our speedup. Structurally it is gemv_kernel from chapter 6 with three changes:

ON THIS PAGE

8.1 Signature and shared state
8.2 The exception warp
8.3 The streaming warps: set-up
8.4 Main loop, version 1 (PIPE 0 and 3)
8.5 Main loop, version 2: software pipelining (PIPE 1 and 2)
8.6 Combining everything
8.7 Launch and instantiation list
Check yourself

It reads the 12-bit format: 48 bytes per lane per 64-wide K group instead of 64.
It decodes those bytes into tensor-core fragments with decz.
One extra warp applies the exact corrections for the rare out-of-window weights while the other warps stream.

This chapter focuses on what is new. For the parts copied from chapter 6, it points back.

8.1 Signature and shared state

submission.py, lines 275–289

// One CTA = one 16-row block x a K-range (blockIdx.y of S). NW streaming warps split the range; warp NW applies
// the exception corrections for the CTA's rows and K-range. PIPE as in gemv_cuda: 0/3 plain (3: minBlocks 1),
// 1/2 software-pipelined (2: minBlocks 2).
template <int NW, int UG, int NB, int PIPE>
__global__ void __launch_bounds__((NW + 1) * 32, (PIPE == 2 ? 2 : PIPE == 3 ? 1 : 0))
gemvz_kernel(const bf16* __restrict__ x, const uint4* __restrict__ w, const int2* __restrict__ esc, int capr,
             bf16* __restrict__ out, float4* __restrict__ ws, int* __restrict__ sem, int M, int N, int K, int S,
             uint32_t TL, uint32_t TH, int dbg) {
  static_assert(NW >= NB, "need at least NB warps for the final reduction");
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3;
  const int rb = blockIdx.x, sidx = blockIdx.y;
  const int KG = K / 64, KGS = KG / S, KGW = KGS / NW;
  __shared__ float4 red[NW][NB * 32];
  __shared__ float ecorr[16][NB * 8 + 1];
  __shared__ int s_last;

275–277The design: NW warps stream the weights, and warp number NW (one extra) handles the exceptions. PIPE works as in chapter 6.
278Template parameters NW, UG (64-wide groups per loop trip), NB and PIPE. There is no DUAL, because the router is never compressed.
279The launch bounds count the extra warp: (NW + 1) × 32 threads.
280–282New parameters:
- esc and capr: the exception table, with capr slots per row;
- TL, TH: the window's lookup words;
- dbg: debug switches. It is always 0 in production; dbg = 1 skips the corrections, so their cost can be measured.
286K in 64-wide groups: KG in total, KGS per CTA slice, KGW per warp.
287The streaming warps' partial sums, as in chapter 6.
288ecorr[16][NB·8 + 1]: the exception warp's results, one per (row, token). The "+1" makes the row length odd, so the 16 lanes writing 16 different rows hit 16 different shared-memory banks. Without it, the row stride would be 8, 16 or 32 floats, and the lanes would queue up on the same banks.

8.2 The exception warp

submission.py, lines 291–314

  if (warp == NW) {
    // exception corrections: lane -> (row r, tokens p, p+2, ...); fixed order, so deterministic
    const int r = lane & 15, p = lane >> 4;
    const int klo = sidx * KGS * 64, khi = klo + KGS * 64;
    float ea[NB * 4];
#pragma unroll
    for (int q = 0; q < NB * 4; q++) ea[q] = 0.f;
    if (!(dbg & 1)) {
      const int2* ep = esc + (size_t)(rb * 16 + r) * capr;
#pragma unroll 4
      for (int j = 0; j < capr; j++) {
        const int2 e = __ldg(ep + j);
        if (e.y != 0 && e.x >= klo && e.x < khi) {
          const float d = __int_as_float(e.y);
#pragma unroll
          for (int q = 0; q < NB * 4; q++) {
            const int tok = p + 2 * q;
            if (tok < M) ea[q] += d * __bfloat162float(x[(size_t)tok * K + e.x]);
          }
        }
      }
    }
#pragma unroll
    for (int q = 0; q < NB * 4; q++) ecorr[r][p + 2 * q] = ea[q];

291Only the last warp takes this branch. All the others go to the streaming code at line 315.
293Its 32 lanes cover 16 rows × 2 token parities: row r = lane & 15, with tokens p, p+2, p+4, … where p = lane >> 4.
294This CTA's slice of K. With split-K, each exception must be applied by exactly one CTA.
295–297One accumulator per token this lane covers.
299Row r's list starts at (16·rb + r) × capr.
300–302Walk all capr slots (unrolled 4 at a time), loading each (k, value bits) entry through the read-only path.
303Skip empty slots (value 0) and corrections that belong to another CTA's K slice.
304Reinterpret the 32 bits as the FP32 correction.
305–309For each of the lane's tokens, add correction × x[token][k]. These x reads are scattered, but a row has only about one exception on average (0.014 % of 8,192 k values), and x sits in L2.
313–314Store the lane's results in ecorr[r][token].

Each (row, token) belongs to exactly one lane and is summed in a fixed order, so the result is deterministic. The warp runs at the same time as the streaming warps and has far less to do, so it finishes long before them. Its cost is roughly one warp's share of issue slots and registers.

8.3 The streaming warps: set-up

submission.py, lines 315–326

  } else {
    const int g0 = sidx * KGS + warp * KGW;
    const uint4* wp = w + ((size_t)rb * KG + g0) * 96 + lane;
    const bf16* xr[NB];
#pragma unroll
    for (int nb = 0; nb < NB; nb++) {
      int tok = min(nb * 8 + g, M - 1);
      xr[nb] = x + (size_t)tok * K + g0 * 64 + t * 8;
    }
    float acc[NB][4];
#pragma unroll
    for (int nb = 0; nb < NB; nb++) acc[nb][0] = acc[nb][1] = acc[nb][2] = acc[nb][3] = 0.f;

316This warp's first group.
317One (row block, group) is 96 sixteen-byte units (1,536 bytes: three sections of 32 lanes). Adding lane makes each section read 512 contiguous bytes per warp.
318–323x pointers, as in chapter 6 but stepping 64 k per group: token nb·8 + g, k = g0·64 + 8t.
324–326The accumulators.

8.4 Main loop, version 1 (PIPE 0 and 3)

submission.py, lines 328–358

    if constexpr (PIPE == 0 || PIPE == 3) {
      for (int i = 0; i < KGW; i += UG) {
        uint4 A0[UG], A1[UG], B[UG];
#pragma unroll
        for (int u = 0; u < UG; u++) {
          const uint4* p = wp + (size_t)(i + u) * 96;
          A0[u] = ldz_stream(p);
          A1[u] = ldz_stream(p + 32);
          B[u] = ldz_stream(p + 64);
        }
#pragma unroll
        for (int u = 0; u < UG; u++) {
          uint4 xa[NB], xb[NB];
#pragma unroll
          for (int nb = 0; nb < NB; nb++) {
            xa[nb] = __ldg((const uint4*)(xr[nb] + (i + u) * 64));
            xb[nb] = __ldg((const uint4*)(xr[nb] + (i + u) * 64 + 32));
          }
          const uint4 f0 = decz(A0[u].x, A0[u].y, B[u].x, TL, TH);
          const uint4 f1 = decz(A0[u].z, A0[u].w, B[u].y, TL, TH);
          const uint4 f2 = decz(A1[u].x, A1[u].y, B[u].z, TL, TH);
          const uint4 f3 = decz(A1[u].z, A1[u].w, B[u].w, TL, TH);
#pragma unroll
          for (int nb = 0; nb < NB; nb++) {
            mmaz(acc[nb], f0, xa[nb].x, xa[nb].y);
            mmaz(acc[nb], f1, xa[nb].z, xa[nb].w);
            mmaz(acc[nb], f2, xb[nb].x, xb[nb].y);
            mmaz(acc[nb], f3, xb[nb].z, xb[nb].w);
          }
        }
      }

329Walk this warp's groups, UG at a time.
330–337All loads first: 3 × 16 bytes per group per lane, for sections 0, 1 and 2 at +0, +32 and +64 units.
339–345For each group, load x: two 16-byte loads per token block. xa covers k 8t…8t+7 of the group's first 32, and xb the same in its second 32. They are issued before the decode, so they arrive while the integer work runs.
346–349Decode the four fragments of the group:
- f0 from the first half of A0 plus code word B.x;
- f1 from the second half of A0 plus B.y;
- f2 and f3 likewise from A1 plus B.z and B.w.
350–356Four mma per token block:
- f0 and f1 are the two steps of the first 32-wide block, with xa's low and high halves;
- f2 and f3 are those of the second block, with xb.
This is the same K relabelling as chapter 5.

8.5 Main loop, version 2: software pipelining (PIPE 1 and 2)

submission.py, lines 359–404

    } else {
      uint4 A0[UG], A1[UG], B[UG];
#pragma unroll
      for (int u = 0; u < UG; u++) {
        const uint4* p = wp + (size_t)u * 96;
        A0[u] = ldz_stream(p);
        A1[u] = ldz_stream(p + 32);
        B[u] = ldz_stream(p + 64);
      }
      for (int i = 0; i < KGW; i += UG) {
        uint4 xa[UG][NB], xb[UG][NB];
#pragma unroll
        for (int u = 0; u < UG; u++)
#pragma unroll
          for (int nb = 0; nb < NB; nb++) {
            xa[u][nb] = __ldg((const uint4*)(xr[nb] + (i + u) * 64));
            xb[u][nb] = __ldg((const uint4*)(xr[nb] + (i + u) * 64 + 32));
          }
        uint4 nA0[UG], nA1[UG], nB[UG];
        if (i + UG < KGW) {
#pragma unroll
          for (int u = 0; u < UG; u++) {
            const uint4* p = wp + (size_t)(i + UG + u) * 96;
            nA0[u] = ldz_stream(p);
            nA1[u] = ldz_stream(p + 32);
            nB[u] = ldz_stream(p + 64);
          }
        }
#pragma unroll
        for (int u = 0; u < UG; u++) {
          const uint4 f0 = decz(A0[u].x, A0[u].y, B[u].x, TL, TH);
          const uint4 f1 = decz(A0[u].z, A0[u].w, B[u].y, TL, TH);
          const uint4 f2 = decz(A1[u].x, A1[u].y, B[u].z, TL, TH);
          const uint4 f3 = decz(A1[u].z, A1[u].w, B[u].w, TL, TH);
#pragma unroll
          for (int nb = 0; nb < NB; nb++) {
            mmaz(acc[nb], f0, xa[u][nb].x, xa[u][nb].y);
            mmaz(acc[nb], f1, xa[u][nb].z, xa[u][nb].w);
            mmaz(acc[nb], f2, xb[u][nb].x, xb[u][nb].y);
            mmaz(acc[nb], f3, xb[u][nb].z, xb[u][nb].w);
          }
        }
#pragma unroll
        for (int u = 0; u < UG; u++) { A0[u] = nA0[u]; A1[u] = nA1[u]; B[u] = nB[u]; }
      }
    }

360–367Prologue: load the first UG groups.
369–376Load x for the current groups.
377–386Issue the next groups' weight loads before decoding the current ones.
387–400Decode and multiply the current groups.
402"Next" becomes "current".

This is the same idea as chapter 6.5. Here it matters even more, because decoding adds integer work between a load and its use. The tuner picked PIPE 2 for the biggest shapes at m = 1 and 8 (qkv_gate and pooler: 16 warps, 2 CTAs per SM), PIPE 1 for qkv_gate at m = 16, and PIPE 3 or 0 for the others.

8.6 Combining everything

submission.py, lines 405–422

#pragma unroll
    for (int nb = 0; nb < NB; nb++) red[warp][nb * 32 + lane] = make_float4(acc[nb][0], acc[nb][1], acc[nb][2], acc[nb][3]);
  }
  __syncthreads();
  float4 r = make_float4(0.f, 0.f, 0.f, 0.f);
  const bool active = tid < NB * 32;
  const int nbo = tid >> 5;
  if (active) {
#pragma unroll
    for (int wi = 0; wi < NW; wi++) {
      float4 v = red[wi][tid];
      r.x += v.x; r.y += v.y; r.z += v.z; r.w += v.w;
    }
    r.x += ecorr[g][nbo * 8 + 2 * t];
    r.y += ecorr[g][nbo * 8 + 2 * t + 1];
    r.z += ecorr[g + 8][nbo * 8 + 2 * t];
    r.w += ecorr[g + 8][nbo * 8 + 2 * t + 1];
  }

406The streaming warps park their partial sums, as in chapter 6.
408One barrier serves both kinds of warp: after it, red and ecorr are both complete.
409–417Add over the streaming warps.
418–421Add the corrections for the thread's four outputs. Value q of the C fragment is (row g or g+8, token nbo·8 + 2t or +1), and ecorr is indexed the same way. From here on, the sums are those of the original weights.

submission.py, lines 423–449

  if (S > 1) {
    if (active) ws[(size_t)(rb * S + sidx) * (NB * 32) + tid] = r;
    __threadfence();
    __syncthreads();
    if (tid == 0) s_last = (atomicAdd(&sem[rb], 1) == S - 1);
    __syncthreads();
    if (!s_last) return;
    __threadfence();
    if (active) {
      r = make_float4(0.f, 0.f, 0.f, 0.f);
      for (int s = 0; s < S; s++) {
        float4 v = __ldcg(ws + (size_t)(rb * S + s) * (NB * 32) + tid);
        r.x += v.x; r.y += v.y; r.z += v.z; r.w += v.w;
      }
    }
    if (tid == 0) sem[rb] = 0;
  }
  if (active) {
    float rr[4] = {r.x, r.y, r.z, r.w};
#pragma unroll
    for (int q = 0; q < 4; q++) {
      int tok = nbo * 8 + 2 * t + (q & 1);
      int row = rb * 16 + g + ((q >> 1) << 3);
      if (tok < M) out[(size_t)tok * N + row] = __float2bfloat16_rn(rr[q]);
    }
  }
}

423–439Split-K across CTAs, line for line the protocol of chapter 6.7. It is used by qkv_gate at m = 16 (S = 4).
440–448Store, as in chapter 6.8. The output is always BF16.

8.7 Launch and instantiation list

submission.py, lines 451–472

template <int NW, int UG, int NB, int PIPE>
static void launchz(torch::Tensor x, torch::Tensor w, torch::Tensor esc, int capr, torch::Tensor out, torch::Tensor ws,
                    torch::Tensor sem, int N, int K, int S, uint32_t TL, uint32_t TH, int dbg) {
  dim3 grid(N / 16, S);
  gemvz_kernel<NW, UG, NB, PIPE><<<grid, (NW + 1) * 32, 0, c10::cuda::getCurrentCUDAStream()>>>(
      (const bf16*)x.data_ptr(), (const uint4*)w.data_ptr(), (const int2*)esc.data_ptr(), capr,
      (bf16*)out.data_ptr(), (float4*)ws.data_ptr(), (int*)sem.data_ptr(), (int)x.size(0), N, K, S, TL, TH, dbg);
}
#define Z3(NW_, U_, NB_, P_) if (NW == NW_ && UG == U_ && NB == NB_ && P == P_) { launchz<NW_, U_, NB_, P_>(x, w, esc, (int)capr, out, ws, sem, (int)N, (int)K, (int)S, (uint32_t)TL, (uint32_t)TH, (int)dbg); return; }
void gemvz(torch::Tensor x, torch::Tensor w, torch::Tensor esc, int64_t capr, torch::Tensor out, torch::Tensor ws,
           torch::Tensor sem, int64_t N, int64_t K, int64_t NW, int64_t S, int64_t UG, int64_t NB, int64_t P,
           int64_t TL, int64_t TH, int64_t dbg) {
  Z3(4, 1, 2, 1)
  Z3(4, 1, 4, 3)
  Z3(4, 2, 1, 3)
  Z3(4, 2, 2, 3)
  Z3(8, 1, 4, 3)
  Z3(8, 2, 2, 0)
  Z3(8, 2, 2, 3)
  Z3(16, 1, 1, 2)
  TORCH_CHECK(false, "unsupported gemvz config");
}

454–455Grid (N/16, S), with (NW + 1) × 32 threads: the +1 is the exception warp.
459The Z3 macro matches (NW, UG, NB, PIPE). S, TL, TH and dbg are runtime arguments.
463–470The 8 combinations compiled in, exactly the ones ZTABLE uses. A new ZTABLE configuration needs a new line here (chapter 10, experiment D).
471Anything else raises "unsupported gemvz config".

Key idea. Keep the hot loop identical for every weight and move rare cases off to the side. A first compression attempt in round 1 applied the corrections after the main loop, which added a serial tail, and it gained only 1–3 %. Round 5's design gives them their own warp running in parallel, and keeps the main loop branch-free. That is what turned the same 25 % byte saving into 13–20 % faster shapes.

Try it · lab3_compress.py — Lab 3: the lossless 12-bit weight format (pack_z + decz), emulated bit for bit in numpy.

cd ~/projects/okc/muse_bf16_proj/learn/labs
python3 lab3_compress.py

The file is already on your Mac at ~/projects/okc/muse_bf16_proj/learn/labs/lab3_compress.py; you can also view or download it here.

Part 3 of the lab is this kernel in miniature. It multiplies the decoded weights by x, which is what the streaming warps compute, and lets an emulated exception warp add its corrections. Then it compares against the original weights. It also runs a judge-style probe row on an exception column: decoded alone the answer is visibly wrong, and with the correction it is exact.

Exercise. Chapter 10's experiment A removes two ZTABLE entries, sending pooler at m = 1 and 8 back to the plain GEMV. Using the harness numbers in chapter 3.3 (17.12 vs 13.63 µs), predict the new score before you run it.

Check yourself

About how much memory does the exception table add for the pooler?

About 4,700 corrections spread over 4,096 rows, padded to capr slots per row (single digits in our samples). At 8 bytes per slot, that is a few hundred KB at most, against about 50 MB of packed weights. The table is read once per call by the exception warps, which costs almost nothing.

Why does the exception warp filter by klo ≤ k < khi?

With split-K (S > 1), every CTA of a row block sees the same exception list, but each CTA covers only its own slice of K. Without the filter, every correction would be added S times.

The final reduction uses threads tid < NB·32. Can the exception warp be one of them?

No. Those threads are the first NB warps. Because NW ≥ NB (the static_assert on line 283), the first NB warps are always streaming warps. The exception warp has the highest index, NW, so its only contribution is what it wrote into ecorr.