enfore all tests to run

Signed-off-by: Max Jeblick <maximilianjeblick@gmail.com>

Author: maxjeblick

.github/PULL_REQUEST_TEMPLATE.md

@@ -4,6 +4,8 @@ Description of your PR. Fixes # (issue) (if applicable)
 
 ## Checklist
 
+Before submitting a PR, please make sure:
+
 - Tests are working (`make test`)
 - Code is formatted correctly (`make style`, on errors try fix with `make format`)
 - Copyright header is included

Makefile

@@ -42,10 +42,16 @@ reports:
 .PHONY: test
 test: reports
 	$(UV) add optimum-quanto
-	$(UV) add flash-attn --no-build-isolation
+	$(UV) add flash-attn
 	PYTHONPATH=. \
 	$(UV) run pytest \
 		--cov-report xml:reports/coverage.xml \
 		--cov=kvpress/ \
 		--junitxml=./reports/junit.xml \
-		tests/
+		-v \
+		tests/ | tee reports/pytest_output.log
+	@if grep -q "SKIPPED" reports/pytest_output.log; then \
+		echo "Error: Tests were skipped. All tests must run."; \
+		grep "SKIPPED" reports/pytest_output.log; \
+		exit 1; \
+	fi

pyproject.toml

@@ -22,6 +22,9 @@ dependencies = [
     "accelerate>=1.0.0,<2",
     "requests>=2.32.3,<3",
     "cachetools>=5.5.2,<6",
+    "optimum-quanto>=0.2.7",
+    "hatch>=1.14.1",
+    "flash-attn>=2.8.2",
 ]
 
 [project.optional-dependencies]
@@ -89,4 +92,4 @@ disable_error_code = ["attr-defined"]
 
 [[tool.mypy.overrides]]
 module = "kvpress.pipeline"
-disable_error_code = ["attr-defined", "assignment", "override"]
+disable_error_code = ["attr-defined", "assignment", "override"]

tests/integration/test_ruler.py

