diff --git a/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu b/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu
index a46f427b998..0473ba6b8f6 100644
--- a/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu
+++ b/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu
@@ -12,9 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "cute/tensor.hpp"
 #include "helper.h"
 #include "paddle/extension.h"
 #include "paddle/phi/core/memory/memcpy.h"
+#include "utils.cuh"
 
 template <int THREADBLOCK_SIZE>
 __global__ void
@@ -142,6 +144,94 @@ __global__ void split_q_block(const int *__restrict__ seq_lens_q,
   }
 }
 
+
+template <uint32_t config_size>
+__global__ void search_chunk_size_for_mla(
+    const int *__restrict__ seq_lens_q,
+    const int *__restrict__ seq_lens_encoder,
+    const int *__restrict__ seq_lens_decoder,
+    int *__restrict__ num_blocks_x,
+    int *__restrict__ res_chunk_size,
+    const int bsz,
+    const int set_chunk_size,
+    const int block_size,
+    const int sm_cout) {
+  const uint32_t conf_id = threadIdx.x;
+  int gridx = 0;
+  if (set_chunk_size > 0 && conf_id == 0) {
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+      if (seq_len == 0 || seq_len_encoder > 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, set_chunk_size);
+      gridx += loop_times;
+    }
+    *num_blocks_x = gridx;
+    *res_chunk_size = set_chunk_size;
+  } else if (conf_id < config_size) {
+    __shared__ int gridx_shared[config_size];
+    // chunk_size is a multiple of 64
+    const int chunk_size = block_size << conf_id;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+      if (seq_len == 0 || seq_len_encoder > 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+      gridx += loop_times;
+    }
+    gridx_shared[conf_id] = gridx;
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      uint32_t res_id = 0;
+      uint32_t max_last_wave_block = 0;
+      for (uint32_t i = 1; i < config_size; i++) {
+        uint32_t last_wave_block = gridx_shared[i] % sm_cout;
+        if (last_wave_block >= max_last_wave_block) {
+            res_id = i;
+            max_last_wave_block = last_wave_block;
+        }
+      }
+      *num_blocks_x = gridx_shared[res_id];
+      *res_chunk_size = block_size << res_id;
+    }
+  }
+}
+
+__global__ void split_block_for_mla(const int *__restrict__ seq_lens_q,
+                                    const int *__restrict__ seq_lens_encoder,
+                                    const int *__restrict__ seq_lens_decoder,
+                                    int *__restrict__ batch_ids,
+                                    int *__restrict__ tile_ids_per_batch,
+                                    const int bsz,
+                                    const int chunk_size) {
+  if (threadIdx.x == 0) {
+    int index = 0;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+
+      if (seq_len == 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+      if (seq_len_encoder > 0) {
+        loop_times = 0;
+      }
+      for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) {
+        batch_ids[index] = bid;
+        tile_ids_per_batch[index++] = tile_id;
+      }
+    }
+  }
+}
+
 __global__ void split_kv_block(const int *__restrict__ seq_lens_decoder,
                                const int *__restrict__ seq_lens_encoder,
                                int *__restrict__ batch_ids,
@@ -230,6 +320,9 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
   paddle::Tensor kv_tile_ids_per_batch;
   paddle::Tensor kv_num_blocks_x_cpu;       /*cpu*/
   paddle::Tensor max_len_kv_cpu;            /*cpu*/
+  paddle::Tensor decoder_num_blocks_x;
+  paddle::Tensor decoder_chunk_size_device;
+  paddle::Tensor decoder_chunk_size_cpu; /*cpu*/
 
   auto max_len_kv =
       GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_decoder.place());
@@ -239,6 +332,59 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
 
   max_len_kv_cpu = max_len_kv.copy_to(paddle::CPUPlace(), false);
 
