From 07bb33866b1ae113f65e7844d66d49a4fccdf5a8 Mon Sep 17 00:00:00 2001
From: Tianxing Wu <chi0chu319@gmail.com>
Date: Wed, 12 Nov 2025 14:00:51 +0000
Subject: [PATCH] block_shape fixes

---
 .../pipeline/unified_attention_pipeline.hpp               | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

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);
             }
         };