KOG LaneFormer-2B & the monokernel — co-designing a model, runtime, and GPU kernel for batch-1 latency.

KOG generate 3,000 output tokens/second on a single request (batch size 1) from a 2B model on an 8× AMD MI300X node — and 2,100 tok/s on 8× NVIDIA H200 — with no quantization, no speculative decoding, no pruning. They do it by treating three things usually owned by three different teams as one co-designed system: the model (LaneFormer-2B), the parallelism scheme (Delayed Tensor Parallelism), and the GPU runtime (a single persistent "monokernel" for the whole decode pass). This is a purely technical walkthrough of how each piece works and why, at batch 1, the bottleneck is memory bandwidth and inter-GPU synchronization — not FLOPs.

The problem: batch-1 decode is a latency problem, not a throughput problem

Almost all production LLM serving optimizes throughput: pack many requests into a big batch so the GPU's matrix engines stay busy and amortize fixed costs across tokens. That maximizes tokens/second/GPU but does nothing for the latency of a single stream. For a real-time agent — one that must read its own output, call a tool, and react inside a tight loop — what matters is how fast one request decodes. That is batch size 1, and it has completely different physics.

At batch 1, autoregressive decoding generates one token at a time, and producing each token requires reading all the active model weights out of HBM once. The math per layer is a matrix–vector product (GEMV), not a matrix–matrix product (GEMM). GEMV has very low arithmetic intensity: a handful of FLOPs per byte loaded. So the GPU's tensor cores sit mostly idle; the limiting resource is how fast you can stream weights through memory. The metric that matters is Memory-Bandwidth Utilization (MBU), not FLOPs utilization.

A worked bandwidth ceiling

Put numbers on it. LaneFormer-2B is ~2.3B parameters at FP16 (2 bytes), so the weights are about \(2.3\text{B}\times 2 = 4.6\) GB. One decoded token must stream all of them once. The model is tensor-parallel across 8 MI300X GPUs, each holding \(\tfrac18\) of the weights (~0.575 GB) and streaming its shard with its own ~4.3 TB/s of bandwidth — so the effective bandwidth for the whole model is roughly \(8\times 4.3 = 34.4\) TB/s:

The measured 3,000 tok/s is therefore around 40% of the naive weight-only ceiling — consistent with the ~36% MBU they quote — with the remaining gap eaten by KV-cache streaming, the residual cross-GPU synchronization, attention math, and sampling. The point of the worked figure is the shape of the problem: you are nowhere near compute-bound; every microsecond not spent moving useful weight bytes is wasted ceiling. That is why the rest of the article is about removing overheads (launches, blocking comm, cache pollution) rather than adding compute tricks.

The thesis: co-design, not a faster kernel bolted onto an existing model

The central claim is that the fastest single-request inference comes from designing the model, the parallelism, and the runtime together. Standard stacks layer a generic model (defined in PyTorch) on top of a generic runtime (vLLM / TensorRT-LLM) on top of generic kernels (cuBLAS, CUTLASS, Triton, NCCL). Each layer is individually good, but the boundaries between them impose costs — kernel launches, intermediate tensors, blocking collectives — that dominate at batch 1. KOG instead build the whole vertical slice in-house and let decisions in one layer simplify another. Three pillars:

Layer	KOG's choice	What it removes
Model	LaneFormer-2B (lane-structured, GQA, sliding-window)	Excess weight & KV streaming; awkward sharding
Parallelism	Delayed Tensor Parallelism (DTP)	Per-layer blocking all-reduce
Runtime	Single persistent monokernel (HIP/CUDA + inline asm)	Kernel launches, CPU scheduling, intermediate round-trips

Pillar 1 — the LaneFormer-2B model

LaneFormer is a Transformer variant whose layout is chosen for 8-way splitting across the 8 GPUs of a node. The unit of split is a "lane," and the model is described in terms of one lane replicated eight times. Headline specs:

Property	Value
Total parameters	~2.3B
Layers	15
Attention	Causal Grouped-Query Attention (GQA)
Heads	32 query heads, 16 key/value heads — evenly sharded across the 8 lanes
Per-lane heads	4 Q, 2 K, 2 V (head dim 96) → 768 QKV outputs per lane
Sliding-window attention	in 10 of the 15 layers
Context length	4,096 (128k extension in progress)
Precision	FP16 (no quantization)
Tokenizer	Llama-2-based
Training	~4T tokens pre-train + ~2T mid-train (code/reasoning) + ~210M post-train; via NVIDIA Nemotron datasets on 256 H100s
Quality	~50% HumanEval (competitive at 2B scale)

