The Pivot: Disabling Async Scheduling to Unlock EAGLE-3 Hidden State Extraction

Introduction

In the complex ecosystem of large language model deployment, few tasks are as delicate as extracting intermediate representations from a running inference engine. Message [msg 2653] represents a critical turning point in an extended debugging session aimed at unblocking the EAGLE-3 training pipeline for the Kimi-K2.5 INT4 model running on 8x Blackwell GPUs. In this message, the assistant synthesizes two previously separate errors into a single coherent diagnosis and proposes a remarkably simple fix: disable async scheduling. This decision, born from deep investigation of vLLM 0.16's internal architecture, demonstrates the kind of system-level understanding required when adapting third-party tools to rapidly evolving inference frameworks.

The Context: A Cascade of API Incompatibilities

To understand the significance of this message, one must first appreciate the debugging journey that preceded it. The assistant had been attempting to use the speculators v0.3.0 library's VllmHiddenStatesGenerator to extract hidden states from intermediate layers of the Kimi-K2.5 model — a prerequisite for training an EAGLE-3 speculative decoding draft model. However, the speculators library was written for an earlier version of vLLM, while the deployment used vLLM 0.16 nightly — a version with substantially rewritten internals.

The first major obstacle was a block hashing assertion failure ([msg 2625]). In vLLM 0.16, the KV cache block pool always expects block_hashes on Request objects, even when prefix caching is disabled. The speculators library was creating Request objects without providing a block_hasher function, causing an AssertionError deep in the scheduler. The assistant methodically traced this through the vLLM source code ([msg 2626]-[msg 2634]), discovering the get_request_block_hasher function in kv_cache_utils.py and the init_none_hash global initialization requirement. A patch was crafted and deployed ([msg 2639]-[msg 2643]), successfully resolving the first error.

With the block hasher fix in place, a second extraction attempt ([msg 2645]-[msg 2648]) progressed past the assertion but immediately hit two new errors: a Boolean value of Tensor with more than one value is ambiguous error from a truthiness check on aux_hidden_states, and a state machine error: sample_tokens() must be called after execute_model() returns None. The assistant investigated the latter by examining vLLM's gpu_model_runner.py ([msg 2650]-[msg 2652]), discovering that vLLM 0.16 introduced a two-phase execution model with async scheduling.

The Insight: Connecting Two Errors into One Diagnosis

Message [msg 2653] opens with a moment of synthesis: "I see — the execute_model_state stores aux_hidden_states as part of the state. This is vLLM's built-in EAGLE-3 support!"

This is the key insight. The assistant realizes that the two errors are not independent — they are symptoms of the same root cause. In vLLM 0.16 with async_scheduling=True (the default), the model runner's execute_model() method no longer returns the model output directly. Instead, it stores intermediate state — including the aux_hidden_states tensor that contains the EAGLE-3 features — in an internal execute_model_state attribute. The actual output is obtained by calling a separate sample_tokens() method.

This architectural change explains both errors:

  1. The sample_tokens() error: The speculators code calls execute_model() and expects the output immediately, but vLLM 0.16 returns None and defers output to sample_tokens().
  2. The ambiguous boolean error: The speculators code checks if not aux_hidden_states: to see if hidden states were captured, but aux_hidden_states is now a list of tensors stored in execute_model_state, not a simple None value. The assistant's recognition that vLLM has "built-in EAGLE-3 support" is particularly insightful. The aux_hidden_states field in execute_model_state is not an accident — it's a deliberate mechanism for supporting speculative decoding architectures like EAGLE that require access to intermediate hidden states. The speculators library was trying to reimplement this capture mechanism, but vLLM 0.16 already has it built in.

The Decision: Disable Async Scheduling

Having identified the root cause, the assistant proposes a solution: "The simplest fix is to disable async scheduling so that execute_model() returns the output directly. We can do this by setting async_scheduling=False in the SchedulerConfig."

This decision reflects a pragmatic trade-off. The assistant could have:

  1. Adapted the speculators code to the two-phase model: Call execute_model(), then call sample_tokens(), and extract hidden states from execute_model_state in between. This would preserve async scheduling benefits but require more extensive code changes.
  2. Directly access execute_model_state: Read the aux_hidden_states from the model runner's internal state without going through the normal output path. This would be fragile and architecture-specific.
  3. Disable async scheduling: The simplest change, requiring only a single configuration parameter. The trade-off is potentially lower throughput for the extraction step, but since this is a one-time data generation step for EAGLE-3 training, throughput is less critical than correctness. The assistant chooses option 3, and the message ends with a grep command to check whether async_scheduling is already referenced in the speculators generator file — a quick reconnaissance to understand the current code before making changes.

Assumptions and Their Validity

The assistant makes several assumptions in this message:

  1. Disabling async scheduling is safe for hidden state extraction: This is a reasonable assumption. The extraction step runs a single prefill pass per batch — there's no iterative decoding, so the performance benefits of async scheduling are minimal. The scheduler's state machine is simpler without async scheduling, reducing the chance of edge cases.
  2. The SchedulerConfig exposes async_scheduling as a settable parameter: The grep command at the end of the message suggests the assistant is about to verify this. In vLLM 0.16, SchedulerConfig does have an async_scheduling field, typically defaulting to True. Setting it to False should restore the synchronous execution behavior that the speculators library expects.
  3. vLLM's built-in EAGLE-3 support is compatible with the speculators' approach: The assistant assumes that by disabling async scheduling, execute_model() will return the model output including aux_hidden_states in a format compatible with what the speculators code expects. This is plausible but not guaranteed — the internal data structures may differ.
  4. The ambiguous boolean error will also be resolved: The assistant doesn't explicitly state this, but disabling async scheduling should cause execute_model() to return the output directly, which means aux_hidden_states will be a proper return value rather than an internal state field. This should resolve the truthiness check issue as well.

Input Knowledge Required

To fully understand this message, one needs:

  1. vLLM 0.16 architecture: Knowledge of the two-phase execution model (execute_model + sample_tokens), the async_scheduling configuration parameter, and the execute_model_state mechanism for passing intermediate state between phases.
  2. EAGLE-3 speculative decoding: Understanding that EAGLE-3 requires hidden states from intermediate layers of the base model to train a lightweight draft model. These hidden states are typically extracted during a prefill pass.
  3. The speculators library: Awareness that VllmHiddenStatesGenerator wraps vLLM's engine to capture hidden states, and that it was written for an earlier vLLM API.
  4. The debugging context: The two errors that led to this message — the sample_tokens() state machine error and the ambiguous boolean tensor error — and their relationship to async scheduling.
  5. Distributed execution on 8 GPUs: The extraction runs with tensor parallelism across 8 GPUs, which adds complexity to the execution flow. The collective_rpc mechanism and worker coordination are relevant background.

Output Knowledge Created

This message creates several valuable pieces of knowledge:

  1. Root cause diagnosis: The connection between async scheduling and the two observed errors is now clearly articulated. Future developers encountering similar issues can immediately suspect async scheduling as the culprit.
  2. vLLM 0.16's built-in EAGLE-3 support: The discovery that execute_model_state already stores aux_hidden_states means that EAGLE-3 support is partially built into vLLM 0.16. This could simplify future integration efforts.
  3. A remediation strategy: The decision to disable async scheduling provides a concrete path forward. The grep command at the end sets up the next step — modifying the speculators code to pass async_scheduling=False.
  4. A deeper understanding of the system: The message demonstrates that the assistant now understands the relationship between vLLM's scheduling architecture and the hidden state capture mechanism. This understanding is essential for any further debugging.

The Thinking Process

The reasoning visible in this message reveals a methodical debugging approach. The assistant:

  1. Synthesizes: Connects the execute_model_state discovery from the previous investigation ([msg 2652]) with the aux_hidden_states storage mechanism, realizing they are part of the same system.
  2. Recognizes intent: Identifies that the aux_hidden_states field in execute_model_state is not incidental but represents "vLLM's built-in EAGLE-3 support" — a deliberate design choice by the vLLM team to support speculative decoding architectures.
  3. Traces causality: Maps the async scheduling behavior to both observed errors, showing how a single architectural change (two-phase execution) produces multiple symptoms.
  4. Evaluates options: Considers the complexity of adapting to the two-phase model versus the simplicity of disabling async scheduling, choosing the path of least resistance.
  5. Prepares for action: The grep command at the end is not just curiosity — it's reconnaissance for the next patch. The assistant is checking whether the speculators code already has any async scheduling awareness before writing the fix.

Conclusion

Message [msg 2653] is a masterclass in diagnostic reasoning. Faced with two seemingly unrelated errors, the assistant traces both back to a single root cause: vLLM 0.16's async scheduling architecture. The proposed fix — disabling async scheduling — is elegant in its simplicity, requiring only a configuration change rather than extensive code restructuring. This decision reflects a deep understanding of the system's architecture and a pragmatic assessment of the trade-offs involved. The message marks the moment when the debugging session pivots from investigation to resolution, setting the stage for the successful hidden state extraction that follows.