[Bug]: IndexError: list index out of range on chunked prefill with speculative decoding

[saidrhs]

Environment

Engine Configuration:

V1 LLM engine (v0.9.1) with config:
- model='/tmp/model/'
- speculative_config=SpeculativeConfig(method='eagle3', model='/tmp/model/eagle_head/', num_spec_tokens=5)
- tensor_parallel_size=8
- pipeline_parallel_size=1
- quantization=compressed-tensors
- max_seq_len=131072

Prefix caching and chunked prefill are enabled by default V1 behavior. This issue also occurred in vLLM v0.8.5.post1 and has been hard to reproduce.

Model: Llama-3.3-70B-Instruct

Hardware: 8 H200 GPUs

🐛 Describe the bug

IndexError with chunked prefill and speculative decoding.

vllm/vllm/v1/worker/gpu_model_runner.py

Line 1612 in 9fb52e5

# Partial prefill (rare case).

Error Logs

Worker Stack Trace (across multiple ranks):

IndexError: list index out of range

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/executor/multiproc_executor.py", line 522, in worker_busy_loop
    output = func(*args, **kwargs)
  
  File "/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/worker/gpu_worker.py", line 293, in execute_model
    output = self.model_runner.execute_model(scheduler_output)
  
  File "/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/worker/gpu_model_runner.py", line 1428, in execute_model
    next_token_id = req_state.get_token_id(seq_len)
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/worker/gpu_input_batch.py", line 53, in get_token_id
    return self.output_token_ids[idx - self.num_prompt_tokens]

Engine Core Stack Trace:

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core.py", line 508, in run_engine_core
    engine_core.run_busy_loop()
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core.py", line 535, in run_busy_loop
    self._process_engine_step()
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core.py", line 560, in _process_engine_step
    outputs, model_executed = self.step_fn()
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core.py", line 231, in step
    model_output = self.execute_model(scheduler_output)
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core.py", line 217, in execute_model
    raise err
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core.py", line 211, in execute_model
    return self.model_executor.execute_model(scheduler_output)
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/executor/multiproc_executor.py", line 163, in execute_model
    (output, ) = self.collective_rpc("execute_model")
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/executor/multiproc_executor.py", line 220, in collective_rpc
    result = get_response(w, dequeue_timeout)
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/executor/multiproc_executor.py", line 207, in get_response
    raise RuntimeError("Worker failed with error 'list index out of range'")

AsyncLLM Stack Trace:

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/async_llm.py", line 379, in output_handler
    outputs = await engine_core.get_output_async()
  
  File "/usr/local/lib/python3.11/dist-packages/vllm/v1/engine/core_client.py", line 790, in get_output_async
    raise self._format_exception(outputs) from None

Final Error:

vllm.v1.engine.exceptions.EngineDeadError: EngineCore encountered an issue. See stack trace (above) for the root cause.

Scheduler State at Time of Error:

SchedulerOutput(
    scheduled_new_reqs=[],
    scheduled_cached_reqs=[
        CachedRequestData(
            req_id='f1f4e063-786b-4cf9-8ebc-548f201ed419',
            resumed_from_preemption=false,
            new_token_ids=[t1],
            new_block_ids=[[]],
            num_computed_tokens=t2
        ),
        CachedRequestData(
            req_id='93ae9bbb-f390-4f02-9d16-ddcba4e3b6d2',
            resumed_from_preemption=false,
            new_token_ids=[t3],
            new_block_ids=[[t4]],
            num_computed_tokens=t5
        )
    ],
    num_scheduled_tokens={
        'f1f4e063-786b-4cf9-8ebc-548f201ed419': 6,
        '93ae9bbb-f390-4f02-9d16-ddcba4e3b6d2': 6
    },
    total_num_scheduled_tokens=12,
    scheduled_spec_decode_tokens={
        'f1f4e063-786b-4cf9-8ebc-548f201ed419': [t6, t7, t8, t9, t10],
        '93ae9bbb-f390-4f02-9d16-ddcba4e3b6d2': [t11, t12, t13, t14, t15]
    }
)

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[chaunceyjiang]

I don't seem to be able to reproduce the issue locally.
Could you share your OpenAI client code?

[saidrhs]

This bug has been challenging to reproduce due to limited client visibility - restricted for security/privacy reasons. The issue persists since v0.8.5.post1, with no consistent traffic pattern identified. All observed in-flight requests in the latest crash have input token counts ≤2100.

Based on the stack trace, three main factors contribute to this issue:

1. Chunked prefill implementation with speculative decoding
2. Empty valid_sampled_token_ids array
3. Token ID boundary violation:
   • Request ID: 93ae9bbb-f390-4f02-9d16-ddcba4e3b6d2
   • Error: TokenId 2236 exceeds output token bounds
   • Potential cause: Optimistic scheduling introduced in [V1][PP] Optimization: continue scheduling prefill chunks #17080

[renjie0]

We are observing similar issue

[gnovack]

@renjie0 is there any more information about your environment that you can share? (e.g. vLLM version, model, spec decoding settings, quantization settings, etc.)?

[ekagra-ranjan]

@saidrhs - could you share how you retrieve the Scheduler State at Time of Error (i believe in production)?

2. Empty valid_sampled_token_ids array

