diff --git a/lit_tests/kernel/wave/attention.py b/lit_tests/kernel/wave/attention.py
index 0968573fd..92498b225 100644
--- a/lit_tests/kernel/wave/attention.py
+++ b/lit_tests/kernel/wave/attention.py
@@ -1171,13 +1171,12 @@ def test_prefill_attention():
         # CHECK-COUNT-1:            vector.store
         # CHECK-COUNT-1:            vector.maskedload
         # CHECK-COUNT-1:            vector.store
-        # CHECK-COUNT-16:           vector.load
-        # CHECK-COUNT-16:           amdgpu.mfma
-        # CHECK-COUNT-4:            gpu.shuffle xor {{.*}}
         # CHECK-COUNT-1:            vector.maskedload
         # CHECK-COUNT-1:            vector.store
         # CHECK-COUNT-1:            vector.maskedload
         # CHECK-COUNT-1:            vector.store
-        # CHECK-COUNT-16:           vector.load
+        # CHECK-COUNT-32:           vector.load
+        # CHECK-COUNT-16:           amdgpu.mfma
+        # CHECK-COUNT-4:            gpu.shuffle xor {{.*}}
         # CHECK-COUNT-16:           amdgpu.mfma
         # CHECK-COUNT-16:      vector.maskedstore