biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix] dtype mismatch in ngram gpu propose (#37246)

Browse Source 
Signed-off-by: PatchouliTaisa <patchychen@tencent.com>
Co-authored-by: PatchouliTaisa <patchychen@tencent.com>
This commit is contained in:
PatchyTIS
2026-03-17 13:19:55 +08:00
committed by GitHub
parent 3e3d320c1b
commit 17c1bdf371
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/v1/spec_decode/ngram_proposer_gpu.py
View File

@@ -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