NewJoin us for our inaugural conference, Forge 2026

Research

Kimi K3 on Fireworks: Fast RL at frontier scale

Train Kimi K3 for your hardest tasks. Fireworks powers faster, more stable reinforcement learning, from serverless experiments to Harvey’s specialized legal model.

TL;DR

  • Fireworks makes reinforcement learning practical on Kimi K3’s 2.8T parameters by optimizing training, rollout generation, and policy updates together for faster, more stable training.
  • In this report, we detail the custom kernels, long-context caching, and compressed weight transfers that reduce iteration overhead, plus numerical alignment and router replay to keep training stable.
  • This infrastructure powered Harvey Tenet, which achieved an all-pass rate 1.8× that of base K3 on the Legal Agent Benchmark and state-of-the-art results on LAB Contracts in its August 2026 evaluation.
  • Fireworks serverless training lets teams test their own tasks and reward functions on Kimi K3 through the Fireworks Training API, with shared training and rollout capacity and no dedicated GPUs to provision.

Specializing a frontier model means adapting its behavior to domain-specific tasks, tools, and success criteria. Reinforcement learning (RL) does this through a repeated loop: the model generates trajectories—sequences of actions and observations—an evaluator assigns rewards, and a trainer updates the policy to make successful behavior more likely.

At Kimi K3’s scale, practical RL requires efficient execution across the entire training loop. Rollout generation and training must sustain high throughput, updated policies must become available quickly, and numerical alignment between training and inference must be maintained. For experimentation, the system must also support independent runs without requiring dedicated infrastructure for each experiment.

This report describes the systems supporting RL on Kimi K3 at Fireworks, covering full-parameter training and rollout generation, training-inference consistency, and serverless LoRA training. We examine how these components reduce iteration overhead and help teams develop and evaluate specialized policies more efficiently.

The same infrastructure powers RL training for Harvey Tenet, Harvey’s Kimi K3-based model for long-horizon legal work. In Harvey’s August 2026 evaluation, Tenet achieved SOTA on LAB Contracts and second on the Legal Agent Benchmark (LAB), outperforming Fable 5 on both:

  • LAB: 19.7% all-pass vs. Fable 5’s 11.5%—a 1.8× improvement over base K3 (10.8%).
  • LAB Contracts: 11.3% all-pass vs. Fable 5’s 9.0%.

Training

Kimi K3 is a 2.8T-parameter model with a 1M-token context window. At this scale, model state and long-sequence activations strain memory capacity, while its hybrid architecture introduces distinct compute and communication bottlenecks.

ComponentKimi K3

Attention

69 Kimi Delta Attention (KDA) + 24 Multi-head Latent Attention (MLA) layers

Latent MoE

896 experts; top-16; latent dimension 3,584

Expert activation

SiTU

Routed-expert quantization

Group-32 MXFP4 weights; group-128 FP8 activations

Attention Residuals

Blockwise attention over layer outputs

Fireworks addresses these bottlenecks with specialized mixture-of-experts (MoE), Kimi Delta Attention (KDA), and Attention Residual kernels, backed by distributed training for long sequences.

MoE Mega Kernel

K3’s LatentMoE distributes FP4 experts across GPUs, requiring all-to-all dispatch and combine around expert computation. DeepEP accelerates these transfers, but exposed communication can still leave tensor cores idle. Similar communication-compute overlap has been published in Mixture-of-Kittens, COMET, and FlashMoE.

Our in-house MoE Mega Kernel adapts these communication-compute overlap techniques to K3’s routed-expert forward and backward passes:

  • Ring buffers bound dispatch staging. Reuse token-buffer slots as work progresses, limiting dispatch workspace.
  • Microbatching enables a continuous pipeline. Overlap dispatch, expert computation, and result return across microbatches; backward similarly overlaps gradient communication with computation.
  • Fuse quantization into staging. Write FP8 activations and scales directly into communication buffers, eliminating intermediate copies while preserving K3’s quantization contract.
MoE Mega Kernel FWD: overlap and buffer reuse
MoE Mega Kernel BWD: overlap and buffer reuse
Forward and backward scheduling schematics; bottom strips show measured 16K-token/rank timings. Gains include reduced staging and data movement, alongside overlap.

Against an optimized native DeepEPv2 baseline, it delivers 1.87-2.61× forward and 1.31-2.13× backward speedup across 1K-128K tokens per rank. The comparison includes Normal EP to show the additional gain beyond an already optimized communication backend.

