Optimizing Kimi K3 on GB300 is a systems problem. Its KDA, Gated MLA, Stable LatentMoE, and AttnRes components interact with communication and persistent state, while long, multi-turn agentic workloads expose overheads that are easy to miss in an isolated decode benchmark.

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. The end-to-end comparison tracks aggregate throughput and per-user generation speed across concurrency levels; the deep dives then isolate the computation, communication, and state-management choices behind that result.

Key takeaways

• EAGLE3 improves the aggregate-throughput versus per-user-generation-speed curve across the evaluated concurrency range.
• Communication-aware sharding of the two LatentMoE projections reclaims approximately 7.7 GiB of parameter storage per GPU while reducing local projection and post-expert tail latency across most measured shapes.
• Replay-based KDA state recovery reduces verify-time workspace from 16.7 GiB to 0.4 GiB per GPU at max_num_seqs = 64.

Kimi K3 optimization series: Part I — TP8 · Part II — expert parallelism (coming soon) · Part III — distributed serving (coming soon)

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

Each point on the serving curve represents one measured concurrency level. Moving right means faster generation for each user; moving up means higher aggregate throughput per GPU. We compare EAGLE3 with non-speculative execution under the same frozen workload and serving configuration.

Benchmark contract

Component	Configuration
GPU	NVIDIA GB300 GPU
Topology	8 GPUs across two 4-GPU compute trays, within one NVL72 NVLink domain
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	CuTe DSL prefill; fused Triton decode/verify
MoE backend	FlashInfer TensorRT-LLM
KV-cache dtype	FP8
Workload	Frozen multi-turn SWE-Smith traces
Context growth	~50K input tokens initially; ~800 additional tokens per turn; up to ~68K tokens
Generation	10–15 turns; up to 500 output tokens per turn
Concurrency	1, 2, 4, 8, and 16
EAGLE3 verification	3 draft steps; 4-position target window

The eight TP ranks span two four-GPU compute trays while remaining inside one NVL72 NVLink domain; this experiment therefore does not include a cross-rack network hop.

SWE-Smith traces approximate 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.

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 time per output token, or TPOT). TPS/user does not include time to the first token. Together, these metrics show how serving more concurrent users changes both system throughput and the generation speed experienced by each user.

Comparison scope

All points use the same trace set, target model, and serving configuration. The intended comparison is EAGLE3 versus non-speculative execution under these matched conditions; results from other systems require matching workload and metric definitions before they are comparable.

Figure 1 plots TPS/user on the x-axis and TPM/GPU on the y-axis. The labels identify the measured concurrency (cc) at each operating point.

Figure 1. Aggregate throughput versus per-user generation speed on multi-turn SWE-Smith traces.

EAGLE3 shifts the measured curve outward: its benefit is visible both at low concurrency and as the workload moves toward higher aggregate throughput. At the reference generation speed of 50 TPS/user, linear interpolation between the measured EAGLE3 points at concurrency 8 and 16 gives approximately 459K TPM/GPU. This is an interpolated point on the curve, not a separate measurement.

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 NVLink SHARP (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-protocol 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.

Figure 2 shows the sharded layout and the assembly path across the eight TP ranks.

Figure 2. Routed-down projection sharding and latent assembly.

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. In Figure 3, the jump between M = 1280 and M = 1281 marks the switch from mailbox assembly to GEMM plus all-gather. Large-M points test projection scalability beyond this workload's decode and verification region.

Figure 3. Routed-down projection latency for replicated and sharded paths.

Routed-up projection: sharding the post-expert tail

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 slice of the prefix tensor 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.

On GB300 systems, this path builds on the Multi-Node NVLink (MNNVL) all-reduce design from TensorRT-LLM, as carried through FlashInfer. Its one- and two-shot Lamport protocols use NVLink multicast memory for small TP reductions. TokenSpeed adds K3-specific epilogues so adjacent post-reduction work can remain fused. The lane-based and NVLink multicast multimem backends implement the same block-injection pattern. For small M, the 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 tensor. This integration was developed with technical guidance and optimization support from NVIDIA DevTech.

Figure 4 summarizes how projection sharding turns the shared all-reduce into both a reduction and an assembly step.

Figure 4. Routed-up projection sharding and post-expert 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.

Starting from the point where routed and shared experts have produced their local partials, the complete tail drops from 32.34 to 18.28 μs at M = 4 (1.77× after subtracting empty-event overhead; 1.69× using raw event intervals), and from 915.77 to 656.22 μs at M = 8192 (1.40×). Figure 5 also shows 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 separate reductions and a replicated up projection, not isolated finalize-fusion gains or end-to-end model speedups.

Figure 5. Routed-up projection and fused-tail latency versus the replicated control.

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.0×) with no observed latency increase when the fused verify and replay kernels are used.

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

Memory implications

Under the cache geometry used here, the 7.7 GiB parameter reduction from projection sharding and the 16.4 GiB verify-workspace reduction from Replay SSM are byte-equivalent to roughly 0.6 million and 1.2 million logical target-plus-draft FP8 cache tokens, respectively. These translations use 14.1 KiB/GPU per logical token; they are not peak allocator measurements or guarantees that the full theoretical capacity is simultaneously allocatable.

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, CuTe DSL skinny GEMM, FlashInfer TGV, and torch.mm. Figure 6 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.

Figure 6. Shape-aware projection routing relative to torch.mm.

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.

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 and page-table preparation, plus KDA replay-metadata preparation. Dedicated launches run per step or cache group rather than per layer.

Local ablation summary

Path	Measured shape	Baseline → optimized	Comparison boundary
MLA query and cache write	M = 4	6.43 → 4.39 μs	Combined query/cache-write stage
AttnRes output RMSNorm	M = 4, 7 historical blocks	10.48 → 8.20 μs	Full-mix epilogue, not the hoisted-partial path
Write-slot preparation	1 request, 4-position window	30.49 → 4.30 μs	Complete metadata chain
KDA committed-state metadata	1 request, 4-position window	21.35 → 4.40 μs	Complete metadata chain

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 GPU, 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

Part I shows why optimizing K3 on TP8 requires system-level co-design. Under the controlled serving workload, EAGLE3 improves the aggregate-throughput versus per-user-generation-speed curve. Inside the target path, LatentMoE projection sharding reclaims approximately 7.7 GiB of parameter storage per GPU, while Replay SSM reduces verify workspace from 16.7 GiB to 0.4 GiB per GPU. The local ablations then show where communication-aware assembly and shape-specific fusion reduce the remaining operator and metadata overhead. These results describe different boundaries and should not be added together as one end-to-end speedup.

Several boundaries remain:

• 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 favor 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, covering wide expert parallelism (WideEP) and prefill–decode disaggregation. Part III steps back to an end-to-end view, with enterprise deployment case studies covering distributed scheduling, multi-tier caching, and other lessons from running long, multi-turn agentic workloads in production.

Kimi K3 optimization series: Part I — TP8 · Part II — expert parallelism (coming soon) · Part III — distributed serving (coming soon)

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: serving launch configuration

The representative serving launch configuration is below. Reproducing the end-to-end result also requires the frozen SWE-Smith traces and the client-side concurrency schedule described in the benchmark contract; those artifacts are not included in this post.

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