Cover custom sdpa + kv cache + quant in CI

Author: Guang Yang

tests/models/test_modeling_llama.py

@@ -17,7 +17,6 @@
 import logging
 import os
 import subprocess
-import sys
 import tempfile
 import unittest
 
@@ -33,9 +32,6 @@
 from ..utils import check_causal_lm_output_quality
 
 
-is_linux_ci = sys.platform.startswith("linux") and os.environ.get("GITHUB_ACTIONS") == "true"
-
-
 @pytest.mark.skipif(
     parse(torchao.__version__) < parse("0.11.0.dev0"),
     reason="Only available on torchao >= 0.11.0.dev0",
@@ -57,10 +53,11 @@ def test_llama3_2_1b_export_to_executorch(self):
                     --model {model_id} \
                     --task {task} \
                     --recipe {recipe} \
-                    --output_dir {tempdir}/executorch \
                     --use_custom_sdpa \
+                    --use_custom_kv_cache \
                     --qlinear \
-                    --qembedding",
+                    --qembedding \
+                    --output_dir {tempdir}/executorch",
                 shell=True,
                 check=True,
             )
@@ -74,15 +71,18 @@ def test_llama3_2_1b_export_to_executorch(self):
 
     @slow
     @pytest.mark.run_slow
-    @pytest.mark.skipif(is_linux_ci, reason="OOM on linux runner")
-    def test_llama3_2_1b_text_generation(self):
-        # TODO: Switch to use meta-llama/Llama-3.2-1B once https://github.com/huggingface/optimum/issues/2127 is fixed
-        # model_id = "lama/Llama-3.2-1B"
+    def test_llama_text_generation_with_custom_sdpa_8da4w_8we(self):
+        # ExecuTorch model + custom sdpa + 8da4w linear quantization + int8 embedding quantization
         model_id = "NousResearch/Llama-3.2-1B"
-        model = ExecuTorchModelForCausalLM.from_pretrained(model_id, recipe="xnnpack")
+        kwargs = {"qlinear": True, "qembedding": True}
+        model = ExecuTorchModelForCausalLM.from_pretrained(
+            model_id,
+            recipe="xnnpack",
+            attn_implementation="custom_sdpa",
+            **kwargs,
+        )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)
         self.assertIsInstance(model.model, ExecuTorchModule)
