address pr feedback

NVIDIA · Dec 11, 2024 · 519222b · 519222b
1 parent 8eb3b3d
commit 519222b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/kvpress/presses/composed_press.py b/kvpress/presses/composed_press.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+
 from kvpress.presses.base_press import BasePress
 
 

diff --git a/tests/integration/test_ruler.py b/tests/integration/test_ruler.py
@@ -4,7 +4,15 @@
 from transformers import DynamicCache, QuantizedCacheConfig, QuantoQuantizedCache
 from transformers.utils import is_flash_attn_2_available, is_optimum_quanto_available
 
-from kvpress import ExpectedAttentionPress, KnormPress, SnapKVPress, StreamingLLMPress, ThinKPress, TOVAPress
+from kvpress import (
+    ExpectedAttentionPress,
+    KnormPress,
+    SimLayerKVPress,
+    SnapKVPress,
+    StreamingLLMPress,
+    ThinKPress,
+    TOVAPress,
+)
 from tests.fixtures import kv_press_llama3_1_flash_attn_pipeline  # noqa: F401
 
 
@@ -18,7 +26,7 @@ def df_ruler():
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU is not available")
 @pytest.mark.skipif(not is_flash_attn_2_available(), reason="flash_attn is not installed")
 @pytest.mark.parametrize(
-    "cls", [KnormPress, ExpectedAttentionPress, StreamingLLMPress, SnapKVPress, TOVAPress, ThinKPress]
+    "cls", [KnormPress, ExpectedAttentionPress, StreamingLLMPress, SnapKVPress, TOVAPress, ThinKPress, SimLayerKVPress]
 )
 @pytest.mark.parametrize("compression_ratio", [0.1, 0.2])
 @pytest.mark.parametrize("cache", ["dynamic", "quantized"])

diff --git a/tests/presses/test_presses.py b/tests/presses/test_presses.py
@@ -12,8 +12,8 @@
     KnormPress,
     ObservedAttentionPress,
     RandomPress,
-    SnapKVPress,
     SimLayerKVPress,
+    SnapKVPress,
     StreamingLLMPress,
     TOVAPress,
 )
Original file line number	Diff line number	Diff line change
		@@ -1,4 +1,5 @@
		from dataclasses import dataclass

		from kvpress.presses.base_press import BasePress


Expand Down