I couldnt relate valid_sampled_token_ids to the issue. The error appears to be in partial prefill which will not have valid_sampled_token_ids which is why calculate the seq_len as (req_state.num_computed_tokens + scheduler_output.num_scheduled_tokens[req_id]).

My current understanding is that since its a partial prefill, the call to get_token_id() should have gone to this if condition but instead goes to the else condition which is supposed to give the error.

[renjie0]

@renjie0 is there any more information about your environment that you can share? (e.g. vLLM version, model, spec decoding settings, quantization settings, etc.)?

We are using vLLM v0.9.1, model is Cohere Command A, spec decoding settings: --speculative_config={"method": "eagle", "model": "$MODEL_PATH/draft-model", "num_speculative_tokens": 3, "draft_tensor_parallel_size": $TENSOR_PARALLEL_SIZE}', quantization: FP8

[ekagra-ranjan]

Another inconsistency I see is that the error is coming from partial prefill (chunk context) case of EAGLE. However, the Scheduler stats shared above show

scheduled_spec_decode_tokens={
        'f1f4e063-786b-4cf9-8ebc-548f201ed419': [t6, t7, t8, t9, t10],
        '93ae9bbb-f390-4f02-9d16-ddcba4e3b6d2': [t11, t12, t13, t14, t15]
    }

which means the 2 req are already in generation phase and the code should not have hit the partial prefill code.

[gnovack]

Thanks for the info @renjie0, we are also seeing this when using Eagle and FP8 quant. In our case, this error appears to coincide with nan values appearing in the model hidden states.

This causes the sampler to sample an OOV token id, then leads to an empty valid_sampled_token_ids array even when not performing a partial prefill.

[ekagra-ranjan]

This causes the sampler to sample an OOV token id, then leads to an empty valid_sampled_token_ids array even when not performing a partial prefill.

@gnovack -

1. Could you share how you narrowed down to this hypothesis?
2. Is this the code pointer to above?
3. IIUC, you mean that if K=3, then we pass 4 tokens (K+1) to Rejection sampler and the logit after doing a fwd pass from the 1st token is nan so we get 0 sampled token from Rejection Sampler after processing?
4. This means that the same target model without EAGLE would have sampled OOV and returned it in the output (or possibly caused error during de-tokenization), right? So the error is not fixable unless target model's quant is improved?

[KungYork]

I’ve observed that the issue only occurs when both CUDA Graphs and torch.compile are enabled at the same time—disabling either one makes the problem disappear. Could you confirm whether you’re seeing the same behavior?

[ekagra-ranjan]

@KungYork - can you share how do you reproduce this error in order to check if disabling CUDA Graphs or torch.compile helped you?

[gnovack]

@ekagra-ranjan - putting the answers to your questions below. hope this helps!

1. Could you share how you narrowed down to this hypothesis?

Lots of trial & error / adding logs / working back from the line that throws the error.

2. Is this the code pointer to above?

This is one of the key points in the code; however, at this point, the sampled_token_ids list is already in a problematic state. Specifically, in my case (using Llama 3.1 8B w/ vocab size 128256 and Spec Length 5), the sampled_token_ids list is [128256, -1, -1, -1, -1, -1]. IIUC, this means that all 5 draft tokens were rejected and the token sampled from the base model is 128256, which is OOV.

3. IIUC, you mean that if K=3, then we pass 4 tokens (K+1) to Rejection sampler and the logit after doing a fwd pass from the 1st token is nan so we get 0 sampled token from Rejection Sampler after processing?

Continuing my answer from (2) above, in my case the rejection sampler is seeing a list of sampled token ids like [128256, -1, -1, -1, -1, -1]. These all get masked out by this valid_mask, since all of the ids are either -1 or >= vocab_size.

4. This means that the same target model without EAGLE would have sampled OOV and returned it in the output (or possibly caused error during de-tokenization), right? So the error is not fixable unless target model's quant is improved?

Potentially. However, I have yet to see this issue occur when using the base model w/o any spec decoding. I am not confident either way yet as to whether this is entirely a base model problem.

[gnovack]

After some further thinking on this, I'm not convinced that the presence of NaNs in the model output is necessarily a new phenomenon (see this recent PR to add NaN presence as a metric); however, the fact that these NaNs can cause the engine to crash does seem to be something new.

When running w/o spec decoding enabled and using the native sampler, hidden state NaNs cause token 0 to be sampled.

i.e.

tensor([nan, ...]).argmax() --> tensor(0)

As many tokenizers map token id 0 to the string !, this can explain the many reported issues related to models producing exclamation marks: https://github.com/search?q=repo%3Avllm-project%2Fvllm+exclamation&type=issues

So why does the engine crash when spec decoding is enabled?

There appears to be a difference in how the native Sampler vs. the RejectionSampler handle NaNs. The Sampler outputs 0 tokens, while the RejectionSampler outputs OOV tokens. This might be because sample_recovered_tokens_kernel kernel here pads the vocab size with -inf to the next power of 2, and tl.argmax([nan, nan, -inf, ...]) --> 2

If we assume that NaNs are always going to be possible, then I think the right solution to this issue right now is to ensure that the rejection sampler behaves the same as the native base model sampler (i.e. by generating in in-vocab token id when NaNs appear in the model output)