8× B300, EP=8; optimized DeepEPv2 = 1×. Routed-expert forward/backward measured separately, including communication; not full-model throughput.

KDA: from recurrent memory to efficient GPU execution

Kimi Delta Attention (KDA) is a linear-attention mechanism that uses channel-wise gating to update a compact recurrent state, giving linear compute scaling with sequence length and fixed-size decoding state. Efficient training must parallelize chunk-local work while limiting stalls from recurrent dependencies and intermediate data movement.

We built an in-house KDA kernel optimized for Blackwell that fuses related work and reduces intermediate memory traffic. Against FLA’s KDA implementation, it delivers 1.49-1.58× forward and 1.49-1.51× backward speedup across sequence lengths of 4K-128K tokens.

B300, no CP; FLA Triton = 1×. KDA core only, including Q/K normalization and gate scan.

Attention Residuals: lower memory, faster training

Attention Residuals combine representations across model depth. At long sequence lengths, intermediate tensors for normalization, weighting, and gradient computation can consume substantial memory and bandwidth, reducing the capacity available for training.

Our kernels fuse normalization and weighted-sum operations to reduce intermediate tensor materialization and data movement during forward and backward.

Against FLA in K3’s layout, this implementation delivers 1.75-2.50× combined forward-and-backward speedup across 1K-256K local tokens. At 128K, measured peak allocated memory falls from 63.02 to 35.03 GiB—44% lower, leaving more memory available for the surrounding training workload.

Shared B300, BF16; operator forward + backward, FLA = 1×. PyTorch results above 64K are omitted due to OOM under GPU sharing.

Distributed training

At 2.8T parameters, full-parameter training requires sharding weights, gradients, and optimizer states across GPUs. Sequences approaching 1M tokens add activation-memory pressure and quadratic attention compute in MLA layers. We combine three parallelism strategies to address both model scale and sequence length:

  • Fully sharded data parallelism (FSDP). Shard parameters, gradients, and optimizer states. Prefetch overlaps upcoming weight all-gathers and gradient reduce-scatter with the current layer’s computation.
  • Context parallelism (CP). Partition long sequences across GPUs, preserving contiguous token ranges for KDA’s recurrent state propagation while balancing MLA’s causal attention workload.
  • Expert parallelism (EP). Distribute routed experts across GPUs and combine EP with FSDP to control the expert-weight working set. Our MoE Mega Kernel keeps expert GEMMs efficient while pipelining token dispatch and combination with computation, reducing exposed all-to-all latency.

Our Kernel Bricks library provides topology- and payload-aware collective implementations that reduce communication latency across training configurations. For example, AllReduce achieves 1.67-3.31× speedups over PyTorch Distributed.

Rollout

Fireworks supports Kimi K3 rollout inference for full-parameter reinforcement learning. The workload is different from ordinary chat serving: thousands of trajectories repeatedly branch from shared prompts, alternate between generation and tool use, and grow into long multi-turn histories. The rollout engine must keep those histories reusable and execute a 2.8-trillion-parameter hybrid MoE model efficiently.

This section explains the rollout-side techniques that make that possible on GB300 systems.

One model, two kinds of cached state

A KV cache stores the attention state needed to continue generation without recomputing the full prompt. K3 combines MLA, which retains token-by-token KV history, with KDA, which summarizes prior tokens in a compact recurrent state. Reusing a conversation prefix requires both histories to represent the same point in the conversation—especially when RL trajectories branch from shared prompts or resume after tool calls.

Fireworks coordinates MLA and KDA history so both cached states always represent the same prefix. This preserves reusable work as conversations grow and more trajectories run concurrently, reducing repeated prefill while each branch maintains independent, consistent state.

MLA KV history and KDA state are reused together, while each trajectory continues independently.

In our long-context, multi-turn rollout benchmarks, Fireworks achieved an approximately 99% cache hit rate, allowing nearly all prompt tokens to be reused from cache.

Speculative decoding without corrupting recurrent state

Speculative decoding uses a smaller draft model to propose tokens and the target model to verify them in parallel. Fireworks trained a speculator for Kimi K3; with both engines using a draft length of 3, our replay benchmark shows up to 12% more tokens per verification step than the SGLang DSpark baseline.

K3’s recurrent KDA state must remain consistent after rejected tokens. We integrate ReplaySSM to replay compact state-update inputs for accepted tokens from a committed checkpoint, avoiding per-candidate full-state snapshots.

ReplaySSM advances KDA state through only the accepted prefix.

