A 297-line kernel file that decodes faster than vLLM
Cloudy ran for 24 hours on one NVIDIA H100 and replaced SimpleLLM’s 5,762 lines of copied kernels with a single Triton file that decodes 11–15% faster than tuned vLLM.
Setup
SimpleLLM is an inference engine for gpt-oss-20b and gpt-oss-120b on one H100. Before this run, 5,762 of its 6,747 lines were a kernels/ folder, mostly a copy of OpenAI’s triton_kernels. We wanted the kernels in one file written like the rest of the repo, and no slower than the copied ones. We sent this in Autoresearch mode with a 24-hour budget:
“In naklecha/simple-llm, replace the copied triton_kernels (44 files, 5,762 lines) with one kernels.py under 300 lines, in the repo’s style. Beat tuned vLLM 0.31.0 decode on gpt-oss-20b and 120b at batch 1 and 64 on one H100. Greedy outputs must match the original. Write up what worked and what didn’t.”
Result · gpt-oss-20b decode on one H100, tok/s
15% faster than vLLM, with 5,762 lines of kernels replaced by ~300 lines of Triton.
Cloudy profiled the original engine and vLLM side by side, rebuilt attention and the MoE first because they took most of each decode step, and benchmarked every change end to end before keeping it. vLLM was re-tuned for each workload: every MoE backend, every attention backend and 14 scheduler and graph options.
Results
kernels.py is 297 lines of Triton that import only torch and triton. Decode throughput against the fastest vLLM 0.31.0 configuration at each point:
| Model | Batch | Original | Tuned vLLM | SimpleLLM | vs vLLM |
|---|---|---|---|---|---|
| gpt-oss-20b | 1 | 264.8 | 330.9 | 381.3 | +15.2% |
| gpt-oss-20b | 64 | 5,346.3 | 8,880.7 | 9,843.5 | +10.8% |
| gpt-oss-120b | 1 | 183.6 | 232.9 | 267.7 | +14.9% |
| gpt-oss-120b | 64 | 3,037.8 | 4,269.5 | 4,780.4 | +12.0% |
Output tok/s at 1,000 tokens, median of 3 runs, each engine in its own process. Differences under 2% count as ties; all four are well outside it.
- Original SimpleLLM
- SimpleLLM with kernels.py
At 131k context, decode is 6–16% faster than vLLM as well, but time to first token is 1.6–2× slower. Every run is in performance.md.
Why the size matters
The engine is now 1,351 lines: llm.py, model/ and kernels.py. That is about 23k tokens, so the whole engine, every kernel included, fits in a single prompt with most of the context window left over. Before the rewrite it was 81k tokens, and 85% of the lines were the kernels/ folder, most of it a general library with code paths for Hopper, Blackwell and AMD GPUs.
| Before | After | |
|---|---|---|
llm.py | 563 | 649 |
model/ | 422 | 405 |
| GPU kernels | 5,762 in 44 files | 297 in 1 file |
| Total | 6,747 | 1,351 |
| Tokens | ~81k | ~23k |
Lines of Python in the engine. Tokens counted with the o200k_base tokenizer.
That makes it a practical base for your own model. Give the repo to Claude Code, Codex or Cloudy and ask for a different MoE layout, attention variant or quantization format: the agent can read every line it needs to change, and it starts from an engine that already beats tuned vLLM on gpt-oss decode.
Where the time went
At batch 64 the original engine spent 3.78 ms of each 12.02 ms step in attention and 1.75 ms in gaps between kernel launches. Those two rows account for 4.3 ms of its 4.8 ms deficit to vLLM. The rewrite ends up a little faster than vLLM in almost every row rather than through one large win.
| Section | Original | Tuned vLLM | SimpleLLM |
|---|---|---|---|
| RMSNorm + residual | 0.11 | 0.11 | 0.11 |
| QKV projection + RoPE | 1.05 | 0.44 | 0.39 |
| Attention | 3.78 | 1.12 | 1.06 |
| Router + top-k | 0.32 | 0.32 | 0.27 |
| MoE gate/up + SwiGLU | 2.66 | 2.72 | 2.59 |
| MoE down + combine | 1.77 | 1.87 | 1.50 |
| Logits + sampling | 0.57 | 0.53 | 0.47 |
| Launch and host gaps | 1.75 | 0.12 | 0.21 |
| Step | 12.02 | 7.22 | 6.59 |
gpt-oss-20b at batch 64, milliseconds per decode step from profiles of each engine. Batch 1 and 120b are in performance.md.
Optimizations
Every change that was kept, in the order it landed:
| Change | Measured effect |
|---|---|
| MXFP4 MoE on tensor cores. Expert weights are permuted once at load into the tensor-core fragment layout, with the bits of 4 values interleaved so dequantizing to bf16 is a few bit ops and one multiply per 2 values. One grouped kernel covers any token count. | MoE 1.33× at batch 1, 1.10× at batch 64 |
| One decode-attention kernel. Split-KV (flash-decoding) with GQA, the sliding window, the KV append and attention sinks in one kernel, replacing four ops. | +33% at batch 1, +36% at batch 64 |
| Fused RoPE. A single kernel at every batch size. | +8% at batch 64 |
| Greedy sampling inside the CUDA graph. The argmax and the next step’s input ids and positions are written inside the captured graph, so there is no host sync per step. | +20% at batch 64 |
| Fixed decode tiles. The autotuner tuned on the graph-warmup batch, where every row routes to the same experts, and picked different tiles in different processes. Decode now pins the measured best tile. | +3.4% at 120b batch 64, no run-to-run variance |
| Per-expert routing tile table. Each expert writes its own tiles instead of the kernel searching a tiles × experts grid (256 × 128 at 120b batch 64). | 24.8 → 3.1 µs per call, +9.6% at 120b batch 64 |
| Ring buffer for sliding-window layers. Half of gpt-oss’s layers only attend to the last 128 tokens but stored the full sequence. They now keep a small ring buffer. | KV per 131k sequence 9.66 → 5.00 GB on 120b |
| CUDA graphs for any set of slots. Slot indices are a graph input and every graph is captured at startup. Before, only 5–20% of serving steps replayed a graph. | Serving 1.56× on 20b, TPOT p50 27.4 → 5.8 ms at 2 req/s |
| Paged KV for full-attention layers. 256-row pages from a shared pool instead of a worst-case 9,216-row reservation per request. | 120b serves 64 requests instead of 19, serving 1.47× |
| MoE token tile follows tokens per expert. 16-token tiles for decode, 32 or 64 for prefill, so a weight tile is unpacked once instead of about 8 times. | 4k-token prefill MoE 236 → 151 ms, serving +12.6% on 120b |
End to end unless a kernel is named. Details in optimizations.md.
Dead ends
- Per-token MXFP4 GEMV with lookup tables, an fp16 bit trick and an fp32 denormal trick. The best was 1.10× over the original MoE at batch 1, not enough to keep.
- Tensor-core MXFP4 with the layout fixed up in the kernel (inline PTX with
tl.dot, andtl.dot_scaled). Slower than permuting the weights once at load. - Tuning decode attention at 131k. About 60 tile, split, warp and stage variants gained under 4%. The kernel already reads the KV cache at 2.5–2.8 TB/s, close to the H100’s memory bandwidth.
Correctness
SimpleLLM loads the official MXFP4 checkpoints and only reorders their layout in memory; nothing is re-quantized. 121 kernel tests compare every kernel against a float32 PyTorch reference and against the original triton_kernels, at 20b and 120b shapes, and deliberately broken kernels fail them.
In a teacher-forced greedy match over 20 prompts × 1,000 steps, the top token disagrees with the original engine on 0.87–0.88% of steps for 20b and 1.47–1.66% for 120b, inside limits set by the original engine’s own batch-to-batch variation. With OpenAI’s eval harness, GPQA Diamond is 56.3% on 20b and 66.8% on 120b against published scores of 56.8 and 67.1; AIME 2025 also matches. Method and the one open item: correctness.md.
Where vLLM is still faster
- Time to first token on long prompts: 9.6 s against 6.1 s at 122k tokens on 20b. SimpleLLM prefills long prompts alone, 4,096 tokens at a time.
- Serving 120b when every request arrives at once: 5.4% slower, because prefill blocks decode. vLLM runs prompt chunks inside decode steps.
- 1,000 concurrent requests: 0.54–0.58× of vLLM. SimpleLLM runs at most 256 slots, because capturing CUDA graphs for more overruns its fixed 4 GiB reserve.
All three are on the repo’s list.
Benchmarking notes
- Over HTTP, vLLM loses about half its throughput at 1,000 concurrent streams: 22.5k tok/s in-process, about 10k behind
vllm serve. The often-quoted 7.8k tok/s on an H100 is an HTTP number. - The best vLLM attention backend depends on context length. Triton is fastest at 1,000 tokens; at 131k, FlashInfer and FlashAttention are 2.1× faster than it.
- Triton autotuning picks kernels per process, so repeated runs in one process understate run-to-run spread.
- The original engine’s per-step
tqdmupdate cost 0.38 ms, about 10% of a batch-1 step.