Why GQA and sliding windows — both cut what you stream

Every architectural choice serves the batch-1 bandwidth floor. Grouped-Query Attention shares each key/value head across several query heads (here 32 Q over 16 KV), which halves the KV cache that must be read every step relative to full multi-head attention. Sliding-Window Attention in 10 of 15 layers bounds how much of the KV cache each of those layers attends to, so KV streaming does not grow without limit as the context fills. Both are "conservative" choices — well-understood, not novel — deliberately so: the only genuinely new idea is the parallelism scheme, and the rest of the model is kept boring to isolate that variable.

What a "lane" actually is

A lane is a vertical slice of the model — its share of the attention heads and MLP — assigned to one GPU. Because heads divide evenly (32 Q / 8 lanes = 4 Q per lane; 16 KV / 8 = 2 KV per lane), each GPU owns a clean, equal partition with no ragged remainder. The QKV projection on a lane therefore emits exactly \(4\times96 + 2\times96 + 2\times96 = 768\) scalars per token. This evenness is not cosmetic: it means the monokernel's compile-time work partition (below) maps cleanly onto the hardware, and DTP's per-lane bookkeeping stays symmetric.

Pillar 2 — Delayed Tensor Parallelism (DTP)

Splitting a model across 8 GPUs with classic tensor parallelism (TP) means each GPU computes a partial result for every layer, and the partials must be summed across GPUs with an all-reduce before the next layer can start. At batch 1 that all-reduce is a disaster: it is a blocking collective on the critical path of every single layer, and once your kernels and weight streaming are already tuned, this communication becomes the dominant latency term.

Standard TP does, per layer \(n\):

The question DTP asks: can we hide that communication behind computation instead of paying it every layer? The answer is to launch the communication early and consume it late. After a lane computes its output, it immediately ships it to the other GPUs — but instead of waiting, it keeps computing the next \(\delta\) layers using only local outputs. By the time those \(\delta\) layers are done, the data from the other lanes has landed, and it gets aggregated then. The schedule has three phases:

In the steady state each lane folds in its own current output \(o_l^{(n)}\) plus the other lanes' outputs from \(\delta\) layers ago, \(o_j^{(n-\delta)}\). The delay \(\delta\) is chosen (calibrated to 2 layers for LaneFormer) to give the interconnect "enough time for communications from each device to land on the other devices." Crucially this is a training-time change as well as an inference one — the model is trained with the delayed aggregation baked in, so there is no accuracy penalty: KOG report training loss comparable to a fully-synchronized baseline.

Pillar 3 — the monokernel runtime

The third pillar is the most extreme. Instead of launching a sequence of kernels (one for QKV projection, one for attention, one for the FFN, one for sampling, repeated per layer), KOG run the entire decode pass in a single persistent GPU kernel — a "monokernel." It launches once, stays resident, and never returns to the host until the token is produced.

The arithmetic that forces this

A kernel launch on the MI300X costs roughly 4.5 µs of overhead, plus the HBM weight stream is interrupted at each boundary (~0.5 µs restart penalty) and intermediate tensors must be written to and read back from memory (>1 µs round-trip). Multiply by the number of kernels. KOG's illustrative figure: a 25-layer model at their working batch would burn about \(25\times(\text{several launches})\approx 1{,}125\,\mu s\) of pure launch overhead per token — capping speed near 890 tokens/s before any other optimization. To beat that ceiling you cannot launch per stage; you must fuse everything into one resident program.

Compile-time work partition (no scheduler)

Inside the monokernel there is no dynamic task scheduling. Work is partitioned across the hardware at compile time by the programmer. The kernel launches with fixed dimensions — on MI300X, gridDim = (256,) and blockDim = (64, 8) — mapping one logical block onto each of the 256 active Compute Units across the chip's 8 chiplets. Because the assignment is static, there is zero runtime overhead for "what should this CU do next" — it already knows.

// conceptual shape of the persistent kernel
__global__ void decode_monokernel(...) {   // gridDim=256, blockDim=(64,8)
  persistent: for each of the 15 layers {
     qkv_projection();      // GEMV via dot2, FP32 accumulate, DPP wave-reduce
     attention();           // 3-stage: per-CU → cross-CU softmax fix → out-proj
     launch_dtp_comm();      // async ship lane output to peers (non-blocking)
     ffn();                 // streams weights while attention of next runs
     consume_dtp_comm(layer-δ);  // fold in peers' outputs from δ layers ago
  }
  sample();                 // argmax(logits + Gumbel noise), no full softmax
}

