Update __init__.py

Update setup.py
Check mtp triton cache (#10065 )
2025-03-11 11:23:29 +08:00 · 2025-03-11 11:21:46 +08:00 · 2025-03-10 20:45:25 +08:00 · 2025-03-10 16:08:18 +08:00 · 2025-03-10 14:59:59 +08:00 · 2025-03-10 14:18:30 +08:00
20 changed files with 656 additions and 275 deletions
--- a/csrc/gpu/mla_attn/batch_mla_with_paged_kv_cache.cu
+++ b/csrc/gpu/mla_attn/batch_mla_with_paged_kv_cache.cu
@@ -135,6 +135,7 @@ void BatchMLAWithPagedKVCacheKernel(
  params.batch_ids = const_cast<int*>(batch_ids.data<int>());
  params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
  params.num_blocks_x = const_cast<int*>(num_blocks_x_device.data<int>());
+  params.num_blocks_x_int = num_blocks_x;
  params.q_stride_bsz = q_head_num * q_head_dim;
  params.q_stride_head_num = q_head_dim;
  params.kv_stride_block_num = block_size * k_head_dim;
--- a/csrc/gpu/mla_attn/mla_hopper.cuh
+++ b/csrc/gpu/mla_attn/mla_hopper.cuh
@@ -99,6 +99,7 @@ struct Params {
    int max_draft_token_num;
    int chunk_size;
    int chunk_num;
+    int num_blocks_x_int;

    float sm_scale;
 };
@@ -118,7 +119,7 @@ struct Params {
    return cudaErrorNotSupported;                            \
  }

-template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, int SM_COUNT = 132>
+template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, int SM_COUNT = 132, bool USE_REG_EALLOC=false, bool USE_FIXED_BLOCK=false>
 __global__ void __launch_bounds__(Ktraits::NUM_WARPS * cutlass::NumThreadsPerWarp, 1)
 MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
                     typename CollectiveMainloop::Params const mainloop_params,
@@ -198,14 +199,63 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
  
  if (warp_group_idx == 0) {
    // producer
-    cutlass::arch::warpgroup_reg_dealloc<72>();
+    if constexpr(USE_REG_EALLOC) {
+      cutlass::arch::warpgroup_reg_dealloc<72>();
+    }
    const uint32_t warp_idx_in_warpgroup = __shfl_sync(0xffffffff, warp_idx % 4, 0);
    
    PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state<MainloopPipelineQ>();
    PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
-    for (int i = blockIdx.x; i < num_blocks_x; i += SM_COUNT) {
-      const int bid = mainloop_params.batch_ids[i];
-      const int tile_id = mainloop_params.tile_ids_per_batch[i];
+    if constexpr(USE_FIXED_BLOCK) {
+      for (int i = blockIdx.x; i < num_blocks_x; i += SM_COUNT) {
+        const int bid = mainloop_params.batch_ids[i];
+        const int tile_id = mainloop_params.tile_ids_per_batch[i];
+        const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+        const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
+        const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+        const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+        cutlass::arch::NamedBarrier::sync(Ktraits::NUM_THREADS,
+                                          /*id=*/static_cast<int>(NamedBarriers::kWG0WG1WG2Sync));
+
+        // load Q
+        collective_mainloop.load_q(
+            mainloop_params,
+            pipeline_q,
+            smem_pipe_write_q,
+            shared_storage,
+            threadIdx.x,
+            bid);
+
+        if constexpr (!use_tma_load_kv) {
+          // load kv
+          collective_mainloop.load_kv(
+              mainloop_params,
+              pipeline_kv,
+              smem_pipe_write_kv,
+              shared_storage,
+              bid,
+              seq_len_decoder_now,
+              tile_id
+          );
+        } else {
+          if (warp_idx_in_warpgroup == 0) {
+            // load kv tma
+            collective_mainloop.load_kv_tma(
+                mainloop_params,
+                pipeline_kv,
+                smem_pipe_write_kv,
+                shared_storage,
+                bid,
+                seq_len_decoder_now,
+                tile_id
+            );
+          }
+        }
+      }
+    } else {
+      const int block_id = blockIdx.x;
+      const int bid = mainloop_params.batch_ids[block_id];
+      const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
      const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
      const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
      const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
@@ -250,7 +300,9 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
    }
  } else {
    // consumer
-    cutlass::arch::warpgroup_reg_alloc<216>(); 
+    if constexpr(USE_REG_EALLOC) {
+      cutlass::arch::warpgroup_reg_alloc<216>(); 
+    }
    PipelineStateQ smem_pipe_read_q;
    PipelineState smem_pipe_read_kv;

@@ -258,11 +310,74 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
    Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_PDV{}));

    auto attention_updater = OnlineSoftmax<2 * size<1>(tOrO), /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
-    for (int i = blockIdx.x; i < num_blocks_x; i += SM_COUNT) {
+    if constexpr(USE_FIXED_BLOCK) {
+      for (int i = blockIdx.x; i < num_blocks_x; i += SM_COUNT) {
+        clear(tOrO);
+        clear(attention_updater.scores_scale);
+        const int bid = mainloop_params.batch_ids[i];
+        const int tile_id = mainloop_params.tile_ids_per_batch[i];
+        const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+        const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
+        const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+        const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+        cutlass::arch::NamedBarrier::sync(Ktraits::NUM_THREADS,
+                                          /*id=*/static_cast<int>(NamedBarriers::kWG0WG1WG2Sync));
+
+        if constexpr (BLOCK_SHAPE_KV == 64) {
+          mma_f16<Ktraits, CAUSAL>(
+            mainloop_params, 
+            pipeline_q, 
+            smem_pipe_read_q,
+            pipeline_kv, 
+            smem_pipe_read_kv,
+            tOrO, 
+            attention_updater,
+            threadIdx.x - NUM_COPY_THREADS,
+            bid,
+            seq_len_decoder_now,
+            seq_len_now,
+            tile_id,
+            shared_storage);
+        } else if (BLOCK_SHAPE_KV == 32) {
+          mma_f16_two_stages<Ktraits, CAUSAL>(
+            mainloop_params, 
+            pipeline_q, 
+            smem_pipe_read_q,
+            pipeline_kv, 
+            smem_pipe_read_kv,
+            tOrO, 
+            attention_updater,
+            threadIdx.x - NUM_COPY_THREADS,
+            bid,
+            seq_len_decoder_now,
+            seq_len_now,
+            tile_id,
+            shared_storage);
+        }
+
+        collective_epilogue.store(
+            epilogue_params, 
+            tOrO, 
+            attention_updater.get_lse(),
+            shared_storage,
+            tiled_mma_pv, 
+            threadIdx.x - NUM_COPY_THREADS,
+            bid,
+            mainloop_params.bsz,
+            seq_len_now,
+            start_token_idx,
+            tile_id,
+            seq_len_decoder_now,
+            mainloop_params.chunk_size,
+            mainloop_params.max_draft_token_num,
+            mainloop_params.o_stride_bsz);
+      }
+    } else {
+      const int block_id = blockIdx.x;
      clear(tOrO);
      clear(attention_updater.scores_scale);
-      const int bid = mainloop_params.batch_ids[i];
-      const int tile_id = mainloop_params.tile_ids_per_batch[i];
+      const int bid = mainloop_params.batch_ids[block_id];
+      const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
      const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
      const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
      const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
@@ -323,7 +438,7 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
 }


-template <typename KernelTraits, bool CAUSAL, typename Params>
+template <typename KernelTraits, bool CAUSAL, typename Params, bool USE_REG_EALLOC=false, bool USE_FIXED_BLOCK=false>
 cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
                                                           cudaStream_t stream) {
  using DTypeQ = typename KernelTraits::DTypeQ;
@@ -385,7 +500,14 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
  int act_blocks_per_sm;
  cudaOccupancyMaxActiveBlocksPerMultiprocessor(
      &act_blocks_per_sm, kernel, KernelTraits::NUM_WARPS * 32, smem_size);
-  dim3 grid_dims = {multiprocessor_count, 1, 1}; // todo: split kv
+
+  int gridx;
+  if constexpr(USE_FIXED_BLOCK) {
+    gridx = multiprocessor_count;
+  } else {
+    gridx = params.num_blocks_x_int;
+  }
+  dim3 grid_dims = {gridx, 1, 1};
  static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32;
  dim3 block_dims(ctaSize, 1, 1);
  kernel<<<grid_dims, block_dims, smem_size, stream>>>(
@@ -420,7 +542,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
  return cudaSuccess;
 }

-template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params>
+template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params, bool USE_REG_EALLOC=false, bool USE_FIXED_BLOCK=false>
 cudaError_t BatchMLAWithPagedKVCacheDispatched(Params& params, cudaStream_t stream) {
  constexpr bool CAUSAL = true;
  if constexpr (HEAD_DIM_QK == 576) {
@@ -438,7 +560,10 @@ cudaError_t BatchMLAWithPagedKVCacheDispatched(Params& params, cudaStream_t stre
                                typename Params::DTypeO,
                                typename Params::IdType, 
                                NV_TYPE>,
-          CAUSAL>(params, stream);)
+          CAUSAL,
+          Params,
+          USE_REG_EALLOC,
+          USE_FIXED_BLOCK>(params, stream);)
  } else {
    return cudaErrorNotSupported;
  }
--- a/csrc/setup.py
+++ b/csrc/setup.py
@@ -1,73 +0,0 @@
-# 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.
-
-""" setup for EfficentLLM """
-
-import os
-
-from setuptools import find_packages, setup
-
-description = "Paddlenlp_ops : inference framework implemented based on PaddlePaddle"
-VERSION = "0.0.0"
-
-
-def read(file: str):
-    """
-    read file and return content
-    """
-    current_dir = os.path.dirname(__file__)
-    path = os.path.join(current_dir, file)
-    with open(path, "r", encoding="utf-8") as f:
-        content = f.read().strip()
-    return content
-
-
-def read_version():
-    """
-    read version and return content
-    """
-    return VERSION
-
-
-def read_readme():
-    """
-    read README.md and return content
-    """
-    return read("README.md")
-
-
-setup(
-    name="paddlenlp_ops",
-    packages=find_packages(),
-    version="0.0.0",
-    author="Paddle Infernce Team",
-    author_email="paddle-inference@baidu.com",
-    description=description,
-    long_description=read_readme(),
-    long_description_content_type="text/markdown",
-    url="",
-    python_requires=">=3.8",
-    package_dir={"paddlenlp_ops": "paddlenlp_ops/"},
-    package_data={"paddlenlp_ops": ["sm70/*", "sm75/*", "sm80/*", "sm86/*", "sm89/*", "sm90/*"]},
-    include_package_data=True,
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "License :: OSI Approved :: Apache Software License",
-        "Operating System :: OS Independent",
-    ],
-    license="Apache 2.0",
-)
--- a/llm/docs/predict/deepseek.md
+++ b/llm/docs/predict/deepseek.md
@@ -30,9 +30,9 @@
 |模型名称|精度|MTP|节点数|静态图下载 model_name|
 |:------|:-:|:-:|:-:|:-:|
 | deepseek-ai/DeepSeek-R1  |weight_only_int4|否|1| deepseek-ai/DeepSeek-R1/weight_only_int4 |
-| deepseek-ai/DeepSeek-R1  |weight_only_int4|是|1| deepseek-ai/DeepSeek-R1-MTP/weight-only-int4 |
+| deepseek-ai/DeepSeek-R1  |weight_only_int4|是|1| deepseek-ai/DeepSeek-R1-MTP/weight_only_int4 |
 | deepseek-ai/DeepSeek-R1  |weight_only_int8|否|2| deepseek-ai/DeepSeek-R1-2nodes/weight_only_int8 |
-| deepseek-ai/DeepSeek-R1  |weight_only_int8|是|2| deepseek-ai/DeepSeek-R1-MTP-2nodes/weight-only-int8 |
+| deepseek-ai/DeepSeek-R1  |weight_only_int8|是|2| deepseek-ai/DeepSeek-R1-MTP-2nodes/weight_only_int8 |
 | deepseek-ai/DeepSeek-R1  |a8w8_fp8|否|2| deepseek-ai/DeepSeek-R1-2nodes/a8w8_fp8|
 | deepseek-ai/DeepSeek-R1  |a8w8_fp8|是|2| deepseek-ai/DeepSeek-R1-MTP-2nodes/a8w8_fp8|
 | deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B |weight_only_int8|-|-| deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B/weight_only_int8 |
@@ -54,7 +54,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"deepseek-ai/DeepSeek-R1/weight_only_int4"}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=8 && start_server $model_name && tail -f /dev/null'
 ```

@@ -80,7 +80,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/weight_only_int8"}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=16 && export MP_NNODE=2 && export POD_0_IP=192.168.0.1 && export POD_IPS=192.168.0.1,192.168.0.2 && start_server $model_name && tail -f /dev/null'

 # node2
@@ -88,7 +88,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/weight_only_int8"}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}"\
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=16 && export MP_NNODE=2 && export POD_0_IP=192.168.0.1 && export POD_IPS=192.168.0.1,192.168.0.2 && start_server $model_name && tail -f /dev/null'
 ```

@@ -100,7 +100,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/a8w8_fp8"}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=16 && export MP_NNODE=2 && export POD_0_IP=192.168.0.1 && export POD_IPS=192.168.0.1,192.168.0.2 && start_server $model_name  && tail -f /dev/null'

 # node2
@@ -108,7 +108,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/a8w8_fp8"}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=16 && export MP_NNODE=2 && export POD_0_IP=192.168.0.1 && export POD_IPS=192.168.0.1,192.168.0.2 && start_server $model_name  && tail -f /dev/null'
 ```

@@ -117,11 +117,12 @@ docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_P
 ### deepseek-ai/DeepSeek-R1-Distill-Qwen-14B

 ```shell
+export MODEL_PATH=${MODEL_PATH:-$PWD}
+export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-Distill-Qwen-14B/weight_only_int8"}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
-v /PATH_TO_MODEL/:/models \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
-c -ex 'model_name=${model_name:-"deepseek-ai/DeepSeek-R1-Distill-Qwen-14B/weight_only_int8"} && cd /opt/output/Serving && bash start_server.sh $model_name && tail -f /dev/null'\
-&& docker exec -it $(docker ps -lq) sh -c "while [ ! -f /opt/output/Serving/log/workerlog.0 ]; do sleep 1; done; tail -f /opt/output/Serving/log/workerlog.0"
+-v /MODEL_PATH/:/models -e "model_name=${model_name}"\
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
+-c -ex 'bash start_server.sh $model_name && tail -f /dev/null'
 ```

 ### 请求服务化
@@ -235,8 +236,17 @@ python -m paddle.distributed.launch \
 ```

 两机 WINT8-TP16 推理
+
 ```shell
-# 动态图推理
+启动2机推理 需要保证2机器节点可以互相ping通
+# 第一个节点(master)
+ping 192.168.0.1
+# 第二个节点(slave)
+ping 192.168.0.2
+```
+
+```shell
+# 动态图推理 node1和node2命令均相同
 export MODEL_TAG=deepseek-ai/DeepSeek-R1
 export QUANT_MODE=weight_only_int8
 export TOTAL_MAX_LENGTH=8192
@@ -244,8 +254,8 @@ export MAX_DEC_LEN=4096
 export FLAGS_mla_use_tensorcore=1 # only support Hopper, Amper shoule be 0
 export FLAGS_cascade_attention_max_partition_size=${TOTAL_MAX_LENGTH}
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
-  --gpus ${CUDA_VISIBLE_DEVICES} \
+python -m paddle.distributed.launch \
+  --gpus ${CUDA_VISIBLE_DEVICES} --ips "192.168.0.1,192.168.0.2"\
  predictor.py \
  --model_name_or_path ${MODEL_TAG} \
  --dtype bfloat16 \
@@ -258,13 +268,13 @@ mpirun python -m paddle.distributed.launch \
  --mla_use_matrix_absorption 1


-# 动转静导出模型
+# 动转静导出模型 node1和node2命令均相同
 export MODEL_TAG=deepseek-ai/DeepSeek-R1
 export OUTPUT_PATH=/path/to/exported_model
 export QUANT_MODE=weight_only_int8
 export TOTAL_MAX_LENGTH=8192
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  export_model.py \
  --model_name_or_path ${MODEL_TAG} \
@@ -277,7 +287,7 @@ mpirun python -m paddle.distributed.launch \
  --mla_use_matrix_absorption 1


-# 静态图推理
+# 静态图推理 node1和node2命令均相同
 export OUTPUT_PATH=/path/to/exported_model
 export QUANT_MODE=weight_only_int8
 export TOTAL_MAX_LENGTH=8192
@@ -285,7 +295,7 @@ export MAX_DEC_LEN=4096
 export FLAGS_mla_use_tensorcore=1 # only support Hopper, Amper shoule be 0
 export FLAGS_cascade_attention_max_partition_size=${TOTAL_MAX_LENGTH}
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  predictor.py \
  --model_name_or_path ${OUTPUT_PATH} \
@@ -300,8 +310,17 @@ mpirun python -m paddle.distributed.launch \
 ```

 两机 FP8-TP16 推理
+
 ```shell
-# 动态图推理
+启动2机推理 需要保证2机器节点可以互相ping通
+# 第一个节点(master)
+ping 192.168.0.1
+# 第二个节点(slave)
+ping 192.168.0.2
+```
+
+```shell
+# 动态图推理 node1和node2命令均相同
 export MODEL_TAG=deepseek-ai/DeepSeek-R1-FP8
 export QUANT_MODE=a8w8_fp8
 export TOTAL_MAX_LENGTH=8192
@@ -309,7 +328,7 @@ export MAX_DEC_LEN=4096
 export FLAGS_mla_use_tensorcore=1 # only support Hopper, Amper shoule be 0
 export FLAGS_cascade_attention_max_partition_size=${TOTAL_MAX_LENGTH}
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  predictor.py \
  --model_name_or_path ${MODEL_TAG} \
@@ -324,13 +343,13 @@ mpirun python -m paddle.distributed.launch \
  --weight_block_size 128 128


-# 动转静导出模型
+# 动转静导出模型 node1和node2命令均相同
 export MODEL_TAG=deepseek-ai/DeepSeek-R1-FP8
 export OUTPUT_PATH=/path/to/exported_model
 export QUANT_MODE=a8w8_fp8
 export TOTAL_MAX_LENGTH=8192
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  export_model.py \
  --model_name_or_path ${MODEL_TAG} \
@@ -344,7 +363,7 @@ mpirun python -m paddle.distributed.launch \
  --weight_block_size 128 128


-# 静态图推理
+# 静态图推理 node1和node2命令均相同
 export OUTPUT_PATH=/path/to/exported_model
 export QUANT_MODE=a8w8_fp8
 export TOTAL_MAX_LENGTH=8192
@@ -352,7 +371,7 @@ export MAX_DEC_LEN=4096
 export FLAGS_mla_use_tensorcore=1 # only support Hopper, Amper shoule be 0
 export FLAGS_cascade_attention_max_partition_size=${TOTAL_MAX_LENGTH}
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  predictor.py \
  --model_name_or_path ${OUTPUT_PATH} \
--- a/llm/docs/predict/inference.md
+++ b/llm/docs/predict/inference.md
@@ -120,15 +120,17 @@ PaddleNLP 提供了多种量化策略，支持 Weight Only INT8及 INT4推理，

 ### 3.4 投机解码参数

- `speculate_method`: 推理解码算法，默认值为`None`，可选的数值有`None`、`inference_with_reference`。为`None`时为正常自回归解码，为`inference_with_reference`时为基于上下文的投机解码[论文地址](https://arxiv.org/pdf/2304.04487)。
+PaddleNLP 提供了多种投机解码方法，具体细节请查阅[投机解码教程](./speculative_decoding.md).

- `speculate_max_draft_token_num`: 投机解码算法中每轮产生的最大 draft tokens 数目，默认值为 1。
+- `speculate_method`: 推理解码算法，默认值为`None`，可选的数值有`None`、`inference_with_reference`、 `mtp`、 `eagle`。为`None`时为正常自回归解码，为`inference_with_reference`时为基于上下文的投机解码[论文地址](https://arxiv.org/pdf/2304.04487)。
+
+- `speculate_max_draft_token_num`: 投机解码算法中每轮产生的最大 draft tokens 数目，默认值为 1，最大支持 5。

 - `speculate_max_ngram_size`: n-gram 匹配 draft tokens 时的最大窗口大小，默认值为`1`。inference_with_reference 算法中会先从 prompt 中使用 ngram 窗口滑动匹配 draft tokens，窗口大小和输入输出重叠程度共同决定了产生 draft tokens 的开销从而影响 inference_with_reference 算法的加速效果。

- `speculate_verify_window`: 投机解码 verify 策略默认采用 TopP + window verify 中的 window 大小，默认值为`2`。更多有关 TopP + window verify 的详细介绍参考[投机解码教程](./speculative_decoding.md)。
+- `speculate_verify_window`(暂时废弃): 投机解码 verify 策略默认采用 TopP + window verify 中的 window 大小，默认值为`2`。更多有关 TopP + window verify 的详细介绍参考[投机解码教程](./speculative_decoding.md)。

- `speculate_max_candidate_len`: 产生的最大候选 tokens 数目，根据候选 tokens 与 draft tokens 比较来进行 verify(仅在 TopP + window verify 时生效)，默认值为`5`。
+- `speculate_max_candidate_len`(暂时废弃): 产生的最大候选 tokens 数目，根据候选 tokens 与 draft tokens 比较来进行 verify(仅在 TopP + window verify 时生效)，默认值为`5`。

 - `draft_model_name_or_path`: 在`MTP`或者`EAGLE`模式下，`Draft Model`的路径。

--- a/llm/docs/predict/speculative_decoding.md
+++ b/llm/docs/predict/speculative_decoding.md
@@ -1,7 +1,29 @@
 # 投机解码教程

 投机解码是一个通过投机性地一次性猜测多个 token 然后进行验证和接收的算法，通过投机解码可以极大地减小推理时延。PaddleNLP 提供了简单、高效的投机解码推理流程。下面提供 PaddleNLP 中各种投机解码算法的使用说明。
+## 高效投机解码框架
+传统的 Draft Token 验证以及 MTP(Eagle/Draft Model) 推理方法，batch_size 一般会膨胀 Draft Token 数量的倍数，我们高效的注意力机制则可以保持原始批次大小，减少计算量，如下图所示：
+![框架图](https://github.com/user-attachments/assets/fd938c8c-fcd5-4ead-ae6e-f455af415e52)
+参考 [Vllm batch_extension](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit?tab=t.0#heading=h.kk7dq05lc6q8)
+## 参数说明
+- `speculate_method`: 推理解码算法，默认值为`None`，可选的数值有`None`、`inference_with_reference`、 `mtp`、 `eagle`。为`None`时为正常自回归解码，为`inference_with_reference`时为基于上下文的投机解码[论文地址](https://arxiv.org/pdf/2304.04487)。

+- `speculate_max_draft_token_num`: 投机解码算法中每轮产生的最大 draft tokens 数目，默认值为 1，最大支持 5。
+
+- `speculate_max_ngram_size`: n-gram 匹配 draft tokens 时的最大窗口大小，默认值为`1`。inference_with_reference 算法中会先从 prompt 中使用 ngram 窗口滑动匹配 draft tokens，窗口大小和输入输出重叠程度共同决定了产生 draft tokens 的开销从而影响 inference_with_reference 算法的加速效果。
+
+- `speculate_verify_window`(暂时废弃): 投机解码 verify 策略默认采用 TopP + window verify 中的 window 大小，默认值为`2`。更多有关 TopP + window verify 的详细介绍参考[投机解码教程](./speculative_decoding.md)。
+
+- `speculate_max_candidate_len`(暂时废弃): 产生的最大候选 tokens 数目，根据候选 tokens 与 draft tokens 比较来进行 verify(仅在 TopP + window verify 时生效)，默认值为`5`。
+
+- `draft_model_name_or_path`: 在`MTP`或者`EAGLE`模式下，`Draft Model`的路径。
+
+- `draft_model_quant_type`: 在`MTP`或者`EAGLE`模式下，`Draft Model`的推理量化精度，参考`--quant_type`。
+
+- `return_full_hidden_states`: 在`MTP`或者`EAGLE`模式下，是否返回全部的隐藏层状态，默认为`False`。
+
+Others：
+1. 投机解码目前支持的最大 batch_size 为 128
 ## Inference with reference

 该算法通过 n-gram 窗口从 prompt 中匹配 draft tokens，适合输入和输出有很大 overlap 的场景如代码编辑、文档查询等，更多信息查看查看[论文地址](https://arxiv.org/pdf/2304.04487)。
@@ -15,7 +37,7 @@ python predictor.py --model_name_or_path meta-llama/Llama-2-7b-chat --inference_

 **Note:**

-1. 该算法目前只支持 llama 系列模型。
+1. 该算法目前只支持 llama/DeepSeek-V3/R1 系列模型。
 2. 投机解码同时支持量化推理，具体命令参考[推理示例](./inference.md)，将 speculate_method 等投机解码参数加上即可。

 ## Multi-Token Prediction(MTP)
@@ -25,15 +47,27 @@ Paper：https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf

 支持 DeepSeek-V3/R1 & MTP 的推理功能，具体原理请查阅论文。

-特性功能：
-1. Base Model 验证阶段使用优化的 Attention，可一次性产出所有 Draft Token 的 logits，无需使用 batch_extension。
-2. 在 Base model 接受部分或所有 Draft Token 的情况下，MTP 均可通过一次 Forward 完成对所有 Token 的计算。
-3. 使用分离框架，即 Base Model 和 MTP 分别加载权重。
+特色功能：
+1. Base Model 验证阶段使用优化的 Attention，可一次性产出所有 Draft Token 的 logits，无需使用 batch_extension 扩大 batch_size，解决了大批次下投机解码无法加速推理性能的问题。
+2. 在 MTP 推理阶段，统一对上轮 Draft Token 接受情况的处理，例如全部拒绝、接受部分、接受全部，同样只需保持原始的批次大小，一次性处理所有输入请求。
+3. 使用分离框架，即 Base Model 和 MTP 分别加载权重，导出后兼容多种解码方式。

-### 使用命令
-**DeepSeek-V3/R1 + MTP, 两机 TP16 划分**
+### MTP 支持量化精度
+| 基础模型 | 部署机器 | 基础模型量化类型 | MTP 量化类型 |
+| --- | --- | --- | --- |
+| DeepSeek-V3 | TP8 | weight_only_int4 | weight_only_int8 |
+| DeepSeek-V3 | TP16(2*TP8) | weight_only_int8 |weight_only_int8 |
+| DeepSeek-V3 | TP16(2*TP8) | a8w8_fp8 | a8w8_fp8 |
+| DeepSeek-R1 | TP8 | weight_only_int4 | weight_only_int8 |
+| DeepSeek-R1 | TP16(2*TP8) | weight_only_int8 |weight_only_int8 |
+| DeepSeek-R1 | TP16(2*TP8) | a8w8_fp8 | a8w8_fp8 |

-1. 一键容器启动 Base Model 静态图 + MTP 动态图推理服务
+ 支持 DeepSeek-V3/R1 与 MTP 的多种推理混合精度，可通过容器部署或脚本的方式进行推理，以下为部分示例
+
+### 方法一：使用容器一键部署
+**DeepSeek-R1(weight_only_int8) + MTP(weight_only_int8), 双机 TP16 划分**
+
+1. 一键容器启动  MTP 推理服务

 ```shell
 需要保证2机器节点可以互相ping通
@@ -41,38 +75,38 @@ Paper：https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
 ping 192.168.0.1
 # 第二个节点(slave)
 ping 192.168.0.2
-model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/weight_only_int8"}
+model_name=${model_name:-"deepseek-ai/DeepSeek-R1-MTP-2nodes/weight_only_int8"}
 export POD_0_IP=master_ip
 export POD_IPS=master_ip,slave_ip # 该环境变量在2机上都需保持一致
 # 服务化默认启动端口，如果冲突可以通过export进行修改
 export SERVICE_HTTP_PORT=${PUSH_MODE_HTTP_PORT:-${SERVICE_HTTP_PORT:-"9965"}}
-# 开启 MTP
-export SPECULATE_METHOD="mtp"
+#注意 SPECULATE_MODEL_QUANT_TYPE应与MTP支持量化精度表格一致
+export SPECULATE_MODEL_QUANT_TYPE="weight_only_int8"
 # /PATH_TO_MODEL # 模型挂载路径
 # /PATH_TO_MTP # MTP 挂载路径
 ```

 ```shell
 # node1
-export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/a8w8_fp8"}
+export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-MTP-2nodes/weight_only_int8"}
 export MODEL_PATH=${MODEL_PATH:-/PATH_TO_MODEL/}
 export MODEL_MTP_PATH=${MODEL_PATH:-/PATH_TO_MTP/}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -v $MODEL_MTP_PATH:/models-mtp -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=16 && export MP_NNODE=2 && export POD_0_IP=192.168.0.1 && export POD_IPS=192.168.0.1,192.168.0.2 && export SPECULATE_MODEL_QUANT_TYPE="weight_only_int8" && export SPECULATE_METHOD="mtp" && export SPECULATE_MODEL_PATH="/models-mtp" && export SPECULATE_MAX_DRAFT_TOKEN_NUM=1 && start_server $model_name && tail -f /dev/null'

 # node2
-export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-2nodes/a8w8_fp8"}
+export model_name=${model_name:-"deepseek-ai/DeepSeek-R1-MTP-2nodes/weight_only_int8"}
 export MODEL_PATH=${MODEL_PATH:-/PATH_TO_MODEL/}
 export MODEL_MTP_PATH=${MODEL_PATH:-/PATH_TO_MTP/}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -v $MODEL_MTP_PATH:/models-mtp -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && export MP_NUM=16 && export MP_NNODE=2 && export POD_0_IP=192.168.0.1 && export POD_IPS=192.168.0.1,192.168.0.2 &&export SPECULATE_MODEL_QUANT_TYPE="weight_only_int8" && export SPECULATE_METHOD="mtp" && export SPECULATE_MODEL_PATH="/models-mtp" && export SPECULATE_MAX_DRAFT_TOKEN_NUM=1 && start_server $model_name  && tail -f /dev/null'
 ```

-### 请求服务化
+**请求服务化**

 curl 请求
 ```shell
@@ -112,19 +146,31 @@ for chunk in response:
 print("\n")
 ```

-2. Base Model 动态图 + MTP 动态图
+### 方法二：使用脚本推理测试
+#### DeepSeek_R1 动态图 + MTP 动态图, 双机 TP16
+1. DeepSeek_R1 使用 weight_only_int8, MTP 使用 weight_only_int8

 ```shell
-export MODEL_TAG=deepseek-ai/DeepSeek-V3
-export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-V3-MTP
+启动2机推理 需要保证2机器节点可以互相ping通
+# 第一个节点(master)
+ping 192.168.0.1
+# 第二个节点(slave)
+ping 192.168.0.2
+```
+
+node1 和 node2 脚本相同
+```shell
+$ cat run_dynamic_mtp.sh
+
+export MODEL_TAG=deepseek-ai/DeepSeek-R1
+export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-R1-MTP
 export QUANT_MODE=weight_only_int8
+export DRAFT_MODEL_QUANT_MODE=weight_only_int8
 export TOTAL_MAX_LENGTH=8192
 export MAX_DEC_LEN=2048
 # 算子加速策略
-export FLAGS_mla_use_tensorcore=1
-export FLAGS_cascade_attention_max_partition_size=${TOTAL_MAX_LENGTH}
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  predictor.py \
  --model_name_or_path ${MODEL_TAG} \
@@ -137,80 +183,29 @@ mpirun python -m paddle.distributed.launch \
  --max_length ${MAX_DEC_LEN} \
  --speculate_method mtp \
  --draft_model_name_or_path ${DRAFT_MODEL_TAG} \
+  --draft_model_quant_type ${DRAFT_MODEL_QUANT_MODE} \
  --speculate_max_draft_token_num 1 \
-  --speculate_max_ngram_size 3 \
  --return_full_hidden_states 1 \
  --mla_use_matrix_absorption 1
+
+$ bash run_dynamic_mtp.sh
 ```

-3.【**推荐**】Base Model 静态图 + MTP 动态图
-
-Base Model 静态图导出
-
-> 注：投机解码导出支持所有方法，因此这里 speculate_method 设为默认的 inference_with_reference 即可
+2. DeepSeek_R1 使用 FP8, MTP 使用 FP8

+node1 和 node2 脚本相同
 ```shell
-export MODEL_TAG=deepseek-ai/DeepSeek-V3
-export OUTPUT_PATH=/path/to/exported_model
-export QUANT_MODE=weight_only_int8
-export TOTAL_MAX_LENGTH=8192
-export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
-  --gpus ${CUDA_VISIBLE_DEVICES} \
-  export_model.py \
-  --model_name_or_path ${MODEL_TAG} \
-  --output_path ${OUTPUT_PATH} \
-  --dtype bfloat16 \
-  --inference_model 1 \
-  --append_attn 1 \
-  --total_max_length ${TOTAL_MAX_LENGTH} \
-  --quant_type ${QUANT_MODE} \
-  --speculate_method inference_with_reference \
-  --return_full_hidden_states 1 \
-  --mla_use_matrix_absorption 1
-```
+$ cat run_dynamic_mtp.sh

-推理脚本
-
-```shell
-export OUTPUT_PATH=/path/to/exported_model
-export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-V3-MTP
-export TOTAL_MAX_LENGTH=8192
-export MAX_DEC_LEN=2048
-export QUANT_MODE=weight_only_int8
-export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
-  --gpus ${CUDA_VISIBLE_DEVICES} \
-  predictor.py \
-  --model_name_or_path ${OUTPUT_PATH} \
-  --dtype bfloat16 \
-  --mode static \
-  --inference_model 1 \
-  --append_attn 1 \
-  --total_max_length ${TOTAL_MAX_LENGTH} \
-  --quant_type ${QUANT_MODE} \
-  --max_length 1024 \
-  --speculate_method mtp \
-  --draft_model_name_or_path ${DRAFT_MODEL_TAG} \
-  --speculate_max_draft_token_num 1 \
-  --speculate_max_ngram_size 3 \
-  --return_full_hidden_states 1 \
-  --mla_use_matrix_absorption 1
-```
-
-4. R1-FP8 Model 动态图 + MTP 动态图
-
-```shell
 export MODEL_TAG=deepseek-ai/DeepSeek-R1-FP8
 export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-R1-MTP-FP8
 export QUANT_MODE=a8w8_fp8
+export DRAFT_MODEL_QUANT_MODE=a8w8_fp8
 export TOTAL_MAX_LENGTH=8192
 export MAX_DEC_LEN=2048
 # 算子加速策略
-export FLAGS_mla_use_tensorcore=1
-export FLAGS_cascade_attention_max_partition_size=${TOTAL_MAX_LENGTH}
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  predictor.py \
  --model_name_or_path ${MODEL_TAG} \
@@ -223,25 +218,136 @@ mpirun python -m paddle.distributed.launch \
  --max_length ${MAX_DEC_LEN} \
  --speculate_method mtp \
  --draft_model_name_or_path ${DRAFT_MODEL_TAG} \
+  --draft_model_quant_type ${DRAFT_MODEL_QUANT_MODE} \
  --speculate_max_draft_token_num 1 \
-  --speculate_max_ngram_size 3 \
  --return_full_hidden_states 1 \
  --mla_use_matrix_absorption 1
+
+$ bash run_dynamic_mtp.sh
+```
+#### DeepSeek_R1 动态图 + MTP 动态图, 单机 TP8
+1. DeepSeek_R1 使用 weight_only_int8, MTP 使用 weight_only_int8
+```shell
+$ cat run_dynamic_mtp.sh
+
+export MODEL_TAG=deepseek-ai/DeepSeek-R1
+export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-R1-MTP
+export QUANT_MODE=weight_only_int4
+export DRAFT_MODEL_QUANT_MODE=weight_only_int8
+export TOTAL_MAX_LENGTH=8192
+export MAX_DEC_LEN=2048
+# 算子加速策略
+export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
+python -m paddle.distributed.launch \
+  --gpus ${CUDA_VISIBLE_DEVICES} \
+  predictor.py \
+  --model_name_or_path ${MODEL_TAG} \
+  --dtype bfloat16 \
+  --mode dynamic \
+  --inference_model 1 \
+  --append_attn 1 \
+  --total_max_length ${TOTAL_MAX_LENGTH} \
+  --quant_type ${QUANT_MODE} \
+  --max_length ${MAX_DEC_LEN} \
+  --speculate_method mtp \
+  --draft_model_name_or_path ${DRAFT_MODEL_TAG} \
+  --draft_model_quant_type ${DRAFT_MODEL_QUANT_MODE} \
+  --speculate_max_draft_token_num 1 \
+  --return_full_hidden_states 1 \
+  --mla_use_matrix_absorption 1
+
+$ bash run_dynamic_mtp.sh
 ```

-5.【**推荐**】R1 Model 静态图 + MTP 动态图
-
-R1 Model 静态图导出
-
-> 注：投机解码导出支持所有方法，因此这里 speculate_method 设为默认的 inference_with_reference 即可
+#### 【**推荐**】Base Model 静态图 + MTP 动态图, 双机 TP16
+> 注：
+1.投机解码导出的模型支持所有方法，因此这里 speculate_method 设为默认的 inference_with_reference 即可.
+2.静态图模型可从 DeepSeek-R1 导出，或直接下载已上传模型

 ```shell
+启动2机推理 需要保证2机器节点可以互相ping通
+# 第一个节点(master)
+ping 192.168.0.1
+# 第二个节点(slave)
+ping 192.168.0.2
+```
+
+1. DeepSeek-R1 使用 weight_only_int8，MTP 使用 weight_only_int8
+
+node1 和 node2 脚本相同
+```shell
+# 导出脚本
+$ cat export.sh
+
+export MODEL_TAG=deepseek-ai/DeepSeek-R1
+export OUTPUT_PATH=/path/to/exported_model
+export QUANT_MODE=weight_only_int8
+export TOTAL_MAX_LENGTH=8192
+export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
+  --gpus ${CUDA_VISIBLE_DEVICES} \
+  export_model.py \
+  --model_name_or_path ${MODEL_TAG} \
+  --output_path ${OUTPUT_PATH} \
+  --dtype bfloat16 \
+  --inference_model 1 \
+  --append_attn 1 \
+  --block_attn 1 \
+  --total_max_length ${TOTAL_MAX_LENGTH} \
+  --quant_type ${QUANT_MODE} \
+  --speculate_method inference_with_reference \
+  --return_full_hidden_states 1 \
+  --mla_use_matrix_absorption 1
+
+$ bash export.sh
+```
+
+node1 和 node2 脚本相同
+```shell
+# 推理脚本
+$ cat run_mtp_infer.sh
+
+export OUTPUT_PATH=/path/to/exported_model
+export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-R1-MTP
+export TOTAL_MAX_LENGTH=8192
+export MAX_DEC_LEN=2048
+export QUANT_MODE=weight_only_int8
+export DRAFT_MODEL_QUANT_MODE=weight_only_int8
+export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
+  --gpus ${CUDA_VISIBLE_DEVICES} \
+  predictor.py \
+  --model_name_or_path ${OUTPUT_PATH} \
+  --dtype bfloat16 \
+  --mode static \
+  --inference_model 1 \
+  --append_attn 1 \
+  --total_max_length ${TOTAL_MAX_LENGTH} \
+  --quant_type ${QUANT_MODE} \
+  --max_length 1024 \
+  --speculate_method mtp \
+  --draft_model_name_or_path ${DRAFT_MODEL_TAG} \
+  --draft_model_quant_type ${DRAFT_MODEL_QUANT_MODE} \
+  --speculate_max_draft_token_num 1 \
+  --return_full_hidden_states 1 \
+  --mla_use_matrix_absorption 1
+
+$ bash run_mtp_infer.sh
+```
+
+2. DeepSeek-R1 使用 FP8, MTP 使用 FP8, 双机 TP16
+
+node1 和 node2 脚本相同
+```shell
+# 导出脚本
+$ cat export.sh
+
 export MODEL_TAG=deepseek-ai/DeepSeek-R1-FP8
 export OUTPUT_PATH=/path/to/exported_model
 export QUANT_MODE=a8w8_fp8
 export TOTAL_MAX_LENGTH=8192
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  export_model.py \
  --model_name_or_path ${MODEL_TAG} \
@@ -254,18 +360,22 @@ mpirun python -m paddle.distributed.launch \
  --speculate_method inference_with_reference \
  --return_full_hidden_states 1 \
  --mla_use_matrix_absorption 1
+
+$ bash export.sh
 ```

-推理脚本
-
+node1 和 node2 脚本相同
 ```shell
+# 推理脚本
+$ cat run_mtp_infer.sh
+
 export OUTPUT_PATH=/path/to/exported_model
 export DRAFT_MODEL_TAG=deepseek-ai/DeepSeek-R1-MTP-FP8
 export TOTAL_MAX_LENGTH=8192
 export MAX_DEC_LEN=2048
 export QUANT_MODE=a8w8_fp8
 export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
-mpirun python -m paddle.distributed.launch \
+python -m paddle.distributed.launch --ips "192.168.0.1,192.168.0.2"\
  --gpus ${CUDA_VISIBLE_DEVICES} \
  predictor.py \
  --model_name_or_path ${OUTPUT_PATH} \
@@ -275,12 +385,13 @@ mpirun python -m paddle.distributed.launch \
  --append_attn 1 \
  --total_max_length ${TOTAL_MAX_LENGTH} \
  --quant_type ${QUANT_MODE} \
-  --max_length 1024 \
+  --max_length ${MAX_DEC_LEN} \
  --speculate_method mtp \
  --draft_model_name_or_path ${DRAFT_MODEL_TAG} \
  --draft_model_quant_type ${QUANT_MODE} \
  --speculate_max_draft_token_num 1 \
-  --speculate_max_ngram_size 3 \
  --return_full_hidden_states 1 \
  --mla_use_matrix_absorption 1
-```
+
+$ bash run_mtp_infer.sh
+```
--- a/llm/server/dockerfiles/Dockerfile_serving_cuda118_cudnn8
+++ b/llm/server/dockerfiles/Dockerfile_serving_cuda118_cudnn8
@@ -24,8 +24,9 @@ RUN cp /opt/source/PaddleNLP/llm/server/server/scripts/start_server.sh start_ser
    && chmod +x stop_server \
    && mv stop_server /usr/local/bin/ 

+RUN cp /opt/source/PaddleNLP/llm/server/server/server/download_model.py download_model.py
+
 ENV PYTHONPATH="/opt/source/PaddleNLP/llm/server/server:/opt/source/PaddleNLP"
-
-
+ENV tag="3.0.0.b4"
 ENV http_proxy=""
 ENV https_proxy=""
--- a/llm/server/dockerfiles/Dockerfile_serving_cuda124_cudnn9
+++ b/llm/server/dockerfiles/Dockerfile_serving_cuda124_cudnn9
@@ -22,9 +22,11 @@ RUN cp /opt/source/PaddleNLP/llm/server/server/scripts/start_server.sh start_ser
    && mv start_server /usr/local/bin/ \
    && cp /opt/source/PaddleNLP/llm/server/server/scripts/stop_server.sh stop_server \
    && chmod +x stop_server \
-    && mv stop_server /usr/local/bin/ 
+    && mv stop_server /usr/local/bin/
+
+RUN cp /opt/source/PaddleNLP/llm/server/server/server/download_model.py download_model.py

 ENV PYTHONPATH="/opt/source/PaddleNLP/llm/server/server:/opt/source/PaddleNLP"
-
+ENV tag="3.0.0.b4"
 ENV http_proxy=""
 ENV https_proxy=""
--- a/llm/server/docs/deploy_usage_tutorial.md
+++ b/llm/server/docs/deploy_usage_tutorial.md
@@ -16,7 +16,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"meta-llama/Meta-Llama-3-8B-Instruct-Block-Attn/float16"}
 docker run  -i --rm  --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \ 
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v2.1 /bin/bash \
 -c -ex 'start_server $model_name && tail -f /dev/null'
 ```

@@ -45,7 +45,7 @@ curl 127.0.0.1:9965/v1/chat/completions \

 为了方便部署，我们提供了 cuda12.4 与 cuda 11.8 的镜像，可以直接拉取镜像，或者使用我们提供的 `Dockerfile` [构建自定义镜像](#基于 dockerfile 创建自己的镜像)
 ```
