Kimi K3 on Fireworks: Fast RL at frontier scale

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.

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.

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.

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.

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.

Open in lightbox

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.

Open in lightbox

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.

Open in lightbox

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.

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.

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.

Open in lightbox

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 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.

Open in lightbox

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.

Open in lightbox

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.

Open in lightbox

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.

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.

Open in lightbox

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.

Open in lightbox

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

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.

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.

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.

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.

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.

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.

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.

Open in lightbox

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.

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

• Kimi K3 architecture, KDA, MLA, Attention Residuals, LatentMoE, SiTU, and quantization.
• MoE routing and communication: Switch Transformers, DeepEP, Mixture-of-Kittens, COMET, and FlashMoE.
• Distributed training: FSDP and unified sequence/context parallelism.
• Rollout serving: PagedAttention, speculative decoding, ReplaySSM, and DistServe.
• Policy synchronization and numerical alignment: XOR checkpoint deltas, Fireworks’ MoE numerics, and Router Replay.
• LLM reinforcement learning and optimization: a learning-from-rewards survey, InstructGPT, DeepSeek-R1, PPO, GSPO, and the Schulman k3 KL estimator.
• Serverless multi-LoRA systems: LoRA, Punica, S-LoRA, mLoRA, and tLoRA.
• Evaluation sources: VisualToolBench and the published customer case study.