From bb49d2e6238c5916f1cdfc2f02b71e2341c26c75 Mon Sep 17 00:00:00 2001
From: mtgu0705 <mtgu@amd.com>
Date: Sat, 10 May 2025 14:33:44 +0800
Subject: [PATCH] fix the bug, functional test passed

---
 .../blockwise_gemm_pipeline_xdlops_b_preshuflle_v1_mx.hpp   | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuflle_v1_mx.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuflle_v1_mx.hpp
index 0b9ab20294..8312f14d0f 100644
--- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuflle_v1_mx.hpp
+++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuflle_v1_mx.hpp
@@ -399,6 +399,9 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v1_mx<BlockGemmPipelineScheduler
                 b_scale_grid_desc, make_multi_index(NWaves * NPerXDL, -ScalesPerKBlockSize));
         });
 
+        b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc,
+                                               make_multi_index(-NPerBlock, ScalesPerKBlockSize));
+
         // Local prefetch A1
         block_sync_lds();
         static_for<0, KRepeat, 1>{}([&](auto k) {
@@ -584,6 +587,9 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v1_mx<BlockGemmPipelineScheduler
                             b_scale_grid_desc,
                             make_multi_index(NWaves * NPerXDL, -ScalesPerKBlockSize));
                     });
+
+                    b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc,
+                                               make_multi_index(-NPerBlock, ScalesPerKBlockSize));
                 };
 
                 LoopFunc(I0, I1);