Promote Q matrix (#273)

This PR provides a short term fix for the failing tests. When the Q matrix is not promoted, scheduling fails. So for now, we promote the Q matrix. Signed-off-by: Harsh Menon <[email protected]>
iree-org · Nov 14, 2024 · 327b84b · 327b84b
1 parent 79ec575
commit 327b84b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/kernel/wave/wave_attention_test.py b/tests/kernel/wave/wave_attention_test.py
@@ -414,7 +414,7 @@ def testAttention(
 
     @tkw.wave(constraints)
     def base_attention(
-        q: tkl.Memory[B, M, K1, GLOBAL_ADDRESS_SPACE, tkl.f16],
+        q: tkl.Memory[B, M, K1, ADDRESS_SPACE, tkl.f16],
         k: tkl.Memory[B, K2, K1, ADDRESS_SPACE, tkl.f16],
         v: tkl.Memory[B, N, K2, ADDRESS_SPACE, tkl.f16],
         c: tkl.Memory[B, M, N, GLOBAL_ADDRESS_SPACE, tkl.f32],