NgramRL: Accelerating RL Post-Training for Structured Generation via Rollout-Aware N-gram Speculative Decoding

Teams doing RL post-training of LLMs, especially on structured outputs like text-to-SQL or code, should read this. It exploits a nearly free source of speedup: the same prompt is sampled again and again during training, so what the model wrote for it last epoch is the best draft for this epoch. In one sentence: keep a small per-prompt n-gram table built from that prompt's recent rollouts, use it as the drafter for speculative decoding, and let the current policy verify each token so the sampling distribution is unchanged. In late training on Spider with Qwen3-8B, it commits 7.77 tokens per forward pass versus 4.15 for SuffixDecoding, the strongest model-free baseline, and its prefetcher raises measured rollout throughput by another 1.38×.

What problem is it solving?

Each step of GRPO (Group Relative Policy Optimization, a common RL algorithm for LLMs) first samples a group of completions per prompt, then scores them and updates the policy. This rollout phase is generated token by token and can exceed 70% of step time.

Speculative decoding speeds up token-by-token generation: a cheap drafter guesses the next few tokens and the large model verifies them in one forward pass. Existing drafters fit RL poorly:

The authors observe that GRPO revisits the same prompts across epochs while the policy changes only slightly, so new completions for a prompt repeat its earlier ones. On Spider, a single earlier rollout covers 48–96% of a rollout one epoch later, with shared spans of 7–23 tokens. In the paper's Figure 2, two SQL answers to the same prompt share 50% early in training, 81% mid-training (only an alias differs), and 100% late.

Overall structure

There are three parts: a per-prompt n-gram table that drafts, standard speculative-sampling verification on the GPU, and a prefetcher that hides CPU lookup behind GPU compute. The table lives in host memory, so it costs no GPU memory and needs no training.

Per-prompt n-gram tables

The table records which token followed a given context and how often. After each step, the prompt's completions are cut into n-grams of order 2 to 6 and inserted. At lookup, the longest matching context proposes its continuation, chaining up to 16 draft tokens; with no match, decoding proceeds normally. Illustration: if "group by billing_country" was followed by "order" three times and "having" once in earlier rollouts, the table proposes "order" and continues the lookup from there. The paper calls the table's probability of proposing a draft q; always taking the most frequent successor means q=1, while sampling by counts gives "order" q=3÷4=0.75.

Prompt keying is where the gain comes from. Merging all prompts into one global table gives 2.41 accepted tokens per step in late training, about the same as SuffixDecoding's 2.43, while per-prompt keys reach 3.05; enlarging SuffixDecoding's cache fourfold to 40k responses only reaches 2.65. History is also bounded to the last w appearances of each prompt. For Qwen3-8B on GSM8K, w=4 gives 2.81 accepted tokens at 1878 MB, w=16 gives 3.38 at 3181 MB, and unbounded history slips to 3.34 at 3634 MB because old rollouts no longer resemble the current policy.

Why verification leaves training unchanged

If the table proposes token x with probability q and the policy assigns π, x is accepted with probability min(1, π÷q); on rejection a token is resampled from the residual where the policy exceeds the draft. With q=0.75 and π=0.9 the token is always accepted; with π=0.3 it is accepted with probability 0.4. If the table always proposes its top token (q=1), the acceptance probability is simply π. This standard rule guarantees every token is distributed exactly as if sampled from the policy, so rollouts, rewards, and gradients are unchanged in distribution and drafters can be compared on speed alone.

The prefetcher

With the table on the CPU, every decode iteration is a round trip: the GPU verifies, syncs the result to the host, and only then does the CPU look up, schedule, and launch the next step. The GPU sits idle about 29% of each iteration, mostly from synchronization and scheduling rather than lookup.

The next lookup depends on the verification outcome (j, b): j accepted draft tokens (17 possibilities when the draft length is 16) and the bonus token b the model samples at position j+1, unknown until verification ends. During the current forward, the prefetcher takes the most frequent successors from the table as candidates for b under each likely j, builds one draft per candidate, and ships a small table keyed by (j, b) to the GPU. A hit needs no CPU wait; a miss makes only that request fall back to normal lookup. A budget of N candidates is split by recent statistics, and after a full accept the bonus is the top candidate about 90% of the time, so that branch needs just one.

With GPU time Tg, hidden CPU time Tc, and hit rate c, speedup is (Tg+Tc)÷(Tg+(1−c)×Tc). Using Figure 3's timeline (26.1 ms synchronous, 19.2 ms prefetched) as an illustration, Tg≈19.2 and Tc≈6.9, so c=0.966 gives about 1.34×. The paper measures 1.38× against a 1.40× ceiling for Qwen3-8B on Spider.

Results

Setup: verl running GRPO on 2 A100s, 64 prompts × 5 completions per step for 15 epochs, Qwen3-8B and DeepSeek-LLM-7B, tasks GSM8K, Spider, and CodeContests. The metric is tokens committed per forward pass.

During epoch 1 every prompt is new, so NgramRL stays at 1, worse than the baselines. On Spider it peaks mid-run and then declines as SQL outputs diversify. On CodeContests the global cache leads early because boilerplate is shared across prompts. The advantage grows with output regularity and prompt specificity. Acceptance is heavy-tailed: on Spider 14.1% of steps accept 12 or more draft tokens and up to 10.1% accept all 16, which is why the draft length is 16.

Limitations I see

The measured wall-clock speedup is the prefetcher's 1.38× over synchronous NgramRL, reached on Spider; GSM8K and CodeContests get 1.05–1.06×. Tokens per forward is not wall-clock, and verifying 16 draft tokens costs more than a plain decode step, so the end-to-end gain over no speculation must be measured in your own setup. Epoch 1 brings no benefit, so short or single-pass training gains little. Host memory is significant: 3181 MB for Qwen3-8B GSM8K at w=16 with only 64 prompts per step, and it grows with prompt pool and answer length.

How an RL training team could try it

Measure the rollout share of step time and how often each prompt recurs. Wire the table into the rollout engine on a structured task with the paper's settings (max order 6, draft length 16, w=4–16) and log tokens per forward without changing training. Then measure wall-clock against no speculation, and finally add the prefetcher with N=8–32 while watching hit rate and GPU idle time.

Verdict

The key lesson: repeated sampling in RL training is itself a free drafter that needs no extra model, no GPU memory, and updates with the policy automatically.