Delete .github/workflows/distribute-v100.yml

* change_save_meta_in_zcc_crossponding_to_paddle

* delete

* fix

* change_save_meta_in_zcc_crossponding_to_paddle and fix args

* fix deadcode

---------

Co-authored-by: liufengwei02 <liufengwei02@baidu.com>

.github/actions/check-bypass/action.yml

@@ -0,0 +1,41 @@
+name: "Check bypass"
+description: "A custom action to encapsulate PFCCLab/ci-bypass"
+inputs:
+  github-token:
+    description: "GitHub token"
+    required: true
+  workflow-name:
+    description: "Workflow name"
+    required: true
+outputs:
+  can-skip:
+    description: "Whether the workflow can be skipped."
+    value: ${{ steps.check-bypass.outputs.can-skip }}
+
+runs:
+  using: "composite"
+  steps:
+    - id: check-bypass
+      name: Check Bypass
+      env:
+        CI_TEAM_MEMBERS: '["tianshuo78520a", "swgu98", "risemeup1", "XieYunshen","luotao1","From00"]'
+      uses: PFCCLab/ci-bypass@v1
+      with:
+        github-token: ${{ inputs.github-token }}
+        non-pull-request-event-strategy: 'never-skipped'
+        type: 'composite'
+        composite-rule: |
+          {
+            "any": [
+              {
+                "type": "labeled",
+                "label": ["skip-ci: ${{ inputs.workflow-name }}", "skip-ci: all"],
+                "username": ${{ env.CI_TEAM_MEMBERS }}
+              },
+              {
+                "type": "commented",
+                "comment-pattern": [".*/skip-ci ${{ inputs.workflow-name }}.*", ".*/skip-ci all.*"],
+                "username": ${{ env.CI_TEAM_MEMBERS }}
+              }
+            ]
+          }

.github/actions/rerun-workflow/action.yml

@@ -0,0 +1,30 @@
+name: 'Rerun Workflow'
+description: 'Re-run GitHub Actions workflow for a given Pull Request'
+inputs:
+  GITHUB_TOKEN:
+    description: 'GitHub token with repo scope'
+    required: true
+  OWNER:
+    description: 'Repository owner'
+    required: true
+  REPO:
+    description: 'Repository name'
+    required: true
+  PR_ID:
+    description: 'Pull Request ID'
+    required: true
+  JOB_NAME:
+    description: 'Job name to rerun'
+    required: true
+
+runs:
+  using: 'composite'
+  steps:
+    - run: bash ./.github/actions/rerun-workflow/rerun.sh
+      shell: bash
+      env:
+        GITHUB_TOKEN: ${{ inputs.GITHUB_TOKEN }}
+        OWNER: ${{ inputs.OWNER }}
+        REPO: ${{ inputs.REPO }}
+        PR_ID: ${{ inputs.PR_ID }}
+        JOB_NAME: ${{ inputs.JOB_NAME }}

.github/actions/rerun-workflow/rerun.sh

@@ -0,0 +1,77 @@
+# Copyright (c) 2025 PaddleNLP 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.
+
+set -e
+
+COMMIT_SHA=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
+  "https://api.github.com/repos/$OWNER/$REPO/pulls/$PR_ID" | jq -r '.head.sha')
+
+echo "Commit SHA: $COMMIT_SHA"
+
+response=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
+  "https://api.github.com/repos/$OWNER/$REPO/actions/runs?head_sha=$COMMIT_SHA&per_page=100")
+
+echo "Response: $response"
+
+run_ids=$(echo "$response" | jq -r '.workflow_runs[].id')
+
+if [ -n "$run_ids" ]; then
+  echo "Found run_ids for commit $COMMIT_SHA: $run_ids"
+
+  for run_id in $run_ids; do
+    if [ "$JOB_NAME" = "all-failed" ]; then
+      echo "Rerunning all failed jobs for run_id: $run_id"
+
+      rerun_response=$(curl -X POST -s -w "%{http_code}" -o /dev/null \
+        -H "Accept: application/vnd.github.v3+json" \
+        -H "Authorization: Bearer $GITHUB_TOKEN" \
+        "https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/rerun-failed-jobs")
+      if [ "$rerun_response" -eq 201 ]; then
+        echo "Successfully requested rerun for all blocked jobs in run_id: $run_id"
+      else
+        echo "Failed to request rerun for run_id: $run_id with status code $rerun_response"
+      fi
+
+    else
+      jobs_response=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
+      "https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/jobs")
+
+      echo "Jobs Response for run_id $run_id: $jobs_response"
+
+      # if [[ "$JOB_NAME" == *"bypass"* ]]; then
+        block_jobs=$(echo "$jobs_response" | jq -r --arg job_name "$JOB_NAME" \
+        '.jobs[] | select(.name == $job_name) | .id')
+      # else
+      #   block_jobs=$(echo "$jobs_response" | jq -r --arg job_name "$JOB_NAME" \
+      #   '.jobs[] | select(.name == $job_name and .conclusion != "success") | .id')
+      # fi
+
+      if [ -n "$block_jobs" ]; then
+        echo "Found block jobs for run_id $run_id: $block_jobs"
+
+        for job_id in $block_jobs; do
+          echo "Rerunning job_id: $job_id"
+          curl -X POST -H "Accept: application/vnd.github.v3+json" \
+            -H "Authorization: token $GITHUB_TOKEN" \
+            "https://api.github.com/repos/$OWNER/$REPO/actions/jobs/$job_id/rerun"
+        done
+      else
+        echo "No block jobs found for run_id $run_id with name $JOB_NAME."
+      fi
+    fi
+  done
+else
+  echo "No matching workflow runs found for commit $COMMIT_SHA."
+  exit 1
+fi

.github/workflows/PaddleFormers-PR-Check.yml

@@ -0,0 +1,47 @@
+name: PaddleFormers PR Check
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+env:
+  BRANCH: ${{ github.base_ref }}
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  token: ${{ vars.ACTION_GITHUB_TOKEN }}
+
+jobs:
+  check-approvers:
+    name: Check approval
+    runs-on: 
+      group: APPROVAL
+    steps:
+      - name: Cleanup
+        run: |
+          rm -rf * .[^.]*
+      - name: Update paddle
+        run: |
+          wget -q --no-proxy https://xly-devops.bj.bcebos.com/PaddleTest/PaddleNLP/PaddleNLP-develop.tar.gz --no-check-certificate
+          tar zxf PaddleNLP-develop.tar.gz --strip-components=1 >/dev/null
+          rm -rf PaddleNLP-develop.tar.gz >/dev/null
+          git fetch origin pull/${PR_ID}/head
+          git checkout -b origin_pr FETCH_HEAD
+          git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+          git fetch upstream
+          git checkout -b test_pr upstream/${BRANCH}
+          git merge --no-edit origin_pr
+          git log --pretty=oneline -10
+
+      - name: Check bypass
+        id: check-bypass
+        uses: ./.github/actions/check-bypass
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          workflow-name: approval
+        
+      - name: Display Required Approvers
+        if: steps.check-bypass.outputs.can-skip != 'true'
+        run: |
+          cd scripts/ci_approval
+          bash run_ci_approval_cherry-pick.sh

.github/workflows/approval.yml

@@ -0,0 +1,47 @@
+name: Approval
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+env:
+  BRANCH: ${{ github.base_ref }}
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  token: ${{ vars.ACTION_GITHUB_TOKEN }}
+
+jobs:
+  check-approvers:
+    name: Check approval
+    runs-on: 
+      group: APPROVAL
+    steps:
+      - name: Cleanup
+        run: |
+          rm -rf * .[^.]*
+      - name: Update paddle
+        run: |
+          wget -q --no-proxy https://xly-devops.bj.bcebos.com/PaddleTest/PaddleNLP/PaddleNLP-develop.tar.gz --no-check-certificate
+          tar zxf PaddleNLP-develop.tar.gz --strip-components=1 >/dev/null
+          rm -rf PaddleNLP-develop.tar.gz >/dev/null
+          git fetch origin pull/${PR_ID}/head
+          git checkout -b origin_pr FETCH_HEAD
+          git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+          git fetch upstream
+          git checkout -b test_pr upstream/${BRANCH}
+          git merge --no-edit origin_pr
+          git log --pretty=oneline -10
+
+      - name: Check bypass
+        id: check-bypass
+        uses: ./.github/actions/check-bypass
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          workflow-name: approval
+        
+      - name: Display Required Approvers
+        if: steps.check-bypass.outputs.can-skip != 'true'
+        run: |
+          cd scripts/ci_approval
+          bash -x run_ci_approval.sh

.github/workflows/check-bypass.yml

@@ -0,0 +1,52 @@
+on:
+  workflow_call:
+    inputs:
+      workflow-name:
+        required: true
+        type: string
+    secrets:
+      github-token:
+        required: true
+    outputs:
+      can-skip:
+        description: "Whether the workflow can be skipped."
+        value: ${{ jobs.check-bypass.outputs.can-skip }}
+
+jobs:
+  check-bypass:
+    name: Check bypass
+    runs-on:
+      group: APPROVAL
+    permissions:
+      contents: read
+    env:
+      CI_TEAM_MEMBERS: '["tianshuo78520a", "swgu98", "risemeup1" , "XieYunshen","luotao1","From00"]'
+    outputs:
+      can-skip: ${{ steps.check-bypass.outputs.can-skip }}
+    steps:
+      - name: Cleanup
+        run: |
+          rm -rf * .[^.]*
+
+      - id: check-bypass
+        name: Check Bypass
+        uses: PFCCLab/ci-bypass@v1
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          non-pull-request-event-strategy: 'never-skipped'
+          type: 'composite'
+          composite-rule: |
+            {
+              "any": [
+                {
+                  "type": "labeled",
+                  "label": ["skip-ci: ${{ inputs.workflow-name }}", "skip-ci: all"],
+                  "username": ${{ env.CI_TEAM_MEMBERS }}
+                },
+                {
+                  "type": "commented",
+                  "comment-pattern": [".*/skip-ci ${{ inputs.workflow-name }}.*", ".*/skip-ci all.*"],
+                  "username": ${{ env.CI_TEAM_MEMBERS }}
+                }
+              ]
+            }

.github/workflows/lint.yml

@@ -1,17 +1,80 @@
-name: Lint
+name: Codestyle Check
 
 on: [push, pull_request]
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+env:
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  TASK: PaddleNLP-CI-Lint-${{ github.event.pull_request.number }}
+
 jobs:
+  check-bypass:
+    name: Check bypass
+    uses: ./.github/workflows/check-bypass.yml
+    with:
+      workflow-name: 'lint'
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+
   Lint:
     name: Lint
-    runs-on: ubuntu-24.04
+    needs: check-bypass
+    if: ${{ needs.check-bypass.outputs.can-skip != 'true' }}
+    runs-on: [self-hosted, ernie-cpu]
     steps:
-      - uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
-      - name: checkout develop
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          python_version: "3.10"
         run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> ${{ github.env }}
+          docker_image="iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu20.04-gcc12.2-cuda12.3-cudnn9.0-nccl2.20.3.1-openmpi4.1.5-latest"
+          docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+            -v $work_dir/../../..:$work_dir/../../.. \
+            -v $work_dir:/workspace \
+            -v /home/.cache/pip:/home/.cache/pip \
+            -e BRANCH \
+            -e PR_ID \
+            -e COMMIT_ID \
+            -e work_dir \
+            -e no_proxy \
+            -e python_version \
+            -w /workspace ${docker_image}
+
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          docker exec -t ${container_name} /bin/bash -c '
+          rm -rf * .[^.]*
+          echo "Downloading PaddleNLP.tar"
+          wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+          echo "Extracting PaddleNLP.tar"
+          tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+          source $work_dir/../../../proxy
+          cd PaddleNLP
+          git config --global user.name "PaddleCI"
+          git config --global user.email "paddle_ci@example.com"
+          git pull
+          git submodule update --init --recursive --force
+          if [ -n "${PR_ID}" ]; then
+            git fetch origin pull/${PR_ID}/head
+            git checkout -b PR_${PR_ID} FETCH_HEAD
+            git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+            git fetch upstream ${BRANCH}
+            git merge ${BRANCH} --no-edit
+            git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+          else
+            echo "Not in a pull_request event. Skipping PR-specific operations."
+          fi
+          git log --pretty=oneline -10
+
           if ! git show-ref --quiet refs/heads/develop; then \
               echo "local develop branch is missing, creating local develop branch that tracks remote develop branch"
               git fetch origin develop
@@ -19,14 +82,32 @@ jobs:
           else
             echo "local develop branch exist, skipping"
           fi
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.10'
-          cache: 'pip' # caching pip dependencies
-      - name: Install dependencies
+          '
+          
+      - name: Setup Environment
         run: |
+          docker exec -t $container_name /bin/bash -c '
+          unlink /usr/local/bin/python
+          ln -sf $(which python${python_version}) /usr/local/bin/python
+          set -e
+          python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+          python -m pip config set global.cache-dir "/home/.cache/pip"
+          source $work_dir/../../../proxy
           python -m pip install --upgrade pip
+          cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
           make install
-      - name: run the command
-        run: make lint
-    
+          '
+
+      - name: Test
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          set -e
+          source $work_dir/../../../proxy
+          cd /workspace/PaddleNLP
+          make lint
+          '
+      
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f $container_name 2>/dev/null || true

.github/workflows/llm.yml

