Kimi K3 presents an interesting inference optimization problem. Its architecture combines KDA, Gated MLA, Stable LatentMoE, and AttnRes, while the agentic workloads we care about place the model in a regime of long, continuously growing contexts and short, latency-sensitive decode segments.

Our Day-0 post described the work required to support these architectural components correctly in TokenSpeed. This three-part series focuses on what came next: optimizing their execution on GB300 systems under realistic serving workloads.

In Part I, we focus on TP8 within a single NVL72 NVLink domain, using multi-turn SWE-Smith traces as the representative workload. We study concurrency from 1 to 16 and organize the optimization analysis around the target-model decode and verification path under EAGLE3 speculative decoding with lightseekorg/kimi-k3-eagle3-mla. Rather than treating one peak decode rate as the result, we examine the trade-off between aggregate throughput and per-user generation speed across concurrency levels.

End-to-end serving on long-context agentic workloads

The end-to-end result is best read as a curve rather than a single peak number. As concurrency increases, each operating point trades per-user generation speed for aggregate system throughput. We compare EAGLE3 with non-speculative execution across the same frozen workload to show how that frontier moves.

Component	Configuration
GPU	NVIDIA GB300
Topology	8 GPUs across two 4-GPU compute trays
Parallelism	Attention TP8, MoE TP8
Target model	nvidia/Kimi-K3-NVFP4
Draft model	lightseekorg/kimi-k3-eagle3-mla
Draft precision	BF16, unquantized
MLA backend	TokenSpeed MLA
KDA backend	CuteDSL prefill; fused Triton decode/verify
MoE backend	FlashInfer TensorRT-LLM
KV-cache dtype	FP8

Workload and metrics. We use SWE-Smith because its traces match the shape of an agentic coding session: long prompts, repeated tool-call traces, and multiple turns that continually extend the same conversation. Each turn reuses a large history, appends tool output, and asks for another relatively short response. The frozen workload begins at roughly 50K input tokens, adds roughly 800 tokens per subsequent turn, runs for 10–15 turns, and generates up to 500 tokens per turn. The longest conversations reach roughly 68K input tokens.

We report aggregate throughput in tokens per minute per GPU (TPM/GPU) and per-user generation speed in tokens per second (TPS/user, derived from TPOT). Together, they capture the trade-off between overall system throughput and generation speed as concurrency increases; TPS/user does not include time to the first token.

An individual TPS/user point is specific to this setup: long, continuously growing histories, this concurrency sweep, and this serving configuration. It should not be compared directly with headline TPS from a different workload, context length, concurrency, or measurement methodology.

The figure plots TPS/user on the x-axis and TPM/GPU on the y-axis, with each point representing a measured concurrency level. EAGLE3 uses three draft steps and a four-position target verification window.

Fig: Throughput–latency frontier on multi-turn SWE-Smith agentic traces.

EAGLE3 shifts the measured throughput–latency curve outward: its benefit is visible both at low concurrency and as the workload moves toward higher aggregate throughput. At an illustrative operating point of 50 TPS/user, interpolation within this benchmark gives approximately 459K TPM/GPU.

On this workload, advancing multiple tokens per target iteration outweighs the additional drafting and verification work. The following deep dives look inside the target-model execution path, using local comparisons to examine the computation, communication, and state-storage choices behind this TP8 implementation.

Optimization deep dives

The profile separates into three recurring sources of overhead: the trade-off between replicated projections and communication-aware sharding in LatentMoE, temporary state storage in KDA, and low-M execution around the main model operators. These costs recur across 93 decoder layers, including 69 KDA layers, 24 MLA layers, and 92 MoE blocks. Multi-token verification makes some of the shapes and state-lifetime constraints more visible, but the focus of the following sections is the K3 target path rather than the speculative decoding algorithm.

Communication-aware sharding for LatentMoE

Our initial implementation kept both routed latent projections replicated: every TP rank stored and evaluated the full hidden-to-latent down projection before the routed experts and the full latent-to-hidden up projection after them. This avoided explicit assembly of projection shards, but duplicated both weight storage and projection work.

Ideally, we would want to shard both projections, but the communication trade-off is different on the two sides of the MoE block. Sharding the down projection introduces a new requirement to assemble the locally computed latent shards into the full routed latent. Sharding the up projection, in contrast, creates an opportunity to restructure the existing post-expert communication path: the routed and shared expert reductions, local projection shard, and final hidden-state assembly can be composed rather than executed as separate stages. The goal is therefore to keep the added assembly cost on the down path small, while exploiting the up-projection sharding to reduce both replicated computation and communication overhead in the post-expert tail.

Routed-down projection: column sharding and latent assembly

At the first boundary, every MoE block projects the hidden state from 7,168 dimensions into the 3,584-dimensional routed latent. In the replicated form, each TP rank stores the full [3584, 7168] BF16 weight and computes every output column. The column-sharded path partitions the output axis across TP8: rank r stores and computes one [448, 7168] block, covering latent columns 448r ... 448r+447. These blocks are then concatenated to restore the full routed latent.

