[generate] ✨ vectorized beam search ✨ (#35802)

2025-03-18 18:39:36 +00:00
parent 12f2ebef63
commit 179d02ffb8
8 changed files with 426 additions and 271 deletions
--- a/tests/models/gemma2/test_modeling_gemma2.py
+++ b/tests/models/gemma2/test_modeling_gemma2.py
@@ -119,10 +119,6 @@ class Gemma2ModelTest(GemmaModelTest, unittest.TestCase):
    def test_generate_continue_from_past_key_values(self):
        pass

-    @unittest.skip("Gemma2 has HybridCache and doesn't support low_memory generation")
-    def test_beam_search_low_memory(self):
-        pass
-
    @unittest.skip("Gemma2 has HybridCache and doesn't support contrastive generation")
    def test_contrastive_generate(self):
        pass