laion/snowball-64k-eagle3-draft-r2egym

EAGLE-3 draft for GrugMoe agent rollouts

A one-layer draft adapted on target-generated agent trajectories. It requires a compatible GrugMoe target and the pinned EAGLE-3 serving fork. It is not a standalone language model.

The benchmark, adaptation and SkyRL recipes that produced and measured it are at https://github.com/lukedhlee/grugmoe-eagle3. The measured target is the 64k-context Snowball checkpoint; the draft was adapted on that model's own R2E-Gym rollouts, so speed on another workload needs its own benchmark and may need re-adaptation on that workload's rollouts.

It was warm-started from the base draft at https://huggingface.co/laion/snowball-64k-eagle3-draft-base; for a workload other than R2E-Gym, adapt from that one. This package omits the target embedding. The serving engine supplies the active target's embedding. The draft weights, output head and vocabulary maps are unchanged from the selected adapted checkpoint. Optimizer state, training logs, target weights and training examples are not included.

Compatibility and use

The measured target was the internal Snowball S3 Nemotron-Terminal step-1888 checkpoint: 26 layers, hidden size 2560, vocabulary 128256. This draft has one Llama-style layer, a 32,000-token draft vocabulary, a 2,048-token attention window, and auxiliary hidden-state slots [2,13,23].

Use the vLLM fork at 00d81ae11404ab68afef4142cdc883f7ba77e9a6 (branch lukedhlee/grugmoe-eagle3) from https://github.com/lukedhlee/vllm. Pass the downloaded draft directory to your target server:

{"method":"eagle3","model":"/path/to/downloaded/draft","num_speculative_tokens":3}

For SkyRL, this dictionary belongs in generator.engine_init_kwargs.speculative_config. Every inference worker needs the patched serving environment and access to the draft directory. The config deliberately has no saved target filesystem path. Select the target explicitly in your serving command.

The embedding-free export has passed tensor-content checks, and on 2026-09-11 it loaded and served on one GH200 node with the measured target (embedding supplied by the engine, pinned fork above): 167.6 decode tokens/s per stream and 2.59 accepted tokens per step on a 32k held-out replay, zero errors, matching a freshly adapted draft on the same node. Do not interpret the original measurements below as measurements of this export. Another target checkpoint needs its own acceptance and speed benchmark. The public 262k Grug model has not been tested, and equal tensor dimensions do not establish tokenizer or positional-encoding compatibility.

Training and measurements

Adaptation used 2,177 prefixes from target-generated R2E-Gym probe trajectories, about 64 million input tokens and 32 million supervised tokens. Training and validation were split by task identity and excluded the serving replay tasks. Training used prefixes up to 32k tokens, three training-time unroll steps, learning rate 5e-5 and six epochs. The selected checkpoint is epoch four, chosen by held-out served speed after validation-loss shortlisting.

On four GH200 GPUs, with 16 concurrent streams and three speculative tokens:

• Adapted draft: 160/158 decode tokens per second per stream, versus 139/140 for the pilot in two same-node repeats.
• At target RL update 36: 155 versus 96 tokens per second without a draft, about 1.61×.
• Mean acceptance length: about 2.5 emitted tokens per verification step, including the bonus/replacement token.
• Paired RL smoke: three updates in 1,859 versus 2,639 seconds, about 1.42×. Measured learning diagnostics were similar; this is not a long-run learning-equivalence result.

Serving was tested up to 64k context in the original experiment. Context beyond 64k and other target families are unvalidated. The original frozen draft held through 36 RL updates on the measured policy trajectory; this does not establish a universal refresh schedule.

Adapt to a new target

Generate that exact target's own outputs, extract its hidden states again, and build a task-disjoint dataset. Keep the included d2t-32000.npy and t2d-32000.npy when warm-starting this draft. The toolkit prepares a separate working copy that names your target stub and restores its embedding before training. Never silently replace the vocabulary mapping of learned draft weights.

The frozen target embedding is intentionally absent; this package still contains the draft output head and all learned draft tensors.