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

@@ -17,13 +17,349 @@
 
 namespace py = pybind11;
 
+std::vector<paddle::Tensor> AppendAttention(
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_batch_ids,
+    const paddle::Tensor& encoder_tile_ids_per_batch,
+    const paddle::Tensor& encoder_num_blocks,
+    const paddle::Tensor& kv_batch_ids,
+    const paddle::Tensor& kv_tile_ids_per_batch,
+    const paddle::Tensor& kv_num_blocks,
+    const paddle::Tensor& decoder_batch_ids,
+    const paddle::Tensor& decoder_tile_ids_per_batch,
+    const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& max_len_kv,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& qkv_bias,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    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,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool speculate_decoder);
+
+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> MultiHeadLatentAttention(
+    const paddle::Tensor& query,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_batch_ids,
+    const paddle::Tensor& encoder_tile_ids_per_batch,
+    const paddle::Tensor& encoder_num_blocks,
+    const paddle::Tensor& kv_batch_ids,
+    const paddle::Tensor& kv_tile_ids_per_batch,
+    const paddle::Tensor& kv_num_blocks,
+    const paddle::Tensor& decoder_batch_ids,
+    const paddle::Tensor& decoder_tile_ids_per_batch,
+    const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& max_len_kv,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& query_bias,
+    const paddle::optional<paddle::Tensor>& query_out_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    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 std::string& compute_dtype,
+    const std::string& cache_quant_type_str,
+    const int nope_size,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_draft_total_token_num,
+    const bool causal,
+    const bool speculate_decoder);
+
+std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cum_offsets,
+    const int group_size,
+    const int block_size,
+    const int decoder_step_token_num);
+
+std::vector<paddle::Tensor> NoauxTc(paddle::Tensor& scores,
+                                    paddle::Tensor& scores_with_bias,
+                                    int n_group,
+                                    int topk_group,
+                                    int topk,
+                                    float routed_scaling_factor);
+
+std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_decoder,
+    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);
+
+std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    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 bool speculate_decoder);
+
+paddle::Tensor cutlass_fp8_fp8_half_block_gemm_fused_func(
+    const paddle::Tensor& x,
+    const paddle::Tensor& y,
+    const paddle::Tensor& x_scale,
+    const paddle::Tensor& y_scale,
+    const paddle::optional<paddle::Tensor>& bias,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype,
+    std::string activation_type);
+
+void GetPositionIdsAndMaskEncoderBatch(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& position_ids,
+    const paddle::Tensor& mask_encoder_batch);
+
+void SetPreidsTokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
+                                      const paddle::Tensor& input_ids,
+                                      const paddle::Tensor& seq_lens_encoder,
+                                      const paddle::Tensor& seq_lens_decoder,
+                                      const paddle::Tensor& step_idx,
+                                      const paddle::Tensor& stop_flags,
+                                      const paddle::Tensor& logits,
+                                      const paddle::Tensor& penalty_scores,
+                                      const paddle::Tensor& frequency_scores,
+                                      const paddle::Tensor& presence_scores,
+                                      const paddle::Tensor& temperatures,
+                                      const paddle::Tensor& bad_tokens,
+                                      const paddle::Tensor& cur_len,
+                                      const paddle::Tensor& min_len,
+                                      const paddle::Tensor& eos_token_id);
+
+void UpdateInputesV2(const paddle::Tensor& stop_flags,
+               const paddle::Tensor& step_idx,
+               const paddle::Tensor& not_need_stop, // cpu
+               const paddle::Tensor& seq_lens_this_time,
+               const paddle::Tensor& seq_lens_encoder,
+               const paddle::Tensor& seq_lens_decoder,
+               const paddle::Tensor& max_dec_len,
+               const paddle::Tensor& input_ids,
+               const paddle::Tensor& stop_nums,
+               const paddle::Tensor& next_tokens,
+               const paddle::Tensor& is_block_step,
+               const paddle::Tensor& end_ids,
+               const paddle::Tensor& kwargs_next_tokens);
+
+paddle::Tensor RebuildPaddingV2Func(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::optional<paddle::Tensor>& output_padding_offset,
+                                    int max_input_length);
+
+std::vector<paddle::Tensor> PerTokenGroupQuant(const paddle::Tensor& x,
+                                        const int group_size,
+                                        const bool transpose_scale,
+                                        const float quant_max_bound,
+                                        const float quant_min_bound);
+
+std::vector<paddle::Tensor> PerTensorQuantFp8(const paddle::Tensor& x, const paddle::optional<paddle::Tensor>& scale);
+
+std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
+                                               const paddle::Tensor& cum_offsets,
+                                               const paddle::Tensor& token_num,
+                                               const paddle::Tensor& seq_len,
+                                               const paddle::optional<paddle::Tensor>& draft_tokens,
+                                               const paddle::optional<paddle::Tensor>& seq_lens_encoder);
+
+void SaveOutMmsg(const paddle::Tensor& x,
+                 const paddle::Tensor& not_need_stop, // cpu
+                 int64_t rank_id);
+
+void GetOutput(const paddle::Tensor& x,
+               int64_t rank_id,
+               bool wait_flag);
+
+void StepPaddle(const paddle::Tensor &stop_flags,
+                const paddle::Tensor &seq_lens_this_time,
+                const paddle::Tensor &ori_seq_lens_encoder,
+                const paddle::Tensor &seq_lens_encoder,
+                const paddle::Tensor &seq_lens_decoder,
+                const paddle::Tensor &block_tables,  // [bsz, block_num_per_seq]
+                const paddle::Tensor &encoder_block_lens,
+                const paddle::Tensor &is_block_step,
+                const paddle::Tensor &step_block_list,
+                const paddle::Tensor &step_lens,
+                const paddle::Tensor &recover_block_list,
+                const paddle::Tensor &recover_lens,
+                const paddle::Tensor &need_block_list,
+                const paddle::Tensor &need_block_len,
+                const paddle::Tensor &used_list_len,
+                const paddle::Tensor &free_list,
+                const paddle::Tensor &free_list_len,
+                const paddle::Tensor &input_ids,
+                const paddle::Tensor &pre_ids,
+                const paddle::Tensor &step_idx,
+                const paddle::Tensor &next_tokens,
+                const paddle::Tensor &first_token_ids,
+                const int block_size,
+                const int encoder_decoder_block_num);
+
+void SaveOutputDygraph(
+    const paddle::Tensor& all_token_ids,
+    const paddle::Tensor& tokens,
+    const paddle::Tensor& result_ids,
+    const paddle::Tensor& step_idx
+);
+
 PYBIND11_MODULE(paddlenlp_ops, m) {
-    /**
-     * all_reduce.cu
-     */
-    m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
-    m.def("all_reduce", &all_reduce, "all reduce function");
-    m.def("dispose", &dispose, "del function for python");
-    m.def("meta_size", &meta_size, "meta_size function for Signal struct");
-    m.def("register_buffer", &register_buffer, "register ipc buffer");
+  /**
+   * all_reduce.cu
+   */
+  m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
+  m.def("all_reduce", &all_reduce, "all reduce function");
+  m.def("dispose", &dispose, "del function for python");
+  m.def("meta_size", &meta_size, "meta_size function for Signal struct");
+  m.def("register_buffer", &register_buffer, "register ipc buffer");
+  m.def("f_append_attention", &AppendAttention, "AppendAttention");
+  m.def("f_fused_rotary_position_encoding", &FusedRotaryPositionEncoding, "FusedRotaryPositionEncoding");
+  m.def("f_multi_head_latent_attention", &MultiHeadLatentAttention, "MultiHeadLatentAttention");
+  m.def("f_noaux_tc", &NoauxTc, "NoauxTc");
+  m.def("f_get_block_shape_and_split_kv_block", &GetBlockShapeAndSplitKVBlock, "GetBlockShapeAndSplitKVBlock");
+  m.def("f_prefill_mla_write_cache", &PrefillMLAWriteCacheKernel, "PrefillMLAWriteCacheKernel");
+  m.def("f_decode_mla_write_cache", &DecodeMLAWriteCacheKernel, "DecodeMLAWriteCacheKernel");
+  m.def("f_get_position_ids_and_mask_encoder_batch", &GetPositionIdsAndMaskEncoderBatch, "GetPositionIdsAndMaskEncoderBatch");
+  m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
+  m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
+  m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
+  m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
+  m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
+  m.def("f_get_output", &GetOutput, "GetOutput");
+  m.def("f_step_paddle", &StepPaddle, "StepPaddle");
+  m.def("f_save_output_dygraph", &SaveOutputDygraph, "SaveOutputDygraph");
+//   m.def("f_cutlass_fp8_fp8_half_block_gemm_fused", &cutlass_fp8_fp8_half_block_gemm_fused_func, "cutlass_fp8_fp8_half_block_gemm_fused_func");
 }
+
+PYBIND11_MODULE(paddlenlp_ops_80, m) {
+  /**
+   * all_reduce.cu
+   */
+  m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
+  m.def("all_reduce", &all_reduce, "all reduce function");
+  m.def("dispose", &dispose, "del function for python");
+  m.def("meta_size", &meta_size, "meta_size function for Signal struct");
+  m.def("register_buffer", &register_buffer, "register ipc buffer");
+  m.def("f_append_attention", &AppendAttention, "AppendAttention");
+  m.def("f_fused_rotary_position_encoding", &FusedRotaryPositionEncoding, "FusedRotaryPositionEncoding");
+  m.def("f_multi_head_latent_attention", &MultiHeadLatentAttention, "MultiHeadLatentAttention");
+  m.def("f_noaux_tc", &NoauxTc, "NoauxTc");
+  m.def("f_get_block_shape_and_split_kv_block", &GetBlockShapeAndSplitKVBlock, "GetBlockShapeAndSplitKVBlock");
+  m.def("f_prefill_mla_write_cache", &PrefillMLAWriteCacheKernel, "PrefillMLAWriteCacheKernel");
+  m.def("f_decode_mla_write_cache", &DecodeMLAWriteCacheKernel, "DecodeMLAWriteCacheKernel");
+  m.def("f_get_position_ids_and_mask_encoder_batch", &GetPositionIdsAndMaskEncoderBatch, "GetPositionIdsAndMaskEncoderBatch");
+  m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
+  m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
+  m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
+  m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
+  m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
+  m.def("f_get_output", &GetOutput, "GetOutput");
+  m.def("f_step_paddle", &StepPaddle, "StepPaddle");
+  m.def("f_save_output_dygraph", &SaveOutputDygraph, "SaveOutputDygraph");
+}
+
+PYBIND11_MODULE(paddlenlp_ops_90, m) {
+  /**
+   * all_reduce.cu
+   */
+  m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
+  m.def("all_reduce", &all_reduce, "all reduce function");
+  m.def("dispose", &dispose, "del function for python");
+  m.def("meta_size", &meta_size, "meta_size function for Signal struct");
+  m.def("register_buffer", &register_buffer, "register ipc buffer");
+  m.def("f_append_attention", &AppendAttention, "AppendAttention");
+  m.def("f_fused_rotary_position_encoding", &FusedRotaryPositionEncoding, "FusedRotaryPositionEncoding");
+  m.def("f_multi_head_latent_attention", &MultiHeadLatentAttention, "MultiHeadLatentAttention");
+  m.def("f_noaux_tc", &NoauxTc, "NoauxTc");
+  m.def("f_get_block_shape_and_split_kv_block", &GetBlockShapeAndSplitKVBlock, "GetBlockShapeAndSplitKVBlock");
+  m.def("f_prefill_mla_write_cache", &PrefillMLAWriteCacheKernel, "PrefillMLAWriteCacheKernel");
+  m.def("f_decode_mla_write_cache", &DecodeMLAWriteCacheKernel, "DecodeMLAWriteCacheKernel");
+  m.def("f_get_position_ids_and_mask_encoder_batch", &GetPositionIdsAndMaskEncoderBatch, "GetPositionIdsAndMaskEncoderBatch");
+  m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
+  m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
+  m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
+  m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
+  m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
+  m.def("f_get_output", &GetOutput, "GetOutput");
+  m.def("f_step_paddle", &StepPaddle, "StepPaddle");
+  m.def("f_save_output_dygraph", &SaveOutputDygraph, "SaveOutputDygraph");
+}

csrc/gpu/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu

@@ -18,7 +18,7 @@
 #include "fp8_gemm_fused/fuse_block_gemm_act_template_3x.h"
 #include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
 
-std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
+paddle::Tensor cutlass_fp8_fp8_half_block_gemm_fused_func(
     const paddle::Tensor& x,
     const paddle::Tensor& y,
     const paddle::Tensor& x_scale,
@@ -145,29 +145,20 @@ std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
                                     x_scale_ptr,
                                     y_scale_ptr};
     fp8_fp8_block_gemm_scale_bias_act(params);