We evaluate 322 replay requests across eight task types, with 16 GB300 GPUs per engine and SGLang v0.5.18. Acceptance length counts tokens advanced per verification step, including the verified target token (maximum 4 at draft length 3); results are weighted by verification steps.

  • Across tasks: general chat reaches 2.80 vs. 2.51 tokens (+12%), competitive programming 3.09 vs. 2.92 (+6%), and repo issue fixing 3.47 vs. 3.29 (+5%). Other agent tasks improve 1-3%, with skill-based agents essentially unchanged (3.24 vs. 3.25).
  • Across context lengths: gains range from 2-9%. The largest bucket, 10-20K tokens (190 requests), reaches 3.37 vs. 3.31 (+2%); 60K+ reaches 3.43 vs. 3.27 (+5%). Buckets with fewer than 30 requests have reported accuracy of approximately ±0.3 tokens.
Warm-cache replay, thinking off, temperature 0, max output 512; updated September 29, 2026. Context buckets aggregate different requests; acceptance gains are not end-to-end throughput gains.

Shortening the communication critical path

Tensor-parallel inference requires GPUs to combine partial results before computation can continue. These repeated synchronizations can limit K3’s generation speed, especially at small decode batch sizes.

Fireworks’ AllReduce library uses topology- and payload-aware implementations to reduce staging and communication latency. This lets dependent computation start sooner and shortens the communication critical path.

Fireworks' AllReduce library selects the collective path according to payload size and hardware topology.

Choosing inline or disaggregated rollout serving

The configurations described here use inline prefill: prefill and decode share the same GPU replica. This fits long-running, multi-turn rollouts because a later turn can reuse the retained history and prefill only its new suffix. Inline serving also avoids handing MLA KV and matching KDA recurrent state to a separate decode pool.

Disaggregated serving isolates decode from prefill and allows each pool to choose its own topology and capacity. It is attractive when uncached prefill volume or strict token latency justifies the separation. For K3, however, each handoff must transfer consistent MLA KV plus KDA recurrent and convolution state, find space in the destination cache, and preserve the active RL policy version.

Inline serving has its own cost: prefill competes with active decoding for GPU time. The right choice depends on prompt volume, generation length, prefix reuse, transfer bandwidth, and latency targets. Comparing the two requires the same workload and total GPU budget; the measured system described here uses inline serving.

Coordinating prefill across a wide expert group

The benchmark below uses a mixed TP4 / DP4 / EP16 layout: tensor parallelism stays within the fast intra-node fabric, four attention groups process requests independently, and expert parallelism spans the 16-GPU group. This is a measurement configuration, not a universal production placement recipe.

That creates a scheduling problem. If one attention group starts prefill, every GPU in the shared expert group must participate, even when the other groups have no useful prefill ready. Uneven RL trajectories can therefore turn many small prefills into repeated padding work and decode interruptions.

Fireworks DP prefill pairing gives sibling groups a short window to align useful prefill work while active requests continue decoding. Pending prefills run once all groups are ready or a timeout expires, trading bounded waiting for larger batches and less padding. This follows the same general bounded-delay approach published in AMD’s ATOM.

Attention groups sharing one expert group coordinate prefill admission; a timeout guarantees progress when arrivals remain asymmetric.

Throughput as histories grow

On a synthetic long-context, multi-turn rollout workload, 16 GB300 GPUs (TP4 / DP4 / EP16) served 128 concurrent trajectories over 50 turns each, averaging 2,834 output tokens/s over the complete timed run, including ramp-up and drain.

Each point shows total output throughput during a one-minute full-concurrency window, plotted against the average prompt length in that window.

Separate ablations measured 49.37% higher complete-run throughput with EAGLE3 and 65.05% with ReplaySSM. ReplaySSM’s benefit includes increased GPU cache capacity.

Putting it together

K3 rollout performance comes from treating the model as a complete system. Coordinated MLA and KDA caching preserves reusable conversation history. ReplaySSM makes speculative decoding safe for recurrent state. Fireworks’ AllReduce library shortens collective communication. Inline serving and DP prefill pairing turn asynchronous RL trajectories into efficient GPU work.

Together, these techniques let Fireworks sustain long, branching, multi-turn K3 rollouts at high throughput while preserving the state correctness required by the model.

RL: Fast iteration, stable training

Weight synchronization with XOR deltas

