Reliable RL Scaling Requires Accounting for Prefill-Decode Kernel Mismatch

Yifan Zhang

Abstract

A language-model rollout is on-policy for a declared target only when the post-sampling-transform behavior distribution equals that target distribution on the histories reached by rollout. Equal checkpoints do not ensure this condition. Prompt prefill, cached decoding, teacher-forced scoring, parallel scans, recurrent state updates, cache precision, quantization, and reduction order can all change token probabilities. For linear attention and state-space models, a small state discrepancy can propagate through all later updates. Moreover, forward probability agreement at one parameter value is not enough to claim optimization of a recurrent deployment policy: the learner must also use the corresponding score function, or a backward-equivalent implementation.

The report studies four responses. First, treat execution mismatch as off-policy data, store the actor-emitted probability, and use an explicit importance ratio, distinguishing the exact local one-step-deviation surrogate from full trajectory correction and stating the required support contract. Second, make one chunkwise or recurrent execution rule canonical across rollout and differentiable learner replay; TTT, Titans, and RWKV-7 illustrate relevant execution semantics, but architecture names alone do not establish parity. Third, use higher precision for recurrent states and sensitive updates as a numerical mitigation, not as an equality proof. Fourth, use a fast path only as a proposal and let a recurrent target perform exact modified rejection sampling; the standard tokenwise correction is exact only when the proposal exposes the actual conditional probabilities of the drafted sequence. The report recommends separating an exact recurrent-target mode from a scalable parallel-target surrogate, with explicit audits for forward probabilities, score functions, support, and state precision.

Read the Report Project Page

Same Checkpoint, Different Policy

Figure 1: The common actor-learner split. For response tokens after the first generated token, rollout probabilities normally come from recurrent decode, while learner probabilities normally come from a parallel teacher-forced pass.

A standard rollout processes the prompt with a parallel prefill kernel, then produces response tokens by one-token recurrent decoding over the cache. The learner typically does something else: it concatenates prompt and response and scores them in one teacher-forced parallel causal pass. After the first generated token, the behavior probability and the learner probability usually come from different kernel families, at the same parameter vector $\theta$.

\text{effective policy} = \text{parameters} + \text{execution path}

The concrete implementation bug is to sample with $\mathsf{K}_{\mathrm{roll}}$ but later reconstruct the "old" probability with $\mathsf{K}_{\mathrm{learn}}$. The correct recordwise ratio is

\rho_t = \frac{\pi_{\theta_n}^{\mathsf{K}_{\mathrm{learn}}}(a_t \mid h_t)}{\pi_{\theta_v}^{\mathsf{K}_{\mathrm{roll}}}(a_t \mid h_t)},

whereas the buggy recomputation replaces the denominator with the learner engine's own probability, which is one by construction at zero parameter lag. Recomputing the denominator with the learner engine therefore hides, rather than fixes, the mismatch. The sampled-token log-ratio separates cleanly into a kernel term and a staleness term, and a fully synchronous system can still be off-policy because the kernel term generally survives at $n=v$.

What perturbs the logits

Tiling and fusion decisions; floating-point accumulation and reduction order; KV-cache layout and prompt-state construction; batch shape and sequence length; tensor- and sequence-parallel communication; precision, quantization, and numerical rounding. The gap may be negligible in one configuration and large enough to affect RL optimization in another; it must be measured rather than inferred from checkpoint identity.

Why Linear Attention and SSMs Are Especially Sensitive

For linear attention and state-space models the distinction is structural. Inference is naturally a recurrence $S_t = F_\theta(S_{t-1}, x_t)$, while efficient training uses a parallel scan, a convolutional dual, or a chunkwise factorization. The chunk size $C$ interpolates between execution schedules: $C=1$ is fully recurrent, $C=L$ is fully parallel. The formulas can be algebraically equivalent yet implement different finite-precision operators: changing the chunk schedule changes the parenthesization of the recurrence, and a small difference in one intermediate state can alter every later state and every later token probability.

Figure 2: Chunking changes the execution schedule. Algebraic equivalence across the line does not imply equality of the induced token distributions.

DeltaNet makes this particularly sharp because the state update is a data-dependent corrective write: $S_t = S_{t-1}(I - \beta_t k_t k_t^{\top}) + \beta_t v_t k_t^{\top}$. An error in $S_{t-1}$ changes the read, which changes the value erased by the delta rule, which changes the next read and corrective write. The state error is not merely observed downstream; it participates in the future update dynamics. Recurrent DeltaNet rollout and learner scoring with the parallelized delta-rule kernel must be regarded as different effective policies unless finite-precision parity is explicitly demonstrated.

Four Complementary Remedies

I

Treat kernel mismatch as off-policy data

Define the behavior distribution by the path that actually sampled each token, preserve that probability in the rollout record, and perform an explicit change of measure at the learner. One time-dependent behavior operator covers both prefill-decode mismatch and chunkwise-parallel-recurrent mismatch.

II

Make one execution rule canonical

Use the same chunkwise or recurrent execution rule for rollout and differentiable learner replay. TTT and Titans illustrate chunkwise memory-update semantics and RWKV-7 illustrates a recurrent architecture, but neither removes mismatch unless the same executable rule actually runs on both sides.

III

Raise precision where it compounds

Retain recurrent states, sensitive parameters, and state-update accumulations in higher precision. This shrinks numerical drift at its source, but it is a mitigation, not an equality proof.

IV

Speculate fast, accept exactly

Use a fast prefill or parallel path only as a proposal and let a recurrent target perform exact modified rejection sampling. The recurrent target defines the accepted policy, and the standard tokenwise correction is exact only when the proposal exposes the actual conditional probabilities of the drafted sequence.

Recommended design

Separate an exact recurrent-target mode from a scalable parallel-target surrogate, with explicit audits for forward probabilities, score functions, support, and state precision. Fast kernels accelerate everything that does not define the policy, and the audits say which mode a run actually earned.

Three Implementation Rules

1. Capture behavior probabilities from the actor path

Store the log-probability emitted by the actual sampling engine at the point where each token is drawn. A learner-side replay is an acceptable substitute only after distributional parity with the actor path has been established.

2. Do not fold kernel mismatch into parameter staleness

Log the rollout-to-learner probability gap at fixed weights separately from version lag. Otherwise actor refresh may appear to solve a mismatch that actually comes from execution kernels.

3. Reserve "on-policy" for a measured execution contract

Literal on-policy learning requires either the same recurrent path for rollout and differentiable learner replay, or a parallel path verified to induce the same token distribution under the declared batch shape, precision, cache construction, chunk schedule, and parallel layout.

When those conditions are not met, off-policy is the correct systems abstraction even if rollout and learner use the same nominal checkpoint. The report includes a minimal fixed-weight parity audit for measuring the gap directly.

Citation

If you find this work useful, please cite:

@misc{zhang2026kernelmismatch,
  title        = {Reliable RL Scaling Requires Accounting for Prefill-Decode Kernel Mismatch},
  author       = {Zhang, Yifan and others},
  year         = {2026},
  howpublished = {\url{https://github.com/yifanzhang-pro/Pretraining-RL-Science}}
}