-docker pull ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0
+docker pull ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1
 ```

 ### 准备模型
@@ -62,12 +62,37 @@ cd /home/workspace/models_dir
 # ├── xxxx.model                 # 词表模型文件
 # ├── special_tokens_map.json    # 词表配置文件
 # ├── tokenizer_config.json      # 词表配置文件
-# ├── rank_mapping.csv           # 多卡模型会有此文件，如为单卡模型，则无此文件（可选，仅在多卡部署模式下需要）
 # └── rank_0                     # 保存模型结构和权重文件的目录
 #     ├── model.pdiparams
 #     └── model.pdmodel 或者 model.json # Paddle 3.0 版本模型为model.json，Paddle 2.x 版本模型为model.pdmodel
 ```

+#### 静态图下载
+
+除了支持通过设置`model_name` 在启动时进行自动下载，服务提供脚本可以进行自行下载
+
+脚本所在路径`/opt/output/download_model.py`
+
+```
+python download_model.py \
+--model_name $model_name \
+--dir $MODEL_PATH \
+--nnodes 2 \
+--mode "master" \
+--speculate_model_path $MODEL_PATH 
+```
+
+**参数说明**
+
+| 字段名 | 字段类型 | 说明 | 是否必填 | 默认值 |
+| :---: | :-----: | :---: | :---: | :-----: |
+| model_name | str | 为指定下载模型名称，具体支持模型可查看[文档](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/server/docs/static_models.md) | 否 | deepseek-ai/DeepSeek-R1/weight_only_int4 |
+| dir | str | 模型存储地址 | 否 | downloads |
+| nnodes | int | 节点个数 | 否 | 1 |
+| mode | str | 下载模式用于区分多机的不同节点 | 否 | 仅支持 master 和 slave 两个值 |
+| speculate_model_path | str | 投机解码模型存储路径 | 否 | None |
+
+
 ### 创建容器

 创建容器之前，请检查 Docker 版本和 GPU 环境，确保 Docker 支持 `--gpus all` 参数。
