The Token Budget Pivot: Precision Over Brute Force in Async Inference Pipelines

In the high-stakes world of large-scale machine learning inference, the difference between a pipeline that finishes in 17 hours and one that runs for 57 hours can come down to a single design decision. Message 3956 in this opencode session captures exactly such a moment: the realization that a sample-count cap is a blunt instrument, and the subsequent implementation of a token-based budget that brings precision and control to a sprawling multi-dataset inference pipeline.

The Context: Generating Training Data for EAGLE-3

The broader project involves training an EAGLE-3 speculative decoding drafter for the Kimi-K2.5 model. To train this drafter, the team needs to generate high-quality response data across eight datasets (B1 through B8), each with different characteristics — different numbers of prompts, different average response lengths, and different domains. The datasets range from B1_glaive (10,000 prompts, ~1,577 tokens average) to B8_sweagent (3,572 prompts, ~8,000 tokens average). The total estimated generation time at the pipeline's optimized throughput of ~930-1350 tokens/second was a staggering 57+ hours for all 88K samples.

The solution was to cap the dataset at a token budget: 10 million tokens per dataset, totaling approximately 92 million tokens across all datasets, which would reduce the estimated runtime to a more manageable 17-26 hours. This was a pragmatic trade-off — enough data to train a good drafter without waiting multiple days for full coverage.

The Blunt Instrument: --max-samples 7000

The first attempt at implementing this cap was straightforward but flawed. In messages 3926-3931, the assistant added a --max-samples argument to run_inference.py. The idea was simple: load at most N prompts per dataset, and the token count would roughly follow. A value of 7000 was chosen as a conservative upper bound.

But this approach made a critical assumption: that all datasets have similar average token lengths. They do not. B2_opencodeinstruct averages 3,793 tokens per response, while B6_ultrachat averages only 1,500. With --max-samples 7000, B2 would generate 7000 × 3793 = 26.6 million tokens — more than 2.5 times the intended 10M budget. Meanwhile, B6 would generate only 7000 × 1500 = 10.5 million tokens, close to the target but only because its average happened to be lower.

The user noticed the problem immediately. In message 3954, they observed: "Still inferecing B2 even tho it's now >10M." B2 had already crossed the 10 million token threshold, but the pipeline kept going because the cap was on samples, not tokens.

The Realization and the Pivot

Message 3955 is the assistant's moment of recognition: "Right — the --max-samples 7000 is a blunt instrument." The assistant calculates the numbers: B2 at 3,793 tokens per response with 7000 samples yields 26.6M tokens. They already have 1,720 responses totaling 6.4M tokens. Only ~950 more samples are needed to hit 10M. The proper fix is clear: cap by token count, not sample count.

This brings us to message 3956, the subject of this article.

Message 3956: Implementing Token Budget Support

The message is concise but dense with reasoning:

I need to add token budget support. The approach: pass max_total_tokens to run_dataset_inference, count tokens from existing + new results, and stop early when budget is hit. 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.

This single paragraph reveals a sophisticated understanding of both the domain problem and the technical architecture. Let us unpack it layer by layer.

The Design Decision: Where to Put the Cap

The assistant chooses to implement the token budget inside run_dataset_inference, the async function that actually dispatches requests to the SGLang server and collects results. This is an architectural decision with important implications. Rather than truncating the prompt list before inference begins (which was the --max-samples approach), the token budget is enforced dynamically during execution. This means the pipeline can account for the actual token counts of generated responses, which may vary significantly from the averages used in planning.

The function signature would need a new parameter: max_total_tokens. The logic would accumulate token counts from both existing results (loaded for resume support) and newly generated responses, checking after each completed request whether the budget has been exhausted.

The Async Challenge

The most revealing part of the message is the acknowledgment of the async complication: "The tricky part is that asyncio.as_completed doesn't let us cancel easily."

The inference pipeline uses asyncio.as_completed to process requests concurrently. This is a standard pattern: you submit a batch of requests, then iterate over them as they complete, processing results one by one. The challenge is that once requests are in flight — sent to the server, awaiting response — you cannot easily cancel them. Python's asyncio provides no mechanism to forcefully terminate a pending coroutine without risking resource leaks or inconsistent state.

The assistant's solution is pragmatic: "we can stop collecting and let in-flight requests drain." Instead of trying to cancel the remaining requests, the pipeline simply stops accepting new results. The in-flight requests complete naturally, their results are collected and written to disk, but no new requests are dispatched. This means the final token count may slightly exceed the budget (by the tokens from the draining requests), but this is a small and acceptable overshoot compared to the 2.6× overshoot from the sample-based cap.

The Edit and the LSP Warning

The message includes an edit action applied to run_inference.py, followed by an LSP diagnostic:

ERROR [27:6] Import "transformers" could not be resolved

This is a false positive from the language server — the transformers library is installed in the remote environment but not in the local development environment where the LSP runs. The assistant correctly ignores this warning, as it does not affect runtime behavior. This pattern appears repeatedly throughout the conversation (messages 3926-3930 all show the same diagnostic), and the assistant has learned to recognize it as harmless.

Assumptions Embedded in the Approach

The token budget implementation makes several assumptions worth examining:

First, it assumes that token counts from existing results can be accurately summed. The raw_responses.jsonl files store completion_tokens as a flat field at the top level of each JSON record (as confirmed in message 3943). This was a change from the earlier OpenAI-compatible format that nested token counts inside a usage dictionary. The stats_collector had to be updated (messages 3944-3945) to handle this new format, but by message 3956 the format is stable and the counts are reliable.

Second, it assumes that the token budget check at the right granularity will not cause excessive overhead. Checking after each completed request is the natural granularity for an asyncio.as_completed loop, and the check itself is a simple integer comparison — negligible cost.

Third, it assumes that draining in-flight requests is acceptable. This is a judgment call about the cost-benefit trade-off. The extra tokens from draining are bounded by the number of concurrent requests (configured at 150 for short prompts, 32 for long prompts) times the maximum tokens per request (10,240 or 16,384). In the worst case, this could add up to ~1.5 million tokens for short requests or ~524,000 tokens for long requests — a small overshoot relative to the 10M budget.

The Broader Significance

This message sits at the intersection of several themes that run through the entire session. The inference pipeline had already undergone a major rewrite to fix a reasoning capture bug (the --reasoning-parser configuration issue that caused thinking content to be embedded in message.content instead of reasoning_content). The server had been tuned for throughput, with KV cache optimization, hierarchical cache configuration, and NCCL settings all adjusted to push from ~600 tok/s to ~930-1350 tok/s. The stats_collector and monitor had been fixed to properly display token counts from the new /generate response format.

The token budget cap is the final piece that makes the pipeline production-ready. Without it, the team would either waste days generating far more data than needed (the 57-hour estimate for full coverage) or manually intervene to stop each dataset when it reaches the budget. With it, the pipeline can run unattended, stopping each dataset automatically when the token budget is exhausted.

Conclusion

Message 3956 is a small but critical moment in a complex engineering effort. It demonstrates the importance of choosing the right abstraction for a problem — sample count vs. token count — and the value of dynamic, runtime enforcement over static, pre-computation limits. It also shows thoughtful async design, acknowledging the limitations of asyncio.as_completed and finding a pragmatic compromise rather than fighting the framework.

The edit itself is straightforward, but the reasoning behind it reveals a deep understanding of the system's architecture, the data characteristics, and the practical constraints of distributed inference. It is a reminder that in machine learning engineering, the most impactful decisions are often not about model architecture or hyperparameters, but about the mundane but essential infrastructure that makes large-scale experimentation possible.