diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
index ad43decc99..b8799ac9f1 100644
--- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
+++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
@@ -850,11 +850,11 @@ struct UnifiedAttentionPipeline
             {
                 gemm_1(o_acc,
                        get_slice_tile(sp(sp_reg_idx).p,
-                                      sequence<0, (k1_loops - 1) * HEAD_SIZE_PADDED>{},
-                                      sequence<BLOCK_M, k1_loops * HEAD_SIZE_PADDED>{}),
+                                      sequence<0, (k1_loops - 1) * BLOCK_SIZE>{},
+                                      sequence<BLOCK_M, k1_loops * BLOCK_SIZE>{}),
                        get_slice_tile(kv_tile.v_tile,
-                                      sequence<0, (k1_loops - 1) * HEAD_SIZE_PADDED>{},
-                                      sequence<BLOCK_SIZE, k1_loops * HEAD_SIZE_PADDED>{}));
+                                      sequence<0, (k1_loops - 1) * BLOCK_SIZE>{},
+                                      sequence<BLOCK_SIZE, k1_loops * BLOCK_SIZE>{}));
                 fmha_alu0(number<1>{} - sp_reg_idx);
             }
         };