[Fix] Fixes speculative decode indexing and unpad condition for attention metadata (#5626)

### What this PR does / why we need it?
This addresses the issue brought up by #5356 and #4963, and we believe
the unnecessary conditions are the root cause.

Change the unpad trigger to be driven by actual size mismatches
(num_reqs vs base_num_reqs or scheduled vs input token counts) rather
than specific speculative-method flags. Then remove brittle workarounds
that forced request counts and sliced query start locations.

This prevents incorrect indexing and length mismatches during
speculative decoding and makes metadata unpadding more robust across
scheduling modes.

### Does this PR introduce _any_ user-facing change?
None.

### How was this patch tested?
Tested by existing cases.

- vLLM version: v0.13.0
- vLLM main:
8be6432bda

---------

Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>
This commit is contained in:
Yizhou
2026-01-08 19:41:08 +08:00
committed by GitHub
parent 503822c56c
commit f4605c2b3c
3 changed files with 9 additions and 15 deletions

View File

@@ -749,13 +749,6 @@ class EagleProposer(VllmEagleProposer):
num_reqs = common_attn_metadata.num_reqs
device = valid_sampled_tokens_count.device
if num_reqs != spec_decode_metadata.cu_num_draft_tokens.shape[0]:
# TODO: This is a serious issue and should be taken care of ASAP
# In short, why input_batch.num_reqs != attn_metadata.num_reqs?
# Previously in #4963, we modified `query_start_loc`, but this
# problem remains unsolved.
num_reqs = spec_decode_metadata.cu_num_draft_tokens.shape[0]
token_indices_to_sample = torch.empty((num_reqs, ),
dtype=torch.int32,
device=device)
@@ -787,9 +780,9 @@ class EagleProposer(VllmEagleProposer):
torch.zeros_like(num_draft_tokens_gpu),
)
query_start_loc = common_attn_metadata.query_start_loc[
1:1 + num_rejected_tokens_gpu.shape[0]]
token_indices_to_sample = query_start_loc - 1 - num_rejected_tokens_gpu
token_indices_to_sample = (
common_attn_metadata.query_start_loc[1:] - 1 -
num_rejected_tokens_gpu)
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu