Optimizing Models to Be Fast at Codegen

How we serve open source models to be the best in class at coding agent workloads.

Tejas Bhakta
Tejas Bhakta
June 19, 20268 min read
Optimizing Models to Be Fast at Codegen

When we built Fast Apply we realized the open inference stack was missing a key piece: a specialized speculator for the task at hand. A generic draft gets 1.93x, a specilized one got 100x faster. Building on this, we trained WarpGrep and some of the best specialized models out there. This muscle is the same core of speculative decoding.

Small Models don't follow the same scaling laws as large models. We train custom speculators for codegen to be best in class at outputting code. We serve open models, Qwen, GLM, DeepSeek, MiniMax, for one workload: the coding agent. Making them fast comes down to three things the open stack won't do for you.

  • A custom speculator. A draft trained on the model's own coding output, not the internet. Generic draft: 1.93x. Trained on the target: 3.07x.
  • Kernels optimized for latency or throughput. FP4 decode that stays 4 bits all the way to the tensor core, searched automatically on the cheap GPUs nobody else tunes. 97 to 162 tok/s on a $7K card.
  • Interconnect optimized for high cache hit rates. A one-shot all-reduce over PCIe, and a prefix cache that crosses NVLink-denied boxes over plain TCP.

Each is a place the general stack stopped and we kept going.

1. Inference that gets faster the more you use it

Speculative decoding: a small draft model guesses the next few tokens, the target checks them in one pass, you keep the run until the first miss. One number decides everything. Acceptance rate, how often the target keeps the guess.

A generic draft is a bad guesser. On Vicuna-13B an off-the-shelf 68M draft gets 1.93x; a draft trained on the target's own output gets 3.07x, same target, same setup. That gap is the section.

