0 train-infer mismatch for Open-weight MoE RL in Open-source code

How to make the trainer and sampler agree

We want the logprobs to match across forward passes, so that's where we start, and we walk backward from there through the trainer and the sampler. At each operation, we record its inputs, rounding points and addition order. For cached state, we compare what prefill writes and every update decode makes. For distributed operations, we record which values arrive and the order in which they are added.

For each operation we run both engines on the same inputs and compare output bytes. When they differ we shrink the problem until it's one row, one reduction, or one state update that still reproduces the mismatch, and then align that operation.

Four ways to align an operation

1. Call the inference kernel in the trainer's forward pass. We call SGLang's forward kernel from XoRL and write a custom backward pass for training.
2. Add one kernel to both engines. We write a new kernel, or extend an existing one, and call it in both XoRL and SGLang.
3. Make the reduction batch-invariant. We pin whatever decides the shape of a reduction, so that a bigger batch just launches more of the same work instead of regrouping the additions.
4. Write two kernels that do the same arithmetic. When the engines need different shapes or layouts, we write a kernel for each, and make them round at the same points and add in the same order.

We want each term to be one. The last factor, the stale state, is one for synchronized rollouts with freshly computed caches. Our job is to align training prefill with inference prefill, and align inference prefill with inference decode.

Anything with a cache needs one more check. Training, prefill, eager decode and CUDA-graph decode all write state, and all four have to write the same bytes at every token boundary.

When training models, we report the train-infer mismatch as K_3,t = ρ_t − 1 − log ρ_tSchulman, “Approximating KL Divergence” (2020). Introduces the K3 estimator. It is non-negative everywhere and zero only when ρ_t = 1.

The guarantee we achieve is stronger than K₃ = 0:Near one it behaves like (ρ_t − 1)²/2, so it reaches zero in FP32 while a small logprob difference remains: a difference of 1e-8 already gives exactly zero. it is that the FP32 selected-token logprob bytes from trainer scoring, inference prefill and inference decode are bitwise-identical. While SGLang has settings to trigger batch-invariant kernels, you can't get 0 train-infer mismatch with SGLang. We will have to build on top of SGLang.

The validation ladder

Ultimately, we train the full model and compare every selected-token logprob at the same weight version. We test in increasing order of difficulty:

1. Run the same program twice and compare its output bytes.
2. Change the batch shape and launch configuration and compare the same rows.
3. Run XoRL and SGLang on the same inputs.
4. Compare them during live sampling at the same weight version.
5. Train the full model and require K3 to remain zero.

We're going to showcase our methodology by solving train-infer mismatch for a series of models, starting with a dense Qwen3-8B.

Qwen3-8B: aligning a dense transformer

In Qwen3-8B, the engines can disagree on how constants are constructed, where values are rounded and how reductions are ordered. Solving mismatch here is fairly straightforward, but will allow us to produce building blocks for more complex models. We'll go through each operation in the model and explain why mismatch occurs and how we can fix it, starting with LoRA.

LoRA: two ways to match the forward

LoRA changes a linear layer from W(x) to W(x) + s B(A(x)). We support two serving modes. The merged mode adds sBA to the base weights before inference. The separate mode keeps A and B apart and adds the adapter activation to W(x). These calculations round at different points: the merged path rounds the updated weight before multiplying by x, while the separate path rounds the base and adapter projections before adding them.

A fixed adapter is faster after merging because inference avoids the extra adapter projections. In multi-LoRA training, XoRL keeps separate parameters and optimizer state for each experiment, while SGLang batches requests from several adapters in one deployment.

In merged mode, we construct W′ = W + sBA in XoRL and SGLang and run the model against identical weight bytes. We cache W′ in XoRL until A or B changes and use a custom autograd operator to map its gradient back to A and B. In separate mode, we call SGLang's LoRA forward from XoRL and implement the backward needed to train A and B. When a quantized expert needs a new LoRA forward, we add the LoRA operations to SGLang first and then call that forward from XoRL.

Code XoRL merged LoRA forward and backward ↗ XoRL LoRA backward ↗ XoRL call to SGLang’s LoRA kernels ↗

RoPE: matching the table and rotation

CPU and GPU tables diverge

RoPE rotates pairs of hidden-state values with position-dependent sine and cosine. Engines can disagree on where to construct the sine/cosine table and when to round during the rotation.

For Qwen, SGLang builds inverse frequencies on the CPU and finishes the position table on the GPU to evaluate the large outer product and sine/cosine operations in parallel. It transfers only the small inverse-frequency vector. XoRL builds the complete table on the CPU because it constructs modules and derived constants on CPU or meta tensors before FSDP assigns each rank its execution GPU.

CPU and CUDA use different implementations of sine and cosine, so the two engines can produce different values. When we decode a sequence with Qwen3-30B-A3B, the BF16 values of the tables agree through position 592. They first differ at position 593, and the first logprob mismatch appears on the next token. This is an important lesson: we need to test mismatch over sufficiently long sequences.

Build the table once on the CPU

We make both engines build the RoPE table the same way: construct the complete FP32 table on the CPU, cache it, and transfer the finished table to the GPU. The table bytes now agree before either engine applies the rotation.

Once the table bytes agree, we can use either of two rotations in both engines. Class A casts sine and cosine to BF16 and rounds throughout the rotation. Class B keeps the rotation in FP32 and casts only its outputs. We use Class B because it is faster1.23× faster. and more accurateAbout 40% lower MAE against FP64..

Code SGLang RoPE implementation ↗ XoRL RoPE kernels ↗ XoRL model-specific RoPE construction ↗

RMSNorm: fixing the reduction order

RMSNorm exposes a new problem: the compiler can change how we do a reduction.

