⌂ Guide home · 4 · The Python side← 3 · The strategy | 5 · Repacking weights →

4 · The Python side, line by line

This chapter reads every line of submission.py that is plain Python: the header, the configuration tables, how the CUDA code gets compiled, prepare(), and the dispatcher kernel_fn(). The CUDA kernels (lines 42–472) and the Triton kernels (475–583) get their own chapters. Each code block below is followed by notes keyed to its line numbers.

ON THIS PAGE

4.1 Header and interface
4.2 The three configuration tables
4.3 Compiling the CUDA code
4.4 prepare(): everything that may be slow
4.5 Launch helpers for Triton
4.6 kernel_fn(): the timed call
Check yourself

4.1 Header and interface

submission.py, lines 1–10

# okc: problem=muse_bf16_proj
"""Dense projections for Muse Spark Open.

m <= 32 (decode and 8-token verify): a CUDA tensor-core GEMV, on losslessly compressed weights (12 bits per BF16
weight: raw low byte + sign/exponent-window code, exact fp32 corrections for out-of-window weights) where that is
faster, else on plain BF16. Weights are repacked in prepare() into mma
A-fragment order so every lane streams its own 16B per fragment; one CTA owns a 16-row block, its warps split
K and reduce through smem (optionally split across CTAs with a semaphore-fused reduction). The router runs on
a BF16 hi+lo split of its FP32 weights with FP32 accumulation, fused into one launch.
m > 32: cuBLAS for the dense layers; the router uses a fused hi+lo Triton GEMM (TMA, semaphore split-K)."""

1A tag the okc command-line tool reads to know which problem this file is for. It must stay the first line.
2–10The module docstring: the whole design in eight lines. After chapter 3 it should read like a summary.

submission.py, lines 12–18

from typing import TypedDict

import torch
import triton
import triton.language as tl
from torch.utils.cpp_extension import load_inline
from triton.tools.tensor_descriptor import TensorDescriptor

12TypedDict is used only to document the shapes of the dictionaries the judge passes in.
14–16PyTorch, plus Triton and its language module tl, for the kernels of chapter 9.
17load_inline compiles C++/CUDA source code held in a Python string into an importable module, at run time (section 4.4).
18TensorDescriptor describes a 2-D tensor plus a block shape for Hopper's TMA copy engine. Triton kernels take it as an argument (chapter 9).

submission.py, lines 21–38

class Weights(TypedDict):
    qkv_gate: torch.Tensor  # bfloat16 [2304, 8192]
    o: torch.Tensor         # bfloat16 [8192, 1024]
    pooler: torch.Tensor    # bfloat16 [4096, 8192]
    unpooler: torch.Tensor  # bfloat16 [8192, 4096]
    router: torch.Tensor    # float32 [256, 8192]


class Config(TypedDict):
    m: int
    n: int      # output width; the unpooler's is 1024 up to 8 tokens (this GPU's slice), else 8192
    k: int
    layer: str  # which weight: "qkv_gate", "o", "pooler", "unpooler" or "router"


class Inputs(TypedDict):
    x: torch.Tensor  # bfloat16 [m, k]
    config: Config

21–26The weights dict: four BF16 matrices plus the FP32 router, each stored as [N, K]. Row n of W holds the weights of output n.
29–33Each call's config. n is the output width: 1024 for the unpooler at ≤ 8 tokens, otherwise the layer's full N.
36–38A call's data: the activations x [m, k] in BF16, plus the config. These classes are documentation; nothing checks them at run time.

submission.py, line 41 · the CUDA source starts here

CUDA_SRC = r'''

41CUDA_SRC is one big raw string (r'''…''', so backslashes stay literal) holding all the C++/CUDA code of lines 42–472. Python never runs it; it hands it to the compiler in _cuda(). The string closes on line 473:

submission.py, line 473

'''

4.2 The three configuration tables

submission.py, lines 586–606 · CTABLE: plain BF16 GEMV

# (layer, m) -> (NW, S, U, P): warps per CTA, K-splits across CTAs, loads in flight per warp, software-pipeline mode
CTABLE: dict = {('o', 1): (4, 1, 4, 2),
 ('o', 8): (4, 1, 2, 3),
 ('o', 16): (8, 1, 2, 0),
 ('o', 32): None,
 ('pooler', 1): (8, 1, 2, 3),
 ('pooler', 8): (8, 1, 2, 3),
 ('pooler', 16): (16, 1, 8, 0),
 ('pooler', 32): (16, 1, 8, 0),
 ('qkv_gate', 1): (8, 1, 4, 2),
 ('qkv_gate', 8): (16, 1, 2, 2),
 ('qkv_gate', 16): (16, 1, 8, 0),
 ('qkv_gate', 32): None,
 ('router', 1): (16, 4, 4, 3),
 ('router', 8): (4, 16, 2, 2),
 ('router', 16): (8, 16, 2, 3),
 ('router', 32): (16, 8, 2, 3),
 ('unpooler', 1): (16, 1, 2, 3),
 ('unpooler', 8): (16, 1, 8, 3),
 ('unpooler', 16): None,
 ('unpooler', 32): None}