-  // if(output_dtype == "bfloat16"){
-  //   dispatch_fuse_block_gemm_c3x<phi::dtype::float8_e4m3fn,
-  //                             phi::dtype::bfloat16,
-  //                             false,
-  //                             cutlass::epilogue::thread::Identity,
-  //                             cute::Shape<_128, _128, _128>,
-  //                             cute::Shape<_1, _2, _1>,
-  //                             cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8GroupBlockScaledAccum<1>,
-  //                             cutlass::epilogue::TmaWarpSpecializedCooperative,
-  //                             cutlass::arch::Sm90>(params);
-  // }else{
-  //   dispatch_fuse_block_gemm_c3x<phi::dtype::float8_e4m3fn,
-  //                             phi::dtype::float16,
-  //                             false,
-  //                             cutlass::epilogue::thread::Identity,
-  //                             cute::Shape<_128, _128, _128>,
-  //                             cute::Shape<_1, _2, _1>,
-  //                             cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8GroupBlockScaledAccum<1>,
-  //                             cutlass::epilogue::TmaWarpSpecializedCooperative,
-  //                             cutlass::arch::Sm90>(params);
-  // }
-  
-  return {out};
+  return out;
+}
+
+std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
+    const paddle::Tensor& x,
+    const paddle::Tensor& y,
+    const paddle::Tensor& x_scale,
+    const paddle::Tensor& y_scale,
+    const paddle::optional<paddle::Tensor>& bias,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype,
+    std::string activation_type) {
+  return {cutlass_fp8_fp8_half_block_gemm_fused_func(x,y,x_scale,y_scale,bias,trans_x,trans_y,output_dtype,activation_type)};
 }
 
 std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfBlockGemmFusedInferShape(

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;
@@ -32,8 +33,20 @@ void GetOutput(const paddle::Tensor& x,
   if (rank_id > 0) return;
 
   static struct msgdata msg_rcv;
-
-  static key_t key = ftok("./", 1);
+  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);
 

csrc/gpu/group_quant.cu

@@ -1,201 +0,0 @@
-// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-// 
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-// 
-//     http://www.apache.org/licenses/LICENSE-2.0
-// 
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#include "helper.h"
-#include<string.h>
-#include <cuda_runtime.h>
-
-template <typename T>
-__inline__ __device__ T WarpReduceAbsMax(T val, unsigned lane_mask) {
-  #pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1){
-    val = max(val, __shfl_xor_sync(lane_mask, val, mask, 32));
-  }
-  return val;
-}
-
-template <typename InType, typename OutType, int GroupSize, int VecSize, bool transpose_scale>
-__global__ void GroupQuantKernel(const InType* input,
-                                   const int64_t numel,
-                                   const int scale_rows,
-                                   const int scale_cols,
-                                   const float quant_max_bound,
-                                   const float quant_min_bound,
-                                   OutType* output,
-                                   float* out_scale_data) {
-
-  const int idx = (blockIdx.x * GroupSize) + threadIdx.x * VecSize;
-
-  if (idx >= numel) return;
-
-  const int scale_row_idx = blockIdx.x / scale_cols;
-  const int scale_col_idx = blockIdx.x % scale_cols;
-  int scale_idx = 0;
-  float abs_max_val = 0.000001;
-
-  if constexpr (transpose_scale) {
-    scale_idx = scale_col_idx * scale_rows + scale_row_idx;
-  }else{
-    scale_idx = scale_row_idx * scale_cols + scale_col_idx;
-  }
-
-  AlignedVector<InType, VecSize> in_vec;
-  AlignedVector<OutType, VecSize> out_vec;
-
-  Load<InType, VecSize>(input + idx, &in_vec);
-  #pragma unroll
-  for (int i = 0; i < VecSize; ++i) {
-    abs_max_val = max(abs_max_val, abs(static_cast<float>(in_vec[i])));
-  }
-  
-  abs_max_val = WarpReduceAbsMax(abs_max_val, 0xffffffff);
-  __syncthreads();
-  float scale = __shfl_sync(0xffffffff, abs_max_val, 0) / quant_max_bound;
-  
-  #pragma unroll
-  for (int i = 0; i < VecSize; ++i) {
-    float quant_val =  static_cast<float>(in_vec[i]) / scale;
-    quant_val = quant_val > quant_max_bound ? quant_max_bound : quant_val;
-    quant_val = quant_val < quant_min_bound ? quant_min_bound : quant_val;
-    out_vec[i] = static_cast<OutType>(quant_val);
-  }
-  Store<OutType, VecSize>(out_vec, output + idx);
-  if (threadIdx.x == 0) {
-    out_scale_data[scale_idx] = scale;
-  }
-}
-
-template <paddle::DataType InType, paddle::DataType OutType>
-std::vector<paddle::Tensor> LaunchGroupQuantKernel(const paddle::Tensor& x,
-                                                   const int group_size,
-                                                   const bool transpose_scale,
-                                                   const float quant_max_bound,
-                                                   const float quant_min_bound) {
-    typedef PDTraits<InType> in_traits;
-    typedef typename in_traits::DataType InDataType;
-    typedef typename in_traits::data_t in_data_t;
-
-    paddle::Tensor out;
-    paddle::Tensor scale_out;
-    auto place = x.place();
-    cudaStream_t stream = x.stream();
-    int rank = x.dims().size();
-    std::vector<int64_t> out_shape = x.shape();
-    std::vector<int64_t> scale_shape = x.shape();
-
-    out = paddle::empty(out_shape, OutType, place);
-    int64_t batch = 1;
-    int64_t m = x.shape()[rank - 2];
-    int64_t n = x.shape()[rank - 1];
-    int64_t scale_n = n / group_size;
-    for(auto i = 0; i < rank - 1; ++i){
-        batch *= x.shape()[i];
-    }
-    if(transpose_scale){
-        scale_shape[rank - 2] = scale_n;
-        scale_shape[rank - 1] = m;
-    }else{
-        scale_shape[rank - 1] = scale_n;
-    }
-    
-    scale_out = paddle::empty(scale_shape, paddle::DataType::FLOAT32, place);
-
-    int64_t numel = x.numel();
-    constexpr int VecSize = 4;
-
-    dim3 threadsPerBlock(group_size / VecSize, 1, 1);
-    dim3 block_per_grid((numel + group_size - 1) / group_size, 1, 1);
-
-    typedef PDTraits<OutType> out_traits;
-    typedef typename out_traits::DataType OutDataType;
-    typedef typename out_traits::data_t out_data_t;
-    
-    if(group_size == 128){
-        GroupQuantKernel<InDataType, OutDataType, 128, VecSize, true><<<block_per_grid, threadsPerBlock, 0, stream>>>(reinterpret_cast<const InDataType*>(x.data<in_data_t>()),
-                            numel,
-                            m,
-                            n / 128,
-                            quant_max_bound,
-                            quant_min_bound,
-                            reinterpret_cast<OutDataType*>(out.data<out_data_t>()),
-                            reinterpret_cast<float*>(scale_out.data<float>()));
-    }else{
-        PD_THROW("group_quant's group_size only support 128.");
-    }
-    
-    return {out, scale_out};
-}
-template <paddle::DataType InType>
-std::vector<paddle::Tensor> LaunchGroupQuant(const paddle::Tensor& x,
-                                             const int group_size,
-                                             const bool transpose_scale,
-                                             const float quant_max_bound,
-                                             const float quant_min_bound) {
-
-    if(fabs(quant_max_bound - 448.0f) < 0.000001){
-        return LaunchGroupQuantKernel<InType, paddle::DataType::FLOAT8_E4M3FN>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
-    }else{
-        PD_THROW("Only supported float8_e4m3fn quantization, please set quant_max_bound=448, quant_min_bound=-448.");
-    }
-    
-}
-
-
-std::vector<paddle::Tensor> GroupQuant(const paddle::Tensor& x,
-                                        const int group_size,
-                                        const bool transpose_scale,
-                                        const float quant_max_bound,
-                                        const float quant_min_bound) {
-    if(x.dtype() == paddle::DataType::FLOAT32){
-        return LaunchGroupQuant<paddle::DataType::FLOAT32>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
-    }else if(x.dtype() == paddle::DataType::FLOAT16){
-        return LaunchGroupQuant<paddle::DataType::FLOAT16>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
-    }else if(x.dtype() == paddle::DataType::BFLOAT16){
-        return LaunchGroupQuant<paddle::DataType::BFLOAT16>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
-    }else{
-        PD_THROW("Unsupported data type.");
-    }
-}
-
-std::vector<std::vector<int64_t>> GroupQuantInferShape(const std::vector<int64_t>& input_shape, const int group_size, const bool transpose_scale, const float quant_max_bound,const float quant_min_bound) {
-    std::vector<int64_t> scale_shape = input_shape;
-    int rank = input_shape.size();
-    if(transpose_scale){
-        scale_shape[rank - 1] = input_shape[rank - 2];
-        scale_shape[rank - 2] = input_shape[rank - 1] / group_size;
-    }else{
-        scale_shape[rank - 1] = input_shape[rank - 1] / group_size;
-    }
-    return {input_shape, scale_shape};
-}
-
-std::vector<paddle::DataType> GroupQuantInferDtype(const paddle::DataType& input_dtype, const int group_size, const bool transpose_scale, const float quant_max_bound,const float quant_min_bound) {
-    
-    if(fabs(quant_max_bound - 448.0f) < 0.000001){
-        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT32};
-    }else{
-        PD_THROW("Only supported attr of quant_max_bound in ['448.0'].");
-    }
-}
-
-PD_BUILD_OP(group_quant)
-    .Inputs({"x"})
-    .Outputs({"output", "scale"})
-    .Attrs({"group_size: int",
-            "transpose_scale: bool",
-            "quant_max_bound: float",
-            "quant_min_bound: float"})
-    .SetKernelFn(PD_KERNEL(GroupQuant))
-    .SetInferShapeFn(PD_INFER_SHAPE(GroupQuantInferShape))
-    .SetInferDtypeFn(PD_INFER_DTYPE(GroupQuantInferDtype));

csrc/gpu/helper.h

@@ -162,6 +162,13 @@ public:
   typedef paddle::float8_e4m3fn data_t;
 };
 