Why is there mismatch?

RMSNorm squares every value in a row, adds the squares, takes the reciprocal square root of their mean, and uses it to rescale the row. RMSNorm does a reduction when it computes the sum-of-squares, so we need to align the order of that reduction across engines.

XoRL adds the residual before returning from the previous layer, so its next RMSNorm receives one BF16 row. SGLang passes the hidden state and residual separately and adds them inside RMSNorm. These two input signatures dispatch to different Triton kernels.

The one-input kernel loads a BF16 row and squares it in registers. The residual kernel adds two inputs in FP32, rounds their sum to BF16 and stores that row as the residual output; it then stores and reloads the FP32 squares. The one-input kernel computes 1 / sqrt, while the residual kernel uses rsqrt.They are mathematically equivalent, but they can round to different FP32 values.

The residual kernel also stores the BF16 residual row for the next layer. That store does not feed the reduction, but its presence changes the layout Triton assigns to the reduction operand from sizePerThread = [4] to [8].sizePerThread is the contiguous run each thread owns, not the total. With 32 threads per warp and four warps, 128 threads cover 1,024 values, so each thread owns eight of them under either layout. Each thread still owns eight values either way. What changes is which eight: under [4] a thread takes two runs of four, 512 apart, and under [8] it takes one contiguous run. The figure below isolates this effect: with the same tl.sum source line, adding only the required BF16 store changes the addition order, and 2,854 of 8,192 outputs change.

// one-input kernel
"tt.reduce"(%vals) <{axis = 0}> : tensor<1024xf32, #blocked>
#blocked = sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [4]
  thread 0 accumulates  x[0]+x[1]+x[2]+x[3] + x[512]+x[513]+x[514]+x[515]

// same line, with the BF16 residual store present
"tt.reduce"(%vals) <{axis = 0}> : tensor<1024xf32, #blocked>
#blocked = sizePerThread = [8], threadsPerWarp = [32], warpsPerCTA = [4]
  thread 0 accumulates  x[0]+x[1]+x[2]+x[3]+x[4]+x[5]+x[6]+x[7]

Efficient exact agreement

Most shapes run in one fused launch. When there are only a few very wide rows, we split the same calculation across three launches so that more GPU blocks can work in parallel.

Code XoRL RMSNorm ↗ SGLang RMSNorm ↗

Matrix multiplication: the same reduction along K

A GEMM computes each output by adding products along K. Training multiplies thousands of packed-token rows, while decode often multiplies one row. A library like cuBLAS will pick a different kernel for each shape, and those kernels can add the same K products in different orders. A Split-K matmul is a clear example: several thread blocks accumulate separate pieces of K, then merge their partial sums. When M changes, the library may pick a different number of splits, which changes the reduction order and gives us mismatch.

We use our Triton kernel as a reference, and because batch-invariant matmul is by now a well-known problem, we don't need to do the work of tuning it to be ultra-performant; there's already great open-source work we can rely on. We use DeepGEMM for supported shapes where it produces the same bits as our Triton kernel and runs faster than it; otherwise, we use our Triton kernel, such as for the shared-expert gate projection in Qwen3.5 MoEs.

Code XoRL fixed-K matrix multiplication ↗ SGLang fixed-K matrix multiplication ↗

SwiGLU: rounding after the multiply

SwiGLU computes SiLU(gate) * up and is another instance of the rounding rule. Training and inference can disagree on whether to round the intermediate SiLU result.

We add a one-round SwiGLU kernel to both codebases. It converts the inputs to FP32, evaluates SiLU and the multiplication, and casts the output to BF16 once. Removing the intermediate BF16 round makes the result more accurate4e-4 MAE vs 6e-4 MAE for SGLang's native implementation.. Fusing both operations into one launch and tiling across the rows and hidden dimension also makes it fasterImproves decode throughput by 2.56%, end-to-end throughput by 2.35%, and prefill throughput by 0.45%..

Code XoRL one-round SwiGLU kernel ↗ SGLang one-round SwiGLU kernel ↗

Dense attention: one streaming reduction

Dense attention is a token-mixing operator that returns a linear combination of low-dimensional projections of the token that we call values, where the weights of the linear combination are the softmax of the dot products of low-dimensional projections of the token that we call keys and queries. Computing this softmax can be numerically unstable when large values exist, but softmax is invariant to translation, so we can stabilize the softmax by subtracting the max from the inputs. Taking this max naively requires a scan over the entire vector. FlashAttention instead performs a streaming softmax: it extracts a maximum m, exponential sum l, and weighted-value sum u from each tile of inputs. Before reducing two tile-level summaries it rescales them to a shared maximum. Reducing many tile-level summaries yields the attention output u / l.

Inference engines can run this through any of several attention backends like FlashAttention or FlashInfer, while the trainer might use FlashAttention or FlexAttention or even SDPA. Each backend tiles and merges the summaries differently.

Training processes many packed queries, while decode often processes one query against a long key/value cache. An inference backend may split that cache so several blocks can work on the query at once, then merge their (m, l, u) summaries in a second reduction. Changing the number of splits changes which values are added together first. Different FlashAttention versions can also tile the first reduction differently.

For dense attention, we use FlashAttention 4 with num_splits=1 in both engines, so neither performs a second merge.

Language-model head: one reduction over the vocabulary

The language-model head takes the final hidden state and computes one logit for every token in the vocabulary. For the selected token y, its logprob is the selected logit minus the log-sum-exp over the entire vocabulary:

The head combines a fixed-K matrix multiplication with a streaming log-sum-exp. Its output depends on the matrix multiplication's accumulation order, the vocabulary-group boundaries and the tree used to merge group summaries.

