diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index 3ff841804..eb24a9c93 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -364,7 +364,9 @@ class NgramProposerGPU: ) token_ids_gpu.scatter_(1, write_positions_long, tokens_to_scatter) - num_tokens_tmp = num_tokens_no_spec + valid_sampled_tokens_count + num_tokens_tmp = (num_tokens_no_spec + valid_sampled_tokens_count).to( + torch.int32 + ) # Compute validity masks. sampled_flags = valid_sampled_tokens_count > 0 @@ -437,7 +439,7 @@ class NgramProposerGPU: ) # Count valid tokens per request. - valid_sampled_tokens_count = valid_mask.sum(dim=1) + valid_sampled_tokens_count = valid_mask.sum(dim=1).to(torch.int32) # Rightmost valid index per row. last_valid_indices = valid_sampled_tokens_count - 1