5 · Repacking weights for the tensor core (_pack)

Four lines of Python make the GEMV fast. _pack reorders each weight matrix once, in prepare(). After that, every lane can fetch its whole tensor-core operand with a single 16-byte load, and every warp reads one contiguous 512-byte chunk. This chapter explains the trick behind it. The lab then runs the real layout through an emulated warp to prove the answer comes out right.

5.1 The problem with the natural layout

Recall from chapter 2 what one lane must hand to mma.sync.m16n8k16 as its A fragment:

• 8 BF16 weights;
• from two rows (g and g+8);
• at four k positions (2t, 2t+1, 2t+8, 2t+9).

In the natural row-major W[n][k], those 8 values sit in four separate places, two rows apart in memory. A lane would need four small loads, and the warp's loads would be scattered. Both are bad for a kernel whose only job is to stream bytes at full speed.

5.2 The trick: the tensor core doesn't know which k is which

A dot product is a sum: Σₖ W[n][k] · x[k]. The order of the terms doesn't matter. Inside one mma, the hardware multiplies A's k-slot j by B's k-slot j, for j = 0…15, and adds the products. It has no idea which real k a slot stands for.

So we may relabel k any way we like, as long as A (weights) and B (x) use the same relabelling. Here is the relabelling we choose. Inside each 32-wide block of K, lane t's slots stand for the real k values 8t, 8t+1, …, 8t+7:

The payoff is on both sides:

• x: a lane needs x[token g][8t … 8t+7], which is 8 consecutive BF16 values, or one 16-byte load. It covers both mma steps of the block.
• W: a lane needs rows g and g+8 at those same 8 k values. _pack stores exactly those 16 values next to each other: 2 × 16 bytes, one 16-byte load per step.

5.3 _pack, decoded index by index

640
641
642
643

def _pack(w: torch.Tensor) -> torch.Tensor:
    """[N, K] bf16 -> [RB, KB, step, lane(g,t), (j,h,e)]: each lane's 16B is one m16n8k16 A fragment."""
    n, k = w.shape
    return w.reshape(n // 16, 2, 8, k // 32, 4, 2, 2, 2).permute(0, 3, 5, 2, 4, 6, 1, 7).contiguous()

• 641The output's axes: [row block, K block, step, lane (g, t), 8 values (j, h, e)]. Each lane's 8 values (16 bytes) form one A fragment.
• 642W is [n, k], for example [4096, 8192] for the pooler.
• 643Two operations. reshape only names pieces of the two indices; permute reorders them; .contiguous() writes the reordered copy to memory.

The reshape(n // 16, 2, 8, k // 32, 4, 2, 2, 2) splits the two indices like this:

So row = 16·RB + 8·h + g and k = 32·KB + 8·t + 4·s + 2·j + e.

permute(0, 3, 5, 2, 4, 6, 1, 7) reorders the axes to [RB, KB, s, g, t, j, h, e]. Read the result from slowest to fastest:

1. For each row block,
2. for each K block,
3. for each mma step,
4. for each of the 32 lanes (lane = 4g + t),
5. come 8 values in the order (j, h, e). That order is - (row g, k+0, k+1): register 0; - (row g+8, k+0, k+1): register 1; - (row g, k+2, k+3): register 2; - (row g+8, k+2, k+3): register 3.

That is exactly the A-fragment register order (a0 a1 | a2 a3 | a4 a5 | a6 a7) from the hardware layout.

5.4 How the kernel walks the packed array

Here are the lines of gemv_kernel (chapter 6) that use this layout:

72
73
74
75
76
77
78
79
80
81
82
83
84

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

• 72g = lane >> 2 and t = lane & 3: the coordinates the layout was built around.
• 74–75K is counted in 32-wide blocks. Each warp starts at its own block kb0.
• 76The offset in 16-byte units. One (row block, K block) occupies 2 steps × 32 lanes = 64 units. Adding lane makes neighbouring lanes read neighbouring 16-byte units, so a warp reads 512 contiguous bytes.
• 79–84x pointers. Lane (g, t) reads token nb·8 + g at k offset 8t: the same 8 real k values its weights cover. Tokens past m are clamped to the last real token. That avoids reading out of bounds, and those results are never stored.

94
95

        a[u][0] = ld_stream(wp + (size_t)(i + u) * 64);
        a[u][1] = ld_stream(wp + (size_t)(i + u) * 64 + 32);

• 94–95Step 0's fragment is at +0; step 1's is 32 units (512 bytes) later. Moving to the next K block jumps 64 units.

105
106
107

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

• 105One 16-byte load of x: values 8t … 8t+7 of token g.
• 106Step 0 multiplies fragment a[u][0] by x values 0–3 (xv.x, xv.y).
• 107Step 1 multiplies a[u][1] by x values 4–7 (xv.z, xv.w). Swap these two and the answer is wrong; lab 2's --break mode shows exactly that.

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

Part 2 of the lab first prints what lanes 0–5 load: the (row, k) of each of their 8 values. Then it emulates one CTA of gemv_kernel, lane by lane, on the real packed layout:

• the 16-byte loads;
• two mma per K block, using the hardware's register layouts;
• the shared-memory reduction over warps;
• the store.

The result is checked against x @ W.T.

Exercises.

1. Run python3 lab2_fragments.py --break. It swaps the two halves of x between the steps, and the error jumps from about 1e-16 to about 1.
2. In part2(), change M to 13: NB becomes 2 and a second 8-token block appears. Change NW to 8 (with K = 512, K/32 = 16 blocks must divide by NW × U). The answer stays right.
3. Change U to 4: 16 blocks ÷ 4 warps = 4 blocks per warp, one trip of 4. Still right. Now think about U = 3: 4 blocks per warp don't split into trips of 3. The real kernel would read past its own slice and count blocks twice; the emulation stops with an IndexError for the same reason. That is why every configuration in the tables divides evenly. Lab 0 checks this for you.

Check yourself