+template <>
+class PDTraits<paddle::DataType::INT8> {
+public:
+  typedef int8_t DataType;
+  typedef int8_t data_t;
+};
+
 template <typename T, int Size>
 struct alignas(sizeof(T) * Size) AlignedVector {
   T val[Size];
@@ -245,3 +252,35 @@ inline bool GetMlaUseTensorcore() {
   const bool mla_use_tensorcore = flags_mla_use_tensorcore && enable_mla_tensorcore;
   return mla_use_tensorcore;
 }
+
+__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
+    float old;
+    old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+                        : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
+    return old;
+}
+
+__device__ __forceinline__ float warpReduceMax(float max_value) {
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 16));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 8));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 4));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 2));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 1));
+    return max_value;
+}
+
+__device__ __forceinline__ float blockReduceMax(float max_value) {
+    static __shared__ float warpLevelMaxs[32];
+    const int laneId = threadIdx.x & 0x1f;;
+    const int warpId = threadIdx.x >> 5;
+
+    max_value = warpReduceMax(max_value);
+
+    if (laneId == 0) warpLevelMaxs[warpId] = max_value;
+        __syncthreads();
+
+    max_value = (threadIdx.x < blockDim.x / 32) ? warpLevelMaxs[laneId] : 0;
+    if (warpId == 0) max_value = warpReduceMax(max_value);
+
+    return max_value;
+}

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/quantization/per_tensor_quant_fp8.cu

@@ -0,0 +1,169 @@
+// 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 "helper.h"
+#include<string.h>
+#include <cuda_runtime.h>
+
+template <typename InType>
+__global__ void
+per_tensor_absmax_kernel(const InType* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) {
+    float max_value = 0.0f;
+    unsigned int tid = threadIdx.x;
+    unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    const int grid_size = blockDim.x * gridDim.x;
+
+    constexpr uint32_t vec_size = 16 / sizeof(InType);
+    using in_vec_t = AlignedVector<InType, vec_size>;
+
+    const int32_t num_vec_elems = num_elements / vec_size;
+
+    for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
+        in_vec_t input_vec;
+        Load<InType, vec_size>(input + i * vec_size, &input_vec);
+
+        #pragma unroll
+        for (uint32_t j = 0; j < vec_size; ++j) {
+            float val = static_cast<float>(input_vec[j]);
+            max_value = fmaxf(max_value, fabsf(val));
+        }
+    }
+
+    const int32_t remaining_start = num_vec_elems * vec_size;
+    for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
+        float val = static_cast<float>(input[idx]);
+        max_value = fmaxf(max_value, fabsf(val));
+    }
+
+    max_value = blockReduceMax(max_value);
+
+    if (tid == 0) {
+        atomicMaxFloat(output_s, max_value / 448);
+    }
+}
+
+template <typename InType, typename OutType>
+__global__ void per_tensor_quant_fp8_kernel(
+    const InType* __restrict__ input,
+    OutType* __restrict__ output,
+    const float* __restrict__ scale,
+    const int64_t num_elements) {
+  const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int grid_size = blockDim.x * gridDim.x;
+  const float scale_val = 1.0f / (*scale);
+
+  constexpr uint32_t vec_size = 16 / sizeof(InType);
+  using in_vec_t = AlignedVector<InType, vec_size>;
+  using out_vec_t = AlignedVector<OutType, vec_size>;
+  in_vec_t input_vec;
+  out_vec_t output_vec;
+  const int32_t num_vec_elems = num_elements / vec_size;
+
+  for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
+    Load<InType, vec_size>(input + i * vec_size, &input_vec);
+    #pragma unroll
+    for (uint32_t j = 0; j < vec_size; ++j) {
+      float val = fmax(fmin(static_cast<float>(input_vec[j]) * scale_val, 448), -448);
+      output_vec[j] = static_cast<OutType>(val);
+    }
+    Store<OutType, vec_size>(output_vec, output + i * vec_size);
+  }
+
+  const int32_t remaining_start = num_vec_elems * vec_size;
+  for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
+    float val = fmax(-448, fmin(static_cast<float>(input[idx]) * scale_val, 448));
+    output[idx] = static_cast<OutType>(val);
+  }
+}
+
+
+template <paddle::DataType InType, paddle::DataType OutType>
+std::vector<paddle::Tensor> LaunchPerTensorQuantFp8Kernel(const paddle::Tensor& x, const paddle::optional<paddle::Tensor>& scale) {
+    typedef PDTraits<InType> in_traits;
+    typedef typename in_traits::DataType InDataType;
+    typedef typename in_traits::data_t in_data_t_pd;
+
+    typedef PDTraits<OutType> out_traits;
+    typedef typename out_traits::DataType OutDataType;
+    typedef typename out_traits::data_t out_data_t_pd;
+
+    paddle::Tensor out;
+    paddle::Tensor scale_out;
+    auto place = x.place();
+    cudaStream_t stream = x.stream();
+    int rank = x.dims().size();
+    std::vector<int64_t> out_shape = x.shape();
+    std::vector<int64_t> scale_shape = {1};
+
+    out = paddle::empty(out_shape, OutType, place);
+    if(scale){
+        scale_out = scale.get();
+    }else{
+        scale_out = paddle::empty(scale_shape, paddle::DataType::FLOAT32, place);
+    }
+
+    const int block_size = 256;
+    const int64_t num_elements = x.numel();
+    const int64_t num_blocks = min((num_elements + block_size - 1) / block_size, static_cast<int64_t>(1024));
+
+    dim3 grid(num_blocks);
+    dim3 block(block_size);
+    if(scale){
+        per_tensor_absmax_kernel<InDataType><<<grid, block, 0, stream>>>(
+          reinterpret_cast<const InDataType*>(x.data<in_data_t_pd>()), reinterpret_cast<float*>(scale_out.data<float>()), num_elements);
+    }
+
+    per_tensor_quant_fp8_kernel<InDataType, OutDataType><<<grid, block, 0, stream>>>(
+        reinterpret_cast<const InDataType*>(x.data<in_data_t_pd>()),
+        reinterpret_cast<OutDataType*>(out.data<out_data_t_pd>()),
+        reinterpret_cast<float*>(scale_out.data<float>()),
+        num_elements);
+    return {out, scale_out};
+}
+template <paddle::DataType InType>
+std::vector<paddle::Tensor> LaunchPerTensorQuantFp8(const paddle::Tensor& x, const paddle::optional<paddle::Tensor>& scale) {
+    return LaunchPerTensorQuantFp8Kernel<InType, paddle::DataType::FLOAT8_E4M3FN>(x, scale);
+}
+
+
+std::vector<paddle::Tensor> PerTensorQuantFp8(const paddle::Tensor& x, const paddle::optional<paddle::Tensor>& scale) {
+    if(x.dtype() == paddle::DataType::FLOAT32){
+        return LaunchPerTensorQuantFp8<paddle::DataType::FLOAT32>(x, scale);
+    }else if(x.dtype() == paddle::DataType::FLOAT16){
+        return LaunchPerTensorQuantFp8<paddle::DataType::FLOAT16>(x, scale);
+    }else if(x.dtype() == paddle::DataType::BFLOAT16){
+        return LaunchPerTensorQuantFp8<paddle::DataType::BFLOAT16>(x, scale);
+    }else{
+        PD_THROW("Unsupported data type.");
+    }
+}
+
+std::vector<std::vector<int64_t>> PerTensorQuantFp8InferShape(const std::vector<int64_t>& input_shape, const paddle::optional<std::vector<int64_t>>& scale_shape) {
+    std::vector<int64_t> scale_out_shape = {1};
+    if(scale_shape){
+        return {input_shape, scale_shape.get()};
+    }
+    return {input_shape, scale_out_shape};
+}
+
+std::vector<paddle::DataType> PerTensorQuantFp8InferDtype(const paddle::DataType& input_dtype, const paddle::optional<paddle::DataType>& scale_dtype) {
+    return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT32};
+}
+
+PD_BUILD_OP(per_tensor_quant_fp8)
+    .Inputs({"x", paddle::Optional("scale")})
+    .Outputs({"output", "scale_out"})
+    .SetKernelFn(PD_KERNEL(PerTensorQuantFp8))
+    .SetInferShapeFn(PD_INFER_SHAPE(PerTensorQuantFp8InferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(PerTensorQuantFp8InferDtype));

csrc/gpu/quantization/per_token_group_quant.cu

@@ -0,0 +1,251 @@
+// 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 "helper.h"
+#include<string.h>
+#include <cuda_runtime.h>
+
+
+__device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
+  unsigned mask = 0xffff;
+
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
+  return val;
+}
+
+template <typename InType, typename OutType>
+__global__ void PerTokenGroupQuantKernel(
+    const InType* __restrict__ input,
+    OutType* __restrict__ output_q,
+    float* __restrict__ output_s,
+    const int group_size,
+    const int num_groups,
+    const int groups_per_block,
+    const float eps,
+    const float quant_min_bound,
+    const float quant_max_bound,
+    bool transpose_scale = false,
+    const int scale_num_rows = 0,
+    const int scale_stride = 0) {
+  const int threads_per_group = 16;
+  const int local_group_id = threadIdx.x / threads_per_group;
+  const int lane_id = threadIdx.x % threads_per_group;
+
+  const int block_group_id = blockIdx.x * groups_per_block;
+  const int global_group_id = block_group_id + local_group_id;
+  const int block_group_offset = global_group_id * group_size;
+
+  float local_absmax = eps;
+
+  const InType* group_input = input + block_group_offset;
+  OutType* group_output = output_q + block_group_offset;
+  float* scale_output;
+
+  if (transpose_scale) {
+    const int row_idx = global_group_id / scale_num_rows;
+    const int col_idx = global_group_id % scale_num_rows;
+    scale_output = output_s + (col_idx * scale_stride + row_idx);
+  } else {
+    scale_output = output_s + global_group_id;
+  }
+
+  constexpr uint32_t vec_size = 16 / sizeof(InType);
+  using vec_t = AlignedVector<InType, vec_size>;
+
+  const int32_t num_vec_elems = group_size / vec_size;
+
+  for (int32_t i = lane_id; i < num_vec_elems; i += 16) {
+    vec_t input_vec;
+    Load<InType, vec_size>(group_input + i * vec_size, &input_vec);
+
+#pragma unroll
+    for (uint32_t j = 0; j < vec_size; ++j) {
+      float val = static_cast<float>(input_vec[j]);
+      float abs_val = fabsf(val);
+      local_absmax = fmaxf(local_absmax, abs_val);
+    }
+  }
+
+  local_absmax = GroupReduceMax(local_absmax, lane_id);
+
+  const float y_s = local_absmax / quant_max_bound;
+
+  if (lane_id == 0) {
+    *scale_output = y_s;
+  }
+
+  for (int32_t i = lane_id; i < num_vec_elems; i += 16) {
+    vec_t input_vec;
+    Load<InType, vec_size>(group_input + i * vec_size, &input_vec);
+
+#pragma unroll
+    for (uint32_t j = 0; j < vec_size; ++j) {
+      float val = static_cast<float>(input_vec[j]);
+      float q_val = fminf(fmaxf(val / y_s, quant_min_bound), quant_max_bound);
+      group_output[i * vec_size + j] = static_cast<OutType>(q_val);
+    }
+  }
+}
+
+
+template <paddle::DataType InType, paddle::DataType OutType>
+std::vector<paddle::Tensor> LaunchPerTokenGroupQuantKernel(const paddle::Tensor& x,
+                                                   const int group_size,
+                                                   const bool transpose_scale,
+                                                   const float quant_max_bound,
+                                                   const float quant_min_bound) {
+    typedef PDTraits<InType> in_traits;
+    typedef typename in_traits::DataType InDataType;
+    typedef typename in_traits::data_t in_data_t;
+
+    paddle::Tensor out;
+    paddle::Tensor scale_out;
+    auto place = x.place();
+    cudaStream_t stream = x.stream();
+    int rank = x.dims().size();
+    std::vector<int64_t> out_shape = x.shape();
+    std::vector<int64_t> scale_shape = x.shape();
+    int64_t m = x.shape()[rank - 2];
+    int64_t k = x.shape()[rank - 1];
+    PD_CHECK(k % group_size == 0);
+    int64_t scale_k = k / group_size;
+
+    out = paddle::empty(out_shape, OutType, place);
+    if(transpose_scale){
+        scale_shape[rank - 2] = scale_k;
+        scale_shape[rank - 1] = m;
+    }else{
+        scale_shape[rank - 1] = scale_k;
+    }
+    scale_out = paddle::empty(scale_shape, paddle::DataType::FLOAT32, place);
+    int64_t numel = x.numel();
+    const int num_groups = numel / group_size;
+    constexpr int THREADS_PER_GROUP = 16;
+    int groups_per_block = 1;
+
+    if (num_groups % 16 == 0) {
+        groups_per_block = 16;
+    } else if (num_groups % 8 == 0) {
+        groups_per_block = 8;
+    } else if (num_groups % 4 == 0) {
+        groups_per_block = 4;
+    } else if (num_groups % 2 == 0) {
+        groups_per_block = 2;
+    }
+
+    const int num_blocks = num_groups / groups_per_block;
+    const int num_threads = groups_per_block * THREADS_PER_GROUP;
+
+    int scale_num_rows = 0;
+    int scale_stride = 0;
+    if (transpose_scale){
+        scale_num_rows = m;
+        scale_stride = scale_k;
+    }
+    
+    dim3 grid(num_blocks);
+    dim3 block(num_threads);
+
+    typedef PDTraits<OutType> out_traits;
+    typedef typename out_traits::DataType OutDataType;
+    typedef typename out_traits::data_t out_data_t;
+    float eps = 0.000001f;
+
+    PerTokenGroupQuantKernel<InDataType, OutDataType><<<grid, block, 0, stream>>>(reinterpret_cast<const InDataType*>(x.data<in_data_t>()),
+    reinterpret_cast<OutDataType*>(out.data<out_data_t>()),
+    reinterpret_cast<float*>(scale_out.data<float>()),
+    group_size,
+    num_groups,
+    groups_per_block,
+    eps,
+    quant_min_bound,
+    quant_max_bound,
+    transpose_scale,
+    scale_num_rows,
+    scale_stride);
+    
+    return {out, scale_out};
+}
+
+template <paddle::DataType InType>
+std::vector<paddle::Tensor> LaunchPerTokenGroupQuant(const paddle::Tensor& x,
+                                             const int group_size,
+                                             const bool transpose_scale,
+                                             const float quant_max_bound,
+                                             const float quant_min_bound) {
+
+    if(fabs(quant_max_bound - 448.0f) < 0.000001){
+        return LaunchPerTokenGroupQuantKernel<InType, paddle::DataType::FLOAT8_E4M3FN>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else if(fabs(quant_max_bound - 127.0f) < 0.000001){
+        return LaunchPerTokenGroupQuantKernel<InType, paddle::DataType::INT8>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else{
+        PD_THROW("Only supported float8_e4m3fn and int8 quantization.");
+    }
+    
+}
+
+
+std::vector<paddle::Tensor> PerTokenGroupQuant(const paddle::Tensor& x,
+                                        const int group_size,
+                                        const bool transpose_scale,
+                                        const float quant_max_bound,
+                                        const float quant_min_bound) {
+    if(x.dtype() == paddle::DataType::FLOAT32){
+        return LaunchPerTokenGroupQuant<paddle::DataType::FLOAT32>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else if(x.dtype() == paddle::DataType::FLOAT16){
+        return LaunchPerTokenGroupQuant<paddle::DataType::FLOAT16>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else if(x.dtype() == paddle::DataType::BFLOAT16){
+        return LaunchPerTokenGroupQuant<paddle::DataType::BFLOAT16>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else{
+        PD_THROW("Unsupported data type.");
+    }
+}
+
+std::vector<std::vector<int64_t>> PerTokenGroupQuantInferShape(const std::vector<int64_t>& input_shape, const int group_size, const bool transpose_scale, const float quant_max_bound,const float quant_min_bound) {
+    std::vector<int64_t> scale_shape = input_shape;
+    int rank = input_shape.size();
+    PD_CHECK(scale_shape[rank-1] % group_size == 0);
+    if(transpose_scale){
+        scale_shape[rank - 1] = input_shape[rank - 2];
+        scale_shape[rank - 2] = input_shape[rank - 1] / group_size;
+    }else{
+        scale_shape[rank - 1] = input_shape[rank - 1] / group_size;
+    }
+    return {input_shape, scale_shape};
+}
+
+std::vector<paddle::DataType> PerTokenGroupQuantInferDtype(const paddle::DataType& input_dtype, const int group_size, const bool transpose_scale, const float quant_max_bound,const float quant_min_bound) {
+    
+    if(fabs(quant_max_bound - 448.0f) < 0.000001){
+        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT32};
+    }else if(fabs(quant_max_bound - 127.0f) < 0.000001){
+        return {paddle::DataType::INT8, paddle::DataType::FLOAT32};
+    }else{
+        PD_THROW("Only supported attr of quant_max_bound in [448.0, 127.0].");
+    }
+}
+
+PD_BUILD_OP(per_token_group_quant)
+    .Inputs({"x"})
+    .Outputs({"output", "scale"})
+    .Attrs({"group_size: int",
+            "transpose_scale: bool",
+            "quant_max_bound: float",
+            "quant_min_bound: float"})
+    .SetKernelFn(PD_KERNEL(PerTokenGroupQuant))
+    .SetInferShapeFn(PD_INFER_SHAPE(PerTokenGroupQuantInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(PerTokenGroupQuantInferDtype));

csrc/gpu/rebuild_padding_v2.cu

@@ -138,7 +138,7 @@ std::vector<paddle::Tensor> rebuild_padding_v2(const paddle::Tensor& tmp_out, //
     return {out};
 }
 
-std::vector<paddle::Tensor> RebuildPaddingV2(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
+paddle::Tensor RebuildPaddingV2Func(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,
@@ -153,7 +153,7 @@ std::vector<paddle::Tensor> RebuildPaddingV2(const paddle::Tensor& tmp_out, // [
                 seq_lens_encoder,
                 output_padding_offset,
                 max_input_length
-            );
+            )[0];
         }
         case paddle::DataType::FLOAT16: {
             return rebuild_padding_v2<paddle::DataType::FLOAT16>(
@@ -163,7 +163,7 @@ std::vector<paddle::Tensor> RebuildPaddingV2(const paddle::Tensor& tmp_out, // [
                 seq_lens_encoder,
                 output_padding_offset,
                 max_input_length
-            );
+            )[0];
         }
         case paddle::DataType::FLOAT32: {
             return rebuild_padding_v2<paddle::DataType::FLOAT32>(
@@ -173,7 +173,7 @@ std::vector<paddle::Tensor> RebuildPaddingV2(const paddle::Tensor& tmp_out, // [
                 seq_lens_encoder,
                 output_padding_offset,
                 max_input_length
-            );
+            )[0];
         }
         default: {
             PD_THROW(
@@ -184,6 +184,22 @@ std::vector<paddle::Tensor> RebuildPaddingV2(const paddle::Tensor& tmp_out, // [
     }
 }
 
+std::vector<paddle::Tensor> RebuildPaddingV2(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::optional<paddle::Tensor>& output_padding_offset,
+                                             int max_input_length) {
+    
+    return {RebuildPaddingV2Func(
+                tmp_out,
+                cum_offsets,
+                seq_lens_decoder,
+                seq_lens_encoder,
+                output_padding_offset,
+                max_input_length
+            )};
+}
 std::vector<std::vector<int64_t>> RebuildPaddingV2InferShape(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,

csrc/gpu/save_output_dygraph.cu

@@ -0,0 +1,76 @@
+// 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 "helper.h"
+
+template<typename T>
+__global__ void SaveOutputDygraphKernel(
+    T* out_pool,
+    const T* cur_out,
+    const int64_t* step_idx,
+    const int* result_idx,
+    const int bsz,
+    const int result_num,
+    const int max_dec_len) {
+        for (int bid = blockIdx.x; bid < bsz; bid += gridDim.x) {
+            auto src_addr = cur_out + bid;
+            auto cur_step_idx = step_idx[bid] - 1;
+            for (int i = threadIdx.x; i < result_num; i += blockDim.x) {
+                auto tgt_id = result_idx[bid * result_num + i];
+                if (tgt_id >= 0) {
+                    auto tgt_addr = out_pool + tgt_id * max_dec_len + cur_step_idx;
+                    *tgt_addr = *src_addr;
+                }
+            }
+        }
+}
+
+void SaveOutputDygraph(
+    const paddle::Tensor& all_token_ids,
+    const paddle::Tensor& tokens,
+    const paddle::Tensor& result_ids,
+    const paddle::Tensor& step_idx
+) {
+    auto cu_stream = all_token_ids.stream();
+    
+    const int bsz = tokens.shape()[0];
+    const int result_num = result_ids.shape()[1];
+    const int max_dec_len = all_token_ids.shape()[1];
+
+    SaveOutputDygraphKernel<<<bsz, 256, 0, cu_stream>>>(
+        const_cast<int64_t*>(all_token_ids.data<int64_t>()),
+        tokens.data<int64_t>(),
+        step_idx.data<int64_t>(),
+        result_ids.data<int>(),
+        bsz,
+        result_num,
+        max_dec_len
+    );
+
+    // SaveOutputDygraphKernel<<<bsz, 256, 0, cu_stream>>>(
+    //     const_cast<float*>(all_scores.data<float>()),
+    //     scores.data<float>(),
+    //     step_idx.data<int64_t>(),
+    //     result_ids.data<int>(),
+    //     bsz,
+    //     result_num,
+    //     max_dec_len
+    // );
+}
+
+PD_BUILD_OP(save_output_dygraph)
+    .Inputs({"all_token_ids", "tokens", "result_ids", "step_idx"})
+    .Outputs({"all_token_ids_out"})
+    .SetInplaceMap({{"all_token_ids", "all_token_ids_out"}})
+    .SetKernelFn(PD_KERNEL(SaveOutputDygraph));

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;
@@ -35,7 +36,25 @@ void SaveOutMmsg(const paddle::Tensor& x,
     auto not_need_stop_data = not_need_stop.data<bool>()[0];
 
     static struct msgdata msg_sed;
-    static key_t key = ftok("./", 1);
+    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;

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_dynamic_per_tensor_quant_fp8.py

@@ -0,0 +1,88 @@
+# 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.
+
+import unittest
+
+import numpy as np
+import paddle
+from paddlenlp_ops import per_tensor_quant_fp8
+
+paddle.seed(1)
+
+
+def create_random_cuda_tensor(shape, dtype, mean: float = 0, std: float = 1):
+    return paddle.empty(shape, dtype=dtype).normal_(mean, std)
+
+
+class DynamicPerTensorQuantFP8Test(unittest.TestCase):
+    def native_per_tensor_quant_fp8(self, x, scale=None):
+        x_fp32 = x.cast("float32")
+        if scale is None:
+            x_s = x_fp32.abs().max() / 448.0
+        else:
+            x_s = scale
+        x_q = x_fp32 / x_s
+        x_q = x_q.clip(min=-448.0, max=448.0)
+        return x_q.cast("float8_e4m3fn"), x_s.reshape([1])
+
+    def test_dynamic_per_tensor_quant_fp8_bf16(self):
+        num_tokens = 1024
+        hidden_size = 1152
+        dtype = paddle.bfloat16
+        x = paddle.rand([num_tokens, hidden_size], dtype=dtype)
+        ref_out, ref_scale = self.native_per_tensor_quant_fp8(x)
+        out, scale = per_tensor_quant_fp8(x, None)
+
+        np.testing.assert_allclose(ref_out.numpy(), out.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(ref_scale.numpy(), scale.numpy(), rtol=1e-3, atol=1e-3)
+
+    def test_dynamic_per_tensor_quant_fp8_fp32(self):
+        num_tokens = 1024
+        hidden_size = 1152
+        dtype = paddle.float32
+        x = paddle.rand([num_tokens, hidden_size], dtype=dtype)
+
+        ref_out, ref_scale = self.native_per_tensor_quant_fp8(x)
+        out, scale = per_tensor_quant_fp8(x, None)
+
+        np.testing.assert_allclose(ref_out.numpy(), out.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(ref_scale.numpy(), scale.numpy(), rtol=1e-3, atol=1e-3)
+
+    def test_dynamic_per_tensor_quant_fp8_fp16(self):
+        num_tokens = 1024
+        hidden_size = 1152
+        dtype = paddle.float16
+        x = paddle.rand([num_tokens, hidden_size], dtype=dtype)
+
+        ref_out, ref_scale = self.native_per_tensor_quant_fp8(x)
+        out, scale = per_tensor_quant_fp8(x, None)
+
+        np.testing.assert_allclose(ref_out.numpy(), out.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(ref_scale.numpy(), scale.numpy(), rtol=1e-3, atol=1e-3)
+
+    def test_static_per_tensor_quant_fp8_bf16(self):
+        num_tokens = 1024
+        hidden_size = 1152
+        dtype = paddle.bfloat16
+        x = paddle.rand([num_tokens, hidden_size], dtype=dtype)
+        scale = paddle.rand([1], dtype=paddle.float32).clip(0.0001)
+        ref_out, ref_scale = self.native_per_tensor_quant_fp8(x, scale)
+        out, scale = per_tensor_quant_fp8(x, scale)
+
+        np.testing.assert_allclose(ref_out.numpy(), out.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(ref_scale.numpy(), scale.numpy(), rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()

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/unittest/test_per_token_group_quant.py

@@ -0,0 +1,119 @@
+# 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.
+
+import unittest
+
+import numpy as np
+import paddle
+from paddlenlp_ops import per_token_group_quant
+
+
+class PerTokenGroupQuantTest(unittest.TestCase):
+    def native_per_token_group_quant(self, x, group_size, quant_max_bound, quant_min_bound):
+        eps = 0.000001
+        x = x.cast(paddle.float32)
+        x_ = x.reshape([x.numel() // group_size, group_size])
+        amax = x_.abs().max(axis=-1, keepdim=True).clip(min=eps)
+        x_s = amax / quant_max_bound
+        x_q = (x_ / x_s).clip(min=quant_min_bound, max=quant_max_bound)
+        x_q = x_q.reshape(x.shape)
+        s_shape = x.shape[:-1] + [x.shape[-1] // group_size]
+        x_s = x_s.reshape(s_shape)
+        if abs(quant_max_bound - 448) < 0.00001:
+            x_q = x_q.cast(paddle.float8_e4m3fn)
+        elif abs(quant_max_bound - 127) < 0.00001:
+            x_q = x_q.cast(paddle.int8)
+        else:
+            assert f"Error quant_max_bound {quant_max_bound}."
+        return x_q, x_s
+
+    def test_per_token_group_quant_fp8_t(self):
+        M = 32
+        K = 1024
+        x = np.random.rand(M, K).astype(np.float32)
+
+        x_tensor = paddle.to_tensor(x).cast(paddle.float16)
+        x_q_ref, x_s_ref = self.native_per_token_group_quant(
+            x_tensor, group_size=128, quant_max_bound=448.0, quant_min_bound=-448.0
+        )
+        x_s_ref = x_s_ref.transpose([1, 0])
+
+        x_q, x_s = per_token_group_quant(
+            x_tensor, group_size=128, transpose_scale=True, quant_max_bound=448.0, quant_min_bound=-448.0
+        )
+
+        x_q_ref = x_q_ref.cast(paddle.float32)
+        x_q = x_q.cast(paddle.float32)
+        np.testing.assert_allclose(x_q.numpy(), x_q_ref.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(x_s.numpy(), x_s_ref.numpy(), rtol=1e-3, atol=1e-3)
+
+    def test_per_token_group_quant_fp8(self):
+        M = 32
+        K = 1024
+        x = np.random.rand(M, K).astype(np.float32)
+
+        x_tensor = paddle.to_tensor(x).cast(paddle.float16)
+        x_q_ref, x_s_ref = self.native_per_token_group_quant(
+            x_tensor, group_size=128, quant_max_bound=448.0, quant_min_bound=-448.0
+        )
+
+        x_q, x_s = per_token_group_quant(
+            x_tensor, group_size=128, transpose_scale=False, quant_max_bound=448.0, quant_min_bound=-448.0
+        )
+
+        x_q_ref = x_q_ref.cast(paddle.float32)
+        x_q = x_q.cast(paddle.float32)
+        np.testing.assert_allclose(x_q.numpy(), x_q_ref.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(x_s.numpy(), x_s_ref.numpy(), rtol=1e-3, atol=1e-3)
+
+    def test_per_token_group_quant_int8_t(self):
+        M = 32
+        K = 1024
+
+        x_tensor = paddle.randn([M, K], dtype=paddle.float16)
+        x_q_ref, x_s_ref = self.native_per_token_group_quant(
+            x_tensor, group_size=128, quant_max_bound=127.0, quant_min_bound=-127.0
+        )
+        x_s_ref = x_s_ref.transpose([1, 0])
+
+        x_q, x_s = per_token_group_quant(
+            x_tensor, group_size=128, transpose_scale=True, quant_max_bound=127.0, quant_min_bound=-127.0
+        )
+        x_q_ref = x_q_ref.cast(paddle.float32)
+        x_q = x_q.cast(paddle.float32)
+        np.testing.assert_allclose(x_q.numpy(), x_q_ref.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(x_s.numpy(), x_s_ref.numpy(), rtol=1e-3, atol=1e-3)
+
+    def test_per_token_group_quant_int8(self):
+        M = 32
+        K = 1024
+        x = np.random.rand(M, K).astype(np.float32)
+
+        x_tensor = paddle.to_tensor(x).cast(paddle.float16)
+        x_q_ref, x_s_ref = self.native_per_token_group_quant(
+            x_tensor, group_size=128, quant_max_bound=127.0, quant_min_bound=-127.0
+        )
+
+        x_q, x_s = per_token_group_quant(
+            x_tensor, group_size=128, transpose_scale=False, quant_max_bound=127.0, quant_min_bound=-127.0
+        )
+
+        x_q_ref = x_q_ref.cast(paddle.float32)
+        x_q = x_q.cast(paddle.float32)
+        np.testing.assert_allclose(x_q.numpy(), x_q_ref.numpy(), rtol=1e-3, atol=1e-3)
+        np.testing.assert_allclose(x_s.numpy(), x_s_ref.numpy(), rtol=1e-3, atol=1e-3)
+
+
+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:
@@ -105,7 +117,6 @@ sources = [
     "./gpu/step.cu",
     "./gpu/quant_int8.cu",
     "./gpu/dequant_int8.cu",
-    "./gpu/group_quant.cu",
     "./gpu/moe/preprocess_for_moe.cu",
     "./gpu/get_position_ids_and_mask_encoder_batch.cu",
     "./gpu/fused_rotary_position_encoding.cu",
@@ -118,12 +129,12 @@ sources = [
     "./gpu/speculate_decoding_kernels/ngram_match.cc",
     "./gpu/speculate_decoding_kernels/speculate_save_output.cc",
     "./gpu/speculate_decoding_kernels/speculate_get_output.cc",
-    "./gpu/cpp_extensions.cu",
+    "./gpu/save_output_dygraph.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()
@@ -151,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"]
@@ -159,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"
@@ -176,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",
@@ -185,21 +201,17 @@ if cc >= 80 and cuda_version >= 12.4:
     ]
     sources += ["./gpu/sage_attn_kernels/sageattn_fused.cu"]
     if cc >= 80 and cc < 89:
-        sources += [
-            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f16_kernel_sm80.cu"
-        ]
-        nvcc_compile_args += ["-gencode", f"arch=compute_80,code=compute_80"]
+        sources += ["./gpu/sage_attn_kernels/sageattn_qk_int_sv_f16_kernel_sm80.cu"]
+        nvcc_compile_args += ["-gencode", "arch=compute_80,code=compute_80"]
     elif cc >= 89 and cc < 90:
-        sources += [
-            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm89.cu"
-        ]
-        nvcc_compile_args += ["-gencode", f"arch=compute_89,code=compute_89"]
+        sources += ["./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm89.cu"]
+        nvcc_compile_args += ["-gencode", "arch=compute_89,code=compute_89"]
     elif cc >= 90:
         sources += [
             "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm90.cu",
-            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_dsk_kernel_sm90.cu"
+            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_dsk_kernel_sm90.cu",
         ]
-        nvcc_compile_args += ["-gencode", f"arch=compute_90a,code=compute_90a"]
+        nvcc_compile_args += ["-gencode", "arch=compute_90a,code=compute_90a"]
 
 if cc >= 90 and cuda_version >= 12.0:
     os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py --cuda_arch 90")
@@ -222,7 +234,10 @@ setup(
     name=ops_name,
     ext_modules=CUDAExtension(
         sources=sources,
-        extra_compile_args={"cxx": ["-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"], "nvcc": nvcc_compile_args},
+        extra_compile_args={
+            "cxx": ["-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"],
+            "nvcc": nvcc_compile_args,
+        },
         libraries=["cublasLt"],
         library_dirs=library_path,
         include_dirs=include_dirs,

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/data_prepare/dataset_list.md

@@ -54,9 +54,9 @@ PaddleNLP provides easy-to-use APIs for the following datasets. Please add **spl
 | [HYP](https://pan.webis.de/semeval19/semeval19-web/) | English political news sentiment classification corpus |
 | Dataset Name | Description | How to Use |
 | ---- | --------- | ------ |
-| [ChnSentiCorp](https://github.com/zhangyics/Chinese-sentiment-analysis-datasets) | Chinese sentiment analysis dataset | `paddlenlp.datasets.load_dataset('chnsenticorp')` |
+| [ChnSentiCorp](https://paddlehub-dataset.bj.bcebos.com/chnsenticorp.tar.gz) | Chinese sentiment analysis dataset | `paddlenlp.datasets.load_dataset('chnsenticorp')` |
 | [LCQMC](http://icrc.hitsz.edu.cn/info/1037/1146.htm) | Chinese question matching corpus, binary classification task | `paddlenlp.datasets.load_dataset('lcqmc', splits=['test', 'dev'])` |
-| [NLPCC-DBQA](http://tcci.ccf.org.cn/conference/2016/pages/page05_eva.html) | Chinese database question answering dataset, binary classification task | `paddlenlp.datasets.load_dataset('nlpcc_dbqa')` |
+| [NLPCC-DBQA](https://github.com/shuishen112/NLPCCDBQA) | Chinese database question answering dataset, binary classification task | `paddlenlp.datasets.load_dataset('nlpcc_dbqa')` |
 
 ## Natural Language Inference
 

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/dataset.md

@@ -37,25 +37,7 @@ Here we use the ziya-llama-13b-v1 model as an example to illustrate the complete
 **2.2.1 Raw Data**
 
 First download sample data:
-```shell
-wget https://paddlenlp.bj.bcebos.com/datasets/PDC_DATASETS/PRETRAIN/zh_sample.txt
-```
 
-Convert to jsonl format:
-```shell
-python trans_to_json.py \
-    --input zh_sample.txt \
-    --output zh_sample.json \
-    --workers 8
-```
-
-For large files (recommended processing method):
-```shell
-split -l 10000 zh_sample.txt -d -a 3 zh_part_
-ls zh_part_* | xargs -I {} -P 8 python trans_to_json.py --input {} --output {}.json
-cat *.json > zh_sample_all.json
-rm zh_part_*
-```
 ```
 mkdir data && cd data
 wget https://bj.bcebos.com/paddlenlp/models/transformers/data_tools/baike.txt

docs/en/llm/devices/dcu/llama/README.md

@@ -120,7 +120,7 @@ PYTHONPATH=.. python run_finetune.py dcu/llama/sft_argument.json
 
 ## 3. Pre-training
 ### Data Preparation
-For detailed data preparation workflow, please refer to [here](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/tools/preprocess/README.md). Example: OpenWebText2 pre-training data preparation refers to [here](https://paddlenlp.readthedocs.io/zh/latest/llm/pretraining/data/OpenWebText2.html).
+For detailed data preparation workflow, please refer to [here](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/tools/preprocess/README.md). Example: OpenWebText2 pre-training data preparation refers to [here](https://paddlenlp.readthedocs.io/zh/latest/llm/pretraining/data/).
 
 For user convenience in testing, this project provides processed 100k doc training samples:
 

docs/en/llm/docs/dpo.md

@@ -140,7 +140,7 @@ Type of DPO loss function, which can be one of:
 - `sigmoid`([DPO](https://arxiv.org/abs/2305.18290)),
 - `hinge`([RSO](https://arxiv.org/abs/2309.06657)),
 - `ipo`([IPO](https://arxiv.org/abs/2310.12036)),
-- `kto_pair` (implementation for preference pairs [KTO](https://github.com/ContextualAI/HALOs/blob/main/assets/report.pdf)),
+- `kto_pair` (implementation for preference pairs [KTO](https://github.com/ContextualAI/HALOs/blob/winnie/research/assets/report.pdf)),
 - `sppo_hard`([SPPO](https://arxiv.org/pdf/2405.00675)),
 - `nca_pair`([NCA](https://arxiv.org/abs/2402.05369)),
 - `dpop`([DPOP](https://arxiv.org/pdf/2402.13228.pdf)),

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/en/llm/tools/preprocess/docs/OpenWebText2.md

@@ -12,7 +12,6 @@ This example uses the cleaned [OpenWebText2 data](https://openwebtext2.readthedo
 After downloading, decompress with the following command:
 
 ```shell
-# wget https://mystic.the-eye.eu/public/AI/pile_preliminary_components/openwebtext2.jsonl.zst.tar
 wget https://paddlenlp.bj.bcebos.com/models/transformers/gpt/openwebtext2.jsonl.zst.tar
 tar -xvf openwebtext2.json.zst.tar -C /path/to/openwebtext
 ```

docs/en/model_zoo/taskflow.md

@@ -1891,7 +1891,7 @@ Word segmentation via this approach can significantly improve prediction speed.
 
 <details><summary><b>Q:</b> Will more tasks be supported in the future?</summary><div>
 
-**A:** Taskflow will continue to expand supported tasks. We'll adjust development priorities based on developer feedback. Please submit your requests via Issues or our [questionnaire](https://wenjuan.baidu-int.com/manage/?r=survey/pageEdit&sid=85827).
+**A:** Taskflow will continue to expand supported tasks. We'll adjust development priorities based on developer feedback.
 
 </div></details>
 

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/alignment/rm/README.md

@@ -0,0 +1 @@
+../../../../../llm/alignment/rm/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