After we use the fixed-K matrix multiplication in both engines, their logits agree, including the selected token's logit z_y. The remaining mismatch can come from the log-sum-exp over the vocabulary. Each vocabulary group produces a maximum and an exponential sum, (m, l), and the head merges those summaries. The trainer can group 256 logits at a time while decode groups 128. Different groups mean different local maxima, so the rescaled sums get added in a different order.

We make both engines use groups of 256 logits and the same explicit reduction tree within each group and when merging the group summaries.

Code XoRL language-model head ↗ SGLang language-model head ↗

Sampling transforms

We have achieved bitwise-identical logprobs, but inference engines don't usually sample directly from this distribution. We may want to sample from only the top-k indices of the distribution to prevent reward hacking, or apply temperature to have more or less exploration. These sampling ops change the distribution we sample from, so we need to relay the temperature parameter, etc. to the trainer and replay the same transformation.

We replay the entire program (temperature, top-k, top-p, min-p), not the mask it produced, so the trainer arrives at the sampler’s support through the same ordering and threshold rules.

Code XoRL sampling transforms ↗ SGLang sampling transforms ↗

Zero mismatch and performance

When we train Qwen3-8B on Wordle, the bitwise-identical forward lowers sampling throughput by 16.2%.

Most of that is the batch-invariant dense GEMMs.

Code XoRL dense-Qwen implementation ↗ SGLang dense-Qwen implementation ↗

Keeping the same arithmetic across parallel layouts

Achieving 0 mismatch when training dense models with dense attention is straightforward. But our goal is to train models like GLM-5.2 with 0 mismatch. Before we get into the specific complexities of different open-weight MoEs, it's important to establish that in order to train and serve models at this scale, we need to employ sharding. The trainer can shard the model differently than the sampler because we need to store activations and optimizer state in training. This changes which device owns each value, where the reductions run, and where cached state crosses a device boundary. So for each kind of parallelism we use, we need to make sure it will produce the same bits as the unsharded computation. XoRL uses CP, PP, EP, and DP (for attention and experts).

Sharding the experts

With expert parallelism, each rank evaluates the routed experts it owns and its shard of the shared expert, and ends with one BF16 partial output per token. A rank may own any number of a token's eight selected experts, so the expert kernel accumulates every routing slot in FP32 in slot order, with slots owned by other ranks contributing exact zeros, and rounds once to BF16. If the trainer and inference use different collectives, they can add the rank outputs with different addition trees.

We write an explicit tree reduction over the rank outputs. We exchange the BF16 partial outputs without reducing them, place them in a fixed logical order and add adjacent pairs. We accumulate each pair in FP64, carry an odd contributor to the next level unchanged, and repeat until one value remains, rounding once to BF16 at the end. The exchange itself does not need to match: the trainer runs an all-to-all over its expert group, while serving gathers over its tensor-parallel group with a collective that a CUDA graph can replay. Each engine only has to deliver the same partial outputs in the same logical order.

This is our general strategy for handling collective communications operations. By default, collectives that perform a reduction will pick their reduction tree based on the topology and message size, leading to mismatch. We instead use collectives that first exchange values without a reduction, and then run explicit reduction trees to ensure we can control the order of reductions. DeepEP kernels accumulate the partial sums in transit, so we can't control the order of additions. We use DeepEP kernels just for communicating the values, and do the reduction on our own outside of DeepEP. We run the explicit reduction trees in FP64 to limit precision loss; it would be great to do this for collectives that reduce in transit as well, but that would increase the communication volume so it's not typically done.

Code XoRL expert-output addition PR ↗ SGLang expert-output addition PR ↗

Sharding the sequence

One of our main goals in RL is to train on long contexts. For large models, that requires sharding over the context. XoRL implements Context Parallelism with Ulysses and RingAttention.

Dense Attention is straightforward. Ulysses shards the attention computations across ranks by heads. Each rank gets every token for a subset of heads via all-to-all. So we get the same output whether or not the trainer is sharded over the sequence.

Different open-weight MoEs tend to have different attention mechanisms, some of which totally break the invariants discussed here. We’ll get to those later on.

Sharding the batch

Data parallelism shards the batch: each rank takes its own rollouts and runs the full computation on them, and the arithmetic does not change. XoRL shards attention over context and data parallelism, and shards the expert layers over expert parallelism and FSDP. If we change the attention sharding, a token’s rows land on different ranks, and we need the expert-output reduction to produce the same bits either way.

We therefore cannot order the expert-output reduction with a rank id. Instead, we give each contributor a tag as data-parallel-major, context-parallel minor. We count the contributors in expert shards. This way, the same leaves of the reduction tree arrive whether attention is running with CP or DP. We then run the reduction over these tags, so that each token gets the same reduction tree as it would have without any sharding, however its rows were actually sharded. Ranks can also own different numbers of rows, so every row carries its own position, and we mark the padding invalid so it never enters the reduction.

Sharding the model by layers

Pipeline parallelism shards the model by layers across different devices. In each decoder layer, we return the BF16 output with the residual already added, so that we can just send this stored tensor across the pipeline boundary without doing any new compute or casting. We get the same output whether or not the trainer is sharded with pipeline parallelism.

Code XoRL head-sharding and pipeline-stage alignment PR ↗ XoRL context-parallel GDN branch ↗ XoRL logical row ownership ↗

After each exchange, the receiving rank gets the per-head BF16 values or the inter-stage BF16 row it would have computed without sharding.

Qwen3.5/3.6: Gated DeltaNet and mixture-of-experts layers