Weight streaming — keeping HBM saturated

Since batch-1 decode is bandwidth-bound, the whole game is keeping the HBM pipe full of useful weights. KOG use two loading paths: weights routed through AMD's 64 KiB Local Data Share (LDS) for flexible thread access and lower register pressure, and direct register loads for the rest. All weight loads carry non-temporal (NT) scope bits, telling the cache "don't bother keeping these" — weights are read once per token and never reused, so caching them would only evict genuinely reusable data. The FFN (almost 80% of each layer's weights) is prefetched during the compute-heavy attention stage, so its bytes are already arriving when it is time to use them. Offline, weights are pre-transformed (RMSNorm folded in, layouts repacked) so no rearrangement happens on the hot path.

Cross-GPU synchronization without atomics

With DTP shipping partials between GPUs, the kernel still needs a fast way to know "has the peer's data arrived?" Conventional grid sync uses arrival counters with atomic increments and epoch numbers — KOG measured that at 7.59–7.88 µs. They replace it with publish-dependent buffers: a destination buffer is initialized to a sentinel (NaN), and the consumer simply polls that memory until a real (non-NaN) value appears. No atomics, no counters. This drops synchronization latency to 0.80–0.93 µs — roughly a 9× reduction. It matters because grid synchronization is about 35% of token-generation time; they further flatten cross-die latency by duplicating tensors per I/O die (topology-aware placement).

The fused operations, briefly

• QKV projection — a GEMV computed on the scalar/vector ALUs with dot2 instructions (not the matrix cores, which are wasted at batch 1), partial products accumulated in FP32 and reduced within a wavefront via Data-Parallel Primitive (DPP) ops.
• Attention — a three-stage pipeline: each CU computes attention over its tile of the KV cache, then a cross-CU aggregation applies the softmax correction (online-softmax style), then the output projection. Partitioning the sequence dimension across CUs avoids any single CU owning a long column.
• FFN — weight-streamed and overlapped with attention as above.
• Sampling — done with the Gumbel-max trick: \(\arg\max_i(\text{logit}_i + g_i)\) with \(g_i\) i.i.d. Gumbel noise is distributionally equivalent to sampling from the softmax, but requires no softmax normalization. Because lanes hold shards of the vocabulary, sharing a partially-reduced argmax instead of full probability vectors cuts cross-device vocabulary traffic by roughly 800%.

Why it needed inline assembly