@@ -27,12 +27,12 @@ def test_ruler_is_correct(kv_press_llama3_1_flash_attn_pipeline, df_ruler, press
     kwargs = press_dict["kwargs"][0]
     press = cls(**kwargs)
     if not hasattr(cls, "compression_ratio"):
-        pytest.skip(reason="Press does not support compression_ratio")
+        return  # "Press does not support compression_ratio"
     # set compression ratio to a small value for testing
     try:
         press.compression_ratio = 0.1
     except AttributeError:
-        pytest.skip(reason="Press does not support setting compression_ratio")
+        return  # "Press does not support setting compression_ratio"
 
     if cache == "dynamic":
         cache = DynamicCache()

tests/presses/test_block_press.py

@@ -33,7 +33,7 @@ def test_block_press_is_streaming_top_k(unit_test_model):  # noqa: F811
     """
     press = HiddenStatesPress(compression_ratio=0.5)
     generator = torch.Generator().manual_seed(0)
-    input_ids = torch.randint(0, 1024, (1, 256), generator=generator)
+    input_ids = torch.randint(0, 1024, (1, 256), generator=generator).to(unit_test_model.device)
     keys_hash = []
     values_hash = []
 

tests/presses/test_finch_press.py

@@ -16,6 +16,6 @@ def test_finch_press(unit_test_model):  # noqa: F811
     ]:
         press.delimiter_token_id = unit_test_model.config.eos_token_id
         with press(unit_test_model):
-            input_ids = torch.arange(10, 20)
+            input_ids = torch.arange(10, 20).to(unit_test_model.device)
             input_ids[8] = press.delimiter_token_id
             unit_test_model(input_ids.unsqueeze(0))

tests/presses/test_head_compression.py

@@ -28,7 +28,7 @@ def test_wrapper_head_compression(unit_test_model, wrapper_press, compression_ra
     p = KnormPress(compression_ratio=compression_ratio)
     press = wrapper_press(press=p)
     with press(unit_test_model):
-        input_ids = torch.randint(0, 1024, (1, 128))
+        input_ids = torch.randint(0, 1024, (1, 128)).to(unit_test_model.device)
         unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values
 
     assert unit_test_model.model.layers[0].self_attn.masked_key_indices is not None
@@ -47,7 +47,7 @@ def test_wrapper_head_compression(unit_test_model, wrapper_press, compression_ra
 def test_head_compression(unit_test_model, press, compression_ratio, layerwise):  # noqa: F811
     press = KVzipPress(compression_ratio=compression_ratio, layerwise=layerwise)
     with press(unit_test_model):
-        input_ids = torch.randint(0, 1024, (1, 128))
+        input_ids = torch.randint(0, 1024, (1, 128)).to(unit_test_model.device)
         unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values
 
     assert unit_test_model.model.layers[0].self_attn.masked_key_indices is not None

tests/presses/test_observed_attention_press.py

@@ -13,18 +13,18 @@
 
 @torch.no_grad()
 def test_observed_drops_attention_output(unit_test_model, unit_test_model_output_attention, caplog):  # noqa: F811
-    input_ids = unit_test_model.dummy_inputs["input_ids"]
+    input_ids = unit_test_model.dummy_inputs["input_ids"].to(unit_test_model.device)
     output = unit_test_model(input_ids, past_key_values=DynamicCache())
     assert output.attentions is None
 
-    input_ids = unit_test_model_output_attention.dummy_inputs["input_ids"]
+    input_ids = unit_test_model_output_attention.dummy_inputs["input_ids"].to(unit_test_model.device)
     attentions = unit_test_model_output_attention(input_ids, past_key_values=DynamicCache()).attentions
     assert all([isinstance(attention, torch.Tensor) for attention in attentions])
 
     with caplog.at_level(logging.DEBUG):
         press = ObservedAttentionPress(compression_ratio=0.4)
         with press(unit_test_model_output_attention):
-            input_ids = unit_test_model_output_attention.dummy_inputs["input_ids"]
+            input_ids = unit_test_model_output_attention.dummy_inputs["input_ids"].to(unit_test_model.device)
             output = unit_test_model_output_attention(input_ids, past_key_values=DynamicCache())
 
             # There's a slight mismatch in outputs when using a model that has output_attentions=True
@@ -36,7 +36,7 @@ def test_observed_drops_attention_output(unit_test_model, unit_test_model_output
 
     press = ObservedAttentionPress(compression_ratio=0.4, output_attentions=True)
     with press(unit_test_model_output_attention):
-        input_ids = unit_test_model_output_attention.dummy_inputs["input_ids"]
+        input_ids = unit_test_model_output_attention.dummy_inputs["input_ids"].to(unit_test_model.device)
         output = unit_test_model_output_attention(input_ids, past_key_values=DynamicCache())
 
         assert all(

tests/presses/test_wrappers.py

@@ -14,7 +14,7 @@ def test_composed_press_qfilter_without_post_init(unit_test_model):  # noqa: F81
     composed_press = ComposedPress([press1, press2])
     with pytest.raises(ValueError, match="post_init_from_model"):
         with composed_press(unit_test_model):
-            input_ids = unit_test_model.dummy_inputs["input_ids"]
+            input_ids = unit_test_model.dummy_inputs["input_ids"].to(unit_test_model.device)
             unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values
 
 
@@ -28,7 +28,7 @@ def test_composed_press_duo_attention_without_post_init(unit_test_model):  # noq
     composed_press = ComposedPress([press1, press2])
     with pytest.raises(ValueError, match="post_init_from_model"):
         with composed_press(unit_test_model):
-            input_ids = unit_test_model.dummy_inputs["input_ids"]
+            input_ids = unit_test_model.dummy_inputs["input_ids"].to(unit_test_model.device)
             unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values
 
 
@@ -45,7 +45,7 @@ def test_composed_qfilter_press_with_post_init(unit_test_model):  # noqa: F811
 
     composed_press = ComposedPress([press1, press2])
     with composed_press(unit_test_model):
-        input_ids = unit_test_model.dummy_inputs["input_ids"]
+        input_ids = unit_test_model.dummy_inputs["input_ids"].to(unit_test_model.device)
         with pytest.raises(RuntimeError, match="The size of tensor"):
             unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values
 
@@ -63,5 +63,5 @@ def test_composed_duo_attention_press_with_post_init(unit_test_model):  # noqa:
 
     composed_press = ComposedPress([press1, press2])
     with composed_press(unit_test_model):
-        input_ids = unit_test_model.dummy_inputs["input_ids"]
+        input_ids = unit_test_model.dummy_inputs["input_ids"].to(unit_test_model.device)
         unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values