[CI] Fix mypy for vllm/attention and vllm/compilation

tools/pre_commit/mypy.py

@@ -26,6 +26,8 @@
 FILES = [
     "vllm/*.py",
     "vllm/assets",
+    "vllm/attention",
+    "vllm/compilation",
     "vllm/distributed",
     "vllm/entrypoints",
     "vllm/executor",
@@ -42,8 +44,6 @@
 # from "skip" to "silent", move the following directories to FILES
 SEPARATE_GROUPS = [
     "tests",
-    "vllm/attention",
-    "vllm/compilation",
     "vllm/engine",
     "vllm/inputs",
     "vllm/lora",

vllm/attention/layer.py

@@ -113,7 +113,9 @@ def maybe_get_vit_flash_attn_backend(
             if use_upstream_fa:
                 from flash_attn import flash_attn_varlen_func
             else:
-                from vllm.vllm_flash_attn import flash_attn_varlen_func
+                from vllm.vllm_flash_attn import (  # type: ignore[attr-defined]
+                    flash_attn_varlen_func,
+                )
     else:
         flash_attn_varlen_func = None
 
@@ -154,6 +156,7 @@ def __init__(
         `self.kv_cache`.
         """
         super().__init__()
+        sliding_window: int | None
         if per_layer_sliding_window is not None:
             # per-layer sliding window
             sliding_window = per_layer_sliding_window
@@ -407,7 +410,7 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
     def get_attn_backend(self) -> type[AttentionBackend]:
         return self.attn_backend
 
-    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
         # Block size may get updated after model loading, refresh it
         block_size = vllm_config.cache_config.block_size
         # Should not be called for enc-dec or encoder-only attention.

vllm/attention/layers/chunked_local_attention.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
-from typing import ClassVar
+from collections.abc import Hashable
+from typing import ClassVar, cast
 
 import torch
 
@@ -24,7 +25,7 @@
 
 @functools.lru_cache
 def create_chunked_local_attention_backend(
-    underlying_attn_backend: AttentionBackend,
+    underlying_attn_backend: type[AttentionBackend],
     attention_chunk_size: int,
     block_size: int,
 ) -> type[AttentionBackend]:
@@ -84,7 +85,9 @@ def __init__(
             )
 
             attn_backend = create_chunked_local_attention_backend(
-                underlying_attn_backend, attention_chunk_size, block_size
+                cast(Hashable, underlying_attn_backend),
+                attention_chunk_size,
+                block_size,
             )
         else:
             # in v0 the local attention is handled inside the backends

vllm/attention/layers/cross_attention.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
+from collections.abc import Hashable
 from copy import copy
+from typing import cast
 
 import numpy as np
 import torch
@@ -80,7 +82,7 @@ def _get_cross_slot_mapping(
 
 @functools.lru_cache
 def create_cross_attention_backend(
-    underlying_attn_backend: AttentionBackend,
+    underlying_attn_backend: type[AttentionBackend],
 ) -> type[AttentionBackend]:
     prefix = "CrossAttention_"
     underlying_builder = underlying_attn_backend.get_builder_cls()
@@ -155,7 +157,9 @@ def __init__(
                 head_size, dtype, kv_cache_dtype, block_size
             )
 
-            attn_backend = create_cross_attention_backend(underlying_attn_backend)
+            attn_backend = create_cross_attention_backend(
+                cast(Hashable, underlying_attn_backend)
+            )
         else:
             # in v0 cross attention is handled inside the backends
             attn_backend = None

vllm/attention/layers/encoder_only_attention.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
+from collections.abc import Hashable
 from copy import copy
+from typing import cast
 
 import torch
 
@@ -19,12 +21,11 @@
     CommonAttentionMetadata,
     subclass_attention_backend,
 )
-from vllm.v1.kv_cache_interface import KVCacheSpec
 
 
 @functools.lru_cache
 def create_encoder_only_attention_backend(
-    underlying_attn_backend: AttentionBackend,
+    underlying_attn_backend: type[AttentionBackend],
 ) -> type[AttentionBackend]:
     prefix = "EncoderOnlyAttention_"
     underlying_builder = underlying_attn_backend.get_builder_cls()
@@ -80,7 +81,7 @@ def __init__(
             )
 
             attn_backend = create_encoder_only_attention_backend(
-                underlying_attn_backend
+                cast(Hashable, underlying_attn_backend)
             )
         else:
             # in v0 encoder only attention is handled inside the backends
@@ -101,6 +102,6 @@ def __init__(
             **kwargs,
         )
 
-    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> None:
         # Does not need KV cache
         return None

vllm/attention/selector.py

@@ -166,11 +166,8 @@ def _cached_get_attn_backend(
     #
     # THIS SELECTION OVERRIDES THE VLLM_ATTENTION_BACKEND
     # ENVIRONMENT VARIABLE.
-    selected_backend = None
-    backend_by_global_setting: _Backend | None = get_global_forced_attn_backend()
-    if backend_by_global_setting is not None:
-        selected_backend = backend_by_global_setting
-    else:
+    selected_backend = get_global_forced_attn_backend()
+    if selected_backend is None:
         # Check the environment variable and override if specified
         backend_by_env_var: str | None = envs.VLLM_ATTENTION_BACKEND
         if backend_by_env_var is not None:

vllm/attention/utils/fa_utils.py

@@ -8,22 +8,22 @@
 logger = init_logger(__name__)
 
 if current_platform.is_cuda():
-    from vllm import _custom_ops as ops
+    from vllm import _custom_ops
 
-    reshape_and_cache_flash = ops.reshape_and_cache_flash
-    from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
+    reshape_and_cache_flash = _custom_ops.reshape_and_cache_flash
+    from vllm.vllm_flash_attn import (  # type: ignore[attr-defined]
+        flash_attn_varlen_func,
+        get_scheduler_metadata,
+    )
 elif current_platform.is_xpu():
-    from vllm._ipex_ops import ipex_ops as ops
+    from vllm._ipex_ops import ipex_ops
 
-    reshape_and_cache_flash = ops.reshape_and_cache_flash
-    flash_attn_varlen_func = ops.flash_attn_varlen_func
-    get_scheduler_metadata = ops.get_scheduler_metadata
+    reshape_and_cache_flash = ipex_ops.reshape_and_cache_flash
+    flash_attn_varlen_func = ipex_ops.flash_attn_varlen_func  # type: ignore[assignment]
+    get_scheduler_metadata = ipex_ops.get_scheduler_metadata
 
 
 def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
-    # import here to avoid circular dependencies
-    from vllm.platforms import current_platform
-
     if current_platform.is_xpu():
         return 2
     try:
@@ -74,24 +74,26 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
 
 
 def flash_attn_supports_fp8() -> bool:
+    device_capability = current_platform.get_device_capability()
     return (
         get_flash_attn_version() == 3
-        and current_platform.get_device_capability().major == 9
+        and device_capability is not None
+        and device_capability.major == 9
     )
 
 
-def flash_attn_supports_mla():
-    from vllm.platforms import current_platform
-
+def flash_attn_supports_mla() -> bool:
     if current_platform.is_cuda():
         try:
             from vllm.vllm_flash_attn.flash_attn_interface import (
                 is_fa_version_supported,
             )
 
+            device_capability = current_platform.get_device_capability()
             return (
                 is_fa_version_supported(3)
-                and current_platform.get_device_capability()[0] == 9
+                and device_capability is not None
+                and device_capability.major == 9
             )
         except (ImportError, AssertionError):
             pass

vllm/compilation/collective_fusion.py

@@ -50,7 +50,7 @@
 
 
 class BasePattern:
-    def __init__(self, dtype: torch.dtype, device: str):
+    def __init__(self, dtype: torch.dtype, device: torch.types.Device):
         self.dtype = dtype
         self.device = device
         self.tp = get_tp_group()
@@ -664,7 +664,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         allreduce_params: FlashInferFusedAllReduceParams,
     ):
         super().__init__(dtype, device)
@@ -719,7 +719,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         allreduce_params: FlashInferFusedAllReduceParams,
     ):
         super().__init__(dtype, device)
@@ -786,7 +786,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         allreduce_params: FlashInferFusedAllReduceParams,
     ):
         super().__init__(dtype, device)
@@ -855,7 +855,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         allreduce_params: FlashInferFusedAllReduceParams,
     ):
         super().__init__(dtype, device)
@@ -929,7 +929,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         allreduce_params: FlashInferFusedAllReduceParams,
     ):
         super().__init__(dtype, device)
@@ -1015,7 +1015,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         allreduce_params: FlashInferFusedAllReduceParams,
     ):
         super().__init__(dtype, device)

vllm/compilation/cuda_graph.py

@@ -125,6 +125,7 @@ def __call__(self, *args, **kwargs):
             # runtime modes.
             return self.runnable(*args, **kwargs)
 
+        assert batch_descriptor is not None
         if batch_descriptor not in self.concrete_cudagraph_entries:
             # create a new entry for this batch descriptor
             self.concrete_cudagraph_entries[batch_descriptor] = CUDAGraphEntry(

vllm/compilation/sequence_parallelism.py

@@ -108,7 +108,7 @@ def __init__(
         self,
         epsilon: float,
         dtype: torch.dtype,
-        device: str,
+        device: torch.types.Device,
         quant_op: torch._ops.OpOverload | None = None,
         **kwargs,
     ):
@@ -257,7 +257,11 @@ def replacement(
 
 class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
     def __init__(
-        self, epsilon: float, dtype: torch.dtype, device: str, op: torch._ops.OpOverload
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: torch.types.Device,
+        op: torch._ops.OpOverload,
     ):
         super().__init__(epsilon, dtype, device, quant_op=op)
 
@@ -313,7 +317,11 @@ def replacement(
 
 class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
     def __init__(
-        self, epsilon: float, dtype: torch.dtype, device: str, op: torch._ops.OpOverload
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: torch.types.Device,
+        op: torch._ops.OpOverload,
     ):
         super().__init__(epsilon, dtype, device, quant_op=op)
 
@@ -373,7 +381,11 @@ def replacement(
 
 class LastAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
     def __init__(
-        self, epsilon: float, dtype: torch.dtype, device: str, op: torch._ops.OpOverload
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: torch.types.Device,
+        op: torch._ops.OpOverload,
     ):
         super().__init__(epsilon, dtype, device, quant_op=op)
 

vllm/compilation/vllm_inductor_pass.py

@@ -43,7 +43,7 @@ def __init__(self, config: VllmConfig):
         )
         self.pass_config = config.compilation_config.pass_config
         self.model_dtype = config.model_config.dtype if config.model_config else None
-        self.device = config.device_config.device if config.device_config else None
+        self.device = config.device_config.device
         self.pass_name = self.__class__.__name__
 
     @staticmethod