Muse Spark Open Kernel Challenge — muse_bf16_proj (dense projections)

Result: 1.33–1.34× vs SGLang (was 1.25×). Round 5 added a lossless 12-bit compressed-weight GEMV for the decode shapes. Scored run 1.331× (95% CI 1.311–1.351×, 3,628.9 µs/step, rank #167, job a9dfcf1c069f4f79); its check run measured 1.343× (CI 1.323–1.363×, 3,597.5 µs/step, 9/9 tests, job 7e998918e84d4ea8). Before round 5 the board showed 1.250× (3,863.2 µs/step, rank ~315). Leader on this problem: 1.637× (2,951.5 µs/step); the top 8 are between 1.55× and 1.64×.

Round 6 built the compressed-weight wgmma GEMM for m=128/256 that round 5 named as the remaining upside. It is correct (bit-exact) but slower than cuBLAS on every shape, so nothing changed in the submission (see Round 6).

Hardware: RunPod H200 SXM (secure cloud). Total spend ≈ $40 over six sessions (sixth session: ~61 min, ≈ $5.3). Pod deleted; nothing running.

What was built

submission.py (generated by build.py from sub_template.py + tuned tables) dispatches on tokens m:

Per-load result (judge numbers, check job 0fdd01df8a014f82; load time = sum of 5 shapes)

Method / tooling

• RunPod pod from lmsysorg/sglang:dev-muse-spark-open (torch 2.14+cu130, Triton 3.8, CUDA 13).
• Local harness bench.py reproduces the judge (CUDA-graph replay, cold L2 via a read of 256 MB, paired alternation): it predicted 1.221× — identical to the official check. (First version flushed L2 with a write and mis-measured; fixed.)
• Measured the practical HBM read ceiling with a hand-written read kernel (cold L2): 16.8 MB → 5.3 µs, 37.7 MB → 9.9 µs, 67 MB → 16.1 µs. pooler is already within ~8% of that; qkv_gate within ~20%.
• Correctness sweep: 420 cases (2 weight seeds × 5 layers × 21 values of m incl. odd/large × normal/outlier inputs) all pass.
• Cold compile (nvcc + Triton) ≈ 53 s in prepare(); official run 1:41 total.

What did not pay off

• Triton-only GEMV: plateaus ≈ 3 TB/s (≈ cuBLAS). TMA helped (12.6 vs 13.9 µs) but not enough.
• Cross-CTA split-K as the default: the reduction tail cost more than it saved; in-CTA K-split won.
• Lossless exponent compression of BF16 weights (allowed by the rules): implemented and correct (low byte raw + 4-bit sign/exponent-window code, escape list for out-of-window values, 0.75× bytes). Idealised test with 25% fewer bytes: qkv 12.3→9.4 µs, pooler 17.3→13.3 µs. Real kernel after decode cost and exception tail: qkv 11.8, pooler 15.7 (m=1), unpooler(m=16) 15.4 vs 17.3 — only ≈ +1–3% overall, so not integrated in round 1. Round 5 redesigned it and it now ships (see Round 5).
• CUDA mma.sync for m=128/256: 4× slower than cuBLAS (needs wgmma); dense verify 16/32 therefore stay at cuBLAS parity.

Round 2 — attempts to close the gap (all measured on the H200; none beat the shipped kernels)

Measured floors first: a hand-written read kernel reaches 16.8 MB in 5.4 µs, 37.7 MB in 10.1 µs, 67 MB in 16.4 µs (cold L2). A TMA bulk-copy (cp.async.bulk) streaming kernel gets the same numbers, so the HBM ceiling is real, not a register-load artefact. Per-layer at m=1 (ours / read floor): qkv_gate 12.2/10.1, o 6.0/5.4, pooler 17.8/16.4, unpooler-slice 4.8/~3.2, router 6.4/~3.3 µs — the plain GEMV kernels are within ~15–20 % of the floor.

Corrected estimate for the wgmma idea: the earlier "1.4–1.6× on m=128/256" was too optimistic. Compute floor at ~700 TF and the cold-L2 weight read put the dense layers at ~54 µs (m=128, now 68.9) and ~69 µs (m=256, now 81.4); with real kernel overheads (≈2–3 µs of start-up and tail per launch measured above) the realistic gain is ≤ ~1.25× on those two loads, i.e. ≤ ~+5 % overall, and the CUTLASS result shows a generic wgmma kernel does not get there.

Not attempted on purpose: pinning weights in L2 (persisting access policy / evict_last) so they survive the judge's cold-L2 flush — that would game the measurement.

Round 3 — profiling + retune (+3.2 % on the scored time: 3,987.8 → 3,863.2 µs/step)

Profiler access. ncu is installed on the pod but every counter query fails with ERR_NVGPUCTRPERM (the host driver restricts GPU performance counters to admins; unfixable from inside the container). So no stall/DRAM-throughput counters. Used instead: SM-clock sampling (nvidia-smi -lms), and in-kernel %globaltimer (32 ns resolution) + clock64 + %smid stamps per CTA, bracketed by two stamp kernels inside the CUDA graph.

What the timeline showed (qkv_gate m=1, 37.7 MB): - The SM clock is 1980 MHz under load (not ~1.1 GHz as I had assumed); measured 1.8–2.0 cycles/ns in-kernel. Clock is not a bottleneck. - Kernel-to-kernel gap inside a graph: ≈0.65 µs before the first CTA starts and ≈0.9 µs from the last CTA's end to the next kernel — fixed, outside our control. - All 144 CTAs start within 0.1 µs; first batch of loads completes ≈2.8 µs after start (bandwidth-bound, 9 MB in flight), CTAs end over a ≈1.6 µs spread. - CTA window (first start → last end) 11.8 µs vs 9.4–9.7 µs for a pure read kernel moving the same bytes. Same kernel with the x-load + mma removed: 10.0 µs. So ≈2 µs of the gap was the per-iteration load → wait → mma round trip (no overlap between batches), not DRAM.

Fix. Added software-pipelined variants to the GEMV (next batch of weight loads issued before the current batch's mma) and a minBlocks register-allocation switch, giving four modes P ∈ {0, 1, 2, 3} per (layer, m), and re-ran the exhaustive per-shape tune (NW × K-split × U × P; top-10 re-measured at 400 reps). The old tune had missed that 8-warp CTAs beat 16-warp ones for most m=1/8 shapes. Same-pod A/B (bench.py, equal-weight geomean): old 1.220× → new 1.264×. Pitfall found on the way: rewriting the P=0 loop changed register allocation and made it 25 % slower for NW=16/NB=2 until minBlocks was left unconstrained for that mode.

Remaining gap at m=1 (µs, ours vs cold-L2 read floor incl. same fixed costs): qkv_gate 11.0 vs ≈10.1, pooler 17.1 vs ≈16.4, o 5.9 vs ≈5.4, unpooler 3.9 vs ≈3.2, router 5.4 vs ≈3.3. Closing all of it would be worth about +3–4 % on the decode loads; beyond that only fewer bytes (lossless compression, which has not paid off here) or faster m=128/256 GEMMs (≤ +5 %) remain.

Round 4 — strategy review, a new wgmma GEMM, and a Triton epilogue fix (official check 1.239× → 1.262×)

Strategy review. The per-load numbers of the best run (judge µs):

If every shape reached its practical floor, the score would be ≈ 1.57×. That is roughly where the leaders are, so they must be near the floor on every load. About half of the remaining headroom is in verify 16 and verify 32 (m=128 and m=256), where we ran cuBLAS for the dense layers.

My earlier report said cuBLAS was close to optimal there. The kernel names in the official profile show it is not: - pooler m=128: cuBLAS runs nvjet 64x64 tiles. Every weight byte goes to 2 CTAs and every x byte to 64, which works out to 2.6 TB/s effective. - qkv m=128: a separate splitKreduce kernel costs 3 µs on its own. - pooler m=256: about 540 TFLOP/s. - router m=128 and m=256: this was our own Triton kernel. It spent about 3–4 µs in a serial split-K reduction, where the last CTA of a tile reads all 16 partials.

New kernel "wgk" (wgk_cuda.py, about 450 lines of CUDA and PTX): - Design: Hopper TMA plus wgmma, warp-specialized. One producer thread issues the TMA loads; 1–3 consumer warpgroups run wgmma.m64n{m}k16. Each CTA holds all m tokens, a tile of 64–256 weight rows, and one K range. The split-K reduction goes either through distributed shared memory inside a thread-block cluster, or through global memory across a co-resident grid. TMA multicast of x is optional. - Correctness: passes on every scored shape, including probe rows and outlier inputs. - Speed: it did not beat cuBLAS on the dense layers (µs, same-pod harness):

What the measurements showed. I used per-CTA %globaltimer stamps plus ablation modes that drop the x loads, the weight loads, or the MMA. - Ingest-bound main loop. With all tokens in one CTA and 64 weight rows, each SM takes in 3× its weight bytes, because x is re-read from L2 by every row tile. Per-SM ingest tops out around 85 GB/s. So at m=128 the main loop is limited by L2→SM traffic rather than HBM: pooler takes 18.7 µs for 195 MB through the SMs. - Bigger row tiles don't fix it. Larger tiles (128 or 256 rows) cut the x traffic, but they need more K-splits. The split-K exchange plus waiting for the slowest CTA of each tile cost 2–5 µs, more than the traffic saved. Per-CTA main-loop times spread by up to 3.6 µs. - Cluster residency limit. On this H200, clusters of 4 can only place 120 CTAs, not 132. A 128-CTA grid in clusters of 4 or more therefore runs as two waves. - Remote arrive cost. A remote mbarrier.arrive.release.cluster costs about 1 µs per arrive. Using the plain .shared::cluster arrive fixed this. - Multicast didn't help. TMA multicast of x across a 2-CTA cluster gave nothing, which is consistent with per-SM ingest, not L2 output, being the limit.

What paid off: fixing the Triton split-K epilogue. Each CTA now does a relaxed atomic add of its fp32 partial into a zeroed buffer. The last CTA of the tile copies the result out and re-zeroes the buffer, so there is no serial S-way read. After a retune (µs, judge):

Official check: 1.262× (verify 16 went from 1.05× to about 1.10×, verify 32 from 1.03× to 1.07×; an earlier version of this report said 1.14× for verify 16, which was wrong). The scored run measured 1.248×. Its decode shapes came out 2–5 % below the previous run, which is noise. The organizers' final reruns use medians, where the gain should show.

Also tried this round, all without a win: - the Triton atomic epilogue for the router at m ≤ 32 (the GEMV is still better: 5.5–5.8 µs); - Triton for the other dense m=128/256 shapes (all slower than cuBLAS); - F.linear vs x @ W.T vs SGLang's dispatch (identical).

A harness fix: bench.py now keeps every captured graph alive. Freeing a graph's memory pool left cuBLAS with a dangling workspace, which caused an illegal-address error when different layers were mixed.

Round 5 — lossless 12-bit compressed weights for the decode shapes (1.262× → 1.343× check, 1.331× scored)

The decode shapes (m ≤ 32) were already within 4–17 % of the cold-HBM read floor, so the only way to go much faster was to read fewer bytes. The rules allow lossless compression, decoded inside the kernel.

Format (gemvz_cuda.py, pack_z). Each BF16 weight becomes 12 bits: - its low byte unchanged (exponent LSB plus 7 mantissa bits); - a 4-bit field: the sign, plus a 3-bit code selecting one of 8 consecutive values of the 7 high exponent bits. The 8 values are chosen per weight matrix in prepare() to cover the most weights. - Weights whose high exponent bits fall outside that window (0.014 % of pooler's weights; very small or zero values) are stored clamped into the window. The exact difference, an exact fp32 value, goes to a per-row exception list. - Decoding is pure byte shuffling: 12 integer instructions (prmt/lop3) per 8 weights. The bit layout puts each sign bit where a single lop3 can merge it, which saved 3 instructions per fragment over the round-1 layout.

Kernel (gemvz). Same structure as the shipped GEMV (one CTA per 16-row block, warps split K, optional split-K across CTAs, the same software-pipelining modes), plus: - one extra warp per CTA applies the exception corrections (out += delta * x[k]) while the other warps stream, so the main loop has no branches and no tail; - every (row, token) correction is summed by a fixed lane in a fixed order, so results are deterministic.

Correctness. Bit-for-bit equal to the original weights on one-hot inputs placed on exception columns (every layer, m = 1/5/8/13/16/31/32, every kernel variant). It also passes the 9 judge tests and the 420-case sweep, which uses two other weight seeds.

Speed (µs, same-pod harness, best config per shape):

Per load (judge µs, check 7e998918):

Tried this round without a win: - 2 or 4 row blocks per CTA, sharing each x fragment, for m=16/32: slower than one row block. - An x-stationary variant (x tile in shared memory, one row block per warp, cp.async weight ring): 35–60 % slower even at m=1. Not resolved. - Interleaving the warps' K ranges for DRAM locality: no gain. One finding: a runtime stride in the inner-loop address arithmetic made the kernel up to 60 % slower. Code generation here is very sensitive. - Compressing the router's hi and lo (merged exception lists): correct, but slower (6.2 vs 5.4 µs at m=1, 10.1 vs 6.0 µs at m=32). The router is too small and latency-bound for the byte saving to pay for the decoding. - m=32: the compressed GEMV wins only on unpooler. qkv_gate and o stay on cuBLAS, and pooler is about even.

Round 6 — compressed-weight wgmma GEMM for m=128/256 (zwg): correct, but slower than cuBLAS everywhere

Goal. verify 16 (m=128) and verify 32 (m=256) were still at 1.11× and 1.06×, with the dense layers on cuBLAS or Triton. At m=128 a BF16 GEMM is memory-bound (pooler: 67 MB of weights, ~9 µs of tensor-core math at peak), so reading the round-5 12-bit weights (−25 % bytes) should have helped there.

Kernel (zwg_cuda.py, tests and tuner zwt.py): - Weights use the round-5 format, repacked into the standard mma k-order and stored tile-major, so one bulk copy per stage fetches a tile's 64-wide K group. - One CTA per SM, warp-specialized. A producer warp issues the TMA x tile (128B swizzle) and the bulk weight copy into a 4–12 stage mbarrier ring. - Two consumer warpgroups decode their 16-row fragments from shared memory into registers (12 prmt/lop3 per 8 weights) and feed them to wgmma.m64nNk16 as the register A operand. x is the shared-memory B operand. Fragment j+1 (or j+2) is decoded while wgmma j runs. - Split-K across S CTAs per tile: plain fp32 partial stores, a per-tile counter, then each CTA reduces 1/S of the tile's rows.

Correctness: bit-exact against the original weights on one-hot inputs on exception columns, and it passes the reference check for all four dense layers at m = 16/32/100/128/200/256.

Exception corrections (0.010–0.014 % of weights, but each needs a whole x column at m=128/256). Three designs:

The L2 gather was slow because of latency, not volume: about 4,600 scattered 4-byte reads per SM against a limited number of outstanding misses. The +4 µs left in the final design is mostly CTA imbalance: the split-K wait grew from 2.6 to 5.9 µs.

Speed (same-pod harness, µs, best config from a sweep; "no exc" skips the corrections, so its output is wrong. It is the floor if corrections were free.)

Why it loses. These are per-CTA timelines and component timings, measured by switching parts of the kernel off, for pooler m=128. - The main loop is the bottleneck, not the weight bytes. It takes 18–19 µs: about 570 ns per 64-wide K group per CTA, against about 280 ns of tensor-core time at peak. - The weight and x copies alone take 16.7 µs. Decode plus wgmma alone, without copies, takes 19.6 µs; wgmma alone 17.4 µs. - So the register-A wgmma pipeline runs at roughly 50 % of peak, and the 25 % byte saving cannot show. - Two wgmmas in flight per warpgroup instead of one changed nothing at m=128. At m=256 it sped up the wgmma-only time (25.3 → 22.3 µs) but not the total. - High fixed cost. An empty pipeline with no copies, decode or math still takes about 10 µs: about 190 ns per stage handshake plus about 4 µs fixed. The split-K wait and reduction add another 4.4 µs. - m=256 is compute-bound. That is 17 GFLOP, about 17 µs at peak, and cuBLAS takes 31 µs. Fewer weight bytes cannot help there.

What it would take, and why I stopped. To win, the main loop would need to be about 1.6× more efficient and the fixed cost halved. Candidate changes: - 128-wide K stages, to halve the handshakes; - a decode warpgroup that writes BF16 tiles to shared memory, so the math uses shared-memory wgmma at CUTLASS efficiency; - a cluster/DSMEM split-K reduction.

Even if all of that worked, verify 16 would drop from about 78 to about 68 µs, which is only about +2 % on the overall score (≈ 1.36×), and verify 32 would not move. That is several more days on this problem for a small gain, so I stopped round 6 here. The submission is unchanged (submission_r5a.py, 1.331× scored).

Where the remaining headroom is

Four of the six loads are now at 1.40–1.90×. The gap to the leaders (1.64×) is now concentrated in two places: - verify 16 and verify 32 (m=128/256), at 1.11× and 1.06×. The dense layers there still run cuBLAS or Triton at parity with SGLang. At m=128 they are still memory-bound: 38 µs of compressed weight reads against about 35 µs of tensor-core math at 700 TFLOP/s, compared with 68 µs today. - decode 32, at 1.20×. The m=32 GEMV is limited by x re-reads and memory latency, not by weight bytes.

Round 6 built that compressed-weight Hopper GEMM (register-A decode into wgmma, as in vLLM's Machete). It is correct but slower than cuBLAS on every m=128/256 shape, because the register-A pipeline runs at about 50 % of tensor-core peak (see Round 6). Closing verify 16 now needs a CUTLASS-grade mainloop for an estimated +2 % overall. That is a poor return compared with the four other challenge problems, which I have not started.

Caveats

• Final standings are medians of organizer reruns on held-out data; run-to-run noise on the judge is ±1.5%.
• Router uses the rules-sanctioned hi+lo split; all other weights are used at native BF16 precision.

Artifacts

~/projects/okc/muse_bf16_proj/ — submission.py (= submission_r5a.py, the scored round-5 build), gemvz_cuda.py + zt.py (compressed GEMV and its tests/tuners), cuda_z.json, logs_r5/, zwg_cuda.py + zwt.py (round-6 compressed wgmma GEMM; earlier versions zwg_cuda_v1..v4.py), logs_r6/, sub_template.py, build.py, bench.py, tune.py, tune_cuda.py, cuda_tuned.json, tuned.json, gemv_cuda.py; helper ~/projects/okc/pod.sh; git repo ~/projects/okc. Notes in shared memory hub_engineering → fact/okc/progress.