diff --git a/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
index 590346877c..ba70a40e6b 100644
--- a/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
+++ b/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
@@ -816,14 +816,21 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
         const auto k_dram = [&]() {
             const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 k_ptr,
-                make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_q),
-                make_tuple(kargs.stride_k, 1),
+                make_tuple(kargs.num_total_pages, kargs.page_block_size, kargs.hdim_q),
+                make_tuple(kargs.batch_stride_k, kargs.stride_k, 1),
                 number<FmhaPipeline::kAlignmentK>{},
                 number<1>{});
 
+            auto k_dram_2d = transform_tensor_view(
+                k_dram_naive,
+                make_tuple(make_merge_transform(make_tuple(kargs.num_total_pages, kargs.page_block_size)),
+                           make_pass_through_transform(kargs.hdim_q)),
+                make_tuple(sequence<0, 1>{}, sequence<2>{}),
+                make_tuple(sequence<0>{}, sequence<1>{}));
+
             constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
             return pad_tensor_view(
-                k_dram_naive,
+                k_dram_2d,
                 make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
                 sequence<kPadSeqLenK_, kPadHeadDimQ>{});
         }();
@@ -832,13 +839,21 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
             {
                 const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                     v_ptr,
-                    make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_v),
-                    make_tuple(kargs.stride_v, 1),
+                    make_tuple(kargs.num_total_pages, kargs.page_block_size, kargs.hdim_v),
+                    make_tuple(kargs.batch_stride_v, kargs.stride_v, 1),
                     number<FmhaPipeline::kAlignmentV>{},
                     number<1>{});
 
-                const auto v_dram_transposed = transform_tensor_view(
+                auto v_dram_2d = transform_tensor_view(
                     v_dram_naive,
+                    make_tuple(make_merge_transform(
+                                   make_tuple(kargs.num_total_pages, kargs.page_block_size)),
+                               make_pass_through_transform(kargs.hdim_v)),
+                    make_tuple(sequence<0, 1>{}, sequence<2>{}),
+                    make_tuple(sequence<0>{}, sequence<1>{}));
+
+                const auto v_dram_transposed = transform_tensor_view(
+                    v_dram_2d,
                     make_tuple(
                         make_pass_through_transform(kargs.hdim_v),
                         make_pass_through_transform(kargs.num_total_pages * kargs.page_block_size)),
@@ -1088,6 +1103,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
                     kargs.kv_page_indices,
                     kargs.stride_k,
                     kargs.stride_v,
+                    kargs.batch_stride_k,
+                    kargs.batch_stride_v,
                     dropout);
             }
             else
@@ -1108,6 +1125,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
                                       kargs.kv_page_indices,
                                       kargs.stride_k,
                                       kargs.stride_v,
+                                      kargs.batch_stride_k,
+                                      kargs.batch_stride_v,
                                       dropout);
             }
         }();
diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp
index f534c6f47e..b586867dce 100644
--- a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp
+++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp
@@ -24,6 +24,7 @@ template <typename OffsetVecType,
           bool kIsKcache>
 CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec,
                                                    const index_t& stride_kv,
+                                                   const index_t& page_stride_kv,
                                                    const CoordVecType& coord_vec,
                                                    OffsetVecType& kv_offset_vec)
 {
@@ -46,7 +47,7 @@ CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec,
                 const index_t page_offset =
                     (thread_coord_start + kLoopStride * k0.value) & kPageMask;
                 kv_offset_vec[k0] =
-                    ((page_vec[kPageId] << kPageShiftSize) + page_offset) * stride_kv;
+                    page_vec[kPageId] * page_stride_kv + page_offset * stride_kv;
             });
         }
         else
@@ -54,11 +55,11 @@ CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec,
             // for v_offset_vec
             const index_t lane0_start   = __builtin_amdgcn_readfirstlane(thread_coord_start);
             const index_t lane0_page_id = (lane0_start + kLoopStart) >> kPageShiftSize;
-            const index_t page_loc      = page_vec[lane0_page_id] << kPageShiftSize;
+            const index_t page_loc      = page_vec[lane0_page_id] * page_stride_kv;
             static_for<0, kLoopCount, 1>{}([&](auto k0) {
                 const index_t page_offset =
                     (thread_coord_start + kLoopStart + k0.value) & kPageMask;
-                kv_offset_vec[k0] = (page_loc + page_offset) * stride_kv;
+                kv_offset_vec[k0] = page_loc + page_offset * stride_kv;
             });
         }
     }
@@ -248,6 +249,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                const index_t* page_idx,
                const index_t stride_k,
                const index_t stride_v,
+               const index_t page_stride_k,
+               const index_t page_stride_v,
                DropoutType& dropout) const
     {
         static_assert(
@@ -386,7 +389,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                                   NRepeat,
                                   kN0 / NRepeat,
                                   kIsSglangLayout,
-                                  true>(page_idx, stride_k, k_coord, k_offsets);
+                                  true>(page_idx, stride_k, page_stride_k, k_coord, k_offsets);
 
         auto k_dram_window = make_tile_scatter_gather(k_dram_block_window.get_bottom_tensor_view(),
                                                       k_dram_block_window.get_window_lengths(),
@@ -429,7 +432,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                                   V_KRepeat,
                                   1,
                                   kIsSglangLayout,
-                                  false>(page_idx, stride_v, v_coord, v_offsets);
+                                  false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
 
         auto v_dram_window =
             make_tile_scatter_gather(v_dram_block_window_tmp.get_bottom_tensor_view(),
@@ -502,7 +505,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                                       V_KRepeat,
                                       1,
                                       kIsSglangLayout,
-                                      false>(page_idx, stride_v, v_coord, v_offsets);
+                                      false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
             v_dram_window.update_page_idx(v_offsets);
 
             __builtin_amdgcn_sched_barrier(0);
@@ -678,7 +681,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                                           V_KRepeat,
                                           1,
                                           kIsSglangLayout,
-                                          false>(page_idx, stride_v, v_coord, v_offsets);
+                                          false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
                 v_dram_window.update_page_idx(v_offsets);
             }
             __builtin_amdgcn_sched_barrier(0);
@@ -815,7 +818,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                                                   V_KRepeat,
                                                   1,
                                                   kIsSglangLayout,
-                                                  false>(page_idx, stride_v, v_coord, v_offsets);
+                                                  false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
                         v_dram_window.update_page_idx(v_offsets);
                     }
                     block_sync_lds();
@@ -877,7 +880,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                                           NRepeat,
                                           kN0 / NRepeat,
                                           kIsSglangLayout,
-                                          true>(page_idx, stride_k, k_coord, k_offsets);
+                                          true>(page_idx, stride_k, page_stride_k, k_coord, k_offsets);
                 k_dram_window.update_page_idx(k_offsets);
                 if constexpr(k1_loops >= 2 &&
                              LdsSeq.at(number<0>{}) == LdsSeq.at(number<k0_loops + k1_loops - 2>{}))
@@ -985,6 +988,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                const index_t* page_idx,
                const index_t stride_k,
                const index_t stride_v,
+               const index_t page_stride_k,
+               const index_t page_stride_v,
                DropoutType& dropout) const
     {
         return operator()(q_dram_block_window_tmp,
@@ -1011,6 +1016,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                           page_idx,
                           stride_k,
                           stride_v,
+                          page_stride_k,
+                          page_stride_v,
                           dropout);
     }
 };