586Each entry maps (layer, m) to (NW, S, U, P):
- NW: warps per CTA;
- S: how many CTAs share the K dimension (split-K);
- U: how many 32-wide K blocks a warp loads per loop trip;
- P: one of four loop styles (chapter 6).
587–60620 entries. None (lines 590, 598, 605, 606) means "don't use the GEMV here", which sends the call to cuBLAS. Many entries are never reached, because ZTABLE is checked first and claims the same key, for example ('o', 1). Today CTABLE actually serves the four router shapes and the unpooler slice at m = 1 and 8.

submission.py, lines 607–619 · ZTABLE: compressed GEMV

# (layer, m) -> (NW, S, UG, P) for the compressed-weight GEMV (UG = 64-wide K groups in flight per warp)
ZTABLE: dict = {('o', 1): (4, 1, 2, 3),
 ('o', 8): (4, 1, 2, 3),
 ('o', 16): (8, 1, 2, 0),
 ('pooler', 1): (16, 1, 1, 2),
 ('pooler', 8): (16, 1, 1, 2),
 ('pooler', 16): (8, 1, 2, 3),
 ('pooler', 32): (8, 1, 1, 3),
 ('qkv_gate', 1): (16, 1, 1, 2),
 ('qkv_gate', 8): (16, 1, 1, 2),
 ('qkv_gate', 16): (4, 4, 1, 1),
 ('unpooler', 16): (4, 1, 2, 3),
 ('unpooler', 32): (4, 1, 1, 3)}

607The same four fields for the compressed kernel. NW counts the warps that stream weights; the kernel launches one extra warp for exceptions. UG is how many 64-wide K groups a warp loads per loop trip.
608–61912 shapes run compressed: qkv_gate, o and pooler at m = 1, 8 and 16; pooler and unpooler at m = 32; unpooler at m = 16.

submission.py, lines 620–623 · TABLE: Triton split-K GEMM

# (layer, m) -> (BN, BM, BK, SPLIT, num_warps, num_stages, tma)
TABLE: dict = {('pooler', 128): (64, 128, 128, 2, 4, 4, 1, 0),
 ('router', 128): (64, 128, 128, 16, 4, 3, 1, 1),
 ('router', 256): (64, 64, 128, 8, 4, 3, 1, 1)}

620The fields:
- BN: weight rows per CTA;
- BM: tokens per CTA;
- BK: K per loop step;
- SPLIT: K-splits;
- Triton's num_warps and num_stages (pipeline depth);
- tma: use TMA descriptors;
- an eighth field the comment doesn't mention: atom = 1 selects the atomic split-K epilogue (chapter 9).
621–623Three shapes: pooler at m = 128, and router at m = 128 and 256.

submission.py, lines 624–626

624
625
626

WS_FLOATS = 8 * 1024 * 1024
_MOD = None
_LAYERS = ("qkv_gate", "o", "pooler", "unpooler", "router")

6248 Mi floats = 32 MB of scratch space for split-K partial sums.
625Cache for the compiled CUDA module, so it is built once per process.
626The five layer names, used by the warm-up loop in prepare().

4.3 Compiling the CUDA code

submission.py, lines 629–637 · _cuda()

def _cuda():
    global _MOD
    if _MOD is None:
        _MOD = load_inline(
            "okc_gemv_proj",
            cpp_sources="#include <torch/extension.h>\nvoid gemv(torch::Tensor x, torch::Tensor w, torch::Tensor w2, torch::Tensor out, torch::Tensor ws, torch::Tensor sem, int64_t N, int64_t K, int64_t NW, int64_t S, int64_t U, int64_t NB, bool dual, int64_t P);\n"
            "void gemvz(torch::Tensor x, torch::Tensor w, torch::Tensor esc, int64_t capr, torch::Tensor out, torch::Tensor ws, torch::Tensor sem, int64_t N, int64_t K, int64_t NW, int64_t S, int64_t UG, int64_t NB, int64_t P, int64_t TL, int64_t TH, int64_t dbg);",
            cuda_sources=CUDA_SRC, functions=["gemv", "gemvz"], extra_cuda_cflags=["-O3", "-arch=sm_90a"], verbose=False)
    return _MOD

629–631Compile on the first call only, and keep the module in _MOD.
632–633load_inline("okc_gemv_proj", …). PyTorch writes the sources to ~/.cache/torch_extensions/…/okc_gemv_proj, runs nvcc and the C++ compiler, and imports the result. If the sources haven't changed, it reuses the earlier build.
634–635cpp_sources: C++ declarations of the two host functions gemv and gemvz, whose definitions are in CUDA_SRC. From these, PyTorch generates the pybind11 glue that lets Python call _cuda().gemv(x, w, …) with tensors and ints.
636
- functions: the names to export.
- -O3: full optimization.
- -arch=sm_90a: compile for Hopper (H100/H200), including Hopper-only features (the "a"). The file therefore only runs on Hopper GPUs.
- Compiling all the template variants takes about 50 s, which is why prepare() calls this.

submission.py, lines 640–647

def _pack(w: torch.Tensor) -> torch.Tensor:
    """[N, K] bf16 -> [RB, KB, step, lane(g,t), (j,h,e)]: each lane's 16B is one m16n8k16 A fragment."""
    n, k = w.shape
    return w.reshape(n // 16, 2, 8, k // 32, 4, 2, 2, 2).permute(0, 3, 5, 2, 4, 6, 1, 7).contiguous()


def nb_of(m):
    return 1 if m <= 8 else 2 if m <= 16 else 4

640–643_pack reorders a weight matrix into the exact order the tensor core consumes. It is short but subtle, so it gets all of chapter 5.
646–647nb_of(m): how many 8-token column blocks (NB) the mma loop must cover. m ≤ 8 gives 1, m ≤ 16 gives 2, and up to 32 gives 4.

Lines 650–708 are pack_z, the 12-bit compressor, explained in chapter 7.

4.4 prepare(): everything that may be slow

submission.py, lines 711–722

def prepare(weights: Weights):
    state = {name: w.contiguous() for name, w in weights.items()}
    r = weights["router"]
    hi = r.to(torch.bfloat16)
    lo = (r - hi.float()).to(torch.bfloat16)
    state["router_hi"] = hi.contiguous()
    state["router_lo"] = lo.contiguous()
    for name in ("qkv_gate", "o", "pooler", "unpooler"):
        state["p_" + name] = _pack(state[name])
    state["p_unpooler_s"] = _pack(state["unpooler"][:1024].contiguous())
    state["p_router_hi"] = _pack(state["router_hi"])
    state["p_router_lo"] = _pack(state["router_lo"])

711Called once with the weights dict (GPU tensors). It may take up to 300 s and is not timed.
712state starts as contiguous copies of all five weights. The cuBLAS fallbacks use these.
713–717The router split.
- hi is the FP32 weight rounded to BF16.
- lo is the part hi missed, rounded to BF16.
- hi + lo reproduces the FP32 weight to about 2⁻¹⁶ relative error (lab 4).
718–719Each dense layer, repacked into fragment order for gemv_kernel, is stored as p_qkv_gate, p_o and so on.
720The unpooler's first 1024 rows, packed separately (p_unpooler_s). At ≤ 8 tokens only that slice is computed.
721–722The router's hi and lo, packed for the dual (hi + lo) GEMV.

submission.py, lines 723–726

    for (layer, m) in ZTABLE:
        key = "z_" + layer + ("_s" if layer == "unpooler" and m <= 8 else "")
        if key not in state:
            state[key] = pack_z(state[layer][:1024].contiguous() if key == "z_unpooler_s" else state[layer])

723–726Compress each matrix that some ZTABLE shape needs, once. The key is z_<layer>, or z_unpooler_s for the slice; no ZTABLE entry uses the slice today. pack_z returns a dict holding the packed bytes, the exception table and the decoder's lookup words.

submission.py, lines 727–731

    dev = r.device
    state["ws"] = torch.empty(WS_FLOATS, dtype=torch.float32, device=dev)
    state["sem"] = torch.zeros(4096, dtype=torch.int32, device=dev)
    state["racc"] = torch.zeros(256 * 8192, dtype=torch.float32, device=dev)  # zeroed split-K accumulator
    state["desc"] = {}

728ws: the 32 MB scratch for split-K partial sums.
729sem: 4,096 integer counters, zeroed once. Every kernel that uses a counter resets it to 0 when it finishes, so the next call (or graph replay) needs no clearing step. Lab 5 shows what happens without that.
730racc: an FP32 accumulator for the atomic split-K epilogue. It is zeroed once, and the kernel re-zeroes the part it used.
731desc: a cache of TMA descriptors for the weight matrices.

submission.py, lines 732–741

    _cuda()
    # compile everything the tables (and defaults) can hit, outside the timed calls
    ms = {m for (_, m) in CTABLE} | {m for (_, m) in TABLE} | {1, 8, 16, 32, 128, 256}
    for layer in _LAYERS:
        k = state[layer].shape[1]
        for m in sorted(ms):
            n = 1024 if layer == "unpooler" and m <= 8 else state[layer].shape[0]
            kernel_fn({"x": torch.zeros(m, k, dtype=torch.bfloat16, device=dev), "config": {"m": m, "n": n, "k": k, "layer": layer}}, state)
    torch.cuda.synchronize()
    return state

732Compile the CUDA module now, while nothing is being timed.
733–739Warm-up. Call kernel_fn once for every layer × m in {1, 8, 16, 32, 128, 256}, with zero inputs.
- Triton compiles a kernel the first time it sees a new set of compile-time parameters, which takes seconds.
- cuBLAS picks algorithms and allocates workspace on first use.
- Doing all of that here keeps it out of the judge's graph capture and timing.
- Line 738 computes n exactly as the judge does.
740Waiting for the GPU is fine here; it is forbidden only inside kernel_fn.
741The judge passes this dict back on every kernel_fn call.

4.5 Launch helpers for Triton

submission.py, lines 744–749 · _desc

def _desc(state, name, t, block):
    key = (name, tuple(block))
    d = state["desc"].get(key)
    if d is None:
        d = state["desc"][key] = TensorDescriptor.from_tensor(t, block_shape=list(block))
    return d

744–749Build a TensorDescriptor for a weight matrix and block shape once, and cache it under (name, block). Weights never change, so one descriptor serves every call.

submission.py, lines 752–771 · _launch

def _launch(x, name, w, w2, out, state, N, K, cfg):
    bn, bm, bk, split, nw, ns, tma = cfg[:7]
    atom = len(cfg) > 7 and cfg[7] and split > 1
    M = x.shape[0]
    grid = (N // bn, split, triton.cdiv(M, bm))
    if tma:
        xd = TensorDescriptor.from_tensor(x, block_shape=[bm, bk])
        wd = _desc(state, name, w, [bn, bk])
        w2d = _desc(state, name + "2", w2, [bn, bk]) if w2 is not None else wd
        _proj_tma_kernel[grid](
            xd, wd, w2d, out, state["racc"] if atom else state["ws"], state["sem"], M, N, K,
            BN=bn, BM=bm, BK=bk, SPLIT=split, DUAL=w2 is not None, OUT_F32=out.dtype == torch.float32, NS=ns,
            ATOM=bool(atom), num_warps=nw, num_stages=ns,
        )
        return
    _proj_kernel[grid](
        x, w, w2, out, state["ws"], state["sem"], M, N, K,
        BN=bn, BM=bm, BK=bk, SPLIT=split, DUAL=w2 is not None, OUT_F32=out.dtype == torch.float32,
        num_warps=nw, num_stages=ns,
    )

753–754Unpack a TABLE entry. atom only means something when K is actually split.
755–756The launch grid: (weight-row tiles, K slices, token tiles), with one Triton "program" (one CTA) per cell. For the router at m = 128 that is (256/64, 16, 1) = 64 CTAs.
757–760The TMA path.
- x is a new tensor on every call, so its descriptor is built every call. That happens on the CPU while the graph is captured.
- Weight descriptors come from the cache.
- Without a second matrix, the w2 descriptor is a dummy that is never used.
761–766Launch _proj_tma_kernel. Every NAME=value argument is a compile-time constant, so each distinct combination is a separately compiled kernel. The scratch buffer is racc for the atomic epilogue, and ws otherwise.
767–771The plain-pointer version of the kernel. Every TABLE entry sets tma=1, so this path is currently unused.

4.6 kernel_fn(): the timed call

submission.py, lines 774–786 · 1. Triton

def kernel_fn(data: Inputs, state) -> torch.Tensor:
    x = data["x"]
    cfg = data["config"]
    layer, n, k = cfg["layer"], cfg["n"], cfg["k"]
    m = x.shape[0]
    router = layer == "router"
    if (layer, m) in TABLE:
        out = torch.empty((m, n), dtype=torch.float32 if router else torch.bfloat16, device=x.device)
        if router:
            _launch(x, "router_hi", state["router_hi"], state["router_lo"], out, state, n, k, TABLE[(layer, m)])
        else:
            _launch(x, layer, state[layer], None, out, state, n, k, TABLE[(layer, m)])
        return out

775–779Unpack the call. m comes from x's shape, and branching on shapes is allowed.
780–786A shape in TABLE goes to Triton.
- The output is allocated as FP32 for the router and BF16 otherwise.
- Allocating inside a captured graph is fine: PyTorch gives the graph its own fixed memory pool.
- The router passes hi and lo (dual); other layers pass their one weight.

submission.py, lines 787–793 · 2. compressed GEMV

    zc = ZTABLE.get((layer, m))
    if zc is not None:
        z = state["z_" + layer + ("_s" if layer == "unpooler" and m <= 8 else "")]
        out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device)
        _cuda().gemvz(x, z["w"], z["esc"], z["capr"], out, state["ws"], state["sem"], n, k, zc[0], zc[1], zc[2],
                      nb_of(m), zc[3], z["TL"], z["TH"], 0)
        return out

787–790A shape in ZTABLE goes to the compressed kernel. The router is never compressed, so the output is always BF16.
791–792Call the C++ host function gemvz with:
- x, the packed weights, the exception table and capr (its row length);
- the output, scratch and counters, and N and K;
- the configuration (NW, S, UG), NB and PIPE;
- the decoder's lookup words TL and TH;
- dbg=0.

submission.py, lines 794–803 · 3. plain GEMV

    if m <= 32 and CTABLE.get((layer, m), 0) is not None:
        nb = 1 if m <= 8 else 2 if m <= 16 else 4
        nw, s, u, p = CTABLE.get((layer, m)) or ((8, 8, 2, 0) if router else (4, 1, 2, 0))
        out = torch.empty((m, n), dtype=torch.float32 if router else torch.bfloat16, device=x.device)
        if router:
            w, w2 = state["p_router_hi"], state["p_router_lo"]
        else:
            w, w2 = state["p_unpooler_s" if layer == "unpooler" and n < 8192 else "p_" + layer], state["ws"]
        _cuda().gemv(x, w, w2, out, state["ws"], state["sem"], n, k, nw, s, u, nb, router, p)
        return out

794For m ≤ 32 and no explicit None. .get(key, 0) returns 0 for a missing key, and 0 is not None, so missing keys come here too.
796The table's configuration, or defaults for shapes without an entry, such as the test shapes m = 5, 7 and 9:
- router: 8 warps, K split 8 ways;
- other layers: 4 warps, no split.
798–801The router passes its packed hi and lo, with dual=True. Other layers pass their packed weights (the slice for the unpooler when n is 1024) plus a dummy second tensor that the kernel ignores.
802Launch through the C++ host function gemv (chapter 6).

submission.py, lines 804–807 · 4. cuBLAS fallbacks

    if router:
        acc = torch.mm(x, state["router_hi"].t(), out_dtype=torch.float32)
        return acc + torch.mm(x, state["router_lo"].t(), out_dtype=torch.float32)
    return x @ state[layer][:n].t()

804–806A router shape that none of the above handles, for example m = 64 in our test sweep: two cuBLAS GEMMs with FP32 output, then an add. This is SGLang's own method.
807Everything else: x @ Wᵀ through PyTorch, which calls cuBLAS. The [:n] trims the unpooler to its 1024-row slice when needed.

Try it · lab0_dispatch.py — Lab 0: which engine runs each shape? The tables are read straight out of submission.py.

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

The file is already on your Mac at ~/projects/okc/muse_bf16_proj/learn/labs/lab0_dispatch.py; you can also view or download it here.

Exercise. Copy submission.py to /tmp/mine.py and delete the line ('pooler', 1): (16, 1, 1, 2), from ZTABLE. Run python3 lab0_dispatch.py /tmp/mine.py. Which engine and configuration does pooler at m = 1 get now, and from which table? Chapter 10's experiment A runs exactly this change on the real H200.

Check yourself

If a (layer, m) is in both ZTABLE and CTABLE, which one is used?

ZTABLE, because kernel_fn checks it first (line 787, before line 794). For example, ('o', 1) appears in both and runs compressed.

What would go wrong if prepare() skipped the warm-up loop?

The first call with each new Triton configuration would compile it, which takes seconds, and cuBLAS would initialize on first use. That work would land inside the judge's capture or timing instead of in the untimed prepare().

Why does every split-K kernel reset its counter to 0 itself, instead of Python calling sem.zero_() before each launch?

zero_() is another kernel launch: about 1–2 µs more per call, inside the timed region. Having the last CTA reset the counter costs nothing and leaves the state ready for the next graph replay.

⌂ Guide home · 4 · The Python side← 3 · The strategy | 5 · Repacking weights →