diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index c6a0632c2d..cf6ab0ac06 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -246,7 +246,7 @@ class MambaHybridModelState(DefaultModelState): # compute them during actual (non-capture) forward execution. num_accepted_tokens = None num_decode_draft_tokens_cpu = None - if not for_capture: + if not for_capture and self.vllm_config.num_speculative_tokens > 0: num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs) num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[ input_batch.idx_mapping