SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[tests] fix static cache implementation is not compatible with attn_implementation==flash_attention_2 (#32039)

Browse Source 
* add flash attention check

* fix

* fix
This commit is contained in:
Fanli Lin
2024-07-26 17:41:27 +08:00
committed by GitHub
parent 5f841c74b6
commit 27c7f971c0
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/utils/test_cache_utils.py
View File

@@ -290,7 +290,7 @@ class CacheIntegrationTest(unittest.TestCase):
self.assertTrue(decoded[0].endswith(last_output)) self.assertTrue(decoded[0].endswith(last_output))
@require_torch_gpu @require_torch_gpu
@parameterized.expand(["eager", "sdpa", "flash_attention_2"]) @parameterized.expand(["eager", "sdpa"])
def test_static_cache_greedy_decoding_pad_left(self, attn_implementation): def test_static_cache_greedy_decoding_pad_left(self, attn_implementation):
EXPECTED_GENERATION = [ EXPECTED_GENERATION = [
"The best color is the one that complements the skin tone of the", "The best color is the one that complements the skin tone of the",
@@ -330,7 +330,7 @@ class CacheIntegrationTest(unittest.TestCase):
self.assertListEqual(decoded, EXPECTED_GENERATION) self.assertListEqual(decoded, EXPECTED_GENERATION)
@require_torch_gpu @require_torch_gpu
@parameterized.expand(["eager", "sdpa", "flash_attention_2"]) @parameterized.expand(["eager", "sdpa"])
def test_static_cache_greedy_decoding_pad_right(self, attn_implementation): def test_static_cache_greedy_decoding_pad_right(self, attn_implementation):
EXPECTED_GENERATION = [ EXPECTED_GENERATION = [
"The best color isЋ the one that complements the skin tone of", "The best color isЋ the one that complements the skin tone of",