The Qwen3.5/3.6 series introduce a new challenge: recurrent layers. Gated DeltaNet is a recurrent layer. It carries an FP32 state from token to token, and any difference in this state will propagate through the rest of the sequence. For Qwen3-8B we just had to align XoRL with SGLang, but now we need to align SGLang prefill with SGLang decode. This is also our first MoE, so we'll discuss the particular challenges of aligning MoEs.

Gated DeltaNet: matching the recurrent state

Gated DeltaNetYang, Kautz, and Hatamizadeh, “Gated Delta Networks” (2024). Introduces the recurrent linear-attention layer used in this architecture. keeps a running FP32 matrix H that summarizes everything before the current token, and updates it once per token. For one head and one token, its recurrent form is:

Here g is the decay gate, β is the update strength, v is the current token’s value, and q̂ and k̂ are its L2-normalized query and key. The update H_t = (I − β_tk̂_tk̂_t^ᵀ)H̄_t + β_tk̂_tv_t^ᵀ is the delta rule.

Training and decode group the recurrence differently

Training, prefill and decode group the recurrence differently. XoRL evaluates 64 tokens at a time. SGLang also uses 64-token chunks during prefill, but decode updates H one token at a time. The 64-token calculation and the token-at-a-time calculation group the FP32 multiplications and additions differently, so they can produce different states at the chunk boundary. The next token reads that state, and the difference continues through the rest of the sequence.

The 64-row block also differs inside. SGLang's own kernels carry the running state as [V, K] with the matching transposed multiply order, and invert the 64-row triangular factor with a separate small-block pass followed by a merge. XoRL carries the state as [K, V] and uses one fused merge. Both evaluate the same sums in a different FP32 grouping, so a small tail of elements disagrees at long context.

Replay the 64-row calculation during decode

Both engines evaluate the q/k L2 normalization with the same reduction order, keep g and β in FP32, and mask future positions before exponentiating the decay. Serving also runs XoRL's [K, V] state kernel and its fused triangular merge, so both engines execute the same kernel bodies for the stages that build the chunk boundary. For XoRL's forward pass, we call SGLang's causal-convolution kernel and implement its backward.

XoRL training and SGLang prefill use one 64-row state update, and prefill stores the FP32 state after every complete chunk. To decode token p we start from that boundary and recompute rows 0 through p in the training order. It recomputes at most 64 rows before reaching the next chunk boundary.

A recurrence is causal, so earlier rows can't depend on a later token. That means we can cache their intermediate values in SGLang and recompute only what the new row touches. After prefill, a batched initializer fills the caches for the prompt's final partial chunk. Each graph replay adds one row.

After an eager decode step, we re-run the cached-row stages so that CUDA-graph decode resumes from the caches it expects. When a chunk completes, we store the next FP32 boundary in SGLang and start a new 64-row slab.

CUDA graphs require fixed tensor addresses, while requests move between scheduler slots and each live chunk grows by one row per step. We reserve 64 rows per slot and pass the current slot IDs and chunk lengths in GPU tensors. We ensure that all of our changes to SGLang still allow us to run with CUDAGraphs.

We can’t shard recurrent layers by heads, because the state runs along the sequence dimension rather than across heads. If we were to shard the sequence in the middle of a chunk, that would regroup the recurrence, changing the reduction order and therefore changing the state, which would propagate into all subsequent tokens. Instead, we only shard the sequence on each document’s 64-token chunk grid. We implement the collator to pad between documents so each cut lands on a boundary. Each rank runs the unmodified 64-row calculation on its shard and passes the FP32 boundary state to the next rank. The causal convolution reads three tokens before the cut, so the previous rank sends those projected inputs and the receiving rank computes with them.

Code XoRL GDN chunk forward and backward ↗ SGLang incremental GDN decode ↗ SGLang full-chunk GDN decode ↗ SGLang Qwen GDN selection ↗

Mixture-of-experts layers

The router selects eight experts and assigns their weights. Each rank evaluates its local routed experts and one shard of the shared expert. We collect the eight BF16 partial outputs and reduce them in a fixed order.

The router: matching expert ids and weights

The router projects each token to one score per expert with the fixed-K matrix multiplication, applies FP32 softmax, selects the top eight probabilities and renormalizes them. The expert choice is discrete: a last-bit difference in the gate projection can swap the eighth and ninth experts. Even when the expert ids agree, adding the eight selected probabilities in a different order changes the BF16 routing weights.

We run the fixed-K projection in both engines. We then add the eight selected FP32 probabilities from left to right, divide each probability by that sum and cast the routing weights to BF16. Both engines return the selected experts sorted by probability, so every later sum that walks the eight slots visits them in the same order.

Code XoRL top-eight routing ↗ XoRL router matrix multiplication ↗

The experts: keeping the route weight inside the accumulator

The trainer can round the down-projection result to BF16 before multiplying by the route weight, while inference multiplies that weight inside the FP32 down-projection accumulator and rounds once at the end.

For XoRL's forward pass, we run SGLang's existing Triton expert kernels and implement a custom backward operator for training. We pass XoRL's weights as transpose views. The route weight remains inside the FP32 accumulator, and the result rounds once at the end.

The shared expert uses the fixed-K projections and SwiGLU from the dense case. Each rank adds its shared-expert shard to its routed-expert output before the eight ranks combine their results.

Code XoRL routed and shared experts ↗

Both engines now hold the same eight BF16 partial outputs for a token, one per rank, and add them with the fixed expert-output reduction; eight contributors take three levels.

Zero mismatch and performance

The bitwise-identical forward lowers sampling throughput by 39.4% and adds 57.9% to sampling wall time. Removing exact GDN decode recovers about half of that added wall time, 49.5%, which makes GDN the largest single cost we can name and leaves the other half spread across the rest of the aligned forward.

