3 · The strategy: one dispatcher, four engines

submission.py is not one clever kernel. It is a dispatcher that sends each (layer, m) to whichever of four engines was fastest when we measured it. This chapter shows the whole map first, then explains why each engine wins where it does, and what we tried that did not work.

3.1 The dispatcher

kernel_fn (lines 774–807) asks four questions in order. The first "yes" decides the engine.

kernel_fn(data, state):          # submission.py lines 774-807
  (layer, m) in TABLE ?  ──yes──▶  Triton split-K GEMM      pooler m=128, router m=128/256
        │ no
  (layer, m) in ZTABLE ? ──yes──▶  gemvz (12-bit weights)   12 shapes with m ≤ 32
        │ no
  m ≤ 32 and CTABLE entry
  is not None ?          ──yes──▶  gemv (BF16 weights)      router m ≤ 32, unpooler slice m ≤ 8,
        │ no                                                 and any other m ≤ 32 (default config)
  router ?               ──yes──▶  cuBLAS × 2 (hi, lo) + add
        │ no
  cuBLAS: x @ W[:n].T                                      everything else

The three tables it consults (TABLE, ZTABLE, CTABLE) are plain Python dictionaries at lines 586–623. They were generated from measurements (section 3.5).

3.2 All 30 scored shapes

A few things stand out:

• Blue (gemvz) covers most shapes with m ≤ 32, and that is where the big wins are (1.07–1.55×).
• Orange (gemv) is the router at m ≤ 32 (2.5–4.5× faster than SGLang) and the small unpooler slice at m ≤ 8 (1.7×).
• Purple (Triton) is the router at m = 128/256 (1.67×) and the pooler at m = 128 (1.09×).
• Gray (cuBLAS) is everything else. These shapes run at 0.98–1.00×, because they are SGLang's kernel.

3.3 Engine by engine: why each wins where it does

gemv: streaming each weight exactly once, at full speed

At small m, cuBLAS's GEMM kernels read weights at well under peak. SGLang's pooler at m = 1 takes 21.3 µs for 67 MB, which is 3.15 TB/s.

Our gemv_kernel gets close to the hardware's limit:

• Each warp streams its own slice of the repacked weights with 16-byte loads that are 100 % useful.
• Partial sums stay in registers until the very end.
• The uncompressed version reads the pooler (m = 1) in 17.1 µs (harness), about 3.9 TB/s.
• A hand-written kernel that only reads the same 67 MB, with the same fixed costs, needs about 16.1–16.4 µs. So the GEMV is within about 5 % of a pure memory copy.

gemvz: read fewer bytes

Once you are at the bandwidth limit, the only way to go faster is to move fewer bytes.

• gemvz_kernel reads weights stored losslessly in 12 bits instead of 16, which is 25 % fewer bytes.
• It rebuilds the BF16 values in registers with about 12 integer instructions per 8 weights. Those instructions hide under the memory time.
• The rare weights that don't fit the format (about 0.014 %) are fixed up exactly by an extra warp.

Compression loses on the tiny shapes, where fixed costs dominate, and at m = 32 for qkv_gate and o. There, the work per byte (decode + 4 mma blocks) and the repeated reads of x start to matter more than the weight bytes.

Triton split-K: one fused pass for the router, bigger tiles for the pooler

• Router at m = 128/256.
• SGLang runs two separate cuBLAS GEMMs (hi, then lo) plus an add.
• We run one kernel that loads a hi tile and a lo tile and accumulates both into the same FP32 accumulator.
• The router has only 256 rows, so K is split 16 or 8 ways to give every SM work.
• The partial sums are combined with atomics (chapter 9). Result: 1.67×.
• Pooler at m = 128.
• cuBLAS picks small 64×64 tiles here, so every weight byte is fetched by 2 CTAs and every x byte by 64.
• Our Triton kernel uses 64 weight rows × 128 tokens per CTA with K split 2 ways: 23.3 vs 25.4 µs.

cuBLAS: where nothing else won

For the dense layers at m = 128/256, and for qkv_gate and o at m = 32, we tried hard to beat cuBLAS and never did (table in 3.6). So those shapes simply call x @ W.T, the same thing SGLang does.

3.4 Where the remaining time is

The two verify loads at m = 128/256 are the biggest gap to the leaders: they sit at cuBLAS parity. Rounds 4, 6 and 7 of the project were attempts to beat cuBLAS there with our own Hopper (wgmma) kernels, and none succeeded.

3.5 How the configurations were chosen

Every kernel has knobs: warps per CTA, K-splits, loads in flight, pipeline mode, and tile sizes for Triton. They were chosen by exhaustive sweeps on a rented H200, timed like the judge: CUDA-graph replay, cold L2, paired against SGLang. build.py then wrote the winners into the tables, using two rules:

• A plain-GEMV entry goes into CTABLE only if it beats SGLang by at least 3 %. Otherwise the entry is None, which means "use cuBLAS".
• A compressed entry goes into ZTABLE only if it beats both the plain GEMV and cuBLAS by at least 2 %.

3.6 What we tried that did not make it

The details of each are in the project report. Chapter 10 shows how to run your own experiments on the real hardware.

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

Before you run the lab, predict which engine handles the test shapes m = 5 unpooler, m = 9 router and m = 33 o. Then check your answers in the second table of the output. The lab also prints each GEMV's launch geometry: CTAs, threads, and how much of K each warp covers.

Check yourself