Kimi K3, a 2.8T-parameter Mixture-of-Experts (MoE) transformer announced on July 16, supports a 1-million-token context window and delivers frontier intelligence across a wide range of knowledge work, with exceptional long-horizon reasoning and coding capabilities.

Working closely with the Moonshot AI team, TokenSpeed enabled Day 0 support for Kimi K3 on today's leading platforms—NVIDIA (G)B200/(G)B300 and AMD MI350X/MI355X—within just one week of the announcement. It is available to try out today; see the TokenSpeed deployment guide to get started. We are also laying the foundations for next-generation platforms such as NVIDIA Vera Rubin and AMD MI455X.

Within one week of announcement, Kimi K3 runs on both NVIDIA and AMD flagship platforms through TokenSpeed, with full prefix caching, speculative decoding, disaggregated serving, and CUDA-graph decode.

This post covers Kimi K3's model architecture, TokenSpeed's KV cache management and disaggregated serving support, portable and specialized kernel optimizations, and some early work-in-progress kernel performance numbers. We continue to optimize every layer of the stack—pushing tokens toward the speed of light.

A Frontier Model with an Innovative Architecture

Moonshot AI presents Kimi K3 as the strongest open model tested, with overall performance behind only Claude Fable 5 and GPT-5.6 Sol. It leads several agentic benchmarks, including AutomationBench (30.8), SpreadsheetBench 2 (34.8), and BrowseComp (91.2), scores 67.5 on DeepSWE, and remains near the top on visual benchmarks. See Kimi K3: Open Frontier Intelligence for details.

TokenSpeed's Day 0 support preserves model quality: benchmark scores on both NVIDIA and AMD platforms match the reference within noise.

Benchmark	Reference	NVIDIA Blackwell	AMD CDNA4
OCRBench	0.89	0.882	0.874
MMMU-Pro	0.82	0.8214	TBD
DeepSWE	67.2	66.7	TBD

Kimi K3's model architecture is captured in the following diagram from the official blog post:

• Kimi Delta Attention (KDA) is a recurrent linear attention mechanism that replaces the growing KV cache with a fixed-size memory state, enabling efficient long-context processing and nearly constant-memory decoding while preserving most of the capabilities of full attention.
• Gated Multi-head Latent Attention (Gated MLA) is a periodic full softmax attention layer that stores compressed latent KV representations and uses a gate to control the attention output, allowing the model to perform accurate global token retrieval that complements KDA's compressed memory.
• Stable LatentMoE is a sparse MoE feed-forward layer that routes each token to only a small subset of experts (16 out of 896) in a compressed latent space, greatly increasing model capacity while keeping computation and memory efficient.
• Attention Residual (AttnRes) is an attention mechanism over previous layer representations rather than previous tokens, allowing each layer to selectively retrieve information from earlier layers and improving information flow through very deep networks beyond what standard residual connections provide.

Kimi K3 introduces a large number of architectural innovations, and each poses unique challenges that require extensive optimization across the entire inference stack to deliver excellent efficiency.

TokenSpeed tackles these challenges through clear modularity and layering, built on generic and unified mechanisms.

A Generic, Unified Flat KV Layout for Cache Management

We debuted flat KV cache management with our Day 0 support for TML Inkling. Inkling has even more heterogeneous cache management needs than Kimi K3, so supporting Kimi K3 was a smooth process—and further validates the design choices behind our generic, unified cache mechanism.

Specifically, Kimi K3 is served with a flat KV layout: one physical page pool shared by every cache family. A page holds either 1,536 tokens of MLA latent history or one KDA recurrent snapshot (conv + recurrent state), so every byte of KV budget is fungible between attention history and recurrent state. The pool is organized into 24 physical slabs. Every slab contains one MLA page, while the first 23 slabs additionally contain three KDA pages, one from each of the three 23-layer KDA groups. A single global page ID indexes the corresponding page across all slabs, enabling prefix caching, copy-on-write state snapshots, and CUDA-graph decode to operate on the same shared page pool.

The same flat KV abstraction extends to disaggregated serving. In P/D deployments, TokenSpeed moves Kimi K3's MLA history and KDA recurrent snapshots through one global page-ID space, keeping the transfer path independent of the underlying cache family. For multimodal workloads, TokenSpeed further separates vision encoding from language prefill and decode: SMG exposes Encoder, Prefill, and Decode as independently routable worker pools, while Mooncake transfers multimodal embeddings from E to P and flat KV pages from P to D. Each stage can scale according to its own compute and memory characteristics, without introducing model-specific transfer protocols.