@@ -81,7 +106,7 @@ docker run --gpus all \
    --network=host \
    --shm-size=5G \
    -v /home/workspace/models_dir:/models/ \
-    -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 bash
+    -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 bash

 # 进入容器，检查GPU环境和模型挂载是否正常
 docker exec -it paddlenlp_serving /bin/bash
--- a/llm/server/docs/general_model_inference.md
+++ b/llm/server/docs/general_model_inference.md
@@ -20,7 +20,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"meta-llama/Meta-Llama-3-8B-Instruct-Append-Attn/bfloat16"}
 docker run  -i --rm  --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'start_server $model_name && tail -f /dev/null'
 ```
 v100
@@ -29,7 +29,7 @@ export MODEL_PATH=${MODEL_PATH:-$PWD}
 export model_name=${model_name:-"meta-llama/Meta-Llama-3-8B-Instruct-Block-Attn/float16"}
 docker run  -i --rm  --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
 -v $MODEL_PATH:/models -e "model_name=${model_name}" \ 
-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v1.0 /bin/bash \
+-dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v2.1 /bin/bash \
 -c -ex 'start_server $model_name && tail -f /dev/null'
 ```

@@ -52,7 +52,7 @@ a100
 ```shell
 export MODEL_PATH=${MODEL_PATH:-$PWD}
 docker run  -i --rm  --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'cd /opt/source/PaddleNLP &&export PYTHONPATH=$PWD:$PYTHONPATH && cd llm && python3 predict/export_model.py --model_name_or_path meta-llama/Meta-Llama-3-8B-Instruct --output_path /models --dtype bfloat16 --inference_model 1 --append_attn 1'\
 && docker logs -f $(docker ps -lq)
 ```
