10 · Measure, experiment, and try your own changes

You now know what every line does. This chapter shows how the numbers in this guide were measured, how to read the judge's profile, and how to run real experiments on the judge's own H200 without owning a GPU. It ends with four ready-made experiments, each with a prediction you can check.

10.1 How we measure: bench.py

bench.py (in the project directory) copies the judge's method as closely as we could. Here is its core:

def capture(fn, data, state):                 # record kernel_fn into a CUDA graph
    for _ in range(3):
        fn(data, state)                       # warm up (compiles, cuBLAS heuristics)
    torch.cuda.synchronize()
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        out = fn(data, state)
    return g, out

def time_pair(ga, gb, flush, reps):           # SGLang's graph vs ours, alternating
    for r in range(reps):
        order = (ga, gb) if r % 2 == 0 else (gb, ga)   # swap who goes first
        for g in order:
            flush.sum()                       # read 256 MB: evicts the 50 MB L2 -> cold weights
            s.record(); g.replay(); e.record(); e.synchronize()
            ...
# speedup = median over repetitions of (sglang_us - 4) / (ours_us - 4)

Before any submission, the routine was: bench.py test (9/9), then sweep.py (420/420), then bench.py time against the previous version on the same machine, then okc check, then okc submit.

10.2 Reading the judge's profile

Every passing okc run saves profiles/<file>.<job>.txt next to your file. It holds the per-shape table and, for each shape, the kernels that ran:

layer=pooler load=decode 1 m=1                              21.31      13.77   1.547x      13.99
...
profile 2: layer=pooler load=decode 1 m=1
kernel                                                           GPU time  share  calls
void gemvz_kernel<16, 1, 1, 2>(__nv_bfloat16 const*, uint...      13.9 µs 100.0%      1

profile 10: layer=qkv_gate load=decode 32 m=32
nvjet_sm90_tst_64x32_64x16_4x1_v_bz_splitK_TNT                     8.9 µs  82.8%      1
void cublasLt::splitKreduce_kernel<32, 16, int, float, __...       1.9 µs  17.2%      1

• gemvz_kernel<16, 1, 1, 2> is our compressed GEMV with NW=16, UG=1, NB=1, PIPE=2, exactly ZTABLE[('pooler', 1)].
• nvjet_sm90_… is cuBLAS's family of Hopper GEMM kernels. The numbers in the name describe its tile configuration. splitK means it split K and needs a separate splitKreduce_kernel launch, which costs 1.9 µs here. That second launch is one reason a fused split-K is attractive.
• _proj_tma_kernel is our Triton kernel (chapter 9).

The ".log" file next to it holds anything your code printed, which is handy for debugging.

10.3 How submission.py was produced

submission.py was generated by build.py, which pastes together:

• sub_template.py (the Python skeleton);
• the CUDA source from gemv_cuda.py and gemvz_cuda.py;
• the Triton code from _triton_part.py;
• the tuned tables from cuda_tuned5.json and cuda_z.json.

It also writes the C3/Z3 instantiation lists.

10.4 Your turn: experiments on the real H200 with okc check

okc check runs a file exactly like a submission: the 9 tests, then timing against SGLang on the judge's H200. But it stays off the leaderboard and doesn't use up a scored run. You get 20 checks per day (reset at 00:00 UTC; okc problems shows what's left). A check takes a few minutes: about 1 minute of compiling, then tests and timing.

export PATH=~/projects/okc/bin:~/.local/bin:$PATH
cd ~/projects/okc/muse_bf16_proj
cp submission.py try_a.py            # keep line 1 ("# okc: problem=muse_bf16_proj") as it is
#   ... edit try_a.py as described below ...
python3 learn/labs/lab0_dispatch.py try_a.py   # optional: confirm the dispatch changed as intended
okc check try_a.py                   # waits and prints the per-shape table and the speedup

Compare the output with the "ours µs" column of chapter 3's table, or with the shipped run's profile profiles/submission.7e998918e84d4ea8.txt. The judge's noise is about ±1.5 %, so ignore smaller differences. Submitting identical code again returns the earlier result, so every check needs a real change.

Experiment A: turn compression off for the pooler at m = 1 and 8

Edit. Delete these two lines from ZTABLE (lines 611–612):

 ('pooler', 1): (16, 1, 1, 2),
 ('pooler', 8): (16, 1, 1, 2),

What happens. kernel_fn falls through to CTABLE[('pooler', 1)] = (8, 1, 2, 3), the plain BF16 GEMV of chapter 6. Its kernel C3(8, 2, 1, false, 3) is already compiled in.

Prediction. pooler m=1 goes from about 13.8 to about 17 µs, and m=8 from about 14.1 to about 17.3 µs. These are the harness numbers of chapter 3.3 (17.12 and 17.34 µs). The score should drop from 1.343 to about 1.31×: lab 1 can compute it if you edit OURS.

What it teaches. This experiment puts a number on the 12-bit format of chapters 7–8: 25 % fewer bytes, about 20 % less time.

Experiment B: take split-K away from the router at m = 1

Edit. In CTABLE, change ('router', 1): (16, 4, 4, 3) to ('router', 1): (16, 1, 4, 3) (line 599).

What happens. The same kernel C3(16, 4, 1, true, 3) runs with S = 1, so 16 CTAs instead of 64. Every CTA streams 4 times more K, and 116 of the 132 SMs sit idle.

Prediction (not measured). Clearly slower. We would guess 8–12 µs instead of 5.4, because a single SM can only pull data in at a limited rate. That would be about −2 % on the overall score.

What it teaches. "Fill the machine" (chapter 6.7): split-K exists for shapes with few rows.

Experiment C: the router at m = 128 with the serial epilogue

Edit. In TABLE, change the last field of ('router', 128) from 1 to 0, giving (64, 128, 128, 16, 4, 3, 1, 0) (line 622).

What happens. The same Triton kernel, but the last of the 16 programs per tile now reads all 16 partial tiles from ws one after another (chapter 9.4).

Prediction. About 12 µs instead of 9.9. Round 4 measured 12.4 µs for this serial version before switching to atomics.

What it teaches. The cost of a serial reduction tail, and why the atomic epilogue exists.

Experiment D: add a kernel variant (and meet the instantiation list)

Edit, part 1. In ZTABLE, change ('qkv_gate', 1): (16, 1, 1, 2) to ('qkv_gate', 1): (8, 1, 2, 3) (line 615).

Edit, part 2. Add a new line Z3(8, 2, 1, 3) just before TORCH_CHECK(false, "unsupported gemvz config"); (line 471). The macro's arguments are (NW, UG, NB, PIPE): 8 warps, 2 groups per trip, NB = 1 for m = 1, PIPE 3.

What happens if you skip part 2. The kernel variant was never compiled. prepare()'s warm-up calls it, and the run fails with "unsupported gemvz config" (chapter 6.9). That wastes a check, so add the line.

Prediction. About 10.1 µs instead of 9.6. Our sweep measured 10.08 vs 9.54 µs (harness) for these two configurations, which is why 16 warps won.

What it teaches.

• Every configuration is a separately compiled template.
• How a tuned table and the compiled list must agree.
• A small taste of the tuning landscape: one config changed, about 5 % slower.

10.5 Optional: running on your own rented GPU

You can also rent an H100 or H200 on RunPod (H200: about $5 per hour) and use bench.py directly. That gives faster turnaround and lets you sweep many configurations, but it costs money and needs set-up.

The project has a helper, ~/projects/okc/pod.sh:

• OKC_POD=<pod id> ./pod.sh sync muse_bf16_proj copies the project to the pod;
• ./pod.sh ssh logs in.

On the pod:

cd /workspace/okc/muse_bf16_proj && source /etc/profile.d/okc.sh
python bench.py test --module try_a            # 9 correctness tests
python bench.py time pooler --module try_a     # pooler at all 6 loads vs SGLang

Always delete the pod when you are done (runpodctl pod delete <id>), and check that runpodctl pod list prints [].

10.6 Where the remaining time is: open problems

These are the gaps between today's numbers and the leaders' (about 1.64×), grouped by the part of the code they would change:

The project report has the full history of what was tried and measured.