Fix OOM in attention kernel test (#1223)

vllm-project · Sep 28, 2023 · 6f88f76 · 6f88f76
1 parent 202351d
commit 6f88f76
Showing 1 changed file with 5 additions and 2 deletions.
diff --git 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))