6 · The GEMV kernel, line by line (gemv_kernel)

This chapter reads lines 41–244: the plain BF16 tensor-core GEMV, its launch helper, and its Python entry point. The kernel runs the router at m ≤ 32 and the unpooler slice at m ≤ 8. It is also the skeleton of the compressed kernel in chapter 8, so understanding it here pays off twice.

The plan in one picture.

• One CTA owns 16 weight rows and a slice of K.
• Its NW warps each stream a sub-slice of K and keep running sums in registers.
• At the end, the warps add their sums through shared memory.
• If K is also split across S CTAs, those CTAs add theirs through global memory.
• Finally the result is rounded and stored.

6.1 Includes and two inline-assembly helpers

42
43
44
45
46
47
48

#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include <stdint.h>

typedef __nv_bfloat16 bf16;

• 42The __nv_bfloat16 type and its conversions, such as __float2bfloat16_rn.
• 43The CUDA runtime API.
• 44Gives access to PyTorch's current CUDA stream. Launching into that stream is what lets PyTorch's graph capture record our kernels.
• 45–46torch::Tensor, TORCH_CHECK, and fixed-width integer types such as uint32_t.
• 48A short name for the BF16 type.

50
51
52
53
54
55

__device__ __forceinline__ uint4 ld_stream(const uint4* p) {
  uint4 r;
  asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];"
               : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) : "l"(p));
  return r;
}

• 50
  • __device__: runs on the GPU and can be called from kernels.
  • __forceinline__: pasted into the caller, so there is no call overhead.
  • It returns a uint4: four 32-bit words, 16 bytes.
• 52The PTX load explained in chapter 2.4: global memory, read-only path, don't keep in L1, 16 bytes. %0…%3 are the four destination registers and [%4] is the address.
• 53The constraints. "=r" means "an output in a 32-bit register", and "l" means "a 64-bit input" (the pointer). volatile stops the compiler from deleting or merging the load; it can still schedule it early.

Why not simply __ldg(p)? That also uses the read-only path, but may keep the line in L1. Weight bytes are used exactly once, so we tell the hardware not to cache them. L1 then stays free for x, which is reused.

56
57
58
59
60

__device__ __forceinline__ void mma16816(float (&d)[4], uint4 a, uint32_t b0, uint32_t b1) {
  asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
               : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
               : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1));
}

• 56Takes the lane's four FP32 accumulators by reference, its A fragment (4 registers, each holding 2 BF16 values, in one uint4) and its two B registers.
• 57m16n8k16 tile shape; row.col means A is row-major and B column-major; the types are D = f32, A = bf16, B = bf16, C = f32. The operand lists are D {%0…%3}, A {%4…%7}, B {%8,%9} and C, which is again {%0…%3}: the instruction accumulates in place.
• 58–59"+f" marks registers that are both read and written (the accumulators). A and B are plain 32-bit inputs.

6.2 The kernel's signature

63
64
65
66
67
68
69
70

