diff --git a/test/test_rangeify.py b/test/test_rangeify.py
index 0febe622d8..b83dff6261 100644
--- a/test/test_rangeify.py
+++ b/test/test_rangeify.py
@@ -116,10 +116,10 @@ class TestRangeify(unittest.TestCase):
     out.realize()
 
   def test_flash_attention(self):
-    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+    #BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
 
     # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 128, 64
+    BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
 
     # llama 8B
     #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128