From 86b64305829782be382482bb72bef376920b1e70 Mon Sep 17 00:00:00 2001
From: cmcamdy <1027740945@qq.com>
Date: Fri, 5 Dec 2025 14:33:17 +0800
Subject: [PATCH] fix split_rope_cache_kv_encoder in mix mtp (#5384)

---
 custom_ops/xpu_ops/src/ops/block_attn.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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*>(