@@ -62,7 +62,7 @@ v100
 ```shell
 export MODEL_PATH=${MODEL_PATH:-$PWD}
 docker run  -i --rm  --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v1.0 /bin/bash \
+-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v2.1 /bin/bash \
 -c -ex 'cd /opt/source/PaddleNLP &&export PYTHONPATH=$PWD:$PYTHONPATH&& cd llm && python3 predict/export_model.py --model_name_or_path meta-llama/Meta-Llama-3-8B-Instruct --output_path /models --dtype float16 --inference_model 1 --block_attn'\
 && docker logs -f $(docker ps -lq)
 ```
@@ -75,7 +75,7 @@ a100
 ```shell
 export MODEL_PATH=${MODEL_PATH:-$PWD}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 /bin/bash \
+-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 /bin/bash \
 -c -ex 'start_server && tail -f /dev/null'
 ```

@@ -83,7 +83,7 @@ v100
 ```shell
 export MODEL_PATH=${MODEL_PATH:-$PWD}
 docker run --gpus all --shm-size 32G --network=host --privileged --cap-add=SYS_PTRACE \
-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v1.0 /bin/bash \
+-v $MODEL_PATH/:/models -dit ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v2.1 /bin/bash \
 -c -ex 'start_server && tail -f /dev/null'
 ```

