fix mypy for worker folder

Signed-off-by: Tobias Pitters <tobias.pitters@gmail.com>
vllm-project · Jan 1, 2025 · a265c6c · a265c6c
1 parent 5d7ced1
commit a265c6c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
@@ -1048,14 +1048,14 @@ def __init__(
         needs_attn_backend = (num_attn_heads != 0
                               or self.model_config.is_attention_free)
 
-        self.attn_backend: AttentionBackend = get_attn_backend(
+        self.attn_backend: Optional[AttentionBackend] = get_attn_backend(
             self.model_config.get_head_size(),
             self.model_config.dtype,
             self.kv_cache_dtype,
             self.block_size,
             self.model_config.is_attention_free,
         ) if needs_attn_backend else None
-        if self.attn_backend:
+        if self.attn_backend is not None:
             self.attn_state = self.attn_backend.get_state_cls()(
                 weakref.proxy(self))
         else:
@@ -1481,6 +1481,7 @@ def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
                         )
                         self.set_active_prompt_adapters(
                             set(), prompt_adapter_mapping)
+                    assert isinstance(self.attn_backend, AttentionBackend)
                     graph_runner = CUDAGraphRunner(
                         self.model, self.attn_backend.get_name(),
                         self.attn_state.graph_clone(batch_size),

diff --git a/vllm/worker/multi_step_model_runner.py b/vllm/worker/multi_step_model_runner.py
@@ -313,6 +313,8 @@ def __init__(self, base_model_runner: GPUModelRunnerBase, *args, **kwargs):
         supported_attention_backends: List[str] = \
             _get_supported_attention_backends(
                 self.scheduler_config.chunked_prefill_enabled)
+
+        assert isinstance(self.attn_backend, AttentionBackend)
         if self.attn_backend.get_name() not in supported_attention_backends:
             ms_config_str: str = "Multi-Step + Chunked-Prefill" \
                 if self.scheduler_config.chunked_prefill_enabled \