The Two-Phase Revelation: Debugging vLLM 0.16's Async Scheduling in Hidden State Extraction

Introduction

In the high-stakes world of large language model deployment, few moments are as simultaneously frustrating and illuminating as watching a twenty-minute model load culminate in a cryptic runtime error. Message [msg 2648] captures exactly such a moment in the development of an EAGLE-3 speculative decoding training pipeline for the Kimi-K2.5 model. After resolving a cascade of API incompatibilities between the speculators v0.3.0 library and the installed vLLM 0.16 nightly, the assistant launched a third attempt at hidden state extraction—only to be met with a new, deeper error that revealed a fundamental architectural shift in vLLM's execution model.

The Message

The message is deceptively brief. It begins with a progress report—"84% (54/64), a few more minutes"—indicating that the model has nearly finished loading its 64 safetensor checkpoint shards across 8 GPUs. The assistant then executes a bash command to check the log after a 600-second sleep:

[bash] sleep 600 && ssh root@10.1.230.174 "tail -30 /root/eagle3-train/extract_test3.log 2>/dev/null" 2>/dev/null
    raise RuntimeError(
RuntimeError: Worker failed with error 'State error: sample_tokens() must be called after execute_model() returns None.', please check the stack trace above for the root cause
(Worker_TP0 pid=256850) INFO 02-21 22:10:49 [multiproc_executor.py:732] Parent process exited, terminating worker
(Worker_TP2 pid=256852) INFO 02-21 22:10:49 [multiproc_executor.py:732] Parent process exited, terminating worker
(Worker_TP3 pid=256853) INFO 02-21 22:10:49 [multiproc_executor.py:732] ...

The output shows a RuntimeError that is both specific and profound: "State error: sample_tokens() must be called after execute_model() returns None." The workers are terminating because the parent process crashed, and the extraction pipeline has failed for the third time.

The Road to This Point

To understand the significance of this message, one must appreciate the journey that led here. The assistant has been building an EAGLE-3 training pipeline for Kimi-K2.5, a 1-trillion-parameter Mixture-of-Experts model deployed on 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The pipeline requires extracting hidden states from the model's internal layers—a process that feeds the EAGLE-3 draft model training.

The first attempt (extract_test2.log, [msg 2624]) failed with an AssertionError related to block hashes. The assistant traced this to vLLM 0.16's Request object requiring a block_hasher function for KV cache management—a requirement absent in the speculators library's VllmHiddenStatesGenerator. The fix involved patching the generator to import get_request_block_hasher and init_none_hash from vLLM's kv_cache_utils, then passing the resulting block hasher to each Request constructor ([msg 2635]-[msg 2641]).

After verifying the patch and cleaning GPUs, the assistant launched attempt 3 ([msg 2645]) with the same extraction parameters: batch size 4, max sequence length 2048, tensor parallelism across 8 GPUs, and target layers 2, 30, 58, and 60. The model began its ~20-minute loading ritual, and the assistant waited patiently, checking progress periodically ([msg 2646], [msg 2647]). Message [msg 2648] is the first check after the model finished loading—and the first indication that the block_hasher fix, while necessary, was not sufficient.

The Error: A Window into vLLM 0.16's Architecture

The error message is remarkably informative: "sample_tokens() must be called after execute_model() returns None." This single sentence reveals that vLLM 0.16 has adopted a two-phase execution model for its GPU model runner. In this design:

  1. Phase 1 (execute_model): The model runner performs the forward pass through the transformer layers, computes logits, and stores intermediate state (including hidden states and KV cache metadata) in an internal execute_model_state attribute. Critically, it returns None rather than the model output.
  2. Phase 2 (sample_tokens): After execute_model completes, the runner must be explicitly called with sample_tokens(), which retrieves the stored state, performs token sampling, and returns the actual ModelRunnerOutput. This is a significant departure from earlier vLLM versions, where execute_model would return the output directly. The change enables async scheduling—a feature where the scheduler can overlap execution and sampling across different requests for improved throughput. When async_scheduling=True (the default in vLLM 0.16), the model runner saves state and returns None, expecting the caller to complete the cycle with sample_tokens(). The speculators library's VllmHiddenStatesGenerator, written for an earlier vLLM API, only calls execute_model() and never calls sample_tokens(). This creates a state machine violation: the second call to execute_model() encounters leftover state from the first call and raises the error.

Assumptions Made and Broken

This message exposes several assumptions that were silently relied upon:

Assumption 1: The execute_model API is stable. The speculators library was written for an earlier version of vLLM where execute_model returned output directly. The assistant had already patched the library for vLLM 0.16's Request and Scheduler APIs but had not anticipated a change in the execution model itself.

Assumption 2: The block_hasher fix would resolve all scheduler-related issues. The assistant correctly identified that the first error (AssertionError on block_hashes) was a missing dependency in Request construction. Fixing this was necessary but insufficient—the deeper architectural change in the execution pipeline remained.

Assumption 3: Hidden state extraction doesn't need sampling. The assistant's custom worker patches capture hidden states during the forward pass. The sampling output is irrelevant for extraction purposes. However, vLLM 0.16's state machine doesn't distinguish between extraction and inference—it enforces the two-phase protocol regardless of whether sampling is needed.

Assumption 4: The speculators library is compatible with vLLM 0.16 nightly. The library's v0.3.0 release targets a specific vLLM version, and the nightly build has diverged significantly. Each API mismatch requires individual discovery and patching.

Knowledge Required to Understand This Message

To fully grasp the significance of message [msg 2648], one needs:

  1. Understanding of the EAGLE-3 training pipeline: EAGLE-3 is a speculative decoding framework that trains a lightweight draft model to predict hidden states from a target LLM. The extraction step runs the target model on training data and captures intermediate layer outputs.
  2. Knowledge of vLLM's architecture: vLLM (Very Large Language Model) is a high-throughput inference engine. Its v1 API introduced significant changes including the Scheduler, Request, and ModelRunner abstractions. The 0.16 nightly represents a bleeding-edge version with further API evolution.
  3. Awareness of tensor parallelism (TP): The model runs across 8 GPUs with TP=8, meaning each GPU processes 1/8 of the hidden dimension. This complicates hidden state capture because each TP rank holds a shard of the full hidden state.
  4. Familiarity with speculative decoding: Speculative decoding accelerates inference by having a small draft model generate candidate tokens that the target model verifies in parallel. EAGLE-3 trains the draft model on the target model's hidden states.
  5. Understanding of async scheduling: The concept of overlapping computation phases (forward pass vs. sampling) to improve GPU utilization. vLLM 0.16's implementation uses a state machine that enforces a strict protocol.

Knowledge Created by This Message

Despite being a failure, message [msg 2648] generates valuable knowledge:

  1. vLLM 0.16 uses a two-phase execution model: execute_model() returns None and stores state; sample_tokens() must be called to complete the cycle. This is a confirmed architectural fact about the nightly build.
  2. The speculators library is incompatible with vLLM 0.16's async scheduling: The VllmHiddenStatesGenerator was designed for a synchronous execution model and cannot run unmodified.
  3. Hidden state extraction cannot bypass the sampling protocol: Even though the extraction pipeline doesn't need sampled tokens, vLLM's state machine requires the full execute_model → sample_tokens cycle to reset internal state for the next iteration.
  4. The block_hasher fix was validated: The first error (assertion on block_hashes) did not recur, confirming that the patch to provide block_hasher to Request objects was correct and sufficient for that specific issue.
  5. Debugging distributed systems requires patience: Each attempt involves a ~20-minute model load followed by a brief execution window. Errors manifest only after the full load completes, making iteration cycles extremely long.

The Thinking Process

The assistant's reasoning in this message is visible in its structure and content. The progress report "84% (54/64), a few more minutes" shows that the assistant is monitoring the load progress and timing the check to occur shortly after completion. The 600-second sleep is calibrated to the remaining load time (approximately 10 minutes based on earlier rates of ~17 seconds per shard for the remaining 10 shards).

The assistant doesn't panic or overreact to the error. It simply records the output, noting the key error message and the worker termination logs. The decision to show the first 30 lines of the tail (rather than the full log) is strategic—the assistant knows the error will appear at the bottom of the log and wants to capture it efficiently.

What's notably absent from this message is an immediate attempted fix. The assistant is in observation mode, gathering information. The next messages ([msg 2649] onward) will show the assistant analyzing the error in depth, examining vLLM's execute_model source code, discovering the async_scheduling parameter, and ultimately crafting a combined patch that both disables async scheduling and adds the sample_tokens() call.

The Broader Context

This message sits at a critical inflection point in the debugging campaign. The assistant has already fixed one API incompatibility (block hashes) and now faces a second, deeper one. The cascade of failures follows a pattern familiar to anyone integrating bleeding-edge ML frameworks: each error reveals a new layer of API divergence between the library and the runtime.

The chunk summary for segment 21 describes the eventual resolution: the assistant patches the generator to set async_scheduling=False in the SchedulerConfig and adds a call to self.executor.sample_tokens(None) after each execute_model() call. These two changes—one disabling the async scheduling feature entirely, the other satisfying the state machine protocol—finally unblock the pipeline. Hidden state extraction runs successfully on 10 test samples at ~2280 tok/s, producing correctly shaped [512, 7168] bfloat16 tensors.

But in message [msg 2648], none of that resolution is visible yet. What we see is the raw moment of discovery—the point where a new obstacle reveals itself, and the assistant must decide how to respond. It's a testament to the iterative, patient nature of systems debugging, where each failure is not a setback but a signal that guides the next step.

Conclusion

Message [msg 2648] captures a pivotal debugging moment in the deployment of an EAGLE-3 training pipeline for Kimi-K2.5. The error—sample_tokens() must be called after execute_model() returns None—reveals a fundamental architectural change in vLLM 0.16's execution model, breaking the speculators library's assumptions about the inference API. The message demonstrates the iterative nature of integrating bleeding-edge ML frameworks, where each API incompatibility must be discovered, understood, and patched in sequence. The assistant's measured response—recording the error without immediate action—reflects a disciplined debugging methodology that prioritizes information gathering over hasty fixes. In the broader narrative of the session, this message marks the moment when the assistant realizes that the vLLM 0.16 integration requires not just surface-level API patches but a deeper understanding of the new execution paradigm.