The compiler could not always express what the design required, so the kernel uses targeted inline assembly: __hip_atomic_load/store with 3-dword (96-bit) data types and global_load_dwordx3 with sc1 scoped-L2-coherence bits (necessary on MI300X's chiplet memory), and manual s_waitcnt placement where the compiler could not track memory-ordering dependencies. They also had to inspect generated assembly to prevent Loop-Invariant Code Motion from causing register spills, and to stop the compiler inserting s_waitcnt vmcnt(0) barriers that would stall the very weight streaming the design depends on.

Why standard frameworks leave the bandwidth on the table

A natural objection: vLLM and TensorRT-LLM are excellent and battle-tested — why can't they hit the same numbers? Because their generality, which is exactly what makes them good throughput servers, imposes batch-1 costs they cannot remove:

Generic-stack cost	Effect at batch 1	Monokernel answer
One kernel per stage/layer	~4.5 µs launch × many → an 890 tok/s-class ceiling	Single persistent kernel, launched once
Intermediate tensors in HBM	>1 µs store/load round-trip per boundary	Stages fused; activations stay in registers/LDS
Blocking NCCL all-reduce per layer	Communication on the critical path	DTP — comm hidden behind δ layers of compute
Atomic-counter grid sync	~7.8 µs per sync, ×35% of token time	Publish-dependent (NaN-poll) buffers, ~0.8 µs
Cache-resident weights	Single-use weights evict reusable data	Non-temporal load hints
Matrix cores for GEMV	Tensor cores idle at low arithmetic intensity	Scalar/vector ALU dot2 + DPP reductions

None of these is a flaw in the generic frameworks; each is the price of supporting arbitrary models, arbitrary batch sizes, and a clean module boundary between model code and runtime. KOG trade that generality away — their stack runs their model on their runtime — and recover the difference. That is the whole bet: a "software ceiling" sits well below the hardware's physical bandwidth limit, and you reach the hardware limit only by collapsing the software layers that created the ceiling. The reported 3.5× speedup over standard frameworks on Instinct is the size of that gap.

The numbers, and what they're bounded by

Node	Tokens/s per request (batch 1)	Precision
8× AMD MI300X	3,000+	FP16
8× NVIDIA H200	2,100	FP16

For reference, the MI300X result is in the same range KOG cite for purpose-built wafer-scale hardware (Cerebras WSE quoted at ~3,000 tok/s on a much larger model), achieved here on standard datacenter GPUs. The MI300X provides ~4.3 TB/s empirical bandwidth, and the entire stack exists to convert as much of that as possible into tokens. KOG also report a 3.5× generation speedup versus standard frameworks on AMD Instinct — the gap they attribute to those frameworks hitting a "software ceiling" (launch overhead, blocking collectives, cache pollution) well below the hardware's physical limit.

What 3,000 tok/s actually buys an agent

The latency framing only matters if the number crosses a useful threshold. 3,000 tokens/s is ~0.33 ms per token; 2,100 tok/s is ~0.48 ms. An agent that must emit, say, a 200-token tool call and then react sees that step complete in ~67 ms at MI300X speed versus several hundred milliseconds on a throughput-tuned server delivering a few hundred single-stream tok/s. For a loop that runs many such steps — plan, call, read result, revise — the difference compounds into the gap between an agent that feels interactive and one that feels like a batch job. This is the workload KOG explicitly target: single-request, low-latency, tight agentic loops, where you genuinely cannot batch because there is only one logical conversation in flight and each step depends on the last. Throughput servers win on cost-per-token at scale; latency engines win here, and the two are not substitutes.

The NVIDIA port — same idea, different floor

The headline 3,000 tok/s is on AMD MI300X, but KOG built the same monokernel design for NVIDIA, reaching 2,100 tok/s/request on an 8× H200 node. The architecture-level ideas port directly — single persistent kernel, DTP, NT weight streaming, polling sync, Gumbel-argmax sampling — because none of them depend on a vendor; they depend on the shape of batch-1 decode. What changes is the low-level substrate: HIP and AMD-specific inline assembly (chiplet-scoped sc1 coherence bits, dot2/DPP, LDS sizing) become CUDA and the H200's own memory hierarchy and bandwidth. The fact that the same co-design wins on both vendors is itself evidence that the gains come from removing software overheads, not from exploiting one chip's quirk — the quirks only decide where the new, higher ceiling sits.

Scaling the idea to frontier models

A 2B model fits comfortably and streams fast; the interesting question is whether the approach survives at frontier scale. KOG's projection applies the same MBU-bound reasoning to large MoE models with quantization: for a DeepSeek-V4-class system (~49B active parameters per token, FP8/MXFP4), they estimate roughly 305–1,410 tokens/s per request depending on GPU generation, assuming a conservative 36% MBU. The active-parameter count of MoE is what makes this plausible — only the active experts' weights stream per token, so a 1.6T-total model can behave, bandwidth-wise, like a ~49B one. The stated target band is 1,000–5,000 tok/s/request on standard datacenter GPUs as hardware improves.

Technical takeaways

• Batch-1 latency is a memory-bandwidth problem. Optimize MBU, not FLOPs; the per-token floor is \(P\cdot b / B\).
• The boundaries between model, runtime, and kernel are where latency hides. Kernel launches (~4.5 µs each), blocking all-reduce, and intermediate round-trips dominate once the math itself is tuned.
• Delayed Tensor Parallelism removes per-layer communication from the critical path by launching transfers early and consuming them \(\delta\) layers later — and it is trained in, so accuracy is preserved.
• The monokernel fuses the whole decode pass into one persistent kernel with compile-time work partition, NT weight streaming, FFN/attention overlap, polling-based sync (0.8 µs vs 7.8 µs), and Gumbel-argmax sampling.
• The model is deliberately conservative (GQA, sliding window, even head sharding) so that the one genuinely new idea — DTP — can be isolated, and so that what must stream per token is as small and as evenly distributed as possible.

References

• KOG — LaneFormer-2B: the latency-first model (Hugging Face) · model architecture, GQA, lanes, training
• KOG — Real-time LLM inference on standard GPUs (3,000 tok/s/request) · the co-design overview & numbers
• KOG — Single-kernel, latency-optimized inference on MI300X · the monokernel deep dive
• KOG — Delayed Tensor Parallelism · the DTP method & schedule
• AMD — KOG reaches 3.5× inference speed on Instinct MI300X · independent corroboration