The optimal assembly path depends on M:

• For M ≤ 8, a static-M SIMT kernel computes the local block and publishes it directly through the NVLS multicast virtual address.
• For 8 < M ≤ 1280, a tensor core kernel writes the local block into a strided view of the same multicast address, avoiding a separate publish kernel.
• In both cases, a lamport gather waits for published fragments from all ranks in the pooled two-slot mailbox and assembles the latent.
• Above the mailbox ceiling, each rank writes its block to ordinary memory and a last-dimension all-gather assembles the latent.

Fig: Routed-down projection sharding.

Sharding the routed-down projection reduces its per-rank weight footprint from 49.0 MiB to 6.1 MiB per MoE block, saving approximately 3.9 GiB per GPU across K3's 92 MoE blocks.

Moreover, the latency benefit persists beyond small decode shapes: at M=4, the complete operation drops from 14.9 to 12.8 μs; at M=8192, from 225.1 to 138.9 μs, including the all-gather. The jump between M=1280 and 1281 marks the switch from mailbox assembly to GEMM plus all-gather. Large-M points test projection scalability beyond this workload's decode/verify region.

Fig: Benchmarking of routed-down projection.

Routed-up projection: sharding the post-expert tail

On GB300 systems, the MNNVL all-reduce path adapts the TensorRT-LLM design carried through FlashInfer. Its one- and two-shot Lamport protocols use NVLink multicast memory for small TP reductions. TokenSpeed adds K3-specific epilogues around that path so adjacent post-reduction work can remain fused. This integration was developed with technical guidance and optimization support from NVIDIA DevTech.

The initial implementation reduces routed and shared expert partials separately, then evaluates the full replicated up projection on every rank. The sharded path computes only hidden_size / TP output elements per rank. Its key idea is block injection: each rank adds its projected block and the corresponding prefix slice into the matching columns of its shared-expert partial. Because these blocks are disjoint, the shared all-reduce both sums the shared contributions and assembles the routed output, eliminating a separate routed all-gather. The lane and multimem backends implement this same pattern. The small-M fused tail instead reduce-scatters the shared branch, adds the local shared shard in the projection epilogue, and multicast-publishes the combined block for a Lamport gather that also adds the prefix.

Fig: Routed-up projection sharding and tail fusion.

The routed-down and routed-up projection weights have the same size, so sharding each saves approximately 3.9 GiB per GPU across the 92 MoE blocks. Together, the two changes reduce parameter storage by approximately 7.7 GiB per GPU relative to the replicated layout. Under the cache geometry used here, this corresponds to roughly 0.6 million additional logical FP8 cache tokens. These figures are parameter-size calculations rather than peak allocator measurements.

From materialized expert partials, the complete tail drops from 32.34 to 18.28 μs at M=4 (1.77× with empty-event correction; 1.69× using raw event intervals), and from 915.77 to 656.22 μs at M=8192 (1.40×). The combined curve retains the M=256 transition, where the staged path is 19.4% slower than the control; the later control-side jump at M=2049 marks its fallback from MNNVL to NCCL. These are local composition comparisons against current separate reductions and a replicated up projection, not isolated finalize-fusion gains or end-to-end speedups.

Fig: Benchmarking of routed-up projection sharding and tail fusion.

Replay SSM: reducing KDA verification-state storage

Every token advances convolution and recurrent state in each KDA layer. Under multi-token verification, the accepted length is unknown until the full candidate window has been evaluated, so snapshotting the complete state at every position makes memory scale with both batch size and verify width. Replay SSM instead keeps the committed state as an anchor and records the smaller projection and gate inputs required for reconstruction. Once sampling returns the accepted length, one batched kernel replays the accepted prefix plus the target-sampled token.

For TP8 K3, a snapshot row is 795 KiB per request and KDA layer: 27 KiB of BF16 convolution state plus a 768 KiB FP32 recurrent matrix. At batch 64, storing the anchor and four verify positions across 69 KDA layers would require 16.7 GiB/GPU. Replay uses 0.4 GiB/GPU, saving 16.4 GiB/GPU (97.8%, or 45.0x) with no observed latency increase when the fused verify and replay kernels are used. At 14.1 KiB/GPU per logical target-plus-draft FP8 cache token, the byte saving corresponds to roughly 1.2 million cache tokens before allocator headroom.

KDA verify workspace at max_num_seqs = 64	Per GPU
Per-position state snapshots	16.7 GiB
Replay SSM workspace	0.4 GiB
Memory saved	16.4 GiB (97.8%)
Equivalent FP8 cache capacity	~1.2M tokens

Kernel fusion and shape-aware routing

Apart from the communication and state-management optimizations above, the remaining decode and verify overhead was no longer dominated by a single large operator. Instead, it came from a collection of short projections, intermediate handoffs, routing steps, and metadata preparation around the four-position verification window. Individually these operations are small, but they recur across the model and become visible at the low token counts typical of decode and verification.

Shape-aware projection routing. No single GEMV or GEMM implementation performs best across all projection shapes. We route each (M, N, K) region among Triton row-per-CTA GEMV, CuteDSL skinny GEMM, FlashInfer TGV, and torch.mm. The figure covers the NVFP4 target's unquantized shared-expert projections and selected BF16 EAGLE3 draft projections; target FP8 attention projections are outside this comparison. Every shape uses torch.mm as the baseline. Routing decisions are based on cold-L2 measurements, which better approximate serving behavior than repeatedly benchmarking the same operands from a hot cache. For selected skinny shapes, we additionally use 256-bit loads when pointer alignment permits.

Fig: Benchmarking of shape-aware projection routing.

KDA fusion. Multi-token verification introduces several small operations around the convolution and recurrent update. We fuse the short convolution and recurrence where beneficial, choose between split and fused gate preparation according to shape, and consume strided slices directly from the merged input projection instead of repacking them into contiguous intermediates. The objective is primarily to remove launches and memory traffic around the recurrent core rather than to change the KDA computation itself.

MLA decode and verify fusion. On the MLA path, several decode-stage handoffs can be collapsed before attention. With FP8 KV caching, the optimized path constructs the FP8 query and writes the latent KV cache in the same launch, eliminating intermediate materialization between these steps. These fusions target the small decode and verification shapes; larger prefill shapes continue to use their existing kernels. At M=4, the combined query/cache-write stage drops from 6.43 to 4.39 μs.

AttnRes and speculative metadata. We extend the existing fused AttnRes path to the smaller multi-token verification shapes. Outside the model operators themselves, we also collapse metadata chains that otherwise become visible at low M: paged-location/page-table preparation and KDA replay-metadata preparation, using dedicated launches per step or cache group rather than per layer. At M=4 with seven historical blocks, fusing output RMSNorm into the complete AttnRes mix reduces 10.48 to 8.20 μs (a full-mix epilogue ablation, not the hoisted-partial path). For one request's four-position window, write-slot preparation drops from 30.49 to 4.30 μs, and KDA committed-state metadata from 21.35 to 4.40 μs.

These paths remain shape-gated; outside their measured regions, the runtime keeps the reference implementations available. The reported timings are local ablations with synthetic inputs at NVFP4-compatible TP8 per-rank shapes on one GB300, using cold-L2 preparation and CUDA Graph replay with in-graph timing over two independent runs—not additive estimates of end-to-end gains.

Limitations and what comes next

In Part I, we show that optimizing K3 on TP8 is a system problem rather than a single-kernel problem. The architecture couples computation, communication, and persistent state in ways that make seemingly local choices interact: sharding LatentMoE changes the communication pattern, multi-token verification changes KDA state lifetime, and short decode segments make small-shape execution and launch overheads increasingly visible.

Moving beyond it changes several of these trade-offs:

• Expert Parallelism. EP changes both MoE routing and communication; it deserves a separate throughput and tail-latency study.
• P/D disaggregation. Long agentic prompts make prefill placement and KV transfer important, but those effects are deliberately excluded here.
• Communication beyond one NVLink domain. Cross-rack execution introduces a network hop and may select a different collective or projection strategy.
• Alternative speculative paths. These are not evaluated here. We recently published a separate blog post comparing different speculative decoding algorithms and draft configurations.

Part II will focus on expert parallel execution. Part III will take a broader end-to-end view, covering distributed scheduling, multi-tier caching, and prefill-decode coordination for long, multi-turn agentic workloads.

Acknowledgements

We thank NVIDIA DevTech for their technical guidance and kernel engineering support on GB300 systems. We appreciate the contributions of the TensorRT-LLM and FlashInfer teams to the MNNVL allreduce kernels, and of the vLLM community to the CuTe DSL kernels underlying our multicast LatentMoE tail. We are grateful to the LongCat IA team for their work on the cache subsystem and production-ready hardening, and to the Mooncake team for their performance tuning and bug fixes.

Appendix

The representative launch configuration is:

tokenspeed serve nvidia/Kimi-K3-NVFP4 \
  --served-model-name kimi-k3 \
  --attn-tp-size 8 \
  --moe-tp-size 8 \
  --mm-encoder-tp-mode data \
  --max-num-seqs 64 \
  --gpu-memory-utilization 0.92 \
  --trust-remote-code \
  --attention-backend tokenspeed_mla \
  --kda-backend cutedsl_kda \
  --moe-backend flashinfer_trtllm \
  --kv-cache-dtype fp8 \
  --speculative-algorithm EAGLE3 \
  --speculative-draft-model-path lightseekorg/kimi-k3-eagle3-mla \
  --speculative-num-steps 3 \
  --drafter-attention-backend tokenspeed_mla