From d79715ba5301a287cf708ddd10603b50869b38f5 Mon Sep 17 00:00:00 2001
From: rocking <ChunYu.Lai@amd.com>
Date: Thu, 24 Oct 2024 11:43:45 +0000
Subject: [PATCH] Fix bug of rmsnorm

---
 .../ck_tile/ops/rmsnorm2d/kernel/rmsnorm2d_fwd_kernel.hpp   | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/include/ck_tile/ops/rmsnorm2d/kernel/rmsnorm2d_fwd_kernel.hpp b/include/ck_tile/ops/rmsnorm2d/kernel/rmsnorm2d_fwd_kernel.hpp
index 6207c24f43..99084a25e4 100644
--- a/include/ck_tile/ops/rmsnorm2d/kernel/rmsnorm2d_fwd_kernel.hpp
+++ b/include/ck_tile/ops/rmsnorm2d/kernel/rmsnorm2d_fwd_kernel.hpp
@@ -134,10 +134,8 @@ struct Rmsnorm2dFwd
                 number<Vector_N>{},
                 number<1>{});
 
-            // NOTE: we don't do any pad in this kernel for loading, assume that inside kernel will
-            // check the max count dynamically
             const auto tmp2_ = pad_tensor_view(
-                tmp_, make_tuple(number<Block_M>{}, number<Block_N>{}), sequence<false, false>{});
+                tmp_, make_tuple(number<Block_M>{}, number<Block_N>{}), sequence<kPadM, kPadN>{});
             return make_tile_window(
                 tmp2_, make_tuple(number<Block_M>{}, number<Block_N>{}), {iM, 0});
         }();
@@ -151,7 +149,7 @@ struct Rmsnorm2dFwd
                 number<1>{});
 
             const auto tmp2_ =
-                pad_tensor_view(tmp_, make_tuple(number<Block_N>{}), sequence<false>{});
+                pad_tensor_view(tmp_, make_tuple(number<Block_N>{}), sequence<kPadM>{});
 
             return make_tile_window(tmp2_, make_tuple(number<Block_N>{}), {0});
         }();