// One CTA = one 16-row block x a K-range; its NW warps split that range, reduced through smem.
// S > 1 additionally splits K across CTAs (semaphore-fused reduction by the last CTA).
// PIPE: 0/3 = load a batch, then compute it (3 also sets minBlocks=1 for the register allocator); 1/2 = software-pipelined (next batch's loads issued before this batch is
// computed); 2 also caps registers so two CTAs fit per SM.
template <int NW, int U, int NB, bool DUAL, int PIPE>
__global__ void __launch_bounds__(NW * 32, (PIPE == 2 ? 2 : PIPE == 3 ? 1 : 0))
gemv_kernel(const bf16* __restrict__ x, const uint4* __restrict__ w, const uint4* __restrict__ w2,
            void* __restrict__ out, float4* __restrict__ ws, int* __restrict__ sem, int M, int N, int K, int S) {

• 63–66The design summary: one CTA = a 16-row block × a K range, warps split it, and four loop styles (PIPE).
• 67The compile-time parameters:
  • NW: warps per CTA;
  • U: K blocks a warp loads per loop trip;
  • NB: 8-token column blocks (1, 2 or 4);
  • DUAL: the router's hi + lo;
  • PIPE: the loop style.
• 68__global__ makes this a kernel. __launch_bounds__(threads, minBlocks) tells the compiler the CTA size:
  • PIPE 2 asks that 2 CTAs fit on one SM. That caps registers at 65,536 / (2 × NW × 32) per thread.
  • PIPE 3 asks for only 1, so the compiler may use as many registers as one CTA can have: 65,536 / (NW × 32), and at most 255.
  • 0 gives no hint.
  These hints change how the compiler allocates registers and schedules loads, sometimes by 25 %, which is why they are tuned per shape.
• 69–70The parameters:
  • x: BF16 [M, K], row-major.
  • w, w2: the packed weights, addressed in 16-byte units. w2 is the router's lo matrix.
  • out: BF16 or FP32, hence void*.
  • ws, sem: scratch space and counters for split-K.
  __restrict__ promises that the pointers don't overlap, so the compiler may reorder loads freely.

6.3 Who does what

71
72
73
74
75
76
77
78

  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 KB = K / 32, KBS = KB / S, KBW = KBS / NW;
  const int kb0 = sidx * KBS + warp * KBW;
  const size_t woff = ((size_t)rb * KB + kb0) * 64 + lane;
  const uint4* wp = w + woff;
  const uint4* wp2 = DUAL ? (w2 + woff) : nullptr;

• 71A compile-time check. The final reduction needs NB × 32 threads, so a CTA must have at least NB warps. An early version without this check gave garbage for 2-warp CTAs at m = 32.
• 72Thread coordinates. g and t are the fragment coordinates of chapter 2.5.
• 73rb: which 16-row block (blockIdx.x, 0 … N/16−1). sidx: which K slice (blockIdx.y, 0 … S−1).
• 74K in 32-wide blocks: KB in total, KBS per CTA slice, KBW per warp. These are integer divisions, so the tuned configurations always divide exactly (lab 0 checks this).
• 75This warp's first block. Warps take consecutive chunks of the slice.
• 76–78Where this lane's first fragment sits in the packed array (chapter 5.4), for w and for the router's w2.

79
80
81
82
83
84
85
86
87

  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 + kb0 * 32 + 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;

• 79–84One x pointer per 8-token block: token nb·8 + g at k = kb0·32 + 8t. Tokens past M are clamped to the last real token: a harmless duplicate whose result is never stored.
• 85–87The accumulators: 4 FP32 values (one C fragment) per token block, zeroed. #pragma unroll turns these tiny loops into straight-line code, which keeps acc in registers.

6.4 The main loop, version 1: load a batch, then compute it (PIPE 0 and 3)

89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

  if constexpr (PIPE == 0 || PIPE == 3) {
    for (int i = 0; i < KBW; i += U) {
      uint4 a[U][2], b[DUAL ? U : 1][2];
#pragma unroll
      for (int u = 0; u < U; u++) {
        a[u][0] = ld_stream(wp + (size_t)(i + u) * 64);
        a[u][1] = ld_stream(wp + (size_t)(i + u) * 64 + 32);
        if (DUAL) {
          b[u][0] = ld_stream(wp2 + (size_t)(i + u) * 64);
          b[u][1] = ld_stream(wp2 + (size_t)(i + u) * 64 + 32);
        }
      }
#pragma unroll
      for (int u = 0; u < U; u++) {
#pragma unroll
        for (int nb = 0; nb < NB; nb++) {
          uint4 xv = __ldg((const uint4*)(xr[nb] + (i + u) * 32));
          mma16816(acc[nb], a[u][0], xv.x, xv.y);
          mma16816(acc[nb], a[u][1], xv.z, xv.w);
          if (DUAL) {
            mma16816(acc[nb], b[u][0], xv.x, xv.y);
            mma16816(acc[nb], b[u][1], xv.z, xv.w);
          }
        }
      }
    }

• 89if constexpr is decided at compile time. The other branch is not even compiled into this instantiation.
• 90Walk this warp's K blocks, U at a time.
• 91Register arrays for U blocks of fragments: two per block, one per mma step. DUAL adds the same again for lo.
• 92–100All loads first: 2U independent 16-byte loads per lane (4U for the router). None depends on another, so they are all in flight together. This is what U tunes.
• 101–113Then compute. For each block and token block:
  • load 16 bytes of x (from L1/L2: every CTA reads the same x);
  • run two mma, one per step.
  With DUAL, two more mma use the same x with the lo fragments and the same accumulators. That is the whole "fused hi + lo" router: Σ x·hi + Σ x·lo in one pass, with one set of sums.

The first mma waits until its weights arrive (about a microsecond). By then the other loads of the batch have arrived too, so the rest of the batch runs without waiting. The weakness is visible too: the next batch's loads only start after this batch's math is done, which leaves a bubble every trip.

6.5 The main loop, version 2: software pipelining (PIPE 1 and 2)

115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

  } else {
    uint4 a[U][2], b[DUAL ? U : 1][2];
#pragma unroll
    for (int u = 0; u < U; u++) {
      a[u][0] = ld_stream(wp + (size_t)u * 64);
      a[u][1] = ld_stream(wp + (size_t)u * 64 + 32);
      if (DUAL) {
        b[u][0] = ld_stream(wp2 + (size_t)u * 64);
        b[u][1] = ld_stream(wp2 + (size_t)u * 64 + 32);
      }
    }
    for (int i = 0; i < KBW; i += U) {
      uint4 xv[U][NB];
#pragma unroll
      for (int u = 0; u < U; u++)
#pragma unroll
        for (int nb = 0; nb < NB; nb++) xv[u][nb] = __ldg((const uint4*)(xr[nb] + (i + u) * 32));
      uint4 na[U][2], nbb[DUAL ? U : 1][2];
      if (i + U < KBW) {
#pragma unroll
        for (int u = 0; u < U; u++) {
          na[u][0] = ld_stream(wp + (size_t)(i + U + u) * 64);
          na[u][1] = ld_stream(wp + (size_t)(i + U + u) * 64 + 32);
          if (DUAL) {
            nbb[u][0] = ld_stream(wp2 + (size_t)(i + U + u) * 64);
            nbb[u][1] = ld_stream(wp2 + (size_t)(i + U + u) * 64 + 32);
          }
        }
      }
#pragma unroll
      for (int u = 0; u < U; u++) {
#pragma unroll
        for (int nb = 0; nb < NB; nb++) {
          mma16816(acc[nb], a[u][0], xv[u][nb].x, xv[u][nb].y);
          mma16816(acc[nb], a[u][1], xv[u][nb].z, xv[u][nb].w);
          if (DUAL) {
            mma16816(acc[nb], b[u][0], xv[u][nb].x, xv[u][nb].y);
            mma16816(acc[nb], b[u][1], xv[u][nb].z, xv[u][nb].w);
          }
        }
      }
#pragma unroll
      for (int u = 0; u < U; u++) {
        a[u][0] = na[u][0]; a[u][1] = na[u][1];
        if (DUAL) { b[u][0] = nbb[u][0]; b[u][1] = nbb[u][1]; }
      }
    }
  }

• 116–125Prologue: load the first batch.
• 126–131Each trip starts by loading x for the current batch.
• 132–143Issue the next batch's weight loads now, before computing the current batch, unless this is the last trip.
• 144–155Compute the current batch: the same mma as in version 1.
• 156–160The loaded "next" batch becomes "current".

Now, while the tensor cores chew on batch i, the loads of batch i+1 are already travelling. In round 3 we put timestamps inside the kernel and found that about 2 µs of qkv_gate's gap to a pure memory-read kernel was the bubble that version 1 leaves. PIPE 2 combines this loop with "2 CTAs per SM".

Which version wins depends on the shape. The tuner tried all four for every shape, and the tables record the winner.

6.6 Adding up the warps

164
165
166
167
168
169
170
171
172
173
174
175
176
177

  __shared__ float4 red[NW][NB * 32];
  __shared__ int s_last;
#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;
  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;
    }
  }