+// decoder
+if (max_dec_len_this_time>0){
+    const bool mla_use_tensorcore = GetMlaUseTensorcore();
+    if (mla_use_tensorcore && group_size <= 64) {
+      const int set_chunk_size = get_mla_dec_chunk_size(bsz);
+      decoder_chunk_size_device = GetEmptyTensor(
+          {1}, paddle::DataType::INT32, seq_lens_encoder.place());
+      decoder_num_blocks_x = GetEmptyTensor(
+          {1}, paddle::DataType::INT32, seq_lens_encoder.place());
+
+      int device;
+      cudaGetDevice(&device);
+      int sm_cout;
+      cudaDeviceGetAttribute(&sm_cout, cudaDevAttrMultiProcessorCount, device);
+      constexpr int config_size =
+          12;  // search space for chunk size:[64, 128, 256, ... 131072]
+
+      search_chunk_size_for_mla<config_size>
+          <<<1, 32, 0, stream>>>(seq_lens_this_time.data<int>(),
+                                 seq_lens_encoder.data<int>(),
+                                 seq_lens_decoder.data<int>(),
+                                 decoder_num_blocks_x.data<int>(),
+                                 decoder_chunk_size_device.data<int>(),
+                                 bsz,
+                                 set_chunk_size,
+                                 block_size,
+                                 sm_cout);
+
+      decoder_chunk_size_cpu =
+          decoder_chunk_size_device.copy_to(paddle::CPUPlace(), false);
+      const int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
+
+      split_block_for_mla<<<1, 32, 0, stream>>>(
+          seq_lens_this_time.data<int>(),
+          seq_lens_encoder.data<int>(),
+          seq_lens_decoder.data<int>(),
+          decoder_batch_ids.data<int>(),
+          decoder_tile_ids_per_batch.data<int>(),
+          bsz,
+          chunk_size);
+    }else{
+      PD_THROW(
+          "Please use FLAGS_mla_use_tensorcore=0 and To ensure num_heads // "
+          "kv_num_heads <= 64");
+    }
+  }else{
+    decoder_chunk_size_cpu =
+        GetEmptyTensor({0}, paddle::DataType::INT32, paddle::CPUPlace());
+    decoder_num_blocks_x =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+  }
+
+// encoder
   if (max_enc_len_this_time > 0) {
     const uint32_t max_tile_size_per_bs_kv =
         div_up(max_enc_dec_len_this_time, block_size);
@@ -315,13 +461,15 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
   }
 
   return {
-    encoder_batch_ids,
-    encoder_tile_ids_per_batch,
-    encoder_num_blocks_x_cpu, /*cpu*/
-    kv_batch_ids,
-    kv_tile_ids_per_batch,
-    kv_num_blocks_x_cpu,      /*cpu*/
-    max_len_kv_cpu,           /*cpu*/
+      encoder_batch_ids,
+      encoder_tile_ids_per_batch,
+      encoder_num_blocks_x_cpu, /*cpu*/
+      kv_batch_ids,
+      kv_tile_ids_per_batch,
+      kv_num_blocks_x_cpu, /*cpu*/
+      decoder_num_blocks_x,
+      decoder_chunk_size_cpu, /*cpu*/
+      max_len_kv_cpu,         /*cpu*/
   };
 }
 
@@ -342,6 +490,8 @@ PD_BUILD_STATIC_OP(get_block_shape_and_split_kv_block)
       paddle::Optional("kv_batch_ids"),
       paddle::Optional("kv_tile_ids_per_batch"),
       paddle::Optional("kv_num_blocks_x_cpu"),