@@ -103,5 +103,5 @@ curl 127.0.0.1:9965/v1/chat/completions \

 |cuda版本| 支持硬件架构|镜像地址|支持的典型设备|
 |:------|:-:|:-:|:-:|
-| cuda11.8 | 70 75 80 86 |ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v1.0 |V100，T4，A100，A30，A10 |
-| cuda12.4 | 80 86 89 90 |ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v1.0 |A100，A30，A10,L20，H20，H100 |
+| cuda11.8 | 70 75 80 86 |ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda118-cudnn8-v2.1 |V100，T4，A100，A30，A10 |
+| cuda12.4 | 80 86 89 90 |ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddlenlp:llm-serving-cuda124-cudnn9-v2.1 |A100，A30，A10,L20，H20，H100 |
--- a/llm/server/docs/static_models.md
+++ b/llm/server/docs/static_models.md
@@ -17,9 +17,9 @@
 |模型名称|精度|MTP|节点数|静态图下载 model_name|
 |:------|:-:|:-:|:-:|:-:|
 | deepseek-ai/DeepSeek-R1  |weight_only_int4|否|1| deepseek-ai/DeepSeek-R1/weight_only_int4 |
-| deepseek-ai/DeepSeek-R1  |weight_only_int4|是|1| deepseek-ai/DeepSeek-R1-MTP/weight-only-int4 |
+| deepseek-ai/DeepSeek-R1  |weight_only_int4|是|1| deepseek-ai/DeepSeek-R1-MTP/weight_only_int4 |
 | deepseek-ai/DeepSeek-R1  |weight_only_int8|否|2| deepseek-ai/DeepSeek-R1-2nodes/weight_only_int8 |
-| deepseek-ai/DeepSeek-R1  |weight_only_int8|是|2| deepseek-ai/DeepSeek-R1-MTP-2nodes/weight-only-int8 |
+| deepseek-ai/DeepSeek-R1  |weight_only_int8|是|2| deepseek-ai/DeepSeek-R1-MTP-2nodes/weight_only_int8 |
 | deepseek-ai/DeepSeek-R1  |a8w8_fp8|否|2| deepseek-ai/DeepSeek-R1-2nodes/a8w8_fp8|
 | deepseek-ai/DeepSeek-R1  |a8w8_fp8|是|2| deepseek-ai/DeepSeek-R1-MTP-2nodes/a8w8_fp8|
 | deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B |weight_only_int8|-|-| deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B/weight_only_int8 |
--- a/llm/server/server/scripts/start_server.sh
+++ b/llm/server/server/scripts/start_server.sh
@@ -31,9 +31,6 @@ export MAX_SEQ_LEN=${MAX_SEQ_LEN:-"8192"}
 export MAX_DEC_LEN=${MAX_DEC_LEN:-"8192"}
 export BATCH_SIZE=${BATCH_SIZE:-"20"}
 export BLOCK_BS=${BLOCK_BS:-"4"}
-export BLOCK_SIZE=${BLOCK_SIZE:-"64"}
-export DTYPE=${DTYPE:-"bfloat16"}
-export USE_CACHE_KV_INT8=${USE_CACHE_KV_INT8:-"0"}  # c8 model requires configuration 1
 export BLOCK_RATIO=${BLOCK_RATIO:-"0.75"}
 export ENC_DEC_BLOCK_NUM=${ENC_DEC_BLOCK_NUM:-"4"}
 export MAX_PREFILL_BATCH=${MAX_PREFILL_BATCH:-"4"}
@@ -59,6 +56,10 @@ export INTER_PROC_PORT=${INTER_QUEUE_PORT:-${INTER_PROC_PORT:-"8813"}}
 export SERVICE_HTTP_PORT=${PUSH_MODE_HTTP_PORT:-${SERVICE_HTTP_PORT:-"9965"}}


+if [ ! -d "llm_model" ];then
+    ln -s /opt/source/PaddleNLP/llm/server/server/llm_model llm_model
+fi
+
 mkdir -p log
 rm -rf console.log log/*
 rm -rf /dev/shm/*
@@ -84,7 +85,7 @@ fi



-tritonserver --exit-timeout-secs 100 --cuda-memory-pool-byte-size 0:0 --cuda-memory-pool-byte-size 1:0 \
+tritonserver --exit-timeout-secs 100000 --cuda-memory-pool-byte-size 0:0 --cuda-memory-pool-byte-size 1:0 \
                 --cuda-memory-pool-byte-size 2:0 --cuda-memory-pool-byte-size 3:0 --cuda-memory-pool-byte-size 4:0 \
                 --cuda-memory-pool-byte-size 5:0 --cuda-memory-pool-byte-size 6:0 --cuda-memory-pool-byte-size 7:0 \
                 --pinned-memory-pool-byte-size 0 --model-repository llm_model/ \
--- a/llm/server/server/server/download_model.py
+++ b/llm/server/server/server/download_model.py
@@ -0,0 +1,149 @@
+import requests
+import os
+from tqdm import tqdm
+import argparse
+import hashlib
+import re
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="download models")
+    parser.add_argument('-m', '--model_name', default='deepseek-ai/DeepSeek-R1/weight_only_int4',
+                       help="model_name")
+    parser.add_argument('-d', '--dir', default='downloads',
+                       help="save dir")
+    parser.add_argument('-n', '--nnodes', type=int, default=1,
+                       help="the number of node")
+    parser.add_argument('-M', '--mode', default="master", choices=["master", "slave"],
+                       help="only support in 2 nodes model. There are two modes, master or slave.")
+    parser.add_argument('-s', '--speculate_model_path', default=None,
+                       help="speculate model path")
+    return parser.parse_args()
+
+def calculate_md5(file_path, chunk_size=8192):
+    hasher = hashlib.md5()
+    with open(file_path, 'rb') as f:
+        for chunk in iter(lambda: f.read(chunk_size), b''):
+            hasher.update(chunk)
+    return hasher.hexdigest()
+
+def download_file(url, save_path, md5sum):
+    """download file"""
+    try:
+        with requests.get(url, stream=True) as response:
+            response.raise_for_status()
+            if os.path.exists(save_path):
+                current_md5sum = calculate_md5(save_path)
+                if md5sum != current_md5sum:
+                    os.remove(save_path)
+                    print("not complete file! start to download again")
+                else:
+                    print(f"{save_path} already exists and md5sum matches")
+                    return save_path
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+            file_name = save_path.split('/')[-1]
+            total_size = int(response.headers.get('content-length', 0))
+            progress_bar = tqdm(
+                total=total_size, 
+                unit='iB', 
+                unit_scale=True,
+                desc=f"download {file_name}"
+            )
+
+            with open(save_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=1024):
+                    if chunk:
+                        f.write(chunk)
+                        progress_bar.update(len(chunk))
+            
+            progress_bar.close()
+            if total_size != 0 and os.path.getsize(save_path) != total_size:
+                raise RuntimeError("not complete")
+                
+            return save_path
+    except Exception as e:
+        if save_path and os.path.exists(save_path):
+            os.remove(save_path)
+        return None
+
+def download_from_txt(base_url, save_dir, model_name=None):
+    txt_url = base_url + "/file_list.txt"
+    print(f"{txt_url}")
+    try:
+        response = requests.get(txt_url)
+        response.raise_for_status()
+        files_name = response.text.splitlines()
+        files_name  = [file.strip() for file in files_name if file.strip()]
+
+        md5sum = [file_name.rsplit(':', 1)[-1] for file_name in files_name]
+        file_name = [file_name.rsplit(':', 1)[0] for file_name in files_name]
+
+        if not files_name:
+            print("No valid files found.")
+            return
+
+        print(f"Found {len(files_name)} files")
+
+        for i in range(len(file_name)): 
+            cur_url = base_url + f"/{file_name[i]}"
+            path = download_file(cur_url, os.path.join(save_dir, file_name[i]), md5sum[i])
+            if path:
+                print(f"[✓] Success: {path}")
+            else:
+                print(f"[×] Failed: {cur_url}")
+    except requests.exceptions.RequestException as e:
+        raise Exception(f"Failed to download file list from {txt_url}: {str(e)}")
+
+def main():
+    args = parse_arguments()
+    print(f"Save Path: {os.path.abspath(args.dir)}")
+
+    # make dir
+    path = os.path.join(args.dir, args.model_name)
+    os.makedirs(path, exist_ok=True)
+
+    model_name = args.model_name
+    env = os.environ
+    # Define supported model patterns
+    supported_patterns = [
+        r".*Qwen.*", 
+        r".+Llama.+",
+        r".+Mixtral.+", 
+        r".+DeepSeek.+",
+    ]
+    
+    # Check if model_name matches any supported pattern
+    if not any(re.match(pattern, model_name) for pattern in supported_patterns):
+        raise ValueError(
+            f"{model_name} is not in the supported list. Currently supported models: Qwen, Llama, Mixtral, DeepSeek. Please check the model name from this document https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/server/docs/static_models.md"
+        )
+    print(f"Start downloading model: {model_name}")
+    tag = env.get("tag")
+    base_url = f"https://paddlenlp.bj.bcebos.com/models/static/{tag}/{model_name}"
+    temp_file = None
+    if args.nnodes == 1:
+        temp_file = "model"
+    elif args.nnodes > 1:
+        if args.mode == "master":
+            temp_file = "node1"
+        elif args.mode == "slave":
+            temp_file = "node2"
+        else:
+            raise ValueError(f"Invalid mode: {args.mode}. Mode must be 'master' or 'slave'.")
+    else:
+        raise ValueError(f"Invalid nnodes: {args.nnodes}. nnodes must be >= 1.")
+
+    if temp_file:
+        model_url = base_url + f"/{temp_file}"
+        download_from_txt(model_url, path)
+    else:
+        print(f"Don't support download the {model_name} in mode {args.mode}")
+
+    if args.speculate_model_path:
+        os.makedirs(args.speculate_model_path, exist_ok=True)
+        print(f"Start downloading mtp model: {model_name}")
+        model_url = base_url + "/mtp"
+        download_from_txt(model_url, args.speculate_model_path)
+
+if __name__ == "__main__":
+    main()
--- a/llm/server/server/server/engine/config.py
+++ b/llm/server/server/server/engine/config.py
@@ -17,7 +17,8 @@ import os
 from datetime import datetime
 import re

-from server.utils import model_server_logger, download_model
+from server.utils import model_server_logger
+from server.download_model import download_from_txt
 from paddlenlp.experimental.transformers import SpeculateArgument
 from paddlenlp.generation import GenerationConfig

@@ -210,14 +211,11 @@ class Config:
            f.close()


-    def _get_download_model(self):
+    def _get_download_model(self, model_type="default"):
        env = os.environ
-        model_name=env.get("model_name")
-        if not model_name:
-            raise Exception(f"Model Dir is empty")
        # Define supported model patterns
        supported_patterns = [
-            r".+Qwen.+", 
+            r".*Qwen.*", 
            r".+Llama.+",
            r".+Mixtral.+", 
            r".+DeepSeek.+",
@@ -233,21 +231,19 @@ class Config:
        base_url=f"https://paddlenlp.bj.bcebos.com/models/static/{tag}/{model_name}"
        if self.nnode == 1:
            # single node model
-            temp_tar = "model.tar"
+            temp_file = "model"
        elif env.get("POD_0_IP", "127.0.0.1") == self.host_ip:
            # Master node model
-            temp_tar = "node1.tar"
+            temp_file = "node1"
        else:
-            temp_tar = "node2.tar"
-        model_url = base_url+f"/{temp_tar}"
-        download_model(model_url, self.model_dir, temp_tar)
-
-        speculate_path = env.get("SPECULATE_MODEL_PATH")
-        if speculate_path:
-            speculate_url = base_url+"/mtp.tar"
-            os.makedirs(speculate_path, exists=True)
-            model_server_logger.info(f"Start Downloading MTP model, save path : {speculate_path}")
-            download_model(model_url, speculate_path, "mtp.tar")
+            temp_file = "node2"
+        if model_type == "default":
+            path = self.model_dir 
+        elif model_type == "speculate":
+            path = os.getenv("SPECULATE_MODEL_PATH")
+            temp_file="mtp"
+        model_url = base_url+f"/{temp_file}"
+        download_from_txt(model_url, path)



@@ -258,15 +254,8 @@ class Config:
        Returns:
            dict: the config file
        """
-        model_config_json = None
-        try:
-            model_config_json = json.load(open(self.model_config_path, "r", encoding="utf-8"))
-        except:
-            try:
-                self._get_download_model()
-                model_config_json = json.load(open(self.model_config_path, "r", encoding="utf-8"))
-            except:
-                raise
+        
+        model_config_json = json.load(open(self.model_config_path, "r", encoding="utf-8"))
        return model_config_json

    def get_speculate_config(self):
@@ -296,6 +285,9 @@ class Config:
                    speculate_model_name_or_path is not None
                ), "[eagle, mtp] method must be set by env SPECULATE_MODEL_PATH"

+            if not os.path.exists(speculate_model_name_or_path):
+                self._get_download_model(model_type="speculate")
+            
            speculate_args = SpeculateArgument.build_from_serving(
                speculate_method=speculate_method,
                speculate_max_draft_token_num=speculate_max_draft_token_num,
@@ -322,12 +314,17 @@ class Config:
        from server.utils import get_logger

        logger = get_logger("model_server", "infer_config.log")
+        env = os.environ
+        model_name=env.get("model_name")
+        if model_name:
+            self._get_download_model()
+
        config = self.get_model_config()

        # check paddle nlp version
        tag = os.getenv("tag")
        if tag not in config["paddlenlp_version"]:
-            raise Exception(f"Current image paddlenlp version {tag} doesn't match the model paddlenlp version {config['paddlenlp_version']} ")
+            logger.warning(f"Current image paddlenlp version {tag} doesn't match the model paddlenlp version {config['paddlenlp_version']} ")

        def reset_value(self, value_name, key, config):
            if key in config:
@@ -338,6 +335,12 @@ class Config:
        reset_value(self, "block_size", "infer_model_block_size", config)
        reset_value(self, "max_seq_len", "infer_model_max_seq_len", config)
        reset_value(self, "return_full_hidden_states", "return_full_hidden_states", config)
+        reset_value(self, "dtype", "infer_model_dtype", config)
+        reset_value(self, "use_cache_kv_int8", "infer_model_cachekv_int8_type", config)
+        if self.use_cache_kv_int8 == None:
+            self.use_cache_kv_int8 = 0
+        else:
+            self.use_cache_kv_int8 = 1
        if self.seq_len_limit > self.max_seq_len:
            self.seq_len_limit = self.max_seq_len
            logger.warning(f"The loading model requires len(input_ids) <= {self.max_seq_len}, now reset MAX_SEQ_LEN.")
--- a/llm/server/server/server/engine/engine.py
+++ b/llm/server/server/server/engine/engine.py
@@ -36,6 +36,10 @@ class Engine(object):
    """
    def __init__(self, cfg, token_processor):
        self.cfg = cfg
+        # Master node only
+        if self.cfg.nnode == 1 or self.cfg.host_ip == os.getenv('POD_0_IP', '127.0.0.1'):
+            self.queue_service = self._start_tasks_queue_service()
+        self.tasks_queue = TaskQueueManager(mp_num=self.cfg.mp_num, port=self.cfg.infer_port)
        self.resource_manager = ResourceManager(self.cfg)
        self.token_processor = token_processor
        self.token_processor.set_resource_manager(self.resource_manager)
@@ -50,11 +54,7 @@ class Engine(object):
        """
        assert not self.is_started, "The engine is already started.!"
        start_time = time.time()
-        # Master node only
-        if self.cfg.nnode == 1 or self.cfg.host_ip == os.getenv('POD_0_IP', '127.0.0.1'):
-            self.queue_service = self._start_tasks_queue_service()
-        self.tasks_queue = TaskQueueManager(mp_num=self.cfg.mp_num, port=self.cfg.infer_port)
-
+        
        self.token_processor.tasks_queue = self.tasks_queue
        self.infer_proc = self._start_infer_service()
        model_server_logger.info("Waitting infer processes ready...")
--- a/paddlenlp/init.py
+++ b/paddlenlp/init.py
@@ -20,7 +20,7 @@ PADDLENLP_STABLE_VERSION = "PADDLENLP_STABLE_VERSION"

 # this version is used for develop and test.
 # release version will be added fixed version by setup.py.
-__version__ = "3.0.0b3.post"
+__version__ = "3.0.0b4.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
    __version__ = __version__.replace(".post", "")
 else:
--- a/paddlenlp/transformers/moe_layer_auto.py
+++ b/paddlenlp/transformers/moe_layer_auto.py
@@ -21,6 +21,20 @@ import copy
 import paddle
 import paddle.distributed as dist
 import paddle.nn.functional as F
+
+try:
+    from paddle.distributed.auto_parallel.local_layer import LocalLayer
+except:
+
+    class LocalLayer(object):
+        """
+        A dummy class for LocalLayer, used when the actual class
+        cannot be imported.
+        """
+
+        pass
+
+
 from paddle import nn

 from .auto_utils import einsum, get_mesh
@@ -91,8 +105,8 @@ def combining(x, combine_weights, scatter_index):
    return paddle.matmul(combine_weights, x).squeeze(1)  # [seq,1,2] @ [seq,2,dim] -> [seq,1,dim]


-class LocalGatePart1(dist.LocalLayer):
-    def __init__(self, config, gate: PretrainedMoEGate, ipp=0):
+class LocalGatePart1(LocalLayer):
+    def __init__(self, config, gate: PretrainedMoEGate, ipp=None):
        mesh = get_mesh(ipp)
        out_dist_attrs = [
            (mesh, [dist.Shard(0)]),  # reshaped_input [b*s, h]
@@ -127,8 +141,8 @@ class LocalGatePart1(dist.LocalLayer):
        return reshaped_input, reshaped_scores, exp_counts, l_aux, l_zloss


-class LocalGateAndDispatch(dist.LocalLayer):
-    def __init__(self, gate: PretrainedMoEGate, ipp=0):
+class LocalGateAndDispatch(LocalLayer):
+    def __init__(self, gate: PretrainedMoEGate, ipp=None):
        mesh = get_mesh(ipp)
        out_dist_attrs = [
            (mesh, [dist.Shard(1)]),  # dispatched_input [e,c,h]
@@ -147,10 +161,10 @@ class LocalGateAndDispatch(dist.LocalLayer):
        return dispatched_input, combine_weights


-class LocalCombine(dist.LocalLayer):
-    def __init__(self, ipp=0):
-        mesh = get_mesh(ipp)
-        out_dist_attrs = [(mesh, [dist.Shard(0)])]
+class LocalCombine(LocalLayer):
+    def __init__(self, ipp=None):
+        self.mesh = get_mesh(ipp)
+        out_dist_attrs = [(self.mesh, [dist.Shard(0)])]
        grad_dist_attrs = [None, None]
        super().__init__(out_dist_attrs, grad_dist_attrs)

--- a/paddlenlp/trl/llm_utils.py
+++ b/paddlenlp/trl/llm_utils.py
@@ -825,6 +825,7 @@ def set_triton_cache(model_name_or_path, mode):
            # del old triton_ops
            shutil.rmtree(triton_kernel_cache_dir)
    elif mode == "static":
+        os.environ["TRITON_KERNEL_CACHE_DIR"] = triton_kernel_cache_dir
        for root, dirs, files in os.walk(triton_kernel_cache_dir):
            for file in files:
                if file.endswith("_package.so"):
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -12,7 +12,7 @@ yacs
 unidecode
 soundfile
 librosa
-numpy==1.23.5
+numpy
 rouge
 tiktoken
 visualdl
@@ -30,4 +30,4 @@ wget
 huggingface_hub>=0.19.2
 tiktoken
 tokenizers<=0.20.3; python_version<="3.8"
-tokenizers>=0.21,<0.22; python_version>"3.8"
+tokenizers>=0.21,<0.22; python_version>"3.8"
--- a/setup.py
+++ b/setup.py
@@ -112,7 +112,7 @@ def show():


 # only use this file to contral the version
-__version__ = "3.0.0b3.post"
+__version__ = "3.0.0b4.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
    __version__ = __version__.replace(".post", "")
 else:
Author	SHA1	Message	Date
Zhong Hui	966bd9aedd	Update __init__.py	2025-03-11 11:23:29 +08:00
Zhong Hui	1367ba4c67	Update setup.py	2025-03-11 11:21:46 +08:00
chen	d67b3b0b21	Check mtp triton cache (#10065 )	2025-03-10 20:45:25 +08:00
ltd0924	688144539e	[LLM] 增加版本号 (#10056 ) Co-authored-by: luotingdan <luotingdan@baidu.com>	2025-03-10 16:08:18 +08:00
bukejiyu	6e832e1e25	[cherry-pick] (PR10034 [server]Add a model download script and fix bugs for the server) (#10035 ) * fix * fix * fix * support download --------- Co-authored-by: ltd0924 <32387785+ltd0924@users.noreply.github.com>	2025-03-10 14:59:59 +08:00
ltd0924	277cdf4961	[Doc] 支持用户自行下载静态图 (#10046 ) * [doc] 支持用户自行下载静态图 * [doc] 支持用户自行下载静态图 * [Doc] fix format * [Doc] fix format --------- Co-authored-by: luotingdan <luotingdan@baidu.com>	2025-03-10 14:18:30 +08:00
Zhong Hui	a0408e7d72	Update requirements-dev.txt (#10051 )	2025-03-10 14:15:46 +08:00
bukejiyu	6d452602cc	update doc (#10043 )	2025-03-09 16:06:58 +08:00
lizhenyun01	15a0675dfd	fix mla nan in mtp (#10041 )	2025-03-08 22:43:22 +08:00
Yichen Zhang	fca10fde8b	[CherryPick] add import check of local_layer (#10038 )	2025-03-07 21:35:24 +08:00
ltd0924	45aca18b27	[LLM] fix llm model path and support download from txt (#10026 ) (#10029 ) * [LLM] fix llm model path and support download from txt * [LLM] fix llm model path and support download from txt	2025-03-07 17:08:49 +08:00
freeliuzc	479f9a1e1e	【Docs】Update speculate decoding docs (#10017 ) * update doc (#10014) * update spec docs	2025-03-07 13:32:08 +08:00