make sizes smaller for test

liangel-02 · liangel-02 · commit 11ec4fd13ea8 · 2025-08-22T10:49:45.000-07:00
diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py
@@ -6003,10 +6003,10 @@ def bias_func(score, b, h, q_idx, kv_idx):
     )
     @torch.compile
     def test_learnable_bias_global_compiled(self, device, params):
-        batch_size = 2
-        num_heads = 16
+        batch_size = 1
+        num_heads = 1
         seq_len = 128
-        head_dim = 64
+        head_dim = 16
         d_model = num_heads * head_dim
 
         query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)
