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 5, 04:45 UTC
[Hybrid] Simplify accepted token counting in spec decode for hybrid models (#38372)
Francesco Fusco committed
Apr 15, 2026 at 00:19 UTC
507df79a29ac8733ea03c624d215ccfffad7bbe9
1 file changed
+4
-19
vllm/v1/worker/gpu_model_runner.py
+4
-19
@@ -1428,26 +1428,11 @@ class GPUModelRunner(
1428
# TODO: Remove .cpu() sync to enable fully async for hybrid model;
1429
# Use num_computed_tokens.gpu instead of req.num_computed_tokens to
1430
# support aligned mamba cache mode.
1431
- # Find the number of accepted tokens for each sequence.
1431
+ # Count the number of accepted tokens for each sequence.
1432
+ # Valid tokens are contiguous from position 0, so counting non-(-1)
1433
+ # tokens gives us the first -1 position (i.e., number of accepted).
1434
num_reqs = output_token_ids.size(0)
1433
- self.num_accepted_tokens.gpu[:num_reqs] = (
1434
- (
1435
- torch.cat(
1436
- [
1437
- output_token_ids,
1438
- torch.full(
1439
- (num_reqs, 1),
1440
- -1,
1441
- device=output_token_ids.device,
1442
- ),
1443
- ],
1444
- dim=1,
1445
- )
1446
- == -1
1447
- )
1448
- .int()
1449
- .argmax(-1)
1450
- )
1435
+ self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1)
1436
1437
if self.cache_config.mamba_cache_mode == "align":
1438
for i, num_tokens in enumerate(