Synchronizing K3’s 2.8T parameters makes full-checkpoint transfers expensive. We encode each update as a bytewise XOR delta, Δ = W_new XOR W_base, and compress it with Zstandard. Unchanged bytes become zeros, reducing the compressed payload. Rollout workers decompress the delta and apply W_base XOR Δ to reconstruct the new checkpoint bit for bit.

In Fireworks’ published workload, a 1,024 GiB checkpoint produced 20.3 GiB average deltas—98% fewer bytes per delta update. Including periodic full checkpoints, total transfer volume fell 94% over 50 steps.

XOR weight sync: a small delta, the same new weights

RL stability & numerics

Determinism

Numerical drift between training and rollout can distort RL importance ratios even when both systems use identical weights. This differs from policy staleness, where rollout data comes from an older policy after the learner has updated its weights. Either effect can destabilize policy updates if it is not measured and controlled.

We enforce deterministic execution across K3’s forward pass, including major kernels and reductions, so fixed weights, inputs, and execution configuration produce matching outputs.

Critical kernels, including RMSNorm, also support batch-invariant execution. For RMSNorm, a fixed reduction order maintains consistent results across batch sizes, while preserving K3’s rounding boundaries reduces training-inference drift.

Router replay (R3)

Small numerical differences can change top-k expert selection even at identical weights, causing training and rollout log-probabilities to diverge. The Fireworks RL platform supports Router Replay (R3): the learner reuses rollout expert selections while recomputing gating weights and expert outputs from current parameters.

With these numerical-alignment measures, synchronous LoRA train-rollout divergence is as low as 0.0008 in this configuration. We report estimated KL divergence using the token-averaged Schulman k3 estimator; results depend on sequence length and execution configuration.

Harvey used the Fireworks RL platform to post-train Kimi K3 into Tenet, achieving state-of-the-art performance on LAB Contracts and second place on LAB in its August 2026 evaluation. On LAB, Tenet completed 1.8× as many held-out tasks as base K3 (source). The gains also generalized to Apex Agents (Corporate Law) and Redline Bench.

Recipe: GSPO + off-policy RL. The agent reads matter files, uses tools, and writes legal deliverables in a sandbox. Rubric-based rewards score substantive work and reward complete solutions. We optimize rank-64 LoRA adapters across K3 with GSPO, using sequence-averaged importance weights, double-sided clipping, and high-importance-ratio token masking.

Harvey off-policy RL: reward and all-pass rate rise while estimated KL divergence (Schulman k3) stays near 0.0053, including policy lag. Same run; reward/all-pass use EMA 0.99, KLD is unsmoothed.

Get started with Kimi K3 serverless training

Experiment faster with serverless Kimi K3

A new reward function, a different dataset, a better learning rate: progress in reinforcement learning comes from trying ideas and measuring what works. Fireworks serverless training makes those experiments practical on Kimi K3, with shared training and rollout capacity available through a Python API.

Your experiment connects to an already-running pool of trainers and rollout workers. There is no dedicated GPU allocation to provision, trainer job to launch, or separate rollout deployment to bring up for each run. Using an existing pool removes infrastructure provisioning from experiment setup, shortening the path from an idea to a training loop.

A multi-LoRA trainer and inference engine let different runs and policy snapshots share the same base weights. On inference, adapters load asynchronously and multi-LoRA kernels reduce repeated preparation and data movement. On training, the scheduler batches compatible requests and balances adapter locality with fairness. Optimizer state is offloaded and asynchronously staged as sessions are scheduled, preserving GPU memory for longer contexts and more active adapters.

CSGMV: efficient adapter computation

Rollout generation often accounts for most RL iteration time. High throughput requires both adapter concurrency and enough cache capacity for long, branching histories. Fireworks uses tiered adapter and cache management: reusable KV and KDA state can move to host memory, active adapters occupy GPU slots, and inactive snapshots load asynchronously. Cache reuse remains scoped to the correct adapter and policy version so ready rollouts keep moving without mixing state across experiments.

Prepare adapters. Keep rollouts moving.

A real-world example: Visual ToolBench

In a serverless Kimi K3 experiment on multi-turn VisualToolBench tasks with image tools, evaluation reward increased from 0.314 to 0.451 after 20 training updates, measured on a fixed 50-task evaluation set. See the complete runnable example in the Fireworks cookbook.

Start a Kimi K3 experiment with Fireworks serverless training. Bring your tasks, environment, and reward function, and use the shared training and rollout workflow to test your next idea.

Background and references

Terms are linked at first use. These sources provide additional background without expanding the main narrative: