[Core]fix type annotation for `swap_blocks` (#4726)

64b77dfd · Kunshang Ji · GitHub · 51d4094f · 64b77dfd
Unverified Commit 64b77dfd authored May 10, 2024 by Kunshang Ji Committed by GitHub May 10, 2024
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

vllm/_custom_ops.py vllm/_custom_ops.py +2 -2

No files found.
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
-from typing import Dict, Optional, Tuple
+from typing import Optional, Tuple

 import torch

@@ -266,7 +266,7 @@ def copy_blocks(key_caches: torch.Tensor, value_caches: torch.Tensor,


 def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
-                block_mapping: Dict[int, int]) -> None:
+                block_mapping: torch.Tensor) -> None:
    vllm_cache_ops.swap_blocks(src, dst, block_mapping)