7 · The 12-bit lossless weight format (pack_z and decz)

At m ≤ 16 the GEMV runs at nearly full memory bandwidth, so it can only go faster by reading fewer bytes. This chapter shows how each BF16 weight is stored in 12 bits with no loss of information. It reads the compressor pack_z (lines 650–708) and the 12-instruction decoder decz (lines 246–273). Lab 3 checks every claim bit for bit.

7.1 Anatomy of a BF16 number

A BF16 value is (−1)^sign × 2^(exponent − 127) × 1.mantissa. Our weights are drawn like the judge's: normal with standard deviation 1/√K ≈ 0.011. Such numbers span only a few powers of two, so their exponents are highly repetitive. The mantissa bits, by contrast, look random and cannot be compressed.

So the format compresses only the high byte: the sign plus the top 7 exponent bits, which we call e7.

Here is the distribution of e7 over 8.4 million weights drawn like the judge's (lab 3, part 1):

  e7= 53  |w| in [2^-21, 2^-19)    0.0098%
  e7= 54  |w| in [2^-19, 2^-17)    0.0416%   <- window
  e7= 55  |w| in [2^-17, 2^-15)    0.1659%   <- window
  e7= 56  |w| in [2^-15, 2^-13)    0.6512%   <- window
  e7= 57  |w| in [2^-13, 2^-11)    2.6323%   <- window
  e7= 58  |w| in [2^-11, 2^-9)    10.4717%   <- window
  e7= 59  |w| in [2^-9, 2^-7)     37.9717%   <- window
  e7= 60  |w| in [2^-7, 2^-5)     47.5746%   <- window
  e7= 61  |w| in [2^-5, 2^-3)      0.4774%   <- window
  window: e7 in [54, 61]  covers 99.9865% of weights

Eight consecutive values of e7 (a 3-bit code) cover 99.987 % of the weights. Only tiny weights, below 2⁻¹⁹, and exact zeros fall outside.

7.2 The format

Decoding never needs a branch: rebuild the high byte from the sign and base + code, and glue it onto the low byte. For the 0.013 % of weights outside the window, that gives a slightly wrong value. A separate warp adds correction × x[k] to the output (chapter 8). The sum is exactly the original weight times x[k], so the format is lossless.

Here is one weight at a time (lab 3, part 1b, base = 54):

  weight +0.0123 -> BF16 bits 0011110001001010 (0x3c4a) = +0.0123291
    sign 0   e7 (bits 14-8)  60 = 0111100   low byte (bits 7-0) 0x4a = 01001010
    stored: low byte 0x4a + nibble (sign 0, code e7-base = 60-54 = 6)  = 12 bits
    rebuilt: sign | (base+code)<<8 | low = 0x3c4a = +0.0123291  exact
  weight +3.1e-07 -> BF16 bits 0011010010100110 (0x34a6) = +3.09199e-07
    sign 0   e7 (bits 14-8)  52 = 0110100   low byte (bits 7-0) 0xa6 = 10100110
    stored: low byte 0xa6 + nibble (sign 0, code e7-base = 52-54 = -2 -> clamped to 0)  = 12 bits
    rebuilt: sign | (base+code)<<8 | low = 0x36a6 = +4.94719e-06  -> needs an exception: -4.63799e-06

How the bytes are laid out. The weights keep exactly the fragment order of chapter 5, but the 16 rows × 64 k of one "group" are split into three 512-byte sections:

Within each fragment's 32-bit code word Q, the bits are placed so the decoder needs as few instructions as possible:

The codes of weights 0–3 fill the low 16 bits and those of weights 4–7 the high 16 bits. The signs of weights 0–3 sit at bit 7 of each byte, which is the BF16 sign position of each decoded high byte.

7.3 pack_z, line by line

650
651
652
653
654
655
656
657

def pack_z(w: torch.Tensor, nwin: int = 8):
    """bf16 [N, K] -> dict(w=packed uint8 [RB, K/64, 3, 32, 16], esc=int32 [N, capr, 2] (k, fp32 delta bits),
    capr, TL, TH, base, nesc). Lossless: decoded + delta == w exactly."""
    assert w.dtype == torch.bfloat16 and nwin == 8
    n, k = w.shape
    assert n % 16 == 0 and k % 64 == 0
    rb, kb = n // 16, k // 32
    dev = w.device

