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 6, 14:30 UTC
[Bugfix] dtype mismatch in ngram gpu propose (#37246)
Signed-off-by: PatchouliTaisa <patchychen@tencent.com> Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
PatchyTIS committed
Mar 17, 2026 at 13:19 UTC
17c1bdf3719d9d8fdf4f13cb1468e5ed5f70d021
1 file changed
+4
-2
vllm/v1/spec_decode/ngram_proposer_gpu.py
+4
-2
@@ -364,7 +364,9 @@ class NgramProposerGPU:
364
)
365
token_ids_gpu.scatter_(1, write_positions_long, tokens_to_scatter)
366
367
- num_tokens_tmp = num_tokens_no_spec + valid_sampled_tokens_count
367
+ num_tokens_tmp = (num_tokens_no_spec + valid_sampled_tokens_count).to(
368
+ torch.int32
369
+ )
370
371
# Compute validity masks.
372
sampled_flags = valid_sampled_tokens_count > 0
@@ -437,7 +439,7 @@ class NgramProposerGPU:
439
)
440
441
# Count valid tokens per request.
440
- valid_sampled_tokens_count = valid_mask.sum(dim=1)
442
+ valid_sampled_tokens_count = valid_mask.sum(dim=1).to(torch.int32)
443
444
# Rightmost valid index per row.
445
last_valid_indices = valid_sampled_tokens_count - 1