The Async Scheduling Revelation: Unlocking vLLM 0.16's Two-Phase Execution for Hidden State Extraction

Introduction

In the sprawling, multi-day effort to deploy EAGLE-3 speculative decoding for the Kimi-K2.5 model on 8x Blackwell GPUs, few moments were as pivotal as message [msg 2652]. This single assistant message represents a breakthrough in understanding — a moment where a cascade of cryptic runtime errors crystallized into a clear architectural insight. The message, deceptively brief, captures the instant when the assistant realized that vLLM 0.16 had fundamentally changed its execution model, introducing a two-phase forward pass that broke every assumption the speculators library had been built upon. This article dissects that message in detail: the reasoning behind it, the context that made it necessary, the assumptions it challenged, and the knowledge it created.

The Message in Full

The subject message reads:

The issue is clear. In vLLM 0.16 with async_scheduling=True (which is the default), execute_model() returns None and stores state internally, then sample_tokens() must be called separately to get the actual output. The speculators code only calls execute_model().

>

But we don't actually need the sampling output — we just need the hidden states. Let me look at what happens: [bash] ssh root@10.1.230.174 "sed -n '3600,3670p' /root/ml-env/lib/python3.12/site-packages/vllm/v1/worker/gpu_model_runner.py" 2>/dev/null if logits is not None: model_output_broadcast_data["logits"] = logits.contiguous()

>

broadcasted = get_pp_group().broadcast_tensor_dict( model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 ) assert broadcasted is not None logits = broadcasted["logits"]

>

self.execute_model_state = ExecuteModelState( scheduler_output, logits, ...

This message is the culmination of a debugging chain that began with an AssertionError about block_hashes and progressed through a sample_tokens() state machine error. The assistant, having already fixed the block hasher issue, now confronts a deeper architectural incompatibility.

The Context: A Debugging Odyssey

To understand why this message matters, one must appreciate the journey that led to it. The assistant had been working for hours to deploy the EAGLE-3 training pipeline for the Kimi-K2.5 model — a massive 1-trillion-parameter Mixture-of-Experts model running on 8x NVIDIA RTX PRO 6000 Blackwell GPUs. The pipeline required extracting hidden states from specific layers of the model to train a draft model for speculative decoding.

The speculators library (v0.3.0) provided a VllmHiddenStatesGenerator class designed to do this extraction, but it was written for an earlier version of vLLM. The installed environment used vLLM 0.16 nightly, which had undergone significant API changes. Previous messages document a series of failures:

The Core Insight: Two-Phase Execution

The key realization in this message is that vLLM 0.16, with async_scheduling=True (the default), has split the model forward pass into two distinct phases:

  1. execute_model(): Runs the transformer layers, computes logits, and stores intermediate state in self.execute_model_state. Returns None.
  2. sample_tokens(): Called afterward, uses the stored state to perform sampling and produce the actual output tokens. This is a significant architectural change from earlier vLLM versions, where execute_model() would return the output directly. The motivation for this split is performance: by separating the computationally expensive forward pass from the lightweight sampling step, the scheduler can overlap communication and computation more efficiently. The model runner can release GPU resources sooner and let the scheduler coordinate the next batch while sampling runs asynchronously. The speculators library, written for an older vLLM, only called execute_model() and expected it to return the model output. In vLLM 0.16, this call returns None, and the subsequent attempt to use the result triggers the state machine error: sample_tokens() must be called after execute_model() returns None.

The Assistant's Reasoning Process

The message reveals a clear chain of reasoning. The assistant states "The issue is clear" — this is not a guess but a conclusion drawn from multiple pieces of evidence gathered in the preceding messages. Let's trace the logic:

  1. Observation: The error message explicitly states sample_tokens() must be called after execute_model() returns None. This is a state machine constraint — the model runner enforces a specific call order.
  2. Hypothesis formation: The assistant hypothesizes that vLLM 0.16 uses a two-phase execution model where execute_model() returns None and stores state internally, requiring a separate sample_tokens() call.
  3. Confirmation via code inspection: The assistant immediately runs a bash command to examine the relevant section of gpu_model_runner.py (lines 3600-3670), looking at how execute_model_state is constructed. The code snippet shows ExecuteModelState(scheduler_output, logits, ...) being stored, confirming the hypothesis.
  4. Pragmatic reframing: Crucially, the assistant then notes: "But we don't actually need the sampling output — we just need the hidden states." This reframes the problem. The two-phase execution is an obstacle only if we need the final sampled tokens. For hidden state extraction, the assistant only needs the intermediate layer outputs, which are computed during execute_model() before the sampling phase. This insight opens up a simpler fix path: rather than adapting the full two-phase flow, the assistant can potentially capture hidden states directly within the execute_model() call, bypassing the need for sample_tokens() entirely.

Assumptions Made and Challenged

This message challenges several assumptions that had been implicit in the earlier work:

Assumption 1: The speculators library's API is compatible with vLLM 0.16. This was the fundamental assumption driving the entire debugging effort. The assistant had been patching the library incrementally — first fixing imports, then the block_hasher, then the custom worker's forward signature. Each fix revealed a deeper incompatibility. Message [msg 2652] reveals that the incompatibility is not just a matter of missing function parameters or changed class names — it's a fundamental architectural change in how the model executes.

Assumption 2: execute_model() returns the model output. This is the assumption that directly caused the sample_tokens() error. The speculators code, like most vLLM client code written before 0.16, treated execute_model() as a synchronous call that returns logits and sampled tokens. The two-phase split breaks this assumption completely.

Assumption 3: Hidden state extraction requires the full execution pipeline. The assistant's reframing — "we don't actually need the sampling output" — challenges this assumption. If the goal is only to capture intermediate hidden states, the sampling phase is unnecessary overhead. The hidden states are computed during the forward pass within execute_model(), and the custom worker's forward function already captures them. The issue is that the generator's control flow expects execute_model() to return something usable, and the vLLM runner's state machine enforces the two-phase protocol even when sampling isn't needed.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

  1. vLLM's architecture: Understanding that vLLM uses a scheduler-worker model where the GPUModelRunner executes model forward passes. Knowledge of the execute_model / sample_tokens split and the async_scheduling configuration parameter.
  2. The speculators library: Understanding that VllmHiddenStatesGenerator wraps vLLM's execution to capture intermediate hidden states, and that it was written for an older vLLM API.
  3. EAGLE-3 training pipeline: The broader goal — extracting hidden states from specific layers of a target model to train a draft model for speculative decoding. The pipeline involves: (1) preparing tokenized data, (2) extracting hidden states via vLLM, (3) training the EAGLE-3 draft model, (4) deploying the combined model.
  4. Distributed execution: The code snippet shows get_pp_group().broadcast_tensor_dict() — this is pipeline-parallel communication, where logits are broadcast from the last pipeline stage to all ranks. Understanding pipeline parallelism helps interpret the code.
  5. Python and PyTorch: The ability to read Python code with PyTorch tensor operations, and to understand state machine patterns in async systems.

Output Knowledge Created

This message creates several pieces of actionable knowledge:

  1. The root cause of the sample_tokens() error: The error is not a bug in the speculators code per se, but a fundamental API mismatch. The speculators code was written for a vLLM version where execute_model() returned output directly. vLLM 0.16's two-phase execution makes this assumption invalid.
  2. The two-phase execution model: The message documents that vLLM 0.16 with async_scheduling=True splits execution into execute_model() (returns None, stores state) and sample_tokens() (uses stored state, returns output). This is critical knowledge for anyone writing custom vLLM workflows.
  3. A path forward: By recognizing that hidden states don't require sampling, the assistant identifies a simpler fix: capture hidden states within the execute_model() call and bypass the sample_tokens() requirement. This insight shapes the subsequent debugging direction.
  4. The specific code location: The bash command pinpoints lines 3600-3670 of gpu_model_runner.py as the relevant code, showing exactly where ExecuteModelState is constructed. This gives the assistant a precise target for further investigation.

The Broader Significance

Message [msg 2652] is a turning point in the debugging effort. Before this message, the assistant was chasing symptoms — assertion errors, ambiguous tensor truthiness, state machine violations. After this message, the assistant understands the underlying cause and can design a targeted fix.

The message also exemplifies a crucial skill in AI engineering: the ability to recognize when a problem is not a bug but a fundamental architectural mismatch. The speculators library wasn't "broken" — it was designed for a different version of vLLM with a different execution model. The fix isn't to patch individual function calls but to adapt the control flow to the new architecture.

This insight also has implications beyond this specific debugging session. Anyone working with vLLM 0.16's async scheduling needs to understand the two-phase execution model. The execute_model / sample_tokens split affects how custom model runners, speculative decoding implementations, and hidden state extraction tools interact with the vLLM runtime. The assistant's discovery in this message is a piece of general knowledge about vLLM 0.16 that would benefit any developer building on top of it.

Conclusion

Message [msg 2652] is a masterclass in diagnostic reasoning. The assistant takes a cryptic runtime error, traces it to a fundamental architectural change in vLLM 0.16, and reframes the problem in a way that opens up a simpler solution path. The message is concise — just a few sentences and a bash command — but it represents the synthesis of hours of debugging across multiple previous messages. It demonstrates how understanding the deep architecture of a system is often the key to solving seemingly unrelated errors, and how a single insight can transform a frustrating debugging session into a clear path forward.