reformat

pytorch-labs · Jan 15, 2025 · 854c968 · 854c968
1 parent c7413e3
commit 854c968
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/attn_gym/paged_attention/latency.py b/attn_gym/paged_attention/latency.py
@@ -58,7 +58,7 @@ def benchmark_layer(
             lambda: paged_foo(batch_idx, input_pos, x, converted_block_mask, converted_score_mod)
         )
         print(
-            f"non_paged_latency: {non_paged_latency} ms, paged_latency: {paged_latency} ms, overhead: {round((paged_latency/non_paged_latency-1.0)*100, 2)}%"
+            f"non_paged_latency: {non_paged_latency} ms, paged_latency: {paged_latency} ms, overhead: {round((paged_latency / non_paged_latency - 1.0) * 100, 2)}%"
         )
 
 

diff --git a/attn_gym/utils.py b/attn_gym/utils.py
@@ -100,9 +100,9 @@ def visualize_attention_scores(
     Returns:
         None
     """
-    assert (
-        score_mod is not None or mask_mod is not None
-    ), "Must provide either score_mod or mask_mod"
+    assert score_mod is not None or mask_mod is not None, (
+        "Must provide either score_mod or mask_mod"
+    )
     query = query[batch_idx, head_idx, :, :]
     key = key[batch_idx, head_idx, :, :]
     scores_viz = create_score_mod(
-Original file line number
+Diff line change
@@ Expand Up / @@ -58,7 +58,7 @@ def benchmark_layer( @@
                 lambda: paged_foo(batch_idx, input_pos, x, converted_block_mask, converted_score_mod)
             )
             print(
-                f"non_paged_latency: {non_paged_latency} ms, paged_latency: {paged_latency} ms, overhead: {round((paged_latency/non_paged_latency-1.0)*100, 2)}%"
+                f"non_paged_latency: {non_paged_latency} ms, paged_latency: {paged_latency} ms, overhead: {round((paged_latency / non_paged_latency - 1.0) * 100, 2)}%"
             )
@@ Expand Down @@