• 650Input: a BF16 matrix [N, K]. nwin = 8 window values means a 3-bit code.
• 651–652The output:
  • the packed bytes, [row blocks, K/64 groups, 3 sections, 32 lanes, 16 bytes];
  • the exception table [N, capr, 2], holding (k, FP32 bits of the correction);
  • the decoder's lookup words.
• 653–657Sanity checks, sizes and the device: N must divide into 16-row blocks, and K into 64-wide groups.

658
659
660
661
662
663
664
665
666

    bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF
    e7 = (bits >> 8) & 0x7F
    h = torch.bincount(e7.flatten(), minlength=128)
    c = torch.cat([torch.zeros(1, dtype=h.dtype, device=dev), torch.cumsum(h, 0)])
    cover = c[nwin:] - c[:-nwin]
    base = int(torch.argmax(cover))
    code = (e7 - base).clamp(0, nwin - 1)
    sign = bits >> 15
    low = bits & 0xFF

• 658The raw 16-bit patterns. View BF16 as int16 (same bits), widen to int32 and mask, so that a set sign bit doesn't make the number negative.
• 659e7 = bits 14–8: the top 7 of the 8 exponent bits.
• 660A histogram of e7 over all weights (128 possible values).
• 661Prefix sums: c[i] = how many weights have e7 < i.
• 662cover[b] = how many weights have b ≤ e7 < b+8, a sliding window computed in one subtraction.
• 663base = the window start that covers the most weights. For the 8.4 M sample above that is 54.
• 664The code, clamped into 0…7. Weights below the window get 0, and weights above it get 7.
• 665–666The sign bit and the raw low byte.

667
668
669
670

    dbits = (sign << 15) | ((base + code) << 8) | low
    dbits = torch.where(dbits >= 32768, dbits - 65536, dbits).to(torch.int16)
    dec = dbits.view(torch.bfloat16)
    del dbits, e7, bits

• 667The bits the decoder will produce: the sign, (base + code) as the top 7 exponent bits, and the low byte. This equals the original exactly whenever e7 was inside the window.
• 668–669Wrap into the signed 16-bit range and reinterpret as BF16. dec is the decoded matrix.
• 670Free the big temporaries: 33 M elements each for the pooler.

671
672
673
674
675
676
677
678
679
680

    # exception lists: per row, (k, fp32 value) entries summing to w - decoded. The fp32 difference is exact
    # unless |w| is ~2^-16 below the window; then a second entry carries the fp32 residual.
    d64 = w.double() - dec.double()
    d32 = d64.float()
    res = (d64 - d32.double()).float()
    i1, i2 = (d32 != 0).nonzero(), (res != 0).nonzero()
    rows = torch.cat([i1[:, 0], i2[:, 0]])
    ks = torch.cat([i1[:, 1], i2[:, 1]])
    vals = torch.cat([d32[i1[:, 0], i1[:, 1]], res[i2[:, 0], i2[:, 1]]])
    del d64, d32, res

• 673The exact difference, in FP64. Both operands are BF16 values, so FP64 holds their difference exactly.
• 674Its nearest FP32 value. This is the correction the kernel adds.
• 675Whatever FP32 could not hold. That is nonzero only if a weight is more than about 2¹⁶ times smaller than the bottom of the window, because then the difference needs more than FP32's 24 significant bits. A second entry then carries the remainder. In the 8.4 M sample, no weight needs one.
• 676–679Collect every nonzero correction as (row, k, value).

681
682
683
684
685
686
687
688
689

    order = torch.argsort(rows * k + ks, stable=True)
    rows, ks, vals = rows[order], ks[order], vals[order]
    cnt = torch.bincount(rows, minlength=n)
    capr = max(1, int(cnt.max()) if len(rows) else 1)
    start = torch.cumsum(cnt, 0) - cnt
    pos = torch.arange(len(rows), device=dev) - start[rows]
    esc = torch.zeros(n, capr, 2, dtype=torch.int32, device=dev)
    esc[rows, pos, 0] = ks.to(torch.int32)
    esc[rows, pos, 1] = vals.contiguous().view(torch.int32)

• 681–682Sort by row, then by k.
• 683–684Count the corrections per row. capr is the longest row's count, at least 1. Every row gets capr slots, so the kernel finds row r's list at r × capr with no extra index. In the 8.4 M sample, capr = 6.
• 685–686Each correction's slot within its row.
• 687–689The table esc[N, capr, 2] of int32: (k, the FP32 bits of the correction). Unused slots stay (0, 0), and the kernel skips entries whose value is 0.

691
692
693

    def perm(a):  # [N, K] -> [RB, G, f(4), lane(32), i(8)] in mma A-fragment order
        p = a.reshape(rb, 2, 8, kb, 4, 2, 2, 2).permute(0, 3, 5, 2, 4, 6, 1, 7)  # RB, KB, j, g, t, h, rh, e
        return p.reshape(rb, kb // 2, 4, 32, 8)

• 692The same reshape/permute as _pack, giving [RB, KB, step, g, t, j, h, e].
• 693Then pairs of 32-wide K blocks merge into 64-wide groups G, each holding 4 fragments: f = (KB mod 2) × 2 + step. The result is [RB, G, f, lane, 8 values].

695
696

    lowp = perm(low.to(torch.uint8))
    A = lowp.permute(0, 1, 3, 2, 4).reshape(rb, kb // 2, 32, 2, 16).permute(0, 1, 3, 2, 4)  # RB, G, sec, lane, 16

• 695Low bytes in fragment order.
• 696Regroup into [RB, G, section, lane, 16 bytes]. Section 0 holds fragments 0 and 1, and section 1 holds fragments 2 and 3. Each lane's 16 bytes are contiguous and the 32 lanes are adjacent, so a warp's load is again 512 contiguous bytes.

697
698
699
700
701
702
703

    codp = perm(code.to(torch.int64))
    sgnp = perm(sign.to(torch.int64))
    csh = torch.tensor([0, 4, 8, 12, 16, 20, 24, 28], device=dev, dtype=torch.int64)
    ssh = torch.tensor([7, 15, 23, 31, 3, 11, 19, 27], device=dev, dtype=torch.int64)
    q = (codp << csh).sum(-1) + (sgnp << ssh).sum(-1)  # [RB, G, 4, 32]
    q = torch.where(q >= 2 ** 31, q - 2 ** 32, q).to(torch.int32)
    B = q.permute(0, 1, 3, 2).contiguous().view(torch.uint8).reshape(rb, kb // 2, 1, 32, 16)

• 697–698Codes and signs in fragment order.
• 699The code of weight i goes to bits 4i … 4i+2 of the fragment's word.
• 700The sign of weight i goes to bit [7, 15, 23, 31, 3, 11, 19, 27][i]: the spare top bit of each nibble, arranged as in the table in 7.2.
• 701One 32-bit word per (row block, group, fragment, lane).
• 702Wrap into the signed int32 range, because torch has no uint32 arithmetic.
• 703The four words of a lane (one per fragment) become its 16 bytes of section 2.

704
705
706
707
708

    packed = torch.cat([A, B], 2).contiguous()
    lut = [base + j for j in range(8)]
    TL = sum(lut[j] << (8 * j) for j in range(4))
    TH = sum(lut[4 + j] << (8 * j) for j in range(4))
    return {"w": packed, "esc": esc.contiguous(), "capr": capr, "TL": TL, "TH": TH, "base": base, "nesc": int(len(rows))}

• 704[RB, G, 3, 32, 16]: 1,536 bytes per 16 rows × 64 k, against 2,048 for BF16.
• 705–707The lookup table: the 8 window values base … base+7, one per byte. TL holds the first four and TH the last four. The kernel receives them as two plain 32-bit arguments.
• 708Return everything. nesc is the exception count, kept for reports.

7.4 The decoder, line by line

246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

__device__ __forceinline__ uint4 ldz_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;
}
__device__ __forceinline__ void mmaz(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));
}
__device__ __forceinline__ uint32_t prmtz(uint32_t a, uint32_t b, uint32_t sel) {
  uint32_t r; asm("prmt.b32 %0, %1, %2, %3;" : "=r"(r) : "r"(a), "r"(b), "r"(sel)); return r;
}
__device__ __forceinline__ uint32_t or_and(uint32_t a, uint32_t b, uint32_t c) {  // a | (b & c)
  uint32_t r; asm("lop3.b32 %0, %1, %2, %3, 0xF8;" : "=r"(r) : "r"(a), "r"(b), "r"(c)); return r;
}

