format

Signed-off-by: Lucas Wilkinson <[email protected]>

Author: LucasWilkinson

tests/v1/attention/test_attention_backends.py

@@ -9,10 +9,10 @@
                                       create_common_attn_metadata,
                                       create_standard_kv_cache_spec,
                                       create_vllm_config,
-                                      get_attention_backend,)
+                                      get_attention_backend)
 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata, set_kv_cache_layout)
+from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
+                                              set_kv_cache_layout)
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 BACKENDS_TO_TEST = [
@@ -421,8 +421,9 @@ def test_backend_correctness(batch_spec_name: str, model: str):
         if backend_name == _Backend.FLASHINFER_VLLM_V1:
             kv_cache_for_backend = kv_cache.transpose(0, 1)
 
-            # For FlashInfer default to HND layout and 
-            kv_cache_for_backend = kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
+            # For FlashInfer default to HND layout and
+            kv_cache_for_backend = kv_cache_for_backend.transpose(
+                2, 3).contiguous().transpose(2, 3)
             set_kv_cache_layout("HND")
 
         backend_output = run_attention_backend(backend_name, kv_cache_spec,

vllm/v1/attention/backends/flashinfer.py

@@ -7,12 +7,12 @@
 from typing import TYPE_CHECKING, Optional
 
 import torch
+
+import vllm.envs as envs
 from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
                         BatchPrefillWithPagedKVCacheWrapper,
                         MultiLevelCascadeAttentionWrapper)
 from flashinfer.decode import trtllm_batch_decode_with_kv_cache
-
-import vllm.envs as envs
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionType)
 from vllm.config import VllmConfig
@@ -673,7 +673,6 @@ def forward(
                     assert block_tables_decode.is_contiguous()
                     assert seq_lens_decode.is_contiguous()
 
-
                     output[:num_decode_tokens] = (
                         trtllm_batch_decode_with_kv_cache(
                             query=decode_query,