1 · The problem, the judge and the speed of light
Before reading any kernel code, you need three things: what is computed, how it is timed, and the fastest time physics allows. This chapter covers all three. It ends with a lab that recomputes our official score from the per-shape times.
1.1 What is computed
Every call to our code is one matrix product: a linear layer with no bias.
out[i, n] = sum over k of x[i, k] * W[n, k] (in PyTorch: out = x @ W.T)
xholds the activations: m tokens × K features, in BF16 (a 16-bit float, explained in chapter 7).Wholds the weights: N output features × K input features, in BF16. The router's weights are FP32.outis m × N, in BF16. The router's output is FP32.
There are five weight matrices. They are generated once by the judge and handed to our prepare() function before any timing starts.
| layer | N × K | weight bytes | role in the model |
|---|---|---|---|
qkv_gate |
2304 × 8192 | 37.7 MB | attention: query, key, value and gate projections fused into one matrix |
o |
8192 × 1024 | 16.8 MB | attention output projection |
pooler |
4096 × 8192 | 67.1 MB | hidden size → MoE width |
unpooler |
8192 × 4096 | 67.1 MB | MoE width → hidden size. At ≤ 8 tokens only this GPU's slice, the first 1024 rows (8.4 MB), is computed. |
router |
256 × 8192 | 8.4 MB (FP32) | one score per expert, 256 experts |
1.2 Six loads, 30 shapes
The judge runs the projections at the batch sizes a serving system actually sees.
| load | m (tokens per call) | what it models |
|---|---|---|
| decode 1 | 1 | one user, generating one token per step |
| decode 16 | 16 | 16 users at once |
| decode 32 | 32 | 32 users at once |
| verify 1 | 8 | one user, checking an 8-token speculative draft in one step |
| verify 16 | 128 | 16 users × 8 draft tokens |
| verify 32 | 256 | 32 users × 8 draft tokens |
Each load runs all five layers, so there are 6 × 5 = 30 scored shapes. There are also 9 extra test shapes with awkward sizes (m = 5, 7, 9, 33, 130, 512, 1000). They are only checked for correctness, but the code must handle them.
1.3 How the score is computed
- A load's time is the sum of its five shapes' times.
- A load's speedup is SGLang's load time divided by ours.
- The score is the geometric mean of the six load speedups.
(Each shape also has a weight of 62, its number of runs per step, but it is the same for every shape, so it cancels.)
Key idea. Every load counts the same. Shaving 10 % off decode 1 (38 µs in total) is worth exactly as much as shaving 10 % off verify 32 (95 µs). Within a load a microsecond is a microsecond: a 6 µs router counts as much per µs saved as a 14 µs pooler.
1.4 How a run is timed
- CUDA graph replay.
- Our
kernel_fnis recorded once into a CUDA graph, and the GPU then replays the recording. - The CPU cost of launching kernels disappears.
- What remains is the GPU's own fixed cost of starting and finishing each kernel, about 1–2 µs per kernel.
- Cold L2.
- Before every replay the judge reads a buffer bigger than the GPU's 50 MB L2 cache.
- So our weights are never in cache: they come from the GPU's main memory (HBM) every single time.
- Keeping them in cache on purpose is against the rules.
- Paired, median.
- Our replays and SGLang's alternate on the same inputs.
- The speedup is the median of the paired ratios.
- A run reports a 95 % interval of at least ±1.5 %, so smaller gains can be noise.
- About 4 µs subtracted. The fixed cost of a graph replay is subtracted from both sides.
1.5 What counts as correct
- The relative L2 error against an FP64 reference must be at most 1 %, or 0.01 % for the router, whose outputs choose experts.
- Probe rows. About one row in eight of
xhas a single nonzero entry, a power of two. That output row is therefore an exact scaled copy of a weight column, and it must match within 0.1 %. - So weights must be used at their own precision.
- That rules out lossy compression for every layer, and plain BF16 for the router.
- No single row may be off by more than 5× the bound, and there must be no NaN or Inf.
1.6 Rules that shaped the code
| rule | what it caused in submission.py |
|---|---|
prepare(weights) runs once, untimed, for up to 300 s |
We repack and compress the weights there, compile the CUDA code (about 50 s), and warm up Triton and cuBLAS. |
| The timed call is replayed as a CUDA graph: no host syncs, no branching on data values | kernel_fn only branches on shapes. Split-K counters are reset on the GPU, not from Python. |
| Fresh inputs every call | Nothing derived from x is ever cached. |
| Lossless compression is allowed; the router may run as BF16 hi + lo | This gave us the 12-bit format (chapter 7) and the fused hi + lo router (chapter 9). |
| 10 minutes per run, compilation included | Only the kernel variants actually used are compiled (chapter 6, the C3 list). |
1.7 The speed of light
A shape can never be faster than the slower of two limits:
- moving the bytes at the H200's memory bandwidth, 4.8 TB/s;
- doing the math at the tensor cores' peak, 989 TFLOP/s for BF16.
floor = max( bytes moved / 4.8 TB/s , FLOPs / 989 TFLOP/s )
Worked example: pooler at m = 1.
| term | calculation | time |
|---|---|---|
| bytes | 4096 × 8192 weights × 2 bytes = 67.1 MB (x and out add only 24 KB) | 67.1 MB ÷ 4.8 TB/s = 13.98 µs |
| math | 2 × 1 × 4096 × 8192 = 67 MFLOP (one multiply and one add per weight per token) | 67 MFLOP ÷ 989 TFLOP/s = 0.07 µs |
The memory term is 200 times larger. In general, every 2-byte weight is used for m multiply-adds, which is 2m FLOPs. So a projection does about m FLOPs per byte; this ratio is its arithmetic intensity.
The H200 can do 989 / 4.8 ≈ 206 FLOPs per byte it reads. Any kernel doing fewer FLOPs per byte than that is limited by memory, not math.
Key idea. For m ≤ 32, a projection is a memory copy with a little math attached. Everything in chapters 5–8 is about two things: streaming bytes at full speed, and streaming fewer of them. Only at m = 256 do the big layers cross the ridge and become math-bound. There, tensor-core efficiency is what matters, and cuBLAS is hard to beat.
1.8 Where we stand
| load | SGLang (µs) | ours (µs) | floor, sum of 5 shapes (µs) | speedup |
|---|---|---|---|---|
| decode 1 | 59.49 | 38.25 | 28.85 | 1.555× |
| decode 16 | 75.94 | 54.03 | 41.42 | 1.406× |
| decode 32 | 76.46 | 63.77 | 41.78 | 1.199× |
| verify 1 | 75.90 | 39.98 | 28.99 | 1.898× |
| verify 16 | 86.34 | 77.91 | 43.90 | 1.108× |
| verify 32 | 100.84 | 94.80 | 51.67 | 1.064× |
The geometric mean of the six speedups is 1.343×.
One surprise is worth noticing: pooler at m = 1 takes 13.77 µs, which is below its 13.99 µs "speed of light". The floor assumes 16 bits per weight; we read only 12 (chapter 7). The same trick puts several other decode shapes close to their floor.
cd ~/projects/okc/muse_bf16_proj/learn/labs
python3 lab1_roofline.py~/projects/okc/muse_bf16_proj/learn/labs/lab1_roofline.py; you can also view or download it here.In the lab output, look at the ours/floor column. It shows where we are already near the limit (pooler, unpooler at m = 16) and where a lot of time is still on the table (router, the unpooler slice, every m = 128/256 shape).
Check yourself
For qkv_gate at m = 16, which term is bigger, bytes or math, and by how much?
Bytes: 37.7 MB (plus a little x and output) ÷ 4.8 TB/s ≈ 7.9 µs.
Math: 2 × 16 × 2304 × 8192 = 0.60 GFLOP ÷ 989 TFLOP/s ≈ 0.61 µs.
Bytes win by about 13×, so the shape is memory-bound.
Why can't we keep the weights in the 50 MB L2 cache between calls?
The judge reads a buffer larger than L2 before every timed replay, so the cache is always cold. Pinning weights in L2 to survive that flush is explicitly forbidden: it would measure something that never happens in real serving, where many other layers run between two calls of this one.
If decode 1 went from 38.25 µs to 35 µs, what would the score become?
The decode 1 speedup goes from 59.49/38.25 = 1.555 to 59.49/35 = 1.700. The score is a geometric mean of six loads, so it is multiplied by (1.700/1.555)^(1/6) ≈ 1.015, giving 1.343 → about 1.363×. You can check this with the lab by editing OURS.