Wall-clock times are not directly comparable because the policies learn different behavior. The bitwise-identical run starts solving games sooner, so it generates a different number of tokens. Sampling is also only part of an RL step, so this is not the overhead a training run pays; we measure that end to end in the conclusion.

GLM-5.2: sparse attention and FP8 experts

GLM-5.2 introduces new complications, including a sparse attention mechanism and FP8 native weights. To train GLM-5.2 properly we need to compose multiple degrees of parallelism: Context Parallelism, Expert Parallelism, Tensor Parallelism, Fully Sharded Data Parallelism and Pipeline Parallelism. XoRL supports full-parameter, LoRA, and QLoRA training for GLM-5.2.

Sparse attention

Some GLM blocks choose earlier positions to attend to, and later blocks reuse those choices. If the selector picks different positions, this mismatch cascades through the sequence. We first match the selector, and then the attention.

Selecting sparse-attention positions

The selector returns the positions that later sparse-attention blocks will read. XoRL and SGLang build its inputs differently. XoRL projects the key and head gate separately and scores BF16 query and key values. SGLang projects the key and head gate together and converts the post-RoPE query and key to FP8, each with its own scale. These differences can change which positions are selected.

We make XoRL follow inference's selector calculation. We project the key and head gate together in BF16 and apply the BF16 round after RoPE. We then cast the query and key separately to FP8 with their own scales. Both engines score the complete history, choose the lower position when scores tie and return positions in ascending order for later sparse-attention blocks.

The prefill kernel writes the prompt keys, and decode appends each generated key. Prefill and decode can place normalization, RoPE and BF16 rounding differently. To reproduce the serving cache in XoRL, we mark the prefill boundary and rebuild each row with the corresponding SGLang prefill or decode calculation.

The engines can store the history in different ways as long as every logical position contains identical values. XoRL stores the history in a contiguous tensor and SGLang stores it in 64-token pages.

We call SGLang's FlashMLA forward from XoRL and implement its backward in TileLang, passing the selected rows straight from XoRL's contiguous tensor instead of routing them through SGLang's page table.

Code XoRL selector inputs ↗ XoRL selected positions ↗ XoRL FlashMLA forward and backward ↗ SGLang sparse selector ↗ SGLang paged sparse selector ↗

Mixture-of-experts layers

In GLM, the router adds a correction bias, the experts run from FP8 weights, and many ranks (we use a minimum EP=16) contribute to each MoE output. We use one FP8 expert forward for adapter training and full-weight training. The two modes differ only in how we update the expert weights after the backward pass.

Choosing GLM's eight experts

GLM adds a correction bias only when choosing experts. Their weights still come from the sigmoid scores before the bias. The trainer can load this bias as BF16 while inference keeps it in FP32, which changes the selected experts. The trainer can also round the normalized routing weights to BF16 before the expert call, while inference keeps them in FP32 and applies the factor 2.5 inside the expert kernel.

To remove these differences, we run the same fixed-K gate projection in both engines and call SGLang's router directly from XoRL. It adds the FP32 correction bias before choosing the top eight experts. It then gathers the corresponding sigmoid scores from before the bias, normalizes them in FP32, and passes them to the expert kernel. The expert kernel applies the factor 2.5.

Code XoRL GLM router ↗ SGLang grouped top-k ↗

LoRA path: running FP8 experts with adapters

The trainer can convert the expert weights to BF16 and compute LoRA separately, while inference reads the checkpoint’s FP8 weights and block scales and applies LoRA inside its expert kernels. The two forwards then round in different places.

We extend SGLang's FP8 expert kernel to apply LoRA after the gate/up projection and again after the down projection. The base weights stay in FP8, and the routing-weight multiply and final BF16 round stay exactly where inference already does them. We call this kernel from XoRL's forward pass and implement a custom backward for the FP32 LoRA matrices.

Code XoRL FP8 experts and backward ↗ SGLang LoRA MoE ↗ SGLang FP8 experts ↗

Full-weight path: training the FP8 expert weights

To train the FP8 experts, the trainer keeps FP32 master weights and quantizes them for its forward pass. If inference re-quantizes on its own side, the two engines have to agree on every rounding point inside the quantizer.

Instead, we quantize XoRL's FP32 master weights into SGLang's block layout once after each optimizer step. Until the next step, every forward pass reads those cached FP8 codes and block scales. We compare both the FP8 codes and the block scales against what we last sent, and send only what changed. A step moves each block's scale very little, so about one percent of the entries change.

After the LoRA or full-weight forward, we place the sixteen partial outputs in logical rank order and apply the fixed expert-output reduction; sixteen contributors take four levels.

The distributed language-model head

Inference can split GLM's 154,880-token vocabulary across sixteen ranks, while the trainer projects the complete vocabulary at once. The differently shaped matrix multiplications can add their products in different orders and produce different logits.

We split XoRL's head across the same sixteen vocabulary shards and send each hidden row and target token id to the ranks that hold those shards. Each rank applies the fixed-K projection, and we gather the logits in rank order before applying the same streaming softmax.

Code XoRL distributed language-model head ↗ SGLang distributed language-model head ↗ SGLang logprob calculation ↗

Zero mismatch and performance

Sampling with the bitwise-identical forward lowers throughput from 147.6 to 107.2 tok/s, a reduction of 27.4%.

Most of the cost is in the expert path. If we turn off the matched router, the FP8 expert forward and the rank-ordered expert-output reduction, throughput goes back up to 125.9 tok/s, so nearly half the alignment cost is in the expert path alone. The exact sparse-attention selector and top-k account for about two points. The cost of batch-invariant addmm, bmm, log_softmax, mean and mm is within noise, because these don't change the expert projections, only the dense matmuls and attention projections, which are a small share of the FLOPS.