• 246–256ldz_stream and mmaz are identical copies of ld_stream and mma16816 from chapter 6. The compressed kernel began life in its own file, and build.py pastes both files into one string; distinct names avoid a duplicate definition.
• 257–259prmt.b32 d, a, b, sel, "permute bytes":
  • Treat {b, a} as 8 bytes (a = bytes 0–3, b = bytes 4–7).
  • Each of the low 4 nibbles of sel picks one of those bytes for one byte of the result.
  • If a nibble's top bit is set, the chosen byte's sign is replicated instead; we always keep that bit clear.
  So one instruction performs any byte shuffle, or a lookup into an 8-entry table.
• 260–262lop3.b32 d, a, b, c, lut computes any bitwise function of three inputs. lut is the function's truth table, written as the function applied to the constants a = 0xF0, b = 0xCC, c = 0xAA. Here 0xF0 | (0xCC & 0xAA) = 0xF0 | 0x88 = 0xF8, so or_and computes a | (b & c) in one instruction.

263
264
265
266
267
268
269
270
271
272
273

// S0/S1: low bytes of weights 0-3 / 4-7 of one A fragment; Q: their codes + signs (layout in the module doc)
__device__ __forceinline__ uint4 decz(uint32_t S0, uint32_t S1, uint32_t Q, uint32_t TL, uint32_t TH) {
  uint32_t H0 = or_and(prmtz(TL, TH, Q & 0x7777u), Q, 0x80808080u);
  uint32_t H1 = or_and(prmtz(TL, TH, (Q >> 16) & 0x7777u), Q << 4, 0x80808080u);
  uint4 a;
  a.x = prmtz(S0, H0, 0x5140u);
  a.y = prmtz(S0, H0, 0x7362u);
  a.z = prmtz(S1, H1, 0x5140u);
  a.w = prmtz(S1, H1, 0x7362u);
  return a;
}

• 264Inputs:
  • S0: the low bytes of weights 0–3; S1: those of weights 4–7;
  • Q: the code/sign word;
  • TL, TH: the window values.
• 265High bytes of weights 0–3.
  • Q & 0x7777 keeps the four 3-bit codes in the low 16 bits.
  • prmt(TL, TH, …) uses each code as a selector, so byte i becomes the window value base + code_i: a 7-bit exponent top.
  • or_and(…, Q, 0x80808080) then copies Q's bits 7, 15, 23 and 31, the signs of weights 0–3, into the top bit of each byte.
  The result H0 holds four complete high bytes.
• 266The same for weights 4–7. Their codes are in the high half (Q >> 16), and their signs sit at bits 3, 11, 19 and 27, so Q << 4 moves them to bits 7, 15, 23 and 31.
• 268prmt(S0, H0, 0x5140). The selector nibbles, from the lowest, are 0, 4, 1, 5, so the result is [low₀, high₀, low₁, high₁]: BF16 weight 0 in the low half and weight 1 in the high half. That is exactly register "a0 a1" of the A fragment.
• 269Selectors 2, 6, 3, 7 give weights 2 and 3.
• 270–271Weights 4–5 and 6–7, from S1 and H1.

The total is about 12 integer instructions per 8 weights per lane: 2 lookups, 2 merges, 4 interleaves and 4 shifts or masks. The memory system needs roughly the same time to deliver those 12 bytes, so the decoding hides under the loads.

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

What the lab does:

• Part 1 prints the histogram above, the window, the exception count, and the size for windows of 2, 4, 8 and 16 values.
• Part 1b encodes four sample weights by hand.
• Part 2 packs a matrix with a numpy port of pack_z, then decodes 32,768 fragments with emulated prmt/lop3 and checks every bit.
• Part 3 is chapter 8's exception warp in miniature.

Exercises.

1. Look at the window-size sweep. Why is 8 the sweet spot? (4 values leave 1.35 % of weights outside, which is far too many exceptions; 16 values cost a 13th bit.)
2. Replace the weights with a heavier-tailed distribution, rng.standard_t(3, (N, K)) / np.sqrt(K), and watch the exception rate change.
3. In decz, change 0x5140 to 0x4051. Which weights come out swapped?

Check yourself