-
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         generated_text = model.text_generation(
             tokenizer=tokenizer,
@@ -101,19 +101,18 @@ def test_llama3_2_1b_text_generation(self):
 
     @slow
     @pytest.mark.run_slow
-    def test_llama_text_generation_with_custom_sdpa_8da4w_8we(self):
-        # ExecuTorch model + custom sdpa + 8da4w linear quantization + int8 embedding quantization
+    def test_llama_text_generation_with_custom_sdpa_and_kv_cache_8da4w_8we(self):
         model_id = "NousResearch/Llama-3.2-1B"
-        kwargs = {"qlinear": True, "qembedding": True}
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_id,
             recipe="xnnpack",
             attn_implementation="custom_sdpa",
-            **kwargs,
+            use_custom_kv_cache=True,
+            **{"qlinear": True, "qembeeding": True},
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)
         self.assertIsInstance(model.model, ExecuTorchModule)
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
         generated_text = model.text_generation(
             tokenizer=tokenizer,
             prompt="Simply put, the theory of relativity states that",

tests/models/test_modeling_olmo.py

@@ -119,3 +119,32 @@ def test_olmo_text_generation_with_custom_sdpa_8da4w_8we(self):
         gc.collect()
 
         self.assertTrue(check_causal_lm_output_quality(model_id, generated_tokens))
+
+    @slow
+    @pytest.mark.run_slow
+    def test_olmo_text_generation_with_custom_sdpa_and_kv_cache_8da4w_8we(self):
+        model_id = "allenai/OLMo-1B-hf"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = ExecuTorchModelForCausalLM.from_pretrained(
+            model_id,
+            recipe="xnnpack",
+            attn_implementation="custom_sdpa",
+            use_custom_kv_cache=True,
+            **{"qlinear": True, "qembeeding": True},
+        )
+        self.assertIsInstance(model, ExecuTorchModelForCausalLM)
+        self.assertIsInstance(model.model, ExecuTorchModule)
+        generated_text = model.text_generation(
+            tokenizer=tokenizer,
+            prompt="My favourite condiment is ",
+            max_seq_len=32,
+        )
+        logging.info(f"\nGenerated text:\n\t{generated_text}")
+        generated_tokens = tokenizer(generated_text, return_tensors="pt").input_ids
+
+        # Free memory before loading eager for quality check
+        del model
+        del tokenizer
+        gc.collect()
+
+        self.assertTrue(check_causal_lm_output_quality(model_id, generated_tokens))

tests/models/test_modeling_qwen3.py

@@ -17,6 +17,7 @@
 import logging
 import os
 import subprocess
+import sys
 import tempfile
 import unittest
 
@@ -35,6 +36,8 @@
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+is_linux_ci = sys.platform.startswith("linux") and os.environ.get("GITHUB_ACTIONS") == "true"
+
 
 class ExecuTorchModelIntegrationTest(unittest.TestCase):
     def __init__(self, *args, **kwargs):
@@ -67,6 +70,7 @@ def test_qwen3_export_to_executorch(self):
 
     @slow
     @pytest.mark.run_slow
+    @pytest.mark.skipif(is_linux_ci, reason="OOM on linux runner")
     def test_qwen3_text_generation(self):
         model_id = "Qwen/Qwen3-0.6B"
         model = ExecuTorchModelForCausalLM.from_pretrained(model_id, recipe="xnnpack")
@@ -91,6 +95,7 @@ def test_qwen3_text_generation(self):
 
     @slow
     @pytest.mark.run_slow
+    @pytest.mark.skipif(is_linux_ci, reason="OOM on linux runner")
     def test_qwen3_text_generation_with_custom_sdpa(self):
         model_id = "Qwen/Qwen3-0.6B"
         prompt = "Give me a short introduction to large language model."
@@ -221,3 +226,37 @@ def test_qwen3_text_generation_with_custom_sdpa_and_kv_cache(self):
         gc.collect()
 
         self.assertTrue(check_causal_lm_output_quality(model_id, generated_tokens))
+
+    @slow
+    @pytest.mark.run_slow
+    @pytest.mark.skipif(
+        parse(transformers.__version__) < parse("4.52.0") or parse(torchao.__version__) < parse("0.11.0"),
+        reason="Only available on transformers >= 4.52.0 and torchao >= 0.11.0",
+    )
+    def test_qwen3_text_generation_with_custom_sdpa_and_kv_cache_8da4w_8we(self):
+        model_id = "Qwen/Qwen3-0.6B"
+        prompt = "Give me a short introduction to large language model."
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = ExecuTorchModelForCausalLM.from_pretrained(
+            model_id,
+            recipe="xnnpack",
+            attn_implementation="custom_sdpa",
+            use_custom_kv_cache=True,
+            **{"qlinear": True, "qembeeding": True},
+        )
+        self.assertIsInstance(model, ExecuTorchModelForCausalLM)
+        self.assertIsInstance(model.model, ExecuTorchModule)
+        generated_text = model.text_generation(
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_seq_len=64,
+        )
+        logging.info(f"\nGenerated text:\n\t{generated_text}")
+        generated_tokens = tokenizer(generated_text, return_tensors="pt").input_ids
+
+        # Free memory before loading eager for quality check
+        del model
+        del tokenizer
+        gc.collect()
+
+        self.assertTrue(check_causal_lm_output_quality(model_id, generated_tokens))

tests/models/test_modeling_smollm.py

@@ -47,7 +47,13 @@ def test_smollm_export_to_executorch(self):
         recipe = "xnnpack"
         with tempfile.TemporaryDirectory() as tempdir:
             subprocess.run(
-                f"optimum-cli export executorch --model {model_id} --task {task} --recipe {recipe} --output_dir {tempdir}/executorch",
+                f"optimum-cli export executorch \
+                    --model {model_id} \
+                    --task {task} \
+                    --recipe {recipe} \
+                    --use_custom_sdpa \
+                    --use_custom_kv_cache \
+                    --output_dir {tempdir}/executorch",
                 shell=True,
                 check=True,
             )
@@ -181,3 +187,37 @@ def test_smollm_text_generation_with_custom_sdpa_8da4w_8we(self):
         gc.collect()
 
         self.assertTrue(check_causal_lm_output_quality(model_id, generated_tokens))
+
+    @slow
+    @pytest.mark.run_slow
+    @pytest.mark.skipif(
+        parse(torchao.__version__) < parse("0.11.0"),
+        reason="Only available on torchao >= 0.11.0",
+    )
+    def test_smollm_text_generation_with_custom_sdpa_and_kv_cache_8da4w_8we(self):
+        model_id = "HuggingFaceTB/SmolLM2-135M"
+        prompt = "My favourite condiment is "
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = ExecuTorchModelForCausalLM.from_pretrained(
+            model_id,
+            recipe="xnnpack",
+            attn_implementation="custom_sdpa",
+            use_custom_kv_cache=True,
+            **{"qlinear": True, "qembeeding": True},
+        )
+        self.assertIsInstance(model, ExecuTorchModelForCausalLM)
+        self.assertIsInstance(model.model, ExecuTorchModule)
+        generated_text = model.text_generation(
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_seq_len=64,
+        )
+        logging.info(f"\nGenerated text:\n\t{generated_text}")
+        generated_tokens = tokenizer(generated_text, return_tensors="pt").input_ids
+
+        # Free memory before loading eager for quality check
+        del model
+        del tokenizer
+        gc.collect()
+
+        self.assertTrue(check_causal_lm_output_quality(model_id, generated_tokens))