DeepSeek-V4: four streams and compressed history

DeepSeek-V4 introduces multiple residual streams, compressed attention history, and MXFP4 expert weights, which we run with LoRA without unpacking.

Four-stream residual mixing

Fusion changes the residual calculation

There are four BF16 residual streams in each block. The model's mHC mixer combines these streams into one BF16 row right before each attention / expert layer, and then adds the layer's output back into the remixed streams. This mixer computes per-token FP32 weights and a 4x4 FP32 mixing matrix. XoRL and SGLang can disagree on how to fuse the mixer, or where to round and add inside the mixer, so there's mismatch.

Share the serving forward

In XoRL's forward pass, we call SGLang's pre-mix and post-mix kernels. Our custom backward returns gradients to all four residual streams. We stop-grad the mixer coefficients during training, so no gradient flows into the mixer's weights.

Code XoRL mHC forward and backward ↗ SGLang fused mHC kernels ↗

Compressed sparse attention

Every layer uses 128-token sliding window attention. C0 layers only attend to this window. C4 layers also attend to overlapping summaries. These are summaries of eight tokens, that are built every four tokens. C128 layers have one summary for each non-overlapping tile of 128 tokens. Each query attends to its summaries and the sliding window in one softmax.

Compressing the prefix

To build a compressed row, the model projects the hidden states in BF16, pools them with FP32 weights, normalizes them, applies RoPE and stores FP8. We use the same batch-invariant matmul for the projection and the same RoPE kernel as the dense layers. The trainer sees the complete sequence, while inference builds this state during prefill and updates it during decode. So the two can create a compressed row at different token boundaries. We create C4 and C128 rows only when their respective token groups are completed, so that we can match the serving boundaries to ensure that later queries read the correct compressed rows.

In XoRL, we use SGLang's summary calculation and FP8 cast. SGLang's sparse prefill kernel isn't byte-stable, so we replay the prompt by looping the decode kernel once per causal query row. We then advance the paged caches one token at a time with the same kernel. Between tokens, we carry the recent-token cache, the FP32 compressor state and the C4/C128 rows, adding a compressed row at the same boundary as serving.

Attending to compressed rows

We build FlashMLA's input in XoRL from the logical positions exposed by SGLang's paged cache, placing the compressed prefix before the most recent 128 tokens. We call SGLang's FlashMLA forward and implement the backward as an FP32 reference of the same attention. When replaying one token, we call the serving decode entry point directly over the paged cache.

Code XoRL DeepSeek-V4 compressor ↗ XoRL C4 indexer ↗ XoRL FP8 attention forward and backward ↗ SGLang DeepSeek-V4 compressor ↗ SGLang C4 indexer and page mapping ↗

Hash and learned expert routing

The trainer can form the router logits with its own matrix multiplication and compute both routing schemes in PyTorch, while inference can form a BF16 gate projection with a persistent kernel, widen it to FP32, and pass it to fused hash or learned-routing kernels. Routing is discrete, so in the learned layers a one-ULP change in a gate logit can change a sqrt-softplus weight, which slots get selected, and in what order. Naturally, different orders lead to mismatch.

To align the two paths, we form the gate logits in XoRL with a batch-invariant matmul. For the first three blocks, a frozen table maps each token id to its six experts, so the logits only set the weights; later blocks use the learned router with the correction bias. We preserve the returned weights, expert ids and slot order through the expert forward.

Code XoRL DeepSeek-V4 routing ↗ XoRL hash and correction-bias selection ↗ SGLang token-to-expert routing ↗ SGLang learned routing ↗

Running MXFP4 experts with LoRA

MXFP4 has a constraint that FP8 doesn't: Marlin can return different bits at different row counts, so both engines run the packed GEMMs one row at a time, the same unpadded launch serving uses during decode. The trainer's backward unpacks each expert weight to BF16 and evaluates LoRA separately, while inference adds LoRA inside its Marlin runner over the packed weights.

As with the FP8 experts, we extend SGLang's runner to apply BF16 LoRA after the packed gate/up and down projections, here without unpacking the base weights at all, and we call it from XoRL's forward pass. We use SGLang's row decomposition, keep its clamps, SwiGLU calculation and final BF16 round in place, and multiply the packed base output and the LoRA delta by the routing weight before adding them.

Code XoRL packed MXFP4 weights and forward ↗ SGLang DeepSeek-V4 LoRA mapping ↗ SGLang Marlin LoRA runner ↗

Each rank receives a different number of rows, and the exchange delivers all eight rank outputs already in rank order, so we just add them with the fixed expert-output reduction.

Code XoRL variable-row exchange and reduction ↗ SGLang expert-output reduction ↗

The final mix and language-model head

After the final decoder block, the mixer produces one final BF16 hidden row and a final RMSNorm normalizes it. The LM head projects that row against the vocabulary weights, gathers the TP shards in rank order, and does log-softmax. Any of these three operations can lead to mismatch if the engines do them differently.

We reuse the fixed-K projection and ordered gathering from the earlier language-model heads. SGLang's LoRA kernels add the rank-1 delta on top of a plain base matmul. We gather the BF16 logits in rank order, call SGLang's BF16 log-softmax from XoRL, and implement the XoRL backward with an FP32 log-softmax.

Code XoRL DeepSeek-V4 language-model head ↗ SGLang DeepSeek-V4 model ↗ SGLang BF16 log-softmax ↗

Zero mismatch and performance

We train DeepSeek-V4 with LoRA on the same Wordle task. For sampling we use TP8/DP8/EP8 on one node, and the bitwise-identical forward lowers decode throughput by 12.7%.