Marrying Portability and Performance in Kernel Management

Kimi K3 brings innovation to every major kernel in the transformer architecture: KDA, Gated MLA, Stable LatentMoE, and AttnRes. Adapting existing top-performing specialized kernels—or developing new ones—for all of them would take far longer than one week.

A core design goal of the TokenSpeed kernel layer is to marry portability with performance. We deliberately maintain in-tree Triton implementations for fast enablement of new models. Combined with customized FLA Triton kernels, Kimi K3 was up and running on both NVIDIA and AMD GPUs very quickly, with decent out-of-the-box performance relative to the PyTorch fallback.

We believe the importance of portable kernels is only increasing—especially with next-generation platforms like NVIDIA Vera Rubin and AMD MI455X on the horizon, where portable kernels will let us shorten bring-up time significantly.

With a default Triton path in place, the rest is principled kernel development: benchmark where the bottleneck is, then add or fuse specialized kernels. All of this happens without touching the vendor-neutral modeling code, thanks to TokenSpeed's modular architecture and unified kernel API.

We are still pushing on these fronts; below we zoom in on a few highlights.

TensorRT-LLM and CuteDSL kernels on NVIDIA GPUs

CuteDSL KDA

The CuteDSL KDA prefill kernel restructures the scan around 16-token chunks. Blocking the recurrence this way exposes the bulk of the per-chunk work as dense matrix multiplications that map efficiently onto the tensor cores, and it allows the running state to advance one chunk at a time instead of one token at a time—shortening the sequential dependency chain by a factor of sixteen.

The decisive property, though, is data residency. A single CTA processes an entire (sequence, head) pair in one fused pass, so the running state lives in registers and shared memory for the full traversal and is never spilled to global memory. Since that state traffic was the dominant term in the naive version's bandwidth budget, removing it brings the memory footprint down to its irreducible floor: read the Q/K/V and gate tensors once, and write the output once.

TensorRT-LLM Stable LatentMoE

On the MoE path, the TensorRT-LLM kernel fuses the entire expert MLP into a single runner call over the routed tokens: FC1 with the SiTU gate as one gated GEMM, followed immediately by FC2, eliminating separate GEMM and activation launches.

TokenSpeed supports both activation paths with the same MXFP4 weights. w4a16 keeps activations in bf16, while w4a8 block-quantizes them to MXFP8 so FC1 runs as an FP8-activation GEMM. w4a16 is the default, and w4a8 targets large-batch prefill.

The better path depends on batch size. At large batch size, the expert GEMMs are compute-bound, so FP8 activations increase arithmetic throughput. At small batch size, both paths are dominated by MXFP4 weight loads, making FC1 latency nearly identical. The difference is that w4a8 additionally quantizes FC1's output to MXFP8 before FC2. This fixed elementwise overhead is not amortized at small batch, making w4a8 slower than w4a16, but is outweighed by the GEMM speedup at large batch.

AttnRes

TokenSpeed also uses an optimized fused kernel for K3's attention-residual mixing, which forms a learned softmax mixture over several residual candidates per token. It runs in fp32 because the result feeds the global residual stream.

The kernel is a single TMA-fed, warp-specialized Blackwell pass. It streams the candidate rows in through TMA, converts each to fp32 once and holds it in tensor memory across the passes it makes over the candidates, and uses an online softmax so the weighted mixture accumulates in fp32 registers in one sweep instead of requiring a second read. The trailing RMSNorm is folded into the epilogue, and small-candidate-count cases are specialized down to a single candidate.

Gluon kernels on AMD GPUs

Gated MLA

For Gated MLA, we developed separate specialized Gluon kernels for prefill and decode, given their different problem characteristics.

For prefill, the kernel launches 256 persistent workgroups with scheduling logic that keeps the workload balanced across XCDs, as described in the TokenSpeed Kernel post. Internally, the kernel double-buffers K/V in shared memory, overlapping asynchronous loads with fused NoPE/RoPE scoring, causal masking, online softmax, and V accumulation.

