Broadcasts attention bias over query dimension

Updates forward/backward equivalence benchmarks to create attention bias with a singleton query dimension so it broadcasts across queries.

Aligns shapes with kernel expectations during cached decoding, reduces memory footprint, and prevents shape mismatches across CUDA, Triton, and Flex paths.

Author: LoserCheems

benchmarks/backward_equivalence.py

@@ -597,7 +597,7 @@ def test_cuda_backward_equivalence(accuracy_threshold=0.95):
             device=device, dtype=dtype, requires_grad=True
         )
         attn_bias = torch.randn(
-            batch_size, num_kv_heads, query_len, key_len,
+            batch_size, num_kv_heads, 1, key_len,
             device=device, dtype=torch.bfloat16
         )
         cache_position = torch.arange(key_len - query_len, key_len, device=device)
@@ -831,7 +831,7 @@ def test_triton_backward_equivalence(accuracy_threshold=0.95):
             device=device, dtype=dtype, requires_grad=True
         )
         attn_bias = torch.randn(
-            batch_size, num_kv_heads, query_len, key_len,
+            batch_size, num_kv_heads, 1, key_len,
             device=device, dtype=torch.bfloat16
         )
         cache_position = torch.arange(key_len - query_len, key_len, device=device)

benchmarks/forward_equivalence.py

@@ -570,7 +570,7 @@ def test_cuda_forward_equivalence(accuracy_threshold=0.95):
             device=device, dtype=torch.bfloat16
         )
         attn_bias = torch.randn(
-            batch_size, num_kv_heads, query_len, key_len,
+            batch_size, num_kv_heads, 1, key_len,
             device=device, dtype=torch.bfloat16
         )
         cache_position = torch.arange(key_len - query_len, key_len, device=device)
@@ -758,7 +758,7 @@ def test_triton_forward_equivalence(accuracy_threshold=0.95):
             device=device, dtype=torch.bfloat16
         )
         attn_bias = torch.randn(
-            batch_size, num_kv_heads, query_len, key_len,
+            batch_size, num_kv_heads, 1, key_len,
             device=device, dtype=torch.bfloat16
         )
         cache_position = torch.arange(key_len - query_len, key_len, device=device)
@@ -963,7 +963,7 @@ def test_flex_forward_equivalence(accuracy_threshold=0.95):
             device=device, dtype=torch.bfloat16
         )
         attn_bias = torch.randn(
-            batch_size, num_kv_heads, query_len, key_len,
+            batch_size, num_kv_heads, 1, key_len,
             device=device, dtype=torch.bfloat16
         )
         cache_position = torch.arange(key_len - query_len, key_len, device=device)