The expert path is again where the money goes: removing the aligned expert geometry and combine recovers 43.8% of the added wall time. This is the cheapest of the three MoEs to make bitwise-identical.

Code XoRL DeepSeek-V4 PR ↗ SGLang DeepSeek-V4 PR ↗

The Price of Progress

We can train Qwen 3.x, GLM-5.y, and DeepSeek-V4 MoEs with 0 train-infer mismatch, and have an easily deployable method to train any model we want with 0 mismatch. But why should we do this?Reader beware, for we have left the orderly kingdom of verifiable kernels and entered the nebulous realm of “empirical reinforcement learning”. What does it cost us? What about async RL? Can't we just handle this with different objectives? Are there really no better alternatives? We're going to walk through a range of ablations on our Wordle task with Qwen3.6-35B-A3B and try to answer some of these questions as best we can. We'll compare to the managed RL APIs Tinker and River, to validate that the levels of mismatch we observe in our baselines, and the final performance of our trained models, are competitive with closed-source RL engines. Of course, we don't have access to the internal details of these engines, so we can't make any definitive claims about their performance.

Why should we eliminate mismatch?

To understand why we should eliminate mismatch, let's take a look at what happened with our runs on Tinker. The run we show in Figure 1 is actually the best of 3 runs; 2 of them diverged on step 62.

Nothing about step 62 is special. There was no system error in Tinker's infra, as far as we can tell. The sampler's mean selected-token logprob stays near −0.20 in all three runs, and the sampled rollouts look reasonable. But in the 2 runs that collapse, the trainer's logprobs are much lower at step 62, and this leads to catastrophic mismatch.

Because we train with importance sampling, the ratios have unbounded influence on the gradient. In Run 1, a single </ with an importance ratio of 19,276.5 carries 65.7% of the squared coefficient mass.Each token's update coefficient is c = ρÂ, its importance ratio times its GRPO advantage. Squared coefficient mass is that token's c² as a share of the batch total. It's not as extreme in Run 3, but the top four tokens do carry 72% of the mass. Not so in Run 2, where the largest single token at step 62 just has 0.39% mass.

Run 1 loses the format and starts emitting invalid actions on most turns. Run 3 keeps playing legal Wordle moves, it's just now making very bad guesses. Run 2 survives, but it doesn't perform particularly well.

Of course, unbounded importance sampling is a bit of a sandbagged baseline, right? Surely we can just fix this in our objective.

What objective mitigates mismatch?

We briefly discussed different objectives and how they handle the importance sampling ratio at the start of this article. Click through to read the definitions of IS, PPO and CISPO again.

Our IS in XoRL trains better than IS in Tinker, but still lags behind the bitwise-identical run. We use the same formulation as Tinker. Late in training, IS generates 44% more tokens per turn than the bitwise-identical run. Most of these tokens aren't legitimate reasoning, they're just the model repeatedly checking formatting, word count, and readiness before guessing. When a large ratio leads to a large gradient, as we discussed above, it stays in the optimizer state and can impact many future steps. We find that the vast majority of the mass comes from tokens with IS ratios greater than one (that is, the trainer assigns higher likelihood than the sampler).

Our PPO run is much worse. We use the same clipping thresholds as Tinker. It drops gradients for too many important tokens, and hardly learns anything. While the measured mismatch is much lower than with IS, this is mostly because the policy is just producing degenerate rollouts.

CISPO is supposed to be the objective that fixes these problems. It doesn't allow unbounded IS ratios, but it also doesn't drop gradients, it just bounds their contribution. We use the default values. We find that CISPO is an effective way to mitigate mismatch. Our CISPO run in XoRL lands at 72.5% held-out solve rate, well above unclipped IS at 63.9% but still short of the bitwise-identical run's 77.4%. CISPO also has a dramatic effect on Tinker, improving the performance by 25% and nearly matching the 0-mismatch run. If we compare the reported K3 values, we see that CISPO in Tinker has much lower K3 than the IS baseline in Tinker.

None of the objectives we tried can perfectly match the performance of the 0-mismatch run, although CISPO comes fairly close, but that doesn't rule out the possibility that a carefully tuned objective can near-totally correct for the mismatch.

Total Router Recall

We do most of our experiments on MoEs because train-infer mismatch can be more catastrophic for MoEs than dense models. For most of the network, mismatch accumulates somewhat smoothly. A few bits here, a few bits there. But the MoE layers make a discrete selection. Once the router top-k flips to select different experts in the trainer than in the sampler, the mismatch is much larger. We conjecture this may be the failure mode that killed our Tinker runs.

One idea proposed in prior work is Router Replay: record the selected expert indices during inference and pass them to the training engine.Group Sequence Policy Optimization describes Qwen's earlier Routing Replay strategy, and Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers proposes Rollout Routing Replay. This approach appears in GSPO and R3. We don't see a benefit to replaying only the router indices (compare the two River runs above). Router Replay does marginally reduce mismatch, but let's think through what would happen if we relied on it to salvage a rollout with very high mismatch. Suppose that our inference engine assigns a high score to expert 0, and our trainer assigns a low score to expert 0, but we replay expert 0 nonetheless. Now there is a severe mismatch: the expert index agrees, but the routing weight does not, because the trainer still recomputes that weight from its own logits. Replaying only the expert indices does nothing when there is little mismatch, is maybe even harmful when there is extreme mismatch, and would only help when the mismatch is in the "moderate" regime; we expect the aggregate effect to nearly cancel out.

Cursor modifies Router Replay when training Composer 2. They filter out replayed experts whose gating scores are implausibly low under the trainer’s own router and replace them with candidates from the trainer’s top-k. They report that this reduces the p99 numerical mismatch left by basic replay. This is a reasonable strategy; let's now consider a third option.

