The Token Budget Pivot: From Blunt Sample Capping to Precise Token-Based Early Stopping in an EAGLE-3 Inference Pipeline

A Single Edit That Fixed a 2.6× Overshoot Problem

In the middle of a sprawling machine learning infrastructure session — spanning GPU driver installation, flash-attn compilation battles, SGLang server tuning, and EAGLE-3 drafter training — there is a moment that perfectly encapsulates the difference between a first-guess solution and a properly engineered one. Message <msg id=3958> is deceptively small: a single edit to a Python file, accompanied by a routine LSP diagnostic about an unresolvable import. But this edit represents the culmination of a rapid diagnostic loop triggered when the user observed that the inference pipeline was overshooting its token budget by a factor of 2.6×. Understanding why this edit exists, what it changed, and the reasoning that led to it reveals a great deal about how real-world ML infrastructure work actually unfolds.

The Subject Message

The message reads in full:

[assistant] Now update the existing results loading to count tokens, and pass to the inference function: [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/datasets/run_inference.py Edit applied successfully.

>

LSP errors detected in this file, please fix: <diagnostics file="/home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/datasets/run_inference.py"> ERROR [27:6] Import "transformers" could not be resolved </diagnostics>

On its surface, this is almost nothing — a one-line description and a tool call confirmation. The real substance lies in the chain of reasoning that produced it.

The Problem: A Blunt Instrument Meets Variable Token Lengths

The story begins roughly thirty messages earlier, in &lt;msg id=3923&gt;. The assistant was running a large-scale inference pipeline to generate synthetic training data for an EAGLE-3 draft model. The pipeline processed ten datasets (A1, A2, B1–B8), each containing thousands of prompts, and used SGLang as the inference server. The total estimated generation time was 57+ hours — far too long. By capping each dataset to roughly 10 million tokens, the assistant calculated the time could be reduced to 17–26 hours.

The first attempt at a cap was --max-samples 7000. This was explicitly acknowledged as a simplification: "it's simpler — I can just truncate the prompts files, or add a --max-samples flag" (&lt;msg id=3923&gt;). The assistant deployed this flag, launched the pipeline, and moved on to other tasks — fixing the stats collector and monitor to display token counts from the new /generate endpoint response format.

Then the user reported the problem in &lt;msg id=3954&gt;: "Still inferecing B2 even tho it's now >10M." The inference was still running on dataset B2_opencodeinstruct even though it had already surpassed the 10 million token target.

The assistant's response in &lt;msg id=3955&gt; is where the real diagnostic work happens:

Right — the --max-samples 7000 is a blunt instrument. B2 averages 3793 tok/response, so 7000 × 3793 = 26.6M tokens — way over the 10M target. We already have 1720 responses × 3693 avg = 6.4M tokens. We only need ~950 more samples to hit 10M.

This calculation reveals the core flaw: different datasets have wildly different average completion lengths. B2 averages ~3,793 tokens per response. A sample cap of 7,000 would produce 26.6 million tokens — more than 2.6 times the intended budget. Conversely, a dataset like B6_ultrachat with an average of ~1,500 tokens per response would only produce 10.5 million tokens under the same cap, barely hitting the target. The sample-based approach was fundamentally mismatched to the problem because it treated all samples as equal when they were anything but.## The Design Decision: Token Budget Instead of Sample Count

The assistant's reasoning in &lt;msg id=3955&gt; and &lt;msg id=3956&gt; reveals a careful design process. The approach was to add a max_total_tokens parameter to run_dataset_inference, which would count tokens from both existing results (loaded for resume support) and newly generated responses, then stop early when the budget was hit. The assistant explicitly noted a subtle concurrency challenge: "The tricky part is that asyncio.as_completed doesn't let us cancel easily, but we can stop collecting and let in-flight requests drain" (&lt;msg id=3956&gt;). This is a practical engineering insight — in an async inference loop using asyncio.as_completed, you cannot forcibly cancel already-dispatched requests. The safe approach is to stop accepting new results from the iterator while allowing in-flight HTTP requests to complete naturally. This prevents wasted computation while avoiding the complexity and potential instability of cancellation.

The implementation was split across three edits. First, in &lt;msg id=3956&gt;, the assistant modified run_dataset_inference to accept and enforce a token budget. Second, in &lt;msg id=3957&gt;, process_dataset was updated to compute the existing token count from loaded raw responses and pass the budget through. Third — and this is the subject message, &lt;msg id=3958&gt; — the existing results loading logic was updated to actually count tokens from the loaded responses, and the token budget was wired through to the inference function. A fourth edit in &lt;msg id=3959&gt; added the CLI argument --max-tokens-per-dataset and plumbed it through the argument parser.

Assumptions and Their Consequences

The initial --max-samples approach rested on a tacit assumption: that sample counts would correlate reasonably well with token counts across datasets. This was a reasonable simplifying assumption for a first pass — it required no per-dataset analysis, no estimation of average completion lengths, and no modification to the inference loop's control flow. It was the kind of shortcut that often works when datasets are homogeneous. But the datasets in this pipeline were far from homogeneous: B1_glaive averaged ~1,577 tokens, B2_opencodeinstruct averaged ~3,793, and A1_deepswekimi contained full agent trajectories averaging over 16,000 tokens per sample. The assumption collapsed as soon as the user checked the actual token count.

A second assumption was that the user would not notice the overshoot immediately. The assistant deployed the capped pipeline and moved on to fix the stats collector and monitor — essential tooling for visibility. But the user was watching the progress and caught the problem before those fixes were even deployed. This highlights the importance of real-time monitoring: without token counts in the display, the user had to infer the problem from the fact that B2 was still running suspiciously long.

Input Knowledge Required

To understand this message, one needs several pieces of context. First, the architecture of the inference pipeline: it processes datasets asynchronously using asyncio, with resume support that loads existing raw responses from JSONL files. Second, the data format: raw responses contain a completion_tokens field at the top level (after the pipeline was rewritten to use SGLang's /generate endpoint instead of the OpenAI-compatible chat completions API). Third, the token economics: each dataset has a target budget of ~10 million tokens, and the total pipeline has ~92 million tokens across all datasets. Fourth, the async concurrency model: asyncio.as_completed returns an iterator of futures that cannot be cancelled mid-flight, which constrains how early stopping can be implemented.

Output Knowledge Created

This message produced a concrete change to run_inference.py: the existing results loading code was modified to count completion tokens from loaded responses, and this count was passed to the inference function as part of the token budget enforcement. The output is not just a code change but a design pattern — a reusable approach to budget-aware async inference that can be applied to any similar pipeline. The pattern is: load existing results, sum their token counts, subtract from the budget, and use the remaining budget to limit new generation. This is fundamentally different from a sample cap because it respects the actual resource consumption (tokens) rather than a proxy (samples).

The Thinking Process

The reasoning visible across messages &lt;msg id=3955&gt; through &lt;msg id=3959&gt; shows a characteristic pattern of ML infrastructure debugging. First, quantitative analysis: "B2 averages 3793 tok/response, so 7000 × 3793 = 26.6M tokens." Second, constraint identification: "asyncio.as_completed doesn't let us cancel easily." Third, incremental implementation: the change is split into four focused edits, each touching a different layer of the code (inference function, dataset processing, results loading, CLI argument). Fourth, pragmatic acceptance of limitations: the LSP error about transformers import is noted but ignored because it's a development environment issue, not a runtime problem.

This message, for all its brevity, is a case study in how real engineering decisions are made under time pressure. The initial shortcut was reasonable; the correction was precise; and the implementation respected the real constraints of async Python and distributed inference. It is the kind of edit that looks trivial in isolation but represents a significant improvement in the pipeline's resource management.