+      paddle::Optional("decoder_num_blocks_x"),
+      paddle::Optional("decoder_chunk_size_cpu"),
       "max_len_kv_cpu"
     })
     .Attrs({
diff --git a/custom_ops/gpu_ops/cpp_extensions.cc b/custom_ops/gpu_ops/cpp_extensions.cc
index b0e6e332b3d..6ec29ef6e4c 100644
--- a/custom_ops/gpu_ops/cpp_extensions.cc
+++ b/custom_ops/gpu_ops/cpp_extensions.cc
@@ -415,6 +415,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
     const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
diff --git a/custom_ops/gpu_ops/env.h b/custom_ops/gpu_ops/env.h
index c7db21ba8f8..7a6dc49d451 100644
--- a/custom_ops/gpu_ops/env.h
+++ b/custom_ops/gpu_ops/env.h
@@ -62,3 +62,12 @@ inline bool get_mla_use_tensorcore() {
             mla_use_tensorcore_env == nullptr ? 1 : std::stoul(std::string(mla_use_tensorcore_env));
     return mla_use_tensorcore != 0 ? true : false;
 }
+inline int get_mla_dec_chunk_size(int bsz) {
+    static const char* mla_dec_chunk_size_env =
+        std::getenv("FLAGS_mla_dec_chunk_size");
+    static const int mla_dec_chunk_size =
+        mla_dec_chunk_size_env == nullptr
+            ? -1
+            : std::stoi(std::string(mla_dec_chunk_size_env));
+    return bsz > 1 ? mla_dec_chunk_size : 64;
+}
\ No newline at end of file
diff --git a/custom_ops/gpu_ops/helper.h b/custom_ops/gpu_ops/helper.h
index ed4efe92702..47bca728483 100644
--- a/custom_ops/gpu_ops/helper.h
+++ b/custom_ops/gpu_ops/helper.h
@@ -556,3 +556,11 @@ inline int GetSMVersion() {
   return sm_version;
 
 }
+
+inline bool GetMlaUseTensorcore() {
+  static const bool flags_mla_use_tensorcore = get_mla_use_tensorcore();
+  static const bool enable_mla_tensorcore = GetSMVersion() >= 90 ? true : false;
+  const bool mla_use_tensorcore =
+      flags_mla_use_tensorcore && enable_mla_tensorcore;
+  return mla_use_tensorcore;
+}
\ No newline at end of file
diff --git a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu
index f7d4b8ae27e..6cc5fbc975d 100644
--- a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu
+++ b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu
@@ -79,6 +79,7 @@ void BatchMLAWithPagedKVCacheKernel(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
+    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
@@ -97,15 +98,20 @@ void BatchMLAWithPagedKVCacheKernel(
   const auto q_head_num = meta_data.q_num_heads;
   const auto max_block_num_per_seq = meta_data.max_blocks_per_seq;
   const auto max_block_num = bsz * max_block_num_per_seq;
-  const uint32_t chunk_size = get_max_partition_size(bsz);
+//   const uint32_t chunk_size = get_max_partition_size(bsz);
 
+#if CUDA_VERSION >= 12080
+    constexpr bool USE_REG_EALLOC = true;
+#else
+    constexpr bool USE_REG_EALLOC = false;
+#endif
 
   int q_head_dim = meta_data.head_dims;
   int k_head_dim = meta_data.head_dims;
   int v_head_dim = meta_data.head_dims_v;
   // int num_chunks = max_dec_len / chunk_size;
   int num_chunks = div_up(max_dec_len, chunk_size);
-
+    std::<< "==================================="num_chunks "==================================="<<std::endl;
   auto *allocator = paddle::GetAllocator(q.place());
   phi::Allocator::AllocationPtr O_tmp, m_tmp, d_tmp;
   O_tmp = allocator->Allocate(
@@ -118,7 +124,8 @@ void BatchMLAWithPagedKVCacheKernel(
       sizeof(float) *
       static_cast<size_t>(num_chunks * bsz * draft_token_num * q_head_num));
 
-  Params<CUTLASS_TYPE, CUTLASS_TYPE, CUTLASS_TYPE, int> params = {};
+  using ParamsType = Params<CUTLASS_TYPE, CUTLASS_TYPE, CUTLASS_TYPE, int>;
+  ParamsType params = {};
   params.Q = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(q.data<T>()));
   params.KV = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(latent_cache.data<T>()));
   params.O = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(out->data<T>()));
@@ -143,6 +150,7 @@ void BatchMLAWithPagedKVCacheKernel(
   params.o_stride_head_num = v_head_dim;
   params.bsz = bsz;
   params.token_num = token_num;
+//   params.max_seq_len = max_seq_len;
   params.max_block_num = max_block_num;
   params.max_block_num_per_seq = max_block_num_per_seq;
   params.q_num_head = q_head_num;
@@ -155,9 +163,11 @@ void BatchMLAWithPagedKVCacheKernel(
   params.chunk_num = num_chunks;
 
   if (q_head_dim == 576) {
-      BatchMLAWithPagedKVCacheDispatched<576, 512, NV_TYPE>(
-          params, stream
-      );
+      BatchMLAWithPagedKVCacheDispatched<576,
+                                         512,
+                                         NV_TYPE,
+                                         ParamsType,
+                                         USE_REG_EALLOC>(params, stream);
   } else {
       PD_THROW("error!!! q_head_dim must be 576 !!!\n");
   }
@@ -185,6 +195,7 @@ template void BatchMLAWithPagedKVCacheKernel<paddle::bfloat16>(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
+    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
@@ -219,6 +230,7 @@ template void BatchMLAWithPagedKVCacheKernel<paddle::float16>(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
+    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
diff --git a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h
index 97fffe39dc3..afd16e2ea83 100644
--- a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h
+++ b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h
@@ -56,6 +56,7 @@ void BatchMLAWithPagedKVCacheKernel(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
+    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
diff --git a/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh b/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh
index ba1f4b4470a..18452a18dcf 100644
--- a/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh
+++ b/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh
@@ -113,6 +113,9 @@ struct Params {
   } else if (group_size == 64) {                             \
     constexpr size_t GROUP_SIZE = 64;                        \
     __VA_ARGS__                                              \
+  } else if (group_size == 128) {                            \
+    constexpr size_t GROUP_SIZE = 128;                       \
+    __VA_ARGS__                                              \
   } else {                                                   \
     PD_THROW("not support the group_size: ", group_size);    \
     return cudaErrorNotSupported;                            \
diff --git a/custom_ops/gpu_ops/multi_head_latent_attention.cu b/custom_ops/gpu_ops/multi_head_latent_attention.cu
index 98a61e83859..c445aa3fde4 100644
--- a/custom_ops/gpu_ops/multi_head_latent_attention.cu
+++ b/custom_ops/gpu_ops/multi_head_latent_attention.cu
@@ -38,6 +38,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
     const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
@@ -66,9 +67,10 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
 
   int decoder_num_blocks_data = decoder_num_blocks_cpu.data<int>()[0];
   int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
+  int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
   int max_len_kv_data = max_len_kv.data<int>()[0];
+  auto mla_use_tensorcore = GetMlaUseTensorcore();
 
-  const bool mla_use_tensorcore = get_mla_use_tensorcore();
   auto sm_version = GetSMVersion();
   if ((speculate_decoder || mla_use_tensorcore) && sm_version < 90) {
     PD_THROW("Please use speculate_decoder=0 and FLAGS_mla_use_tensorcore=0 when sm < 90.");
@@ -105,6 +107,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
                                              decoder_num_blocks,
                                              cache_quant_type_str,
                                              decoder_num_blocks_data,
+                                             chunk_size,
                                              max_input_length,
                                              max_len_kv_data,
                                              softmax_scale,
@@ -161,6 +164,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
     const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
@@ -224,6 +228,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           decoder_tile_ids_per_batch,
           decoder_num_blocks,
           decoder_num_blocks_cpu,
+          decoder_chunk_size_cpu,
           max_enc_len_this_time,
           max_dec_len_this_time,
           max_len_kv,
@@ -270,6 +275,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           decoder_tile_ids_per_batch,
           decoder_num_blocks,
           decoder_num_blocks_cpu,
+          decoder_chunk_size_cpu,
           max_enc_len_this_time,
           max_dec_len_this_time,
           max_len_kv,
@@ -323,6 +329,7 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
     const std::vector<int64_t>& decoder_tile_ids_per_batch_shape,
     const std::vector<int64_t>& decoder_num_blocks_shape,
     const std::vector<int64_t>& decoder_num_blocks_cpu_shape,
+    const std::vector<int64_t>& decoder_chunk_size_cpu_shape,
     const std::vector<int64_t>& max_enc_len_this_time_shape,
     const std::vector<int64_t>& max_dec_len_this_time_shape,
     const std::vector<int64_t>& max_len_kv_shape,
@@ -377,6 +384,7 @@ std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const paddle::DataType& decoder_tile_ids_per_batch_dtype,
     const paddle::DataType& decoder_num_blocks_dtype,
     const paddle::DataType& decoder_num_blocks_cpu_dtype,
+    const paddle::DataType& decoder_chunk_size_cpu_dtype,
     const paddle::DataType& max_enc_len_this_time_dtype,
     const paddle::DataType& max_dec_len_this_time_dtype,
     const paddle::DataType& max_len_kv_dtype,
@@ -431,6 +439,7 @@ PD_BUILD_STATIC_OP(multi_head_latent_attention)
              "decoder_tile_ids_per_batch",
              "decoder_num_blocks",
              "decoder_num_blocks_cpu",
+             "decoder_chunk_size_cpu",
              "max_enc_len_this_time",
              "max_dec_len_this_time",
              "max_len_kv",
diff --git a/fastdeploy/demo/offline_demo.py b/fastdeploy/demo/offline_demo.py
index c02bdb45c41..50a8888cd29 100644
--- a/fastdeploy/demo/offline_demo.py
+++ b/fastdeploy/demo/offline_demo.py
@@ -1,27 +1,52 @@
-"""
-# Copyright (c) 2025  PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License"
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-
-from fastdeploy.engine.sampling_params import SamplingParams
-from fastdeploy.entrypoints.llm import LLM
-
-model_name_or_path = "./models/llama-7b"
-
-# 超参设置
-sampling_params = SamplingParams(temperature=0.1, max_tokens=30)
-llm = LLM(model=model_name_or_path, tensor_parallel_size=1)
-output = llm.generate(prompts="who are you？", use_tqdm=True, sampling_params=sampling_params)
-
-print(output)
+# """
+# # Copyright (c) 2025  PaddlePaddle Authors. All Rights Reserved.
+# #
+# # Licensed under the Apache License, Version 2.0 (the "License"
+# # you may not use this file except in compliance with the License.
+# # You may obtain a copy of the License at
+# #
+# #     http://www.apache.org/licenses/LICENSE-2.0
+# #
+# # Unless required by applicable law or agreed to in writing, software
+# # distributed under the License is distributed on an "AS IS" BASIS,
+# # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# # See the License for the specific language governing permissions and
+# # limitations under the License.
+# """
+# # export CUTLASS_GEMM_STREAM_K=1
+# # model_name_or_path = "/root/paddlejob/workspace/env_run/output/45t04_wint2_perm_2card_all_safetensors"
+# # model_name_or_path = "/cwb/models/DeepSeek-V3-0324"
+# # model_name_or_path = "/cwb/ernie-4_5-300b-a47b-bf16-paddle"
+# model_name_or_path = "/cwb/models/Qwen3-30B-A3B"
+# prompts = ["解析三首李白的诗?"]
+# from fastdeploy import LLM, SamplingParams
+# sampling_params = SamplingParams(temperature=0.7, top_p=0, max_tokens=128)
+# llm = LLM(model=model_name_or_path, tensor_parallel_size=1,enable_custom_all_reduce=False,quantization="wint4", use_cudagraph=False,)
+# # outputs = llm.chat(messages=[{"role": "user", "content": "给我三首李白的诗"}]*20, sampling_params=sampling_params,)
+# outputs = llm.generate(prompts, sampling_params)
+
+# print(outputs)
+
+
+
+
+from fastdeploy import LLM, SamplingParams
+
+prompts = 20 * ["勒布朗詹姆斯是谁?"]
+
+
+# 采样参数
+sampling_params = SamplingParams(temperature=0.7, top_p=0, max_tokens=100)
+
+llm = LLM(model="/cwb/DeepSeekV3-0324-5layers", tensor_parallel_size=8,engine_worker_queue_port=8008, 
+quantization="wint4", use_cudagraph=False, enable_custom_all_reduce=True)
+
+# 批量进行推理（llm内部基于资源情况进行请求排队、动态插入处理）
+outputs = llm.generate(prompts, sampling_params)
+print("==== The outputs are:", outputs)
+# 输出结果
+for output in outputs:
+    prompt = output.prompt
+    #print("===The prompt is : ", prompt)
+    generated_text = output.outputs.text
+    print("===The generated_text is : ", generated_text)
\ No newline at end of file
diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py
index 5279b68f6f6..fa3562fba07 100644
--- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py
+++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py
@@ -24,6 +24,11 @@
 import paddle
 from paddle.nn.functional.flash_attention import flash_attn_unpadded
 
+try:
+    from paddle.nn.functional.flash_attention import flash_attention_v3_varlen
+except:
+    flash_attention_v3_varlen = None
+
 from fastdeploy.model_executor.layers.attention.ops import (
     get_block_shape_and_split_kv_block,
     init_kv_signal_per_query,
@@ -91,6 +96,7 @@ class MLAAttentionBackend(AttentionBackend):
     """
 
     __infer_dynamic_dims_fields__ = ["attention_metadata"]
+    flash_attn_func: callable = None
     attention_metadata: MLAAttentionMetadata
 
     def __init__(
@@ -147,6 +153,21 @@ def __init__(
         self.device_id: int = os.getenv("CUDA_VISIBLE_DEVICES", None)
 
         self.rank, self.device_id = init_rank_and_device_id(fd_config)
+        if self.flash_attn_func is None:
+            prop = paddle.device.cuda.get_device_properties()
+            cc = prop.major * 10 + prop.minor
+            is_current_sm_supported = cc >= 90
+            is_paddle_supported = any(num >= 90 for num in paddle.version.cuda_archs())
+            if is_current_sm_supported and is_paddle_supported:
+                self.flash_attn_func = flash_attention_v3_varlen
+                print("The current platform supports Flash Attention V3.")
+                self.flash_attn_kwargs = {}
+            else:
+                self.flash_attn_func = flash_attn_unpadded
+                self.flash_attn_kwargs = {"scale": self.head_dim**-0.5, "training": False}
+                print(
+                    "The current platform does not support Flash Attention V3, so Flash Attention V2 will be used instead."
+                )
 
     def init_attention_metadata(self, forward_meta: ForwardMeta):
         """Initialize attention metadata hence all layers in the forward pass can reuse it."""
@@ -173,6 +194,8 @@ def init_attention_metadata(self, forward_meta: ForwardMeta):
             metadata.kv_batch_ids,
             metadata.kv_tile_ids_per_batch,
             metadata.kv_num_blocks,
+            metadata.decoder_num_blocks,
+            metadata.decoder_chunk_size_cpu,
             metadata.max_len_kv,
         ) = get_block_shape_and_split_kv_block(
             forward_meta.seq_lens_encoder,
@@ -269,17 +292,16 @@ def forward_extend(
         )
 
         # Flash注意力计算
-        fmha_out = flash_attn_unpadded(
+        fmha_out = self.flash_attn_func(
             q,
             k,
             v,
             forward_meta.cu_seqlens_q,
             forward_meta.cu_seqlens_k,
-            metadata.max_enc_len_this_time,
-            metadata.max_enc_len_this_time,
-            self.attn_softmax_scale,
-            causal=True,
-            training=False,
+            max_seqlen_q=forward_meta.max_len_tensor_cpu[0],
+            max_seqlen_k=forward_meta.max_len_tensor_cpu[3],
+            causal=self.causal,
+            **self.flash_attn_kwargs,
         )[0]
 
         return fmha_out
@@ -346,8 +368,9 @@ def forward_decode(
             metadata.kv_num_blocks,
             forward_meta.decoder_batch_ids,
             forward_meta.decoder_tile_ids_per_batch,
+            metadata.decoder_num_blocks, #decoder_num_blocks_cpu
             forward_meta.decoder_num_blocks_cpu,
-            forward_meta.decoder_num_blocks_cpu,
+            metadata.decoder_chunk_size_cpu,
             metadata.max_enc_len_this_time,
             metadata.max_dec_len_this_time,
             metadata.max_len_kv,
@@ -418,17 +441,16 @@ def forward_mixed(
             )
 
             # FA
-            fmha_out = flash_attn_unpadded(
+            fmha_out = self.flash_attn_func(
                 q,
                 k,
                 v,
                 forward_meta.cu_seqlens_q,
                 forward_meta.cu_seqlens_k,
-                metadata.max_enc_len_this_time,
-                metadata.max_enc_len_this_time,
-                self.attn_softmax_scale,
-                causal=True,
-                training=False,
+                max_seqlen_q=forward_meta.max_len_tensor_cpu[0],
+                max_seqlen_k=forward_meta.max_len_tensor_cpu[3],
+                causal=self.causal,
+                **self.flash_attn_kwargs,
             )[0]
 
             return fmha_out
@@ -468,8 +490,9 @@ def forward_mixed(
                 metadata.kv_num_blocks,
                 forward_meta.decoder_batch_ids,
                 forward_meta.decoder_tile_ids_per_batch,
+                metadata.decoder_num_blocks,
                 forward_meta.decoder_num_blocks_cpu,
-                forward_meta.decoder_num_blocks_cpu,
+                metadata.decoder_chunk_size_cpu,
                 metadata.max_enc_len_this_time,
                 metadata.max_dec_len_this_time,
                 metadata.max_len_kv,
diff --git a/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py b/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py
index dd57b525937..6f800b61e6c 100644
--- a/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py
+++ b/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py
@@ -49,6 +49,8 @@ def get_block_shape_and_split_kv_block(
             kv_batch_ids,
             kv_tile_ids_per_batch,
             kv_num_blocks,
+            decoder_num_blocks_x,
+            decoder_chunk_size_cpu,
             max_len_kv_cpu,
         ) = get_block_shape_and_split_kv_block_cuda(
             seq_lens_encoder,
@@ -71,6 +73,8 @@ def get_block_shape_and_split_kv_block(
             kv_batch_ids,
             kv_tile_ids_per_batch,
             kv_num_blocks,
+            decoder_num_blocks_x,
+            decoder_chunk_size_cpu,
             max_len_kv_cpu,
         )
     else:
diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py
index 03f6cea76c1..74bcee28b9a 100644
--- a/fastdeploy/model_executor/models/deepseek_v3.py
+++ b/fastdeploy/model_executor/models/deepseek_v3.py
@@ -316,30 +316,23 @@ def forward(
         mask_encoder_batch: paddle.Tensor,
     ):
         """ """
-        layernorm_out = hidden_states
-        fmha_out = paddle.zeros(
-            shape=[
-                layernorm_out.shape[0],
-                self.num_attention_heads_tp * self.v_head_dim,
-            ],
-            dtype=layernorm_out.dtype,
-        )
-
-        if forward_meta.max_len_tensor_cpu[1]:  # max_enc_len_this_time
-            query = self.q_a_proj(layernorm_out)
-            query = self.q_a_layernorm(query)
-            query = self.q_b_proj(query)
+        fmha_out = None
+        # NOTE: (changwenbin) Bring out the public calculation in PD MIX to avoid repeated calculation.
+        query = self.q_a_proj(hidden_states)
+        query = self.q_a_layernorm(query)
+        query = self.q_b_proj(query)
+        query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim])
+        query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1)
 
-            query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim])
-            query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1)
+        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+        compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1)
+        key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim])
+        compressed_kv = self.kv_a_layernorm(compressed_kv)
 
-            compressed_kv = self.kv_a_proj_with_mqa(layernorm_out)
-            compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1)
-            key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim])
-            compressed_kv = self.kv_a_layernorm(compressed_kv)
-
-            query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe)
+        query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe)
 
+        if forward_meta.max_len_tensor_cpu[1]:  # max_enc_len_this_time
+            # NOTE: (changwenbin) We will take the public part
             key_value = self.kv_b_proj(compressed_kv)
             key_value = key_value.reshape(
                 [
@@ -371,23 +364,10 @@ def forward(
             fmha_out_prefill = fmha_out_prefill.reshape([-1, self.num_attention_heads_tp * self.v_head_dim])
             fmha_out_prefill = fmha_out_prefill * mask_encoder_batch.cast(fmha_out_prefill.dtype)
 
-            fmha_out = fmha_out + fmha_out_prefill
-        if forward_meta.max_len_tensor_cpu[2]:  # max_dec_len_this_time
-            query = self.q_a_proj(layernorm_out)
-            query = self.q_a_layernorm(query)
-            ln_out_or_q_c = query
-
-            compressed_kv = self.kv_a_proj_with_mqa(layernorm_out)
-            compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1)
-            key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim])
-            compressed_kv = self.kv_a_layernorm(compressed_kv)
-
-            query = self.q_b_proj(ln_out_or_q_c)
-            query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim])
-
-            query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1)
-            query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe)
+            fmha_out = fmha_out_prefill
 
+        if forward_meta.max_len_tensor_cpu[2]:  # max_dec_len_this_time
+            # NOTE: (changwenbin) We will take the public part
             q_nope_out = self.kv_b_proj_bmm(query_nope.transpose([1, 0, 2]), proj_type="k").transpose([1, 0, 2])
 
             q_input = paddle.concat([q_nope_out, query_pe], axis=-1)
@@ -416,7 +396,10 @@ def forward(
                 .transpose([1, 0, 2])
                 .reshape([-1, self.num_attention_heads_tp * self.v_head_dim])
             )
-            fmha_out = fmha_out + fmha_out_decode
+            if fmha_out is None:
+                fmha_out = fmha_out_decode
+            else:
+                fmha_out = fmha_out + fmha_out_decode
 
         output = self.o_proj(fmha_out)
         return output