We extend Router Replay to record and transmit the routing weights too.The R3 authors specifically only transmit the indices so that the router can remain trainable, which is an admirable goal. SGLang returns two payloads per request: the selected expert ids as ints, and the normalized routing weights as FP32, one per token per MoE layer per selected expert. XoRL uses the transmitted weights directly instead of recomputing them.The payload grows with tokens × MoE layers × selected experts, so we move it by handle through Mooncake or the filesystem rather than inline in the rollout record.

Our extended Router Replay, which we call Total Router Recall, works quite well. It gets almost the same held-out accuracy as the bitwise-identical run, and its K3 stays at 3e-4 for the whole run instead of drifting up like the baseline's does.

But Total Router Recall does come at a cost. We can no properly longer train the router. We can technically still get a gradient by using the stop-gradient trick. If we pass forward w − sg[w] + w*, the forward value is still the sampler's weight w*, and the backward pass is ∂w/∂θ. But subsequent layers run on w*, so ∂L/∂w is evaluated with the replayed router experts while ∂w/∂θ is evaluated with the trainer's logits. The two factors come from different router selections, so we wind up optimizing a linearization of the loss around the sampler's weight, where the gradient error term is first order in the gap between the two routers. The trainer is also coupled to the sampler's internals now: it needs the sampler's expert ids and its normalized weights, in the sampler's layout, for every MoE layer of every token. This is a huge amount of data to transmit from the sampler to the trainer for every request, and Mooncake makes it somewhat manageable but it's still nontrivial overhead to capture this metadata and send it across devices.

We run our experiments in XoRL through a Tinker-compatible interface, so we start training and sampling servers and just send them requests to sample and train from a lightweight client. We certainly don't want to send the massive router metadata payload from the sampler to the client back to the trainer for every request. But that means that if we want to modify the tokens / logprobs at all in the client, such as truncating the tokens past the end of the first answer, then we would need to somehow mirror that on the routing payload.

You may think by now that this all seems a little kludgey, and it is. Tinker doesn't support even basic Router Replay. Still, Total Router Recall is the best alternative to our bitwise-identical system that we've come up with so far. If you are looking for an easy drop-in way to mitigate mismatch and don't care about training the router, this is a good option. Any alternative to 0-k3 is attractive because it may not have to pay the cost that the 0-mismatch solution has to pay.

What does 0-mismatch RL cost?

We've already covered the overhead of enforcing 0 mismatch in sampling, and it is significant. But at the step-level, the entire cost of RL with 0 mismatch is just around 20%.

20% is a significant overhead to pay, but recall that this tax is actually steepest for Qwen, and that for Qwen the lion's share of the tax is in aligning prefill and decode within SGLang with exact GDN decode.

Doubtless there are many future optimizations that can reduce the step-level overhead, and we are optimistic that soon it will not cost much to do 0-mismatch RL. And we do believe that it's worth it, for the peace of mind to know that our runs aren't plateauing due to mismatch, but there is another cost to enforcing 0 mismatch.

As much as we have worked to achieve zero train-infer mismatch, async RL is still very compelling for production-scale runs. However, it costs us a sizeable mismatch term with stale weights and the stale cache.

Even when running the 0-mismatch stack with CISPO, which was the best objective we found for mitigating mismatch, asynchrony craters performance. The mismatch from stale weights and a stale cache spikes quickly and the model plateaus early. But incorporating Total Router Recall recovers a significant chunk of performance, while reaping the speedups of Async RL. While the reported K3 still spikes early in our TRR + CISPO + Async run, it spends the rest of the run trending downward.

The goal of streaming / async RL is to have all trainer GPUs and inference GPUs be fully utilized at all times. One way to partially overlap training and inference that doesn't induce any mismatch is Streaming RL.

We have SGLang return every completed trajectory immediately, and the client assembles completed groups and passes them to XoRL. XoRL does a forward-backward pass on a completed group, because we only need to normalize the advantages across the group, and only does one optimizer step at the very end. The gradient update is therefore exactly the same as if we had sent all the datums in one big batch. Because the waiting distribution is geometric, we get pretty good overlap between the trainer and sampler without needing any asynchrony.

In our Wordle runs, we can save 46% of serial wall time with Streaming RL. However, the sampler utilization goes down as more trajectories finish, so this will by no means be as good as async RL.

Async RL fills the bubbles left by heterogeneous sequence lengths with new requests from a stale policy, in exchange for mismatch. But we can also fill these bubbles with requests from different policies, by training with multi-LoRA. If we want to run multiple jobs at once, then when one job is down to the last few stragglers, another job can fill the sampling queue with a fresh batch. Every request is still on-policy for its own adapter, so this overlap costs no mismatch. We only get it when the jobs share a base model. Of course, this doesn't speed up training for a single job, but if we want to spin up a XoRL training service and an SGLang sampling service and run multiple experiments, multi-LoRA can effectively saturate GPU utilization.

Parting Thoughts

We believe that training open-weight MoEs with 0 mismatch is a very achievable goal for any RL stack with very real benefits. All of these techniques are used in XoRL today, but can just as easily be ported to other stacks. While XoRL leverages SGLang for inference today, we see no reason why the same methodology could not be used for other inference engines.

Acknowledgements: I wrote this article in the usual academic "we" phrasing, but allow me to break character here to thank some folks. I'd like to thank friends from Core Automation, Recursive, Ricursive, and River AI for providing feedback on earlier versions of this article. Huge thanks to Qingyang Wu and Zhongzhu Zhou for their contributions to XoRL and their encouragement on my pursuit for 0-mismatch RL. And of course none of this work would be possible without GPU hours and tokens from TogetherAI.

Appendix A: Detailed operation map