Fig: Compared to the portable Triton kernel, the gain grows as longer sequences spend more time in the double-buffered main loop.

For decode, the Gluon kernel partitions along the context length rather than across query heads. It divides the paged history across 64 workgroups, processes 128 KV tokens per tile, and combines the partial softmax states in a secondary kernel. This keeps the GPU fully occupied even with a single query token.

Fig: The portable Triton kernel launches only 12 workgroups—one per local query head—and each scans the full context serially in 16-token tiles. Gluon handles the heads together and parallelizes the scan across 64 workgroups, so the gap widens rapidly with context length. At 4K, Gluon's 17.95 µs consists of 6.08 µs of attention and 11.87 µs of split-state reduction. The nearly fixed reduction cost dominates at short contexts; parallel KV traversal drives the speedup to 134.78x over our in-tree portable Triton baseline at 16K.

Stable LatentMoE

For Stable LatentMoE, we extend the warp-decode Gluon kernel. The core idea is to organize around outputs rather than experts: each wave owns one output column and reduces the entire K dimension for it, pulling from whichever expert that column needs. Route sorting, row padding, and gather buffers are never built. We first introduced this structure for GPT-OSS in the TokenSpeed kernel post.

The Kimi K3 version likewise consists of two staged kernels: one fusing the gate and up-projection GEMV with SiTU, the other handling the down-projection GEMV, together completing the warp-decode MoE kernel. They operate on the compressed 3,584-wide latent space rather than the 7,168-wide hidden state. The activations are in BF16 while weights are in MXFP4, expanded through CDNA4's native scaled upcast inside the K loop to save memory bandwidth.

Fig: The specialized Gluon path is up to 8.65x faster than the in-tree Triton baseline at small batches, because it avoids the routing, padding, metadata, and grouped-GEMM overhead that dominates when each EP8 rank processes only a few local routes. As batch size grows, those fixed costs are amortized and both implementations spend more time in the projection math and memory traffic, narrowing the advantage to roughly 1.5–2x, with Gluon remaining faster across all tested shapes.

Iris and communication

TP8 and EP8 communication runs on Iris, AMD's multi-GPU communication library written entirely in Python and Triton. Iris exposes a symmetric heap in which every rank allocates the same buffer at the same logical offset and reads a peer's copy directly.

Our collectives write their partials straight into the symmetric buffer, and the readiness handshake costs seven flag reads per collective rather than one per block per peer. We also merge the two reductions each MoE layer issues: the shared-expert 7,168-wide partial is not consumed until after the routed result has been projected back to full width, so both reduce as segments of one collective addressed through disjoint masks. At 14 KiB and 7 KiB, neither payload is bandwidth-bound—what the merge reduces is the number of moments where eight ranks have to agree, 184 times per token. Two ordinary calls at a combined 46.18 µs become one call at 31.68 µs.

Instead of launching a separate norm kernel after the reduction, we run the routed latent's RMSNorm inside the collective. RMSNorm needs a sum over an entire row, and a collective normally splits that row across workgroups. To avoid paying for partial sums plus another pass to combine them, we assign the routed row to a single workgroup: it gathers the row from all eight ranks and holds every value, so the sum stays local. The two-segment reduction went from 8.768 ms to 7.144 ms of summed GPU time while also absorbing the norm; the single-tensor reduction went from 2.718 ms to 0.783 ms.

Looking Forward

We envision—and continue to push—TokenSpeed as a speed-of-light LLM inference engine for agentic workloads. The rapid Day 0 enablement of recent models within one to two weeks, validates our design choices of a modular architecture and a unified kernel API.

The strong momentum in open-weight model development is helping us solidify the core pieces of the stack: KV cache management, disaggregated serving, and key kernel development. There is, as always, plenty left to optimize—and looking ahead, we are also pushing on system- and rack-level fundamentals, including distributed KV cache management and more. Stay tuned.

Acknowledgments

This work would not have been possible without close collaboration across teams and companies. We are deeply grateful to the Moonshot AI team for the close partnership and early access that made Day 0 enablement possible, to NVIDIA DevTech for their kernel engineering support on Blackwell, to the AMD Triton team for their work on the CDNA4 compiler stack and Gluon kernel development, and to the Mooncake team for their support on the disaggregated serving transfer path. Thank you to everyone in the open-source community who helped along the way.