• 164Shared memory, one float4 per (warp, token block, lane). With 16 warps and NB = 1 that is 16 × 32 × 16 B = 8 KB.
• 165A one-word flag used by the split-K step below.
• 166–167Every warp parks its accumulators there.
• 168__syncthreads(): wait until every warp has written.
• 169–177The first NB × 32 threads each add one position over all NW warps. Thread tid ends up holding the summed C fragment of lane tid & 31 for token block tid >> 5.

6.7 Adding up CTAs: split-K in the same launch

178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194

  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;
  }

• 178Only when K is split across CTAs: the router (S = 4, 8 or 16) and the default router configuration.
• 179Write this CTA's partial sums to its own slot in the global scratch buffer ws.
• 180__threadfence(): the write must become visible to the whole GPU before anything after it, in particular the counter increment.
• 181Wait until every thread of the CTA has written and fenced.
• 182–183One thread increments this row block's counter. atomicAdd returns the old value, so exactly one CTA sees S−1: the last to finish. The flag is passed to the other threads through shared memory.
• 184Every CTA except the last is done.
• 185Another fence, this time on the reading side, so the partials we are about to read are the fresh ones.
• 186–192The last CTA adds all S partials in a fixed order (s = 0, 1, …). So the result does not depend on which CTA happened to finish last. __ldcg reads from L2, bypassing L1: L1 is private to an SM and would not see other SMs' writes.
• 193Reset the counter for the next call. The CUDA graph is replayed again and again, so whatever we leave behind is the next call's starting state.

6.8 Storing the result

