From 6f88f762bf907bccdadc52948001b21ccb616a01 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 28 Sep 2023 14:33:24 -0700 Subject: [PATCH] Fix OOM in attention kernel test (#1223) --- tests/kernels/test_attention.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/kernels/test_attention.py b/tests/kernels/test_attention.py index 813f6fdb..59d8b0a5 100644 --- a/tests/kernels/test_attention.py +++ b/tests/kernels/test_attention.py @@ -247,8 +247,11 @@ def test_multi_query_kv_attention( torch.random.manual_seed(seed) torch.cuda.manual_seed(seed) - seq_lens = random.sample(range(1, MAX_SEQ_LEN), num_seqs) - seq_lens[-1] = MAX_SEQ_LEN + # MAX_SEQ_LEN sometimes causes OOM in the reference implementation. + # As the xformers library is already tested with its own tests, we can use + # a smaller MAX_SEQ_LEN here. + max_len = min(MAX_SEQ_LEN, 4096) + seq_lens = random.sample(range(1, max_len), num_seqs) num_tokens = sum(seq_lens) scale = float(1.0 / (head_size**0.5))