@@ -0,0 +1,207 @@
+# Disabled LLM CI workflow - manually enabled when needed
+name: LLM CI (DISABLED)
+
+#on:
+#  pull_request:
+#    types: [opened, synchronize, reopened]
+#    branches: [develop]
+#    schedule:
+#    - cron: "2 0 * * *"
+#  workflow_call:
+#    inputs:
+#      run_downstream:
+#        required: true
+#        type: string
+#      image_name:
+#        required: true
+#        type: string
+
+#concurrency:
+#  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+#  cancel-in-progress: true
+
+#env:
+#  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  TASK: paddlenlp-CI-${{ github.event.pull_request.number }}-llm
+  ci_scripts: /workspace/PaddleNLP/scripts/regression
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  AGILE_COMPILE_BRANCH: ${{ github.event.pull_request.base.ref }}
+  CI_name: llm-ci
+  no_proxy: "localhost,bj.bcebos.com,su.bcebos.com,bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /ssd1/paddlenlp
+  HF_DATASETS_CACHE: /ssd1/paddlenlp/huggingface/datasets
+  TRANSFORMERS_CACHE: /ssd1/paddlenlp/huggingface
+  CCACHE_DIR: /home/data/gzcfs/.ccache/gpubox
+  RUN_DOWNSTREAM: ${{ inputs.run_downstream }}
+
+#defaults:
+#  run:
+#    shell: bash
+
+#jobs:
+#  llm-ci:
+#    name: llm-ci (DISABLED)
+    runs-on: [self-hosted, ernie-8gpu]
+    steps:
+      - name: Determine Image Name
+        env:
+          IMAGE_NAME: ${{ inputs.image_name }}
+        run: |
+          if [[ -n "${IMAGE_NAME}" ]]; then
+            echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITHUB_ENV"
+          else
+            echo "IMAGE_NAME=iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu18.04-gcc8.2-cuda11.8-cudnn8.6-nccl2.15.5-paddlenlp-latest" >> "$GITHUB_ENV"
+          fi
+
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          CACHE_DIR: /home/data/cfs/.cache
+          FLAGS_dynamic_static_unified_comm: "True"
+          python_version: "3.10"
+          paddle_whl: https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/latest/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> "$GITHUB_ENV"
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+              -v $work_dir/../../..:$work_dir/../../.. \
+              -v $work_dir:/workspace \
+              -v /home/.cache/pip:/home/.cache/pip \
+              -v /ssd1/paddlenlp:/ssd1/paddlenlp \
+              -v /home/data/gzcfs/.ccache/gpubox:/home/data/gzcfs/.ccache/gpubox \
+              -e BRANCH \
+              -e AGILE_COMPILE_BRANCH \
+              -e PR_ID \
+              -e COMMIT_ID \
+              -e work_dir \
+              -e ci_scripts \
+              -e no_proxy \
+              -e CI_name \
+              -e paddle_whl \
+              -e HF_ENDPOINT \
+              -e STUDIO_GIT_HOST \
+              -e PPNLP_HOME \
+              -e HF_DATASETS_CACHE \
+              -e TRANSFORMERS_CACHE \
+              -e CACHE_DIR \
+              -e FLAGS_dynamic_static_unified_comm \
+              -e python_version \
+              -w /workspace --runtime=nvidia $IMAGE_NAME
+          fi
+
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping.."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            rm -rf * .[^.]*
+            echo "Downloading PaddleNLP.tar.gz"
+            wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+            echo "Extracting PaddleNLP.tar.gz"
+            tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+            source $work_dir/../../../proxy
+            cd PaddleNLP
+            git config --global user.name "PaddleCI"
+            git config --global user.email "paddle_ci@example.com"
+            git pull
+            git submodule update --init --recursive --force
+            if [ -n "${PR_ID}" ]; then
+              git fetch origin pull/${PR_ID}/head
+              git checkout -b PR_${PR_ID} FETCH_HEAD
+              git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+              git fetch upstream ${BRANCH}
+              git merge ${BRANCH} --no-edit
+              git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+            else
+              echo "Not in a pull_request event. Skipping PR-specific operations."
+            fi
+            git log --pretty=oneline -10
+            '
+          fi
+
+      - name: Skip For Bug
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            cd /workspace/PaddleNLP
+            git revert f2477c07272d04244cd3287d1f21c70482a4a85f --no-edit  # 套件PR#10413引入bug-待修复
+            git revert d74c950e15a35b7a100d1688c89318195cc83bca --no-edit  # 框架升级后配合修改，当前因框架bug固定commit适配
+            '
+          fi
+      
+      - name: Test
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            ldconfig
+            unlink /usr/bin/python3
+            ln -sf $(which python${python_version}) /usr/bin/python3
+            pip config set global.cache-dir "/home/.cache/pip"
+            set -e
+            source $work_dir/../../../proxy
+            cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+            echo "Paddle PR#73283  PR#74484 import bug, set paddle_commit=8ae742 to skip "
+            export paddle_whl=https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/8ae7423e99b2ea96e410968a0ebb3f1795e37205/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
+            timeout 2h bash scripts/regression/run_ci.sh python${python_version} ${paddle_whl}
+            '
+          fi
+
+      - name: Upload Allure-reports & Logs
+        if: always()
+        env:
+          home_path: ${{ github.workspace }}/../../..
+          bos_file: ${{ github.workspace }}/../../../bos/BosClient.py
+          allure_file: ${{ github.workspace }}/../../../allure-2.19.0/bin/allure
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            unset http_proxy && unset https_proxy
+            if [ ! -f "${{ env.bos_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/bos_new.tar.gz https://xly-devops.bj.bcebos.com/home/bos_new.tar.gz --no-check-certificate
+              mkdir ${{ env.home_path }}/bos
+              tar xf ${{ env.home_path }}/bos_new.tar.gz -C ${{ env.home_path }}/bos
+            fi
+            if [ ! -f "${{ env.allure_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/allure-2.19.0.zip https://xly-devops.bj.bcebos.com/tools/allure-2.19.0.zip --no-check-certificate
+              unzip -q ${{ env.home_path }}/allure-2.19.0.zip -d ${{ env.home_path }}/
+            fi
+            if [[ "${{ env.RUN_DOWNSTREAM }}" == "" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}"
+            elif [[ "${{ env.RUN_DOWNSTREAM }}" == "true" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}/test_build"
+            else
+              bos_prefix="schedule/$(date +%Y%m%d)"
+            fi
+            cd /workspace/PaddleNLP/model_logs
+            for FILE in /workspace/PaddleNLP/model_logs/*; do
+              file=$(basename "$FILE")
+              python ${{ env.bos_file }} $file paddle-github-action/PR/PaddleNLP/llm/${bos_prefix}/logs
+              echo "$file: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/llm/${bos_prefix}/logs/$file"
+            done
+            cd /workspace/PaddleNLP/
+            ${{ env.allure_file }} generate result -o report
+            tar -czf products.tar.gz report model_logs
+            python ${{ env.bos_file }} products.tar.gz paddle-github-action/PR/PaddleNLP/llm/${bos_prefix}/logs
+            echo "products: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/llm/${bos_prefix}/logs/products.tar.gz"
+            '
+          fi
+
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f ${{ env.container_name }} 2>/dev/null || true

.github/workflows/rerun.yml

@@ -0,0 +1,73 @@
+name: Re-run
+
+on:
+  issue_comment:
+    types: [created]
+
+jobs:
+  re-run:
+    if: ${{ github.event.issue.pull_request && contains(github.event.comment.body, '/re-run')  && github.event.comment.user.login == github.event.issue.user.login }}
+    runs-on: [self-hosted, ernie-cpu]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Rerun all failed jobs
+        if: ${{ contains(github.event.comment.body, 'all-failed') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'all-failed'
+
+      - name: Rerun LLM
+        if: ${{ contains(github.event.comment.body, 'LLM') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'LLM CI / llm-ci'
+
+      - name: Unittest GPU
+        if: ${{ contains(github.event.comment.body, 'Unittest GPU') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Unittest GPU CI / unittest-gpu-ci'
+
+      - name: Rerun Unittest CPU
+        if: ${{ contains(github.event.comment.body, 'Unittest CPU') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Unittest CPU CI / unit-cpu-ci'
+
+      - name: Rerun Distribute (V100)
+        if: ${{ contains(github.event.comment.body, 'Distribute (V100)') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Distribute CI (V100)/ distribute-v100-ci'
+
+      - name: Rerun Codestyle Check
+        if: ${{ contains(github.event.comment.body, 'Codestyle Check') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Codestyle Check / Lint'

.github/workflows/tests.yml

@@ -1,29 +0,0 @@
-name: Test
-
-on: [push, pull_request]
-
-jobs:
-  Test:
-    name: Test
-    runs-on: ubuntu-24.04
-    permissions:
-      pull-requests: write
-      contents: read
-      id-token: write
-    steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.10'
-          cache: 'pip' # caching pip dependencies
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -r tests/requirements.txt
-          make install
-      - name: run the command
-        run: make test
-      - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v4
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/unittest-cpu.yml

@@ -0,0 +1,120 @@
+# Disabled unittest CPU workflow - manually enabled when needed
+name: Unittest CPU CI (DISABLED)
+
+#on: [pull_request]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+env:
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  TASK: PaddleNLP-CI-${{ github.event.pull_request.number }}-unittest-cpu
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /home/disk1/cache
+  HF_DATASETS_CACHE: /home/disk1/cache/huggingface/datasets
+  TRANSFORMERS_CACHE: /home/disk1/cache/huggingface
+
+#jobs:
+#  unittest-cpu-ci:
+#    name: unittest-cpu-ci (DISABLED)
+    runs-on: [self-hosted, ernie-cpu]
+    permissions:
+      pull-requests: write
+      contents: read
+      id-token: write
+    steps:
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          python_version: "3.10"
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> ${{ github.env }}
+          docker_image="iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu20.04-gcc12.2-cuda12.3-cudnn9.0-nccl2.20.3.1-openmpi4.1.5-latest"
+          docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+            -v $work_dir/../../..:$work_dir/../../.. \
+            -v $work_dir:/workspace \
+            -v /home/.cache/pip:/home/.cache/pip \
+            -v /home/disk1/cache:/home/disk1/cache \
+            -e BRANCH \
+            -e PR_ID \
+            -e COMMIT_ID \
+            -e work_dir \
+            -e no_proxy \
+            -e python_version \
+            -e HF_ENDPOINT \
+            -e STUDIO_GIT_HOST \
+            -e PPNLP_HOME \
+            -e HF_DATASETS_CACHE \
+            -e TRANSFORMERS_CACHE \
+            -w /workspace ${docker_image}
+
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          docker exec -t ${container_name} /bin/bash -c '
+          rm -rf * .[^.]*
+          echo "Downloading PaddleNLP.tar"
+          wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+          echo "Extracting PaddleNLP.tar"
+          tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+          source $work_dir/../../../proxy
+          cd PaddleNLP
+          git config --global user.name "PaddleCI"
+          git config --global user.email "paddle_ci@example.com"
+          git pull
+          git submodule update --init --recursive --force
+          if [ -n "${PR_ID}" ]; then
+            git fetch origin pull/${PR_ID}/head
+            git checkout -b PR_${PR_ID} FETCH_HEAD
+            git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+            git fetch upstream ${BRANCH}
+            git merge ${BRANCH} --no-edit
+            git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+          else
+            echo "Not in a pull_request event. Skipping PR-specific operations."
+          fi
+          git log --pretty=oneline -10
+          '
+
+      - name: Setup Environment
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          unlink /usr/bin/python3
+          ln -sf $(which python${python_version}) /usr/bin/python3
+          set -e
+          python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+          python -m pip config set global.cache-dir "/home/.cache/pip"
+          source $work_dir/../../../proxy
+          python -m pip install --upgrade pip
+          cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+          pip install -r tests/requirements.txt
+          make install
+          '
+
+      - name: Test
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          source $work_dir/../../../proxy
+          cd /workspace/PaddleNLP
+          set -e
+          make test
+          '
+
+      # - name: Upload Coverage To Codecov
+      #   uses: codecov/codecov-action@v3
+      #   env:
+      #     CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+      #   with:
+      #     token: ${{ secrets.CODECOV_TOKEN }}
+      #     files: ./PaddleNLP/coverage.xml
+      
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f $container_name 2>/dev/null || true

.github/workflows/unittest-gpu.yml

@@ -0,0 +1,204 @@
+# Disabled unittest GPU workflow - manually enabled when needed
+name: Unittest GPU CI (DISABLED)
+
+#on:
+#  pull_request:
+#    types: [opened, synchronize, reopened]
+#    branches: [develop]
+#  schedule:
+#    - cron: "3 0 * * *"
+#  workflow_call:
+#    inputs:
+#      run_downstream:
+#        required: true
+#        type: string
+#      image_name:
+#        required: true
+#        type: string
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+env:
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  TASK: paddlenlp-CI-${{ github.event.pull_request.number }}-unit-gpu
+  ci_scripts: /workspace/PaddleNLP/scripts/unit_test
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  AGILE_COMPILE_BRANCH: ${{ github.event.pull_request.base.ref }}
+  CI_name: unittest-gpu-ci
+  no_proxy: "localhost,bj.bcebos.com,su.bcebos.com,bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /ssd1/paddlenlp
+  HF_DATASETS_CACHE: /ssd1/paddlenlp/huggingface/datasets
+  TRANSFORMERS_CACHE: /ssd1/paddlenlp/huggingface
+  CCACHE_DIR: /home/data/gzcfs/.ccache/gpubox
+  RUN_DOWNSTREAM: ${{ inputs.run_downstream }}
+
+defaults:
+  run:
+    shell: bash
+
+#jobs:
+#  unittest-gpu-ci:
+#    name: unittest-gpu-ci (DISABLED)
+    runs-on: [self-hosted, ernie-8gpu]
+    steps:
+      - name: Determine Image Name
+        env:
+          IMAGE_NAME: ${{ inputs.image_name }}
+        run: |
+          if [[ -n "${IMAGE_NAME}" ]]; then
+            echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITHUB_ENV"
+          else
+            echo "IMAGE_NAME=iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu18.04-gcc8.2-cuda11.8-cudnn8.6-nccl2.15.5-paddlenlp-latest" >> "$GITHUB_ENV"
+          fi
+      
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          CACHE_DIR: /home/data/cfs/.cache
+          FLAGS_dynamic_static_unified_comm: "True"
+          python_version: "3.10"
+          paddle_whl: https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/latest/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> "$GITHUB_ENV"
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+              -v $work_dir/../../..:$work_dir/../../.. \
+              -v $work_dir:/workspace \
+              -v /home/.cache/pip:/home/.cache/pip \
+              -v /ssd1/paddlenlp:/ssd1/paddlenlp \
+              -v /home/data/gzcfs/.ccache/gpubox:/home/data/gzcfs/.ccache/gpubox \
+              -e BRANCH \
+              -e AGILE_COMPILE_BRANCH \
+              -e PR_ID \
+              -e COMMIT_ID \
+              -e work_dir \
+              -e ci_scripts \
+              -e no_proxy \
+              -e CI_name \
+              -e paddle_whl \
+              -e HF_ENDPOINT \
+              -e STUDIO_GIT_HOST \
+              -e PPNLP_HOME \
+              -e HF_DATASETS_CACHE \
+              -e TRANSFORMERS_CACHE \
+              -e CACHE_DIR \
+              -e FLAGS_dynamic_static_unified_comm \
+              -e python_version \
+              -w /workspace --runtime=nvidia $IMAGE_NAME
+          fi
+          
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping.."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            rm -rf * .[^.]*
+            echo "Downloading PaddleNLP.tar.gz"
+            wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+            echo "Extracting PaddleNLP.tar.gz"
+            tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+            source $work_dir/../../../proxy
+            cd PaddleNLP
+            git config --global user.name "PaddleCI"
+            git config --global user.email "paddle_ci@example.com"
+            git pull
+            git submodule update --init --recursive --force
+            if [ -n "${PR_ID}" ]; then
+              git fetch origin pull/${PR_ID}/head
+              git checkout -b PR_${PR_ID} FETCH_HEAD
+              git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+              git fetch upstream ${BRANCH}
+              git merge ${BRANCH} --no-edit
+              git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+            else
+              echo "Not in a pull_request event. Skipping PR-specific operations."
+            fi
+            git log --pretty=oneline -10
+            '
+          fi
+      
+      - name: Skip For Bug
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            cd /workspace/PaddleNLP
+            echo "no skip for bug"
+            '
+          fi
+      
+      - name: Test
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            ldconfig
+            unlink /usr/bin/python3
+            ln -sf $(which python${python_version}) /usr/bin/python3
+            pip config set global.cache-dir "/home/.cache/pip"
+            source $work_dir/../../../proxy
+            set -e
+            cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+            timeout 50m bash scripts/unit_test/ci_unit.sh ${paddle_whl}
+            '
+          fi
+
+      - name: Upload Allure-reports & Logs
+        if: always()
+        env:
+          home_path: ${{ github.workspace }}/../../..
+          bos_file: ${{ github.workspace }}/../../../bos/BosClient.py
+          allure_file: ${{ github.workspace }}/../../../allure-2.19.0/bin/allure
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            unset http_proxy && unset https_proxy
+            if [ ! -f "${{ env.bos_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/bos_new.tar.gz https://xly-devops.bj.bcebos.com/home/bos_new.tar.gz --no-check-certificate
+              mkdir ${{ env.home_path }}/bos
+              tar xf ${{ env.home_path }}/bos_new.tar.gz -C ${{ env.home_path }}/bos
+            fi
+            if [ ! -f "${{ env.allure_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/allure-2.19.0.zip https://xly-devops.bj.bcebos.com/tools/allure-2.19.0.zip --no-check-certificate
+              unzip -q ${{ env.home_path }}/allure-2.19.0.zip -d ${{ env.home_path }}/
+            fi
+            if [[ "${{ env.RUN_DOWNSTREAM }}" == "" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}"
+            elif [[ "${{ env.RUN_DOWNSTREAM }}" == "true" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}/test_build"
+            else
+              bos_prefix="schedule/$(date +%Y%m%d)"
+            fi
+            cd /workspace/PaddleNLP/unittest_logs
+            for FILE in /workspace/PaddleNLP/unittest_logs/*; do
+              file=$(basename "$FILE")
+              python ${{ env.bos_file }} $file paddle-github-action/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs
+              echo "$file: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs/$file"
+            done
+            cd /workspace/PaddleNLP/
+            ${{ env.allure_file }} generate result -o report
+            tar -czf products.tar.gz report unittest_logs
+            python ${{ env.bos_file }} products.tar.gz paddle-github-action/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs
+            echo "report: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs/products.tar.gz"
+            '
+          fi
+
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f $container_name 2>/dev/null || true

.gitignore

@@ -140,4 +140,4 @@ autogen/
 #fp8
 ops/csrc/fp8/deep_gemm/include/cutlass
 ops/csrc/fp8/deep_gemm/include/cute
-
+.ccls-cache

.pre-commit-config.yaml

@@ -62,3 +62,12 @@ repos:
         language: python
         files: \.(md|markdown|rst)$
         pass_filenames: true
+
+-   repo: local
+    hooks:
+    -   id: clang-format
+        name: clang-format
+        description: Format files with ClangFormat.
+        entry: bash ./tools/codestyle/clang_format.sh -i
+        language: system
+        files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|xpu|kps)$

CONTRIBUTING.md

@@ -72,7 +72,7 @@ PaddleNLP 使用 [Git 分支模型](http://nvie.com/posts/a-successful-git-branc
    类似于以下内容：
 
    ```text
-    ➜  (my-virtual-env) git commit -m "commiting my cool feature"
+    ➜  (my-virtual-env) git commit -m "committing my cool feature"
     black....................................................................Passed
     isort....................................................................Passed
     flake8...................................................................Passed
@@ -91,7 +91,7 @@ PaddleNLP 使用 [Git 分支模型](http://nvie.com/posts/a-successful-git-branc
    但大多数时候事情并没有那么顺利。当您的代码或文档不符合标准时，`pre-commit` 检查将失败。
 
    ```text
-    ➜  (my-virtual-env) git commit -m "commiting my cool feature"
+    ➜  (my-virtual-env) git commit -m "committing my cool feature"
     black....................................................................Passed
     isort....................................................................Failed
     - hook id: isort

README.md

@@ -33,17 +33,20 @@
 
 ## News 📢
 
-* **2025.03.17 《DeepSeek-R1满血版单机部署实测》** 🔥🔥🔥 飞桨框架3.0大模型推理部署全面升级，支持多款主流大模型，DeepSeek-R1满血版实现单机部署，吞吐提升一倍！欢迎广大用户开箱体验～现已开启有奖活动：完成 DeepSeek-R1-MTP 单机部署任务、提交高质量测评 blog，即可实时赢取奖金！💰💰💰
-报名[地址](https://www.wjx.top/vm/OlzzmbG.aspx#)， 活动详情：https://github.com/PaddlePaddle/PaddleNLP/issues/10166 ， 参考文档：https://github.com/PaddlePaddle/PaddleNLP/issues/10157 。
+* **2025.04.29 PaddleNLP 现已支持 Qwen3 系列模型**: Qwen3 系列模型支持持两种思考模式，预训练约 36 万亿个 token、119 种语言和方言。包括六个 Dense 模型, Qwen3-32B、Qwen3-14B、Qwen3-8B、Qwen3-4B、Qwen3-1.7B 和 Qwen3-0.6B。两个 MoE 模型的权重：Qwen3-235B-A22B，Qwen3-30B-A3B。
 
 * **2025.03.12 [PaddleNLP v3.0 Beta4](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta4)**：全面支持 DeepSeek V3/R1/R1-Distill, 及 QwQ-32B 等热门思考模型。**DeepSeek V3/R1完整版支持 FP8、INT8、4-bit 量化推理，MTP 投机解码**。单机 FP8推理输出超**1000 tokens/s**; 4-bit 推理输出超**2100 tokens/s**! 发布新版推理部署镜像，热门模型[一键部署](https://paddlenlp.readthedocs.io/zh/latest/llm/server/docs/general_model_inference.html)。推理部署[使用文档](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/predict/index.html)全面更新，体验全面提升！自研下一代通用信息抽取模型 PP-UIE [全新发布](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/application/information_extraction)，支持8K 长度信息抽取。新增大模型 Embedding 训练，支持 INF-CL 超大 batch size 训练。新增[MergeKit](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/mergekit.html)模型融合工具，缓解对齐代价。低资源训练全面优化，16G 小显存可以流畅训练。
 
-* **2025.03.06 PaddleNLP 现已支持 Qwen/QwQ-32B 模型**: 其模型参数仅有 32B，但其数学推理、编程能力和通用能力可与具备 671B 参数（其中 37B 被激活）的 DeepSeek-R1 媲美。借助 PaddleNLP 3.0套件，现可实现多种并行策略[微调训练](./llm/README.md)、[高性能推理、低比特量化](./llm/docs/predict/qwen.md)和[服务化部署](./llm/server/README.md)。
 
 * **2025.02.10 PaddleNLP 现已支持 DeepSeek-R1系列模型，[在线使用](https://aistudio.baidu.com/projectdetail/8775758)**：依托全新的 PaddleNLP 3.0套件，DeepSeek-R1系列模型现已全面支持。凭借数据并行、数据分组切分并行、模型并行、流水线并行以及专家并行等一系列先进的分布式训练能力，结合 Paddle 框架独有的列稀疏注意力掩码表示技术——FlashMask 方法，DeepSeek-R1系列模型在训练过程中显著降低了显存消耗，同时取得了卓越的训练性能提升。
 
 <details><summary> <b>点击展开</b> </summary><div>
 
+* **2025.03.17 《DeepSeek-R1满血版单机部署实测》** 🔥🔥🔥 飞桨框架3.0大模型推理部署全面升级，支持多款主流大模型，DeepSeek-R1满血版实现单机部署，吞吐提升一倍！欢迎广大用户开箱体验～现已开启有奖活动：完成 DeepSeek-R1-MTP 单机部署任务、提交高质量测评 blog，即可实时赢取奖金！💰💰💰
+报名[地址](https://www.wjx.top/vm/OlzzmbG.aspx#)， 活动详情：https://github.com/PaddlePaddle/PaddleNLP/issues/10166 ， 参考文档：https://github.com/PaddlePaddle/PaddleNLP/issues/10157 。
+
+* **2025.03.06 PaddleNLP 现已支持 Qwen/QwQ-32B 模型**: 其模型参数仅有 32B，但其数学推理、编程能力和通用能力可与具备 671B 参数（其中 37B 被激活）的 DeepSeek-R1 媲美。借助 PaddleNLP 3.0套件，现可实现多种并行策略[微调训练](./llm/README.md)、[高性能推理、低比特量化](./llm/docs/predict/qwen.md)和[服务化部署](./llm/server/README.md)。
+
 * **2025.02.20 🔥🔥《PP-UIE 信息抽取智能引擎全新升级》** 强化零样本学习能力，支持极少甚至零标注数据实现高效冷启动与迁移学习，显著降低数据标注成本；具备处理长文本能力，支持 8192 个 Token 长度文档信息抽取，实现跨段落识别关键信息，形成完整理解；提供完整可定制化的训练和推理全流程，训练效率相较于 LLama-Factory 实现了1.8倍的提升。
 2月26日（周三）19：00为您深度解析全新 PP-UIE 技术方案及在部署方面的功能、优势与技巧。报名链接：https://www.wjx.top/vm/mBKC6pb.aspx?udsid=606418
 
@@ -119,6 +122,7 @@
 |           [Qwen2.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-7B-Instruct-1M, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-14B-Instruct-1M, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, Qwen/Qwen2.5-72B, Qwen/Qwen2.5-72B-Instruct          |
 |         [Qwen2.5-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)         | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct, Qwen/Qwen2.5-Math-72B, Qwen/Qwen2.5-Math-72B-Instruct, Qwen/Qwen2.5-Math-RM-72B                                                                                                                                                                                                 |
 |        [Qwen2.5-Coder](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)         | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                                                                                                                                              |
+|           [Qwen3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)              | Qwen/Qwen3-0.6B, Qwen/Qwen3-1.7B, Qwen/Qwen3-4B, Qwen/Qwen3-8B, Qwen/Qwen3-14B, Qwen/Qwen3-32B, Qwen/Qwen3-30B-A3B, Qwen/Qwen3-235B-A22B, Qwen/Qwen3-0.6B-Base, Qwen/Qwen3-1.7B-Base, Qwen/Qwen3-4B-Base, Qwen/Qwen3-8B-Base, Qwen/Qwen3-14B-Base, Qwen/Qwen3-30B-A3B-Base                                                                                                                    |
 |             [QwQ](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)              | Qwen/QwQ-32B, Qwen/QwQ-32B-Preview                                                                                                                                                                                                                                                                                                                                                            |
 |            [Yuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/yuan/)             | IEITYuan/Yuan2-2B, IEITYuan/Yuan2-51B, IEITYuan/Yuan2-102B                                                                                                                                                                                                                                                                                                                                    |
 

csrc/gpu/all_reduce.cu

@@ -72,12 +72,14 @@ void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
                           reinterpret_cast<half*>(out.data()), out.numel());
       break;
     }
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
     case phi::DataType::BFLOAT16: {
       fa->allreduce<nv_bfloat16>(
           stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
           reinterpret_cast<nv_bfloat16*>(out.data()), out.numel());
       break;
     }
+#endif
     default:
       throw std::runtime_error(
           "custom allreduce only supports float32, float16 and bfloat16");

csrc/gpu/all_reduce.cuh

@@ -98,7 +98,7 @@ DINLINE half& assign_add(half& a, half b) {
 }
 DINLINE float& assign_add(float& a, float b) { return a += b; }
 
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
 DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
 template <>
 DINLINE nv_bfloat16 downcast_s(float val) {

csrc/gpu/append_attention.cu

@@ -53,6 +53,7 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
     const paddle::optional<paddle::Tensor>& cache_v_zp,
     const paddle::optional<paddle::Tensor>& out_linear_shifts,
     const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
     const int max_input_length,
@@ -140,6 +141,7 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
           cache_v_quant_scales,
           cache_k_zp,
           cache_v_zp,
+          excess_blocks,
           cache_quant_type_str,
           kv_num_blocks_data,
           max_input_length,
@@ -167,6 +169,7 @@ std::vector<paddle::Tensor> AppendAttentionKernel(
           cache_v_quant_scales,
           cache_k_zp,
           cache_v_zp,
+          excess_blocks,
           cache_quant_type_str,
           kv_num_blocks_data,
           max_input_length,
@@ -568,6 +571,7 @@ std::vector<paddle::Tensor> AppendAttention(
     const paddle::optional<paddle::Tensor>& cache_v_zp,
     const paddle::optional<paddle::Tensor>& out_linear_shifts,
     const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& compute_dtype,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
@@ -632,6 +636,7 @@ std::vector<paddle::Tensor> AppendAttention(
           cache_v_zp,
           out_linear_shifts,
           out_linear_smooths,
+          excess_blocks,
           cache_quant_type_str,
           use_neox_rotary_style,
           max_input_length,
@@ -679,6 +684,7 @@ std::vector<paddle::Tensor> AppendAttention(
           cache_v_zp,
           out_linear_shifts,
           out_linear_smooths,
+          excess_blocks,
           cache_quant_type_str,
           use_neox_rotary_style,
           max_input_length,
@@ -727,6 +733,7 @@ std::vector<paddle::Tensor> AppendAttention(
             cache_v_zp,
             out_linear_shifts,
             out_linear_smooths,
+            excess_blocks,
             cache_quant_type_str,
             use_neox_rotary_style,
             max_input_length,
@@ -773,6 +780,7 @@ std::vector<paddle::Tensor> AppendAttention(
             cache_v_zp,
             out_linear_shifts,
             out_linear_smooths,
+            excess_blocks,
             cache_quant_type_str,
             use_neox_rotary_style,
             max_input_length,
@@ -831,7 +839,8 @@ std::vector<std::vector<int64_t>> AppendAttentionInferShape(
     const paddle::optional<std::vector<int64_t>>& cache_k_zp_shape,
     const paddle::optional<std::vector<int64_t>>& cache_v_zp_shape,
     const paddle::optional<std::vector<int64_t>>& out_linear_shifts_shape,
-    const paddle::optional<std::vector<int64_t>>& out_linear_smooths_shape) {
+    const paddle::optional<std::vector<int64_t>>& out_linear_smooths_shape,
+    const paddle::optional<std::vector<int64_t>>& excess_blocks_shape) {
   const int token_num = qkv_shape[0];
   const int kv_num_heads = key_cache_shape[1];
   const int head_dim_qk = key_cache_shape[3];
@@ -876,6 +885,7 @@ std::vector<paddle::DataType> AppendAttentionInferDtype(
     const paddle::optional<paddle::DataType>& cache_v_zp_dtype,
     const paddle::optional<paddle::DataType>& out_linear_shifts_dtype,
     const paddle::optional<paddle::DataType>& out_linear_smooths_dtype,
+    const paddle::optional<paddle::DataType>& excess_blocks_dtype,
     const std::string& compute_dtype,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
@@ -949,7 +959,8 @@ PD_BUILD_OP(append_attention)
              paddle::Optional("cache_k_zp"),
              paddle::Optional("cache_v_zp"),
              paddle::Optional("out_linear_shifts"),
-             paddle::Optional("out_linear_smooths")})
+             paddle::Optional("out_linear_smooths"),
+             paddle::Optional("excess_blocks")})
     .Outputs({"fmha_out", "qkv_out", "key_cache_out", "value_cache_out"})
     .SetInplaceMap({{"key_cache", "key_cache_out"},
                     {"value_cache", "value_cache_out"}})

csrc/gpu/append_attn/encoder_write_cache_with_rope_impl.cuh

@@ -691,6 +691,99 @@ __global__ void cache_kernel(
   }
 }
 
+template <typename T, int VecSize = 1>
+__global__ void cache_use_excess_kernel(
+    const T *__restrict__ qkv,    // [num_tokens, num_heads + 2 * kv_num_heads,
+                                  // head_size]
+    T *__restrict__ key_cache,    // [num_blocks, kv_num_heads, block_size,
+                                  // head_size]
+    T *__restrict__ value_cache,  // [num_blocks, kv_num_heads, block_size,
+                                  // head_size]
+    const int *__restrict__ block_tables,      // [bsz, max_blocks_per_seq]
+    const int *__restrict__ padding_offsets,   // [num_tokens]
+    const int *__restrict__ cum_offsets,
+    const int *__restrict__ seq_lens,          // [bsz]
+    const int *__restrict__ seq_lens_decoder,  // [bsz]
+    const int *__restrict__ excess_blocks,  // [bsz, excess_num]
+    const int max_seq_len,
+    const int max_blocks_per_seq,
+    const int num_heads,
+    const int head_size,
+    const int block_size,
+    const uint32_t elem_cnt,
+    const int kv_num_heads,
+    const int token_num,
+    const int excess_num) {
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec;
+
+  uint32_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
+  const uint32_t hidden_size = kv_num_heads * head_size;
+  const uint32_t offset = 2 * hidden_size;
+  for (uint32_t linear_index = global_thread_idx * VecSize,
+                step = gridDim.x * blockDim.x * VecSize;
+       linear_index < elem_cnt;
+       linear_index += step) {
+    uint32_t token_idx = linear_index / offset;
+    const uint32_t bias = linear_index % offset;
+    const uint32_t qkv_id = bias / hidden_size;  // skip q
+    const uint32_t qkv_bias = bias % hidden_size;
+    const uint32_t hi = qkv_bias / head_size;
+    const uint32_t h_bias = qkv_bias % head_size;
+
+    uint32_t block_idx, block_offset;
+
+    if (token_idx < token_num) {
+      const uint32_t ori_token_idx = token_idx + padding_offsets[token_idx];
+      const uint32_t ori_bi = ori_token_idx / max_seq_len;
+      const uint32_t last_offset = seq_lens[ori_bi] % block_size;
+      if (seq_lens[ori_bi] == 0) continue;
+
+      const int32_t *block_table_now = nullptr;
+      block_table_now = block_tables + ori_bi * max_blocks_per_seq;
+      const uint32_t ori_seq_id =
+          ori_token_idx % max_seq_len + seq_lens_decoder[ori_bi];
+      if (ori_seq_id >= seq_lens[ori_bi] - last_offset) continue;
+
+      block_idx = block_table_now[ori_seq_id / block_size];
+      block_offset = ori_seq_id % block_size;
+    } else {
+      const uint32_t excess_token_id = token_idx - token_num;
+      const uint32_t ori_bi = excess_token_id / (excess_num * block_size);
+      const uint32_t last_offset = seq_lens[ori_bi] % block_size;
+      if (seq_lens[ori_bi] == 0) continue;
+
+      const uint32_t excess_id =
+          (excess_token_id % (excess_num * block_size)) / block_size;
+      const uint32_t excess_token_offset = excess_token_id % block_size;
+
+      if (excess_token_offset < last_offset) {
+        token_idx = ori_bi * max_seq_len - cum_offsets[ori_bi] +
+                    seq_lens[ori_bi] - last_offset + excess_token_offset;
+      } else {
+        continue;
+      }
+
+      block_idx = excess_blocks[ori_bi * excess_num + excess_id];
+      block_offset = excess_token_offset;
+    }
+
+    const uint32_t tgt_idx =
+        block_idx * kv_num_heads * block_size * head_size +
+        hi * block_size * head_size + block_offset * head_size + h_bias;
+
+    const uint32_t ori_idx =
+        token_idx * (num_heads + 2 * kv_num_heads) * head_size +
+        num_heads * head_size + qkv_id * hidden_size + hi * head_size + h_bias;
+
+    Load<T, VecSize>(&qkv[ori_idx], &src_vec);
+    if (qkv_id == 0) {
+      Store<T, VecSize>(src_vec, &key_cache[tgt_idx]);
+    } else {
+      Store<T, VecSize>(src_vec, &value_cache[tgt_idx]);
+    }
+  }
+}
 
 template <typename T,
           uint32_t num_frags_y,
@@ -1463,9 +1556,12 @@ void CascadeAppendWriteCacheKVQKV(
                // kv_num_heads, head_dim] if GQA)
     const paddle::Tensor &block_table,
     const paddle::Tensor &padding_offsets,
+    const paddle::Tensor &cum_offsets,
     const paddle::Tensor &seq_lens_encoder,
     const paddle::Tensor &seq_lens_decoder,
     const int max_seq_len,
+    const int bsz,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     cudaStream_t &stream,
     paddle::Tensor *key_cache_out,
     paddle::Tensor *value_cache_out) {
@@ -1477,29 +1573,58 @@ void CascadeAppendWriteCacheKVQKV(
   auto head_dim_v = meta_data.head_dims_v;
   auto block_size = meta_data.block_size;
 
-  const uint32_t elem_nums =
-      num_tokens * kv_num_heads * (head_dim_qk + head_dim_v);
+  int excess_block_num = 0;
+  int *excess_blocks_ptr = nullptr;
+  if (excess_blocks) {
+    excess_block_num = excess_blocks.get().dims()[1];
+    excess_blocks_ptr =const_cast<int*>(excess_blocks.get().data<int>());
+  }
+  uint32_t elem_nums = (num_tokens + bsz * excess_block_num * block_size) * kv_num_heads * (head_dim_qk + head_dim_v);
+  // 额外每个bid 多分配excess_block_num * block_size 个
+
   constexpr int PackSize = 16 / sizeof(T);
   const int pack_num = elem_nums / PackSize;
   const int blocksize = 128;
   int grid_size = 1;
   GetNumBlocks<128>(pack_num, &grid_size);
-  cache_kernel<T, PackSize><<<grid_size, blocksize, 0, stream>>>(
-      reinterpret_cast<T *>(const_cast<T *>(qkv.data<T>())),
-      reinterpret_cast<T *>(key_cache_out->data<T>()),
-      reinterpret_cast<T *>(value_cache_out->data<T>()),
-      block_table.data<int>(),
-      padding_offsets.data<int>(),
-      seq_lens_encoder.data<int>(),
-      seq_lens_decoder.data<int>(),
-      max_seq_len,
-      max_blocks_per_seq,
-      num_heads,
-      head_dim_qk,
-      head_dim_v,
-      block_size,
-      elem_nums,
-      kv_num_heads);
+  if (excess_blocks_ptr) {
+      cache_use_excess_kernel<T, PackSize><<<grid_size, blocksize, 0, stream>>>(
+        reinterpret_cast<T *>(const_cast<T *>(qkv.data<T>())),
+        reinterpret_cast<T *>(key_cache_out->data<T>()),
+        reinterpret_cast<T *>(value_cache_out->data<T>()),
+        block_table.data<int>(),
+        padding_offsets.data<int>(),
+        cum_offsets.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        excess_blocks_ptr,
+        max_seq_len,
+        max_blocks_per_seq,
+        num_heads,
+        head_dim_qk,
+        block_size,
+        elem_nums,
+        kv_num_heads,
+        num_tokens,
+        excess_block_num);
+  } else {
+    cache_kernel<T, PackSize><<<grid_size, blocksize, 0, stream>>>(
+        reinterpret_cast<T *>(const_cast<T *>(qkv.data<T>())),
+        reinterpret_cast<T *>(key_cache_out->data<T>()),
+        reinterpret_cast<T *>(value_cache_out->data<T>()),
+        block_table.data<int>(),
+        padding_offsets.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        max_seq_len,
+        max_blocks_per_seq,
+        num_heads,
+        head_dim_qk,
+        head_dim_v,
+        block_size,
+        elem_nums,
+        kv_num_heads);
+  }
 }
 
 template <typename T, uint32_t HEAD_DIM, uint32_t BLOCK_SIZE>

csrc/gpu/append_attn/encoder_write_cache_with_rope_kernel.h

@@ -36,6 +36,7 @@ void EncoderWriteCacheWithRopeKernel(
     const paddle::optional<paddle::Tensor>& cache_v_scale,
     const paddle::optional<paddle::Tensor>& cache_k_zp,
     const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& cache_quant_type_str,
     const int num_blocks,
     const int max_seq_len,
@@ -48,6 +49,7 @@ void EncoderWriteCacheWithRopeKernel(
   auto num_heads = meta_data.q_num_heads;
   auto kv_num_heads = meta_data.kv_num_heads;
   auto head_dim = meta_data.head_dims;
+  int bsz = cum_offsets.dims()[0];
   if (rotary_embs) {
     if (num_heads == kv_num_heads) {
       rotary_qk_variable(
@@ -93,9 +95,12 @@ void EncoderWriteCacheWithRopeKernel(
                                     *qkv_out,
                                     block_tables,
                                     padding_offsets,
+                                    cum_offsets,
                                     seq_lens_encoder,
                                     seq_lens_decoder,
                                     max_seq_len,
+                                    bsz,
+                                    excess_blocks,
                                     stream,
                                     key_cache_out,
                                     value_cache_out);

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_bfloat16_bfloat16_kernel.cu

@@ -34,6 +34,7 @@ EncoderWriteCacheWithRopeKernel<paddle::bfloat16, paddle::bfloat16>(
     const paddle::optional<paddle::Tensor>& cache_v_scale,
     const paddle::optional<paddle::Tensor>& cache_k_zp,
     const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& cache_quant_type_str,
     const int num_blocks,
     const int max_seq_len,

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_bfloat16_int_kernel.cu

@@ -33,6 +33,7 @@ template void EncoderWriteCacheWithRopeKernel<paddle::bfloat16, int>(
     const paddle::optional<paddle::Tensor>& cache_v_scale,
     const paddle::optional<paddle::Tensor>& cache_k_zp,
     const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& cache_quant_type_str,
     const int num_blocks,
     const int max_seq_len,

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_float16_float16_kernel.cu

@@ -33,6 +33,7 @@ template void EncoderWriteCacheWithRopeKernel<paddle::float16, paddle::float16>(
     const paddle::optional<paddle::Tensor>& cache_v_scale,
     const paddle::optional<paddle::Tensor>& cache_k_zp,
     const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& cache_quant_type_str,
     const int num_blocks,
     const int max_seq_len,

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_float16_int_kernel.cu

@@ -33,6 +33,7 @@ template void EncoderWriteCacheWithRopeKernel<paddle::float16, int>(
     const paddle::optional<paddle::Tensor>& cache_v_scale,
     const paddle::optional<paddle::Tensor>& cache_k_zp,
     const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& cache_quant_type_str,
     const int num_blocks,
     const int max_seq_len,

csrc/gpu/cpp_extensions.cu

@@ -51,6 +51,7 @@ std::vector<paddle::Tensor> AppendAttention(
     const paddle::optional<paddle::Tensor>& cache_v_zp,
     const paddle::optional<paddle::Tensor>& out_linear_shifts,
     const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
     const std::string& compute_dtype,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
@@ -237,11 +238,9 @@ std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
 
 void SaveOutMmsg(const paddle::Tensor& x,
                  const paddle::Tensor& not_need_stop, // cpu
-                 const paddle::Tensor& queue_id,      // cpu
                  int64_t rank_id);
 
 void GetOutput(const paddle::Tensor& x,
-               const paddle::Tensor& queue_id, // cpu
                int64_t rank_id,
                bool wait_flag);
 

csrc/gpu/get_output.cc

@@ -20,6 +20,7 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
+// #define GET_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -27,14 +28,25 @@ struct msgdata {
 };
 
 void GetOutput(const paddle::Tensor& x,
-               const paddle::Tensor& queue_id,
                int64_t rank_id,
                bool wait_flag) {
   if (rank_id > 0) return;
 
   static struct msgdata msg_rcv;
-  int queue_id_val = queue_id.data<int>()[0];
-  static key_t key = ftok("./", queue_id_val);
+  int msg_queue_id = 1;
+  if (const char* inference_msg_queue_id_env_p =
+          std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+      std::string inference_msg_queue_id_env_str(
+          inference_msg_queue_id_env_p);
+      int inference_msg_queue_id_from_env =
+          std::stoi(inference_msg_queue_id_env_str);
+#ifdef GET_OUTPUT_DEBUG
+      std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                << inference_msg_queue_id_from_env << std::endl;
+#endif
+      msg_queue_id = inference_msg_queue_id_from_env;
+  }
+  static key_t key = ftok("./", msg_queue_id);
 
   static int msgid = msgget(key, IPC_CREAT | 0666);
 
@@ -62,7 +74,7 @@ void GetOutput(const paddle::Tensor& x,
 }
 
 PD_BUILD_OP(get_output)
-    .Inputs({"x", "queue_id"})
+    .Inputs({"x"})
     .Attrs({"rank_id: int64_t",
             "wait_flag: bool"})
     .Outputs({"x_out"})

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h

@@ -69,7 +69,8 @@ void generic_moe_gemm_kernelLauncher(const T* A,
                                      cudaStream_t stream,
                                      int* kernel_occupancy = nullptr) {
   if (gemm_config.split_k_style != SplitKStyle::NO_SPLIT_K) {
-    PADDLE_FATAL("[MoeGemm] Grouped gemm does not support split-k");
+    throw std::runtime_error(
+        "[MoeGemm] Grouped gemm does not support split-k");
   }
 
 #ifdef PADDLE_CUDA_BF16
@@ -169,7 +170,7 @@ void generic_moe_gemm_kernelLauncher(const T* A,
   int occupancy = std::min(2, GemmGrouped::maximum_active_blocks());
 
   if (occupancy == 0) {
-    PADDLE_FATAL(
+    throw std::runtime_error(
         "[MoE Runner] GPU lacks the shared memory resources to run "
         "GroupedGEMM kernel");
   }
@@ -197,7 +198,7 @@ void generic_moe_gemm_kernelLauncher(const T* A,
   if (can_implement != cutlass::Status::kSuccess) {
     std::string err_msg = "MoEFC kernel will fail for params. Error: " +
                           std::string(cutlassGetStatusString(can_implement));
-    PADDLE_FATAL("[MoE Runner] " + err_msg);
+    throw std::runtime_error("[MoE Runner] " + err_msg);
   }
 
   auto init_status = gemm.initialize(args);
@@ -243,7 +244,7 @@ struct dispatch_stages {
     std::string err_msg = "Cutlass fpA_intB gemm. Not instantiates for arch " +
                           std::to_string(arch::kMinComputeCapability) +
                           " with stages set to " + std::to_string(Stages);
-    PADDLE_FATAL("[dispatch_stages::dispatch] " + err_msg);
+    throw std::runtime_error("[dispatch_stages::dispatch] " + err_msg);
   }
 };
 
@@ -394,7 +395,8 @@ void dispatch_gemm_config(const T* A,
     default:
       std::string err_msg = "dispatch_gemm_config does not support stages " +
                             std::to_string(gemm_config.stages);
-      PADDLE_FATAL("[MoE][dispatch_gemm_config] " + err_msg);
+      throw std::runtime_error(
+          "[MoE][dispatch_gemm_config] " + err_msg);
       break;
   }
 }
@@ -452,15 +454,16 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
     dispatch_gemm_config_macro(64, 128, 64, 32, 64, 64);
     dispatch_gemm_config_macro(128, 128, 64, 64, 32, 64);
     case CutlassTileConfig::Undefined:
-      PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
       break;
     case CutlassTileConfig::ChooseWithHeuristic:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass] gemm config should have "
           "already been set by heuristic.");
       break;
     default:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass] Config is invalid for same "
           "type MoE tensorop GEMM.");
       break;
@@ -497,15 +500,16 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
       dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
       dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
       case CutlassTileConfig::Undefined:
-        PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
         break;
       case CutlassTileConfig::ChooseWithHeuristic:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] gemm config should have "
             "already been set by heuristic.");
         break;
       default:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
             "mixed type tensorop GEMM.");
         break;
@@ -513,22 +517,27 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
   } else {
     switch (gemm_config.tile_config) {
       dispatch_gemm_config_macro(16, 128, 64, 16, 32, 64);
+      dispatch_gemm_config_macro(16, 256, 64, 16, 64, 64);
+      dispatch_gemm_config_macro(64, 64, 64, 32, 32, 64);
       dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+      dispatch_gemm_config_macro(128, 64, 64, 64, 32, 64);
       dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
       dispatch_gemm_config_macro(128, 128, 64, 64, 64, 64);
       dispatch_gemm_config_macro(128, 128, 64, 128, 32, 64);
       dispatch_gemm_config_macro(128, 256, 64, 64, 64, 64);
       dispatch_gemm_config_macro(64, 128, 64, 64, 32, 64);
+      dispatch_gemm_config_macro(256, 128, 64, 64, 64, 64);
       case CutlassTileConfig::Undefined:
-        PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
         break;
       case CutlassTileConfig::ChooseWithHeuristic:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] gemm config should have "
             "already been set by heuristic.");
         break;
       default:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
             "mixed type tensorop GEMM.");
         break;
@@ -561,17 +570,17 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
   switch (gemm_config.tile_config) {
     dispatch_gemm_config_macro(128, 128, 8, 64, 64, 8);
     case CutlassTileConfig::Undefined:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config "
           "undefined.");
       break;
     case CutlassTileConfig::ChooseWithHeuristic:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config should "
           "have already been set by heuristic.");
       break;
     default:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass][SIMT] Unsupported config "
           "for float MoE gemm.");
       break;
@@ -725,7 +734,7 @@ void MoeGemmRunner<T, WeightType>::run_gemm<EpilogueTag>(
       gemmConfigManager.addBestConfig(gemmId, profile_total_rows, best_config);
       chosen_config = best_config;
     } else {
-      PADDLE_FATAL("[MoE Configure Search] find no one avaliable config.");
+      PADDLE_FATAL("[MoE Configure Search] find no one available config.");
     }
   }
   dispatch_to_arch<EpilogueTag>(A,

csrc/gpu/moe/preprocess_for_moe.cu

@@ -140,12 +140,14 @@ std::vector<paddle::Tensor> preprocess_for_moe_kernel(const paddle::Tensor& topk
     
     if (num_experts == 8) {
       run_align_kernel(8);
+    } else if (num_experts == 64) {
+      run_align_kernel(64);
     } else if (num_experts == 256) {
       run_align_kernel(256);
     } else if (num_experts == 2) {
       run_align_kernel(2);
     } else {
-      printf("errors");
+      printf("【ERROR】, please add the current number of experts in the case");
     }
 
     const int block_threads = 256;

csrc/gpu/multi_head_latent_attention.cu

@@ -205,6 +205,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
   meta_data.batch_size = cum_offsets.dims()[0];
 
   switch (query.dtype()) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
     case paddle::DataType::BFLOAT16: {
       return MultiHeadLatentAttentionKernel<paddle::DataType::BFLOAT16>(
           meta_data,
@@ -253,6 +254,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           causal,
           speculate_decoder);
     }
+#endif
     case paddle::DataType::FLOAT16: {
       return MultiHeadLatentAttentionKernel<paddle::DataType::FLOAT16>(
           meta_data,

csrc/gpu/save_with_output_msg.cc

@@ -20,6 +20,7 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
+// #define SAVE_WITH_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -28,7 +29,6 @@ struct msgdata {
 
 void SaveOutMmsg(const paddle::Tensor& x,
                  const paddle::Tensor& not_need_stop, // cpu
-                 const paddle::Tensor& queue_id,      // cpu
                  int64_t rank_id) {
     if (rank_id > 0) return;
     auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
@@ -36,8 +36,25 @@ void SaveOutMmsg(const paddle::Tensor& x,
     auto not_need_stop_data = not_need_stop.data<bool>()[0];
 
     static struct msgdata msg_sed;
-    int queue_id_val = queue_id.data<int>()[0];
-    static key_t key = ftok("./", queue_id_val);
+    int msg_queue_id = 1;
+    if (const char* inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+        std::string inference_msg_queue_id_env_str(
+            inference_msg_queue_id_env_p);
+        int inference_msg_queue_id_from_env =
+            std::stoi(inference_msg_queue_id_env_str);
+        msg_queue_id = inference_msg_queue_id_from_env;
+#ifdef SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                  << inference_msg_queue_id_from_env << std::endl;
+#endif
+    } else {
+#ifdef SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
+                  << std::endl;
+#endif
+    }
+    static key_t key = ftok("./", msg_queue_id);
     static int msgid = msgget(key, IPC_CREAT | 0666);
 
     msg_sed.mtype = 1;
@@ -54,7 +71,7 @@ void SaveOutMmsg(const paddle::Tensor& x,
 }
 
 PD_BUILD_OP(save_output)
-    .Inputs({"x", "not_need_stop", "queue_id"})
+    .Inputs({"x", "not_need_stop"})
     .Attrs({"rank_id: int64_t"})
     .Outputs({"x_out"})
     .SetInplaceMap({{"x", "x_out"}})

csrc/gpu/set_preids_token_penalty_multi_scores.cu

@@ -41,11 +41,11 @@ __global__ void set_preids_token_penalty_multi_scores_kernel(const bool *stop_fl
     T *logits_now = logits + bi * length;
     int tid = threadIdx.x;
 
-    if (tid < bs && !stop_flags[tid]) {
-        int64_t *pre_ids_now = pre_ids + tid * length_id;
-        const int64_t *input_ids_now = input_ids + tid * length_input_ids;
-        const int seq_len_dec = seq_lens_decoder[tid];
-        const int seq_len_enc = seq_lens_encoder[tid];
+    if (bi < bs && !stop_flags[bi]) {
+        int64_t *pre_ids_now = pre_ids + bi * length_id;
+        const int64_t *input_ids_now = input_ids + bi * length_input_ids;
+        const int seq_len_dec = seq_lens_decoder[bi];
+        const int seq_len_enc = seq_lens_encoder[bi];
         if (seq_len_dec == 0 && seq_len_enc == 0) return; // stopped
 
         const int step_idx_now = step_idx[bi];

csrc/gpu/unittest/test_get_padding_offset_v2.py

@@ -64,6 +64,5 @@ class GetPaddingOffsetV2Test(unittest.TestCase):
         assert sum(ref_cu_seqlens_q - cu_seqlens_q) == 0, "Check cu_seqlens_q failed."
         assert sum(ref_cu_seqlens_k - cu_seqlens_k) == 0, "Check cu_seqlens_k failed."
 
-
 if __name__ == "__main__":
     unittest.main()

csrc/gpu/update_inputs_v2.cu

@@ -42,8 +42,10 @@ __global__ void update_inputs_kernel_v2(
     const int bsz,
     const int max_bsz,
     const int input_ids_stride,
-    const int end_length) {
+    const int end_length,
+    const int Flag_truncated_return_eos) {
   int thread_idx = threadIdx.x;
+  bool output_len_truncated = false;
   // update step_idx and stop_flags
   if (thread_idx < max_bsz) {
     bool stop_flag = stop_flags[thread_idx];
@@ -52,6 +54,7 @@ __global__ void update_inputs_kernel_v2(
     }
     if (step_idx[thread_idx] >= max_dec_len[thread_idx]) {
       stop_flags[thread_idx] = true;
+      output_len_truncated = true;
     }
   }
   __syncthreads();
@@ -61,8 +64,13 @@ __global__ void update_inputs_kernel_v2(
       if (seq_lens_this_time[thread_idx] == 0) {
         next_tokens[thread_idx] = -1;
       } else {
-        next_tokens[thread_idx] = end_ids[0];
-        kwargs_next_tokens[thread_idx] = end_ids[0];
+        if (!Flag_truncated_return_eos && output_len_truncated) {
+          // output len truncated will not return eos for rl.
+          kwargs_next_tokens[thread_idx] = next_tokens[thread_idx];
+        }else{
+          next_tokens[thread_idx] = end_ids[0];
+          kwargs_next_tokens[thread_idx] = end_ids[0];
+        }
       }
     } else {
       kwargs_next_tokens[thread_idx] = next_tokens[thread_idx];
@@ -127,6 +135,15 @@ void UpdateInputesV2(const paddle::Tensor& stop_flags,
   const int end_length = end_ids.shape()[0];
 
   auto not_need_stop_gpu = not_need_stop.copy_to(stop_flags.place(), false);
+  int Flag_truncated_return_eos = 1;
+  if (const char* inference_truncated_return_eos_env_p =
+          std::getenv("INFERENCE_TRUNCATED_RETURN_EOS")) {
+      std::string inference_truncated_return_eos_env_str(
+          inference_truncated_return_eos_env_p);
+      int inference_truncated_return_eos_from_env =
+          std::stoi(inference_truncated_return_eos_env_str);
+      Flag_truncated_return_eos = inference_truncated_return_eos_from_env;
+  }
 
   update_inputs_kernel_v2<1024><<<1, 1024, 0, input_ids.stream()>>>(
     const_cast<bool*>(not_need_stop_gpu.data<bool>()),
@@ -145,7 +162,8 @@ void UpdateInputesV2(const paddle::Tensor& stop_flags,
     now_bsz,
     max_bsz,
     input_ids_stride,
-    end_length
+    end_length,
+    Flag_truncated_return_eos
   );
 
   auto not_need_stop_cpu = not_need_stop_gpu.copy_to(not_need_stop.place(), false);

csrc/setup_cuda.py

@@ -15,12 +15,24 @@
 import os
 import shutil
 import subprocess
+from packaging.version import parse, Version
 
 import paddle
 from paddle.utils.cpp_extension import CUDAExtension, setup
 
 sm_version = int(os.getenv("CUDA_SM_VERSION", "0"))
 
+def get_nvcc_cuda_version(cuda_dir: str) -> Version:
+    """Get the CUDA version from nvcc.
+
+    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
+    """
+    nvcc_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"],
+                                          universal_newlines=True)
+    output = nvcc_output.split()
+    release_idx = output.index("release") + 1
+    nvcc_cuda_version = parse(output[release_idx].split(",")[0])
+    return nvcc_cuda_version
 
 def update_git_submodule():
     try:
@@ -118,14 +130,11 @@ sources = [
     "./gpu/speculate_decoding_kernels/speculate_save_output.cc",
     "./gpu/speculate_decoding_kernels/speculate_get_output.cc",
     "./gpu/save_output_dygraph.cu",
-    "./gpu/cpp_extensions.cu",
     "./gpu/all_reduce.cu",
     "./gpu/quantization/per_token_group_quant.cu",
     "./gpu/quantization/per_tensor_quant_fp8.cu",
 ]
 sources += find_end_files("./gpu/speculate_decoding_kernels", ".cu")
-sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
-sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
 
 nvcc_compile_args = gencode_flags
 update_git_submodule()
@@ -153,6 +162,7 @@ include_dirs = [
 ]
 cc = get_sm_version()
 cuda_version = float(paddle.version.cuda())
+nvcc_version = get_nvcc_cuda_version(os.environ.get("CUDA_HOME", "/usr/local/cuda"))
 
 if cc >= 80:
     sources += ["gpu/int8_gemm_with_cutlass/gemm_dequant.cu"]
@@ -161,6 +171,9 @@ if cc >= 80:
 
     sources += find_end_files("./gpu/append_attn", ".cu")
     sources += find_end_files("./gpu/append_attn/template_instantiation", ".cu")
+    sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
+    sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
+    sources += "./gpu/cpp_extensions.cu",
 
 
 fp8_auto_gen_directory = "gpu/cutlass_kernels/fp8_gemm_fused/autogen"
@@ -178,7 +191,8 @@ if cc == 89 and cuda_version >= 12.4:
         "gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
     ]
 
-if cc >= 80 and cuda_version >= 12.4:
+if cc >= 80 and nvcc_version >= Version("12.4"):
+    os.environ.pop('PADDLE_CUDA_ARCH_LIST', None)
     nvcc_compile_args += [
         "-std=c++17",
         "--use_fast_math",

csrc/tools/build_wheel.sh

@@ -61,7 +61,7 @@ function generate_sm_version(){
         sm_versions=($SM_VERSION )
     elif [ "$ARCHITECTURE" = "all" ]; then
         if awk -v version="$cuda_version" 'BEGIN { exit !(version >= 12.0) }'; then
-          sm_versions=(70 75 80 80 86 89 90 )
+          sm_versions=(70 75 80 86 89 90 )
         else
           sm_versions=(70 75 80 86 89 ) 
         fi 
@@ -72,10 +72,12 @@ function generate_sm_version(){
 }
 
 function create_directories(){
-  mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops
-  touch $OPS_SRC_DIR/tmp/setup.py
-  touch $OPS_SRC_DIR/tmp/paddlenlp_ops/__init__.py
-  echo '# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+  for sm_version in "${sm_versions[@]}"; do
+    echo "create sm$sm_version"
+    mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops
+    touch $OPS_SRC_DIR/tmp/setup.py
+    touch $OPS_SRC_DIR/tmp/paddlenlp_ops/__init__.py
+    echo '# 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.
@@ -93,6 +95,7 @@ function create_directories(){
 
 import os
 from datetime import datetime
+import paddle
 
 from setuptools import find_packages, setup
 
@@ -109,14 +112,19 @@ def read(file: str):
         content = f.read().strip()
     return content
 
-
 def read_version():
     """
     read version and return content
     """
     __version__ = "3.0.0b4.post"
+
     formatted_date = datetime.now().date().strftime("%Y%m%d")
-    __version__ = __version__.replace(".post", ".post{}".format(formatted_date))
+    cuda_version = float(paddle.version.cuda())
+    sm_version=80
+    paddle_commit = paddle.__git_commit__[:7]
+    build_tag = "{}+cuda{}sm{}paddle{}".format(formatted_date, cuda_version, sm_version, paddle_commit)
+
+    __version__ = __version__.replace(".post", ".post{}".format(build_tag))
     
     return __version__
 
@@ -184,12 +192,9 @@ try:
 except ImportError:
     logger.WARNING(f"No {module_name} ")
 ' > $OPS_SRC_DIR/tmp/paddlenlp_ops/__init__.py
-
-  for sm_version in "${sm_versions[@]}"; do
-        echo "create sm$sm_version"
-        mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}
-        touch $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}/__init__.py
-        echo '# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+    mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}
+    touch $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}/__init__.py
+    echo '# 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.
@@ -210,6 +215,7 @@ try:
 except ImportError:
     logger.WARNING("No paddlenlp_ops_'${sm_version}' ops")
 ' > $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}/__init__.py
+    build_ops
   done
 }
 
@@ -228,11 +234,11 @@ function init() {
 }
 
 function build_ops() {
-    for sm_version in "${sm_versions[@]}"; do
-        echo "Building and installing for sm_version: $sm_version"
-        build_and_install_ops $sm_version
-    done
-    return 
+    echo "Building and installing for sm_version: $sm_version"
+    build_and_install_ops $sm_version
+    build_and_install_whl
+    unittest
+    cleanup
 }
 
 function copy_ops(){
@@ -269,6 +275,7 @@ function build_and_install_whl() {
   echo -e "${BLUE}[build]${NONE} building paddlenlp_ops wheel..."
   rm -rf ./dist
   cd ${TMP_DIR}
+  sed -i "s/sm_version=80/sm_version=${sm_version}/g" setup.py
   ${python} setup.py bdist_wheel --dist-dir ./$DIST_DIR
   if [ $? -ne 0 ]; then
     echo -e "${RED}[FAIL]${NONE} build paddlenlp_ops wheel failed !"
@@ -286,7 +293,8 @@ function build_and_install_whl() {
   fi
   echo -e "${BLUE}[install]${NONE} ${GREEN}paddlenlp_ops install success\n"
   cd ..
-  mv $DIST_DIR ../
+  mkdir -p ../$DIST_DIR
+  mv $DIST_DIR/* ../$DIST_DIR/
   cd ..
 }
 
@@ -321,10 +329,6 @@ trap 'abort' 0
 set -e
 
 init
-build_ops
-build_and_install_whl
-unittest
-cleanup
 
 # get Paddle version
 PADDLE_VERSION=`${python} -c "import paddle; print(paddle.version.full_version)"`

csrc/xpu/src/absorb_mla_block_mha_decoder_xpu.cc

@@ -0,0 +1,219 @@
+// Copyright (c) 2024 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+
+#include "paddle/extension.h"
+#include "paddle/phi/core/enforce.h"
+#include "xpu/plugin.h"
+#include <core/ctx_manager.h>
+#include <core/xft_check.h>
+#include <core/xft_event.h>
+#include <core/xft_params.h>
+#include <xft/xdnn_plugin.h>
+#include <xft/operation/page_attn.h>
+#include <xft/operation/fmha.h>
+#include <flash_api.h> // link xfa
+#include "ops.h"
+
+namespace xftkernel = baidu::xpu::xftkernel;
+
+template <typename T>
+struct kl3_pa_TL_trait {
+    using TL = T;
+};
+template <>
+struct kl3_pa_TL_trait<bfloat16> {
+    using TL = float;
+};
+std::vector<paddle::Tensor> MlaDeAttn(
+    const paddle::Tensor& q,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& decoder_context_len,
+    const paddle::Tensor& decoder_batch_map,
+    const paddle::Tensor& decoder_context_len_cpu,
+    const paddle::Tensor& decoder_batch_map_cpu,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const float softmax_scale,
+    const int block_size,
+    const int num_head,
+    const int kv_lora_rank, 
+    const int rope_head_dim,
+    const int dim_qk,
+    const int dim_v) {  
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] MlaDeAttn" );
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
+
+  using QType = typename XPUTypeTrait<bfloat16>::Type;
+  using CacheType = typename XPUTypeTrait<bfloat16>::Type;
+  typedef paddle::bfloat16 qdata_t, cache_t;
+  const auto& input_dims = q.dims();
+  const int bsz = cum_offsets.dims()[0];
+  const int token_num = input_dims[0];
+  const int block_batch = block_tables.dims()[0]; // TODO参数含义 block_batch_  PageParam page_param_
+  const int max_block_per_seq = block_tables.dims()[1];
+  const int max_seq_len = block_size * max_block_per_seq;
+  int dec_batch = dec_batch_tensor.data<int32_t>()[0];
+  // 初始化输入：q k v
+  auto q_xft = baidu::xpu::xft::xftTensor<QType, 3>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(q.data<qdata_t>())),
+      std::array<int64_t, 3>{q.shape()[0],
+                             q.shape()[1],
+                             q.shape()[2]});  
+  // 初始化输入：k cache
+  auto kv_cache_xft = baidu::xpu::xft::xftTensor<CacheType, 4>(
+  reinterpret_cast<CacheType*>(const_cast<paddle::bfloat16*>(kv_cache.data<cache_t>())),
+  std::array<int64_t, 4>{kv_cache.shape()[0],
+                          kv_cache.shape()[1],
+                          kv_cache.shape()[2],
+                          kv_cache.shape()[3]});                            
+  // 初始化输入：block table
+  auto block_tables_xft = baidu::xpu::xft::xftTensor<int, 2>(
+  reinterpret_cast<int*>(const_cast<int*>(block_tables.data<int>())),
+  std::array<int64_t, 2>{block_tables.shape()[0],
+                          block_tables.shape()[1]}); 
+  // 初始化输出tensor
+  auto fmha_out = paddle::full({q.shape()[0], num_head * kv_lora_rank}, -2, q.type(), q.place()); 
+  auto fmha_out_xft = baidu::xpu::xft::xftTensor<QType, 2>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(fmha_out.data<qdata_t>())),
+      std::array<int64_t, 2>{fmha_out.shape()[0],
+                             fmha_out.shape()[1]});
+
+  // decoder
+  if(dec_batch > 0){
+    // context_len
+    baidu::xpu::api::VectorParam<int32_t> context_len_vp{const_cast<int32_t*>(decoder_context_len_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_context_len.data<int32_t>())};
+    // real batch     
+    baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(decoder_batch_map_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_map.data<int32_t>())};
+  
+    // multi_latent_attention
+    using TQ = bfloat16; 
+    using TKVCACHE = bfloat16; 
+    using TO = TQ; 
+    using TGEMM = float; 
+    using TEW = float;
+    using TID = int;
+    constexpr int quant_mode = 0;
+    // xpu_ctx->x_context().set_debug_level(0xa1);
+    int ret = baidu::xpu::xfa::multi_latent_attention<
+            TQ, 
+            TKVCACHE, 
+            TO, 
+            TGEMM,
+            TEW, 
+            TID, 
+            quant_mode>(
+            xpu_ctx->x_context(),
+            fmha_out_xft.data(),
+            q_xft.data(),
+            kv_cache_xft.data(),
+            block_tables_xft.data(),
+            context_len_vp,
+            valid_batch_vp,
+            block_batch,
+            max_seq_len,
+            num_head,
+            kv_lora_rank,
+            rope_head_dim,
+            nullptr, // attn_mask
+            softmax_scale, // 0.13523377478122711f, // scale
+            block_size,
+            max_block_per_seq,
+            -1,
+            nullptr,
+            nullptr,
+            nullptr);
+  } 
+    baidu::xpu::api::plugin::print_times("[TIME END] MlaDeAttn");
+    return {fmha_out};   
+}
+
+std::vector<std::vector<int64_t>> MlaDeAttnInferShape(
+    const std::vector<int64_t>& q_shape,
+    const std::vector<int64_t>& kv_cache_shape,
+    const std::vector<int64_t>& decoder_context_len_shape,
+    const std::vector<int64_t>& decoder_batch_map_shape,
+    const std::vector<int64_t>& decoder_context_len_cpu_shape,
+    const std::vector<int64_t>& decoder_batch_map_cpu_shape,
+    const std::vector<int64_t>& dec_batch_tensor_shape,
+    const std::vector<int64_t>& padding_offsets_shape,
+    const std::vector<int64_t>& cum_offsets_shape,
+    const std::vector<int64_t>& block_tables_shape,
+    const float softmax_scale,
+    const int block_size,
+    const int num_head,
+    const int kv_lora_rank, 
+    const int rope_head_dim,
+    const int dim_qk,
+    const int dim_v) {  
+  return {{q_shape[0], num_head * kv_lora_rank}};
+}
+
+std::vector<paddle::DataType> MlaDeAttnInferDtype(
+    const paddle::DataType& q_dtype,
+    const paddle::DataType& kv_cache_dtype,
+    const paddle::DataType& decoder_context_len_dtype,
+    const paddle::DataType& decoder_batch_map_dtype, 
+    const paddle::DataType& decoder_context_len_cpu_dtype,
+    const paddle::DataType& decoder_batch_map_cpu_dtype,
+    const paddle::DataType& dec_batch_tensor_dtype,
+    const paddle::DataType& padding_offsets_dtype,
+    const paddle::DataType& cum_offsets_dtype,
+    const paddle::DataType& block_tables_dtype,
+    const float softmax_scale,
+    const int block_size,
+    const int num_head,
+    const int kv_lora_rank, 
+    const int rope_head_dim,
+    const int dim_qk,
+    const int dim_v) {  
+    if (q_dtype == paddle::DataType::FLOAT16) {
+        return {paddle::DataType::FLOAT16};
+    } else if(q_dtype == paddle::DataType::BFLOAT16){
+        return {paddle::DataType::BFLOAT16};
+    } 
+    else {
+    PD_THROW("Only supported attr of compute_dtype in ['fp16','bfp16'].");
+    }
+}
+
+PD_BUILD_OP(absorb_mla_block_mha_decoder_xpu)
+    .Inputs({"q",
+             "kv_cache",
+             "decoder_context_len",
+             "decoder_batch_map",
+             "decoder_context_len_cpu",
+             "decoder_batch_map_cpu",
+             "dec_batch_tensor",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables"})
+    .Outputs({"fmha_out"})
+    .Attrs({"softmax_scale: float",
+            "block_size: int",
+            "num_head: int",
+            "kv_lora_rank: int",
+            "rope_head_dim: int",
+            "dim_qk: int",
+            "dim_v: int"})
+    .SetKernelFn(PD_KERNEL(MlaDeAttn))
+    .SetInferShapeFn(PD_INFER_SHAPE(MlaDeAttnInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MlaDeAttnInferDtype));
+

csrc/xpu/src/absorb_mla_block_mha_encoder_xpu.cc

@@ -0,0 +1,251 @@
+// Copyright (c) 2024 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+
+#include "paddle/extension.h"
+#include "paddle/phi/core/enforce.h"
+#include "xpu/plugin.h"
+#include <core/ctx_manager.h>
+#include <core/xft_check.h>
+#include <core/xft_event.h>
+#include <core/xft_params.h>
+#include <xft/xdnn_plugin.h>
+#include <xft/operation/page_attn.h>
+#include <xft/operation/fmha.h>
+#include <flash_api.h> // link xfa
+#include "ops.h"
+
+namespace xftkernel = baidu::xpu::xftkernel;
+
+template <typename T>
+struct kl3_pa_TL_trait {
+    using TL = T;
+};
+template <>
+struct kl3_pa_TL_trait<bfloat16> {
+    using TL = float;
+};
+std::vector<paddle::Tensor> MlaEnAttn(
+    const paddle::Tensor& q,
+    const paddle::Tensor& k,
+    const paddle::Tensor& v,
+    const paddle::Tensor& encoder_seq_lod,
+    const paddle::Tensor& encoder_batch_map,
+    const paddle::Tensor& encoder_seq_lod_cpu,
+    const paddle::Tensor& encoder_batch_map_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const float softmax_scale,
+    const int block_size,
+    const int num_head,
+    const int dim_qk,
+    const int dim_v) {
+        
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] MlaEnAttn");
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+
+  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
+  using QType = typename XPUTypeTrait<bfloat16>::Type;
+  using CacheType = typename XPUTypeTrait<bfloat16>::Type;
+  typedef paddle::bfloat16 qdata_t, cache_t;
+  const auto& input_dims = q.dims();
+  const int bsz = cum_offsets.dims()[0];
+  const int token_num = input_dims[0];
+  const int block_batch = block_tables.dims()[0]; // TODO参数含义 block_batch_  PageParam page_param_
+  const int max_block_per_seq = block_tables.dims()[1];
+  const int max_seq_len = block_size * max_block_per_seq;
+  int enc_batch = enc_batch_tensor.data<int32_t>()[0]; 
+  // 初始化输入：q k v
+  auto q_xft = baidu::xpu::xft::xftTensor<QType, 3>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(q.data<qdata_t>())),
+      std::array<int64_t, 3>{q.shape()[0],
+                             q.shape()[1],
+                             q.shape()[2]});
+  auto k_xft = baidu::xpu::xft::xftTensor<QType, 3>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(k.data<qdata_t>())),
+      std::array<int64_t, 3>{k.shape()[0],
+                             k.shape()[1],
+                             k.shape()[2]});
+  auto v_xft = baidu::xpu::xft::xftTensor<QType, 3>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(v.data<qdata_t>())),
+      std::array<int64_t, 3>{v.shape()[0],
+                             v.shape()[1],
+                             v.shape()[2]});       
+                                                   
+  // 初始化输出tensor
+  auto fmha_out = paddle::full({q.shape()[0], num_head * dim_v}, -2, q.type(), q.place()); 
+  auto fmha_out_xft = baidu::xpu::xft::xftTensor<QType, 2>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(fmha_out.data<qdata_t>())),
+      std::array<int64_t, 2>{fmha_out.shape()[0],
+                             fmha_out.shape()[1]});
+
+  // encoder
+  if(max_enc_len_this_time.data<int>()[0] > 0){
+    // q_lod
+    baidu::xpu::api::VectorParam<int32_t> context_len_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
+    // real batch（encoder阶段 kv cache写需要）
+    baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(encoder_batch_map_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_map.data<int32_t>())};
+    // prefix (not support)
+    baidu::xpu::api::VectorParam<int32_t> prefix_lens_vp = baidu::xpu::api::VectorParam<int32_t>();
+    // kv_lod (非prefix cache情况与q_lod一致)
+    baidu::xpu::api::VectorParam<int32_t> encoder_kv_lods_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
+    // page_param_
+    baidu::xpu::xft::PageParam page_param_(block_batch, // block_batch_
+                                            max_enc_len_this_time.data<int>()[0], // max_enc_len_this_time.data<int>()[0],max_seq_len // max_context_len_
+                                            block_size, // block_size_
+                                            max_block_per_seq, // max_num_blocks_per_seq_
+                                            false); // v_trans  
+    // attn_param
+    baidu::xpu::xft::PageAttnParam attn_param_(dim_qk, // head_dim(q,k)
+                                              num_head, // q_head_num
+                                              num_head, // kv_head_num
+                                              false, // vp_lod_flag
+                                              context_len_vp, // context_len_vp
+                                              valid_batch_vp, // valid_batch_vp
+                                              prefix_lens_vp, // prefix_lens_vp
+                                              encoder_kv_lods_vp); // encoder_kv_lods_vp
+    // fmha op
+    using FMHA_Type = typename baidu::xpu::xft::FMHA_QBF16_KVBF16;
+    using TQ = typename FMHA_Type::Qtype;
+    using TK = typename FMHA_Type::Ktype;
+    using TV = typename FMHA_Type::Vtype;
+    using TBIAS = typename FMHA_Type::Btype;
+    using TO = typename FMHA_Type::Otype;
+    using TMASK = typename FMHA_Type::Mtype;
+    int ret = baidu::xpu::xfa::flash_attention_context_vllm<
+            TQ,
+            TK,
+            TV,
+            TBIAS,
+            TO,
+            float, // T_GEMM0
+            float, // T_GEMM1
+            float, // T_EW
+            int,
+            0>(
+            xpu_ctx->x_context(),
+            q_xft.data(), // q
+            std::is_same<TK, int8_t>::value ? nullptr : reinterpret_cast<TK*>(k_xft.data()), // k
+            std::is_same<TK, int8_t>::value ? nullptr : reinterpret_cast<TV*>(v_xft.data()), // v
+            nullptr, // p_bias == nullptr ? nullptr : p_bias->data(), // mask_bias
+            fmha_out_xft.data(), //o
+            enc_batch, // batch_size
+            page_param_.max_context_len_,      // max_seq_q
+            page_param_.max_context_len_,  // max_seq_kv, （prefix cache是会不一样）
+            attn_param_.q_head_num_,
+            attn_param_.head_dim_,
+            attn_param_.kv_head_num_,
+            nullptr,
+            attn_param_.context_len_vp_,
+            attn_param_.encoder_kv_lods_vp_,
+            {}, // mask_shape_
+            {}, // alibi_slopes_shape
+            true, // is_causal_mask_
+            false, // is_qkv_fusion_
+            0x0010,  //     param.qkv_layout_ = AttnQKVLayout_t::ATTN_BLHD
+            0x10, // vsl_flag_ = VslConverter::FMHA_LVSL
+            nullptr,  // q_maxptr
+            nullptr,  // k_maxptr
+            nullptr,  // v_maxptr
+            nullptr,
+            dim_v, // vo_head_dim
+            nullptr, // p_cache_k->data(),
+            nullptr, // p_cache_v->data(),
+            nullptr, // std::is_same<TK, int8_t>::value ? p_kcache_perhead_scale->data() : p_cache_k->max_data(),
+            nullptr, // std::is_same<TK, int8_t>::value ? p_vcache_perhead_scale->data() : p_cache_v->max_data(),
+            block_size, // block_size
+            max_block_per_seq, // max_blocks_per_seq (prefix cache)
+            page_param_.max_context_len_, // prefill_len
+            nullptr,
+            softmax_scale * sqrt(dim_qk)); // block_tables (prefix cache)
+
+    // std::cout << "fmha kernel done " <<std::endl;
+  }
+  baidu::xpu::api::plugin::print_times("[TIME END] MlaEnAttn:");
+    return {fmha_out};   
+}
+
+std::vector<std::vector<int64_t>> MlaEnAttnInferShape(
+    const std::vector<int64_t>& q_shape,
+    const std::vector<int64_t>& k_shape,
+    const std::vector<int64_t>& v_shape,
+    const std::vector<int64_t>& encoder_seq_lod_shape,
+    const std::vector<int64_t>& encoder_batch_map_shape,
+    const std::vector<int64_t>& encoder_seq_lod_cpu_shape,
+    const std::vector<int64_t>& encoder_batch_map_cpu_shape,
+    const std::vector<int64_t>& enc_batch_tensor_shape,
+    const std::vector<int64_t>& padding_offsets_shape,
+    const std::vector<int64_t>& cum_offsets_shape,
+    const std::vector<int64_t>& block_tables_shape,
+    const std::vector<int64_t>& max_enc_len_this_time_shape,
+    const std::vector<int64_t>& max_dec_len_this_time_shape) {
+
+  return {v_shape};
+}
+
+std::vector<paddle::DataType> MlaEnAttnInferDtype(
+    const paddle::DataType& q_dtype,
+    const paddle::DataType& k_dtype,
+    const paddle::DataType& v_dtype,
+    const paddle::DataType& encoder_seq_lod_dtype,
+    const paddle::DataType& encoder_batch_map_dtype,
+    const paddle::DataType& encoder_seq_lod_cpu_dtype,
+    const paddle::DataType& encoder_batch_map_cpu_dtype,
+    const paddle::DataType& enc_batch_tensor_dtype,
+    const paddle::DataType& padding_offsets_dtype,
+    const paddle::DataType& cum_offsets_dtype,
+    const paddle::DataType& block_tables_dtype,
+    const paddle::DataType& max_enc_len_this_time_dtype,
+    const paddle::DataType& max_dec_len_this_time_dtype) {
+    if (q_dtype == paddle::DataType::FLOAT16) {
+        return {paddle::DataType::FLOAT16};
+    } else if(q_dtype == paddle::DataType::BFLOAT16){
+        return {paddle::DataType::BFLOAT16};
+    } 
+    else {
+    PD_THROW("Only supported attr of compute_dtype in ['fp16','bfp16'].");
+    }
+}
+
+PD_BUILD_OP(absorb_mla_block_mha_encoder_xpu)
+    .Inputs({"q",
+             "k",
+             "v",
+             "encoder_seq_lod",
+             "encoder_batch_map",
+             "encoder_seq_lod_cpu",
+             "encoder_batch_map_cpu",
+             "enc_batch_tensor",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables",
+             "max_enc_len_this_time",
+             "max_dec_len_this_time"})
+    .Outputs({"fmha_out"})
+    .Attrs({"softmax_scale: float",
+            "block_size: int",
+            "num_head: int",
+            "dim_qk: int",
+            "dim_v: int"})
+    .SetKernelFn(PD_KERNEL(MlaEnAttn))
+    .SetInferShapeFn(PD_INFER_SHAPE(MlaEnAttnInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MlaEnAttnInferDtype));
+

csrc/xpu/src/adjust_batch.cc

@@ -19,63 +19,35 @@
 namespace xftkernel = baidu::xpu::xftkernel;
 std::vector<paddle::Tensor> AdjustBatch(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
                                              const paddle::Tensor& cum_offsets, // [bsz, 1]
-                                             const paddle::Tensor& seq_lens_decoder,
-                                             const paddle::Tensor& seq_lens_encoder,
+                                             const paddle::Tensor& encoder_seq_lod,
+                                             const paddle::Tensor& encoder_batch_idx,
+                                             const paddle::Tensor& decoder_batch_idx,
+                                             const paddle::Tensor& encoder_seq_lod_cpu,
+                                             const paddle::Tensor& encoder_batch_idx_cpu,
+                                             const paddle::Tensor& decoder_batch_idx_cpu,
+                                             const paddle::Tensor& enc_batch_tensor,
+                                             const paddle::Tensor& dec_batch_tensor,
                                              const paddle::optional<paddle::Tensor>& output_padding_offset,
                                              int max_input_length) {
+
   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
-  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
   using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
   typedef paddle::bfloat16 data_t;
   const int token_num = tmp_out.dims()[0]; 
   const int dim = tmp_out.dims()[1]; 
-  const int bsz = cum_offsets.shape()[0];
+  const int bsz = cum_offsets.shape()[0];  
+  int enc_batch = enc_batch_tensor.data<int32_t>()[0];
+  int dec_batch = dec_batch_tensor.data<int32_t>()[0];
+
+  baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
+  baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp{const_cast<int32_t*>(encoder_batch_idx_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_idx.data<int32_t>())};
+  baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp{const_cast<int32_t*>(decoder_batch_idx_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_idx.data<int32_t>())};
 
-  std::vector<int> seq_lens_encoder_cpu(bsz, 0);
-  std::vector<int> seq_lens_decoder_cpu(bsz, 0);
-  std::vector<int> encoder_batch_idx; // 去除空隙的batch map
-  std::vector<int> decoder_batch_idx; // 去除空隙的batch map
-  std::vector<int> encoder_seq_lod;
-  int r = xpu_memcpy(seq_lens_encoder_cpu.data(),
-                 seq_lens_encoder.data<int>(),
-                 sizeof(int32_t) * bsz,
-                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
-  r = xpu_memcpy(seq_lens_decoder_cpu.data(),
-                 seq_lens_decoder.data<int>(),
-                 sizeof(int32_t) * bsz,
-                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
-  int enc_batch = 0, dec_batch = 0;
-  int batch_offset = 0;
-  encoder_seq_lod.push_back(0);
-  for(int i = 0; i < bsz; ++i){
-    if(seq_lens_encoder_cpu[i] > 0){
-      enc_batch++;
-      encoder_batch_idx.push_back(i - batch_offset);
-      encoder_seq_lod.push_back(seq_lens_encoder_cpu[i]);
-      encoder_seq_lod[enc_batch] += encoder_seq_lod[enc_batch - 1];
-    }
-    else if(seq_lens_decoder_cpu[i] > 0){
-      dec_batch++;
-      decoder_batch_idx.push_back(i - batch_offset);
-    }
-    else{
-        batch_offset++;
-    }
-  }         
-  baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp =
-      baidu::xpu::api::VectorParam<int32_t>{encoder_seq_lod.data(), enc_batch + 1, nullptr}
-          .to_xpu(RAII_GUARD);
-  baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp =
-      baidu::xpu::api::VectorParam<int32_t>{encoder_batch_idx.data(), enc_batch, nullptr}
-          .to_xpu(RAII_GUARD);
-  baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp =
-      baidu::xpu::api::VectorParam<int32_t>{decoder_batch_idx.data(), dec_batch, nullptr}
-          .to_xpu(RAII_GUARD);
   auto out = paddle::full({token_num, dim}, -2, tmp_out.type(), tmp_out.place()); 
 
-  r = xftkernel::xft_eb_adjust_batch<XPUType, XPUType>(
+  int r = xftkernel::xft_eb_adjust_batch<XPUType, XPUType>(
           xpu_ctx->x_context(),
           reinterpret_cast<const XPUType*>(tmp_out.data<data_t>()),
           reinterpret_cast<XPUType*>(out.data<data_t>()),
@@ -88,8 +60,14 @@ std::vector<paddle::Tensor> AdjustBatch(const paddle::Tensor& tmp_out, // [token
 
 std::vector<std::vector<int64_t>> AdjustBatchInferShape(const std::vector<int64_t>& tmp_out_shape,
                                                              const std::vector<int64_t>& cum_offsets_shape,
-                                                             const std::vector<int64_t>& seq_lens_decoder_shape,
-                                                             const std::vector<int64_t>& seq_lens_encoder_shape,
+                                                             const std::vector<int64_t>& encoder_seq_lod_shape,
+                                                             const std::vector<int64_t>& encoder_batch_idx_shape,
+                                                             const std::vector<int64_t>& decoder_batch_idx_shape,
+                                                             const std::vector<int64_t>& encoder_seq_lod_cpu_shape,
+                                                             const std::vector<int64_t>& encoder_batch_idx_cpu_shape,
+                                                             const std::vector<int64_t>& decoder_batch_idx_cpu_shape, 
+                                                             const std::vector<int64_t>& enc_batch_tensor_shape,
+                                                             const std::vector<int64_t>& dec_batch_tensor_shape,                                                       
                                                              const paddle::optional<std::vector<int64_t>>& output_padding_offset_shape) {
     if (output_padding_offset_shape) {
       PD_THROW("speculative decoding is not supported in XPU.");
@@ -101,14 +79,20 @@ std::vector<std::vector<int64_t>> AdjustBatchInferShape(const std::vector<int64_
 
 std::vector<paddle::DataType> AdjustBatchInferDtype(const paddle::DataType& tmp_out_dtype,
                                                          const paddle::DataType& cum_offsets_dtype,
-                                                         const paddle::DataType& seq_lens_decoder_dtype,
-                                                         const paddle::DataType& seq_lens_encoder_dtype,
+                                                         const paddle::DataType& encoder_seq_lod_dtype,
+                                                         const paddle::DataType& encoder_batch_idx_dtype,
+                                                         const paddle::DataType& decoder_batch_idx_dtype,
+                                                         const paddle::DataType& encoder_seq_lod_cpu_dtype,
+                                                         const paddle::DataType& encoder_batch_idx_cpu_dtype,
+                                                         const paddle::DataType& decoder_batch_idx_cpu_dtype,  
+                                                         const paddle::DataType& enc_batch_tensor_dtype,
+                                                         const paddle::DataType& dec_batch_tensor_dtype,
                                                          const paddle::optional<paddle::DataType>& output_padding_offset_dtype) {
     return {tmp_out_dtype};
 }
 
 PD_BUILD_OP(adjust_batch)
-    .Inputs({"tmp_out", "cum_offsets", "seq_lens_decoder", "seq_lens_encoder", paddle::Optional("output_padding_offset")})
+    .Inputs({"tmp_out", "cum_offsets", "encoder_seq_lod", "encoder_batch_idx", "decoder_batch_idx", "encoder_seq_lod_cpu", "encoder_batch_idx_cpu", "decoder_batch_idx_cpu", "enc_batch_tensor", "dec_batch_tensor", paddle::Optional("output_padding_offset")})
     .Outputs({"out"})
     .Attrs({"max_input_length: int"})
     .SetKernelFn(PD_KERNEL(AdjustBatch))

csrc/xpu/src/batch_matmul_xpu.cc

@@ -0,0 +1,76 @@
+// Copyright (c) 2024 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include "paddle/extension.h"
+#include "xpu/plugin.h"
+#include <xft/xdnn_plugin.h>
+#include "cublasLt.h"
+#include "ops.h"
+
+namespace xftkernel = baidu::xpu::xftkernel;
+std::vector<paddle::Tensor> Bmm(const paddle::Tensor& input, 
+                                             const paddle::Tensor& weight) {
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
+  using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
+  typedef paddle::bfloat16 data_t;
+  const int batch = input.dims()[0];
+  const int m = input.dims()[1];
+  const int n = weight.dims()[2];
+  const int k = weight.dims()[1];
+
+  auto out = paddle::full({batch, m, n}, -2, input.type(), input.place());           
+  int ret = baidu::xpu::xblas::fc_batched<XPUType, XPUType, XPUType, float, float, 0>(
+          xpu_ctx->x_context(),
+          batch,
+          false,
+          false,
+          m,
+          n,
+          k,
+          1.0f,
+          reinterpret_cast<const XPUType*>(input.data<data_t>()),
+          m * k,
+          reinterpret_cast<const XPUType*>(weight.data<data_t>()),
+          k * n,
+          0.0f,
+          reinterpret_cast<XPUType*>(out.data<data_t>()),
+          m * n,
+          nullptr,
+          nullptr);
+  return {out};
+}
+
+std::vector<std::vector<int64_t>> BmmInferShape(const std::vector<int64_t>& input_shape,
+                                                             const std::vector<int64_t>& weight_shape) {
+    int64_t batch = input_shape[0];
+    int64_t m = input_shape[1];
+    int64_t n = weight_shape[2];
+    return {{batch, m, n}};
+}
+
+std::vector<paddle::DataType> BmmInferDtype(const paddle::DataType& input_dtype,
+                                                         const paddle::DataType& weight_dtype) {
+    return {input_dtype};
+}
+
+PD_BUILD_OP(batch_matmul_xpu)
+    .Inputs({"imput", "weight"})
+    .Outputs({"out"})
+    .SetKernelFn(PD_KERNEL(Bmm))
+    .SetInferShapeFn(PD_INFER_SHAPE(BmmInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(BmmInferDtype));

csrc/xpu/src/block_attn.cc

@@ -0,0 +1,448 @@
+// 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include "paddle/extension.h"
+#include "paddle/phi/core/enforce.h"
+#include "xpu/plugin.h"
+#include "ops.h"
+
+#include <core/ctx_manager.h>
+#include <core/xft_check.h>
+#include <core/xft_event.h>
+#include <core/xft_params.h>
+#include <xft/xdnn_plugin.h>
+#include <xft/operation/page_attn.h>
+#include <xft/operation/fmha.h>
+#include <flash_api.h> // link xfa
+
+namespace xftkernel = baidu::xpu::xftkernel;
+namespace api = baidu::xpu::api;
+
+std::vector<paddle::Tensor> MlaBlockAttn(
+    const paddle::Tensor& attn_input, // [m, 7168]  
+    const paddle::Tensor& q_a_proj_weights, // self.q_a_proj_weights[i] [7168, 1536] 【enc】【dec】
+    const paddle::Tensor& q_b_proj_weights, // self.q_b_proj_weights[i] [1536, 24576] 【enc】【dec】
+    const paddle::Tensor& q_a_layernorm_weights, // self.q_a_layernorm_weights[i]  [1536] 【enc】【dec】
+    const paddle::Tensor& kv_a_proj_with_mqa_weights, // self.kv_a_proj_with_mqa_weights[i] [7168, 576] 【enc】
+    const paddle::Tensor& kv_a_layernorm_weights, // self.kv_a_layernorm_weights[i] [512] 【enc】【dec】
+    const paddle::Tensor& kv_b_proj_weights, // self.kv_b_proj_weights[i] [512, 32768] 【enc】
+    const paddle::Tensor& linear_weights, // self.linear_weights[i] [16384, 7168] 【enc】【dec】
+    const paddle::Tensor& k_b_proj_weights, // self.k_b_proj_weights[i] [128, 128, 512] 【dec】
+    const paddle::Tensor& v_b_proj_weights, // self.v_b_proj_weights[i] [128, 512, 128] 【dec】
+    const paddle::Tensor& cos_sin_cache, // self.config.rotary_emb.cos_sin_cache [163840, 64]
+    const paddle::Tensor& position_ids, // self.position_ids [m]
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_seq_lod, 
+    const paddle::Tensor& encoder_batch_map,
+    const paddle::Tensor& start_token_raw,
+    const paddle::Tensor& decoder_context_len_cache,
+    const paddle::Tensor& decoder_context_len,
+    const paddle::Tensor& decoder_batch_map,
+    const paddle::Tensor& kv_seq_lod_raw,
+    const paddle::Tensor& encoder_seq_lod_cpu, 
+    const paddle::Tensor& encoder_batch_map_cpu,
+    const paddle::Tensor& start_token_raw_cpu,
+    const paddle::Tensor& decoder_context_len_cache_cpu,
+    const paddle::Tensor& decoder_context_len_cpu,
+    const paddle::Tensor& decoder_batch_map_cpu,
+    const paddle::Tensor& kv_seq_lod_raw_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& total_enc_len_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const std::string& cache_quant_type_str,
+    const float softmax_scale,
+    const int block_size,
+    const int num_heads, // self.num_heads
+    const int q_lora_rank, // 传 self.config.mla_config.q_lora_rank = 1536
+    const int qk_head_dim, // self.config.mla_config.qk_head_dim 192
+    const int qk_nope_head_dim, // self.config.mla_config.qk_nope_head_dim 128
+    const int qk_rope_head_dim, // self.config.mla_config.qk_rope_head_dim 64
+    const int kv_lora_rank, // self.config.mla_config.kv_lora_rank 512
+    const int v_head_dim // self.config.mla_config.v_head_dim 128
+    ) { 
+
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
+  using CType = typename XPUTypeTrait<bfloat16>::Type;
+  using CacheType = typename XPUTypeTrait<bfloat16>::Type;
+  typedef paddle::bfloat16 data_t, cache_t;
+
+  const int bsz = encoder_batch_map.dims()[0]; // batch map tensor长度都是max_batch_size, 有效长度根据
+  // const int block_batch = block_tables.dims()[0]; // TODO参数含义 block_batch_  PageParam page_param_
+  const int max_block_per_seq = block_tables.dims()[1];
+  const int max_seq_len = block_size * max_block_per_seq;
+  const int token_num = attn_input.dims()[0];
+  const int hidden_dim = attn_input.dims()[1];
+  int enc_batch = enc_batch_tensor.data<int32_t>()[0]; 
+  int dec_batch = dec_batch_tensor.data<int32_t>()[0];
+  int total_enc_len = total_enc_len_tensor.data<int32_t>()[0];
+   
+  // 输出tensor
+  auto block_attn_out = paddle::full({token_num, hidden_dim}, -1, attn_input.type(), attn_input.place()); 
+
+  // encoder
+  if(enc_batch > 0){
+    // split and copy
+    auto ln_out_encoder = paddle::experimental::slice(attn_input, {0}, {0}, {total_enc_len}, {0}, {});
+    auto position_ids_enc = paddle::experimental::slice(position_ids, {0}, {0}, {total_enc_len}, {0}, {});
+    // encoder no absorb fc
+    if (q_lora_rank <= 0) {
+        assert("Not supported in block attn");
+    }
+    auto query_0 = paddle::matmul(ln_out_encoder, q_a_proj_weights);
+    auto query_1 = std::get<0>(paddle::experimental::rms_norm(query_0, NULL, NULL, q_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0));
+    auto query_2 = paddle::matmul(query_1, q_b_proj_weights);
+
+    // auto query_nope = paddle::full({total_enc_len, num_heads, qk_nope_head_dim}, -1, query_2.type(), query_2.place());
+    auto query_pe = paddle::full({total_enc_len, num_heads, qk_rope_head_dim}, -1, query_2.type(), query_2.place()); 
+    int r = api::split<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(query_2.data<data_t>()),
+      {nullptr, reinterpret_cast<CType*>(query_pe.data<data_t>())},
+      {total_enc_len, num_heads, qk_head_dim},
+      {qk_nope_head_dim, qk_rope_head_dim},
+      2); // axis = 2
+    PD_CHECK(r == 0, "api::split failed.");
+
+    auto compressed_kv = paddle::matmul(ln_out_encoder, kv_a_proj_with_mqa_weights);
+
+    auto compressed_kv_0 = paddle::full({total_enc_len, kv_lora_rank}, -1, compressed_kv.type(), compressed_kv.place()); 
+    auto key_pe = paddle::full({total_enc_len, qk_rope_head_dim}, -1, compressed_kv.type(), compressed_kv.place()); 
+    r = api::split<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(compressed_kv.data<data_t>()),
+      {reinterpret_cast<CType*>(compressed_kv_0.data<data_t>()), reinterpret_cast<CType*>(key_pe.data<data_t>())},
+      {total_enc_len, kv_lora_rank + qk_rope_head_dim},
+      {kv_lora_rank, qk_rope_head_dim},
+      1); // axis = 1
+    PD_CHECK(r == 0, "api::split failed.");
+
+    auto compressed_kv_1 = std::get<0>(paddle::experimental::rms_norm(compressed_kv_0, NULL, NULL, kv_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0));
+    FusedRotaryPositionEncoding(query_pe, key_pe, position_ids_enc, cos_sin_cache, qk_rope_head_dim, false);
+    PrefillMLAWriteCacheKernel(
+      compressed_kv_1,
+      key_pe,
+      kv_cache,
+      encoder_seq_lod,
+      encoder_batch_map,
+      start_token_raw,
+      encoder_seq_lod_cpu,
+      encoder_batch_map_cpu,
+      start_token_raw_cpu,
+      enc_batch_tensor,
+      dec_batch_tensor,
+      padding_offsets,
+      cum_offsets,
+      block_tables,
+      cache_quant_type_str,
+      max_seq_len,
+      num_heads);
+
+    auto key_value = paddle::matmul(compressed_kv_1, kv_b_proj_weights);
+
+    auto key_nope = paddle::full({total_enc_len, num_heads, qk_nope_head_dim}, -1, key_value.type(), key_value.place()); 
+    auto value = paddle::full({total_enc_len, num_heads, v_head_dim}, -1, key_value.type(), key_value.place()); 
+    r = api::split<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(key_value.data<data_t>()),
+      {reinterpret_cast<CType*>(key_nope.data<data_t>()), reinterpret_cast<CType*>(value.data<data_t>())},
+      {total_enc_len, num_heads, qk_nope_head_dim + v_head_dim},
+      {qk_nope_head_dim, v_head_dim},
+      2); // axis = 2
+    PD_CHECK(r == 0, "api::split failed.");
+
+    r = api::strided_slice_view_update<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<CType*>(query_pe.data<data_t>()),
+      reinterpret_cast<CType*>(query_2.data<data_t>()),
+      {total_enc_len, num_heads, qk_rope_head_dim},
+      {total_enc_len, num_heads, qk_head_dim},
+      {0, 0, qk_nope_head_dim}, // start
+      {total_enc_len, num_heads, qk_head_dim}, // end
+      {1, 1, 1}); // stride
+    PD_CHECK(r == 0, "api::strided_slice_view_update failed.");
+
+    auto key = paddle::full({total_enc_len, num_heads, qk_head_dim}, -1, query_2.type(), query_2.place());
+    r = api::strided_slice_view_update<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<CType*>(key_nope.data<data_t>()),
+      reinterpret_cast<CType*>(key.data<data_t>()),
+      {total_enc_len, num_heads, qk_nope_head_dim},
+      {total_enc_len, num_heads, qk_head_dim}, 
+      {0, 0, 0}, // start
+      {total_enc_len, num_heads, qk_nope_head_dim}, // end
+      {1, 1, 1}); // stride
+    PD_CHECK(r == 0, "api::strided_slice_view_update failed.");
+
+    auto key_pe_broadcast = paddle::full({total_enc_len, num_heads, qk_rope_head_dim}, -1, query_2.type(), query_2.place());
+    r = api::broadcast<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<CType*>(key_pe.data<data_t>()),
+      reinterpret_cast<CType*>(key_pe_broadcast.data<data_t>()),
+      {total_enc_len, 1, qk_rope_head_dim},
+      {total_enc_len, num_heads, qk_rope_head_dim});
+    PD_CHECK(r == 0, "api::broadcast failed.");
+
+    r = api::strided_slice_view_update<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<CType*>(key_pe_broadcast.data<data_t>()),
+      reinterpret_cast<CType*>(key.data<data_t>()),
+      {total_enc_len, num_heads, qk_rope_head_dim},
+      {total_enc_len, num_heads, qk_head_dim},
+      {0, 0, qk_nope_head_dim}, // start
+      {total_enc_len, num_heads, qk_head_dim}, // end
+      {1, 1, 1}); // stride
+    PD_CHECK(r == 0, "api::strided_slice_view_update failed.");
+
+    // attn
+    auto fmha_out_prefill = MlaEnAttn(
+      query_2,
+      key,
+      value,
+      encoder_seq_lod,
+      encoder_batch_map,
+      encoder_seq_lod_cpu,
+      encoder_batch_map_cpu,
+      enc_batch_tensor,
+      padding_offsets,
+      cum_offsets,
+      block_tables,
+      max_enc_len_this_time,
+      max_dec_len_this_time,
+      softmax_scale,
+      block_size,
+      num_heads,
+      qk_head_dim,
+      v_head_dim)[0];    
+
+    auto out_linear_out_prefill = paddle::matmul(fmha_out_prefill, linear_weights);
+
+    r = api::copy<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(out_linear_out_prefill.data<data_t>()),
+      reinterpret_cast<CType*>(block_attn_out.data<data_t>()),
+      {total_enc_len * hidden_dim});
+    PD_CHECK(r == 0, "api::copy failed.");
+
+  }
+
+  // decoder
+  if(dec_batch > 0){
+    // split
+    auto ln_out_decoder = paddle::experimental::slice(attn_input, {0}, {total_enc_len}, {token_num}, {0}, {});
+    auto position_ids_dec = paddle::experimental::slice(position_ids, {0}, {total_enc_len}, {token_num}, {0}, {});
+    // decoder absorb fc
+    auto compressed_kv = paddle::matmul(ln_out_decoder, kv_a_proj_with_mqa_weights);
+
+    auto compressed_kv_0 = paddle::full({dec_batch, kv_lora_rank}, -1, compressed_kv.type(), compressed_kv.place()); 
+    auto key_pe = paddle::full({dec_batch, qk_rope_head_dim}, -1, compressed_kv.type(), compressed_kv.place()); 
+    int r = api::split<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(compressed_kv.data<data_t>()),
+      {reinterpret_cast<CType*>(compressed_kv_0.data<data_t>()), reinterpret_cast<CType*>(key_pe.data<data_t>())},
+      {dec_batch, kv_lora_rank + qk_rope_head_dim},
+      {kv_lora_rank, qk_rope_head_dim},
+      1); // axis = 1
+    PD_CHECK(r == 0, "api::split failed.");
+
+    auto compressed_kv_1 = std::get<0>(paddle::experimental::rms_norm(compressed_kv_0, NULL, NULL, kv_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0)); // TODO inplace
+
+    if (q_lora_rank <= 0) {
+        assert("Not supported in block attn");
+    }
+    auto query_0 = paddle::matmul(ln_out_decoder, q_a_proj_weights);
+    auto query_1 = std::get<0>(paddle::experimental::rms_norm(query_0, NULL, NULL, q_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0));
+    auto query_2 = paddle::matmul(query_1, q_b_proj_weights);
+
+    auto query_nope = paddle::full({dec_batch, num_heads, qk_nope_head_dim}, -1, query_2.type(), query_2.place()); 
+    auto query_pe = paddle::full({dec_batch, num_heads, qk_rope_head_dim}, -1, query_2.type(), query_2.place()); 
+    r = api::split<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(query_2.data<data_t>()),
+      {reinterpret_cast<CType*>(query_nope.data<data_t>()), reinterpret_cast<CType*>(query_pe.data<data_t>())},
+      {dec_batch, num_heads, qk_head_dim},
+      {qk_nope_head_dim, qk_rope_head_dim},
+      2); // axis = 2
+    PD_CHECK(r == 0, "api::split failed.");
+
+    auto query_nope_trans = paddle::full({num_heads, dec_batch, qk_nope_head_dim}, -1, query_2.type(), query_2.place()); 
+    r = api::transpose<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(query_nope.data<data_t>()),
+      reinterpret_cast<CType*>(query_nope_trans.data<data_t>()),
+      {dec_batch, num_heads, qk_nope_head_dim},
+      {1, 0, 2});
+    PD_CHECK(r == 0, "api::transpose failed.");
+
+    auto q_nope_bmm = Bmm(query_nope_trans, k_b_proj_weights)[0];
+
+    auto q_nope_out = paddle::full({dec_batch, num_heads, kv_lora_rank}, -1, query_2.type(), query_2.place());
+    r = api::transpose<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(q_nope_bmm.data<data_t>()),
+      reinterpret_cast<CType*>(q_nope_out.data<data_t>()),
+      {num_heads, dec_batch, kv_lora_rank},
+      {1, 0, 2});
+    PD_CHECK(r == 0, "api::transpose failed.");
+
+    FusedRotaryPositionEncoding(query_pe, key_pe, position_ids_dec, cos_sin_cache, qk_rope_head_dim, false);
+    DecodeMLAWriteCacheKernel(
+      compressed_kv_1,
+      key_pe,
+      kv_cache,
+      decoder_context_len_cache,
+      decoder_batch_map,
+      kv_seq_lod_raw,
+      decoder_context_len_cache_cpu,
+      decoder_batch_map_cpu,
+      kv_seq_lod_raw_cpu,
+      enc_batch_tensor,
+      dec_batch_tensor,
+      padding_offsets,
+      cum_offsets,
+      block_tables,
+      cache_quant_type_str,
+      max_seq_len,
+      num_heads,
+      false); // self.config.speculate_config.speculate_method is not None
+
+    auto q_input = paddle::full({dec_batch, num_heads * (kv_lora_rank + qk_rope_head_dim)}, -1, query_2.type(), query_2.place()); // TODO inplace
+    r = api::concat<CType>(
+      xpu_ctx->x_context(),
+      {reinterpret_cast<const CType*>(q_nope_out.data<data_t>()), reinterpret_cast<const CType*>(query_pe.data<data_t>())},
+      reinterpret_cast<CType*>(q_input.data<data_t>()),
+      {{dec_batch, num_heads, kv_lora_rank}, {dec_batch, num_heads, qk_rope_head_dim}},
+      2); // axis = 2
+
+    // attn 
+    auto fmha_out_decode = MlaDeAttn( // [dec_batch, 128*512]
+      q_input,
+      kv_cache,
+      decoder_context_len,
+      decoder_batch_map,
+      decoder_context_len_cpu,
+      decoder_batch_map_cpu,
+      dec_batch_tensor,
+      padding_offsets,
+      cum_offsets,
+      block_tables,
+      softmax_scale,
+      block_size,
+      num_heads,
+      kv_lora_rank,
+      qk_rope_head_dim,
+      qk_head_dim,
+      v_head_dim)[0];    
+  
+    auto fmha_out_decode_trans = paddle::full({num_heads, dec_batch, kv_lora_rank}, -1, query_2.type(), query_2.place()); // shape要跟transpose后的对齐，下面的BMM会读取shape做计算
+    r = api::transpose<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(fmha_out_decode.data<data_t>()),
+      reinterpret_cast<CType*>(fmha_out_decode_trans.data<data_t>()),
+      {dec_batch, num_heads, kv_lora_rank},
+      {1, 0, 2});
+    PD_CHECK(r == 0, "api::transpose failed.");
+
+    auto fmha_out_decode_bmm = Bmm(fmha_out_decode_trans, v_b_proj_weights)[0]; // fmha_out_decode_bmm = [dec_batch, 128*128]
+
+    auto fmha_out_decode_out = paddle::full({dec_batch, num_heads * v_head_dim}, -1, query_2.type(), query_2.place());
+    r = api::transpose<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(fmha_out_decode_bmm.data<data_t>()),
+      reinterpret_cast<CType*>(fmha_out_decode_out.data<data_t>()),
+      {num_heads, dec_batch, v_head_dim},
+      {1, 0, 2});
+    PD_CHECK(r == 0, "api::transpose failed.");
+
+    auto out_linear_out_decode = paddle::matmul(fmha_out_decode_out, linear_weights);
+
+    r = api::copy<CType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const CType*>(out_linear_out_decode.data<data_t>()),
+      reinterpret_cast<CType*>(block_attn_out.data<data_t>()) + total_enc_len * hidden_dim,
+      {dec_batch * hidden_dim});
+    PD_CHECK(r == 0, "api::copy failed.");
+  }
+  return {block_attn_out};
+    
+}
+
+
+std::vector<std::vector<int64_t>> MlaBlockAttnInferShape(const std::vector<int64_t>& attn_input_shape) {  
+  return {attn_input_shape};
+}
+
+std::vector<paddle::DataType> MlaBlockAttnInferDtype(const paddle::DataType& attn_input_dtype) {  
+  return {attn_input_dtype};
+}
+
+PD_BUILD_OP(absorb_mla_block_attention_xpu)
+    .Inputs({"attn_input", 
+             "q_a_proj_weights", 
+             "q_b_proj_weights", 
+             "q_a_layernorm_weights", 
+             "kv_a_proj_with_mqa_weights", 
+             "kv_a_layernorm_weights", 
+             "kv_b_proj_weights", 
+             "linear_weights", 
+             "k_b_proj_weights", 
+             "v_b_proj_weights", 
+             "cos_sin_cache",  
+             "position_ids", 
+             "kv_cache",
+             "block_tables",
+             "encoder_seq_lod", 
+             "encoder_batch_map",
+             "start_token_raw",
+             "decoder_context_len_cache",
+             "decoder_context_len",
+             "decoder_batch_map",
+             "kv_seq_lod_raw",
+             "encoder_seq_lod_cpu", 
+             "encoder_batch_map_cpu",
+             "start_token_raw_cpu",
+             "decoder_context_len_cache_cpu",
+             "decoder_context_len_cpu",
+             "decoder_batch_map_cpu",
+             "kv_seq_lod_raw_cpu",
+             "enc_batch_tensor",
+             "dec_batch_tensor",
+             "total_enc_len_tensor",
+             "padding_offsets",
+             "cum_offsets",
+             "max_enc_len_this_time",
+             "max_dec_len_this_time",})
+    .Outputs({"block_attn_out"})
+    // .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
+    .Attrs({"cache_quant_type_str: std::string",
+            "softmax_scale: float",
+            "block_size: int",
+            "num_heads: int",
+            "q_lora_rank: int",
+            "qk_head_dim: int",
+            "qk_nope_head_dim: int",
+            "qk_rope_head_dim: int",
+            "kv_lora_rank: int",
+            "v_head_dim: int"})
+    .SetKernelFn(PD_KERNEL(MlaBlockAttn))
+    .SetInferShapeFn(PD_INFER_SHAPE(MlaBlockAttnInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MlaBlockAttnInferDtype));

csrc/xpu/src/dynamic_set_split_value.cc

@@ -0,0 +1,102 @@
+// Copyright (c) 2024 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.
+#include <cassert> 
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include "paddle/extension.h"
+#include "xpu/plugin.h"
+#include <xft/xdnn_plugin.h>
+namespace xftkernel = baidu::xpu::xftkernel;
+std::vector<paddle::Tensor> DynamicSetValue(const paddle::Tensor& input,
+                                            const paddle::Tensor& index,
+                                            paddle::Tensor* output,
+                                            bool is_before_index) { // true: [ : index, :] ; false: [index : , :]
+
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
+  typedef paddle::bfloat16 data_t;
+
+  int index_num = index.data<int64_t>()[0];
+  int total_len = output->dims()[0];
+  int dim = output->dims()[1];
+  int copy_num = input.numel();
+  assert(dim == input.dims()[1] && "input.dim != output.dim! ");
+  assert((index_num >= 0 && index_num <= total_len) && "index out of range! ");
+  if (is_before_index) {
+    assert((index_num - 0) * dim == copy_num && "copy_num wrong! ");
+  }
+  else{
+    assert((total_len - index_num) * dim == copy_num && "copy_num wrong! ");
+  }
+  int r = baidu::xpu::api::copy<XPUType>(
+      xpu_ctx->x_context(),
+      reinterpret_cast<const XPUType*>(input.data<data_t>()),
+      reinterpret_cast<XPUType*>(output->data<data_t>() + (is_before_index ? 0 : index_num * dim)),
+      copy_num);
+
+  return {};
+}
+
+PD_BUILD_OP(dynamic_set_value)
+    .Inputs({"input", "index", "output"})
+    .Outputs({"output_"})
+    .SetInplaceMap({{"output_", "output"}})
+    .Attrs({"is_before_index: bool"})
+    .SetKernelFn(PD_KERNEL(DynamicSetValue));
+
+
+
+// std::vector<paddle::Tensor> DynamicSplitValue(const paddle::Tensor& input,
+//                                             const paddle::Tensor& index,
+//                                             bool is_before_index) { // true: [ : index, :] ; false: [index : , :]
+
+//   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+//   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+//   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+//   using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
+//   typedef paddle::bfloat16 data_t;
+
+
+
+  
+//   int index_num = index.data<int64_t>()[0];
+//   int total_len = input.dims()[0];
+//   int dim = input.dims()[1];
+//   int copy_len = 0;
+
+//   if (is_before_index) {
+//     copy_len = index_num
+//   }
+//   else{
+//     copy_len = total_len - index_num
+//   }  
+
+//   auto output = paddle::full({copy_len, dim}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+//   assert(dim == input.dims()[1] && "input.dim != output.dim! ");
+//   assert((index_num >= 0 && index_num <= total_len) && "index out of range! ");
+//   if (is_before_index) {
+//     assert((index_num - 0) * dim == copy_len && "copy_len wrong! ");
+//   }
+//   else{
+//     assert((total_len - index_num) * dim == copy_len && "copy_len wrong! ");
+//   }
+//   int r = baidu::xpu::api::copy<XPUType>(
+//       xpu_ctx->x_context(),
+//       reinterpret_cast<const XPUType*>(input.data<data_t>()),
+//       reinterpret_cast<XPUType*>(output->data<data_t>() + (is_before_index ? 0 : index_num * dim)),
+//       copy_len);
+
+//   return {};
+// }

csrc/xpu/src/fused_rotary_position_encoding.cc

@@ -0,0 +1,72 @@
+// 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include "paddle/extension.h"
+#include "xpu/plugin.h"
+#include "ops.h"
+#include <xft/xdnn_plugin.h>
+
+
+
+void FusedRotaryPositionEncoding(
+    paddle::Tensor& query,  // [num_tokens, num_heads, head_size] or
+                            // [num_tokens, num_heads * head_size]
+    paddle::Tensor& key,
+    // [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads *
+    // head_size]
+    const paddle::Tensor& position_ids,   // [num_tokens]
+    const paddle::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    int head_size,
+    bool is_neox) {
+
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] FusedRotaryPositionEncoding" );
+
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
+
+  typedef paddle::bfloat16 data_t;
+  using XPUType = typename XPUTypeTrait<data_t>::Type;
+
+  int64_t num_tokens = query.dims()[0];
+  int num_heads = query.numel() / num_tokens / head_size;
+  int num_kv_heads = key.numel() / num_tokens / head_size;
+  int rot_dim = cos_sin_cache.dims()[1];
+  int64_t query_stride = num_heads * head_size;
+  int64_t key_stride = num_kv_heads * head_size;
+
+  baidu::xpu::api::plugin::rotary_embedding_neox<XPUType, XPUType>(xpu_ctx->x_context(),
+  position_ids.data<int>(),
+  reinterpret_cast<XPUType*>(query.data<data_t>()),
+  reinterpret_cast<XPUType*>(key.data<data_t>()),
+  reinterpret_cast<const XPUType*>(cos_sin_cache.data<data_t>()),
+  rot_dim,
+  num_heads,
+  num_kv_heads,
+  head_size,
+  num_tokens);
+  baidu::xpu::api::plugin::print_times("[TIME END] FusedRotaryPositionEncoding" );
+
+}
+
+
+PD_BUILD_OP(fused_rotary_position_encoding)
+    .Inputs({"query", "key", "position_ids", "cos_sin_cache"})
+    .Outputs({"query_out", "key_out"})
+    .Attrs({"head_size: int", "is_neox: bool"})
+    .SetInplaceMap({{"query", "query_out"}, {"key", "key_out"}})
+    .SetKernelFn(PD_KERNEL(FusedRotaryPositionEncoding));
+    

csrc/xpu/src/gather_next_token.cc

@@ -19,62 +19,33 @@
 namespace xftkernel = baidu::xpu::xftkernel;
 std::vector<paddle::Tensor> GatherNextToken(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
                                              const paddle::Tensor& cum_offsets, // [bsz, 1]
-                                             const paddle::Tensor& seq_lens_decoder,
-                                             const paddle::Tensor& seq_lens_encoder,
+                                             const paddle::Tensor& encoder_seq_lod,
+                                             const paddle::Tensor& encoder_batch_map,
+                                             const paddle::Tensor& decoder_batch_map,
+                                             const paddle::Tensor& encoder_seq_lod_cpu,
+                                             const paddle::Tensor& encoder_batch_map_cpu,
+                                             const paddle::Tensor& decoder_batch_map_cpu,
+                                             const paddle::Tensor& enc_batch_tensor,
+                                             const paddle::Tensor& dec_batch_tensor,                                                                                          
                                              const paddle::optional<paddle::Tensor>& output_padding_offset,
                                              int max_input_length) {
   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
-  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
   using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
   typedef paddle::bfloat16 data_t;
   const int dim = tmp_out.dims()[1]; 
   const int bsz = cum_offsets.shape()[0];
+  int enc_batch = enc_batch_tensor.data<int32_t>()[0];
+  int dec_batch = dec_batch_tensor.data<int32_t>()[0];
+    
+  baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
+  baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp{const_cast<int32_t*>(encoder_batch_map_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_map.data<int32_t>())};
+  baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp{const_cast<int32_t*>(decoder_batch_map_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_map.data<int32_t>())};
 
-  std::vector<int> seq_lens_encoder_cpu(bsz, 0);
-  std::vector<int> seq_lens_decoder_cpu(bsz, 0);
-  std::vector<int> encoder_batch_map; // 去除空隙的batch map
-  std::vector<int> decoder_batch_map; // 去除空隙的batch map
-  std::vector<int> encoder_seq_lod;
-  int r = xpu_memcpy(seq_lens_encoder_cpu.data(),
-                 seq_lens_encoder.data<int>(),
-                 sizeof(int32_t) * bsz,
-                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
-  r = xpu_memcpy(seq_lens_decoder_cpu.data(),
-                 seq_lens_decoder.data<int>(),
-                 sizeof(int32_t) * bsz,
-                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
-  int enc_batch = 0, dec_batch = 0;
-  int batch_offset = 0;
-  encoder_seq_lod.push_back(0);
-  for(int i = 0; i < bsz; ++i){
-    if(seq_lens_encoder_cpu[i] > 0){
-      enc_batch++;
-      encoder_batch_map.push_back(i);
-      encoder_seq_lod.push_back(seq_lens_encoder_cpu[i]);
-      encoder_seq_lod[enc_batch] += encoder_seq_lod[enc_batch - 1];
-    }
-    else if(seq_lens_decoder_cpu[i] > 0){
-      dec_batch++;
-      decoder_batch_map.push_back(i);
-    }
-    else{
-        batch_offset++;
-    }
-  }         
-  baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp =
-      baidu::xpu::api::VectorParam<int32_t>{encoder_seq_lod.data(), enc_batch + 1, nullptr}
-          .to_xpu(RAII_GUARD);
-  baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp =
-      baidu::xpu::api::VectorParam<int32_t>{encoder_batch_map.data(), enc_batch, nullptr}
-          .to_xpu(RAII_GUARD);
-  baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp =
-      baidu::xpu::api::VectorParam<int32_t>{decoder_batch_map.data(), dec_batch, nullptr}
-          .to_xpu(RAII_GUARD);
   auto out = paddle::full({bsz, dim}, -2, tmp_out.type(), tmp_out.place()); 
 
-  r = xftkernel::xft_eb_gather_next_token<XPUType, XPUType>(
+  int r = xftkernel::xft_eb_gather_next_token<XPUType, XPUType>(
           xpu_ctx->x_context(),
           reinterpret_cast<const XPUType*>(tmp_out.data<data_t>()),
           reinterpret_cast<XPUType*>(out.data<data_t>()),
@@ -87,8 +58,14 @@ std::vector<paddle::Tensor> GatherNextToken(const paddle::Tensor& tmp_out, // [t
 
 std::vector<std::vector<int64_t>> GatherNextTokenInferShape(const std::vector<int64_t>& tmp_out_shape,
                                                              const std::vector<int64_t>& cum_offsets_shape,
-                                                             const std::vector<int64_t>& seq_lens_decoder_shape,
-                                                             const std::vector<int64_t>& seq_lens_encoder_shape,
+                                                             const std::vector<int64_t>& encoder_seq_lod_shape,
+                                                             const std::vector<int64_t>& encoder_batch_map_shape,
+                                                             const std::vector<int64_t>& decoder_batch_map_shape,
+                                                             const std::vector<int64_t>& encoder_seq_lod_cpu_shape,
+                                                             const std::vector<int64_t>& encoder_batch_map_cpu_shape,
+                                                             const std::vector<int64_t>& decoder_batch_map_cpu_shape,
+                                                             const std::vector<int64_t>& enc_batch_tensor_shape,
+                                                             const std::vector<int64_t>& dec_batch_tensor_shape, 
                                                              const paddle::optional<std::vector<int64_t>>& output_padding_offset_shape) {
     if (output_padding_offset_shape) {
       PD_THROW("speculative decoding is not supported in XPU.");
@@ -100,14 +77,20 @@ std::vector<std::vector<int64_t>> GatherNextTokenInferShape(const std::vector<in
 
 std::vector<paddle::DataType> GatherNextTokenInferDtype(const paddle::DataType& tmp_out_dtype,
                                                          const paddle::DataType& cum_offsets_dtype,
-                                                         const paddle::DataType& seq_lens_decoder_dtype,
-                                                         const paddle::DataType& seq_lens_encoder_dtype,
+                                                         const paddle::DataType& encoder_seq_lod_dtype,
+                                                         const paddle::DataType& encoder_batch_map_dtype,
+                                                         const paddle::DataType& decoder_batch_map_dtype,
+                                                         const paddle::DataType& encoder_seq_lod_cpu_dtype,
+                                                         const paddle::DataType& encoder_batch_map_cpu_dtype,
+                                                         const paddle::DataType& decoder_batch_map_cpu_dtype,
+                                                         const paddle::DataType& enc_batch_tensor_dtype,
+                                                         const paddle::DataType& dec_batch_tensor_dtype,
                                                          const paddle::optional<paddle::DataType>& output_padding_offset_dtype) {
     return {tmp_out_dtype};
 }
 
 PD_BUILD_OP(gather_next_token)
-    .Inputs({"tmp_out", "cum_offsets", "seq_lens_decoder", "seq_lens_encoder", paddle::Optional("output_padding_offset")})
+    .Inputs({"tmp_out", "cum_offsets", "encoder_seq_lod", "encoder_batch_map", "decoder_batch_map","encoder_seq_lod_cpu", "encoder_batch_map_cpu", "decoder_batch_map_cpu", "enc_batch_tensor", "dec_batch_tensor", paddle::Optional("output_padding_offset")})
     .Outputs({"out"})
     .Attrs({"max_input_length: int"})
     .SetKernelFn(PD_KERNEL(GatherNextToken))

csrc/xpu/src/get_infer_param.cc

@@ -0,0 +1,196 @@
+// Copyright (c) 2024 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include "paddle/extension.h"
+#include "xpu/plugin.h"
+#include <xft/xdnn_plugin.h>
+namespace xftkernel = baidu::xpu::xftkernel;
+std::vector<paddle::Tensor> GetInferParam(const paddle::Tensor& seq_lens_encoder,
+                                          const paddle::Tensor& seq_lens_decoder) {
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] GetInferParam" );
+
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  const int bsz = seq_lens_encoder.dims()[0];
+  // 判断逻辑
+  std::vector<int32_t> seq_lens_encoder_vec(bsz, 0); // input
+  std::vector<int32_t> seq_lens_decoder_vec(bsz, 0); // input
+
+  std::vector<int32_t> encoder_batch_map_vec(bsz, 0); // 
+  std::vector<int32_t> decoder_batch_map_vec(bsz, 0); // 
+  std::vector<int32_t> encoder_batch_idx_vec(bsz, 0); // 去除空隙的batch map
+  std::vector<int32_t> decoder_batch_idx_vec(bsz, 0); // 去除空隙的batch map
+  std::vector<int32_t> encoder_seq_lod_vec(bsz + 1, 0);
+  std::vector<int32_t> decoder_context_len_vec(bsz, 0);
+  std::vector<int32_t> decoder_context_len_cache_vec(bsz, 0);
+  xpu_wait(xpu_ctx->x_context()->xpu_stream); // 是否需要！！！！TODO
+  int r = xpu_memcpy(seq_lens_encoder_vec.data(),
+                 seq_lens_encoder.data<int32_t>(),
+                 sizeof(int32_t) * bsz,
+                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
+  r = xpu_memcpy(seq_lens_decoder_vec.data(),
+                 seq_lens_decoder.data<int32_t>(),
+                 sizeof(int32_t) * bsz,
+                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
+
+  int enc_batch = 0, dec_batch = 0;
+  int total_enc_len = 0;
+  int batch_offset = 0;
+  for(int i = 0; i < bsz; ++i){
+    if(seq_lens_encoder_vec[i] > 0){
+      enc_batch++;
+      total_enc_len += seq_lens_encoder_vec[i];
+      encoder_batch_map_vec[enc_batch - 1] = i;
+      encoder_batch_idx_vec[enc_batch - 1] = i - batch_offset;
+      encoder_seq_lod_vec[enc_batch] = seq_lens_encoder_vec[i] + encoder_seq_lod_vec[enc_batch - 1];
+    }
+    else if(seq_lens_decoder_vec[i] > 0){
+      dec_batch++;
+      decoder_batch_map_vec[dec_batch - 1] = i;
+      decoder_batch_idx_vec[dec_batch - 1] = i - batch_offset;
+      decoder_context_len_vec[dec_batch - 1] = seq_lens_decoder_vec[i] + 1;
+      decoder_context_len_cache_vec[dec_batch - 1] = seq_lens_decoder_vec[i];
+    }
+    else{
+        batch_offset++;
+    }
+  }
+
+  auto encoder_batch_map_xpu = paddle::full({encoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+  auto decoder_batch_map_xpu = paddle::full({decoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+  auto encoder_batch_idx_xpu = paddle::full({encoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+  auto decoder_batch_idx_xpu = paddle::full({decoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+  auto encoder_seq_lod_xpu = paddle::full({encoder_seq_lod_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+  auto decoder_context_len_xpu = paddle::full({decoder_context_len_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+  auto decoder_context_len_cache_xpu = paddle::full({decoder_context_len_cache_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+
+  auto encoder_batch_map_cpu = paddle::full({encoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto decoder_batch_map_cpu = paddle::full({decoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto encoder_batch_idx_cpu = paddle::full({encoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto decoder_batch_idx_cpu = paddle::full({decoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto encoder_seq_lod_cpu = paddle::full({encoder_seq_lod_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto decoder_context_len_cpu = paddle::full({decoder_context_len_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto decoder_context_len_cache_cpu = paddle::full({decoder_context_len_cache_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  
+  int ret = 0;
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(encoder_batch_map_xpu.data<int32_t>())), encoder_batch_map_vec.data(), sizeof(int32_t) * encoder_batch_map_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_batch_map_xpu.data<int32_t>())), decoder_batch_map_vec.data(), sizeof(int32_t) * decoder_batch_map_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(encoder_batch_idx_xpu.data<int32_t>())), encoder_batch_idx_vec.data(), sizeof(int32_t) * encoder_batch_idx_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_batch_idx_xpu.data<int32_t>())), decoder_batch_idx_vec.data(), sizeof(int32_t) * decoder_batch_idx_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(encoder_seq_lod_xpu.data<int32_t>())), encoder_seq_lod_vec.data(), sizeof(int32_t) * encoder_seq_lod_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_context_len_xpu.data<int32_t>())), decoder_context_len_vec.data(), sizeof(int32_t) * decoder_context_len_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_context_len_cache_xpu.data<int32_t>())), decoder_context_len_cache_vec.data(), sizeof(int32_t) * decoder_context_len_cache_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  
+  std::memcpy(encoder_batch_map_cpu.data<int32_t>(), encoder_batch_map_vec.data(), sizeof(int32_t) * encoder_batch_map_vec.size());
+  std::memcpy(decoder_batch_map_cpu.data<int32_t>(), decoder_batch_map_vec.data(), sizeof(int32_t) * decoder_batch_map_vec.size());
+  std::memcpy(encoder_batch_idx_cpu.data<int32_t>(), encoder_batch_idx_vec.data(), sizeof(int32_t) * encoder_batch_idx_vec.size());
+  std::memcpy(decoder_batch_idx_cpu.data<int32_t>(), decoder_batch_idx_vec.data(), sizeof(int32_t) * decoder_batch_idx_vec.size());
+  std::memcpy(encoder_seq_lod_cpu.data<int32_t>(), encoder_seq_lod_vec.data(), sizeof(int32_t) * encoder_seq_lod_vec.size());
+  std::memcpy(decoder_context_len_cpu.data<int32_t>(), decoder_context_len_vec.data(), sizeof(int32_t) * decoder_context_len_vec.size());
+  std::memcpy(decoder_context_len_cache_cpu.data<int32_t>(), decoder_context_len_cache_vec.data(), sizeof(int32_t) * decoder_context_len_cache_vec.size());
+  
+  // auto batch_tensor = paddle::full({3}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
+  // auto batch_tensor_data = const_cast<int32_t*>(batch_tensor.data<int32_t>());
+  // batch_tensor_data[0] = enc_batch;
+  // batch_tensor_data[1] = dec_batch;
+  // batch_tensor_data[2] = encoder_seq_lod_vec[enc_batch];
+  auto enc_batch_tensor = paddle::full({1}, enc_batch, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto dec_batch_tensor = paddle::full({1}, dec_batch, seq_lens_encoder.type(), paddle::CPUPlace());
+  auto total_enc_len_tensor = paddle::full({1}, total_enc_len, seq_lens_encoder.type(), paddle::CPUPlace());
+
+
+  baidu::xpu::api::plugin::print_times("[TIME END] GetInferParam" );
+  return {encoder_batch_map_xpu, 
+          decoder_batch_map_xpu, 
+          encoder_batch_idx_xpu, 
+          decoder_batch_idx_xpu, 
+          encoder_seq_lod_xpu, 
+          decoder_context_len_xpu, 
+          decoder_context_len_cache_xpu,
+          encoder_batch_map_cpu,
+          decoder_batch_map_cpu,
+          encoder_batch_idx_cpu,
+          decoder_batch_idx_cpu,
+          encoder_seq_lod_cpu,
+          decoder_context_len_cpu,
+          decoder_context_len_cache_cpu,
+          enc_batch_tensor,
+          dec_batch_tensor,
+          total_enc_len_tensor};
+}
+
+std::vector<std::vector<int64_t>> GetInferParamInferShape(const std::vector<int64_t>& seq_lens_encoder_shape,
+                                                             const std::vector<int64_t>& seq_lens_decoder_shape) {
+    return {seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            {seq_lens_encoder_shape[0] + 1}, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape, 
+            seq_lens_encoder_shape,
+            {1},{1},{1}};
+}
+
+std::vector<paddle::DataType> GetInferParamInferDtype(const paddle::DataType& seq_lens_encoder_dtype,
+                                                         const paddle::DataType& seq_lens_decoder_dtype) {
+    return {seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype,
+            seq_lens_encoder_dtype};
+}
+
+PD_BUILD_OP(get_infer_param)
+    .Inputs({"seq_lens_encoder", "seq_lens_decoder"})
+    .Outputs({"encoder_batch_map_xpu", 
+              "decoder_batch_map_xpu", 
+              "encoder_batch_idx_xpu", 
+              "decoder_batch_idx_xpu", 
+              "encoder_seq_lod_xpu", 
+              "decoder_context_len_xpu", 
+              "decoder_context_len_cache_xpu", 
+              "encoder_batch_map_cpu", 
+              "decoder_batch_map_cpu", 
+              "encoder_batch_idx_cpu", 
+              "decoder_batch_idx_cpu", 
+              "encoder_seq_lod_cpu",
+              "decoder_context_len_cpu",
+              "decoder_context_len_cache_cpu",
+              "enc_batch_tensor",
+              "dec_batch_tensor",
+              "total_enc_len_tensor"})
+    .SetKernelFn(PD_KERNEL(GetInferParam))
+    .SetInferShapeFn(PD_INFER_SHAPE(GetInferParamInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GetInferParamInferDtype));

csrc/xpu/src/get_output.cc

@@ -0,0 +1,70 @@
+// Copyright (c) 2024 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.
+
+#include <stdio.h>
+#include <string.h>
+#include <sys/ipc.h>
+#include <sys/msg.h>
+#include <sys/types.h>
+#include "paddle/extension.h"
+
+#define MAX_BSZ 512
+
+struct msgdata {
+    long mtype;
+    int mtext[MAX_BSZ + 2];   // stop_flag, bsz, tokens
+};
+
+void GetOutput(const paddle::Tensor& x,
+               const paddle::Tensor& msg_queue_id,
+               int64_t rank_id,
+               bool wait_flag) {
+  if (rank_id > 0) return;
+
+  static struct msgdata msg_rcv;
+  int msg_queue_id_val = msg_queue_id.data<int>()[0];
+  static key_t key = ftok("./", msg_queue_id_val);
+
+  static int msgid = msgget(key, IPC_CREAT | 0666);
+
+  int64_t *out_data = const_cast<int64_t*>(x.data<int64_t>());
+  int ret = -1;
+  if (!wait_flag) {
+    ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, IPC_NOWAIT);
+  } else {
+    ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, 0);
+  }
+  if(ret == -1)
+	{
+    // read none
+    out_data[0] = -2;
+    out_data[1] = 0;
+		return;
+	}
+
+  int bsz = msg_rcv.mtext[1];
+
+  for (int64_t i = 0; i < bsz + 2; i++) {
+    out_data[i] = (int64_t)msg_rcv.mtext[i];
+  }
+  return;
+}
+
+PD_BUILD_OP(get_output)
+    .Inputs({"x", "msg_queue_id"})
+    .Attrs({"rank_id: int64_t",
+            "wait_flag: bool"})
+    .Outputs({"x_out"})
+    .SetInplaceMap({{"x", "x_out"}})
+    .SetKernelFn(PD_KERNEL(GetOutput));

csrc/xpu/src/get_position_ids_v2.cc

@@ -28,6 +28,8 @@ std::vector<paddle::Tensor> GetPositionIdsKernelV2(
     const paddle::Tensor& seq_lens_decoder,
     const paddle::Tensor& seq_lens_this_time) {
 
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] GetPositionIdsKernelV2" );
+
   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
@@ -54,7 +56,9 @@ std::vector<paddle::Tensor> GetPositionIdsKernelV2(
         const_cast<int32_t*>(position_ids_out.data<int32_t>()),
         bs
   );
-  PD_CHECK(ret == 0, "api::plugin::get_position_ids failed");
+  PD_CHECK(ret == 0, "api::plugin::get_position_ids failed"); 
+  
+  baidu::xpu::api::plugin::print_times("[TIME END] GetPositionIdsKernelV2" );
   return {position_ids_out};
 }
 

csrc/xpu/src/get_token_penalty_multi_scores_v2.cc

@@ -27,6 +27,7 @@ void TokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
                              const paddle::Tensor& cur_len,
                              const paddle::Tensor& min_len,
                              const paddle::Tensor& eos_token_id) {
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] TokenPenaltyMultiScores");
   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
@@ -90,8 +91,11 @@ void TokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
           "Only float16 and float32 are supported. ");
       break;
   }
+
+  baidu::xpu::api::plugin::print_times("[TIME END] TokenPenaltyMultiScores");
 }
 
+
 PD_BUILD_OP(get_token_penalty_multi_scores_v2)
     .Inputs({"pre_ids",
              "logits",

csrc/xpu/src/mla_block_multihead_attention_xpu.cc

@@ -323,7 +323,8 @@ std::vector<paddle::Tensor> MlaAttn(
             block_size, // block_size
             max_block_per_seq, // max_blocks_per_seq (prefix cache)
             page_param_.max_context_len_, // prefill_len
-            nullptr); // block_tables (prefix cache)
+            nullptr,
+            softmax_scale * sqrt(dim_qk)); // block_tables (prefix cache)
 
     // std::cout << "fmha kernel done " <<std::endl;
   }
@@ -425,7 +426,7 @@ std::vector<paddle::Tensor> MlaAttn(
             const_cast<TKV*>(key_cache_xft.data()),
             const_cast<TKV*>(value_cache_xft.data()),
             attn_param_.kv_head_num_,
-            attn_param_.scale_,
+            softmax_scale,
             block_tables_xft.data(),
             attn_param_.context_len_vp_,
             attn_param_.valid_batch_vp_,

csrc/xpu/src/mla_cache_kernel_xpu.cc

@@ -0,0 +1,477 @@
+// Copyright (c) 2024 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.
+#include <core/ctx_manager.h>
+#include <core/xft_check.h>
+#include <core/xft_event.h>
+#include <core/xft_params.h>
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include <xft/operation/xft_fc_helper.h>
+#include <xft/xdnn_plugin.h>
+
+#include "paddle/extension.h"
+#include "paddle/phi/core/enforce.h"
+#include "xpu/plugin.h"
+#include "ops.h"
+
+namespace xftkernel = baidu::xpu::xftkernel;
+
+template <paddle::DataType D>
+class PaddleTypeToXPUType;
+
+template <>
+class PaddleTypeToXPUType<paddle::DataType::FLOAT32> {
+public:
+  typedef float DataType;
+  typedef float data_t;
+};
+
+template <>
+class PaddleTypeToXPUType<paddle::DataType::FLOAT16> {
+public:
+  typedef float16 DataType;
+  typedef paddle::float16 data_t;
+};
+
+template <>
+class PaddleTypeToXPUType<paddle::DataType::BFLOAT16> {
+public:
+  typedef bfloat16 DataType;
+  typedef paddle::bfloat16 data_t;
+};
+
+template <paddle::DataType T>
+std::vector<paddle::Tensor> PrefillMLAWriteCache(
+                    const paddle::Tensor& kv_nope,
+                    const paddle::Tensor& kv_pe,
+                    const paddle::Tensor& encoder_seq_lod,
+                    const paddle::Tensor& encoder_batch_map,
+                    const paddle::Tensor& start_token_raw, // encoder cache写时 start token为0 （不算prefix cache）
+                    const paddle::Tensor& encoder_seq_lod_cpu,
+                    const paddle::Tensor& encoder_batch_map_cpu,
+                    const paddle::Tensor& start_token_raw_cpu, // encoder cache写时 start token为0 （不算prefix cache）
+                    const paddle::Tensor& enc_batch_tensor,
+                    const paddle::Tensor& dec_batch_tensor,
+                    const paddle::Tensor& padding_offsets,
+                    const paddle::Tensor& cum_offsets,
+                    const paddle::Tensor& block_tables,
+                    const int max_seq_len,
+                    const int kv_num_heads,
+                    paddle::Tensor* kv_cache) {
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] PrefillMLAWriteCache");
+
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  typedef PaddleTypeToXPUType<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  const auto& kv_nope_dims = kv_nope.dims();
+  const auto& kv_pe_dims = kv_pe.dims();
+  const auto& kv_cache_dims = kv_cache->dims();
+  auto num_tokens = kv_nope_dims[0];
+  auto block_size = kv_cache_dims[2];
+  auto block_batch = block_tables.dims()[0];
+  auto max_blocks_per_seq = block_tables.dims()[1];
+  auto kv_lora_rank = kv_nope_dims[kv_nope_dims.size() - 1];
+  auto rope_head_dim = kv_pe_dims[kv_pe_dims.size() - 1];
+  auto bsz = cum_offsets.dims()[0];
+  int enc_batch = enc_batch_tensor.data<int32_t>()[0];
+  int dec_batch = dec_batch_tensor.data<int32_t>()[0];
+  // 初始化输入
+  auto kv_nope_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
+      reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+      std::array<int64_t, 2>{num_tokens,
+                             kv_lora_rank});
+  auto kv_pe_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
+      reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+      std::array<int64_t, 2>{num_tokens,
+                             rope_head_dim});
+  // 初始化输入 cache
+  auto cache_xft = baidu::xpu::xft::xftTensor<DataType_, 4>(
+  reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_cache->data<data_t>())),
+  std::array<int64_t, 4>{kv_cache->shape()[0],
+                          kv_cache->shape()[1],
+                          kv_cache->shape()[2],
+                          kv_cache->shape()[3]});  
+  // 初始化输入：block table
+  auto block_tables_xft = baidu::xpu::xft::xftTensor<int, 2>(
+  reinterpret_cast<int*>(const_cast<int*>(block_tables.data<int>())),
+  std::array<int64_t, 2>{block_tables.shape()[0],
+                          block_tables.shape()[1]}); 
+  // 初始化concat
+  std::vector<const DataType_*> concat_input;
+  concat_input.push_back(kv_nope_xft.data());
+  concat_input.push_back(kv_pe_xft.data());
+  auto concat_out = paddle::full({num_tokens, kv_lora_rank + rope_head_dim}, -2, kv_nope.type(), kv_nope.place()); 
+  auto concat_out_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
+      reinterpret_cast<DataType_*>(const_cast<data_t*>(concat_out.data<data_t>())),
+      std::array<int64_t, 2>{concat_out.shape()[0],
+                             concat_out.shape()[1]});
+
+  // 拼接kv_nope和kv_pe
+  baidu::xpu::api::concat<DataType_>(xpu_ctx->x_context(),
+                                      concat_input,
+                                      concat_out_xft.data(),
+                                      {{num_tokens, kv_lora_rank},{num_tokens, rope_head_dim}},
+                                      1);
+
+
+  baidu::xpu::api::VectorParam<int32_t> context_len_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
+  baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(encoder_batch_map_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_map.data<int32_t>())};
+  baidu::xpu::api::VectorParam<int32_t> start_tokens_vp{const_cast<int32_t*>(start_token_raw_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(start_token_raw.data<int32_t>())};
+  int ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
+          xpu_ctx->x_context(),
+          reinterpret_cast<float16*>(concat_out_xft.data()),
+          reinterpret_cast<float16*>(cache_xft.data()),
+          block_tables_xft.data(),
+          context_len_vp,
+          start_tokens_vp,
+          valid_batch_vp,
+          enc_batch,
+          1, // num_head == 1 吸收版
+          kv_lora_rank + rope_head_dim,
+          block_batch,
+          block_size,
+          max_blocks_per_seq,
+          "BLHD", // qkv_layout,
+          "HLD", // page_layout,
+          nullptr, // scale,
+          nullptr, // batch_max_ptrs, 
+          nullptr); // max_ptrs
+
+  baidu::xpu::api::plugin::print_times("[TIME END] PrefillMLAWriteCache");
+
+  return {};
+}
+
+std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& encoder_seq_lod,
+    const paddle::Tensor& encoder_batch_map,
+    const paddle::Tensor& start_token_raw, // encoder cache写时 start token为0 （不算prefix cache）
+    const paddle::Tensor& encoder_seq_lod_cpu,
+    const paddle::Tensor& encoder_batch_map_cpu,
+    const paddle::Tensor& start_token_raw_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len,
+    const int kv_num_heads) {
+
+  switch (kv_pe.dtype()) {
+    case paddle::DataType::BFLOAT16: {
+      return PrefillMLAWriteCache<paddle::DataType::BFLOAT16>(
+                              kv_nope,
+                              kv_pe,
+                              encoder_seq_lod,
+                              encoder_batch_map,
+                              start_token_raw,
+                              encoder_seq_lod_cpu,
+                              encoder_batch_map_cpu,                            
+                              start_token_raw_cpu,
+                              enc_batch_tensor,
+                              dec_batch_tensor,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              kv_num_heads,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+    case paddle::DataType::FLOAT16: {
+      return PrefillMLAWriteCache<paddle::DataType::FLOAT16>(
+                              kv_nope,
+                              kv_pe,
+                              encoder_seq_lod,
+                              encoder_batch_map,
+                              start_token_raw,
+                              encoder_seq_lod_cpu,
+                              encoder_batch_map_cpu,                            
+                              start_token_raw_cpu,
+                              enc_batch_tensor,
+                              dec_batch_tensor,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              kv_num_heads,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+  }
+  return {};
+}
+
+template <paddle::DataType T>
+std::vector<paddle::Tensor> DecodeMLAWriteCache(
+                    const paddle::Tensor& kv_nope,
+                    const paddle::Tensor& kv_pe,
+                    const paddle::Tensor& decoder_context_len_cache,
+                    const paddle::Tensor& decoder_batch_map,
+                    const paddle::Tensor& kv_seq_lod_raw, // decoder cache写时，lod为0，1，2，3，……
+                    const paddle::Tensor& decoder_context_len_cache_cpu,
+                    const paddle::Tensor& decoder_batch_map_cpu,
+                    const paddle::Tensor& kv_seq_lod_raw_cpu, // decoder cache写时，lod为0，1，2，3，……
+                    const paddle::Tensor& enc_batch_tensor,
+                    const paddle::Tensor& dec_batch_tensor,
+                    const paddle::Tensor& padding_offsets,
+                    const paddle::Tensor& cum_offsets,
+                    const paddle::Tensor& block_tables,
+                    const int max_seq_len,
+                    const int kv_num_heads,
+                    const bool speculate_decoder,
+                    paddle::Tensor* kv_cache) {
+ 
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] DecodeMLAWriteCache");
+                       
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
+  typedef PaddleTypeToXPUType<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  const auto& kv_nope_dims = kv_nope.dims();
+  const auto& kv_pe_dims = kv_pe.dims();
+  const auto& kv_cache_dims = kv_cache->dims();
+  auto num_tokens = kv_nope_dims[0];
+  auto block_size = kv_cache_dims[2];
+  auto block_batch = block_tables.dims()[0];
+  auto max_blocks_per_seq = block_tables.dims()[1];
+  auto kv_lora_rank = kv_nope_dims[kv_nope_dims.size() - 1];
+  auto rope_head_dim = kv_pe_dims[kv_pe_dims.size() - 1];
+  auto bsz = cum_offsets.dims()[0];
+  int enc_batch = enc_batch_tensor.data<int32_t>()[0];
+  int dec_batch = dec_batch_tensor.data<int32_t>()[0];
+
+  // 初始化输入
+  auto kv_nope_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
+      reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+      std::array<int64_t, 2>{num_tokens,
+                             kv_lora_rank});
+  auto kv_pe_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
+      reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+      std::array<int64_t, 2>{num_tokens,
+                             rope_head_dim});
+  // 初始化输入 cache
+  auto cache_xft = baidu::xpu::xft::xftTensor<DataType_, 4>(
+  reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_cache->data<data_t>())),
+  std::array<int64_t, 4>{kv_cache->shape()[0],
+                          kv_cache->shape()[1],
+                          kv_cache->shape()[2],
+                          kv_cache->shape()[3]});  
+  // 初始化输入：block table
+  auto block_tables_xft = baidu::xpu::xft::xftTensor<int, 2>(
+  reinterpret_cast<int*>(const_cast<int*>(block_tables.data<int>())),
+  std::array<int64_t, 2>{block_tables.shape()[0],
+                          block_tables.shape()[1]}); 
+  // 初始化concat
+  // std::vector<DataType_*> concat_input = {kv_nope_xft.data(), kv_pe_xft.data()};
+  std::vector<const DataType_*> concat_input;
+  concat_input.push_back(kv_nope_xft.data());
+  concat_input.push_back(kv_pe_xft.data());
+  auto concat_out = paddle::full({num_tokens, kv_lora_rank + rope_head_dim}, -2, kv_nope.type(), kv_nope.place()); 
+  auto concat_out_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
+      reinterpret_cast<DataType_*>(const_cast<data_t*>(concat_out.data<data_t>())),
+      std::array<int64_t, 2>{concat_out.shape()[0],
+                             concat_out.shape()[1]});
+
+  // 拼接kv_nope和kv_pe
+  baidu::xpu::api::concat<DataType_>(xpu_ctx->x_context(),
+                                      concat_input,
+                                      concat_out_xft.data(),
+                                      {{num_tokens, kv_lora_rank},{num_tokens, rope_head_dim}},
+                                      1);
+
+
+
+
+  // context_len
+  baidu::xpu::api::VectorParam<int32_t> context_len_vp_cache{const_cast<int32_t*>(decoder_context_len_cache_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_context_len_cache.data<int32_t>())};
+  // real batch     
+  baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(decoder_batch_map_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_map.data<int32_t>())};
+  // baidu::xpu::api::VectorParam<int32_t> kv_seq_lod_vp{const_cast<int32_t*>(kv_seq_lod_raw_cpu.data<int32_t>()), dec_batch + 1, const_cast<int32_t*>(kv_seq_lod_raw.data<int32_t>())};
+  
+  // baidu::xpu::api::VectorParam<int32_t> kv_seq_lod_vp = baidu::xpu::api::VectorParam<int32_t>{const_cast<int32_t*>(kv_seq_lod_raw_cpu.data<int32_t>()), dec_batch + 1, const_cast<int32_t*>(kv_seq_lod_raw.data<int32_t>())};
+  
+  std::vector<int> kv_seq_lod(dec_batch + 1);
+  std::iota(kv_seq_lod.begin(), kv_seq_lod.end(), 0);
+  baidu::xpu::api::VectorParam<int32_t> kv_seq_lod_vp =
+      baidu::xpu::api::VectorParam<int32_t>{kv_seq_lod.data(), dec_batch + 1, nullptr}.to_xpu(RAII_GUARD);
+
+
+    // std::cout << "Tensor dtype: " << kv_seq_lod_raw_cpu.dtype() << std::endl;
+    // std::cout << "Tensor place: " << kv_seq_lod_raw_cpu.place() << std::endl;
+    
+    // // 打印形状信息
+    // std::cout << "Tensor shape: [";
+    // for (int i = 0; i < kv_seq_lod_raw_cpu.dims().size(); ++i) {
+    //     if (i > 0) std::cout << ", ";
+    //     std::cout << kv_seq_lod_raw_cpu.dims()[i];
+    // }
+    // std::cout << "]" << std::endl;
+
+    // std::cout << "Tensor dtype: " << kv_seq_lod_raw.dtype() << std::endl;
+    // std::cout << "Tensor place: " << kv_seq_lod_raw.place() << std::endl;
+    
+    // // 打印形状信息
+    // std::cout << "Tensor shape: [";
+    // for (int i = 0; i < kv_seq_lod_raw.dims().size(); ++i) {
+    //     if (i > 0) std::cout << ", ";
+    //     std::cout << kv_seq_lod_raw.dims()[i];
+    // }
+    // std::cout << "]" << std::endl;
+
+
+  int ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
+          xpu_ctx->x_context(),
+          reinterpret_cast<float16*>(concat_out_xft.data()),
+          reinterpret_cast<float16*>(cache_xft.data()),
+          block_tables_xft.data(),
+          kv_seq_lod_vp,
+          context_len_vp_cache,
+          valid_batch_vp,
+          dec_batch,
+          1, //  num_head == 1
+          kv_lora_rank + rope_head_dim,
+          block_batch,
+          block_size,
+          max_blocks_per_seq,
+          "BLHD", // qkv_layout,
+          "HLD", // page_layout,
+          nullptr, // scale,
+          nullptr, // batch_max_ptrs, 
+          nullptr); // max_ptrs
+
+ 
+  baidu::xpu::api::plugin::print_times("[TIME END] DecodeMLAWriteCache");
+
+  return {};
+}
+
+std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& decoder_context_len_cache,
+    const paddle::Tensor& decoder_batch_map,
+    const paddle::Tensor& kv_seq_lod_raw,
+    const paddle::Tensor& decoder_context_len_cache_cpu,
+    const paddle::Tensor& decoder_batch_map_cpu,
+    const paddle::Tensor& kv_seq_lod_raw_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len,
+    const int kv_num_heads,
+    const bool speculate_decoder) {
+  switch (kv_pe.dtype()) {
+    case paddle::DataType::BFLOAT16: {
+      return DecodeMLAWriteCache<paddle::DataType::BFLOAT16>(
+                              kv_nope,
+                              kv_pe,
+                              decoder_context_len_cache,
+                              decoder_batch_map,
+                              kv_seq_lod_raw,
+                              decoder_context_len_cache_cpu,
+                              decoder_batch_map_cpu,
+                              kv_seq_lod_raw_cpu,
+                              enc_batch_tensor,
+                              dec_batch_tensor,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              kv_num_heads,
+                              speculate_decoder,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+    case paddle::DataType::FLOAT16: {
+      return DecodeMLAWriteCache<paddle::DataType::FLOAT16>(
+                              kv_nope,
+                              kv_pe,
+                              decoder_context_len_cache,
+                              decoder_batch_map,
+                              kv_seq_lod_raw,
+                              decoder_context_len_cache_cpu,
+                              decoder_batch_map_cpu,
+                              kv_seq_lod_raw_cpu,
+                              enc_batch_tensor,
+                              dec_batch_tensor,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              kv_num_heads,
+                              speculate_decoder,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+  }
+  return {};
+}
+
+
+PD_BUILD_OP(prefill_mla_write_cache_xpu)
+    .Inputs({"kv_nope",
+             "kv_pe",
+             "kv_cache",
+             "encoder_seq_lod",
+             "encoder_batch_map",
+             "start_token_raw",
+             "encoder_seq_lod_cpu",
+             "encoder_batch_map_cpu",
+             "start_token_raw_cpu",
+             "enc_batch_tensor",
+             "dec_batch_tensor",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables"})
+    .Outputs({"kv_cache_out"})
+    .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
+    .Attrs({"cache_quant_type_str: std::string",
+            "max_seq_len: int",
+            "kv_num_heads: int"})
+    .SetKernelFn(PD_KERNEL(PrefillMLAWriteCacheKernel));
+
+PD_BUILD_OP(decode_mla_write_cache_xpu)
+    .Inputs({"kv_nope",
+             "kv_pe",
+             "kv_cache",
+             "decoder_context_len_cache",
+             "decoder_batch_map",
+             "kv_seq_lod_raw",
+             "decoder_context_len_cache_cpu",
+             "decoder_batch_map_cpu",
+             "kv_seq_lod_raw_cpu",
+             "enc_batch_tensor",
+             "dec_batch_tensor",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables"})
+    .Outputs({"kv_cache_out"})
+    .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
+    .Attrs({"cache_quant_type_str: std::string",
+            "max_seq_len: int",
+            "kv_num_heads: int",
+            "speculate_decoder: bool"})
+    .SetKernelFn(PD_KERNEL(DecodeMLAWriteCacheKernel));

csrc/xpu/src/moe_fused.cc

@@ -0,0 +1,173 @@
+// Copyright (c) 2024 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.
+
+#include <core/ctx_manager.h>
+#include <core/xft_tensor.h>
+#include <core/xft_check.h>
+#include <core/xft_event.h>
+#include <core/xft_params.h>
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include <xft/operation/ffn_moe_vllm.h>
+#include <xft/xdnn_plugin.h>
+
+#include "paddle/extension.h"
+#include "paddle/phi/core/enforce.h"
+#include "xpu/plugin.h"
+
+namespace xftkernel = baidu::xpu::xftkernel;
+namespace xft = baidu::xpu::xft;
+
+inline uint32_t get_xpu_moe_speed_layers() {
+    static const char* xpu_moe_speed_layers_env = std::getenv("FLAGS_xpu_moe_speed_layers");
+    static const uint32_t xpu_moe_speed_layers =
+            xpu_moe_speed_layers_env == nullptr ? 0 : std::stoul(std::string(xpu_moe_speed_layers_env));
+    return xpu_moe_speed_layers;
+}
+
+std::vector<paddle::Tensor> MoeFusedKernel(
+    const paddle::Tensor& input, //[980. 7168]
+    const paddle::Tensor& gate_weight,//[7168. 256] ->[256,7168]
+    const paddle::Tensor& ffn_inter_weights,//[256, 7168,4096] reshape
+    const paddle::Tensor& ffn_output_weights,//[256, 2048, 7168] reshape
+    const paddle::Tensor& ffn1_weights_scale,// [256, 4096]
+    const paddle::Tensor& ffn2_weights_scale,// [256, 7168]
+    const paddle::Tensor& score_bias, // 256
+    int moe_top_k,
+    int expert_group_num,
+    int moe_topk_group,
+    int layer) {
+
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] MoeFusedKernel");
+
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+
+  using XPUType = typename XPUTypeTrait<bfloat16>::Type;
+  typedef paddle::bfloat16 data_t;
+
+  const int64_t m = input.shape()[0];
+  const int64_t hidden_size = input.shape()[1];
+  const int64_t expert_num = gate_weight.shape()[0];
+  const int64_t ffn_inter_hidden_size = ffn_inter_weights.shape()[2];
+  const int64_t ffn_output_hidden_size = ffn_output_weights.shape()[2];
+
+  xft::Tensor input_tensor(
+      reinterpret_cast<XPUType*>(const_cast<data_t*>(input.data<data_t>())),
+      xft::DataType::DT_BFLOAT16,
+      input.shape());
+  xft::Tensor gate_weight_tensor(
+      const_cast<float*>(gate_weight.data<float>()),
+      xft::DataType::DT_FLOAT,
+      gate_weight.shape());
+  xft::Tensor ffn_inter_weights_tensor(
+      const_cast<int8_t*>(ffn_inter_weights.data<int8_t>()),
+      nullptr,
+      const_cast<float*>(ffn1_weights_scale.data<float>()),
+      xft::DataType::DT_INT8,
+      {expert_num, ffn_inter_hidden_size, hidden_size});
+  xft::Tensor ffn_output_weights_tensor(
+      const_cast<int8_t*>(ffn_output_weights.data<int8_t>()),
+      nullptr,
+      const_cast<float*>(ffn2_weights_scale.data<float>()),
+      xft::DataType::DT_INT8,
+      {expert_num, ffn_output_hidden_size, ffn_inter_hidden_size / 2});
+  xft::Tensor score_bias_tensor(const_cast<float*>(score_bias.data<float>()),
+                                xft::DataType::DT_FLOAT,
+                                score_bias.shape());
+
+  auto output = paddle::full({m, hidden_size}, -1, input.type(), input.place());
+  xft::Tensor output_tensor(reinterpret_cast<XPUType*>(output.data<data_t>()),
+                            xft::DataType::DT_BFLOAT16,
+                            output.shape());
+
+  xft::MoeFFNWeight moe_weight{&gate_weight_tensor,
+                               &ffn_inter_weights_tensor,
+                               nullptr,
+                               &ffn_output_weights_tensor,
+                               nullptr,
+                               nullptr,
+                               nullptr,
+                               &score_bias_tensor};
+  xft::MoeFFNParam moe_param{
+    expert_num,
+    moe_top_k,
+    false,
+    true,
+    expert_group_num,
+    moe_topk_group,
+    true,
+    false,
+    1,
+    0,
+    "sigmoid"
+  };
+
+    if (layer >= get_xpu_moe_speed_layers()) {
+        xft::xft_moe_ffn_block_sorted<XPUType, int8_t, XPUType, float>(
+            xpu_ctx->x_context(), &input_tensor, &output_tensor, moe_weight, moe_param);
+    } else {
+        xft::xft_moe_ffn_block_sorted<XPUType, int8_t, XPUType, int8_wo_t>(
+            xpu_ctx->x_context(), &input_tensor, &output_tensor, moe_weight, moe_param);
+    }
+  
+  baidu::xpu::api::plugin::print_times("[TIME END] MoeFusedKernel");
+
+  return {
+      output,
+  };
+}
+
+
+std::vector<std::vector<int64_t>> MoeFusedInferShape(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<int64_t>& gate_weight_shape,
+    const std::vector<int64_t>& ffn_inter_weights_shape,
+    const std::vector<int64_t>& ffn_output_weights_shape,
+    const std::vector<int64_t>& ffn1_weights_scale_shape,
+    const std::vector<int64_t>& ffn2_weights_scale_shape,
+    const std::vector<int64_t>& score_bias_shape) {
+  return {input_shape};
+}
+
+std::vector<paddle::DataType> MoeFusedInferDtype(
+    const paddle::DataType& input_type,
+    const paddle::DataType& gate_weight_type,
+    const paddle::DataType& ffn_inter_weights_type,
+    const paddle::DataType& ffn_output_weights_type,
+    const paddle::DataType& ffn1_weights_scale_type,
+    const paddle::DataType& ffn2_weights_scale_type,
+    const paddle::DataType& score_bias_type,
+    const int moe_top_k,
+    const int expert_group_num,
+    const int moe_topk_group,
+    const int layer) {
+        return {input_type};
+}
+
+PD_BUILD_OP(moe_fused_xpu)
+    .Inputs({
+        "input",
+        "gate_weight",
+        "ffn_inter_weights",
+        "ffn_outer_weights",
+        "ffn1_weights_scale",
+        "ffn2_weights_scale",
+        "score_bias",
+    })
+    .Outputs({"output"})
+    .Attrs({"moe_top_k: int", "expert_group_num: int", "moe_topk_group: int", "layer: int"})
+    .SetKernelFn(PD_KERNEL(MoeFusedKernel))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeFusedInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeFusedInferDtype));

csrc/xpu/src/ops.h

@@ -0,0 +1,113 @@
+// 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.
+
+#include <paddle/phi/backends/xpu/xpu_context.h>
+#include "paddle/extension.h"
+#include "xpu/plugin.h"
+#include <xft/xdnn_plugin.h>
+
+void FusedRotaryPositionEncoding(
+    paddle::Tensor& query,  // [num_tokens, num_heads, head_size] or
+                            // [num_tokens, num_heads * head_size]
+    paddle::Tensor& key,
+    // [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads *
+    // head_size]
+    const paddle::Tensor& position_ids,   // [num_tokens]
+    const paddle::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    int head_size,
+    bool is_neox);
+
+
+std::vector<paddle::Tensor> MlaDeAttn(
+    const paddle::Tensor& q,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& decoder_context_len,
+    const paddle::Tensor& decoder_batch_map,
+    const paddle::Tensor& decoder_context_len_cpu,
+    const paddle::Tensor& decoder_batch_map_cpu,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const float softmax_scale,
+    const int block_size,
+    const int num_head,
+    const int kv_lora_rank, 
+    const int rope_head_dim,
+    const int dim_qk,
+    const int dim_v);
+
+
+std::vector<paddle::Tensor> MlaEnAttn(
+    const paddle::Tensor& q,
+    const paddle::Tensor& k,
+    const paddle::Tensor& v,
+    const paddle::Tensor& encoder_seq_lod,
+    const paddle::Tensor& encoder_batch_map,
+    const paddle::Tensor& encoder_seq_lod_cpu,
+    const paddle::Tensor& encoder_batch_map_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const float softmax_scale,
+    const int block_size,
+    const int num_head,
+    const int dim_qk,
+    const int dim_v);
+
+
+std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& decoder_context_len_cache,
+    const paddle::Tensor& decoder_batch_map,
+    const paddle::Tensor& kv_seq_lod_raw,
+    const paddle::Tensor& decoder_context_len_cache_cpu,
+    const paddle::Tensor& decoder_batch_map_cpu,
+    const paddle::Tensor& kv_seq_lod_raw_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len,
+    const int kv_num_heads,
+    const bool speculate_decoder);
+
+std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& encoder_seq_lod,
+    const paddle::Tensor& encoder_batch_map,
+    const paddle::Tensor& start_token_raw, // encoder cache写时 start token为0 （不算prefix cache）
+    const paddle::Tensor& encoder_seq_lod_cpu,
+    const paddle::Tensor& encoder_batch_map_cpu,
+    const paddle::Tensor& start_token_raw_cpu,
+    const paddle::Tensor& enc_batch_tensor,
+    const paddle::Tensor& dec_batch_tensor,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len,
+    const int kv_num_heads);
+
+std::vector<paddle::Tensor> Bmm(const paddle::Tensor& input, 
+                                             const paddle::Tensor& weight);

csrc/xpu/src/plugin/include/xpu/plugin.h

@@ -17,11 +17,20 @@
 
 #pragma once
 #include "xpu/xdnn.h"
-
+#include <chrono>
 namespace baidu {
 namespace xpu {
 namespace api {
 namespace plugin {
+static std::string print_times(std::string str){
+#ifdef XPU_PRINT_OP_TIME
+  auto now = std::chrono::system_clock::now();
+  auto duration = now.time_since_epoch();
+  auto milliseconds = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
+  std::cout<<str<<":" <<milliseconds<<std::endl;
+#endif
+  return "";
+} 
 
 template <typename T>
 DLL_EXPORT int set_stop_value_multi_ends(Context* ctx,
@@ -84,20 +93,6 @@ DLL_EXPORT int get_position_ids(Context *ctx,
                        int *position_ids,
                        const int bs);
 
-DLL_EXPORT int update_inputs(Context* ctx,
-                             bool* not_need_stop,
-                             int* seq_lens_this_time,
-                             int* seq_lens_encoder,
-                             int* seq_lens_decoder,
-                             int64_t* input_ids,
-                             const int64_t* stop_nums,
-                             const bool* stop_flags,
-                             const bool* is_block_step,
-                             const int64_t* next_tokens,
-                             const int bsz,
-                             const int max_bsz,
-                             const int input_ids_stride);
-
 template <typename T>
 DLL_EXPORT int rebuild_padding(Context *ctx,
                     T *output_data, // [bs, dim_embed]
@@ -172,6 +167,19 @@ DLL_EXPORT int free_and_dispatch_block(Context *ctx,
                             const int block_size,
                             const int block_num_per_seq,
                             const int max_decoder_block_num);
+
+template <typename T, typename TR>
+DLL_EXPORT int rotary_embedding_neox(
+        Context* ctx,
+        const int* positions,  // [num_tokens]
+        T* query,                  // [num_tokens, num_heads, head_size]
+        T* key,                    // [num_tokens, num_kv_heads, head_size]
+        const TR* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens);
 }  // namespace plugin
 }  // namespace api
 }  // namespace xpu

csrc/xpu/src/plugin/src/kernel/kunlun3cpp/absorb_mla_cache_write.xpu

@@ -0,0 +1,96 @@
+#include "xpu/kernel/cluster.h"
+#include "xpu/kernel/cluster_partition.h"
+#include "xpu/kernel/cluster_primitive.h"
+
+namespace xpu3 {
+namespace plugin {
+
+static inline __device__ void batch_load(
+        _global_ptr_ const int32_t* kv_seq_lod,
+        _global_ptr_ const int32_t* start_token,
+        _global_ptr_ const int32_t* real_batch,
+        int* kv_seq_lod_lm,
+        int* start_token_lm,
+        int* real_batch_lm,
+        int batch_begin,
+        int read_batch_len) {
+    GM2LM_ASYNC(kv_seq_lod + batch_begin, kv_seq_lod_lm, (read_batch_len + 1) * sizeof(int));
+    GM2LM_ASYNC(start_token + batch_begin, start_token_lm, read_batch_len * sizeof(int));
+    GM2LM_ASYNC(real_batch + batch_begin, real_batch_lm, read_batch_len * sizeof(int));
+    mfence();
+}
+
+template <typename TI, typename TO, typename TID>
+__global__ void absorb_mla_cache_write(
+        const TI* src,
+        TO* dst,
+        const TID* block_table,
+        int32_t* kv_seq_lod,
+        int32_t* start_token,
+        int32_t* real_batch,
+        int64_t batch_size,
+        int64_t kv_head_num,
+        int64_t head_dim,
+        int64_t max_seq_num,
+        int64_t block_size,
+        int64_t max_num_blocks_per_seq) {
+    constexpr int loop_batch_size = 32;
+    int cid = core_id();
+    int ncores = core_num();
+    int tid = cid * cluster_num() + cluster_id();
+    int nthreads = ncores * cluster_num();
+    // prepare lm buffer
+    __local__ TI local_scr[2048];
+    __local__ int lod_lm[loop_batch_size];
+    __local__ int kv_seq_lod_lm[loop_batch_size];
+    __local__ int start_token_lm[loop_batch_size];
+    __local__ int real_batch_lm[loop_batch_size];
+    int table_cnt = max_seq_num * max_num_blocks_per_seq;
+    batch_load(
+            kv_seq_lod,
+            start_token,
+            real_batch,
+            kv_seq_lod_lm,
+            start_token_lm,
+            real_batch_lm,
+            0,
+            batch_size);    
+    int token_num = kv_seq_lod[batch_size];
+    for(int i = tid; i < token_num; i += nthreads){
+        int batch = 0;
+        for(int b = 0; b < batch_size; b++){
+            if(i < lod_lm[b + 1]){
+                batch = b;
+                break;
+            }
+            GM2LM(src + i * kv_head_num * head_dim, local_scr, kv_head_num * head_dim * sizeof(TI));
+            TID block_id = block_table[real_batch_lm[batch] * max_num_blocks_per_seq + (start_token_lm[batch] + i - lod_lm[b]) / block_size];
+            LM2GM(local_scr, dst + block_id * kv_head_num * block_size * head_dim + (start_token_lm[batch] + i - lod_lm[b]) % block_size * head_dim, kv_head_num * head_dim * sizeof(TO));
+
+        }
+    }
+}
+
+
+
+
+#define _XPU_DEF__ABSORB_MLA_CACHE_WRITE(TI, TO, TID)             \
+    template __global__ void absorb_mla_cache_write<TI, TO, TID>( \
+            const TI*,                                                                     \
+            TO*,                                                                           \
+            const TID*,                                                                    \
+            int32_t*,                                                                      \
+            int32_t*,                                                                      \
+            int32_t*,                                                                      \
+            int64_t,                                                                       \
+            int64_t,                                                                       \
+            int64_t,                                                                       \
+            int64_t,                                                                       \
+            int64_t,                                                                       \
+            int64_t);
+
+_XPU_DEF__ABSORB_MLA_CACHE_WRITE(float16, float16, int32_t);
+_XPU_DEF__ABSORB_MLA_CACHE_WRITE(bfloat16, bfloat16, int32_t);
+
+}  // namespace plugin
+}  // namespace xpu3

csrc/xpu/src/plugin/src/kernel/kunlun3cpp/fused_rotary_position_encoding.xpu

@@ -0,0 +1,238 @@
+// 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.
+
+#include "xpu/kernel/cluster.h"
+#include "xpu/kernel/cluster_partition.h"
+#include "xpu/kernel/cluster_primitive.h"
+#include "xpu/kernel/cluster_debug.h"
+
+namespace xpu3 {
+namespace plugin {
+
+static __device__ void apply_rotary_pos_emb_nonexo(
+        float* qk,
+        float* cos_sin,
+        int rot_dim,
+        int head_size,
+        int element_num) {
+    const int emb_dim = rot_dim / 2;
+    for(int i = 0; i < emb_dim; i++){
+        float cos_ind = cos_sin[i];
+        float sin_ind = cos_sin[i + emb_dim];
+        for (int j = 0; j < element_num; j++) {
+            // if(i == 0){
+            //     printf("input %f %f   %f %f  ",qk[j * head_size],qk[j * head_size + 1],cos_ind , sin_ind);
+            // }
+            float x_ind = qk[j * head_size + i * 2];
+            float y_ind = qk[j * head_size + i * 2 + 1];
+            qk[j * head_size + i * 2] = x_ind * cos_ind - y_ind * sin_ind;
+            qk[j * head_size + i * 2 + 1] =  y_ind * cos_ind + x_ind * sin_ind;
+            // if(i == 0){
+            //     printf("output %f %f \n ",qk[j * head_size],qk[j * head_size + 1]);
+            // }
+        }   
+    }
+
+    mfence();
+}
+
+
+static __device__ void apply_rotary_pos_emb_nexo(
+        float* qk,
+        float* cos_sin,
+        int rot_dim,
+        int head_size,
+        int element_num) {
+    float32x16_t vqx_0;
+    float32x16_t vqx_1;
+    float32x16_t vqy_0;
+    float32x16_t vqy_1;
+    float32x16_t vcos_0;
+    float32x16_t vcos_1;
+    float32x16_t vsin_0;
+    float32x16_t vsin_1;
+    float32x16_t v_x_cos;
+    float32x16_t v_y_cos;
+    const float var = -1.0f;
+    const int emb_dim = rot_dim / 2;
+    // TODO: fixme. It assumes that rot_dim is the multiplier of 32.
+    for (int i = 0; i < emb_dim; i += 32) {
+        vload2_lm(cos_sin + i, vcos_0, vcos_1);
+        vload2_lm(cos_sin + i + emb_dim, vsin_0, vsin_1);
+        for (int j = 0; j < element_num; j++) {
+            int offset = j * head_size;
+            vload2_lm(qk + offset + i, vqx_0, vqx_1);
+            vload2_lm(qk + offset + i + emb_dim, vqy_0, vqy_1);
+            // qx
+            v_x_cos = vvmul_float32x16(vqx_0, vcos_0);
+            vqx_0 = vvmul_float32x16(vqx_0, vsin_0);
+            // qy * cons
+            v_y_cos = vvmul_float32x16(vqy_0, vcos_0);
+            // -qy
+            vqy_0 = svmul_float32x16(var, vqy_0);
+            // -qy * sin
+            vqy_0 = vvmul_float32x16(vqy_0, vsin_0);
+            // add
+            vqy_0 = vvadd_float32x16(v_x_cos, vqy_0);  // out_left
+            vqx_0 = vvadd_float32x16(vqx_0, v_y_cos);  // out_right
+            // qy
+            v_x_cos = vvmul_float32x16(vqx_1, vcos_1);
+            vqx_1 = vvmul_float32x16(vqx_1, vsin_1);
+            // qy * cons
+            v_y_cos = vvmul_float32x16(vqy_1, vcos_1);
+            // -qy
+            vqy_1 = svmul_float32x16(var, vqy_1);
+            // -qy * sin
+            vqy_1 = vvmul_float32x16(vqy_1, vsin_1);
+            // add
+            vqy_1 = vvadd_float32x16(v_x_cos, vqy_1);  // out_left
+            vqx_1 = vvadd_float32x16(vqx_1, v_y_cos);  // out_right
+            vstore2_lm(qk + offset + i, vqy_0, vqy_1);
+            vstore2_lm(qk + offset + i + emb_dim, vqx_0, vqx_1);
+            mfence();
+        }
+    }
+}
+
+template <typename TX, typename TY>
+static __device__ void fast_cast(TX* x, TY* y, int64_t n) {
+    primitive_cast<TX, TY>(x, y, n);
+}
+
+template <>
+__device__ void fast_cast<bfloat16, float>(bfloat16* x, float* y, int64_t n) {
+    int start = (n - 1) / 32 * 32;
+    x = x + start;
+    y = y + start;
+    float32x16_t xl;
+    float32x16_t xh;
+    for (int i = start; i >= 0; i -= 32) {
+        vload2_lm(x, xl, xh);
+        vstore2_lm(y, xl, xh);
+        x -= 32;
+        y -= 32;
+    }
+    mfence_lm();
+}
+template <>
+__device__ void fast_cast<float, bfloat16>(float* x, bfloat16* y, int64_t n) {
+    float32x16_t xl;
+    float32x16_t xh;
+    for (int i = 0; i < n; i += 32) {
+        vload2_lm(x, xl, xh);
+        vstore2_lm(y, xl, xh);
+        x += 32;
+        y += 32;
+    }
+    mfence_lm();
+}
+
+template <typename T, typename TR>
+__global__ void rotary_embedding_neox(
+        const int* positions,  // [num_tokens]
+        T* query,                  // [num_tokens, num_heads, head_size]
+        T* key,                    // [num_tokens, num_kv_heads, head_size]
+        const TR* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens) {
+    int cid = core_id();
+    int ncores = core_num();
+    if (cid >= ncores) {
+        return;
+    }
+    int tid = cluster_id() * ncores + cid;
+    int nthreads = ncores * cluster_num();
+    constexpr int sm_buffer_size = 64 * 512;
+    int head_num = 2048 / sizeof(T) / head_size;  // max DMA is 2048 Bytes
+    int lm_buffer_size = head_num * head_size;
+    __simd__ float data_lm[lm_buffer_size];
+    __simd__ float cos_sin_lm[rot_dim];
+    __simd__ __shared__ int pos_sm[sm_buffer_size];
+
+    int32_t cluster_start_idx = 0;
+    int32_t cluster_end_idx = 0;
+    int32_t start_idx = 0;
+    int32_t end_idx = 0;
+    partition(cluster_id(), cluster_num(), num_tokens, 1, &cluster_start_idx, &cluster_end_idx);
+    if(cid == 0 && cluster_end_idx - cluster_start_idx > 0){
+        GM2SM(positions + cluster_start_idx, pos_sm, (cluster_end_idx - cluster_start_idx) * sizeof(int));
+    }
+    partition(cid, ncores, cluster_end_idx - cluster_start_idx, 1, &start_idx, &end_idx);
+    start_idx = start_idx + cluster_start_idx;
+    end_idx = end_idx + cluster_start_idx;
+    mfence();
+    sync_all();
+    // printf("id %d %d  %d %d\n ",cluster_id(), cluster_start_idx, cluster_end_idx, pos_sm[0]);
+    for (int64_t token_idx = start_idx; token_idx < end_idx; token_idx += 1) {
+        int64_t pos = pos_sm[token_idx - cluster_start_idx];
+        // printf("poss %lld %lld \n ",token_idx, pos);
+
+        _global_ptr_ const TR* cache_ptr = cos_sin_cache + pos * rot_dim;
+        int64_t offset_q = token_idx * num_heads * head_size;
+        int64_t offset_k = token_idx * num_kv_heads * head_size;
+        // get cos & sin
+        if (num_tokens > token_idx) {
+            GM2LM(cache_ptr, cos_sin_lm, rot_dim * sizeof(TR));
+            fast_cast<TR, float>((TR*)cos_sin_lm, cos_sin_lm, rot_dim);
+        }
+        // generate query
+        for (int64_t i = 0; i < num_heads; i += head_num) {
+            // int loop_size = min<int64_t>(n - nstart, lm_buffer_size);
+            int64_t element_num = min<int64_t>(num_heads - i, head_num);
+            GM2LM(query + offset_q, data_lm, element_num * head_size * sizeof(T));
+            fast_cast<T, float>((T*)data_lm, data_lm, element_num * head_size);
+            // calculation
+            apply_rotary_pos_emb_nonexo(data_lm, cos_sin_lm, rot_dim, head_size, element_num);
+            // write back to GM
+            fast_cast<float, T>(data_lm, (T*)data_lm, element_num * head_size);
+            LM2GM(data_lm, query + offset_q, element_num * head_size * sizeof(T));
+            offset_q += (head_size * element_num);
+        }
+        // generate key
+        for (int64_t j = 0; j < num_kv_heads; j += head_num) {
+            int64_t element_num = min<int64_t>(num_kv_heads - j, head_num);
+            GM2LM(key + offset_k, data_lm, element_num * head_size * sizeof(T));
+            fast_cast<T, float>((T*)data_lm, data_lm, element_num * head_size);
+            // calculation
+            apply_rotary_pos_emb_nonexo(data_lm, cos_sin_lm, rot_dim, head_size, element_num);
+            // write back to GM
+            fast_cast<float, T>(data_lm, (T*)data_lm, element_num * head_size);
+            LM2GM(data_lm, key + offset_k, element_num * head_size * sizeof(T));
+            offset_k += (head_size * element_num);
+        }
+    }
+}
+
+#define _XPU_DEF_(T, TR)                                        \
+    template __global__ void rotary_embedding_neox<T, TR>( \
+            const int* positions,                           \
+            T* query,                                           \
+            T* key,                                             \
+            const TR* cos_sin_cache,                            \
+            const int rot_dim,                                  \
+            const int num_heads,                                \
+            const int num_kv_heads,                             \
+            const int head_size,                                \
+            const int32_t num_tokens);
+
+_XPU_DEF_(float, float);
+_XPU_DEF_(float16, float16);
+_XPU_DEF_(bfloat16, float);
+_XPU_DEF_(bfloat16, bfloat16);
+
+}
+}

csrc/xpu/src/plugin/src/kernel/kunlun3cpp/get_position_ids.xpu

@@ -6,10 +6,10 @@ namespace xpu3 {
 namespace plugin {
 
 __global__ void get_position_ids(const int *seq_lens_encoder,
-                                  const int *seq_lens_decoder,
-                                  const int *seq_lens_this_time,
-                                  int *position_ids,
-                                  const int bs) {
+                                 const int *seq_lens_decoder,
+                                 const int *seq_lens_this_time,
+                                 int *position_ids,
+                                 const int bs) {
   int cid = core_id();
   int ncores = core_num();
   int clusterid = cluster_id();
@@ -26,31 +26,39 @@ __global__ void get_position_ids(const int *seq_lens_encoder,
   GM2LM_ASYNC(seq_lens_encoder, seq_lens_encoder_lm, bs * sizeof(int));
   GM2LM_ASYNC(seq_lens_decoder, seq_lens_decoder_lm, bs * sizeof(int));
   GM2LM(seq_lens_this_time, seq_lens_this_time_lm, bs * sizeof(int));
-
-  int offset = 0;
-  for(int ind = tid; ind < bs; ind += ncores * nclusters){
+  int encode_sum = 0;
+  int is_encode = 0;
+  for (int i = 0; i < bs; i++) {
+    if (seq_lens_encoder_lm[i] > 0) {
+      encode_sum = encode_sum + seq_lens_encoder_lm[i];
+    }
+  }
+  for (int ind = tid; ind < bs; ind += ncores * nclusters) {
+    int decode_offset = encode_sum;
+    int encode_offset = 0;
     for (int i = 0; i < ind; i++) {
-      offset += seq_lens_encoder_lm[i];
-      if (seq_lens_decoder_lm[i] > 0) {
-        offset += seq_lens_this_time_lm[i];
+      if (seq_lens_encoder_lm[i] > 0) {
+        encode_offset += seq_lens_encoder_lm[i];
+      } else {
+        decode_offset += seq_lens_this_time_lm[i];
       }
     }
     int encoder_len = seq_lens_encoder_lm[ind];
     int decoder_len = seq_lens_decoder_lm[ind];
     int seq_len_this_time = decoder_len > 0 ? seq_lens_this_time_lm[ind] : 0;
+    int offset = decoder_len > 0 ? decode_offset : encode_offset;
 
-    int begin_offset = offset;
-    for (int i = 0; i < encoder_len; i++) {
-      position_ids_lm[offset + i] = i;
-    }
-    offset += encoder_len;
-    if (decoder_len > 0) {
-      for (int i = 0; i < seq_len_this_time; i++) {
-        position_ids_lm[offset + i] = decoder_len + i;  // 使用 decoder 长度本身
+    for (int i = 0; i < encoder_len + seq_len_this_time; i += buf_len) {
+      int read_len = min(encoder_len + seq_len_this_time - i, buf_len);
+      for (int j = 0; j < buf_len; j++) {
+        int ind_id = i + j;
+        if (ind_id < encoder_len) {
+          position_ids_lm[j] = ind_id;
+        } else {
+          position_ids_lm[j] = decoder_len + ind_id;
+        }
       }
-    }
-    if (encoder_len + seq_len_this_time > 0) {
-      LM2GM(position_ids_lm + begin_offset, position_ids + begin_offset, (encoder_len + seq_len_this_time) * sizeof(int));
+      LM2GM(position_ids_lm, position_ids + offset + i, read_len * sizeof(int));
     }
   }
 }

csrc/xpu/src/plugin/src/kernel/kunlun3cpp/update_repeat_times.xpu

@@ -27,7 +27,7 @@ __global__ void update_repeat_times(const int64_t *pre_ids,
   int nclusters = cluster_num();
   int tid = clusterid * ncores + cid;
 
-  const int max_sm_len = 256 * 1024 / sizeof(int);
+  const int max_sm_len = 128 * 1024 / sizeof(int);
   __shared__ int repeated_times_sm[max_sm_len];
   int64_t pre_id_lm;
   int n_length = (length + max_sm_len - 1) / max_sm_len;

csrc/xpu/src/plugin/src/wrapper/absorb_mla_cache_write.cpp

@@ -0,0 +1,168 @@
+// Copyright (c) 2024 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.
+
+#include <algorithm>
+#include <numeric>
+#include "xpu/plugin.h"
+#include "xpu/refactor/impl_public/wrapper_check.h"
+
+
+namespace xpu3 {
+namespace plugin {
+
+template <typename TI, typename TO, typename TID>
+__attribute__((global)) void absorb_mla_cache_write(const TI* src,
+                                                    TO* dst,
+                                                    const TID* block_table,
+                                                    int32_t* kv_seq_lod,
+                                                    int32_t* start_token,
+                                                    int32_t* real_batch,
+                                                    int64_t batch_size,
+                                                    int64_t kv_head_num,
+                                                    int64_t head_dim,
+                                                    int64_t max_seq_num,
+                                                    int64_t block_size,
+                                                    int64_t max_num_blocks_per_seq);
+
+}  // namespace plugin
+}  // namespace xpu3
+
+namespace baidu {
+namespace xpu {
+namespace api {
+namespace plugin {
+
+template <typename TI, typename TO, typename TID>
+static int xpu3_wrapper(
+        Context* ctx,
+        const TI* x,
+        TO* y,
+        const TID* block_table,
+        const VectorParam<int32_t>& kv_seq_lod_xpu,
+        const VectorParam<int32_t>& start_tokens_xpu,
+        const VectorParam<int32_t>& real_batch_xpu,
+        int64_t batch_size,
+        int64_t kv_head_num,
+        int64_t head_dim,
+        int64_t max_seq_num,
+        int64_t block_size,
+        int64_t max_num_blocks_per_seq) {
+
+        auto func = &xpu3::plugin::absorb_mla_cache_write<TI, TO, TID>;
+        func<<<ctx->ncluster(), 64, ctx->xpu_stream>>>(
+                x,
+                y,
+                block_table,
+                kv_seq_lod_xpu.xpu,
+                start_tokens_xpu.xpu,
+                real_batch_xpu.xpu,
+                batch_size,
+                kv_head_num,
+                head_dim,
+                max_seq_num,
+                block_size,
+                max_num_blocks_per_seq);
+
+  return api::SUCCESS;
+
+}
+
+template <typename TI, typename TO, typename TID>
+int absorb_mla_cache_write_xpu(Context *ctx,
+        const TI* x,
+        TO* y,
+        const TID* block_table,
+        const VectorParam<int32_t>& kv_seq_lod,
+        const VectorParam<int32_t>& start_tokens,
+        const VectorParam<int32_t>& real_batch,
+        int64_t batch_size,
+        int64_t kv_head_num,
+        int64_t head_dim,
+        int64_t max_seq_num,
+        int64_t block_size,
+        int64_t max_num_blocks_per_seq) {
+    WRAPPER_CHECK_CTX(ctx);
+    WRAPPER_DUMP_FUNCTION_T3(ctx, "absorb_mla_cache_write", TI, TO, TID);
+    WRAPPER_DUMP_PARAM6(ctx, x, y, block_table, kv_seq_lod, start_tokens, real_batch);
+    WRAPPER_DUMP_PARAM6(ctx, batch_size, kv_head_num, head_dim, max_seq_num, block_size, max_num_blocks_per_seq);
+    WRAPPER_DUMP(ctx);
+    // check vector param size
+    WRAPPER_ASSERT_EQ(ctx, batch_size + 1, kv_seq_lod.len);
+    WRAPPER_ASSERT_EQ(ctx, batch_size, start_tokens.len);
+    WRAPPER_ASSERT_EQ(ctx, batch_size, real_batch.len);
+    WRAPPER_ASSERT_EQ(ctx, kv_head_num, 1);
+    WRAPPER_ASSERT_LE(ctx, batch_size, 32);
+    // check input ptr shape
+    int64_t kv_seqlen_sum = kv_seq_lod.cpu[batch_size];
+    int64_t hidden_dim = kv_head_num * head_dim;
+    // check kv_cache and block_table ptr shape
+    int64_t table_count = max_seq_num * max_num_blocks_per_seq;        // block number in a block table
+    int64_t kv_cache_size = table_count * kv_head_num * head_dim * block_size; 
+    // x 'shape is supposed to be [kv_seqlen_sum, hidden_dim]
+    WRAPPER_CHECK_PTR(ctx, TI, kv_seqlen_sum * hidden_dim, x);
+    // y is kv cache ptr, which is supposed to be [table_count, kv_head_num, block_size, head_dim]
+    WRAPPER_CHECK_PTR(ctx, TO, kv_cache_size, y);
+
+    // int64_t single_block_occupy = kv_head_num * head_dim * block_size; // single_block_occupy represent how many bytes
+                                                                       // occupied in a block
+                                                                       // which is different from block size
+                                                                       // block_size means the seqlen a block can hold
+    // kv_cache 'shape is supposed to be [block_num, hidden_dim, block_size]
+    // FIXME: in real cases, kv cache may not be full
+    // WRAPPER_CHECK_PTR(ctx, TO, table_count * single_block_occupy, y);
+    // block_table 'shape is supposed to be [max_seq_num, max_num_blocks_per_seq]
+    WRAPPER_CHECK_PTR(ctx, TID, table_count, block_table);
+
+  if (ctx->dev().type() == api::kXPU3) {
+    return xpu3_wrapper<TI, TO, TID>(ctx,
+                                        x,
+                                        y,
+                                        block_table,
+                                        kv_seq_lod,
+                                        start_tokens,
+                                        real_batch,
+                                        batch_size,
+                                        kv_head_num,
+                                        head_dim,
+                                        max_seq_num,
+                                        block_size,
+                                        max_num_blocks_per_seq);
+  }
+  WRAPPER_UNIMPLEMENTED(ctx);
+}
+
+#define INSTANTIATION_ABSORB_MLA_CACHE_WRITE(TI, TO, TID)                                                                \
+    template int absorb_mla_cache_write_xpu<TI, TO, TID>(api::Context*,                                                     \
+                                                const TI*,                                                             \
+                                                TO*,                                                                   \
+                                                const TID*,                                                            \
+                                                const api::VectorParam<int32_t>&,                                      \
+                                                const api::VectorParam<int32_t>&,                                      \
+                                                const api::VectorParam<int32_t>&,                                      \
+                                                int64_t,                                                               \
+                                                int64_t,                                                               \
+                                                int64_t,                                                               \
+                                                int64_t,                                                               \
+                                                int64_t,                                                               \
+                                                int64_t);
+
+INSTANTIATION_ABSORB_MLA_CACHE_WRITE(float16, float16, int32_t);
+INSTANTIATION_ABSORB_MLA_CACHE_WRITE(bfloat16, bfloat16, int32_t);
+
+
+
+}  // namespace plugin
+}  // namespace api
+}  // namespace xpu
+}  // namespace baidu

csrc/xpu/src/plugin/src/wrapper/fused_rotary_position_encoding.cpp

@@ -0,0 +1,166 @@
+// Copyright (c) 2024 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.
+
+#include <algorithm>
+#include <numeric>
+
+#include "xpu/plugin.h"
+#include "xpu/refactor/impl_public/wrapper_check.h"
+
+namespace xpu3 {
+namespace plugin {
+template <typename T, typename TR>
+__attribute__((global)) void rotary_embedding_neox(
+        const int* positions,  // [num_tokens]
+        T* query,                  // [num_tokens, num_heads, head_size]
+        T* key,                    // [num_tokens, num_kv_heads, head_size]
+        const TR* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens);
+}  // namespace plugin
+}  // namespace xpu3
+
+namespace baidu {
+namespace xpu {
+namespace api {
+namespace plugin {
+
+// template <typename T>
+// static int cpu_wrapper() {
+// }
+
+template <typename T, typename TR>
+static int xpu3_wrapper(Context* ctx,
+        const int* positions,  // [num_tokens]
+        T* query,                  // [num_tokens, num_heads, head_size]
+        T* key,                    // [num_tokens, num_kv_heads, head_size]
+        const TR* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens) {
+  using XPU_TID = typename XPUIndexType<T>::type;
+  using XPU_TRID = typename XPUIndexType<TR>::type;
+  using XPU_INT64 = typename XPUIndexType<int64_t>::type;
+  auto rotary_embedding_neox = xpu3::plugin::rotary_embedding_neox<XPU_TID, XPU_TRID>;
+  rotary_embedding_neox<<<ctx->ncluster(), 64, ctx->xpu_stream>>>(
+      positions,
+      reinterpret_cast<XPU_TID *>(query),
+      reinterpret_cast<XPU_TID *>(key),
+      reinterpret_cast<const XPU_TRID *>(cos_sin_cache),
+      rot_dim,
+      num_heads,
+      num_kv_heads,
+      head_size,
+      num_tokens);
+  return api::SUCCESS;
+}
+
+// template <typename T>
+// int set_stop_value_multi_ends(Context* ctx,
+//                               bool* stop_flags,
+//                               T* topk_ids,
+//                               T* next_tokens,
+//                               const T* end_ids,
+//                               const int* seq_lens,
+//                               const int bs,
+//                               const int end_length,
+//                               const bool beam_search) 
+                              
+                              
+template <typename T, typename TR>
+int rotary_embedding_neox(
+        Context* ctx,
+        const int* positions,  // [num_tokens]
+        T* query,                  // [num_tokens, num_heads, head_size]
+        T* key,                    // [num_tokens, num_kv_heads, head_size]
+        const TR* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens)
+{
+  WRAPPER_CHECK_CTX(ctx);
+  WRAPPER_DUMP_FUNCTION_T2(ctx, "rotary_embedding_neox", T, TR);
+  WRAPPER_DUMP_PARAM5(
+      ctx, positions, query, key, cos_sin_cache, rot_dim);
+  WRAPPER_DUMP_PARAM4(ctx, num_heads, num_kv_heads, head_size,num_tokens);
+  WRAPPER_DUMP(ctx);
+
+  WRAPPER_CHECK_PTR(ctx, int, num_tokens, positions);
+  WRAPPER_CHECK_PTR(ctx, T, num_tokens * num_heads * head_size, query);
+  WRAPPER_CHECK_PTR(ctx, T, num_tokens * num_kv_heads * head_size, key);
+  WRAPPER_CHECK_PTR(ctx, TR, num_tokens * rot_dim, cos_sin_cache);
+
+  if (ctx->dev().type() == api::kXPU2 || ctx->dev().type() == api::kXPU3) {
+    return xpu3_wrapper<T, TR>(ctx,
+                           positions,
+                           query,
+                           key,
+                           cos_sin_cache,
+                           rot_dim,
+                           num_heads,
+                           num_kv_heads,
+                           head_size,
+                           num_tokens);
+  }
+  WRAPPER_UNIMPLEMENTED(ctx);
+}
+
+
+template int rotary_embedding_neox<float, float>(
+        Context* ctx,
+        const int* positions,  // [num_tokens]
+        float* query,                  // [num_tokens, num_heads, head_size]
+        float* key,                    // [num_tokens, num_kv_heads, head_size]
+        const float* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens);
+
+template int rotary_embedding_neox<float16, float16>(
+        Context* ctx,
+        const int* positions,  // [num_tokens]
+        float16* query,                  // [num_tokens, num_heads, head_size]
+        float16* key,                    // [num_tokens, num_kv_heads, head_size]
+        const float16* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens);
+
+template int rotary_embedding_neox<bfloat16, bfloat16>(
+        Context* ctx,
+        const int* positions,  // [num_tokens]
+        bfloat16* query,                  // [num_tokens, num_heads, head_size]
+        bfloat16* key,                    // [num_tokens, num_kv_heads, head_size]
+        const bfloat16* cos_sin_cache,   // [max_position, 2, rot_dim // 2]
+        const int rot_dim,
+        const int num_heads,
+        const int num_kv_heads,
+        const int head_size,
+        const int32_t num_tokens);
+
+}  // namespace plugin
+}  // namespace api
+}  // namespace xpu
+}  // namespace baidu

csrc/xpu/src/plugin/src/wrapper/nn_set_value_by_flags_and_idx.cpp

@@ -163,7 +163,7 @@ int set_value_by_flags_and_idx(Context* ctx,
   WRAPPER_CHECK_SHAPE(ctx, &seq_lens_encoder_len, {bs});
   WRAPPER_CHECK_SHAPE(ctx, &seq_lens_decoder_len, {bs});
   WRAPPER_CHECK_SHAPE(ctx, &step_idx_len, {bs});
-  WRAPPER_CHECK_PTR(ctx, int64_t, stop_flags_len, stop_flags);
+  WRAPPER_CHECK_PTR(ctx, bool, stop_flags_len, stop_flags);
   WRAPPER_CHECK_PTR(ctx, int64_t, pre_ids_all_len, pre_ids_all);
   WRAPPER_CHECK_PTR(ctx, int64_t, input_ids_len, input_ids);
   WRAPPER_CHECK_PTR(ctx, int, seq_lens_encoder_len, seq_lens_encoder);

csrc/xpu/src/plugin/src/wrapper/update_inputs.cpp

@@ -1,192 +0,0 @@
-// Copyright (c) 2024 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.
-
-#include <algorithm>
-#include <numeric>
-#include "xpu/plugin.h"
-#include "xpu/refactor/impl_public/wrapper_check.h"
-
-namespace xpu2 {
-namespace plugin {
-
-__attribute__((global)) void update_inputs(bool *not_need_stop,
-                                           int *seq_lens_this_time,
-                                           int *seq_lens_encoder,
-                                           int *seq_lens_decoder,
-                                           int64_t *input_ids,
-                                           const int64_t *stop_nums,
-                                           const bool *stop_flags,
-                                           const bool *is_block_step,
-                                           const int64_t *next_tokens,
-                                           const int bsz,
-                                           const int max_bsz,
-                                           const int input_ids_stride);
-
-}  // namespace plugin
-}  // namespace xpu2
-
-namespace xpu3 {
-namespace plugin {
-
-__attribute__((global)) void update_inputs(bool *not_need_stop,
-                                           int *seq_lens_this_time,
-                                           int *seq_lens_encoder,
-                                           int *seq_lens_decoder,
-                                           int64_t *input_ids,
-                                           const int64_t *stop_nums,
-                                           const bool *stop_flags,
-                                           const bool *is_block_step,
-                                           const int64_t *next_tokens,
-                                           const int bsz,
-                                           const int max_bsz,
-                                           const int input_ids_stride);
-
-}  // namespace plugin
-}  // namespace xpu3
-
-namespace baidu {
-namespace xpu {
-namespace api {
-namespace plugin {
-
-static int cpu_wrapper(Context *ctx,
-                       bool *not_need_stop,
-                       int *seq_lens_this_time,
-                       int *seq_lens_encoder,
-                       int *seq_lens_decoder,
-                       int64_t *input_ids,
-                       const int64_t *stop_nums,
-                       const bool *stop_flags,
-                       const bool *is_block_step,
-                       const int64_t *next_tokens,
-                       const int bsz,
-                       const int max_bsz,
-                       const int input_ids_stride) {
-  std::vector<int64_t> stop_flag_now_int(max_bsz, 1);
-  for (int i = 0; i < bsz; i++) {
-    bool stop_flags_now = stop_flags[i];
-    stop_flag_now_int[i] = is_block_step[i] ? 0 : stop_flags_now;
-    const int seq_len_encoder = seq_lens_encoder[i];
-    const int seq_len_decoder = seq_lens_decoder[i];
-
-    seq_lens_decoder[i] =
-        stop_flags[i]
-            ? 0
-            : (seq_len_decoder == 0 ? seq_len_encoder : seq_len_decoder + 1);
-
-    seq_lens_this_time[i] = stop_flags[i] ? 0 : 1;
-    seq_lens_encoder[i] = 0;
-    int64_t *input_ids_now = input_ids + i * input_ids_stride;
-    input_ids_now[0] = next_tokens[i];
-  }
-  int64_t stop_sum = 0;
-  for (size_t i = 0; i < stop_flag_now_int.size(); i++) {
-    stop_sum += stop_flag_now_int[i];
-  }
-  not_need_stop[0] = stop_sum < stop_nums[0];
-  return api::SUCCESS;
-}
-
-static int xpu2or3_wrapper(Context *ctx,
-                        bool *not_need_stop,
-                        int *seq_lens_this_time,
-                        int *seq_lens_encoder,
-                        int *seq_lens_decoder,
-                        int64_t *input_ids,
-                        const int64_t *stop_nums,
-                        const bool *stop_flags,
-                        const bool *is_block_step,
-                        const int64_t *next_tokens,
-                        const int bsz,
-                        const int max_bsz,
-                        const int input_ids_stride) {
-  using XPU_INT64 = typename XPUIndexType<int64_t>::type;
-  auto update_inputs = xpu2::plugin::update_inputs;
-  update_inputs<<<ctx->ncluster(), 64, ctx->xpu_stream>>>(
-      not_need_stop,
-      seq_lens_this_time,
-      seq_lens_encoder,
-      seq_lens_decoder,
-      reinterpret_cast<XPU_INT64 *>(input_ids),
-      reinterpret_cast<const XPU_INT64 *>(stop_nums),
-      stop_flags,
-      is_block_step,
-      reinterpret_cast<const XPU_INT64 *>(next_tokens),
-      bsz,
-      max_bsz,
-      input_ids_stride);
-  return api::SUCCESS;
-}
-
-int update_inputs(Context *ctx,
-                  bool *not_need_stop,
-                  int *seq_lens_this_time,
-                  int *seq_lens_encoder,
-                  int *seq_lens_decoder,
-                  int64_t *input_ids,
-                  const int64_t *stop_nums,
-                  const bool *stop_flags,
-                  const bool *is_block_step,
-                  const int64_t *next_tokens,
-                  const int bsz,
-                  const int max_bsz,
-                  const int input_ids_stride) {
-  WRAPPER_CHECK_CTX(ctx);
-  WRAPPER_DUMP_FUNCTION_T1(ctx, "update_inputs", int);
-  WRAPPER_DUMP_PARAM5(ctx,
-                      not_need_stop,
-                      seq_lens_this_time,
-                      seq_lens_encoder,
-                      seq_lens_decoder,
-                      input_ids);
-  WRAPPER_DUMP_PARAM4(ctx, stop_nums, stop_flags, is_block_step, next_tokens);
-  WRAPPER_DUMP_PARAM3(ctx, bsz, max_bsz, input_ids_stride);
-  WRAPPER_DUMP(ctx);
-  if (ctx->dev().type() == api::kCPU) {
-    return cpu_wrapper(ctx,
-                       not_need_stop,
-                       seq_lens_this_time,
-                       seq_lens_encoder,
-                       seq_lens_decoder,
-                       input_ids,
-                       stop_nums,
-                       stop_flags,
-                       is_block_step,
-                       next_tokens,
-                       bsz,
-                       max_bsz,
-                       input_ids_stride);
-  }
-  if (ctx->dev().type() == api::kXPU2 || ctx->dev().type() == api::kXPU3) {
-    return xpu2or3_wrapper(ctx,
-                        not_need_stop,
-                        seq_lens_this_time,
-                        seq_lens_encoder,
-                        seq_lens_decoder,
-                        input_ids,
-                        stop_nums,
-                        stop_flags,
-                        is_block_step,
-                        next_tokens,
-                        bsz,
-                        max_bsz,
-                        input_ids_stride);
-  }
-  WRAPPER_UNIMPLEMENTED(ctx);
-}
-
-}  // namespace plugin
-}  // namespace api
-}  // namespace xpu
-}  // namespace baidu

csrc/xpu/src/plugin/src/wrapper/update_inputs_v2.cpp

@@ -191,17 +191,35 @@ int update_inputs_v2(Context* ctx,
                       stop_flags,
                       seq_lens_this_time,
                       seq_lens_encoder);
+  
   WRAPPER_DUMP_PARAM5(ctx,
                       seq_lens_decoder,
                       next_tokens,
                       kwargs_next_tokens,
                       input_ids,
                       end_ids);
+
   WRAPPER_DUMP_PARAM3(ctx, stop_nums, is_block_step, max_dec_len);
   WRAPPER_DUMP_PARAM4(ctx, now_bsz, max_bsz, input_ids_stride, end_length);
+
+  
   WRAPPER_DUMP(ctx);
-  WRAPPER_ASSERT_LE(ctx, max_bsz, 1024);
-  WRAPPER_ASSERT_LE(ctx, now_bsz, max_bsz);
+  WRAPPER_CHECK_PTR(ctx, bool, 1, not_need_stop);
+  WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, step_idx);
+  WRAPPER_CHECK_PTR(ctx, bool, max_bsz, stop_flags);
+  WRAPPER_CHECK_PTR(ctx, int, max_bsz, seq_lens_this_time);
+  WRAPPER_CHECK_PTR(ctx, int, max_bsz, seq_lens_encoder);
+  WRAPPER_CHECK_PTR(ctx, int, max_bsz, seq_lens_decoder);
+  WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, next_tokens);
+  WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, kwargs_next_tokens);
+  WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, input_ids);
+  WRAPPER_CHECK_PTR(ctx, int64_t, end_length, end_ids);
+  WRAPPER_CHECK_PTR(ctx, int64_t, 1, stop_nums);
+  WRAPPER_CHECK_PTR(ctx, bool, max_bsz, is_block_step);
+  WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, max_dec_len);
+
+
+  // std::cout << "wht --- ctx debug level is " << ctx->debug_level() << std::endl;
   // TODO(mayang02): check ptrs
   if (ctx->dev().type() == api::kCPU) {
     return cpu_wrapper(ctx,

csrc/xpu/src/save_with_output.cc

@@ -0,0 +1,166 @@
+// Copyright (c) 2023 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.
+
+#include <iostream>
+#include <fstream>
+#include <memory>
+#include <stdexcept>
+#include <string>
+
+#include "stdlib.h"
+#include <stdio.h>
+#include <dlfcn.h>  // dladdr
+#include <sys/time.h>
+#include <sys/stat.h>
+#include "paddle/extension.h"
+
+constexpr char kSEP = '/';
+
+std::string DirName(const std::string &filepath) {
+  auto pos = filepath.rfind(kSEP);
+  if (pos == std::string::npos) {
+    return "";
+  }
+  return filepath.substr(0, pos);
+}
+
+bool FileExists(const std::string &filepath) {
+  struct stat buffer;
+  return (stat(filepath.c_str(), &buffer) == 0);
+}
+
+void MkDir(const char *path) {
+  std::string path_error(path);
+  path_error += " mkdir failed!";
+  if (mkdir(path, 0755)) {
+    if (errno != EEXIST) {
+      throw std::runtime_error(path_error);
+    }
+  }
+}
+
+void MkDirRecursively(const char *fullpath) {
+  if (*fullpath == '\0') return;  // empty string
+  if (FileExists(fullpath)) return;
+  MkDirRecursively(DirName(fullpath).c_str());
+  MkDir(fullpath);
+}
+
+
+template<typename data_t>
+void saveToFile(std::ostream & os, const void* x_data, std::vector<int64_t> shape, int64_t x_numel, const char type_id) {
+  // 1.type
+  os.write(reinterpret_cast<const char *>(&type_id),sizeof(type_id));
+  // 2.data
+  uint64_t size = x_numel * sizeof(data_t);
+  os.write(static_cast<const char*>(x_data),static_cast<std::streamsize>(size));
+
+}
+
+template<typename data_t>
+void save_with_output_kernel(const paddle::Tensor& x,
+                             const paddle::Tensor& batch_idx,
+                             const paddle::Tensor& step_idx,
+                             std::string file_path,
+                             int64_t rank_id,
+                             char type_id) {
+  std::vector<int64_t> x_shape = x.shape();
+
+  if(rank_id >= 0) {
+      file_path += "_rank_" + std::to_string(rank_id);
+  }
+
+  int batch_idx_data = -1, step_idx_data = -1;
+
+  if(batch_idx.is_gpu()) {
+    paddle::Tensor batch_idx_cpu = batch_idx.copy_to<int32_t>(paddle::CPUPlace());
+    batch_idx_data = batch_idx_cpu.data<int32_t>()[0];
+  } else {
+    batch_idx_data = batch_idx.data<int32_t>()[0];
+  }
+  if(step_idx.is_gpu()) {
+    paddle::Tensor step_idx_cpu = step_idx.copy_to<int64_t>(paddle::CPUPlace());
+    step_idx_data = step_idx_cpu.data<int64_t>()[0];
+  } else {
+    step_idx_data = step_idx.data<int64_t>()[0];
+  }
+  auto x_data = x.data<data_t>();
+
+  if(batch_idx_data >= 0) {
+    file_path += "_batch_" + std::to_string(batch_idx_data);
+  }
+  if(step_idx_data >= 0) {
+    file_path += "_step_" + std::to_string(step_idx_data);
+  }
+  MkDirRecursively(DirName(file_path).c_str());
+  std::ofstream fout(file_path, std::ios::binary);
+  fout.write("0",1);
+  saveToFile<data_t>(fout, x_data, x_shape, x.numel(),type_id);
+  fout.seekp(std::ios::beg);
+  fout.write("1",1);
+  fout.close();
+
+}
+
+void print_shape(const paddle::Tensor& tmp, char *tmp_str){
+    std::vector<int64_t> shape = tmp.shape();
+    printf("%s's shape: \n", tmp_str);
+    for(int i=0; i < shape.size(); i++) {
+        printf("%d ", (int)shape[i]);
+    }
+    printf("\n");
+}
+
+std::vector<paddle::Tensor> SaveWithOutputForward(const paddle::Tensor& x,
+                                                  const paddle::Tensor& batch_idx,
+                                                  const paddle::Tensor& step_idx,
+                                                  std::string file_path,
+                                                  int64_t rank_id) {
+    auto out = x.copy_to(paddle::CPUPlace(), false);
+    switch(x.type()) {
+      case paddle::DataType::FLOAT32:
+         save_with_output_kernel<float>(out, batch_idx, step_idx, file_path, rank_id, '0');
+         break;
+      case paddle::DataType::INT64:
+        save_with_output_kernel<int64_t>(out, batch_idx, step_idx, file_path, rank_id,'1');
+         break;
+      case paddle::DataType::INT32:
+        save_with_output_kernel<int32_t>(out, batch_idx, step_idx, file_path, rank_id, '2');
+         break;
+      default:
+        PD_THROW("function SaveWithOutputForward is not implemented for data type");
+    }
+   return {out};
+}
+
+std::vector<std::vector<int64_t>> SaveWithOutputInferShape(const std::vector<int64_t>& x_shape,
+                                                           const std::vector<int64_t>& batch_idx_shape,
+                                                           const std::vector<int64_t>& step_idx_shape) {
+    return {x_shape};
+}
+
+std::vector<paddle::DataType> SaveWithOutputInferDtype(const paddle::DataType& x_dtype,
+                                                       const paddle::DataType& batch_idx_dtype,
+                                                       const paddle::DataType& step_idx_dtype) {
+    return {x_dtype};
+}
+
+PD_BUILD_OP(save_with_output)
+    .Inputs({"x", "batch_idx", "step_idx"})
+    .Attrs({"file_path: std::string",
+            "rank_id: int64_t"})
+    .Outputs({"out"})
+    .SetKernelFn(PD_KERNEL(SaveWithOutputForward))
+    .SetInferShapeFn(PD_INFER_SHAPE(SaveWithOutputInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(SaveWithOutputInferDtype));

csrc/xpu/src/save_with_output_msg.cc

@@ -1,11 +1,11 @@
-// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
-//
+// Copyright (c) 2024 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.
@@ -18,41 +18,54 @@
 #include <sys/msg.h>
 #include <sys/types.h>
 #include "paddle/extension.h"
+#include "xpu/plugin.h"
 
-#define MAX_BSZ 256
+#define MAX_BSZ 512
 
 struct msgdata {
-  long mtype;              // NOLINT
-  int mtext[MAX_BSZ + 2];  // stop_flag, bsz, tokens
+    long mtype;
+    int mtext[MAX_BSZ + 2];   // stop_flag, bsz, tokens
 };
 
 void SaveOutMmsg(const paddle::Tensor& x,
-                 const paddle::Tensor& not_need_stop,
+                 const paddle::Tensor& not_need_stop, // cpu
+                 const paddle::Tensor& msg_queue_id,      // cpu
                  int64_t rank_id) {
-  if (rank_id > 0) return;
-  auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
-  int64_t* x_data = x_cpu.data<int64_t>();
-  static struct msgdata msg_sed;
-  static key_t key = ftok("./", 1);
-  static int msgid = msgget(key, IPC_CREAT | 0666);
 
-  msg_sed.mtype = 1;
-  bool not_need_stop_data = not_need_stop.data<bool>()[0];
-  msg_sed.mtext[0] = not_need_stop_data ? 1 : -1;
-  int bsz = x.shape()[0];
-  msg_sed.mtext[1] = bsz;
-  for (int i = 2; i < bsz + 2; i++) {
-    msg_sed.mtext[i] = static_cast<int>(x_data[i - 2]);
-  }
-  if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, IPC_NOWAIT)) == -1) {
-    printf("full msg buffer\n");
-  }
-  return;
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] SaveOutMmsg");
+    if (rank_id > 0) return;
+    auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
+    int64_t *x_data = x_cpu.data<int64_t>();
+    auto not_need_stop_data = not_need_stop.data<bool>()[0];
+
+    std::cout << "not_need_stop_data: " << not_need_stop_data << std::endl;
+
+    static struct msgdata msg_sed;
+    int msg_queue_id_val = msg_queue_id.data<int>()[0];
+
+    std::cout << "msg_queue_id_val: " << msg_queue_id_val << std::endl;
+
+    static key_t key = ftok("./", msg_queue_id_val);
+    static int msgid = msgget(key, IPC_CREAT | 0666);
+
+    msg_sed.mtype = 1;
+    msg_sed.mtext[0] = not_need_stop_data ? 1 : -1;
+    int bsz = x.shape()[0];
+    msg_sed.mtext[1] = bsz;
+    for (int i = 2; i < bsz + 2; i++) {
+        msg_sed.mtext[i] = (int)x_data[i - 2];
+    }
+    if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, 0)) == -1) {
+    //   printf("full msg buffer\n");
+    }
+
+  baidu::xpu::api::plugin::print_times("[TIME END] SaveOutMmsg");
+    return;
 }
 
 PD_BUILD_OP(save_output)
-    .Inputs({"x", "not_need_stop"})
+    .Inputs({"x", "not_need_stop", "msg_queue_id"})
     .Attrs({"rank_id: int64_t"})
     .Outputs({"x_out"})
     .SetInplaceMap({{"x", "x_out"}})
-    .SetKernelFn(PD_KERNEL(SaveOutMmsg));
+    .SetKernelFn(PD_KERNEL(SaveOutMmsg));

csrc/xpu/src/set_preids_token_penalty_multi_scores.cc

@@ -56,7 +56,7 @@ void SetPreidsTokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
       seq_lens_decoder.data<int>(),
       step_idx.data<int64_t>(),
       bs,
-      length,
+      length_id,
       length_input_ids);
   PD_CHECK(r == 0, "xpu::plugin::set_value_by_flags_and_idx failed.");
 

csrc/xpu/src/setup.py

@@ -29,9 +29,15 @@ PADDLE_PATH = os.path.dirname(paddle.__file__)
 PADDLE_INCLUDE_PATH = os.path.join(PADDLE_PATH, "include")
 PADDLE_LIB_PATH = os.path.join(PADDLE_PATH, "libs")
 
-XPU_LIB = os.getenv("XPU_LIB")
+BKCL_PATH = os.getenv("BKCL_PATH")
+if BKCL_PATH is None:
+    BKCL_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xpu")
+    BKCL_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libbkcl.so")
+else:
+    BKCL_INC_PATH = os.path.join(BKCL_PATH, "include")
+    BKCL_LIB_PATH = os.path.join(BKCL_PATH, "so", "libbkcl.so")
 
-XFT_PATH = f"{XPU_LIB}/xft_output"
+XFT_PATH = os.getenv("XFT_PATH")
 if XFT_PATH is None:
     XFT_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xft")
     XFT_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxft.so")
@@ -39,18 +45,29 @@ else:
     XFT_INC_PATH = os.path.join(XFT_PATH, "include")
     XFT_LIB_PATH = os.path.join(XFT_PATH, "so", "libxft.so")
 
+XRE_PATH = os.getenv("XRE_PATH")
+if XRE_PATH is None:
+    XRE_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xre")
+    XRE_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpucuda.so")
+else:
+    XRE_INC_PATH = os.path.join(XRE_PATH, "include")
+    XRE_LIB_PATH = os.path.join(XRE_PATH, "so", "libxpucuda.so")
 
-XRE_PATH = f"{XPU_LIB}/xre"
-XRE_INC_PATH = os.path.join(XRE_PATH, "include")
-XRE_LIB_PATH = os.path.join(XRE_PATH, "so", "libcudart.so")
+XFA_PATH = os.getenv("XFA_PATH")
+if XFA_PATH is None:
+    XFA_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xhpc", "xfa")
+    XFA_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpu_flash_attention.so")
+else:
+    XFA_INC_PATH = os.path.join(XFA_PATH, "include")
+    XFA_LIB_PATH = os.path.join(XFA_PATH, "so", "libxpu_flash_attention.so")
 
-XFA_PATH = f"{XPU_LIB}/xhpc/xfa"
-XFA_INC_PATH = os.path.join(XFA_PATH, "include")
-XFA_LIB_PATH = os.path.join(XFA_PATH, "so", "libxpu_flash_attention.so")
-
-XBLAS_PATH = f"{XPU_LIB}/xhpc/xblas"
-XBLAS_INC_PATH = os.path.join(XBLAS_PATH, "include")
-XBLAS_LIB_PATH = os.path.join(XBLAS_PATH, "so", "libxpu_blas.so")
+XBLAS_PATH = os.getenv("XBLAS_PATH")
+if XBLAS_PATH is None:
+    XBLAS_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xhpc", "xblas")
+    XBLAS_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpu_blas.so")
+else:
+    XBLAS_INC_PATH = os.path.join(XBLAS_PATH, "include")
+    XBLAS_LIB_PATH = os.path.join(XBLAS_PATH, "so", "libxpu_blas.so")
 
 setup(
     name="paddlenlp_ops",
@@ -63,24 +80,46 @@ setup(
                 "./set_value_by_flags_and_idx_v2.cc",
                 "./get_token_penalty_multi_scores_v2.cc",
                 "./get_padding_offset_v2.cc",
-                # "./update_inputs.cc",
                 "./rebuild_padding_v2.cc",
-                "../../gpu/save_with_output.cc",
-                "../../gpu/save_with_output_msg.cc",
-                "../../gpu/get_output.cc",
+                "./save_with_output.cc",
+                "./save_with_output_msg.cc",
+                "./get_output.cc",
                 "./moe_dispatch.cc",
                 "./moe_ffn.cc",
                 "./moe_reduce.cc",
+                "./moe_fused.cc",
                 "./mla_block_multihead_attention_xpu.cc",
-                "./weight_only_linear.cc",
+                "./absorb_mla_block_mha_encoder_xpu.cc",
+                "./absorb_mla_block_mha_decoder_xpu.cc",
                 "./get_position_ids.cc",
                 "./get_position_ids_v2.cc",
                 "./adjust_batch.cc",
                 "./gather_next_token.cc",
-                "step.cc",
+                "./step.cc",
+                "./fused_rotary_position_encoding.cc",
+                "./mla_cache_kernel_xpu.cc",
+                "./batch_matmul_xpu.cc",
+                "./get_infer_param.cc",
+                "./block_attn.cc",
+                # "./dynamic_set_split_value.cc"
+            ],
+            include_dirs=[
+                ".",
+                "./plugin/include",
+                XRE_INC_PATH,
+                XFT_INC_PATH,
+                BKCL_LIB_PATH,
+                XFA_INC_PATH,
+                XBLAS_INC_PATH,
+            ],
+            extra_objects=[
+                "./plugin/build/libxpuplugin.a",
+                XRE_LIB_PATH,
+                XFT_LIB_PATH,
+                BKCL_LIB_PATH,
+                XFA_LIB_PATH,
+                XBLAS_LIB_PATH,
             ],
-            include_dirs=[".", "./plugin/include", XRE_INC_PATH, XFT_INC_PATH, XFA_INC_PATH, XBLAS_INC_PATH],
-            extra_objects=["./plugin/build/libxpuplugin.a", XRE_LIB_PATH, XFT_LIB_PATH, XFA_LIB_PATH, XBLAS_LIB_PATH],
             extra_compile_args={"cxx": ["-D_GLIBCXX_USE_CXX11_ABI=1", "-DPADDLE_WITH_XPU"]},
         )
     ],

csrc/xpu/src/step.cc

@@ -41,7 +41,10 @@ void StepPaddle(const paddle::Tensor& stop_flags,
                 const paddle::Tensor& first_token_ids,
                 const int block_size,
                 const int encoder_decoder_block_num) {
-  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+
+   baidu::xpu::api::plugin::print_times("[TIME BEGIN] StepPaddle");
+
+      phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
 
@@ -81,7 +84,7 @@ void StepPaddle(const paddle::Tensor& stop_flags,
   auto recover_lens_cpu = recover_lens.copy_to(paddle::CPUPlace(), false);
   int recover_lens_cpu_data = recover_lens_cpu.data<int>()[0];
   if (recover_lens_cpu_data > 0) {
-    r = baidu::xpu::api::plugin::recover_block(
+        r = baidu::xpu::api::plugin::recover_block(
         xpu_ctx->x_context(),
         const_cast<int*>(recover_block_list.data<int>()),
         const_cast<int*>(recover_lens.data<int>()),
@@ -105,7 +108,8 @@ void StepPaddle(const paddle::Tensor& stop_flags,
         length,
         pre_id_length);
     PD_CHECK(r == 0, "recover_block failed.");
-  }
+    }
+  baidu::xpu::api::plugin::print_times("[TIME END] StepPaddle");
 }
 
 PD_BUILD_OP(step_paddle)

csrc/xpu/src/update_inputs.cc

@@ -26,6 +26,9 @@ void UpdateInputes(const paddle::Tensor& stop_flags,
                    const paddle::Tensor& stop_nums,
                    const paddle::Tensor& next_tokens,
                    const paddle::Tensor& is_block_step) {
+  baidu::xpu::api::plugin::print_times("[TIME BEGIN] UpdateInputes");
+
+
   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
   auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
   auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
@@ -53,7 +56,11 @@ void UpdateInputes(const paddle::Tensor& stop_flags,
       max_bsz,
       input_ids_stride);
   PD_CHECK(r == 0, "baidu::xpu::api::plugin::update_inputs failed.");
-}
+
+
+  baidu::xpu::api::plugin::print_times("[TIME END] UpdateInputes");
+  }
+
 
 PD_BUILD_OP(update_inputs)
     .Inputs({"stop_flags",

csrc/xpu/src/weight_only_linear.cc

@@ -1,142 +0,0 @@
-// Copyright (c) 2024 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.
-
-#include <paddle/phi/backends/xpu/xpu_context.h>
-#include "paddle/extension.h"
-#include "paddle/phi/core/enforce.h"
-#include "xpu/plugin.h"
-#include <xft/xdnn_plugin.h>
-#include "xblas_legacy_api.h"
-
-namespace xftkernel = baidu::xpu::xftkernel;
-namespace api = baidu::xpu::api;
-// namespace xblas = baidu::xpu::xblas;
-
-std::vector<paddle::Tensor> WeightOnlyLinearKernel(
-    const paddle::Tensor& x,
-    const paddle::Tensor& weight,
-    const paddle::optional<paddle::Tensor>& weight_scale,
-    const paddle::optional<paddle::Tensor>& bias,
-    const std::string& weight_dtype="int8",
-    const int32_t arch=0,
-    const int32_t group_size=64
-    ) {
-
-  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
-  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
-  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
-
-  using XPUType = typename XPUTypeTrait<bfloat16>::Type;
-  typedef paddle::bfloat16 data_t;
-
-  int64_t n = weight.dims()[0];
-  int64_t k = weight.dims()[1];
-  int64_t m = x.numel() / k;
-  if(weight_dtype == "int4") {
-    n = n * 2;
-  }
-
-  paddle::Tensor bias_fp32;
-  if (bias.is_initialized() && bias.get().dtype() == phi::DataType::FLOAT16) {
-    bias_fp32 = paddle::full(bias->shape(), -1, paddle::DataType::FLOAT32, bias->place());
-    int r = baidu::xpu::api::cast<XPUType, float>(
-        xpu_ctx->x_context(),
-        reinterpret_cast<const XPUType*>(
-            bias.get().data<phi::dtype::float16>()),
-        bias_fp32.data<float>(),
-        n);
-  }
-
-  paddle::Tensor weight_scale_fp32;
-  if (weight_scale.is_initialized() && weight_scale.get().dtype() == phi::DataType::FLOAT16) {
-    weight_scale_fp32 = paddle::full(weight_scale->shape(), -1, paddle::DataType::FLOAT32, weight_scale->place());
-    int r = baidu::xpu::api::cast<XPUType, float>(
-        xpu_ctx->x_context(),
-        reinterpret_cast<const XPUType*>(
-            weight_scale.get().data<phi::dtype::float16>()),
-        weight_scale_fp32.data<float>(),
-        n);
-  }
-
-  auto output = paddle::full({m, n}, -1, x.type(), x.place());
-
-  auto input_x = reinterpret_cast<const XPUType*>(x.data<data_t>());
-  auto input_y = reinterpret_cast<XPUType*>(output.data<data_t>());
-
-  baidu::xpu::xblas::FcFusionTensor<const XPUType>
-          tensor_x{input_x, nullptr, m, k, k, false};
-  baidu::xpu::xblas::FcFusionTensor<const XPUType> tensor_y_const{
-      input_y, nullptr, m, n, n, false};
-  baidu::xpu::xblas::FcFusionTensor<XPUType> tensor_y{
-      input_y, nullptr, m, n, n, false};
-
-  baidu::xpu::xblas::FcFusionEpilogue<float, float> epilogue{
-      api::Activation_t::LINEAR,
-      bias ? bias_fp32.data<float>() : nullptr,
-      nullptr,
-      weight_scale ? weight_scale_fp32.data<float>() : nullptr,  // weight_scale_fp32.data<float>(), 
-      0,
-      1,
-      nullptr};
-  if(weight_dtype == "int4") {
-    baidu::xpu::xblas::FcFusionDesc<int4_wo_int15, float, XPUType> desc{1.0f, 0.0f};
-    baidu::xpu::xblas::FcFusionTensor<const int4_t> tensor_w{
-        reinterpret_cast<const int4_t*>(weight.data<int8_t>()), nullptr, n, k, k, true};
-    int r1 =  baidu::xpu::xblas::fc_fusion<XPUType,
-                                        int4_t,
-                                        XPUType,
-                                        XPUType,
-                                        int4_wo_int15, // int8_wo_t
-                                        float,
-                                        XPUType,
-                                        float,
-                                        float>(xpu_ctx->x_context(),
-                                                tensor_x,
-                                                tensor_w,
-                                                tensor_y_const,
-                                                tensor_y,
-                                                desc,
-                                                epilogue);
-    PD_CHECK(r1 == 0, "xblas::fc_fusion failed");
-  } else {
-    baidu::xpu::xblas::FcFusionDesc<float, float, float> desc{1.0f, 0.0f};
-    baidu::xpu::xblas::FcFusionTensor<const XPUType> tensor_w{
-        // reinterpret_cast<const XPUType*>(weight.data<phi::dtype::bfloat16>()), nullptr, k, n, n, true};
-        reinterpret_cast<const XPUType*>(weight.data<phi::dtype::bfloat16>()), nullptr, n, k, k, true};
-    int r1 =  baidu::xpu::xblas::fc_fusion<XPUType,
-                                        bfloat16,
-                                        XPUType,
-                                        XPUType,
-                                        float, 
-                                        float,
-                                        float,
-                                        float,
-                                        float,
-                                        XPUType>(xpu_ctx->x_context(),
-                                                tensor_x,
-                                                tensor_w,
-                                                tensor_y_const,
-                                                tensor_y,
-                                                desc,
-                                                epilogue);
-    PD_CHECK(r1 == 0, "xblas::fc_fusion failed");
-  }
-  return {output};
-}
-
-PD_BUILD_OP(weight_only_linear_kernel)
-    .Inputs({"x", "weight", paddle::Optional("weight_scale"), paddle::Optional("bias")})
-    .Outputs({"output"})
-    .Attrs({"weight_dtype: std::string", "arch: int", "group_size: int"})
-    .SetKernelFn(PD_KERNEL(WeightOnlyLinearKernel));

csrc/xpu/test/python/test_blha_get_max_len.py

@@ -20,24 +20,20 @@ import paddle
 np.random.seed(2023)
 
 
-class GetUpdateInputsTest(unittest.TestCase):
-    def test_update_inputs(self):
+def test_update_inputs():
+    seq_lens_encoder = paddle.to_tensor(
+        [[0], [0], [0], [0]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
+    )
+    seq_lens_decoder = paddle.to_tensor(
+        [[27], [29], [31], [31]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
+    )
+    batch_size = paddle.to_tensor(
+        [0, 8191, 16382, 24573], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
+    )
 
-        seq_lens_encoder = paddle.to_tensor(
-            [[0], [0], [0], [0]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
-        )
-        seq_lens_decoder = paddle.to_tensor(
-            [[27], [29], [31], [31]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
-        )
-        batch_size = paddle.to_tensor(
-            [0, 8191, 16382, 24573], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
-        )
+    a, b = paddle.incubate.nn.functional.blha_get_max_len(seq_lens_encoder, seq_lens_decoder, batch_size)
 
-        a, b = paddle.incubate.nn.functional.blha_get_max_len(seq_lens_encoder, seq_lens_decoder, batch_size)
+    print(a)
+    print(b)
 
-        print(a)
-        print(b)
-
-
-if __name__ == "__main__":
-    unittest.main()
+test_update_inputs()

csrc/xpu/test/python/test_get_position_ids.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2024 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.
+
+import numpy as np
+import paddle
+import unittest
+from paddlenlp_ops import update_inputs
+
+np.random.seed(2023)
+class GetUpdateInputsTest(unittest.TestCase):
+    def test_update_inputs(self):
+        bs = 48
+        max_bs = 64
+        max_input_length = 6144
+
+        stop_flags = np.random.randint(0, 2, max_bs).astype("bool")
+        not_need_stop = np.array([1], "bool")
+        seq_lens_this_time = np.zeros([bs], "int32")
+        seq_lens_encoder = np.zeros([max_bs], "int32")
+        seq_lens_decoder = np.zeros([max_bs], "int32")
+        for i in range(bs):
+            if i % 2 == 0:
+                seq_lens_encoder[i] = i
+                seq_lens_this_time[i] = i
+            else:
+                seq_lens_decoder[i] = i
+                seq_lens_this_time[i] = 1
+        input_ids_np = np.random.randint(1, 10, [max_bs, max_input_length], "int64")
+        stop_nums = np.array([max_bs], "int64")
+        next_tokens = np.random.randint(1, 10, [max_bs], "int64")
+        is_block_step = np.random.randint(0, 2, [max_bs]).astype("bool")
+
+        stop_flags = paddle.to_tensor(stop_flags)
+        not_need_stop = paddle.to_tensor(not_need_stop, place=paddle.CPUPlace())
+        seq_lens_this_time = paddle.to_tensor(seq_lens_this_time)
+        seq_lens_encoder = paddle.to_tensor(seq_lens_encoder)
+        seq_lens_decoder = paddle.to_tensor(seq_lens_decoder)
+        input_ids = paddle.to_tensor(input_ids_np)
+        stop_nums = paddle.to_tensor(stop_nums)
+        next_tokens = paddle.to_tensor(next_tokens)
+        is_block_step = paddle.to_tensor(is_block_step)
+
+        print("stop_flags:\n", stop_flags)
+        print("not_need_stop:\n", not_need_stop)
+        print("seq_lens_this_time:\n", seq_lens_this_time)
+        print("seq_lens_encoder:\n", seq_lens_encoder)
+        print("seq_lens_decoder:\n", seq_lens_decoder)
+        print("input_ids:\n", input_ids)
+        print("stop_nums:\n", stop_nums)
+        print("next_tokens:\n", next_tokens)
+        print("is_block_step:\n", is_block_step)
+
+        update_inputs(
+            stop_flags,
+            not_need_stop,
+            seq_lens_this_time,
+            seq_lens_encoder,
+            seq_lens_decoder,
+            input_ids,
+            stop_nums,
+            next_tokens,
+            is_block_step
+        )
+
+        print("-" * 50)
+        print("stop_flags:\n", stop_flags)
+        print("not_need_stop:\n", not_need_stop)
+        print("seq_lens_this_time:\n", seq_lens_this_time)
+        print("seq_lens_encoder:\n", seq_lens_encoder)
+        print("seq_lens_decoder:\n", seq_lens_decoder)
+        print("input_ids:\n", input_ids)
+        print("stop_nums:\n", stop_nums)
+        print("next_tokens:\n", next_tokens)
+
+        ref_not_need_stop_out = np.array([True])
+        ref_seq_lens_this_time_out = np.array([0, 0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1,
+                                            1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 1, 1], "int32")
+        ref_seq_lens_encoder_out = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "int32")
+        ref_seq_lens_decoder_out = np.array([0, 0, 2, 0, 0, 6, 0, 8, 8, 10, 0, 12, 12, 0, 0, 0, 0, 0, 0, 0, 20, 22, 0, 24,
+                                            24, 0, 26, 28, 0, 0, 0, 32, 32, 0, 34, 0, 0, 38, 0, 40, 0, 0, 42, 0, 0, 46, 46, 48,
+                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "int32")
+        input_ids_np[:, 0] = np.array([6, 5, 9, 8, 6, 2, 8, 1, 3, 1, 3, 6, 9, 8, 1, 9, 1, 8, 8, 6, 7, 6, 5, 3,
+                                    5, 9, 3, 6, 3, 9, 8, 8, 8, 8, 4, 8, 7, 4, 2, 3, 5, 8, 4, 2, 5, 6, 8, 9,
+                                    6, 7, 4, 2, 4, 6, 2, 3, 4, 9, 7, 2, 1, 8, 7, 8], "int64")
+
+        assert not_need_stop.numpy() == ref_not_need_stop_out, 'Check not_need_stop failed.'
+        assert np.all(seq_lens_this_time.numpy()
+                    == ref_seq_lens_this_time_out), 'Check seq_lens_this_time failed.'
+        assert np.all(seq_lens_encoder.numpy()
+                    == ref_seq_lens_encoder_out), 'Check seq_lens_encoder failed.'
+        assert np.all(seq_lens_decoder.numpy()
+                    == ref_seq_lens_decoder_out), 'Check seq_lens_decoder failed.'
+        assert np.all(input_ids.numpy()
+                    == input_ids_np), 'Check input_ids failed.'
+
+if __name__ == '__main__':
+    unittest.main()

docs/en/llm/alignment/ppo/README.md

@@ -91,7 +91,7 @@ The configuration files we use are located in `llm/config/llama/ppo_argument.jso
 - `top_p`: Generation decoding hyperparameter
 - `temperature`: Generation decoding hyperparameter
 - `repetition_penalty`: Generation decoding hyperparameter
-- `num_return_sequences`: Generation decoding hyperparameter
+- `rollout_n`: Generation decoding hyperparameter
 - `min_learning_rate`: Minimum learning rate for Actor model
 - `critic_learning_rate`: Minimum learning rate for Critic model
 - `recompute`: Whether to use recomputation strategy for Actor model, enables training memory saving

docs/en/llm/docs/pretrain.rst

@@ -72,7 +72,7 @@ pip install fast_dataindex visualdl==2.5.3`` and other missing whl packages.
 6. During multi-machine training, if all machines use training data files from the same location (e.g., mounted shared storage), specify ``--share_folder true`` to have global rank 0 create cached data. Otherwise, rank 0 on each machine will create cached data independently by default.
 7. If the default cache folder ``index-cache/`` exists in the dataset directory, the additionally specified ``--data_cache`` will not take effect, and training will prioritize loading content from the default cache folder.
 
-The pre-training uses PaddleNLP's Trainer module. For related distributed strategy usage, please refer to the `Large Model Trainer Hybrid Parallel Training Tutorial <./llm_trainer.rst>`.
+The pre-training uses PaddleNLP's Trainer module. For related distributed strategy usage, please refer to the `Large Model Trainer Hybrid Parallel Training Tutorial <./llm_trainer.rst>`_.
 
 .. _model_capability:
 Overview of Distributed Capabilities Supported in Model Pre-training

docs/en/llm/docs/rlhf.md

@@ -239,7 +239,7 @@ Most parameters in `ppo_argument.json` share the same definitions as [LLM Fine-t
 - `reward_critic_model_name_or_path`: Model name or path for the critic-model. If not provided, `reward_model_name_or_path` will be used to initialize the critic-model.
 - `per_device_prompt_batch_size`: Batch size per device for prompt-only dataset loading during rollout generation in training.
 - `per_device_train_batch_size`: Batch size per device for generation and training based on prompts.
-- `num_return_sequences`: Number of responses generated per prompt during generation.
+- `rollout_n`: Number of responses generated per prompt during generation.
 `GenerationConfig.num_return_sequences`, all responses will be used for training.
 - `temperature`: The `temperature` used in generation sampling, i.e., `GenerationConfig.temperature`.
 - `top_p`: The top-p-filtering threshold used in generation sampling, i.e., `GenerationConfig.top_p`.

docs/generate_model_list.py

@@ -21,42 +21,34 @@ MODEL_ROOT = "/xx/bos/community/"
 URL_BASE = "https://paddlenlp.bj.bcebos.com/models/community/"
 OUTPUT_DIR = "./website"
 
+# Markdown templates
 MAIN_TEMPLATE = """
-<!DOCTYPE html>
-<html>
-<head>
-    <title>Model Downloads</title>
-</head>
-<body>
-    <h1>Available Models</h1>
-    <ul>
-        {% for model in models %}
-        <li><a href="{{ model }}/index.html">{{ model }}</a></li>
-        {% endfor %}
-    </ul>
-</body>
-</html>
+# Model Downloads
+
+## Available Models
+
+{% for model in models %}
+- [{{ model }}]({{ model }}/index.md)
+{% endfor %}
 """
 
 MODEL_TEMPLATE = """
-<!DOCTYPE html>
-<html>
-<head>
-    <title>{{ model_name }}</title>
-</head>
-<body>
-    <h1>{{ model_name }} Files</h1>
-    <ul>
-        {% for file in files %}
-        <li>
-            <a href="{{ model_path }}/{{ file.name }}" download>{{ file.name }}</a>
-            <span>({{ file.size }})</span>
-        </li>
-        {% endfor %}
-    </ul>
-    <p><a href="../../">Back to Main</a></p>
-</body>
-</html>
+# {{ model_name }}
+---
+
+{% if readme_content %}
+## README([From Huggingface]({{ huggingface_url }}))
+
+{{ readme_content }}
+
+{% endif %}
+
+## Model Files
+{% for file in files %}
+- [{{ file.name }}]({{ model_path }}/{{ file.name }}) ({{ file.size }})
+{% endfor %}
+
+[Back to Main]({{back_to_main_path}})
 """
 
 
@@ -69,38 +61,86 @@ def convert_size(size_bytes):
     return f"{size_bytes:.1f} {units[unit_index]}"
 
 
+def process_image_links(text, model_path):
+    image_link_pattern = re.compile(r"!\[.*?\]\((.*?)\)")
+
+    image_links = image_link_pattern.findall(text)
+    for i, link in enumerate(image_links):
+        if not link.startswith(("http://", "https://", "/")):
+            prefix = f"https://huggingface.co/{model_path}/resolve/main/"
+            image_links[i] = prefix + link
+
+    def replace_link(match):
+        original_link = match.group(1)
+        new_link = next((new_link for new_link in image_links if original_link in new_link), original_link)
+        return f'![{match.group(0).split("](")[0]}]({new_link})'
+
+    processed_text = image_link_pattern.sub(replace_link, text)
+    return processed_text
+
+
+def process_license(text):
+    license_pattern = re.compile(r"---\nlicense:.*?---", re.DOTALL)
+    processed_text = license_pattern.sub("", text)
+    return processed_text
+
+
+def get_back_to_main_path(model_path):
+    # calculate the level by counting the number of slashes
+    level = model_path.count("/") + 1
+
+    # back_to_main_path = '../' * (level - 1)
+    back_to_main_path = "../" * level
+    return back_to_main_path
+
+
 def generate_model_page(model_path, model_name):
     full_path = os.path.join(MODEL_ROOT, model_path)
     files = []
+    readme_content = False
 
     for root, _, filenames in os.walk(full_path):
         for f in filenames:
-            if f.endswith("index.html"):
+            if f.endswith("index.md"):
                 continue
             file_path = os.path.join(root, f)
             rel_path = os.path.relpath(file_path, full_path)
 
+            if f == "README.md":
+                with open(file_path, "r", encoding="utf-8") as rf:
+                    readme_content = rf.read()
+
             size = os.path.getsize(file_path)
             files.append({"name": rel_path, "size": convert_size(size)})
 
     output_path = os.path.join(OUTPUT_DIR, model_path)
     os.makedirs(output_path, exist_ok=True)
 
+    if readme_content:
+        readme_content = process_image_links(readme_content, model_path)
+        readme_content = process_license(readme_content)
+    huggingface_url = os.path.join("https://huggingface.co", model_path)
+    back_to_main_path = get_back_to_main_path(model_path)
     template = Template(MODEL_TEMPLATE)
-    html = template.render(
-        model_name=model_name, model_path=URL_BASE + model_path, files=sorted(files, key=lambda x: x["name"])
+    markdown_content = template.render(
+        model_name=model_name,
+        huggingface_url=huggingface_url,
+        model_path=URL_BASE + model_path,
+        files=sorted(files, key=lambda x: x["name"]),
+        readme_content=readme_content,
+        back_to_main_path=back_to_main_path,
     )
 
-    with open(os.path.join(output_path, "index.html"), "w") as f:
-        f.write(html)
+    with open(os.path.join(output_path, "index.md"), "w") as f:
+        f.write(markdown_content)
 
 
 def generate_main_page(models):
     template = Template(MAIN_TEMPLATE)
-    html = template.render(models=sorted(models))
+    markdown_content = template.render(models=sorted(models))
 
-    with open(os.path.join(OUTPUT_DIR, "index.html"), "w") as f:
-        f.write(html)
+    with open(os.path.join(OUTPUT_DIR, "index.md"), "w") as f:
+        f.write(markdown_content)
 
 
 def is_model_directory(path):
@@ -120,7 +160,17 @@ def is_model_directory(path):
     return len(model_files) > 0 or len(sharded_files) > 0
 
 
-ommit_paths = ["_internal_", "hf-internal", "zhuweiguo"]
+ommit_paths = [
+    "_internal_",
+    "hf-internal",
+    "zhuweiguo",
+    "ziqingyang",
+    "yuhuili",
+    "westfish",
+    "junnyu",
+    "Yang-Changhui",
+    "baicai",
+]
 
 
 def find_models():

docs/zh/conf.py

@@ -110,24 +110,31 @@ html_css_files = [
 html_static_path = ["_static"]
 html_logo = "paddle.png"
 
-import tarfile
-from io import BytesIO
-
-import requests
+import shutil
+import subprocess
 
 
 def extract_tar(app, config):
     try:
-        url = "https://paddlenlp.bj.bcebos.com/datasets/website.tar"
-        response = requests.get(url)
-        response.raise_for_status()  # 检查请求是否成功
+        repo_url = "https://github.com/ZHUI/paddlenlp_model_readme.git"
+        clone_dir = "./paddlenlp_model_readme"
 
-        file_stream = BytesIO(response.content)
+        # 清理已存在的目录
+        if os.path.exists(clone_dir):
+            shutil.rmtree(clone_dir)
 
-        with tarfile.open(fileobj=file_stream, mode="r") as tar:
-            tar.extractall("_static/")
-    except:
-        pass
+        # 执行克隆
+        subprocess.run(["git", "clone", "--depth=1", repo_url], check=True)
+        if os.path.exists("./website"):
+            shutil.rmtree("./website")
+        shutil.move(clone_dir, "./website")
+        print("Repository cloned successfully!")
+    except subprocess.CalledProcessError as e:
+        print(f"Git操作失败: {e}")
+    except FileNotFoundError:
+        print("未找到git命令，请确保已安装Git并配置环境变量")
+    except Exception as e:
+        print(f"发生未知错误: {e}")
 
 
 def setup(app):

docs/zh/index.rst

@@ -38,7 +38,7 @@
    :caption: 模型库列表
    
    热门模型介绍 <hot_model.md>
-   模型库列表 <model_list.rst>
+   模型库列表 <website/index.md>
    
 
 .. toctree::
@@ -128,6 +128,10 @@
    :caption: 实践教程
 
    AI Studio Notebook <tutorials/overview>
+   大模型预训练新手指南 <llm/docs/pretrain_tutorial.md>
+   大模型精调新手指南 <llm/docs/finetune_tutorial.md>
+   大模型对齐新手指南 <llm/docs/alignment_tutorial.md>
+   大模型量化新手指南 <llm/docs/quantization_tutorial.md>
 
 .. toctree::
    :maxdepth: 1

docs/zh/llm/alignment/ppo/README.md

@@ -1 +0,0 @@
-../../../../../llm/alignment/ppo/README.md

docs/zh/llm/alignment/rl/README.md

@@ -0,0 +1 @@
+../../../../../llm/alignment/rl/README.md

docs/zh/llm/devices/intel_hpu/tests/README.md

@@ -0,0 +1 @@
+../../../../../../llm/devices/intel_hpu/tests/README.md

docs/zh/llm/docs/alignment_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/alignment_tutorial.md

docs/zh/llm/docs/finetune_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/finetune_tutorial.md

docs/zh/llm/docs/pretrain.md

@@ -0,0 +1 @@
+../../../../llm/docs/pretrain.md

docs/zh/llm/docs/pretrain_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/pretrain_tutorial.md

docs/zh/llm/docs/quantization_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/quantization_tutorial.md

docs/zh/locale/en/LC_MESSAGES/source/paddlenlp.taskflow.utils.po

@@ -276,7 +276,7 @@ msgid "word of current node."
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.BurkhardKellerTree:1
-msgid "Implementataion of BK-Tree"
+msgid "Implementation of BK-Tree"
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.BurkhardKellerTree.add:1
@@ -300,7 +300,7 @@ msgid "similar words."
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.TriedTree:1
-msgid "Implementataion of TriedTree"
+msgid "Implementation of TriedTree"
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.TriedTree.add_word:1

docs/zh/locale/en/LC_MESSAGES/source/paddlenlp.transformers.squeezebert.tokenizer.po

@@ -50,7 +50,7 @@ msgid ""
 msgstr ""
 
 #: of paddlenlp.transformers.squeezebert.tokenizer.SqueezeBertTokenizer:11
-msgid "The special token for unkown words. Default: \"[UNK]\"."
+msgid "The special token for unknown words. Default: \"[UNK]\"."
 msgstr ""
 
 #: of paddlenlp.transformers.squeezebert.tokenizer.SqueezeBertTokenizer:13

docs/zh/model_list.rst

@@ -1,10 +0,0 @@
-模型库列表
----------
-
-.. raw:: html
-
-   <iframe 
-     src="_static/website/index.html" 
-     style="width: 100%; height: 1000px; border: none;"
-   ></iframe>
-

llm/README.md

@@ -13,27 +13,9 @@
 - **大模型无损量化**。大模型套件预先集成了 PaddleSlim LLM.PTQ 量化算法，以及业界广泛采用的 GPTQ 和 AWQ 的 W4量化方法，成功实现了对主流大模型的无损量化处理，显著加速了模型的推理速度。
 - **高性能推理**。大模型套件的高性能推理模块内置了动态插入和全环节算子融合的高级策略，这极大地提升了并行推理的速度。同时，该模块隐藏了底层技术细节，为用户提供了开箱即用的高性能并行推理功能。
 
-
 ## 🛠️ 支持模型列表 🛠️
 
-| Model                                  | Pretrain | SFT | LoRA | Prefix Tuning | DPO/SimPO/ORPO/KTO | RLHF | Mergekit | Quantization | Torch convert |
-|----------------------------------------|----------|-----|------|---------------|----------------|------|-------|--------------|---------------|
-| [LLaMA](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)                | ✅        | ✅   | ✅    | ✅             | ✅             | ✅    | ✅    | ✅            | ✅             |
-| [Qwen](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen)                  | ✅        | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | 🚧           | ✅             |
-| [Mixtral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mixtral)            | ✅        | ✅   | ✅    | ❌             | ✅             | 🚧   | ✅    | 🚧           | 🚧            |
-| [Mistral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mistral)            | ✅         | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | 🚧           | ✅             |
-| [Baichuan/Baichuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)   | ✅        | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | ✅            | ✅             |
-| [ChatGLM-6B](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm)         | ✅        | ✅   | ✅    | ✅             | 🚧            | 🚧   | ✅    | ✅            | ❌             |
-| [ChatGLM2/ChatGLM3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm2) | ✅        | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | ✅            | ✅             |
-| [Bloom](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/bloom)                | ✅        | ✅   | ✅    | ✅             | 🚧            | 🚧   | ✅    | ✅            | ✅             |
-| [GPT-3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/gpt-3)                | ✅        | ✅   | 🚧   | 🚧            | 🚧            | 🚧   | ✅    | 🚧           | ✅             |
-| [OPT](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/opt)                    | ✅       | ✅   | ✅    | 🚧            | 🚧            | 🚧   | ✅    | 🚧           | ✅             |
-| [Gemma](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/gemma)                | ✅       | ✅   | ✅    | 🚧            | ✅            | 🚧   | ✅    | 🚧           | 🚧             |
-| [Yuan](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/yuan)                  | ✅       | ✅   | ✅    | 🚧            | ✅            | 🚧   | ✅    | 🚧           | 🚧             |
-
-- ✅: Supported
-- 🚧: In Progress
-- ❌: Not Supported
+参考 [PaddleNLP 支持模型列表](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/README.md#模型支持)
 
 ## 🚀 快速开始 🚀
 
@@ -186,7 +168,7 @@ python run_finetune.py ./config/qwen/pt_argument.json
 
 ### 3. 对齐
 
-我们支持 DPO、KTO、RLHF 等偏好对齐策略。DPO、KTO 策略采用 zero_padding 策略，结合 FlashMask 策略，有效提升模型训练效率。
+我们支持 DPO、KTO、RL 等偏好对齐策略。DPO、KTO 策略采用 zero_padding 策略，结合 FlashMask 策略，有效提升模型训练效率。
 
 #### 3.1 DPO
 
@@ -272,8 +254,8 @@ python -u  ./alignment/dpo/run_dpo.py ./config/llama/dpo_lora_argument.json
 为了方便测试，我们也提供了偏好数据集可以直接使用：
 
 ```bash
-wget https://bj.bcebos.com/paddlenlp/datasets/examples/ultrafeedback_binarized_pointwise.tar.gz
-tar -zxvf ultrafeedback_binarized_pointwise.tar.gz
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/ultrafeedback_binarized_pointwise.tar
+tar -xvf ultrafeedback_binarized_pointwise.tar.gz
 ```
 
 ##### 全参 KTO
@@ -289,20 +271,18 @@ python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./alignment/
 python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./alignment/kto/run_kto.py ./config/llama/kto_lora_argument.json
 ```
 
-#### 3.3 RLHF
+#### 3.3 RL
 
-飞桨大模型套件提供了提供了基于强化学习 PPO 算法对 LLM 进行人类偏好对齐的代码及完整使用示例，支持**3D 分布式并行训练以及 rollout 阶段使用预测优化进行生成加速**。详细使用教程详见[RLHF 文档](./docs/rlhf.md)。
+飞桨大模型套件提供了提供了基于强化学习 GRPO、Reinforce++、PPO 等 算法对 LLM 进行人类偏好对齐的代码及完整使用示例，支持**3D 分布式并行训练以及 rollout 阶段使用预测优化进行生成加速**。详细使用教程详见[RL 文档](./alignment/rl/README.md)。
 
 ### 4. 模型融合
 PadlleNLP 支持多种模型融合方法，包括**Linear、Slerp、Ties、DARE、DELLA**，并支持模型参数稀疏化方法与模型融合算法的灵活组合使用。
 ```shell
 # 模型融合启动命令参考
-# cd PaddleNLP/llm/tools
-python mergekit.py \
-    --tensor_type np \
-    --n_process 2 \
+python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./llm/tools/mergekit.py \
+    --tensor_type pd \
     --merge_method linear \
-    --model_path_list ../checkpoints/model1 ../checkpoints/model \
+    --model_path_str "../checkpoints/model1,../checkpoints/model2" \
     --output_path ../checkpoints/model_merge
 
 ```

llm/alignment/ppo/README.md

@@ -1,141 +0,0 @@
-# PPO && GRPO
-
-PPO（Proximal Policy Optimization，近端策略优化）是一种强化学习算法，旨在通过优化策略来最大化累积奖励。PPO 算法结合了 Policy Gradient 和‌TRPO 的优点，通过使用随机梯度上升优化一个“替代”目标函数，实现小批量更新，而不是每个数据样本只进行一次梯度更新。
-GRPO（Group Relative Policy Optimization，组相对策略优化）是 PPO（Proximal Policy Optimization，近端策略优化）算法的一种变体。与 PPO 不同，GRPO 省略了价值函数估计器。在 GRPO 中，对于每个状态 \(s\)，算法会从当前策略 \(\pi_{\theta_{t}}\) 中采样多个动作 \(a_{1}, \dots, a_{G}\)。然后，GRPO 计算这些动作相对于组内其他动作的“组相对优势”（group-relative advantage），以此作为优化策略的依据。
-
-以下是详细的使用文档和示例：
-
-## 环境依赖
-
-* 训练环境：
-1. 参考 Paddle 官网安装 PaddlePaddle-GPU
-2. clone 并安装 PaddleNLP
-```shell
-git clone https://github.com/PaddlePaddle/PaddleNLP.git
-```
-3. 安装 paddlenlp_ops，参考 PaddleNLP/csrc 进行安装（必需）
-```shell
-cd your_PaddleNLP_path/csrc
-python setup_cuda.py install
-```
-
-## 支持模型
-
-|   模型系列    | 模型名称                                                                                                                                                                                                                                                                      |
-|:-------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-|   Llama3.1    | meta-llama/Meta-Llama-3.1-8B, meta-llama/Meta-Llama-3.1-8B-Instruct, meta-llama/Meta-Llama-3.1-70B, meta-llama/Meta-Llama-3.1-70B-Instruct, meta-llama/Meta-Llama-3.1-405B, meta-llama/Meta-Llama-3.1-405B-Instruct, meta-llama/Llama-Guard-3-8B                              |
-|   Llama3.2    | meta-llama/Llama-3.2-1B, meta-llama/Llama-3.2-1B-Instruct, meta-llama/Llama-3.2-3B, meta-llama/Llama-3.2-3B-Instruct                                                                                                                                                          |
-|    Qwen1.5    | Qwen/Qwen1.5-0.5B, Qwen/Qwen1.5-0.5B-Chat, Qwen/Qwen1.5-1.8B, Qwen/Qwen1.5-1.8B-Chat, Qwen/Qwen1.5-4B, Qwen/Qwen1.5-4B-Chat, Qwen/Qwen1.5-7B, Qwen/Qwen1.5-7B-Chat, Qwen/Qwen1.5-14B, Qwen/Qwen1.5-14B-Chat, Qwen/Qwen1.5-32B, Qwen/Qwen1.5-32B-Chat                          |
-|     Qwen2     | Qwen/Qwen2-0.5B, Qwen/Qwen2-0.5B-Instruct, Qwen/Qwen2-1.5B, Qwen/Qwen2-1.5B-Instruct, Qwen/Qwen2-7B, Qwen/Qwen2-7B-Instruct, Qwen/Qwen2-72B, Qwen/Qwen2-72B-Instruct, Qwen/Qwen2-57B-A14B, Qwen/Qwen2-57B-A14B-Instruct                                                       |
-|  Qwen2-Math   | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct                                                                                                                                                                          |
-|    Qwen2.5    | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, |
-| Qwen2.5-Math  | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct                                                                                                                                                                  |
-| Qwen2.5-Coder | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                              |
-
-## 数据协议
-
-### 字段说明
-
-- src (list(str)): 用户对话内容，可能会包含 markup 内容，如 [<search-res>]；
-- tgt (list(str)): 除了最后一轮的系统多轮回复内容，以对话轮次排列，可能会包含 markup 内容，如 [<search>]；注意：len(tgt)==len(src)-1
-
-### 数据示例
-
-```json
-{
-    "src": [
-        "需要你帮我写几个有创意的广告语来打开市场。",
-        "目标用户是年轻人，追求时尚、个性和自我。"
-    ],
-    "tgt": [
-        "当然！我很乐意帮助你创作几个有创意的广告语来推广你的新洗发露。请告诉我一些关于你的产品的特点，目标受众以及你希望传达的核心信息，我会根据这些信息为你提供几个创意的广告语。"
-    ]
-}
-```
-
-
-### PPO & GRPO 数据准备
-
-```
-wget https://paddlenlp.bj.bcebos.com/datasets/examples/ppo-kk.tgz && tar zxf ppo-kk.tgz
-```
-
-## 训练
-
-### 训练配置
-
-我们采用的配置文件在放置在`llm/config/llama/ppo_argument.json`和`llm/config/llama/grpo_argument.json`中，同时我们提供了详细参数释义如下：
-
-- `train_task_config`: 训练数据 config, 请以`config/task_ppo.json`为例
-- `eval_task_config`: 评估数据 config, 请以`config/task_ppo.json`为例
-- `ptx_task_config`: SFT 辅助数据, 请以`config/task_sft.json`为例，默认为""
-- `actor_model_name_or_path`: PPO 中 actor-model 和 reference-model 模型本地的模型路径
-- `reward_model_name_or_path`: PPO 中 reward-model 和 critic-model 模型本地的模型路径
-- `use_fusemt`: 是否通过 FustMT 加速生成，默认为 True
-- `use_flash_attention`: 是否启用 FlashAttention-2，默认为 False
-- `output_dir`: 模型参数保存目录
-- `max_seq_len`: 输入数据的最大长度，默认为 4096
-- `max_dec_len`: 最大生成长度
-- `min_dec_len`: 最小生成长度
-- `top_p`: 生成解码超参数
-- `temperature`: 生成解码超参数
-- `repetition_penalty`: 生成解码超参数
-- `num_return_sequences`: 生成解码超参数
-- `min_learning_rate`: Actor 模型的最小学习率
-- `critic_learning_rate`: Critic 模型的最小学习率
-- `recompute`: Actor 模型是否使用重计算策略，开启后可节省训练显存
-- `critic_recompute`: Critic 模型是否使用重计算策略，开启后可节省训练显存
-- `recompute_granularity` Actor 模型的重计算的粒度，可选项为`core_attn`和`full`. `core_attn`速度快但是显存占用，`full`速度慢但是显存占用低
-- `critic_recompute_granularity` Critic 模型重计算的粒度，可选项为`core_attn`和`full`. `core_attn`速度快但是显存占用，`full`速度慢但是显存占用低
-- `warmup_ratio`: Actor 模型用于从 0 到 `learning_rate` 的线性 warmup 的总训练步骤的比例
-- `critic_warmup_ratio`: Critic 模型用于从 0 到 `critic_learning_rate` 的线性 warmup 的总训练步骤的比例
-- `lr_scheduler_type`: Actor 模型要使用的学习率调度策略。 (`str`, 可选, 默认为 `"linear"`)
-- `critic_lr_scheduler_type`: Critic 模型要使用的学习率调度策略。 (`str`, 可选, 默认为 `"linear"`)
-- `weight_decay`: Actor 模型除了所有 bias 和 LayerNorm 权重之外，应用于所有层的权重衰减数值。（`float`，可选，默认为 0.0）
-- `critic_weight_decay`: Critic 模型除了所有 bias 和 LayerNorm 权重之外，应用于所有层的权重衰减数值。（`float`，可选，默认为 0.0）
-- `max_prompt_len`: 生成样本时的最大生成长度， max_length 调大会增加生成时间，并且增加显存占用。注意：
-max_dec_len + max_prompt_len 应当小于 max_seq_len。
-- `per_device_prompt_batch_size`: PPO 生成样本时的批处理大小，同 micro batch size，即满足 global_batch_size = dp（data parallel）* sharding * micro batch size。batch_size 调大会增加生成时间，并且增加显存占用
-- `per_device_train_batch_size`: 训练 batch 大小, 当前为了优化性能设为1，请避免更改
-- `per_device_eval_batch_size`: 评估 batch 大小。
-- `max_steps`: 总的训练步数
-- `eval_steps`: 模型评估的间隔步数
-- `max_evaluate_steps`: 模型单次评估的最大步数
-- `logging_steps`: 训练日志打印的间隔步数
-- `save_steps`: 模型参数保存的间隔步数
-- `weight_decay`: 权重衰减数值
-- `do_train`: 是否进行训练任务
-- `do_eval`: 是否进行评估任务
-- `fp16`: 使用 float16 精度进行模型训练和推理。
-- `bf16`: 使用 bfloat16 精度进行模型训练和推理。
-- `fp16_opt_level`: float16 精度训练模式，`O2`表示纯 float16 训练
-
-
-<!-- ### PPO 训练命令
-
-```shell
-python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7"  run_ppo.py llm/config/llama/ppo_argument.json
-``` -->
-
-### GRPO 训练命令
-```shell
-cd your_PaddleNLP_path/llm/alignment/ppo
-```
-
-```shell
-# 启动 reward server
-python reward_server.py
-```
-
-```shell
-export PYTHONPATH=your_PaddleNLP_path/:$PYTHONPATH
-export PYTHONPATH=your_PaddleNLP_path/llm:$PYTHONPATH
-python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_ppo.py ../../config/qwen/grpo_argument.json
-# python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_ppo.py ../../config/llama/grpo_argument.json
-```
-
-### 在线监控
-在`grpo_argument.json`中设置的输出目录为`"logging_dir": "vdl_log"`, 可以通过以下命令查看训练过程
-```shell
-visualdl --logdir vdl_log --host 0.0.0.0
-```

llm/alignment/rl/README.md

@@ -0,0 +1,214 @@
+# GRPO && REINFORCE++
+
+GRPO（Group Relative Policy Optimization，组相对策略优化）是 PPO（Proximal Policy Optimization，近端策略优化）算法的一种变体。与 PPO 不同，GRPO 省略了价值函数估计器。在 GRPO 中，对于每个状态 $s$，算法会从当前策略 $\pi_{\theta_{t}}$ 中采样多个动作 $a_{1}, \dots, a_{G}$。然后，GRPO 计算这些动作相对于组内其他动作的“组相对优势”（group-relative advantage），以此作为优化策略的依据。
+REINFORCE++ 是经典 REINFORCE 算法的改进版本，通过融合 PPO 的关键优化技术并移除 Critic Model，实现了更加简洁高效的策略优化。相比于传统的 REINFORCE，REINFORCE++ 在 Token-Level KL 惩罚、PPO-Clip、优势标准化、奖励裁剪和奖励标准化等关键技术上进行了改进，从而提高了训练过程的效率和稳定性。
+
+以下是详细的使用文档和示例：
+
+## 环境依赖
+
+* 训练环境：
+1. 参考 [Paddle 官网](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/develop/install/pip/linux-pip.html)安装 PaddlePaddle-GPU, 要求 PaddlePaddle>=3.0
+2. clone 并安装 PaddleNLP
+```shell
+git clone https://github.com/PaddlePaddle/PaddleNLP.git
+```
+3. 安装 paddlenlp_ops 推理算子，参考 PaddleNLP/csrc 进行安装（必需）
+```shell
+cd your_PaddleNLP_path/csrc
+python setup_cuda.py install
+```
+4. 安装 fused_ln 和 fast_ln 训练算子，参考 PaddleNLP/slm/model_zoo/gpt-3/external_ops (必须)
+```shell
+cd your_PaddleNLP_path/slm/model_zoo/gpt-3/external_ops
+python setup.py install
+```
+
+## 支持模型
+
+|   模型系列    | 模型名称                                                                                                                                                                                                                                                                      |
+|:-------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+|    Qwen1.5    | Qwen/Qwen1.5-0.5B, Qwen/Qwen1.5-0.5B-Chat, Qwen/Qwen1.5-1.8B, Qwen/Qwen1.5-1.8B-Chat, Qwen/Qwen1.5-4B, Qwen/Qwen1.5-4B-Chat, Qwen/Qwen1.5-7B, Qwen/Qwen1.5-7B-Chat, Qwen/Qwen1.5-14B, Qwen/Qwen1.5-14B-Chat, Qwen/Qwen1.5-32B, Qwen/Qwen1.5-32B-Chat                          |
+|     Qwen2     | Qwen/Qwen2-0.5B, Qwen/Qwen2-0.5B-Instruct, Qwen/Qwen2-1.5B, Qwen/Qwen2-1.5B-Instruct, Qwen/Qwen2-7B, Qwen/Qwen2-7B-Instruct, Qwen/Qwen2-72B, Qwen/Qwen2-72B-Instruct, Qwen/Qwen2-57B-A14B, Qwen/Qwen2-57B-A14B-Instruct                                                       |
+|  Qwen2-Math   | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct                                                                                                                                                                          |
+|    Qwen2.5    | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, |
+| Qwen2.5-Math  | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct                                                                                                                                                                  |
+| Qwen2.5-Coder | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                              |
+
+## 数据协议
+
+### 字段说明
+
+- src (list(str)): 经过 chat_template 处理后的 prompt 输入；或者根据需要自己拼接构造 prompt；
+- tgt (list(str)): 标签内容；
+
+### 数据示例
+
+```json
+{
+    "src": ["<|im_start|>system\nYou are a helpful assistant. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and<answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>.  Now the user asks you to solve a logical reasoning problem. After thinking, when you finally reach a conclusion, clearly state the identity of each character within <answer> </answer> tags. i.e., <answer> (1) Zoey is a knight\n(2) ... </answer>.\n<|im_end|>\n<|im_start|>user\nA very special island is inhabited only by knights and knaves. Knights always tell the truth, and knaves always lie. You meet 3 inhabitants: Michael, Zoey, and Ethan. Michael was heard saying, \"Ethan is a knight if and only if Michael is a knight\". \"Zoey is a knight or Ethan is a knight,\" Zoey mentioned. Ethan asserted: \"Michael is a knave if and only if Zoey is a knave\". So who is a knight and who is a knave?\n<|im_end|>\n<|im_start|>assistant\n<think>"],
+    "tgt": ["(1) Michael is a knight\n(2) Zoey is a knight\n(3) Ethan is a knight"]
+}
+```
+
+
+### GRPO & REINFORCE++ 数据准备
+我们提供了一版使用 `Qwen/Qwen2.5-7B-Instruct-1M` 的`chat template`预处理后的[KK 数据集](https://hf-mirror.com/datasets/K-and-K/knights-and-knaves)。
+```
+wget https://paddlenlp.bj.bcebos.com/datasets/examples/ppo-kk.tgz && tar zxf ppo-kk.tgz
+```
+
+## 训练
+
+### GRPO && REINFORCE++ 训练配置
+
+我们采用的配置文件放置在`llm/config/qwen/grpo_argument.yaml`中，同时我们提供了详细参数释义如下：
+- `rl_algorithm`: 使用的强化学习算法，支持`grpo`、`reinforce_plus_plus`
+- `actor_model_name_or_path`: actor-model 和 reference-model 模型本地的模型路径
+- `reward_model_name_or_path`: reward 模型的名称或本地路径
+- `use_rm_server`: 是否使用 reward model server，设置为`False`时需要提供`reward_model_name_or_path`
+- `reward_server`: reward model server 的 URL 地址, 比如`http://127.0.0.1:8731`
+- `logging_dir`: 日志保存的文件夹
+- `logging_steps`: 训练日志打印的间隔步数
+- `output_dir`: 模型参数保存目录
+- `report_to`: 训练可视化工具，支持 "all"、"wandb"、"tensorboard"、"visualdl"、"none"
+- `wandb_http_proxy`: 连接 wandb 使用的 HTTP 代理
+- `run_name`: 实验名称
+- `train_datasets`: 训练集路径
+- `eval_datasets`: 验证集路径
+- `prompt_key`: 数据集中 query 对应的字段名
+- `response_key`: 数据集中 response 对应的字段名
+- `dataloader_drop_last`: dataloader 是否丢弃最后不完整的 batch
+- `balance_batch`: 该参数用于指定是否在数据并行场景下，对批次内的 token 数量进行均衡分配。若设置为 True，系统将尝试在不同并行设备间平衡 token 的分布；若设置为 False（默认值），则不进行此类均衡操作。
+- `use_remove_padding`: 此参数决定是否在训练过程中去除输入数据中的 padding 部分。启用该选项（设置为 True）可有效提高训练过程中有效 token 的占比，从而提升训练效率；若设置为 False（默认值），则保留输入数据中的 padding。
+- `tensor_parallel_degree`: 张量并行度
+- `sequence_parallel`: 是否启用序列并行
+- `sharding_parallel_degree`: sharding 并行度
+- `sharding`: 分片策略，支持 "stage1" 或 "stage2"
+- `sharding_parallel_config`: sharding 并行配置
+- `pipeline_parallel_degree`: 流水线并行度
+- `virtual_pp_degree`: 虚拟流水线并行度
+- `max_prompt_len`: 生成样本时的最大生成长度， max_length 调大会增加生成时间，并且增加显存占用。注意：
+max_dec_len + max_prompt_len 应当小于 max_seq_len。
+- `max_dec_len`: 最大生成长度
+- `min_dec_len`: 最小生成长度
+- `top_p`: 生成解码超参数
+- `temperature`: 生成解码超参数
+- `repetition_penalty`: 生成解码超参数
+- `rollout_max_num_seqs`: 单次推理可以处理的最大序列数
+- `rollout_quant_type`: 量化类型，例如 "weight_only_int8"
+- `seed`: 随机种子
+- `global_batch_size`: 一次（一个 step）推理（rollout)采样的 prompt 数量
+- `global_mini_batch_size`: actor model 更新一次参数训练的 prompt 数量
+- `rollout_n`: 一个 prompt 采样的 response 数量
+- `update_iters`: 同一批数据训练次数
+- `per_device_logprob_batch_size`: 计算 log_probs 时，一个 batch 的样本数量
+- `per_device_reward_batch_size`: critic model 计算 loss 与反向传播时，一个 batch 的的样本数量
+- `per_device_value_batch_size`: critic model 前向计算 values 时，一个 batch 的的样本数量
+- `per_device_train_batch_size`: actor model 计算 loss 与反向传播时，一个 batch 的样本数量
+- `num_train_epochs`: 训练的 epoch 数
+- `max_length`: 训练时的最大长度，应大于 `max_prompt_len` 和 `max_dec_len` 之和
+- `learning_rate`: 学习率
+- `lr_scheduler_type`: Actor 模型要使用的学习率调度策略。 (`str`, 可选, 默认为`linear`)
+- `weight_decay`: AdamW 优化器的权重衰减
+- `adam_beta1`: AdamW 优化器的 beta1
+- `adam_beta2`: AdamW 优化器的 beta2
+- `adam_epsilon`: AdamW 优化器的 epsilon
+- `max_grad_norm`: 梯度裁剪的最大值
+- `max_steps`: 总的训练步数
+- `save_steps`: 模型参数保存的间隔步数
+- `ignore_save_lr_and_optim`: 是否忽略保存学习率和优化器状态
+- `kl_coeff`: KL 惩罚系数
+- `kl_loss_coeff`: KL Loss 系数
+- `pg_loss_coeff`: 策略梯度损失系数
+- `entropy_coeff`: entropy loss 系数
+- `clip_range_ratio`: PPO-Clip 裁剪阈值
+- `clip_range_ratio_low`: PPO-Clip 裁剪下限阈值
+- `clip_range_ratio_high`: PPO-Clip 裁剪上限阈值
+- `clip_range_score`: reward 的剪切范围，reward 会被限制在 [-clip_range_score, clip_range_score] 范围内
+- `clip_range_value`: value 模型输出的剪切范围，value 会被限制在 [-clip_range_value, clip_range_value] 范围内
+- `normalize_reward`: 是否使用 reward 标准化
+- `normalize_advantage`: 是否使用 advantage 标准化
+- `use_fp32_compute`: 是否使用 fp32 来计算 log_prob、reward、advantage 和 loss
+- `do_eval`: 是否进行评估
+- `per_device_eval_batch_size`: 估 batch 大小
+- `evaluation_strategy`: 评估策略，例如 `steps`
+- `eval_steps`: 模型评估的间隔步数
+- `use_flash_attention`: 是否启用 FlashAttention-2，默认为 False
+- `use_fused_rms_norm`: 是否使用融合的 RMSNorm 算子，需安装 fused_ln
+- `recompute`: Actor 模型是否使用重计算策略，开启后可节省训练显存
+- `recompute_granularity`: Actor 模型的重计算的粒度，可选项为`core_attn`和`full`. `core_attn`速度快但是显存占用，`full`速度慢但是显存占用低
+- `bf16`: 使用 bfloat16 精度进行模型训练和推理。
+- `fp16_opt_level`: float16 精度训练模式，`O2`表示纯 float16 训练
+- `amp_custom_black_list`: 自定义 AMP 黑名单
+- `amp_custom_white_list`: 自定义 AMP 白名单
+
+
+
+### GRPO 训练命令
+```shell
+cd your_PaddleNLP_path/llm/alignment/rl
+```
+
+```shell
+# 启动 reward server
+python reward_server.py
+```
+
+```shell
+export PYTHONPATH=your_PaddleNLP_path/:$PYTHONPATH
+export PYTHONPATH=your_PaddleNLP_path/llm:$PYTHONPATH
+
+export FLAGS_set_to_1d=False
+export NVIDIA_TF32_OVERRIDE=0
+export FLAGS_dataloader_use_file_descriptor=False
+export HF_DATASETS_DOWNLOAD_TIMEOUT=1
+export FLAGS_gemm_use_half_precision_compute_type=False
+export FLAGS_force_cublaslt_no_reduced_precision_reduction=True
+
+export FLAGS_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+python -u -m paddle.distributed.launch --devices "0,1,2,3" run_rl.py ../../config/qwen/grpo_argument.yaml
+
+# QWEN32B 2k prompt + 30k response 9台8x80G 显卡训练命令如下：
+# python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_rl.py ../../config/qwen/grpo_32b_argument.yaml
+```
+我们提供根据上述脚本可复现的[wandb 日志](https://api.wandb.ai/links/junyu/5jiulhem)。
+
+
+### Reinforce++ 训练命令
+```shell
+cd your_PaddleNLP_path/llm/alignment/rl
+```
+
+```shell
+# 启动 reward server
+python reward_server.py
+```
+
+```shell
+export PYTHONPATH=your_PaddleNLP_path/:$PYTHONPATH
+export PYTHONPATH=your_PaddleNLP_path/llm:$PYTHONPATH
+
+export FLAGS_set_to_1d=False
+export NVIDIA_TF32_OVERRIDE=0
+export FLAGS_dataloader_use_file_descriptor=False
+export HF_DATASETS_DOWNLOAD_TIMEOUT=1
+export FLAGS_gemm_use_half_precision_compute_type=False
+export FLAGS_force_cublaslt_no_reduced_precision_reduction=True
+
+export FLAGS_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+python -u -m paddle.distributed.launch --devices "0,1,2,3" run_rl.py ../../config/qwen/reinforce_plus_plus_argument.yaml
+```
+我们提供根据上述脚本可复现的[wandb 日志](https://api.wandb.ai/links/ainlp66-netflix/injcw3ra)。
+
+### 在线监控
+在`grpo_argument.yaml`和`reinforce_plus_plus_argument.yaml`中设置的输出目录为`"logging_dir": "vdl_log"`, 可以通过以下命令查看训练过程
+```shell
+visualdl --logdir vdl_log --host 0.0.0.0
+```
+
+也支持 wandb 等多种监控，可设置`"logging_dir": "wandb"`，需要提前安装好 wandb 依赖并登录。

llm/alignment/rl/gsm8k_processor.py

@@ -0,0 +1,74 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+Preprocess the GSM8k dataset to parquet format
+"""
+
+import argparse
+import os
+import re
+
+import datasets
+
+
+def extract_solution(solution_str):
+    solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str)
+    assert solution is not None
+    final_solution = solution.group(0)
+    final_solution = final_solution.split("#### ")[1].replace(",", "")
+    return final_solution
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_dir", default="./gsm8k")
+
+    args = parser.parse_args()
+
+    data_source = "openai/gsm8k"
+
+    dataset = datasets.load_dataset(data_source, "main")
+
+    train_dataset = dataset["train"]
+    test_dataset = dataset["test"]
+
+    instruction_following = 'Let\'s think step by step and output the final answer after "####".'
+
+    # add a row to each data item that represents a unique id
+    def make_map_fn(split):
+        def process_fn(example, idx):
+            question_raw = "<|im_start|>user\n" + example.pop("question")
+
+            system_raw = (
+                "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n"
+            )
+            question = system_raw + question_raw + " " + instruction_following + "<|im_end|>\n<|im_start|>assistant\n"
+
+            answer_raw = example.pop("answer")
+            solution = extract_solution(answer_raw)
+            data = {
+                "src": question,
+                "tgt": solution,
+            }
+            return data
+
+        return process_fn
+
+    train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True)
+    test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True)
+
+    local_dir = args.local_dir
+
+    train_dataset.to_json(os.path.join(local_dir, "train.jsonl"), orient="records", lines=True)
+    test_dataset.to_json(os.path.join(local_dir, "test.jsonl"), orient="records", lines=True)

llm/alignment/rl/reward/fake_reward_server.py

@@ -0,0 +1,112 @@
+# 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.
+
+"""Launch Reward HTTP Server."""
+
+import argparse
+import json
+import logging
+import threading
+import traceback
+from typing import List
+
+import uvicorn
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+
+class Request(BaseModel):
+    """The request for RM server."""
+
+    src: List[str]
+    tgt: List[str]
+    response: List[str]
+
+
+class Response(BaseModel):
+    """The response for RM server."""
+
+    error_code: int = 0
+    error_msg: str = "Success"
+    score: List[float] = None
+
+
+def compute_score(
+    solution_str: str, ground_truth: str, query=None, format_reward: int = 1, answer_reward: float = 1.0
+):
+    score = float(1.0)
+    print(
+        f"==============================================================={ground_truth}=========================================================================="
+    )
+    print(f"score {score}, solution_str\n", solution_str)
+    print(
+        "================================================================================================================================================="
+    )
+    return score
+
+
+def setup_args():
+    """Setup inerance server arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8731)
+    parser.add_argument("--log_file", type=str, default="rm_server.log")
+    args = parser.parse_args()
+    return args
+
+
+def server(args):
+    """Launch RM server."""
+    app = FastAPI()
+    lock = threading.Lock()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        filename=args.log_file,
+        filemode="w",
+        format="%(asctime)s - %(message)s",
+    )
+
+    @app.post("/")
+    async def _server(request: Request) -> Response:
+        lock.acquire()
+        logging.info(f"Request: {request}")
+        try:
+            all_result = []
+            if len(request.tgt) != len(request.response) or len(request.tgt) != len(request.src):
+                raise ValueError("The length of response, tgt, and src should be equal.")
+            for i in range(len(request.response)):
+                reward = compute_score(request.response[i], request.tgt[i], request.src[i])
+                all_result.append(reward)
+            output = {
+                "error_code": 0,
+                "error_msg": "Success",
+                "score": all_result,
+            }
+        except Exception as err:
+            logging.error(f"Server error: when process {request}\n{traceback.format_stack()}")
+            output = {
+                "error_code": 500,
+                "error_msg": f"{err}",
+                "score": [0] * len(request.tgt),
+            }
+        logging.info(f"Response: {json.dumps(output, indent=2, ensure_ascii=False)}")
+        lock.release()
+        return output
+
+    uvicorn.run(app, host="0.0.0.0", port=args.port)
+
+
+if __name__ == "__main__":
+    args = setup_args()
+    server(args)

llm/alignment/rl/run_rl.py

@@ -23,6 +23,7 @@ import paddle
 from paddle.distributed import fleet
 
 from paddlenlp.datasets.rlhf_datasets import RLHFDataset, collate_fn
+from paddlenlp.generation import GenerationConfig
 from paddlenlp.rl.models.score_model import AutoModelForScore
 from paddlenlp.rl.trainer.ppo_trainer import PPOTrainer
 from paddlenlp.rl.utils.config_utils import (
@@ -31,7 +32,7 @@ from paddlenlp.rl.utils.config_utils import (
     TrainingArguments,
 )
 from paddlenlp.rl.utils.offload_utils import offload_tensor_to_cpu
-from paddlenlp.rl.utils.reshard_utils import init_rollout_env
+from paddlenlp.rl.utils.reshard_utils import ReshardController
 from paddlenlp.rl.utils.timer_utils import timers_scope_runtimer
 from paddlenlp.trainer import (
     EarlyStoppingCallback,
@@ -41,9 +42,11 @@ from paddlenlp.trainer import (
 from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForTokenClassification,
     AutoTokenizer,
     PretrainedConfig,
 )
+from paddlenlp.transformers.configuration_utils import LlmMetaConfig
 from paddlenlp.trl import llm_utils
 from paddlenlp.utils.log import logger
 
@@ -79,6 +82,7 @@ def create_actor_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
+    reshard_controller: ReshardController = None,
 ):
     with timers_scope_runtimer("Actor model loading time"):
         # actor model
@@ -87,22 +91,21 @@ def create_actor_models(
             tensor_parallel_output=training_args.tensor_parallel_output,
             tensor_parallel_degree=training_args.tensor_parallel_degree,
             tensor_parallel_rank=training_args.tensor_parallel_rank,
-            recompute_granularity=model_args.recompute_granularity,
+            recompute_granularity=training_args.recompute_granularity,
             dtype=training_args.model_dtype,
             recompute=training_args.recompute,
             recompute_use_reentrant=training_args.recompute_use_reentrant,
             **common_config,
         )
+        LlmMetaConfig.set_llm_config(actor_model_config, training_args)
 
         actor_model_config.use_fused_head_and_loss_fn = training_args.use_fused_head_and_loss_fn
         actor_model_config.set_attn_func = True
         actor_model_config.max_position_embeddings = data_args.max_length
         actor_model_config.use_sparse_head_and_loss_fn = False
-        actor_model_config.fused_linear = model_args.fused_linear
-        actor_model_config.use_fused_rms_norm = training_args.use_fused_rms_norm
         actor_model_config.seq_length = data_args.max_length
         actor_model_config.max_sequence_length = data_args.max_length
-        print(f"Loading Actor model with config:\n\t{actor_model_config}\n")
+        logger.info(f"Loading Actor model with config:\n\t{actor_model_config}\n")
 
         if not training_args.autotuner_benchmark:
             actor_model = AutoModelForCausalLM.from_pretrained(
@@ -112,18 +115,16 @@ def create_actor_models(
             actor_model = AutoModelForCausalLM.from_config(actor_model_config)
 
     with timers_scope_runtimer("Actor eval model loading time"):
-        if (
-            training_args.rollout_tensor_parallel_degree != training_args.tensor_parallel_degree
-            or training_args.pipeline_parallel_degree > 1
-        ):
+        if reshard_controller is not None:
+            reshard_controller.set_rollout_env("[create actor eval model]")
             actor_eval_model_config = copy.deepcopy(actor_model_config)
             actor_eval_model_config.use_fused_head_and_loss_fn = False
-            with init_rollout_env(training_args.rollout_tensor_parallel_degree):
-                hcg = fleet.get_hybrid_communicate_group()
-                actor_eval_model_config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
-                actor_eval_model_config.tensor_parallel_rank = hcg.get_model_parallel_rank()
-                # TODO(gongenlei): lazy load lazy guard
-                actor_eval_model = AutoModelForCausalLM.from_config(actor_eval_model_config)
+            hcg = fleet.get_hybrid_communicate_group()
+            actor_eval_model_config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
+            actor_eval_model_config.tensor_parallel_rank = hcg.get_model_parallel_rank()
+            # TODO(gongenlei): lazy load lazy guard
+            actor_eval_model = AutoModelForCausalLM.from_config(actor_eval_model_config)
+            reshard_controller.set_train_env("[after create actor eval model]")
         else:
             actor_eval_model = None
 
@@ -134,7 +135,6 @@ def create_actor_models(
         )
         if not training_args.autotuner_benchmark:
             reference_model.set_state_dict(actor_model.state_dict())
-
     actor_tokenizer = AutoTokenizer.from_pretrained(
         model_args.actor_model_name_or_path,
         model_max_length=data_args.max_length,
@@ -167,10 +167,10 @@ def create_reward_models(
             recompute_use_reentrant=training_args.recompute_use_reentrant,
             **common_config,
         )
+        LlmMetaConfig.set_llm_config(reward_model_config, training_args)
         reward_model_config.max_position_embeddings = data_args.max_length
         reward_model_config.use_sparse_head_and_loss_fn = False
-        reward_model_config.fused_linear = model_args.fused_linear
-        print(f"Loading Reward model with config:\n\t{reward_model_config}\n")
+        logger.info(f"Loading Reward model with config:\n\t{reward_model_config}\n")
 
         config = copy.deepcopy(reward_model_config)
         if training_args.eval_mode is not None:
@@ -210,46 +210,43 @@ def create_critic_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
-    reward_model,
 ):
     with timers_scope_runtimer("Critic model loading time"):
-        reward_model_config = reward_model.config
-        if model_args.critic_model_name_or_path is None:
-            model_args.critic_model_name_or_path = model_args.reward_model_name_or_path
-            critic_model = AutoModelForScore.from_config(
-                reward_model_config,
-                dtype=training_args.model_dtype,
-                score_type="critic",
-                do_normalize=False,
-                clip_range_value=training_args.clip_range_value,
-                **common_config,
+        critic_model_config = AutoConfig.from_pretrained(
+            model_args.critic_model_name_or_path,
+            tensor_parallel_output=training_args.tensor_parallel_output,
+            tensor_parallel_degree=training_args.tensor_parallel_degree,
+            tensor_parallel_rank=training_args.tensor_parallel_rank,
+            dtype=training_args.model_dtype,
+            recompute=training_args.critic_recompute,
+            recompute_granularity=model_args.critic_recompute_granularity,
+            recompute_use_reentrant=training_args.recompute_use_reentrant,
+            **common_config,
+        )
+        LlmMetaConfig.set_llm_config(critic_model_config, training_args)
+
+        critic_model_config.max_position_embeddings = data_args.max_length
+        critic_model_config.use_sparse_head_and_loss_fn = False
+        critic_model_config.num_labels = 1
+        critic_model_config.classifier_dropout = 0.0
+        critic_model_config.hidden_dropout = 0.0
+        logger.info(f"Loading Critic model with config:\n\t{critic_model_config}\n")
+
+        if not training_args.autotuner_benchmark:
+            critic_model = AutoModelForTokenClassification.from_pretrained(
+                model_args.critic_model_name_or_path,
+                config=critic_model_config,
             )
-            if not training_args.autotuner_benchmark:
-                critic_model.set_state_dict(reward_model.state_dict())
         else:
-            if not training_args.autotuner_benchmark:
-                critic_model = AutoModelForScore.from_pretrained(
-                    model_args.critic_model_name_or_path,
-                    config=reward_model_config,
-                    score_type="critic",
-                    do_normalize=False,
-                    clip_range_value=training_args.clip_range_value,
-                    **common_config,
-                )
-            else:
-                critic_model = AutoModelForScore.from_config(
-                    reward_model_config,
-                    score_type="critic",
-                    do_normalize=False,
-                    clip_range_value=training_args.clip_range_value,
-                    **common_config,
-                )
+            critic_model = AutoModelForTokenClassification.from_config(
+                critic_model_config,
+            )
 
     critic_tokenizer = AutoTokenizer.from_pretrained(
         model_args.critic_model_name_or_path,
         model_max_length=data_args.max_length,
         padding_side="left",
-        tokenizer_alpha=model_args.reward_critic_tokenizer_alpha,
+        tokenizer_alpha=model_args.critic_tokenizer_alpha,
         use_fast=True,
     )
     if critic_tokenizer.pad_token_id is None:
@@ -261,8 +258,8 @@ def create_critic_models(
         if training_args.eval_mode == "single":
             config.tensor_parallel_degree = -1
             config.tensor_parallel_rank = 0
-        with timers_scope_runtimer("Reward critic eval model loading time"):
-            critic_eval_model = AutoModelForScore.from_config(config)
+        with timers_scope_runtimer("Critic eval model loading time"):
+            critic_eval_model = AutoModelForTokenClassification.from_config(config)
     else:
         critic_eval_model = None
 
@@ -270,13 +267,14 @@ def create_critic_models(
 
 
 def create_rl_dataset(data_args, training_args, tokenizer):
-    requires_label = True if training_args.use_rm_server else False
+    requires_label = True if training_args.use_rm_server or training_args.use_rule_reward else False
     train_ds = RLHFDataset(
         dataset_name_or_path=data_args.train_datasets,
         tokenizer=tokenizer,
         max_prompt_len=data_args.max_prompt_len,
         requires_label=requires_label,
-        label_key=data_args.label_key,
+        prompt_key=data_args.prompt_key,
+        response_key=data_args.response_key,
         splits="train",
     )
     dev_ds = RLHFDataset(
@@ -284,7 +282,8 @@ def create_rl_dataset(data_args, training_args, tokenizer):
         tokenizer=tokenizer,
         max_prompt_len=data_args.max_prompt_len,
         requires_label=requires_label,
-        label_key=data_args.label_key,
+        prompt_key=data_args.prompt_key,
+        response_key=data_args.response_key,
         splits="dev",
     )
     return train_ds, dev_ds
@@ -295,6 +294,8 @@ def main():
     parser = PdArgumentParser((ModelArgument, DataArgument, TrainingArguments))
     if len(sys.argv) >= 2 and sys.argv[1].endswith(".json"):
         model_args, data_args, training_args = parser.parse_json_file_and_cmd_lines()
+    elif len(sys.argv) >= 2 and sys.argv[1].endswith(".yaml"):
+        model_args, data_args, training_args = parser.parse_yaml_file_and_cmd_lines()
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
@@ -312,24 +313,33 @@ def main():
             )
 
     common_config = dict(
-        use_flash_attention=model_args.use_flash_attention,
+        use_flash_attention=training_args.use_flash_attention,
         sequence_parallel=training_args.sequence_parallel,
         fused_rotary=False,
         max_sequence_length=data_args.max_length,
     )
 
-    actor_model, actor_eval_model, reference_model, actor_tokenizer = create_actor_models(
-        model_args, data_args, training_args, common_config
-    )
+    if (
+        training_args.rollout_tensor_parallel_degree != training_args.tensor_parallel_degree
+        or training_args.pipeline_parallel_degree > 1
+    ):
+        reshard_controller = ReshardController(tensor_parallel_degree=training_args.rollout_tensor_parallel_degree)
+    else:
+        reshard_controller = None
 
-    if not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
+    actor_model, actor_eval_model, reference_model, actor_tokenizer = create_actor_models(
+        model_args, data_args, training_args, common_config, reshard_controller
+    )
+    if training_args.use_rule_reward:
+        reward_model, reward_tokenizer = None, actor_tokenizer
+    elif not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
         reward_model, reward_tokenizer = create_reward_models(model_args, data_args, training_args, common_config)
     else:
         reward_model, reward_tokenizer = model_args.reward_server, actor_tokenizer
 
     if training_args.rl_algorithm == "ppo":
         critic_model, critic_eval_model, critic_tokenizer = create_critic_models(
-            model_args, data_args, training_args, common_config, reward_model
+            model_args, data_args, training_args, common_config
         )
     else:
         critic_model, critic_eval_model, critic_tokenizer = None, None, None
@@ -343,7 +353,8 @@ def main():
         offload_tensor_to_cpu((reference_model, "freeze_model"))
 
         if training_args.rl_algorithm == "ppo":
-            offload_tensor_to_cpu((reward_model, "freeze_model"))
+            if not training_args.use_rm_server and not training_args.use_rule_reward:
+                offload_tensor_to_cpu((reward_model, "freeze_model"))
             if critic_eval_model is not None:
                 offload_tensor_to_cpu((critic_eval_model, "freeze_model"))
 
@@ -351,9 +362,22 @@ def main():
         paddle.device.cuda.empty_cache()
 
     def compute_metrics(eval_preds):
-        accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
+        '''
+        If "use_rm_server" is TRUE, the score ranges from -3 to 3, with 3 being the only correct score (format + result).
+        If using the "Regularized Matching Function (use_rule_reward=True)" (currently only implemented for the gsm8k dataset), the score ranges from 0 to 1.
+        '''
+        if training_args.use_rule_reward:
+            accuracy = (eval_preds.predictions == 1).astype("float32").mean().item()
+        else:
+            accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
         return {"accuracy": accuracy}
 
+    try:
+        generation_config = GenerationConfig.from_pretrained(model_args.actor_model_name_or_path)
+    except:
+        logger.warning("Can't find generation config, so it will not use generation_config field in the model config")
+        generation_config = None
+
     trainer = PPOTrainer(
         actor_model=actor_model,
         reference_model=reference_model,
@@ -371,9 +395,12 @@ def main():
         data_collator=partial(
             collate_fn,
             pad_token_id=actor_tokenizer.pad_token_id,
-            requires_label=True if training_args.use_rm_server else False,
-        ),
+            requires_label=True if training_args.use_rm_server or training_args.use_rule_reward else False,
+            max_prompt_len=data_args.max_prompt_len if training_args.balance_batch else None,
+        ),  # NOTE: enforce prompt padding to max_prompt_len when using balance_batch
         compute_metrics=compute_metrics,  # TODO: only used for grpo (kk datasets)
+        generation_config=generation_config,
+        reshard_controller=reshard_controller,
     )
 
     # TODO(gongenlei) resume_from_checkpoint is not ready

llm/application/distill/README.md

@@ -195,26 +195,26 @@ def process_data_en(example):
     return {"src": src, "tgt": tgt}
 
 # construct Chinese sft dataset
-paddlenlp_datatset = deepcopy(dataset)
-paddlenlp_datatset["train"] = paddlenlp_datatset["train"].map(
-    process_data_zh, remove_columns=paddlenlp_datatset["train"].column_names
+paddlenlp_dataset = deepcopy(dataset)
+paddlenlp_dataset["train"] = paddlenlp_dataset["train"].map(
+    process_data_zh, remove_columns=paddlenlp_dataset["train"].column_names
 )
-paddlenlp_datatset["test"] = paddlenlp_datatset["test"].map(
-    process_data_zh, remove_columns=paddlenlp_datatset["test"].column_names
+paddlenlp_dataset["test"] = paddlenlp_dataset["test"].map(
+    process_data_zh, remove_columns=paddlenlp_dataset["test"].column_names
 )
-paddlenlp_datatset["train"].to_json("data/gsm8k_distilled_zh_sft/train.json", force_ascii=False)
-paddlenlp_datatset["test"].to_json("data/gsm8k_distilled_zh_sft/dev.json", force_ascii=False)
+paddlenlp_dataset["train"].to_json("data/gsm8k_distilled_zh_sft/train.json", force_ascii=False)
+paddlenlp_dataset["test"].to_json("data/gsm8k_distilled_zh_sft/dev.json", force_ascii=False)
 
 # construct English sft dataset
-paddlenlp_datatset = deepcopy(dataset)
-paddlenlp_datatset["train"] = paddlenlp_datatset["train"].map(
-    process_data_en, remove_columns=paddlenlp_datatset["train"].column_names
+paddlenlp_dataset = deepcopy(dataset)
+paddlenlp_dataset["train"] = paddlenlp_dataset["train"].map(
+    process_data_en, remove_columns=paddlenlp_dataset["train"].column_names
 )
-paddlenlp_datatset["test"] = paddlenlp_datatset["test"].map(
-    process_data_en, remove_columns=paddlenlp_datatset["test"].column_names
+paddlenlp_dataset["test"] = paddlenlp_dataset["test"].map(
+    process_data_en, remove_columns=paddlenlp_dataset["test"].column_names
 )
-paddlenlp_datatset["train"].to_json("data/gsm8k_distilled_en_sft/train.json", force_ascii=False)
-paddlenlp_datatset["test"].to_json("data/gsm8k_distilled_en_sft/dev.json", force_ascii=False)
+paddlenlp_dataset["train"].to_json("data/gsm8k_distilled_en_sft/train.json", force_ascii=False)
+paddlenlp_dataset["test"].to_json("data/gsm8k_distilled_en_sft/dev.json", force_ascii=False)
 
 ```
 最终我们将会得到如下字段的数据集

llm/application/distill/prepare_data.py

@@ -47,23 +47,23 @@ def process_data_en(example):
 
 
 # construct Chinese sft dataset
-paddlenlp_datatset = deepcopy(dataset)
-paddlenlp_datatset["train"] = paddlenlp_datatset["train"].map(
-    process_data_zh, remove_columns=paddlenlp_datatset["train"].column_names
+paddlenlp_dataset = deepcopy(dataset)
+paddlenlp_dataset["train"] = paddlenlp_dataset["train"].map(
+    process_data_zh, remove_columns=paddlenlp_dataset["train"].column_names
 )
-paddlenlp_datatset["test"] = paddlenlp_datatset["test"].map(
-    process_data_zh, remove_columns=paddlenlp_datatset["test"].column_names
+paddlenlp_dataset["test"] = paddlenlp_dataset["test"].map(
+    process_data_zh, remove_columns=paddlenlp_dataset["test"].column_names
 )
-paddlenlp_datatset["train"].to_json("data/gsm8k_distilled_zh_sft/train.json", force_ascii=False)
-paddlenlp_datatset["test"].to_json("data/gsm8k_distilled_zh_sft/dev.json", force_ascii=False)
+paddlenlp_dataset["train"].to_json("data/gsm8k_distilled_zh_sft/train.json", force_ascii=False)
+paddlenlp_dataset["test"].to_json("data/gsm8k_distilled_zh_sft/dev.json", force_ascii=False)
 
 # construct English sft dataset
-paddlenlp_datatset = deepcopy(dataset)
-paddlenlp_datatset["train"] = paddlenlp_datatset["train"].map(
-    process_data_en, remove_columns=paddlenlp_datatset["train"].column_names
+paddlenlp_dataset = deepcopy(dataset)
+paddlenlp_dataset["train"] = paddlenlp_dataset["train"].map(
+    process_data_en, remove_columns=paddlenlp_dataset["train"].column_names
 )
-paddlenlp_datatset["test"] = paddlenlp_datatset["test"].map(
-    process_data_en, remove_columns=paddlenlp_datatset["test"].column_names
+paddlenlp_dataset["test"] = paddlenlp_dataset["test"].map(
+    process_data_en, remove_columns=paddlenlp_dataset["test"].column_names
 )
-paddlenlp_datatset["train"].to_json("data/gsm8k_distilled_en_sft/train.json", force_ascii=False)
-paddlenlp_datatset["test"].to_json("data/gsm8k_distilled_en_sft/dev.json", force_ascii=False)
+paddlenlp_dataset["train"].to_json("data/gsm8k_distilled_en_sft/train.json", force_ascii=False)
+paddlenlp_dataset["test"].to_json("data/gsm8k_distilled_en_sft/dev.json", force_ascii=False)