[Bugfix][Spec Decode] Fix SWA block freeing during async scheduling

vllm/model_executor/layers/attention/attention.py

@@ -658,6 +658,11 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
                 dtype=self.kv_cache_torch_dtype,
                 kv_quant_mode=quant_mode,
                 sliding_window=self.sliding_window,
+                num_speculative_tokens=(
+                    vllm_config.speculative_config.num_speculative_tokens
+                    if vllm_config.speculative_config
+                    else 0
+                ),
                 page_size_padded=shared_page,
             )
         elif self.kv_cache_dtype.startswith("turboquant_"):

vllm/v1/core/single_type_kv_cache_manager.py

@@ -670,6 +670,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
     def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None:
         super().__init__(kv_cache_spec, **kwargs)
         self.sliding_window = kv_cache_spec.sliding_window
+        self.num_speculative_tokens = kv_cache_spec.num_speculative_tokens
 
     @classmethod
     def _contiguous_blocks_for_hit(
@@ -861,7 +862,15 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
         Returns:
             The number of tokens that will be skipped for attention computation.
         """
-        return max(0, num_computed_tokens - self.sliding_window + 1)
+        # NOTE: During async scheduling, num_computed_tokens contains the number of
+        # tokens drafted during the previous step, some of which may be rejected later.
+        # To avoid overestimating the sequence length (and freeing blocks that are
+        # actually needed) num_speculative_tokens is subtracted from
+        # num_computed_tokens.
+        return max(
+            0,
+            num_computed_tokens - self.num_speculative_tokens - self.sliding_window + 1,
+        )
 
     def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
         """

vllm/v1/kv_cache_interface.py

@@ -518,6 +518,7 @@ def is_uniform_with_collection(
 class SlidingWindowSpec(AttentionSpec):
     sliding_window: int
     head_size_v: int = None  # type: ignore[assignment]
+    num_speculative_tokens: int = 0
 
     def __post_init__(self):
         if self.head_size_v is None: