A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix][Spec Decode] Fix wrong valid_mask for padded speculation when chunked prefill occurs (#26231)
Signed-off-by: seven-mile <i@7li.moe> Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
7mile committed
Oct 7, 2025 at 02:24 UTC
b2ea5ba6776fdf7343f24acfee9de8604e1d5d6d
1 file changed
+3
-7
vllm/v1/spec_decode/eagle.py
+3
-7
@@ -522,13 +522,9 @@ class EagleProposer:
522
)
523
524
# Generate a mask for all valid tokens within those requests
525
- max_gen_len = sampled_token_ids.shape[-1]
526
- if max_gen_len == 1:
527
- valid_mask = torch.ones_like(valid_sampled_token_ids_gpu, dtype=torch.bool)
528
- else:
529
- valid_mask = (valid_sampled_token_ids_gpu != -1) & (
530
- valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
531
- )
525
+ valid_mask = (valid_sampled_token_ids_gpu != -1) & (
526
+ valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
527
+ )
528
529
# Count the number of valid tokens in each request
530
valid_sampled_tokens_count = valid_mask.sum(dim=1)