Fix OOM in attention kernel test (#1223)

6f88f762 · Woosuk Kwon · GitHub · 202351d5 · 6f88f762
Unverified Commit 6f88f762 authored Sep 28, 2023 by Woosuk Kwon Committed by GitHub Sep 28, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 2 deletions

tests/kernels/test_attention.py tests/kernels/test_attention.py +5 -2

No files found.
--- 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)
+    # MAX_SEQ_LEN sometimes causes OOM in the reference implementation.
-    seq_lens[-1] = MAX_SEQ_LEN
+    # 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))