diff --git a/custom_ops/xpu_ops/src/ops/block_attn.cc b/custom_ops/xpu_ops/src/ops/block_attn.cc
index c9e3313f2..d3dd4c726 100644
--- a/custom_ops/xpu_ops/src/ops/block_attn.cc
+++ b/custom_ops/xpu_ops/src/ops/block_attn.cc
@@ -585,7 +585,8 @@ std::vector<paddle::Tensor> BlockAttnKernel(
                                                      int,
                                                      E_Scale>(
             xpu_ctx->x_context(),
-            reinterpret_cast<const XPU_XType*>(qkv.data<data_t>()),  // qkv
+            reinterpret_cast<const XPU_XType*>(qkv.data<data_t>()) +
+                total_enc_len * qkv_shape[qkv_shape.size() - 1],  // qkv
             reinterpret_cast<const float*>(
                 rotary_embs.data<float>()),  // rotary_pos_emb
             reinterpret_cast<const int*>(