⌂ Guide home · Start here1 · The problem and the judge →

How our muse_bf16_proj kernels work: a hands-on guide

This guide walks through our Open Kernel Challenge submission for muse_bf16_proj one line at a time. It is written for someone who knows Python and roughly what a GPU kernel is. Every other idea is explained the first time it comes up. Each chapter ends with something to run yourself: most of it on your laptop, and some of it on the real H200 through okc check.

The submission currently scores 1.343× against SGLang on the official check run, and 1.331× on the scored run.

The whole thing in one minute

The job.

The physics.

The code. submission.py is a dispatcher. For each (layer, m) it picks one of four engines:

engine what it is where it runs
gemv_kernel A hand-written CUDA tensor-core matrix-vector kernel (GEMV). The weights are pre-shuffled so every thread reads exactly what the tensor core needs. router at m ≤ 32; the small unpooler slice at m ≤ 8
gemvz_kernel The same idea, reading weights stored losslessly in 12 bits instead of 16. most other shapes with m ≤ 32
Triton split-K GEMM A Triton kernel that splits the long K dimension across many thread blocks. router at m = 128/256; pooler at m = 128
cuBLAS NVIDIA's library, the same one SGLang uses. everywhere nothing we wrote was faster

The router has FP32 weights. We run it as two BF16 matrices ("hi" and "lo") whose sum is accurate enough, fused into a single kernel.

Chapters

# chapter you will learn code covered hands-on
1 The problem and the judge what is measured, how, and the "speed of light" none yet lab 1
2 Just enough GPU SMs, warps, memory bandwidth, tensor-core fragments none yet lab 2, part 1
3 The strategy which engine runs each shape, and why the dispatch lab 0
4 The Python side prepare(), kernel_fn(), the tables, compiling CUDA from Python lines 1–40, 586–637, 711–807 lab 0
5 Repacking weights the K-permutation trick behind _pack lines 640–643 lab 2, part 2
6 The GEMV kernel gemv_kernel: loads, mma, reductions, split-K lines 41–244 lab 5
7 The 12-bit format pack_z and the prmt/lop3 decoder lines 246–273, 646–708 lab 3
8 The compressed GEMV gemvz_kernel and its exception warp lines 275–472 lab 3, part 3
9 Router and Triton the hi + lo split, _proj_tma_kernel, split-K epilogues lines 475–583 labs 4 and 5
10 Measure and experiment bench.py, profiles, four real-hardware experiments none okc check
Glossary every term in one place

Chapters 1–10 together cover all 807 lines of submission.py.

How to use this guide

  1. Read in order. Chapters 1–3 are the "why", chapters 4–9 the "how, line by line", and chapter 10 the "now change something".
  2. Keep submission.py open. All line numbers refer to ~/projects/okc/muse_bf16_proj/submission.py. That file is identical to submission_r5a.py, the build that was scored.
  3. Run the labs as you go. They are small Python programs that emulate the GPU code exactly, so you can poke at it on a laptop:
cd ~/projects/okc/muse_bf16_proj/learn/labs
python3 -m pip install numpy   # labs 2, 3 and 4 need numpy; labs 0, 1 and 5 need only Python
python3 lab1_roofline.py
lab what it shows chapter
lab0_dispatch.py which engine and launch geometry each shape gets, read from submission.py 3, 4
lab1_roofline.py the floor of each shape, and the official score recomputed from per-shape times 1
lab2_fragments.py tensor-core fragment layouts; _pack plus one CTA of gemv_kernel emulated 2, 5
lab3_compress.py the 12-bit format: histogram, pack_z, a bit-exact decz, exceptions 7, 8
lab4_router.py why the router needs hi + lo 9
lab5_splitk.py the split-K "last CTA reduces" and atomic protocols 6, 9
  1. Try changes on the real GPU (chapter 10). okc check runs a modified copy of submission.py on the judge's own H200 and returns per-shape times. You don't need a GPU of your own. A check uses 1 of your 20 daily checks and never touches the leaderboard.

Where the numbers come from. Unless a number is marked otherwise, it comes from the official check run 7e998918 of the shipped submission. Numbers marked "harness" come from our own bench.py on a rented H200; they are measured the same way, but are not the judge's. The round-by-round project log is in the project report.

⌂ Guide home · Start here1 · The problem and the judge →