195
196
197
198
199
200
201
202
203
204
205
206
207
208

  if (active) {
    const int nb = tid >> 5;
    float rr[4] = {r.x, r.y, r.z, r.w};
#pragma unroll
    for (int q = 0; q < 4; q++) {
      int tok = nb * 8 + 2 * t + (q & 1);
      int row = rb * 16 + g + ((q >> 1) << 3);
      if (tok < M) {
        if (DUAL) ((float*)out)[(size_t)tok * N + row] = rr[q];
        else ((bf16*)out)[(size_t)tok * N + row] = __float2bfloat16_rn(rr[q]);
      }
    }
  }
}

• 196–197The thread's token block and its four sums.
• 199–201The C-fragment mapping from chapter 2.5. Value q of lane (g, t) is token 2t + (q & 1) (plus 8·nb) and row g + 8·(q >> 1) (plus 16·rb).
• 202Skip the clamped tokens past M.
• 203–204The router stores FP32. Every other layer rounds to BF16 (round to nearest even) and stores. These are scattered 2- or 4-byte writes, which is fine: the output is tiny compared with the weights.

6.9 Launching from C++, and the instantiation list

210
211
212
213
214
215
216
217

template <int NW, int U, int NB, bool DUAL, int PIPE>
static void launch(torch::Tensor x, torch::Tensor w, torch::Tensor w2, torch::Tensor out, torch::Tensor ws,
                   torch::Tensor sem, int N, int K, int S) {
  dim3 grid(N / 16, S);
  gemv_kernel<NW, U, NB, DUAL, PIPE><<<grid, NW * 32, 0, c10::cuda::getCurrentCUDAStream()>>>(
      (const bf16*)x.data_ptr(), (const uint4*)w.data_ptr(), DUAL ? (const uint4*)w2.data_ptr() : nullptr,
      out.data_ptr(), (float4*)ws.data_ptr(), (int*)sem.data_ptr(), (int)x.size(0), N, K, S);
}

• 210–212A host function template, one per kernel instantiation.
• 213The grid: N/16 row blocks × S K slices.
• 214The <<<grid, threads, dynamic shared bytes, stream>>> launch: NW × 32 threads, no dynamic shared memory, and PyTorch's current stream.
• 215–216Raw device pointers from the tensors. M is x's first dimension.

220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244

#define C3(NW_, U_, NB_, D_, P_) if (NW == NW_ && U == U_ && NB == NB_ && dual == D_ && P == P_) { launch<NW_, U_, NB_, D_, P_>(x, w, w2, out, ws, sem, N, K, S); return; }
void gemv(torch::Tensor x, torch::Tensor w, torch::Tensor w2, torch::Tensor out, torch::Tensor ws, torch::Tensor sem,
          int64_t N, int64_t K, int64_t NW, int64_t S, int64_t U, int64_t NB, bool dual, int64_t P) {
  C3(4, 2, 1, false, 0)
  C3(4, 2, 1, false, 3)
  C3(4, 2, 1, true, 2)
  C3(4, 2, 2, false, 0)
  C3(4, 2, 4, false, 0)
  C3(4, 4, 1, false, 2)
  C3(8, 2, 1, false, 3)
  C3(8, 2, 1, true, 0)
  C3(8, 2, 2, false, 0)
  C3(8, 2, 2, true, 0)
  C3(8, 2, 2, true, 3)
  C3(8, 2, 4, true, 0)
  C3(8, 4, 1, false, 2)
  C3(16, 2, 1, false, 2)
  C3(16, 2, 1, false, 3)
  C3(16, 2, 4, true, 3)
  C3(16, 4, 1, true, 3)
  C3(16, 8, 1, false, 3)
  C3(16, 8, 2, false, 0)
  C3(16, 8, 4, false, 0)
  TORCH_CHECK(false, "unsupported gemv config");
}

• 220The C3 macro: "if the runtime arguments equal this compile-time combination, launch that instantiation and return". S is not part of it, because S is an ordinary runtime argument.
• 221–222The function exported to Python and declared on line 634.
• 223–242The 20 combinations compiled into the file. build.py generated this list from the tables, plus the defaults used for odd test sizes, such as C3(4, 2, 1, false, 0) for m = 5 and 7. Each line instantiates one kernel, which costs compile time but is needed so everything stays in registers.
• 243Any other combination raises "unsupported gemv config" in Python. Chapter 10's experiment D runs into exactly this and shows the fix: add a line.

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

Exercises.

1. In the lab, set RESET = False and watch replay 1 lose its output.
2. Router at m = 1 uses CTABLE[('router', 1)] = (16, 4, 4, 3). With lab 0, work out its grid (64 CTAs). How many bytes does each lane have in flight right after the load phase? (4 blocks × 4 loads × 16 B = 256 B.)
3. Chapter 10's experiment B sets S = 1 for that shape on the real H200: 16 CTAs instead of 64. Predict the effect before you run it.

Check yourself