Speculative decodingdraft proposes, target verifies in one pass, keep the run until the first miss
Generic draft
68M, off-the-shelf
const x = fetch(✗ verify stops
1.93×
Trained on code
drafts diffs it has seen
const user = await db.get(✗ verify stops
3.07×

More accepted tokens per step means fewer target passes. A draft trained on the model's own coding output keeps a longer run than a generic one on the same target.

The architectures are public and good. EAGLE-3 lets the draft train on raw data instead of copying the target's features, and acceptance length climbs from 3.96 to 6.62. DFlash, SGLang's Spec V2 since June 2026, drafts a whole block in one pass: over 6x lossless, 3.2x on HumanEval where EAGLE-3 gets 2.2x.

But an architecture is an empty socket. Nobody hands you a drafter trained on your target, for your workload. You train it, or you run the generic one and eat the 1.93x.

Training a good drafter is small-model training, and that is the part we are good at. Fast Apply and Compact made us one of the best teams in the world at it. The thing you learn under 30B: the frontier scaling laws stop applying. Chinchilla says ~20 tokens per parameter is compute-optimal, but that assumes training is the cost. For a model you train once and serve billions of times, it isn't, and the optimum slides hard toward small and overtrained.

A speculator lives exactly there. Small, overtrained, shaped to one distribution.

So we train one per open model, on coding output instead of web text. Generated code reuses templates and the symbols already on screen, and an edit is mostly a copy of the file it edits. A draft that has read a million diffs predicts those tokens. One that read the internet doesn't, which is why code is the highest-speedup task for every speculation method. For Fast Apply, that specialized speculator is what gets apply to 10,500 tok/s and compaction to 33,000. Same Qwen weights you can download. Ours is faster because the speculator riding it was trained, by us, on the work.

2. Kernels optimized for latency or throughput

Decode is memory-bound. At batch 1 the matmul is a GEMV: it reads every weight once per token and does almost no arithmetic per byte, so tokens per second is bandwidth divided by bytes touched. The only lever is the bytes. Store the weights in FP4 (e2m1) and you move 4 bits where bf16 moves 16, a 4x cut in HBM traffic. So the whole job is keeping them at 4 bits until they hit the tensor core.

The default kernels quit right here. They dequantize FP4 to bf16 in HBM, hand the 4x straight back, and call a stock GEMM. Do it right and you never materialize bf16 at all. You stream the packed weights through cp.async/TMA, unpack e2m1 and the per-16 block scale in registers, and feed Blackwell's FP4 tensor cores straight out of shared memory. On an 80B MoE that's a grouped GEMM over only the active experts, plus a persistent kernel so every SM stays resident at batch 1, plus split-K to fill a machine the GEMV is too skinny to saturate alone. Miss any one of those and you fall back to the bf16 bound.

Tuning that across architectures is brutal, and the defaults are written for the cards frontier labs buy. Port a kernel without retuning and it runs at 7% of optimal; reaching state of the art on AMD's MI250 took rewriting 40% of a flash-attention kernel by hand. So we don't hand-write the search space. A kernel is verifiable, correct against a reference output or not, and that makes it a search you automate.

Autoresearch loopa kernel is correct or it isn't, so the search is automatable
01
Propose
Generate a kernel for the target GPU.
02
Verify
Check output against a reference. Wrong → discard.
03
Benchmark
Time it on real production traces. Slower than baseline → discard.
04
Ship
Correct and faster → into the fleet.
162 tok/swarp-decode on a $7K card1.67× TRT-LLM CUTLASSlow-demand NVIDIA + AMD, hardware nobody else tunes for

The harness proposes a kernel, verifies it against production traces, benchmarks it on the low-demand NVIDIA and AMD boxes nobody else tunes, and ships the winners. KernelBench shows why it has to run automated: scored on correct-and-faster, frontier models clear under 20% of tasks cold. Volume through a tight verify loop is the only way through.

One output: our warp-decode kernels run an 80B MoE at 162 tok/s on a $7K card, up from 97, past a $25K H100's 120. No accuracy loss, code open. It only pays because compute is scarce, which put a price on the cracks the general stack stepped over.

3. Interconnect optimized for high cache hit rates

Programming traffic shares 97% of its prefix tokens, with prompts 37x to 2,494x longer than the outputs. Cache the prefix and the next request pays only for the new tokens. Hit rate is the cost. The cache abstraction is open and we use it: RadixAttention holds prefixes in a tree, a cache-aware router takes hit rate from 20% to 75%, HiCache spills the tree to host RAM and remote storage and, on Qwen3-Coder-480B, moves hit rate from 40% to 80% and doubles throughput.

Cheap GPUs come with a catch. No NVLink.

NVLink moves 900 GB/s between GPUs. PCIe Gen5, the bus on the affordable boxes, moves 64 GB/s per direction. 14x less. Invisible until you shard a model across GPUs, then it is everything: tensor parallelism fires an all-reduce after the attention and the MLP of every layer, and that all-reduce costs 8-11% of the step on NVLink and 40-75% on PCIe. No fast fabric, and communication eats most of the forward pass.

The standard fix is to buy NVLink. We wrote the kernel instead.

KV cache hierarchy on NVLink-denied boxes
L1
GPU HBM
local, fastest, smallest
L2
Host RAM
local spillover
L3
Neighbor over TCP
pulled instead of recomputed
Interconnect
NVLink900 GB/s
PCIe Gen564 GB/s · 14× less
All-reduce share of step
on NVLink8–11%
on PCIe40–75%, hidden by our kernels

NCCL's ring all-reduce is built for big payloads on a fast fabric. At batch 1 the messages are a few hundred KB and latency-bound, so the ring's 2(N-1) hops dominate. We run a one-shot all-reduce instead: each GPU pushes its shard straight into peer memory over PCIe P2P and reduces locally, one round trip. We launch it on a separate stream so it overlaps the next layer's GEMM, fuse the reduce into the down-projection epilogue rather than launching a second kernel, and quantize the payload to FP8 to halve what crosses the bus. CUDA graphs erase the per-layer launch overhead a collective would otherwise pay every layer. That hides most of the 14x gap. The other kernel is a prefix cache that crosses machines over plain TCP.

HiCache already defines a remote L3 tier behind a backend that is three functions: get, exist, set. That runs over any transport. The catch is that its published wins are over RDMA, where a transfer is sub-millisecond and costs under 0.1% of request latency. Plain TCP is an order of magnitude slower. On a PCIe-only box the open stack quietly falls over, because the number it quoted you assumed hardware you don't have.

So the TCP win can't come from the transport. It comes from the hit rate. The trained speculator and the autoresearched kernels drive the rate high enough that a prefix which misses on the GPU and in host RAM gets pulled from a neighbor over TCP instead of recomputed, and skipping a prefill beats the slow fetch. Against full recompute that fetch cuts time-to-first-token 84%.

The fast fabric everyone buys to avoid this, we replaced with kernels. We run the GPUs the market wrote off, at hit rates that are supposed to require the hardware we didn't buy.

One workload

Three things, one loop:

  • The speculator drafts the model's own coding output.
  • The kernels keep the cache hot on hardware nobody else supports.
  • The network shares that cache across boxes never wired to share anything.

None of it is general. All of it points at the coding agent, the highest-volume workload in AI Why all agents will become coding agents:

If you're shipping an agent that codes, the stack is one import away.

The fastest deployments of everything above are private. A speculator trained on your traffic instead of ours, the cache tuned to your workload, dedicated capacity, pricing under the public per-token rates. Over 100 billion tokens a day run through Morph this way. If that's the scale you're heading toward, talk to us.

If this work interests you, consider joining us. https://www.workatastartup.com/jobs/92617