format

vllm-project · Sep 19, 2024 · 0838a3f · 0838a3f
1 parent 0d0c331
commit 0838a3f
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/kernels/test_attention.py b/tests/kernels/test_attention.py
@@ -247,7 +247,8 @@ def test_paged_attention(
                      key_cache, value_cache, num_kv_heads, scale, block_tables,
                      seq_lens, block_size, max_seq_len, alibi_slopes,
                      kv_cache_dtype, k_scale, v_scale, 0, 0, 0, 64, 0),
-                    cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]))
+                    cond=(head_size == HEAD_SIZES[0]
+                          and block_size == BLOCK_SIZES[0]))
 
         else:
             ops.paged_attention_rocm(
@@ -275,7 +276,8 @@ def test_paged_attention(
                      key_cache, value_cache, num_kv_heads, scale, block_tables,
                      seq_lens, block_size, max_seq_len, alibi_slopes,
                      kv_cache_dtype, k_scale, v_scale),
-                    cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]))
+                    cond=(head_size == HEAD_SIZES[0]
+                          and block_size == BLOCK_SIZES[0]))
 
     else:
         raise AssertionError(f"Unknown version: {version}")