Attention Perf: Transpose blocked K right before QK instead of pre-transposing before the kernel

stack-info: PR: #2374, branch: AmesingFlank/stack/50

Author: AmesingFlank

examples/attention.py

@@ -61,7 +61,7 @@ def attention(
     assert head_dim == k_in.size(-1) == v_in.size(-1)
     q_view = q_in.reshape([-1, m_dim, head_dim])
     v_view = v_in.reshape([-1, n_dim, head_dim])
-    k_view = k_in.reshape([-1, n_dim, head_dim]).transpose(1, 2)
+    k_view = k_in.reshape([-1, n_dim, head_dim])
     out = torch.empty_like(q_view)
     sm_scale = 1.0 / math.sqrt(head_dim)
     qk_scale = sm_scale * 1.44269504  # 1/log(2)
@@ -73,10 +73,10 @@ def attention(
         for tile_n in hl.tile(v_view.size(1)):
             # scaling Q in-loop on-demand reduces spillage, faster than keeping pre-scaled Q
             q_scaled = q * qk_scale
-            k = k_view[tile_b, :, tile_n]
+            k = k_view[tile_b, tile_n, :]
             # Keep scores in fp32 to match SDPA tolerances on bf16/fp16 inputs.
             # same as hl.dot(q, k, out_dtype=torch.float32)
-            qk = torch.bmm(q_scaled, k, torch.float32)
+            qk = torch.bmm(q_scaled, k.transpose(1, 2), torch.float32)
             m_ij = torch.maximum(m_i, torch.amax(qk, -1))
             qk = qk - m_ij[:, :, None]
             p = torch.exp2(qk)

test/test_pallas.py

@@ -306,7 +306,7 @@ def pallas_attention(
     head_dim = hl.specialize(q_in.size(-1))
     assert head_dim == k_in.size(-1) == v_in.size(-1)
     q_view = q_in.reshape([-1, m_dim, head_dim])
-    k_view = k_in.reshape([-1, n_dim, head_dim]).transpose(1, 2)
+    k_view = k_in.reshape([-1, n_dim, head_dim])
     v_view = v_in.reshape([-1, n_dim, head_dim])
     out = torch.empty_like(q_view)
     sm_scale = 1.0 / math.sqrt(head_dim)
@@ -319,9 +319,10 @@ def pallas_attention(
         for tile_n in hl.tile(v_view.size(1)):
             # scaling Q in-loop on-demand reduces spillage, faster than keeping pre-scaled Q
             q_scaled = q * qk_scale
-            k = k_view[tile_b, :, tile_n]
+            k = k_view[tile_b, tile_n, :]
+            # Keep scores in fp32 to match SDPA tolerances on bf16/fp16 inputs.
             # same as hl.dot(q, k, out_dtype=torch.float32)
-            qk = torch.bmm(q_scaled, k, torch.float32)
+            qk = torch.bmm(q_scaled, k.transpose(1, 2), torch.float32)
             m_ij = torch.maximum(m_i, torch.amax(qk, -1))
             qk = qk - m_ij[:, :, None]
             p = torch.exp2(qk)
@@ -1059,7 +1060,7 @@ def test_attention_fori_loop_correctness_head_dim_256(self) -> None:
             "((4, 128, 128), 'jnp.float32', 'vmem'), "
             "((4, 128, 128), 'jnp.float32', 'vmem'), "
             "((4, 128, 256), 'jnp.float32', 'vmem'), "
-            "((4, 256, 128), 'jnp.float32', 'vmem'), "
+            "((4, 128, 256), 'jnp.float32', 'vmem'), "
             "((), None, 'dma_semaphore'), "
             "((4, 128, 256), 'jnp.float32', 'vmem'), "
             "((), None, 'dma_semaphore')]",