The 100K Dataset Plan: A Pivotal Moment in EAGLE-3 Speculative Decoding

Introduction

In the long and arduous journey of deploying EAGLE-3 speculative decoding for the Kimi-K2.5 language model, there are moments that define the trajectory of the entire project. Message 3656 in this conversation represents one such moment — a carefully constructed plan to scale the training dataset by an order of magnitude, from 10,000 samples to 100,000, in a determined effort to make speculative decoding finally beat the non-speculative baseline.

This message is not merely a to-do list or a project plan. It is the crystallization of days of debugging, benchmarking, dataset research, and strategic decision-making. It reflects the team's diagnosis of the core problem — insufficient training data — and their proposed cure: a massive, multi-phase data pipeline that would consume days of compute time and terabytes of disk space. To understand this message fully, one must understand the painful debugging journey that preceded it, the critical bug that was fixed just hours earlier, and the sobering benchmark results that revealed EAGLE-3 was actually slower than running the model without any speculative decoding at all.

The Context: A Bug Fixed, A Problem Revealed

Just before this message was written, the team had been wrestling with a devastating bug. The EAGLE-3 draft model, despite being trained on 10,000 samples of hidden states extracted from Kimi-K2.5, was achieving a zero percent acceptance rate. Every draft token was being rejected by the target model. The root cause, traced through layers of code, was a single flag mismatch: the SGLang server had been started with --speculative-algorithm EAGLE instead of --speculative-algorithm EAGLE3. This seemingly trivial difference had catastrophic consequences.

The is_eagle3() check in the SGLang codebase is strict — only the string "EAGLE3" triggers the target model to capture and concatenate intermediate layer hidden states from layers [2, 30, 58]. With EAGLE, the draft model received only the final-layer hidden states (7168 dimensions) instead of the expected three-layer concatenated states (21504 dimensions). The draft model's fusion layer was silently bypassed, and all trained weights became useless.

After fixing this flag and restarting the server, the team benchmarked the corrected configuration. The results were sobering: the best EAGLE-3 configuration achieved 82.3 tokens per second with CUDA graphs and 5 draft tokens, compared to the non-speculative baseline of 90.0 tokens per second. EAGLE-3 was not just failing to provide a speedup — it was actually making the system 9% slower. The acceptance length of approximately 2.1 tokens was simply insufficient to overcome the overhead of speculative decoding.

The EAGLE-3 paper's scaling curves suggested a clear path forward: more training data. The paper had used approximately 530,000 samples; SpecForge (SGLang's official EAGLE trainer) had used 320,000. The team's 10,000 samples were a rounding error by comparison. The user's response was decisive: scale up by 10×, find better datasets, and retrain.

The Message: A Comprehensive Six-Phase Plan

Message 3656 is the assistant's response to the user's directive. It opens with a crucial piece of news: "Disk resized to 11TB — that solves the storage constraint entirely." This single sentence resolves a tension that had been building throughout the previous discussions. The hidden state extraction pipeline, which saves per-token hidden states to disk for training, consumes approximately 92 gigabytes per 1,000 samples. At 10,000 samples, this was manageable at just under a terabyte. At 100,000 samples, the requirement would balloon to approximately 9.2 terabytes — far exceeding the previously available 1.8 terabytes on the /data partition. The disk resize to 11 terabytes removes this bottleneck entirely, making the entire plan feasible.

The assistant then lays out a six-phase plan, each phase with specific datasets, resource requirements, and timeline estimates. The plan is structured as follows:

Phase 1: Download and Prepare Prompt Datasets. Ten datasets are selected, totaling approximately 100,000 prompts. The selection reflects a deliberate strategy: an "agentic-heavy mix" that prioritizes coding, tool use, and reasoning over general chat. The datasets range from nebius/SWE-agent-trajectories (10,000 samples of full coding agent trajectories) to crownelius/KimiK2.5-2000x (2,000 samples of Kimi-K2.5's own outputs). Crucially, the assistant notes that most datasets will be used for their prompts only — the responses will be regenerated through Kimi-K2.5 to ensure the training data matches the target model's actual token distribution.

Phase 2: Response Generation via Kimi-K2.5. This is the most time-intensive phase. Each of the 100,000 prompts will be fed through the Kimi-K2.5 model running on the SGLang server to generate responses. At an estimated 2,000 output tokens per sample, this amounts to 200 million tokens of generation. The assistant calculates that at the server's peak throughput of approximately 2,300 tokens per second (at high concurrency), this would take roughly 24 hours.

Phase 3: Hidden State Extraction. The extracted hidden states from the SGLang server's patched dump mechanism would consume approximately 9.2 terabytes of disk space — now feasible with the 11-terabyte partition. At the previously observed extraction rate of 1.48 samples per second, this phase would take approximately 18.8 hours.

Phase 4: Vocabulary Mapping. This phase reuses existing mappings between the target model's vocabulary and the reduced 32,000-token draft vocabulary, requiring no new work.

Phase 5: Training. The EAGLE-3 draft model would be trained from scratch on the full 100,000 samples for 3-5 epochs, using the same architecture as before: one decoder layer, three auxiliary hidden state layers from layers [2, 30, 58]. The assistant estimates approximately 12 hours for this phase.

Phase 6: Deploy and Benchmark. The trained draft model would be deployed with the correct --speculative-algorithm EAGLE3 flag, CUDA graphs enabled, and benchmarked against the 90 tokens per second baseline. The target is an acceptance length greater than 3.0, which the assistant calculates would be sufficient to achieve a real speedup.

The total estimated timeline is approximately 2.5 days.

The Reasoning Behind the Dataset Selection

The dataset selection in Phase 1 reveals a sophisticated understanding of what makes good training data for speculative decoding. The assistant is not simply grabbing the largest datasets available; it is curating a mix designed to match the expected production use case of Kimi-K2.5.

The emphasis on agentic coding datasets is deliberate. The user's application involves agentic coding tasks — the model writing code, using tools, debugging, and reasoning through software engineering problems. The draft model needs to learn to predict tokens in precisely these scenarios. The inclusion of SWE-Factory/DeepSWE-Agent-Kimi-K2-Trajectories-2.8K is particularly clever: these trajectories were generated by Kimi-K2 itself, meaning their token distribution closely matches what the target model would produce. These samples can be used as-is without response regeneration.

The reasoning datasets (open-r1/Mixture-of-Thoughts, open-thoughts/OpenThoughts-114k) serve a different purpose. Kimi-K2.5 is a reasoning model that generates <think> blocks — long chains of internal reasoning before producing a final answer. The draft model must learn to predict these reasoning patterns, including the distinctive token distributions of chain-of-thought text.

The general chat datasets (ultrachat_200k, ShareGPT_Vicuna_unfiltered) provide diversity. The EAGLE-3 paper itself used these datasets, and they help prevent the draft model from overfitting to the narrow patterns of coding and reasoning data.

The assistant explicitly notes that most datasets will be used only for their prompts, with responses regenerated through Kimi-K2.5. This is a critical insight: the draft model must learn to predict Kimi-K2.5's token distribution, not the distribution of GPT-4, DeepSeek, or any other model that generated the original dataset responses. Response regeneration is the mechanism that aligns the training distribution with the target distribution.

Assumptions Embedded in the Plan

The plan rests on several key assumptions, some explicit and some implicit.

The most fundamental assumption is that more data will improve acceptance rate. The EAGLE-3 paper's scaling curves show a clear relationship between training data volume and acceptance rate, but this relationship was demonstrated on different model architectures (LLaMA, Qwen) and different data distributions. Whether the same scaling law holds for Kimi-K2.5, a Mixture-of-Experts model with a unique token distribution, is an open question. The assistant implicitly assumes that the scaling behavior is transferable.

A second assumption is that the draft model architecture is not the bottleneck. The current architecture — one decoder layer with three auxiliary hidden state layers — was chosen based on the EAGLE-3 paper's recommendations for models of this size. But if the architecture itself is insufficient to capture the complexity of Kimi-K2.5's token distribution, then more data alone may not help. The plan does not include any experimentation with larger draft model architectures.

A third assumption is that response regeneration through the target model is sufficient to align the distributions. The assistant assumes that generating responses from Kimi-K2.5 using the prompts from these datasets will produce a training distribution that matches the model's behavior in production. But the prompts themselves come from different sources — SWE-agent issues, function-calling examples, general chat — and the model's response distribution may vary significantly across these domains. The draft model might learn domain-specific patterns that don't generalize.

A fourth assumption concerns timeline estimates. The assistant estimates 24 hours for response generation, 19 hours for hidden state extraction, and 12 hours for training, totaling approximately 2.5 days. These estimates assume no failures, no server crashes, no disk space issues, and no unexpected bugs. In a project that has already seen multiple server crashes, build failures, and configuration bugs, this assumption of smooth execution is optimistic.

The Thinking Process Visible in the Message

The message reveals a structured, analytical thinking process. The assistant begins by acknowledging the resolved constraint (disk space), then immediately moves to the plan structure. The use of a table to present the dataset mix shows an attempt to make complex information digestible and to justify each selection with a clear rationale.

The assistant's thinking is also visible in the "Key Design Decisions" section of the previous message (message 3655), which feeds directly into this plan. The assistant had considered alternatives — using mlabonne/open-perfectblend directly (as SpecBundle does), or using SpecForge's online training mode to avoid disk storage entirely — and had posed these as questions to the user. The user's answers (agentic-heavy mix, disk resized to 11TB) shaped the final plan.

The timeline estimates reveal an awareness of the system's actual performance characteristics. The assistant knows the extraction rate (1.48 samples per second), the server's peak throughput (~2,300 tok/s), and the disk consumption per sample (~92 GB per 1K samples). These numbers come from actual measurements during the previous 10K-sample run, not from theoretical calculations.

The plan also shows an understanding of what can be reused and what must be redone. The vocabulary mappings (t2d.pt, d2t.pt) can be reused because the draft vocabulary hasn't changed. The training script (04_train.py) needs only path updates. But the hidden states must be re-extracted because they depend on the specific responses generated by Kimi-K2.5, which will be different for the new prompts.

What This Message Creates

Message 3656 creates several things of value:

1. A shared understanding between the assistant and the user about what the next 2.5 days of work will look like. The plan serves as a contract — the user can approve it, modify it, or reject it before any work begins.
2. A concrete dataset selection with specific sample counts and rationales. This selection can be debated, refined, or replaced based on further research.
3. Resource requirement estimates that allow the team to plan infrastructure. The disk space requirement is now known to be ~9.2 TB; the compute time is known to be ~2.5 days.
4. A dependency graph showing the order of operations: datasets must be downloaded before response generation, which must complete before hidden state extraction, which must complete before training.
5. Success criteria: an acceptance length greater than 3.0, which would translate to a real speedup over the 90 tok/s baseline.

Potential Weaknesses and Unaddressed Questions

The plan, while comprehensive, leaves several questions unaddressed. How will the team handle dataset samples that are too short or too long for the 2048-token sequence length? What happens if some datasets fail to download or are restructured on HuggingFace? How will the team validate that the response regeneration is producing high-quality outputs before proceeding to the expensive hidden state extraction phase?

The plan also assumes that training from scratch on 100K samples will yield better results than continued training on the existing 10K samples. There is a case to be made for incremental training — starting from the existing checkpoint and fine-tuning on additional data — which might converge faster and preserve whatever knowledge the current draft model has already acquired.

The acceptance length target of 3.0 is somewhat arbitrary. The assistant calculates that 2.5 would be needed to break even and 3.0+ for real speedup, but these numbers depend on the specific overhead of the speculative decoding implementation, which may change with CUDA graph optimizations or different configuration parameters.

Conclusion

Message 3656 represents a pivotal moment of strategic planning in the EAGLE-3 project. After days of debugging, benchmarking, and dataset research, the team has converged on a clear diagnosis (insufficient data) and a concrete treatment plan (10× scale-up with carefully curated datasets). The plan is ambitious — 2.5 days of continuous computation, 200 million tokens of generation, 9.2 terabytes of hidden states — but it is grounded in real measurements and informed by the latest research.

The message itself is a model of technical communication: it acknowledges constraints, presents options with clear rationales, estimates resources honestly, and defines success criteria. Whether the plan succeeds or fails — whether 100,000 samples is enough to push EAGLE-3 past the 90 tok/s baseline — the thinking behind it is sound, and the path forward is clear. The next 2.5 days will determine whether speculative decoding becomes a valuable optimization for Kimi-K2.5 or remains an interesting experiment that couldn't quite deliver.