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-20.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/pipelines.yml

@@ -12,7 +12,7 @@ on:
 jobs:
   Pipelines-Test:
     name: Pipelines-Test
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4

.github/workflows/release.yml

@@ -16,7 +16,7 @@ on:
 jobs:
   Release:
     name: Release
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v1

.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/stale.yml

@@ -13,7 +13,7 @@ permissions:
 
 jobs:
   stale:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
 
     permissions:
       issues: write

.github/workflows/tests.yml

@@ -1,29 +0,0 @@
-name: Test
-
-on: [push, pull_request]
-
-jobs:
-  Test:
-    name: Test
-    runs-on: ubuntu-20.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)$

.readthedocs.yaml

@@ -15,7 +15,7 @@ submodules:
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
-   configuration: docs/conf.py
+   configuration: docs/zh/conf.py
 
 # Optionally build your docs in additional formats such as PDF
 #formats:

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,14 +33,20 @@
 
 ## News 📢
 
-* **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.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
 
@@ -63,7 +69,7 @@
 ### <a href=#多硬件训推一体> 🔧 多硬件训推一体 </a>
 
 支持英伟达 GPU、昆仑 XPU、昇腾 NPU、燧原 GCU 和海光 DCU 等多个硬件的大模型和自然语言理解模型训练和推理，套件接口支持硬件快速切换，大幅降低硬件切换研发成本。
-当前支持的自然语言理解模型：[多硬件自然语言理解模型列表](./docs/model_zoo/model_list_multy_device.md)
+当前支持的自然语言理解模型：[多硬件自然语言理解模型列表](./docs/zh/model_zoo/model_list_multy_device.md)
 
 ### <a href=#高效易用的预训练> 🚀 高效易用的预训练 </a>
 
@@ -116,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                                                                                                                                                                                                                                                                                                                                    |
 
@@ -192,7 +199,7 @@ pip install --upgrade paddlenlp==3.0.0b4
 pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/paddlenlp.html
 ```
 
-更多关于 PaddlePaddle 和 PaddleNLP 安装的详细教程请查看[Installation](./docs/get_started/installation.rst)。
+更多关于 PaddlePaddle 和 PaddleNLP 安装的详细教程请查看[Installation](./docs/zh/get_started/installation.rst)。
 
 ------------------------------------------------------------------------------------------
 
@@ -203,13 +210,14 @@ pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/pad
 PaddleNLP 提供了方便易用的 Auto API，能够快速的加载模型和 Tokenizer。这里以使用 `Qwen/Qwen2-0.5B` 模型做文本生成为例：
 
 ```python
->>> from paddlenlp.transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
->>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B", dtype="float16")
->>> input_features = tokenizer("你好！请自我介绍一下。", return_tensors="pd")
->>> outputs = model.generate(**input_features, max_length=128)
->>> print(tokenizer.batch_decode(outputs[0], skip_special_tokens=True))
-['我是一个AI语言模型，我可以回答各种问题，包括但不限于：天气、新闻、历史、文化、科学、教育、娱乐等。请问您有什么需要了解的吗？']
+from paddlenlp.transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
+# if using CPU, please change float16 to float32
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B", dtype="float16")
+input_features = tokenizer("你好！请自我介绍一下。", return_tensors="pd")
+outputs = model.generate(**input_features, max_new_tokens=128)
+print(tokenizer.batch_decode(outputs[0], skip_special_tokens=True))
+# ['我是一个AI语言模型，我可以回答各种问题，包括但不限于：天气、新闻、历史、文化、科学、教育、娱乐等。请问您有什么需要了解的吗？']
 ```
 
 ### 大模型预训练

csrc/cpu/setup.sh

@@ -47,9 +47,11 @@ fi
 cd xFasterTransformer
 git apply paddle.patch
 
-# #4. build xFasterTransformer
-sh ./3rdparty/prepare_oneccl.sh
-source ./3rdparty/oneccl/build/_install/env/setvars.sh
+#4. build xFasterTransformer
+cd 3rdparty
+bash prepare_oneccl.sh
+source ./oneccl/build/_install/env/setvars.sh
+cd ..
 
 rm -rf build
 mkdir build && cd build

csrc/cpu/src/avx_weight_only.cc

@@ -15,7 +15,7 @@
 #include "matmul_helper.h"
 #include "my_types.h"
 #include "paddle/extension.h"
-git adtemplate <typename T>
+template <typename T>
 void AvxCompute(const paddle::Tensor &x,
                 const paddle::Tensor &weight,
                 bool trans,

csrc/cpu/src/stop_generation_multi_ends.cc

@@ -14,9 +14,20 @@
 
 #include <stdlib.h>
 #include <string.h>
-
-#include "helper.h"
 #include <stdio.h>
+#include <vector>
+#include "paddle/extension.h"
+
+
+bool is_in_end(const int64_t id, const int64_t *end_ids, int length) {
+    bool flag = false;
+    for (int i = 0; i < length; i++) {
+        if (id == end_ids[i]) {
+            return true;
+        }
+    }
+    return flag;
+}
 
 void set_value_by_flags(const bool* stop_flags,
                         const int64_t* end_ids,

csrc/gpu/all_reduce.cu

@@ -0,0 +1,125 @@
+// 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 "all_reduce.cuh"
+
+// Fake pointer type, must match fptr_t type in ops.h.
+// We use this type alias to indicate when pointers are passed in as int64_t.
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+fptr_t init_custom_all_reduce(const std::vector<fptr_t>& fake_ipc_ptrs,
+                      paddle::Tensor& rank_data, int64_t rank,
+                      bool full_nvlink) {
+  int world_size = fake_ipc_ptrs.size();
+  if (world_size > 8)
+    throw std::invalid_argument("world size > 8 is not supported");
+  if (world_size % 2 != 0)
+    throw std::invalid_argument("Odd num gpus is not supported for now");
+  if (rank < 0 || rank >= world_size)
+    throw std::invalid_argument("invalid rank passed in");
+
+  paddle::Signal* ipc_ptrs[8];
+  for (int i = 0; i < world_size; i++) {
+    ipc_ptrs[i] = reinterpret_cast<paddle::Signal*>(fake_ipc_ptrs[i]);
+  }
+  return (fptr_t) new paddle::CustomAllreduce(ipc_ptrs, rank_data.data(),
+                                            rank_data.numel(), rank, world_size,
+                                            full_nvlink);
+}
+
+/**
+ * Performs an out-of-place allreduce and stores result in out.
+ *
+ * If _reg_buffer is null, assumes inp.data() is already IPC-registered.
+ * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
+ * copied into _reg_buffer.
+ */
+void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
+                fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  auto stream = inp.stream();
+
+  auto input_size = inp.numel() * 2;
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+  if (reg_buffer) {
+    cudaMemcpyAsync(reg_buffer, inp.data(), input_size,
+                                  cudaMemcpyDeviceToDevice, stream);
+  } else {
+    reg_buffer = inp.data();
+  }
+  switch (out.dtype()) {
+    case phi::DataType::FLOAT32: {
+      fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
+                           reinterpret_cast<float*>(out.data()),
+                           out.numel());
+      break;
+    }
+    case phi::DataType::FLOAT16: {
+      fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
+                          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");
+  }
+}
+
+void dispose(fptr_t _fa) {
+  delete reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+}
+
+int64_t meta_size() { return sizeof(paddle::Signal); }
+
+void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  void* ipc_ptrs[8];
+  for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
+    ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
+  }
+  fa->register_buffer(ipc_ptrs);
+}
+
+// Use vector<int64_t> to represent byte data for python binding compatibility.
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  auto [handle, offsets] = fa->get_graph_buffer_ipc_meta();
+  std::vector<int64_t> bytes(handle.begin(), handle.end());
+  return std::make_tuple(bytes, offsets);
+}
+
+// Use vector<int64_t> to represent byte data for python binding compatibility.
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  std::vector<std::string> bytes;
+  bytes.reserve(handles.size());
+  for (int i = 0; i < handles.size(); i++) {
+    bytes.emplace_back(handles[i].begin(), handles[i].end());
+  }
+  bytes.reserve(handles.size());
+  fa->register_graph_buffers(bytes, offsets);
+}

csrc/gpu/all_reduce.cuh

@@ -0,0 +1,528 @@
+// 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.
+#pragma once
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <iostream>
+#include <array>
+#include <limits>
+#include <map>
+#include <unordered_map>
+#include <vector>
+
+#define CUDACHECK(cmd)                                              \
+  do {                                                              \
+    cudaError_t e = cmd;                                            \
+    if (e != cudaSuccess) {                                         \
+      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
+             cudaGetErrorString(e));                                \
+      exit(EXIT_FAILURE);                                           \
+    }                                                               \
+  } while (0)
+
+namespace paddle {
+
+constexpr int kMaxBlocks = 36;
+// Counter may overflow, but it's fine since unsigned int overflow is
+// well-defined behavior.
+using FlagType = uint32_t;
+struct Signal {
+  alignas(128) FlagType self_counter[kMaxBlocks][8];
+  // Two sets of peer counters are needed for two syncs. The reason is that
+  // it's possible for peer GPU block to arrive at the second sync point while
+  // the current GPU block haven't passed the first sync point. Thus, peer GPU
+  // may write counter+1 while current GPU is busy waiting for counter. We use
+  // alternating counter array to avoid this possibility.
+  alignas(128) FlagType peer_counter[2][kMaxBlocks][8];
+};
+
+struct __align__(16) RankData {
+  const void* __restrict__ ptrs[8];
+};
+
+struct __align__(16) RankSignals {
+  Signal* signals[8];
+};
+
+// like std::array, but aligned
+template <typename T, int sz>
+struct __align__(alignof(T) * sz) array_t {
+  T data[sz];
+  using type = T;
+  static constexpr int size = sz;
+};
+
+// use packed type to maximize memory efficiency
+// goal: generate ld.128 and st.128 instructions
+template <typename T>
+struct packed_t {
+  // the (P)acked type for load/store
+  using P = array_t<T, 16 / sizeof(T)>;
+  // the (A)ccumulator type for reduction
+  using A = array_t<float, 16 / sizeof(T)>;
+};
+
+#define DINLINE __device__ __forceinline__
+
+// scalar cast functions
+DINLINE float upcast_s(half val) { return __half2float(val); }
+
+template <typename T>
+DINLINE T downcast_s(float val);
+template <>
+DINLINE half downcast_s(float val) {
+  return __float2half(val);
+}
+
+// scalar add functions
+// for some reason when compiling with Paddle, the + operator for half and
+// bfloat is disabled so we call the intrinsics directly
+DINLINE half& assign_add(half& a, half b) {
+  a = __hadd(a, b);
+  return a;
+}
+DINLINE float& assign_add(float& a, float b) { return a += b; }
+
+#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) {
+  return __float2bfloat16(val);
+}
+DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
+  a = __hadd(a, b);
+  return a;
+}
+#endif
+
+template <typename T, int N>
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
+#pragma unroll
+  for (int i = 0; i < N; i++) {
+    assign_add(a.data[i], b.data[i]);
+  }
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
+  if constexpr (std::is_same<T, float>::value) {
+    return val;
+  } else {
+    array_t<float, N> out;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      out.data[i] = upcast_s(val.data[i]);
+    }
+    return out;
+  }
+}
+
+template <typename O>
+DINLINE O downcast(array_t<float, O::size> val) {
+  if constexpr (std::is_same<typename O::type, float>::value) {
+    return val;
+  } else {
+    O out;
+#pragma unroll
+    for (int i = 0; i < O::size; i++) {
+      out.data[i] = downcast_s<typename O::type>(val.data[i]);
+    }
+    return out;
+  }
+}
+
+static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+#else
+  asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+#endif
+}
+
+static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
+  FlagType flag;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+#else
+  asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
+               : "=r"(flag)
+               : "l"(flag_addr));
+#endif
+  return flag;
+}
+
+static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
+  asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
+}
+
+static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
+  FlagType flag;
+  asm volatile("ld.volatile.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  return flag;
+}
+
+// is_start: whether this is the very first synchronization barrier.
+// need_fence: whether a memory fence is needed. If true, a release-acquire
+// semantic is used to enforce memory access order before and after this
+// barrier.
+template <int ngpus, bool is_start, bool need_fence = false>
+DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg,
+                               int rank) {
+  if constexpr (!is_start) __syncthreads();
+  static_assert(
+      !(is_start && need_fence));  // Start barrier shouldn't need fence.
+  if (threadIdx.x < ngpus) {
+    // Increment the counter. Technically we only need one counter, but we use
+    // multiple per block to eliminate the need to share the counter via smem.
+    auto val = self_sg->self_counter[blockIdx.x][threadIdx.x] += 1;
+    // Write the expected counter value to peer and wait for correct value from
+    // peer.
+    auto peer_counter_ptr =
+        &sg.signals[threadIdx.x]->peer_counter[val % 2][blockIdx.x][rank];
+    auto self_counter_ptr =
+        &self_sg->peer_counter[val % 2][blockIdx.x][threadIdx.x];
+    if constexpr (need_fence) {
+      st_flag_release(peer_counter_ptr, val);
+      while (ld_flag_acquire(self_counter_ptr) != val);
+    } else {
+      st_flag_volatile(peer_counter_ptr, val);
+      while (ld_flag_volatile(self_counter_ptr) != val);
+    }
+  }
+  if constexpr (is_start || need_fence) __syncthreads();
+}
+
+template <typename P, int ngpus, typename A>
+DINLINE P packed_reduce(const P* ptrs[], int idx) {
+  A tmp = upcast(ptrs[0][idx]);
+#pragma unroll
+  for (int i = 1; i < ngpus; i++) {
+    packed_assign_add(tmp, upcast(ptrs[i][idx]));
+  }
+  return downcast<P>(tmp);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(512, 1)
+    cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg,
+                               T* __restrict__ result, int rank, int size) {
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  // note: we don't reorder the address so the accumulation order is the same
+  // for all ranks, ensuring bitwise identical results
+  auto dp = *_dp;
+  multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
+  // do the actual reduction
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+       idx += gridDim.x * blockDim.x) {
+    ((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
+  }
+  multi_gpu_barrier<ngpus, false>(sg, self_sg, rank);
+}
+
+template <typename P>
+DINLINE P* get_tmp_buf(Signal* sg) {
+  return (P*)(((Signal*)sg) + 1);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(512, 1)
+    cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg,
+                               T* __restrict__ result, int rank, int size) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  int part = size / ngpus;
+  int start = rank * part;
+  int end = rank == ngpus - 1 ? size : start + part;
+  int largest_part = part + size % ngpus;
+  const P* ptrs[ngpus];
+  P* tmps[ngpus];
+#pragma unroll
+  for (int i = 0; i < ngpus; i++) {
+    int target = (rank + i) % ngpus;
+    ptrs[i] = (const P*)_dp->ptrs[target];
+    tmps[i] = get_tmp_buf<P>(sg.signals[target]);
+  }
+  auto tmp_out = tmps[0];
+  multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
+  // stage 1: reduce scatter
+  for (int idx = start + tid; idx < end; idx += stride) {
+    tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
+  }
+  multi_gpu_barrier<ngpus, false, true>(sg, self_sg, rank);
+
+  // stage 2: allgather. Note: it's important to match the tid between
+  // the two stages, because visibility across devices is only guaranteed
+  // between threads that have the same tid. If thread i computes the sum of
+  // start + i in the first stage, then thread i also gathers start + i from all
+  // ranks.
+  for (int idx = tid; idx < largest_part; idx += stride) {
+#pragma unroll
+    for (int i = 0; i < ngpus; i++) {
+      int gather_from_rank = ((rank + i) % ngpus);
+      if (gather_from_rank == ngpus - 1 || idx < part) {
+        int dst_idx = gather_from_rank * part + idx;
+        ((P*)result)[dst_idx] = tmps[i][idx];
+      }
+    }
+  }
+}
+
+using IPC_KEY = std::array<uint8_t, sizeof(cudaIpcMemHandle_t)>;
+static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t));
+static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t));
+
+class CustomAllreduce {
+ public:
+  int rank_;
+  int world_size_;
+  bool full_nvlink_;
+
+  RankSignals sg_;
+  // Stores an map from a pointer to its peer pointters from all ranks.
+  std::unordered_map<void*, RankData*> buffers_;
+  Signal* self_sg_;
+
+  // Stores rank data from all ranks. This is mainly for cuda graph purposes.
+  // For cuda graph to work, all kernel arguments must be fixed during graph
+  // capture time. However, the peer pointers are not known during graph capture
+  // time. Therefore, during capture, we increment the rank data pointer and use
+  // that as the argument to the kernel. The kernel arguments are stored in
+  // graph_unreg_buffers_. The actual peer pointers will be filled in at the
+  // memory pointed to by the pointers in graph_unreg_buffers_ when
+  // the IPC handles are exchanged between ranks.
+  //
+  // The overall process looks like this:
+  // 1. Graph capture.
+  // 2. Each rank obtains the IPC handles for each addresses used during cuda
+  // graph capture using get_graph_buffer_ipc_meta.
+  // 3. (In Python) all gather the IPC handles.
+  // 4. Obtain the peer pointers by opening the IPC handles, and store them in
+  // the rank data array at corresponding positions.
+  RankData *d_rank_data_base_, *d_rank_data_end_;
+  std::vector<void*> graph_unreg_buffers_;
+  // a map from IPC handles to opened IPC pointers
+  std::map<IPC_KEY, char*> ipc_handles_;
+
+  /**
+   * Signals are an array of ipc-enabled buffers from all ranks.
+   * For each of the buffer, the layout is as follows:
+   * | -- sizeof(Signal) -- | ------ a few MB ----- |
+   * The first section is for allreduce synchronization, and the second section
+   * is for storing the intermediate results required by some allreduce algos.
+   *
+   * Note: this class does not own any device memory. Any required buffers
+   * are passed in from the constructor.
+   */
+  CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz,
+                  int rank, int world_size, bool full_nvlink = true)
+      : rank_(rank),
+        world_size_(world_size),
+        full_nvlink_(full_nvlink),
+        self_sg_(signals[rank]),
+        d_rank_data_base_(reinterpret_cast<RankData*>(rank_data)),
+        d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) {
+    for (int i = 0; i < world_size_; i++) {
+      sg_.signals[i] = signals[i];
+    }
+  }
+
+  char* open_ipc_handle(const void* ipc_handle) {
+    auto [it, new_handle] =
+        ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr});
+    if (new_handle) {
+      char* ipc_ptr;
+      CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptr,
+                                     *((const cudaIpcMemHandle_t*)ipc_handle),
+                                     cudaIpcMemLazyEnablePeerAccess));
+      it->second = ipc_ptr;
+    }
+    return it->second;
+  }
+
+  std::pair<std::string, std::vector<int64_t>> get_graph_buffer_ipc_meta() {
+    auto num_buffers = graph_unreg_buffers_.size();
+    auto handle_sz = sizeof(cudaIpcMemHandle_t);
+    std::string handles(handle_sz * num_buffers, static_cast<char>(0));
+    std::vector<int64_t> offsets(num_buffers);
+    for (int i = 0; i < num_buffers; i++) {
+      auto ptr = graph_unreg_buffers_[i];
+      void* base_ptr;
+      // note: must share the base address of each allocation, or we get wrong
+      // address
+      if (cuPointerGetAttribute(&base_ptr,
+                                CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+                                (CUdeviceptr)ptr) != CUDA_SUCCESS)
+        throw std::runtime_error("failed to get pointer attr");
+      CUDACHECK(cudaIpcGetMemHandle(
+          (cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr));
+      offsets[i] = ((char*)ptr) - ((char*)base_ptr);
+    }
+    return std::make_pair(handles, offsets);
+  }
+
+  void check_rank_data_capacity(size_t num = 1) {
+    if (d_rank_data_base_ + num > d_rank_data_end_)
+      throw std::runtime_error(
+          "Rank data buffer is overflowed by " +
+          std::to_string(d_rank_data_base_ + num - d_rank_data_end_));
+  }
+
+  /**
+   * Register already-shared IPC pointers.
+   */
+  void register_buffer(void** ptrs) {
+    check_rank_data_capacity();
+    RankData data;
+    for (int i = 0; i < world_size_; i++) {
+      data.ptrs[i] = ptrs[i];
+    }
+    auto d_data = d_rank_data_base_++;
+    CUDACHECK(
+        cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice));
+    buffers_[ptrs[rank_]] = d_data;
+  }
+
+  // Note: when registering graph buffers, we intentionally choose to not
+  // deduplicate the addresses. That means if the allocator reuses some
+  // addresses, they will be registered again. This is to account for the remote
+  // possibility of different allocation patterns between ranks. For example,
+  // rank 1 may get the same input address for the second allreduce, but rank 2
+  // got a different address. IPC handles have internal reference counting
+  // mechanism so overhead should be small.
+  void register_graph_buffers(
+      const std::vector<std::string>& handles,
+      const std::vector<std::vector<int64_t>>& offsets) {
+    auto num_buffers = graph_unreg_buffers_.size();
+    check_rank_data_capacity(num_buffers);
+    std::vector<RankData> rank_data(num_buffers);
+    for (int i = 0; i < num_buffers; i++) {
+      auto self_ptr = graph_unreg_buffers_[i];
+      auto& rd = rank_data[i];
+      for (int j = 0; j < world_size_; j++) {
+        if (j != rank_) {
+          char* handle =
+              open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]);
+          handle += offsets[j][i];
+          rd.ptrs[j] = handle;
+        } else {
+          rd.ptrs[j] = self_ptr;
+        }
+      }
+    }
+    CUDACHECK(cudaMemcpy(d_rank_data_base_, rank_data.data(),
+                         sizeof(RankData) * num_buffers,
+                         cudaMemcpyHostToDevice));
+    d_rank_data_base_ += num_buffers;
+    graph_unreg_buffers_.clear();
+  }
+
+  /**
+   * Performs allreduce, assuming input has already been registered.
+   *
+   * Block and grid default configs are results after careful grid search. Using
+   * 36 blocks give the best or close to the best runtime on the devices I
+   * tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also only
+   * take a small amount of SMs. Not quite sure the underlying reason, but my
+   * guess is that too many SMs will cause contention on NVLink bus.
+   */
+  template <typename T>
+  void allreduce(cudaStream_t stream, T* input, T* output, int size,
+                 int threads = 512, int block_limit = 36) {
+    // std::cout<<"input "<< input <<" output "<< output<<std::endl;
+    auto d = packed_t<T>::P::size;
+    if (size % d != 0)
+      throw std::runtime_error(
+          "custom allreduce currently requires input length to be multiple "
+          "of " +
+          std::to_string(d));
+    if (block_limit > kMaxBlocks)
+      throw std::runtime_error("max supported block limit is " +
+                               std::to_string(kMaxBlocks) + ". Got " +
+                               std::to_string(block_limit));
+
+    RankData* ptrs;
+    cudaStreamCaptureStatus status;
+    CUDACHECK(cudaStreamIsCapturing(stream, &status));
+    if (status == cudaStreamCaptureStatusActive) {
+      ptrs = d_rank_data_base_ + graph_unreg_buffers_.size();
+      graph_unreg_buffers_.push_back(input);
+    } else {
+      auto it = buffers_.find(input);
+      if (it == buffers_.end())
+        throw std::runtime_error(
+            "buffer address " +
+            std::to_string(reinterpret_cast<uint64_t>(input)) +
+            " is not registered!");
+      ptrs = it->second;
+    }
+
+    size /= d;
+    auto bytes = size * sizeof(typename packed_t<T>::P);
+    int blocks = std::min(block_limit, (size + threads - 1) / threads);
+#define KL(ngpus, name)                                                       \
+  name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, \
+                                                 rank_, size);
+    // TODO(hanzhi713): Threshold is different for A100 and H100.
+    // Add per device threshold.
+#define REDUCE_CASE(ngpus)                            \
+  case ngpus: {                                       \
+    if (world_size_ == 2) {                           \
+      KL(ngpus, cross_device_reduce_1stage);          \
+    } else if (full_nvlink_) {                        \
+      if ((world_size_ <= 4 && bytes < 512 * 1024) || \
+          (world_size_ <= 8 && bytes < 256 * 1024)) { \
+        KL(ngpus, cross_device_reduce_1stage);        \
+      } else {                                        \
+        KL(ngpus, cross_device_reduce_2stage);        \
+      }                                               \
+    }                                                 \
+    break;                                            \
+  }
+
+    switch (world_size_) {
+      REDUCE_CASE(2)
+      REDUCE_CASE(4)
+      REDUCE_CASE(6)
+      REDUCE_CASE(8)
+      default:
+        throw std::runtime_error(
+            "custom allreduce only supports num gpus in (2,4,6,8). Actual num "
+            "gpus = " +
+            std::to_string(world_size_));
+    }
+#undef REDUCE_CASE
+#undef KL
+  }
+
+  ~CustomAllreduce() {
+    for (auto [_, ptr] : ipc_handles_) {
+      CUDACHECK(cudaIpcCloseMemHandle(ptr));
+    }
+  }
+};
+}  // namespace paddle

csrc/gpu/all_reduce.h

@@ -0,0 +1,30 @@
+
+// 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"
+
+using fptr_t = int64_t;
+fptr_t init_custom_all_reduce(const std::vector<int64_t>& fake_ipc_ptrs,
+                      paddle::Tensor& rank_data, int64_t rank, bool full_nvlink);
+void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
+                fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
+void dispose(fptr_t _fa);
+int64_t meta_size();
+void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa);
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets);

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/decode_attention_kernel.cu

@@ -146,8 +146,11 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
                                                     T * __restrict__ tmp_m, // [batch_size, num_chunks, num_heads]
                                                     T * __restrict__ tmp_d, // [batch_size, num_chunks, num_heads]
                                                     OutT * __restrict__ out) {
+  constexpr uint32_t q_block_num = GROUP_SIZE / bdy;
   const uint32_t bidx = blockIdx.x, kv_head_idx = blockIdx.z;
-  const uint32_t bid = bidx, gid = threadIdx.y;
+  const uint32_t bid = bidx / q_block_num, q_block_idx = bidx % q_block_num;
+  const uint32_t bidy = threadIdx.y;
+  const uint32_t gid = q_block_idx * bdy + threadIdx.y;
   const uint32_t tidx = threadIdx.x;
   constexpr uint32_t num_vec_per_head_qk = HEAD_DIM_QK / VEC_SIZE;
   constexpr uint32_t num_vec_per_head_v = HEAD_DIM_V / VEC_SIZE;
@@ -184,11 +187,10 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   extern __shared__ uint8_t smem[];
   const T *q_now = q + (q_start_idx * q_num_heads + q_head_idx) * HEAD_DIM_QK;
   T *q_smem = reinterpret_cast<T*>(smem); // [HEAD_DIM_QK * sizeof(T)]
-  T *cu_q_smem = q_smem + gid * HEAD_DIM_QK;
+  T *cu_q_smem = q_smem + bidy * HEAD_DIM_QK;
 #pragma unroll
   for(uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
     ((float4*)(&cu_q_smem[vid * VEC_SIZE]))[0] = ((float4*)(&q_now[vid * VEC_SIZE]))[0];
-
   }
   __syncthreads();
   using VecT = AlignedVector<T, VEC_SIZE>;
@@ -203,14 +205,14 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   }
 
 
-  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + GROUP_SIZE * HEAD_DIM_QK * sizeof(CacheT)); 
+  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + bdy * HEAD_DIM_QK * sizeof(CacheT)); 
   uint32_t stage_idx = 0;  
   constexpr int loop_times = DEAL_EACH_TIME / bdy;
 #pragma unroll
   for (int i = 0; i < NUM_STAGES; ++i) {
 #pragma unroll
     for (int j = 0; j < loop_times; ++j) {
-      const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + gid;
+      const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + bidy;
       const uint32_t k_seq_id = chunk_start + k_seq_offset;
       produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
         kv_smem,
@@ -267,7 +269,7 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
 
 #pragma unroll
     for (int j = 0; j < loop_times; ++j) {
-      const uint32_t k_seq_offset = j * bdy + gid;
+      const uint32_t k_seq_offset = j * bdy + bidy;
       produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
         kv_smem,
         cache_k,
@@ -342,10 +344,10 @@ void MultiQueryDecoderAttention(
   constexpr int num_vec_per_head = HEAD_DIM_QK / vec_size;
   constexpr int blockx = num_vec_per_head < 32 ? num_vec_per_head : 32;
 
-  constexpr int blocky = GROUP_SIZE;
-  const int gridx = bsz;
+  constexpr int blocky = GROUP_SIZE > 16 ? 16 : GROUP_SIZE;
+  const int gridx = bsz * (GROUP_SIZE / blocky);
   
-  constexpr int num_threads = blockx * blocky;  
+  constexpr int num_threads = blockx * blocky; 
 
   auto splitkv_kernel = multi_query_decode_attention_kernel<true, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V,
                                                                         BLOCK_SIZE, vec_size, cache_vec_size, blockx, blocky>;
@@ -357,7 +359,7 @@ void MultiQueryDecoderAttention(
   
   const uint32_t chunk_size = get_max_partition_size(bsz);
   const int num_chunks = div_up(max_dec_len, chunk_size);
-  size_t smem_size = cache_smem_bytes + GROUP_SIZE * HEAD_DIM_QK * sizeof(T);
+  size_t smem_size = cache_smem_bytes + blocky * HEAD_DIM_QK * sizeof(T);
 
   if (smem_size >= 48 * 1024) {
     cudaFuncSetAttribute(

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/get_block_shape_and_split_kv_block.cu

@@ -224,7 +224,7 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
   int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
   if (max_dec_len_this_time_data > 0) {
     const bool mla_use_tensorcore = GetMlaUseTensorcore();
-    if (mla_use_tensorcore) {
+    if (mla_use_tensorcore && group_size <=64) {
       const int set_chunk_size = get_mla_dec_chunk_size(bsz);
       decoder_chunk_size_device =
           GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
@@ -296,6 +296,8 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
                                           group_size);
       decoder_num_blocks_x_cpu =
         decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+      decoder_chunk_size_cpu =
+        paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace());
     }
   } else {
     decoder_batch_ids =
@@ -303,7 +305,7 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
     decoder_tile_ids_per_batch =
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
     decoder_chunk_size_cpu =
-        paddle::full({64}, -1, paddle::DataType::INT32, paddle::CPUPlace());
+        paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace());
     decoder_num_blocks_x = 
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
     decoder_num_blocks_x_cpu =
@@ -401,6 +403,7 @@ std::vector<paddle::DataType> GetBlockShapeAndSplitKVBlockInferDtype(
           paddle::DataType::INT32,
           paddle::DataType::INT32,
           paddle::DataType::INT32,
+          paddle::DataType::INT32,
           paddle::DataType::INT32};
 }
 
@@ -423,6 +426,7 @@ std::vector<std::vector<int64_t>> GetBlockShapeAndSplitKVBlockInferShape(
           dynamic_shape,
           {1},
           {1},
+          {1},
           {1}};
 }
 

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

@@ -0,0 +1,365 @@
+// 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/extension.h"
+#include "all_reduce.h"
+
+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");
+  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

@@ -45,6 +45,7 @@ namespace cub = hipcub;
 #include "paddle/phi/backends/gpu/gpu_info.h"
 #include "nlohmann/json.hpp"
 
+
 using json = nlohmann::json;
 
 #define CUDA_CHECK(call)                           \
@@ -161,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];
@@ -244,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/mla_attn/batch_mla_with_paged_kv_cache.cu

@@ -144,7 +144,11 @@ void BatchMLAWithPagedKVCacheKernel(
           params.cumsum_q_seqlens = const_cast<int*>(cu_seqlens_q.data<int>());
           params.padding_offsets = const_cast<int*>(padding_offsets.data<int>());
           params.batch_ids = const_cast<int*>(batch_ids.data<int>());
-          params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+          if (q_head_num <=64) {
+            params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+          } else {
+            params.q_tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+          }
           params.num_blocks_x = const_cast<int*>(num_blocks_x_device.data<int>());
           params.num_blocks_x_int = num_blocks_x;
           params.q_stride_bsz = q_head_num * q_head_dim;
@@ -186,7 +190,11 @@ void BatchMLAWithPagedKVCacheKernel(
         params.cumsum_q_seqlens = const_cast<int*>(cu_seqlens_q.data<int>());
         params.padding_offsets = const_cast<int*>(padding_offsets.data<int>());
         params.batch_ids = const_cast<int*>(batch_ids.data<int>());
-        params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+        if (q_head_num <=64) {
+          params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+        } else {
+          params.q_tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+        }
         params.num_blocks_x = const_cast<int*>(num_blocks_x_device.data<int>());
         params.num_blocks_x_int = num_blocks_x;
         params.q_stride_bsz = q_head_num * q_head_dim;
@@ -217,6 +225,7 @@ void BatchMLAWithPagedKVCacheKernel(
   }
 }
 
+
 template void BatchMLAWithPagedKVCacheKernel<paddle::bfloat16>(
     const AppendAttnMetaData& meta_data,
     const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]

csrc/gpu/mla_attn/epilogue.cuh

@@ -133,7 +133,7 @@ struct CollectiveEpilogue {
                             const int kv_len,
                             const int chunk_size,
                             const int draft_total_token_num,
-                            const int o_stride_bsz) {
+                            const int o_stride_head) {
     const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
     const int num_chunks = cute::ceil_div(kv_len, chunk_size);
     Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
@@ -152,7 +152,7 @@ struct CollectiveEpilogue {
     cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS,
                                       /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
     TiledCopyO gmem_tiled_copy_O;
-    auto O_ptr = num_chunks == 1 ? epilogue_params.O_ptr + (start_token_idx + q_group_offset) * o_stride_bsz : epilogue_params.O_ptr_tmp + (tile_idx * bsz + bid) * draft_total_token_num * o_stride_bsz;
+    auto O_ptr = num_chunks == 1 ? epilogue_params.O_ptr + (start_token_idx * Ktraits::GROUP_SIZE + q_group_offset) * o_stride_head : epilogue_params.O_ptr_tmp + (((tile_idx * bsz + bid) * draft_total_token_num) * Ktraits::GROUP_SIZE + q_group_offset) * o_stride_head;
     Tensor mO = make_tensor(make_gmem_ptr(O_ptr), epilogue_params.layout_O);
     Tensor gO = local_tile(mO, select<0, 1>(TileShape_PDV{}), make_coord(_, _0{}))(_, _, _0{});
     Tensor cO = make_identity_tensor(gO.shape());  // (O, D) -> (o_idx, d_idx)

csrc/gpu/mla_attn/mainloop_load.cuh

@@ -242,8 +242,8 @@ struct CollectiveMainloop {
                              const int bid,
                              const int q_tile_idx) {
     const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
-    int start_q_token_idx = mainloop_params.cumsum_q_seqlens[bid] + q_group_offset;
-    int offset_Q = mainloop_params.q_stride_bsz * start_q_token_idx;
+    int start_q_head_idx = mainloop_params.cumsum_q_seqlens[bid] * Ktraits::GROUP_SIZE + q_group_offset;
+    int offset_Q = mainloop_params.q_stride_head_num * start_q_head_idx;
     Tensor mQ = make_tensor(make_gmem_ptr(mainloop_params.Q_ptr + offset_Q), mainloop_params.layout_Q);
     Tensor gQ =
         local_tile(mQ, select<0, 2>(TileShape_QKD{}), make_coord(_, _0{}))(_, _, _0{});

csrc/gpu/mla_attn/mla_hopper.cuh

@@ -77,7 +77,7 @@ struct Params {
 
     alignas(16) IdType *batch_ids;
     alignas(16) IdType *q_tile_ids_per_batch = nullptr;
-    alignas(16) IdType *tile_ids_per_batch;
+    alignas(16) IdType *tile_ids_per_batch = nullptr;
     alignas(16) IdType *num_blocks_x;
 
 
@@ -117,6 +117,9 @@ struct Params {
   } else if (group_size == 64) {                             \
     constexpr size_t GROUP_SIZE = 64;                        \
     __VA_ARGS__                                              \
+  } else if (group_size == 128) {                            \
+    constexpr size_t GROUP_SIZE = 128;                       \
+    __VA_ARGS__                                              \
   } else {                                                   \
     PD_THROW("not support the group_size: ", group_size);    \
     return cudaErrorNotSupported;                            \
@@ -210,7 +213,7 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
     PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
     const int block_id = blockIdx.x;
     const int bid = mainloop_params.batch_ids[block_id];
-    const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
     const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
     const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
     const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
@@ -268,7 +271,7 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
     clear(tOrO);
     clear(attention_updater.scores_scale);
     const int bid = mainloop_params.batch_ids[block_id];
-    const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
     const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
     const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
     const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
@@ -325,7 +328,7 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
         seq_len_decoder_now,
         mainloop_params.chunk_size,
         mainloop_params.draft_total_token_num,
-        mainloop_params.o_stride_bsz);
+        mainloop_params.o_stride_head_num);
   }
 }
 

csrc/gpu/mla_attn/mla_hopper_wg4.cuh

@@ -149,7 +149,7 @@ MLAWithKVCacheWG4Kernel(CUTE_GRID_CONSTANT
     PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
     const int block_id = blockIdx.x;
     const int bid = mainloop_params.batch_ids[block_id];
-    const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
     const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
     const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
     const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
@@ -202,7 +202,7 @@ MLAWithKVCacheWG4Kernel(CUTE_GRID_CONSTANT
     const int block_id = blockIdx.x;
     clear(attention_updater.scores_scale);
     const int bid = mainloop_params.batch_ids[block_id];
-    const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
     const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
     const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
     const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
@@ -290,7 +290,7 @@ MLAWithKVCacheWG4Kernel(CUTE_GRID_CONSTANT
       seq_len_decoder_now,
       mainloop_params.chunk_size,
       mainloop_params.draft_total_token_num,
-      mainloop_params.o_stride_bsz);
+      mainloop_params.o_stride_head_num);
   }
 }
 

csrc/gpu/moe/fused_moe/cutlass_extensions/gemm/kernel/gemm_moe_problem_visitor.h

@@ -0,0 +1,93 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief Scheduler for grouped GEMM
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/kernel/gemm_grouped_problem_visitor.h"
+#include "cutlass/matrix_coord.h"
+
+#include "cutlass_extensions/gemm/kernel/moe_problem_visitor.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/// Visitor class to abstract away the algorithm for iterating over tiles
+template <typename ThreadblockShape,
+          GroupScheduleMode GroupScheduleMode_,
+          int PrefetchTileCount,
+          int ThreadCount,
+          bool Transposed = false>
+struct GemmMoeProblemVisitor
+    : public MoeProblemVisitor<
+          detail::GemmGroupedProblemSizeHelper<ThreadblockShape, Transposed>,
+          ThreadblockShape,
+          GroupScheduleMode_,
+          PrefetchTileCount,
+          ThreadCount> {
+  static bool const kTransposed = Transposed;
+
+  using ProblemSizeHelper =
+      detail::GemmGroupedProblemSizeHelper<ThreadblockShape, Transposed>;
+  using Base = MoeProblemVisitor<ProblemSizeHelper,
+                                 ThreadblockShape,
+                                 GroupScheduleMode_,
+                                 PrefetchTileCount,
+                                 ThreadCount>;
+  using Params = typename Base::Params;
+  using SharedStorage = typename Base::SharedStorage;
+
+  //
+  // Methods
+  //
+  CUTLASS_DEVICE
+  GemmMoeProblemVisitor(Params const& params_,
+                        SharedStorage& shared_storage_,  // NOLINT
+                        int32_t block_idx)
+      : Base(params_, shared_storage_, block_idx) {}
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

csrc/gpu/moe/fused_moe/cutlass_extensions/gemm/kernel/moe_problem_visitor.h

@@ -0,0 +1,346 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief Base scheduler for grouped problems, using MoE
+*/
+
+#pragma once
+
+#include "cutlass/gemm/kernel/grouped_problem_visitor.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Visitor class to abstract away the algorithm for iterating over tiles
+template <typename ProblemSizeHelper, typename ThreadblockShape_>
+struct BaseMoeProblemVisitor {
+  using ThreadblockShape = ThreadblockShape_;
+
+  struct ProblemInfo {
+    static int32_t const kNoPrefetchEntry = -1;
+    int32_t problem_idx;
+    int32_t problem_start;
+
+    CUTLASS_DEVICE
+    ProblemInfo()
+        : problem_idx(kNoPrefetchEntry), problem_start(kNoPrefetchEntry) {}
+
+    CUTLASS_DEVICE
+    ProblemInfo(int32_t problem_idx_, int32_t problem_start_)
+        : problem_idx(problem_idx_), problem_start(problem_start_) {}
+  };
+
+  struct Params {
+    int64_t const* last_row_for_problem;
+    int64_t gemm_n;
+    int64_t gemm_k;
+    int32_t problem_count;
+    void const* workspace;
+    int32_t tile_count;
+
+    //
+    // Methods
+    //
+
+    /// Ctor
+    CUTLASS_HOST_DEVICE
+    Params()
+        : last_row_for_problem(nullptr),
+          gemm_n(0),
+          gemm_k(0),
+          problem_count(0),
+          workspace(nullptr),
+          tile_count(0) {}
+
+    /// Ctor
+    CUTLASS_HOST_DEVICE
+    Params(int64_t const* last_row_for_problem,
+           int64_t gemm_n,
+           int64_t gemm_k,
+           int32_t problem_count,
+           void const* workspace = nullptr,
+           int32_t tile_count = 0)
+        : last_row_for_problem(last_row_for_problem),
+          gemm_n(gemm_n),
+          gemm_k(gemm_k),
+          problem_count(problem_count),
+          workspace(workspace),
+          tile_count(tile_count) {}
+  };
+
+  Params const& params;
+  int32_t tile_idx;
+  int32_t problem_tile_start;
+  int32_t problem_idx;
+
+  //
+  // Methods
+  //
+  CUTLASS_DEVICE
+  BaseMoeProblemVisitor(Params const& params_, int32_t block_idx)
+      : params(params_),
+        tile_idx(block_idx),
+        problem_tile_start(0),
+        problem_idx(0) {}
+
+  /// Get the grid shape
+  CUTLASS_HOST_DEVICE
+  static cutlass::gemm::GemmCoord grid_shape(
+      const cutlass::gemm::GemmCoord& problem) {
+    return cutlass::gemm::GemmCoord(
+        ((problem.m() - 1 + ThreadblockShape::kM) / ThreadblockShape::kM),
+        ((problem.n() - 1 + ThreadblockShape::kN) / ThreadblockShape::kN),
+        1);
+  }
+
+  /// Gets the global tile index
+  CUTLASS_HOST_DEVICE
+  int32_t tile_index() const { return tile_idx; }
+
+  /// Gets the index of the problem
+  CUTLASS_HOST_DEVICE
+  int32_t problem_index() const { return problem_idx; }
+
+  CUTLASS_HOST_DEVICE
+  int32_t threadblock_idx() const { return tile_idx - problem_tile_start; }
+
+  CUTLASS_DEVICE
+  void advance(int32_t grid_size) { tile_idx += grid_size; }
+
+  CUTLASS_HOST_DEVICE
+  static void possibly_transpose_problem(
+      cutlass::gemm::GemmCoord& problem) {  // NOLINT
+    ProblemSizeHelper::possibly_transpose_problem(problem);
+  }
+
+  /// Returns the problem size for the current problem
+  CUTLASS_HOST_DEVICE
+  cutlass::gemm::GemmCoord problem_size() const {
+    return problem_size(problem_idx);
+  }
+
+  CUTLASS_HOST_DEVICE
+  cutlass::gemm::GemmCoord problem_size(int idx) const {
+    const int64_t prev_problem_row =
+        idx == 0 ? 0 : params.last_row_for_problem[idx - 1];
+    const int64_t current_problem_row = params.last_row_for_problem[idx];
+    const int64_t gemm_m = current_problem_row - prev_problem_row;
+    GemmCoord problem(GemmCoord::Index(gemm_m),
+                      GemmCoord::Index(params.gemm_n),
+                      GemmCoord::Index(params.gemm_k));
+    ProblemSizeHelper::possibly_transpose_problem(problem);
+    return problem;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static int32_t tile_count(const cutlass::gemm::GemmCoord& grid) {
+    return ProblemSizeHelper::tile_count(grid);
+  }
+
+  static int32_t group_tile_count(
+      const cutlass::gemm::GemmCoord* host_problem_sizes_ptr,
+      int32_t problem_count) {
+    int32_t total_tiles = 0;
+    for (int32_t i = 0; i < problem_count; ++i) {
+      auto problem = host_problem_sizes_ptr[i];
+      possibly_transpose_problem(problem);
+      auto grid = grid_shape(problem);
+      total_tiles += tile_count(grid);
+    }
+
+    return total_tiles;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename ProblemSizeHelper,
+          typename ThreadblockShape,
+          GroupScheduleMode GroupScheduleMode_,
+          int PrefetchTileCount,
+          int ThreadCount>
+struct MoeProblemVisitor;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// ProblemVisitor that performs all scheduling on device
+//
+template <typename ProblemSizeHelper,
+          typename ThreadblockShape,
+          int PrefetchTileCount,
+          int ThreadCount>
+struct MoeProblemVisitor<ProblemSizeHelper,
+                         ThreadblockShape,
+                         GroupScheduleMode::kDeviceOnly,
+                         PrefetchTileCount,
+                         ThreadCount>
+    : public BaseMoeProblemVisitor<ProblemSizeHelper, ThreadblockShape> {
+  using Base = BaseMoeProblemVisitor<ProblemSizeHelper, ThreadblockShape>;
+  using Params = typename Base::Params;
+  static int const kThreadCount = ThreadCount;
+  static bool const kRequiresPrecomputation = false;
+  static int const kThreadsPerWarp = 32;
+
+  struct SharedStorage {};
+
+  // Final tile of the problem loaded by this thread. Each thread will hold
+  // a separate value.
+  int32_t problem_ending_tile;
+
+  SharedStorage& shared_storage;
+
+  //
+  // Methods
+  //
+  CUTLASS_DEVICE
+  MoeProblemVisitor(Params const& params_,
+                    SharedStorage& shared_storage_,  // NOLINT
+                    int32_t block_idx)
+      : Base(params_, block_idx),
+        problem_ending_tile(0),
+        shared_storage(shared_storage_) {
+    this->problem_idx = -1 * kThreadsPerWarp;
+    this->problem_tile_start = 0;
+  }
+
+  CUTLASS_DEVICE
+  bool next_tile() {
+    // Check whether the tile to compute is within the range of the current
+    // problem.
+    int32_t problem_tile_end = __shfl_sync(
+        0xffffffff, problem_ending_tile, this->problem_idx % kThreadsPerWarp);
+    if (this->tile_idx < problem_tile_end) {
+      return true;
+    }
+
+    // Check whether the tile to compute is within the current group of problems
+    // fetched by the warp. The last tile for this group is the final tile of
+    // the problem held by the final thread in the warp.
+    int32_t group_tile_end =
+        __shfl_sync(0xffffffff, problem_ending_tile, kThreadsPerWarp - 1);
+
+    // Keep the starting problem for this group in `problem_idx`. This is done
+    // to reduce register pressure. The starting problem for this group is
+    // simply the first problem in the group most recently fetched by the warp.
+    int32_t& group_problem_start = this->problem_idx;
+    group_problem_start =
+        (this->problem_idx / kThreadsPerWarp) * kThreadsPerWarp;
+
+    // Keep the starting tile for this group in `problem_tile_start`. This is
+    // done to reduce register pressure.
+    int32_t& group_tile_start = this->problem_tile_start;
+
+    // Each thread in the warp processes a separate problem to advance until
+    // reaching a problem whose starting tile is less less than tile_idx.
+    while (group_tile_end <= this->tile_idx) {
+      group_problem_start += kThreadsPerWarp;
+      if (group_problem_start > this->params.problem_count) {
+        return false;
+      }
+
+      // Since `group_tile_start` is a reference to `this->problem_tile_start`,
+      // this also sets `this->problem_tile_start`. The fact that
+      // `this->problem_tile_start` is also set here is used later in
+      // `next_tile`.
+      group_tile_start = group_tile_end;
+
+      int lane_idx = threadIdx.x % kThreadsPerWarp;
+      int32_t lane_problem = group_problem_start + lane_idx;
+
+      // Compute the number of tiles in the problem assigned to each thread.
+      problem_ending_tile = 0;
+      if (lane_problem < this->params.problem_count) {
+        cutlass::gemm::GemmCoord problem = this->problem_size(lane_problem);
+        cutlass::gemm::GemmCoord grid = this->grid_shape(problem);
+        problem_ending_tile = this->tile_count(grid);
+      }
+
+      // Compute a warp-wide inclusive prefix sum to compute the ending tile
+      // index of each thread's problem.
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 1; i < kThreadsPerWarp; i <<= 1) {
+        int32_t val = __shfl_up_sync(0xffffffff, problem_ending_tile, i);
+        if (lane_idx >= i) {
+          problem_ending_tile += val;
+        }
+      }
+
+      // The total tile count for this group is now in the final position of the
+      // prefix sum
+      int32_t tiles_in_group =
+          __shfl_sync(0xffffffff, problem_ending_tile, kThreadsPerWarp - 1);
+
+      problem_ending_tile += group_tile_start;
+      group_tile_end += tiles_in_group;
+    }
+
+    // The next problem to process is the first one that does not have ending
+    // tile position that is greater than or equal to tile index.
+    int32_t problem_idx_in_group = __popc(
+        __ballot_sync(0xffffffff, problem_ending_tile <= this->tile_idx));
+
+    this->problem_idx = group_problem_start + problem_idx_in_group;
+
+    // The starting tile for this problem is the ending tile of the previous
+    // problem. In cases where `problem_idx_in_group` is the first problem in
+    // the group, we do not need to reset `problem_tile_start`, because it is
+    // set to the previous group's ending tile in the while loop above.
+    if (problem_idx_in_group > 0) {
+      this->problem_tile_start = __shfl_sync(
+          0xffffffff, problem_ending_tile, problem_idx_in_group - 1);
+    }
+
+    return true;
+  }
+
+  static size_t get_workspace_size(
+      const cutlass::gemm::GemmCoord* host_problem_sizes_ptr,
+      int32_t problem_count,
+      int32_t block_count) {
+    return 0;
+  }
+
+  static void host_precompute(
+      const cutlass::gemm::GemmCoord* host_problem_sizes_ptr,
+      int32_t problem_count,
+      int32_t block_count,
+      void* host_workspace_ptr) {}
+};
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass

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

@@ -0,0 +1,611 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ *modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *notice, this list of conditions and the following disclaimer in the
+ *documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its
+ *contributors may be used to endorse or promote products derived from this
+ *software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
+ *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+ *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+ *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief
+*/
+
+#pragma once
+
+#include "cutlass/complex.h"
+#include "cutlass/cutlass.h"
+#include "cutlass/fast_math.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/matrix_coord.h"
+#include "cutlass/semaphore.h"
+
+#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/trace.h"
+
+#include "cutlass_extensions/gemm/kernel/gemm_moe_problem_visitor.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/tile_interleaved_layout.h"
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// This section exists to that we can use the same kernel code for regular gemm
+// and dequantizing gemms. It will dispatch to the dequantizing gemm if the Mma
+// type has an Iterator for scales in global.
+template <typename...>
+using void_t = void;
+
+template <typename Mma, typename = void>
+struct use_dq_gemm : platform::false_type {};
+
+template <typename Mma>
+struct use_dq_gemm<Mma, void_t<typename Mma::IteratorScale>>
+    : platform::true_type {};
+
+// SFINAE overload for dequantizing gemm
+template <
+    typename Mma,
+    typename ElementScale,
+    typename platform::enable_if<use_dq_gemm<Mma>::value, bool>::type = true>
+CUTLASS_DEVICE static void run_mma(Mma mma,
+                                   int gemm_k_iterations,
+                                   typename Mma::FragmentC& accum,  // NOLINT
+                                   typename Mma::IteratorA iterator_A,
+                                   typename Mma::IteratorB iterator_B,
+                                   typename Mma::FragmentC const& src_accum,
+                                   ElementScale* weight_scale_ptr,
+                                   MatrixCoord scale_extent,
+                                   const int thread_idx,
+                                   MatrixCoord tb_offset_scale) {
+  typename Mma::IteratorScale iterator_scale(
+      Mma::IteratorScale::Layout(scale_extent.column()),
+      weight_scale_ptr,
+      scale_extent,
+      thread_idx,
+      tb_offset_scale);
+
+  mma(gemm_k_iterations,
+      accum,
+      iterator_A,
+      iterator_B,
+      iterator_scale,
+      src_accum);
+}
+
+// SFINAE overload for normal gemm. This completely ignores the scale parameters
+template <
+    typename Mma,
+    typename ElementScale,
+    typename platform::enable_if<!use_dq_gemm<Mma>::value, bool>::type = true>
+CUTLASS_DEVICE static void run_mma(Mma mma,
+                                   int gemm_k_iterations,
+                                   typename Mma::FragmentC& accum,  // NOLINT
+                                   typename Mma::IteratorA iterator_A,
+                                   typename Mma::IteratorB iterator_B,
+                                   typename Mma::FragmentC const& src_accum,
+                                   ElementScale* weight_scale_ptr,
+                                   MatrixCoord scale_extent,
+                                   const int thread_idx,
+                                   MatrixCoord tb_offset_scale) {
+  mma(gemm_k_iterations, accum, iterator_A, iterator_B, src_accum);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
+          typename Epilogue_,            ///! Epilogue
+          typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
+          typename KernelArch,  ///! The Architecture this kernel is compiled
+                                /// for. Used since SIMT kernels lose top-level
+                                /// arch.
+          GroupScheduleMode GroupScheduleMode_  ///! Type of scheduling to //
+                                                /// NOLINT perform
+          >
+struct MoeFCGemm {
+ public:
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueOutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_;
+  static bool const kTransposed = false;
+
+  // Optional transpose
+  using MapArguments =
+      kernel::detail::MapArguments<typename Mma::IteratorA::Element,
+                                   typename Mma::IteratorA::Layout,
+                                   Mma::kTransformA,
+                                   Mma::IteratorA::AccessType::kElements,
+                                   typename Mma::IteratorB::Element,
+                                   typename Mma::IteratorB::Layout,
+                                   Mma::kTransformB,
+                                   Mma::IteratorB::AccessType::kElements,
+                                   typename Mma::LayoutC,
+                                   kTransposed>;
+
+  // Public-facing type definitions related to operand element type, layout, and
+  // complex conjugate operation. Must interact with the 'kTransposed' notion.
+  static_assert(!kTransposed, "Transpose problem not supported");
+  using ElementA = typename MapArguments::ElementA;
+  using LayoutA = typename MapArguments::LayoutA;
+  using ElementB = typename MapArguments::ElementB;
+  using LayoutB = typename MapArguments::LayoutB;
+  using ElementC = typename Epilogue::OutputTileIterator::Element;
+  using LayoutC = typename MapArguments::LayoutC;
+  using ElementScale = ElementC;
+
+  static ComplexTransform const kTransformA = MapArguments::kTransformA;
+  static ComplexTransform const kTransformB = MapArguments::kTransformB;
+
+  // Type definitions about the mainloop.
+  using Operator = typename Mma::Operator;
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
+
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = MapArguments::kAlignmentA;
+  static int const kAlignmentB = MapArguments::kAlignmentB;
+  static int const kAlignmentC =
+      Epilogue::OutputTileIterator::kElementsPerAccess;
+
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
+
+  using ProblemVisitor = GemmMoeProblemVisitor<ThreadblockShape,
+                                               kGroupScheduleMode,
+                                               kThreadCount,
+                                               kThreadCount,
+                                               kTransposed>;
+
+  //
+  // Structures
+  //
+
+  /// Argument structure
+  struct Arguments {
+    //
+    // Data members
+    //
+
+    int problem_count;
+    int threadblock_count;
+
+    typename EpilogueOutputOp::Params output_op;
+
+    ElementA* ptr_A;
+    ElementB* ptr_B;
+    ElementScale* weight_scales;
+    ElementC* ptr_C;
+    ElementC* ptr_D;
+
+    int64_t* total_rows_before_expert;
+    int64_t gemm_n;
+    int64_t gemm_k;
+
+    // Only used by device-level operator
+    GemmCoord* host_problem_sizes;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments()
+        : problem_count(0),
+          threadblock_count(0),
+          ptr_A(nullptr),
+          ptr_B(nullptr),
+          weight_scales(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr),
+          total_rows_before_expert(nullptr),
+          gemm_n(0),
+          gemm_k(0),
+          host_problem_sizes(nullptr) {}
+
+    /// Ctor
+    CUTLASS_HOST_DEVICE
+    Arguments(int problem_count,
+              int threadblock_count,
+              typename EpilogueOutputOp::Params output_op,
+              const ElementA* ptr_A,
+              const ElementB* ptr_B,
+              const ElementScale* weight_scales,
+              const ElementC* ptr_C,
+              ElementC* ptr_D,
+              int64_t* total_rows_before_expert,
+              int64_t gemm_n,
+              int64_t gemm_k,
+              GemmCoord* host_problem_sizes = nullptr)
+        : problem_count(problem_count),
+          threadblock_count(threadblock_count),
+          output_op(output_op),
+          ptr_A(const_cast<ElementA*>(ptr_A)),
+          ptr_B(const_cast<ElementB*>(ptr_B)),
+          weight_scales(const_cast<ElementScale*>(weight_scales)),
+          ptr_C(const_cast<ElementC*>(ptr_C)),
+          ptr_D(ptr_D),
+          total_rows_before_expert(total_rows_before_expert),
+          gemm_n(gemm_n),
+          gemm_k(gemm_k),
+          host_problem_sizes(nullptr) {
+      if (platform::is_same<uint8_t, ElementB>::value ||
+          platform::is_same<uint4b_t, ElementB>::value) {
+        assert(weight_scales);
+      }
+    }
+  };
+
+  //
+  // Structure for precomputing values in host memory and passing to kernels
+  //
+
+  /// Parameters structure
+  struct Params {
+    typename ProblemVisitor::Params problem_visitor;
+    int threadblock_count;
+
+    typename EpilogueOutputOp::Params output_op;
+
+    ElementA* ptr_A;
+    ElementB* ptr_B;
+    ElementScale* weight_scales;
+    ElementC* ptr_C;
+    ElementC* ptr_D;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : ptr_A(nullptr),
+          ptr_B(nullptr),
+          weight_scales(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(Arguments const& args,
+           void* workspace = nullptr,
+           int tile_count = 0)  // NOLINT
+        : problem_visitor(args.total_rows_before_expert,
+                          args.gemm_n,
+                          args.gemm_k,
+                          args.problem_count,
+                          workspace,
+                          tile_count),
+          threadblock_count(args.threadblock_count),
+          output_op(args.output_op),
+          ptr_A(args.ptr_A),
+          ptr_B(args.ptr_B),
+          weight_scales(args.weight_scales),
+          ptr_C(args.ptr_C),
+          ptr_D(args.ptr_D) {}
+
+    CUTLASS_HOST_DEVICE
+    void update(Arguments const& args,
+                void* workspace = nullptr,
+                int tile_count = 0) {
+      problem_visitor =
+          typename ProblemVisitor::Params(args.total_rows_before_expert,
+                                          args.gemm_n,
+                                          args.gemm_k,
+                                          args.problem_count,
+                                          workspace,
+                                          tile_count);
+      threadblock_count = args.threadblock_count;
+      output_op = args.output_op;
+      ptr_A = args.ptr_A;
+      ptr_B = args.ptr_B;
+      weight_scales = args.weight_scales;
+      ptr_C = args.ptr_C;
+      ptr_D = args.ptr_D;
+    }
+  };
+
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename ProblemVisitor::SharedStorage problem_visitor;
+    typename Mma::SharedStorage main_loop;
+    typename Epilogue::SharedStorage epilogue;
+  };
+
+ public:
+  //
+  // Methods
+  //
+
+  CUTLASS_DEVICE
+  MoeFCGemm() {}
+
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
+    return Status::kSuccess;
+  }
+
+  static Status can_implement(Arguments const& args) {
+    if (platform::is_same<uint8_t, ElementB>::value ||
+        platform::is_same<uint4b_t, ElementB>::value) {
+      if (args.weight_scales == nullptr) {
+        CUTLASS_TRACE_HOST(
+            "MoeFCGemm::can_implement() - weight scales are required for "
+            "uint8_t and uint4b_t");
+        return Status::kInvalid;
+      }
+    } else if (args.weight_scales != nullptr) {
+      CUTLASS_TRACE_HOST(
+          "MoeFCGemm::can_implement() - weight scales are ignored for all "
+          "types except uint8_t and uint4b_t");
+      return Status::kInvalid;
+    }
+    return Status::kSuccess;
+  }
+
+  static size_t get_extra_workspace_size(
+      Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape) {
+    return 0;
+  }
+
+  // The dummy template parameter is not used and exists so that we can compile
+  // this code using a standard earlier than C++17. Prior to C++17, fully
+  // specialized templates HAD to exists in a namespace
+  template <bool B, typename dummy = void>
+  struct KernelRunner {
+    CUTLASS_DEVICE
+    static void run_kernel(Params const& params,
+                           SharedStorage& shared_storage) {  // NOLINT
+      CUTLASS_NOT_IMPLEMENTED();
+    }
+  };
+
+  template <typename dummy>
+  struct KernelRunner<true, dummy> {
+    CUTLASS_DEVICE
+    static void run_kernel(Params const& params,
+                           SharedStorage& shared_storage) {  // NOLINT
+      //
+      // These types shadow the type-level definitions and support the ability
+      // to implement a 'transposed' GEMM that computes the transposed problems.
+      //
+      using ElementA = typename Mma::IteratorA::Element;
+      using LayoutA = typename Mma::IteratorA::Layout;
+      using ElementB = typename Mma::IteratorB::Element;
+      using LayoutB = typename Mma::IteratorB::Layout;
+      using ElementC = typename Epilogue::OutputTileIterator::Element;
+      using LayoutC = typename Epilogue::OutputTileIterator::Layout;
+      static constexpr int kInterleave =
+          Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
+      static_assert(
+          platform::is_same<LayoutB, layout::RowMajor>::value &&
+                  kInterleave == 1 ||
+              platform::is_same<LayoutB, layout::ColumnMajor>::value &&
+                  kInterleave >= 1,
+          "B must be row major/col major OR col major interleaved.");
+
+      //
+      // Problem visitor.
+      //
+      ProblemVisitor problem_visitor(
+          params.problem_visitor, shared_storage.problem_visitor, blockIdx.x);
+
+      const int64_t gemm_k = params.problem_visitor.gemm_k;
+      const int64_t gemm_n = params.problem_visitor.gemm_n;
+      int64_t bytes_per_expert_matrix =
+          (gemm_k * gemm_n / 8) * cutlass::sizeof_bits<ElementB>::value;
+
+      // Outer 'persistent' loop to iterate over tiles
+      while (problem_visitor.next_tile()) {
+        GemmCoord problem_size = problem_visitor.problem_size();
+        int32_t problem_idx = problem_visitor.problem_index();
+        int32_t cta_idx = int32_t(problem_visitor.threadblock_idx());
+
+        GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
+
+        cutlass::gemm::GemmCoord threadblock_offset(
+            int(cta_idx / grid_shape.n()) * Mma::Shape::kM,  // NOLINT
+            int(cta_idx % grid_shape.n()) * Mma::Shape::kN,  // NOLINT
+            0);
+
+        // Load element pointers. Exchange pointers and strides if working on
+        // the transpose
+        const int64_t rows_to_jump =
+            problem_idx == 0
+                ? 0
+                : params.problem_visitor.last_row_for_problem[problem_idx - 1];
+        ElementA* ptr_A =
+            reinterpret_cast<ElementA*>(params.ptr_A) + rows_to_jump * gemm_k;
+        typename LayoutA::LongIndex ldm_A = gemm_k;
+
+        char* byte_ptr_B = ((char*)params.ptr_B) +                 // NOLINT
+                           problem_idx * bytes_per_expert_matrix;  // NOLINT
+        ElementB* ptr_B = reinterpret_cast<ElementB*>(byte_ptr_B);
+        typename LayoutB::LongIndex ldm_B =
+            platform::is_same<layout::RowMajor, LayoutB>::value
+                ? gemm_n
+                : gemm_k * kInterleave;
+
+        // Compute initial location in logical coordinates
+        cutlass::MatrixCoord tb_offset_A{
+            threadblock_offset.m(),
+            0,
+        };
+
+        cutlass::MatrixCoord tb_offset_B{0,
+                                         threadblock_offset.n() / kInterleave};
+
+        cutlass::MatrixCoord tb_offset_scale{0, threadblock_offset.n()};
+
+        // Compute position within threadblock
+        int thread_idx = threadIdx.x;
+
+        // Construct iterators to A and B operands
+        typename Mma::IteratorA iterator_A(LayoutA(ldm_A),
+                                           ptr_A,
+                                           {problem_size.m(), problem_size.k()},
+                                           thread_idx,
+                                           tb_offset_A);
+
+        typename Mma::IteratorB iterator_B(
+            LayoutB(ldm_B),
+            ptr_B,
+            {problem_size.k() * kInterleave, problem_size.n() / kInterleave},
+            thread_idx,
+            tb_offset_B);
+
+        typename Mma::FragmentC accumulators;
+
+        accumulators.clear();
+
+        // Broadcast the warp_id computed by lane 0 to ensure dependent code
+        // is compiled as warp-uniform.
+        int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+
+        int lane_idx = threadIdx.x % 32;
+
+        //
+        // Matrix multiply phase
+        //
+
+        // Construct thread-scoped matrix multiply
+        auto CreateMMA = [&]() {
+          if constexpr (use_dq_gemm<Mma>::value)
+            return Mma(
+                shared_storage.main_loop, -1, thread_idx, warp_idx, lane_idx);
+          else
+            return Mma(
+                shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+        };
+        Mma mma = CreateMMA();
+
+        // Compute threadblock-scoped matrix multiply-add
+        int gemm_k_iterations =
+            (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK;
+
+        // Wait for all threads to finish their epilogue phases from the
+        // previous tile.
+        __syncthreads();
+
+        // Compute threadblock-scoped matrix multiply-add
+        ElementScale* weight_scale_ptr =
+            params.weight_scales + problem_idx * problem_size.n();
+        run_mma<Mma>(mma,
+                     gemm_k_iterations,
+                     accumulators,
+                     iterator_A,
+                     iterator_B,
+                     accumulators,
+                     weight_scale_ptr,
+                     {1, problem_size.n()},
+                     thread_idx,
+                     tb_offset_scale);
+
+        //
+        // Epilogue
+        //
+
+        EpilogueOutputOp output_op(params.output_op);
+
+        ElementC* ptr_C =
+            reinterpret_cast<ElementC*>(params.ptr_C) + problem_idx * gemm_n;
+        ElementC* ptr_D =
+            reinterpret_cast<ElementC*>(params.ptr_D) + rows_to_jump * gemm_n;
+
+        LayoutC layout_C(0);
+        LayoutC layout_D(gemm_n);
+
+        typename Epilogue::OutputTileIterator::Params params_C(layout_C);
+        typename Epilogue::OutputTileIterator::Params params_D(layout_D);
+
+        // Tile iterator loading from source tensor.
+        typename Epilogue::OutputTileIterator iterator_C(
+            params_C,
+            ptr_C,
+            problem_size.mn(),
+            thread_idx,
+            threadblock_offset.mn());
+
+        // Tile iterator writing to destination tensor.
+        typename Epilogue::OutputTileIterator iterator_D(
+            params_D,
+            ptr_D,
+            problem_size.mn(),
+            thread_idx,
+            threadblock_offset.mn());
+
+        Epilogue epilogue(
+            shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
+
+        // Execute the epilogue operator to update the destination tensor.
+        epilogue(output_op, iterator_D, accumulators, iterator_C);
+
+        // Next tile
+        problem_visitor.advance(gridDim.x);
+      }
+    }
+  };
+
+  /*
+    To improve compilation speed, we do not compile the device operator if the
+    CUDA_ARCH does not correspond to the ArchTag of the cutlass kernel operator.
+  */
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const& params,
+                  SharedStorage& shared_storage) {  // NOLINT
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700) && (__CUDA_ARCH__ < 750)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm70>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750) && (__CUDA_ARCH__ < 800)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm75>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && (__CUDA_ARCH__ < 900)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm80>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 910)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm80>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#else
+    CUTLASS_NOT_IMPLEMENTED();
+#endif
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

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

@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+
+#pragma once
+#include <cuda_runtime_api.h>
+#include <string>
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/ft_gemm_configs.h"
+using namespace phi;
+
+template <typename T, /*The type used for activations/scales/compute*/
+          typename WeightType /* The type for the MoE weights */>
+class MoeGemmRunner {
+ public:
+  MoeGemmRunner();
+
+  void moe_gemm_bias_act(const T* A,
+                         const WeightType* B,
+                         const T* weight_scales,
+                         const T* biases,
+                         T* C,
+                         int64_t* total_rows_before_expert,
+                         int64_t total_rows,
+                         int64_t gemm_n,
+                         int64_t gemm_k,
+                         int num_experts,
+                         std::string activation_type,
+                         cudaStream_t stream);
+
+  void moe_gemm(const T* A,
+                const WeightType* B,
+                const T* weight_scales,
+                T* C,
+                int64_t* total_rows_before_expert,
+                int64_t total_rows,
+                int64_t gemm_n,
+                int64_t gemm_k,
+                int num_experts,
+                cudaStream_t stream);
+
+ private:
+  template <typename EpilogueTag>
+  void dispatch_to_arch(const T* A,
+                        const WeightType* B,
+                        const T* weight_scales,
+                        const T* biases,
+                        T* C,
+                        int64_t* total_rows_before_expert,
+                        int64_t total_rows,
+                        int64_t gemm_n,
+                        int64_t gemm_k,
+                        int num_experts,
+                        CutlassGemmConfig gemm_config,
+                        cudaStream_t stream,
+                        int* occupancy = nullptr);
+
+  template <typename EpilogueTag>
+  void run_gemm(const T* A,
+                const WeightType* B,
+                const T* weight_scales,
+                const T* biases,
+                T* C,
+                int64_t* total_rows_before_expert,
+                int64_t total_rows,
+                int64_t gemm_n,
+                int64_t gemm_k,
+                int num_experts,
+                cudaStream_t stream);
+
+ private:
+  int sm_;
+  int multi_processor_count_;
+};

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_bf16.cu

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+ 
+#pragma once
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+#ifdef PADDLE_CUDA_BF16
+template class MoeGemmRunner<__nv_bfloat16, __nv_bfloat16>;
+#endif

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int4.cu

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+#pragma once
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+#ifdef PADDLE_CUDA_BF16
+template class MoeGemmRunner<__nv_bfloat16, cutlass::uint4b_t>;
+#endif
+

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int8.cu

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+
+#pragma once
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+#ifdef PADDLE_CUDA_BF16
+template class MoeGemmRunner<__nv_bfloat16, uint8_t>;
+#endif
+

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_fp16.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+
+#pragma once
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+template class MoeGemmRunner<half, half>;

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int4.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+
+#pragma once
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+template class MoeGemmRunner<half, cutlass::uint4b_t>;

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int8.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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.
+ */
+
+#pragma once
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+template class MoeGemmRunner<half, uint8_t>;

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

@@ -0,0 +1,819 @@
+/*
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.  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.
+ */
+
+// Ignore CUTLASS warnings about type punning
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <math.h>
+
+#include <optional>
+#include <sstream>
+
+#include "cutlass/array.h"
+#include "cutlass/gemm/device/gemm_grouped.h"
+#include "cutlass/gemm/kernel/default_gemm_grouped.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_cutlass_kernel.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/compute_occupancy.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/epilogue_helpers.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/threadblock/default_mma.h"
+
+#pragma GCC diagnostic pop
+
+#include "helper.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_kernels/cutlass_heuristic.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_kernels/gemm_config_manager.h"
+
+using namespace phi;
+// ============================= Variable batched Gemm things
+// ===========================
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages>
+void generic_moe_gemm_kernelLauncher(const T* A,
+                                     const WeightType* B,
+                                     const T* weight_scales,
+                                     const T* biases,
+                                     T* C,
+                                     int64_t* total_rows_before_expert,
+                                     int64_t gemm_n,
+                                     int64_t gemm_k,
+                                     int num_experts,
+                                     CutlassGemmConfig gemm_config,
+                                     const int multi_processor_count,
+                                     cudaStream_t stream,
+                                     int* kernel_occupancy = nullptr) {
+  if (gemm_config.split_k_style != SplitKStyle::NO_SPLIT_K) {
+    throw std::runtime_error(
+        "[MoeGemm] Grouped gemm does not support split-k");
+  }
+
+#ifdef PADDLE_CUDA_BF16
+  static_assert(cutlass::platform::is_same<T, __nv_bfloat16>::value ||
+                    cutlass::platform::is_same<T, half>::value ||
+                    cutlass::platform::is_same<T, float>::value,
+                "Specialized for bfloat16, half, float");
+#else
+  static_assert(cutlass::platform::is_same<T, half>::value ||
+                    cutlass::platform::is_same<T, float>::value,
+                "Specialized for half, float");
+#endif
+
+  static_assert(
+      cutlass::platform::is_same<T, WeightType>::value ||
+          cutlass::platform::is_same<WeightType, uint8_t>::value ||
+          cutlass::platform::is_same<WeightType, cutlass::uint4b_t>::value,
+      "");
+
+  // The cutlass type for the input elements. This is needed to convert to
+  // cutlass::half_t if necessary.
+  using ElementType_ = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<T, half>::value,
+      cutlass::half_t,
+      T>::type;
+#ifdef PADDLE_CUDA_BF16
+  using ElementType = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<ElementType_, __nv_bfloat16>::value,
+      cutlass::bfloat16_t,
+      ElementType_>::type;
+#else
+  using ElementType = ElementType_;
+#endif
+
+  using CutlassWeightType_ = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<WeightType, half>::value,
+      cutlass::half_t,
+      WeightType>::type;
+#ifdef PADDLE_CUDA_BF16
+  using CutlassWeightType = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<CutlassWeightType_, __nv_bfloat16>::value,
+      cutlass::bfloat16_t,
+      CutlassWeightType_>::type;
+#else
+  using CutlassWeightType = CutlassWeightType_;
+#endif
+
+  // We need separate config for each architecture since we will target
+  // different tensorcore instructions. For float, we do not target TCs.
+  using MixedGemmArchTraits = cutlass::gemm::kernel::
+      MixedGemmArchTraits<ElementType, CutlassWeightType, arch>;
+  using ElementAccumulator = typename MixedGemmArchTraits::AccType;
+
+  using EpilogueOp = typename Epilogue<ElementType,
+                                       MixedGemmArchTraits::ElementsPerAccessC,
+                                       ElementAccumulator,
+                                       EpilogueTag>::Op;
+
+  // Finally, set up the kernel.
+  using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemmGrouped<
+      ElementType,
+      cutlass::layout::RowMajor,
+      cutlass::ComplexTransform::kNone,
+      MixedGemmArchTraits::ElementsPerAccessA,
+      CutlassWeightType,
+      typename MixedGemmArchTraits::LayoutB,
+      cutlass::ComplexTransform::kNone,
+      MixedGemmArchTraits::ElementsPerAccessB,
+      ElementType,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      typename MixedGemmArchTraits::OperatorClass,
+      arch,
+      ThreadblockShape,
+      WarpShape,
+      typename MixedGemmArchTraits::InstructionShape,
+      EpilogueOp,
+      cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
+      Stages,
+      cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly,
+      typename MixedGemmArchTraits::Operator>::GemmKernel;
+
+  using GemmKernel =
+      cutlass::gemm::kernel::MoeFCGemm<typename GemmKernel_::Mma,
+                                       typename GemmKernel_::Epilogue,
+                                       typename GemmKernel_::ThreadblockSwizzle,
+                                       arch,  // Ensure top level arch is used
+                                              // for dispatch
+                                       GemmKernel_::kGroupScheduleMode>;
+
+  using GemmGrouped = cutlass::gemm::device::GemmGrouped<GemmKernel>;
+
+  if (kernel_occupancy != nullptr) {
+    *kernel_occupancy = compute_occupancy_for_kernel<GemmKernel>();
+    return;
+  }
+  int occupancy = std::min(2, GemmGrouped::maximum_active_blocks());
+
+  if (occupancy == 0) {
+    throw std::runtime_error(
+        "[MoE Runner] GPU lacks the shared memory resources to run "
+        "GroupedGEMM kernel");
+  }
+  const int threadblock_count = multi_processor_count * occupancy;
+
+  typename EpilogueOp::Params epilogue_op(ElementAccumulator(1.f),
+                                          ElementAccumulator(0.f));
+
+  typename GemmGrouped::Arguments args(
+      num_experts,
+      threadblock_count,
+      epilogue_op,
+      reinterpret_cast<const ElementType*>(A),
+      reinterpret_cast<const CutlassWeightType*>(B),
+      reinterpret_cast<const ElementType*>(weight_scales),
+      reinterpret_cast<const ElementType*>(biases),
+      reinterpret_cast<ElementType*>(C),
+      total_rows_before_expert,
+      gemm_n,
+      gemm_k);
+
+  GemmGrouped gemm;
+
+  auto can_implement = gemm.can_implement(args);
+  if (can_implement != cutlass::Status::kSuccess) {
+    std::string err_msg = "MoEFC kernel will fail for params. Error: " +
+                          std::string(cutlassGetStatusString(can_implement));
+    throw std::runtime_error("[MoE Runner] " + err_msg);
+  }
+
+  auto init_status = gemm.initialize(args);
+  if (init_status != cutlass::Status::kSuccess) {
+    std::string err_msg =
+        "Failed to initialize cutlass variable batched gemm. Error: " +
+        std::string(cutlassGetStatusString(init_status));
+    throw std::runtime_error("[MoE Runner] " + err_msg);
+  }
+
+  auto run_status = gemm.run(stream);
+  if (run_status != cutlass::Status::kSuccess) {
+    std::string err_msg =
+        "Failed to run cutlass variable batched gemm. Error: " +
+        std::string(cutlassGetStatusString(run_status));
+    throw std::runtime_error("[MoE Runner] " + err_msg);
+  }
+}
+
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages,
+          typename Enable = void>
+struct dispatch_stages {
+  static void dispatch(const T* A,
+                       const WeightType* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int64_t* total_rows_before_expert,
+                       int64_t gemm_n,
+                       int64_t gemm_k,
+                       int num_experts,
+                       CutlassGemmConfig gemm_config,
+                       int multi_processor_count,
+                       cudaStream_t stream,
+                       int* occupancy = nullptr) {
+    // FT_LOG_DEBUG(__PRETTY_FUNCTION__);
+    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);
+    throw std::runtime_error("[dispatch_stages::dispatch] " + err_msg);
+  }
+};
+
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape>
+struct dispatch_stages<T,
+                       WeightType,
+                       arch,
+                       EpilogueTag,
+                       ThreadblockShape,
+                       WarpShape,
+                       2> {
+  static void dispatch(const T* A,
+                       const WeightType* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int64_t* total_rows_before_expert,
+                       int64_t gemm_n,
+                       int64_t gemm_k,
+                       int num_experts,
+                       CutlassGemmConfig gemm_config,
+                       int multi_processor_count,
+                       cudaStream_t stream,
+                       int* occupancy = nullptr) {
+    generic_moe_gemm_kernelLauncher<T,
+                                    WeightType,
+                                    arch,
+                                    EpilogueTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    2>(A,
+                                       B,
+                                       weight_scales,
+                                       biases,
+                                       C,
+                                       total_rows_before_expert,
+                                       gemm_n,
+                                       gemm_k,
+                                       num_experts,
+                                       gemm_config,
+                                       multi_processor_count,
+                                       stream,
+                                       occupancy);
+  }
+};
+
+template <typename T,
+          typename WeightType,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages>
+struct dispatch_stages<T,
+                       WeightType,
+                       cutlass::arch::Sm80,
+                       EpilogueTag,
+                       ThreadblockShape,
+                       WarpShape,
+                       Stages,
+                       typename std::enable_if<(Stages > 2)>::type> {
+  static void dispatch(const T* A,
+                       const WeightType* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int64_t* total_rows_before_expert,
+                       int64_t gemm_n,
+                       int64_t gemm_k,
+                       int num_experts,
+                       CutlassGemmConfig gemm_config,
+                       int multi_processor_count,
+                       cudaStream_t stream,
+                       int* occupancy = nullptr) {
+    generic_moe_gemm_kernelLauncher<T,
+                                    WeightType,
+                                    cutlass::arch::Sm80,
+                                    EpilogueTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    Stages>(A,
+                                            B,
+                                            weight_scales,
+                                            biases,
+                                            C,
+                                            total_rows_before_expert,
+                                            gemm_n,
+                                            gemm_k,
+                                            num_experts,
+                                            gemm_config,
+                                            multi_processor_count,
+                                            stream,
+                                            occupancy);
+  }
+};
+
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape>
+void dispatch_gemm_config(const T* A,
+                          const WeightType* B,
+                          const T* weight_scales,
+                          const T* biases,
+                          T* C,
+                          int64_t* total_rows_before_expert,
+                          int64_t gemm_n,
+                          int64_t gemm_k,
+                          int num_experts,
+                          CutlassGemmConfig gemm_config,
+                          int multi_processor_count,
+                          cudaStream_t stream,
+                          int* occupancy = nullptr) {
+#define dispatch_stages_macro(STAGE)                           \
+  case STAGE:                                                  \
+    dispatch_stages<T,                                         \
+                    WeightType,                                \
+                    arch,                                      \
+                    EpilogueTag,                               \
+                    ThreadblockShape,                          \
+                    WarpShape,                                 \
+                    STAGE>::dispatch(A,                        \
+                                     B,                        \
+                                     weight_scales,            \
+                                     biases,                   \
+                                     C,                        \
+                                     total_rows_before_expert, \
+                                     gemm_n,                   \
+                                     gemm_k,                   \
+                                     num_experts,              \
+                                     gemm_config,              \
+                                     multi_processor_count,    \
+                                     stream,                   \
+                                     occupancy);               \
+    break;
+
+  switch (gemm_config.stages) {
+    dispatch_stages_macro(2);
+    dispatch_stages_macro(3);
+    dispatch_stages_macro(4);
+    dispatch_stages_macro(5);
+    default:
+      std::string err_msg = "dispatch_gemm_config does not support stages " +
+                            std::to_string(gemm_config.stages);
+      throw std::runtime_error(
+          "[MoE][dispatch_gemm_config] " + err_msg);
+      break;
+  }
+}
+
+#define dispatch_gemm_config_macro(AA, BB, CC, DD, EE, FF)      \
+  case CutlassTileConfig::                                      \
+      CtaShape##AA##x##BB##x##CC##_WarpShape##DD##x##EE##x##FF: \
+    dispatch_gemm_config<T,                                     \
+                         WeightType,                            \
+                         arch,                                  \
+                         EpilogueTag,                           \
+                         cutlass::gemm::GemmShape<AA, BB, CC>,  \
+                         cutlass::gemm::GemmShape<DD, EE, FF>>( \
+        A,                                                      \
+        B,                                                      \
+        weight_scales,                                          \
+        biases,                                                 \
+        C,                                                      \
+        total_rows_before_expert,                               \
+        gemm_n,                                                 \
+        gemm_k,                                                 \
+        num_experts,                                            \
+        gemm_config,                                            \
+        multi_processor_count,                                  \
+        stream,                                                 \
+        occupancy);                                             \
+    break;
+
+// This overload will handle tensorop gemms. It is disabled via SFINAE for fp32.
+// This overload is only enabled when T == WeightType.
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename std::enable_if<!std::is_same<T, float>::value &&
+                                  std::is_same<T, WeightType>::value>::type* =
+              nullptr>
+void dispatch_moe_gemm_to_cutlass(const T* A,
+                                  const WeightType* B,
+                                  const T* weight_scales,
+                                  const T* biases,
+                                  T* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
+  switch (gemm_config.tile_config) {
+    dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+    dispatch_gemm_config_macro(64, 128, 64, 32, 64, 64);
+    dispatch_gemm_config_macro(128, 128, 64, 64, 32, 64);
+    case CutlassTileConfig::Undefined:
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+      break;
+    case CutlassTileConfig::ChooseWithHeuristic:
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass] gemm config should have "
+          "already been set by heuristic.");
+      break;
+    default:
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass] Config is invalid for same "
+          "type MoE tensorop GEMM.");
+      break;
+  }
+}
+
+// Tensorop GEMM overload
+// Overload for quantize MoE GEMMs. We disable some warp configs here since they
+// will not be used and we can improve compile time
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename std::enable_if<!std::is_same<T, float>::value &&
+                                  !std::is_same<T, WeightType>::value>::type* =
+              nullptr>
+void dispatch_moe_gemm_to_cutlass(const T* A,
+                                  const WeightType* B,
+                                  const T* weight_scales,
+                                  const T* biases,
+                                  T* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
+  if constexpr (std::is_same<arch, cutlass::arch::Sm70>::value) {
+    switch (gemm_config.tile_config) {
+      dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+      dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
+      case CutlassTileConfig::Undefined:
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        break;
+      case CutlassTileConfig::ChooseWithHeuristic:
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config should have "
+            "already been set by heuristic.");
+        break;
+      default:
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
+            "mixed type tensorop GEMM.");
+        break;
+    }
+  } 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:
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        break;
+      case CutlassTileConfig::ChooseWithHeuristic:
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config should have "
+            "already been set by heuristic.");
+        break;
+      default:
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
+            "mixed type tensorop GEMM.");
+        break;
+    }
+  }
+}
+
+// This overload will handle simt gemms. It is disabled via SFINAE for tensorop.
+template <
+    typename T,
+    typename WeightType,
+    typename arch,
+    typename EpilogueTag,
+    typename std::enable_if<std::is_same<T, float>::value>::type* = nullptr>
+void dispatch_moe_gemm_to_cutlass(const T* A,
+                                  const WeightType* B,
+                                  const T* weight_scales,
+                                  const T* biases,
+                                  T* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
+  switch (gemm_config.tile_config) {
+    dispatch_gemm_config_macro(128, 128, 8, 64, 64, 8);
+    case CutlassTileConfig::Undefined:
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config "
+          "undefined.");
+      break;
+    case CutlassTileConfig::ChooseWithHeuristic:
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config should "
+          "have already been set by heuristic.");
+      break;
+    default:
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass][SIMT] Unsupported config "
+          "for float MoE gemm.");
+      break;
+  }
+}
+
+template <typename T, typename WeightType>
+MoeGemmRunner<T, WeightType>::MoeGemmRunner() {
+  int device{-1};
+  check_cuda_error(cudaGetDevice(&device));
+  sm_ = getSMVersion();
+  check_cuda_error(cudaDeviceGetAttribute(
+      &multi_processor_count_, cudaDevAttrMultiProcessorCount, device));
+}
+
+template <typename T, typename WeightType>
+template <typename EpilogueTag>
+void MoeGemmRunner<T, WeightType>::dispatch_to_arch<EpilogueTag>(
+    const T* A,
+    const WeightType* B,
+    const T* weight_scales,
+    const T* biases,
+    T* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    int num_experts,
+    CutlassGemmConfig gemm_config,
+    cudaStream_t stream,
+    int* occupancy) {
+#define dispatch_moe_gemm_to_cutlass_macro(ARCH)                  \
+  dispatch_moe_gemm_to_cutlass<T, WeightType, ARCH, EpilogueTag>( \
+      A,                                                          \
+      B,                                                          \
+      weight_scales,                                              \
+      biases,                                                     \
+      C,                                                          \
+      total_rows_before_expert,                                   \
+      total_rows,                                                 \
+      gemm_n,                                                     \
+      gemm_k,                                                     \
+      num_experts,                                                \
+      gemm_config,                                                \
+      sm_,                                                        \
+      multi_processor_count_,                                     \
+      stream,                                                     \
+      occupancy);
+
+  if (sm_ >= 70 && sm_ < 75) {
+    dispatch_moe_gemm_to_cutlass_macro(cutlass::arch::Sm70);
+  } else if (sm_ >= 75 && sm_ < 80) {
+    dispatch_moe_gemm_to_cutlass_macro(cutlass::arch::Sm75);
+  } else if (sm_ >= 80 && sm_ < 91) {
+    dispatch_moe_gemm_to_cutlass_macro(cutlass::arch::Sm80);
+  } else {
+    PADDLE_FATAL("[MoE][GEMM Dispatch] Arch unsupported for MoE GEMM");
+  }
+}
+
+template <typename T, typename WeightType>
+template <typename EpilogueTag>
+void MoeGemmRunner<T, WeightType>::run_gemm<EpilogueTag>(
+    const T* A,
+    const WeightType* B,
+    const T* weight_scales,
+    const T* biases,
+    T* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    int num_experts,
+    cudaStream_t stream) {
+  static constexpr bool is_weight_only = !std::is_same<T, WeightType>::value;
+  static constexpr bool only_simt_configs = std::is_same<T, float>::value;
+  std::vector<CutlassGemmConfig> candidate_configs =
+      get_candidate_configs(sm_, -1, is_weight_only, only_simt_configs, true);
+  static constexpr int warm_time = 5;
+  static constexpr int test_time = 10;
+  auto& gemmConfigManager = GemmConfigManager::Instance();
+  constexpr GemmDataType dtype = getGemmDataType<T>();
+  constexpr GemmDataType wdtype = getGemmDataType<WeightType>();
+  GemmIDType gemmId{
+      gemm_n, gemm_k, GemmType::MOEGEMM, dtype, wdtype, num_experts};
+  CutlassGemmConfig chosen_config;
+  auto chosen_config_optional =
+      gemmConfigManager.getBestConfig(gemmId, total_rows);
+  if (chosen_config_optional != std::nullopt) {
+    chosen_config = chosen_config_optional.value();
+  } else {
+    float best_time = std::numeric_limits<float>::max();
+    CutlassGemmConfig best_config;
+    int profile_total_rows =
+        std::min(gemmConfigManager.nextPowerOfTwo(total_rows),
+                 gemmConfigManager.getMaxProfileM());
+    bool find_one = false;
+    for (size_t ii = 0; ii < candidate_configs.size(); ++ii) {
+      try {
+        for (int i = 0; i < warm_time; i++) {
+          dispatch_to_arch<EpilogueTag>(A,
+                                        B,
+                                        weight_scales,
+                                        biases,
+                                        C,
+                                        total_rows_before_expert,
+                                        total_rows,
+                                        gemm_n,
+                                        gemm_k,
+                                        num_experts,
+                                        candidate_configs[ii],
+                                        stream);
+        }
+        cudaEvent_t start;
+        cudaEvent_t stop;
+        check_cuda_error(cudaEventCreate(&start));
+        check_cuda_error(cudaEventCreate(&stop));
+        check_cuda_error(cudaStreamSynchronize(stream));
+        check_cuda_error(cudaEventRecord(start, stream));
+        for (int i = 0; i < test_time; i++) {
+          dispatch_to_arch<EpilogueTag>(A,
+                                        B,
+                                        weight_scales,
+                                        biases,
+                                        C,
+                                        total_rows_before_expert,
+                                        total_rows,
+                                        gemm_n,
+                                        gemm_k,
+                                        num_experts,
+                                        candidate_configs[ii],
+                                        stream);
+        }
+        check_cuda_error(cudaEventRecord(stop, stream));
+        check_cuda_error(cudaEventSynchronize(stop));
+        float elapsed;
+        check_cuda_error(cudaEventElapsedTime(&elapsed, start, stop));
+        check_cuda_error(cudaEventDestroy(start));
+        check_cuda_error(cudaEventDestroy(stop));
+        if (elapsed < best_time) {
+          best_time = elapsed;
+          best_config = candidate_configs[ii];
+        }
+        find_one = true;
+      } catch (const std::exception& e) {
+        std::cerr << "MOE config[" << ii << "]  Caught exception: " << e.what()
+                  << std::endl;
+      }
+    }
+    if (find_one) {
+      gemmConfigManager.addBestConfig(gemmId, profile_total_rows, best_config);
+      chosen_config = best_config;
+    } else {
+      PADDLE_FATAL("[MoE Configure Search] find no one available config.");
+    }
+  }
+  dispatch_to_arch<EpilogueTag>(A,
+                                B,
+                                weight_scales,
+                                biases,
+                                C,
+                                total_rows_before_expert,
+                                total_rows,
+                                gemm_n,
+                                gemm_k,
+                                num_experts,
+                                chosen_config,
+                                stream);
+}
+
+template <typename T, typename WeightType>
+void MoeGemmRunner<T, WeightType>::moe_gemm_bias_act(
+    const T* A,
+    const WeightType* B,
+    const T* weight_scales,
+    const T* biases,
+    T* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    int num_experts,
+    std::string activation_type,
+    cudaStream_t stream) {
+  if (activation_type == "none") {
+    if (biases) {
+      run_gemm<EpilogueOpBias>(A,
+                               B,
+                               weight_scales,
+                               biases,
+                               C,
+                               total_rows_before_expert,
+                               total_rows,
+                               gemm_n,
+                               gemm_k,
+                               num_experts,
+                               stream);
+    } else {
+      run_gemm<EpilogueOpNoBias>(A,
+                                 B,
+                                 weight_scales,
+                                 nullptr,
+                                 C,
+                                 total_rows_before_expert,
+                                 total_rows,
+                                 gemm_n,
+                                 gemm_k,
+                                 num_experts,
+                                 stream);
+    }
+  }
+}
+
+template <typename T, typename WeightType>
+void MoeGemmRunner<T, WeightType>::moe_gemm(const T* A,
+                                            const WeightType* B,
+                                            const T* weight_scales,
+                                            T* C,
+                                            int64_t* total_rows_before_expert,
+                                            int64_t total_rows,
+                                            int64_t gemm_n,
+                                            int64_t gemm_k,
+                                            int num_experts,
+                                            cudaStream_t stream) {
+  run_gemm<EpilogueOpNoBias>(A,
+                             B,
+                             weight_scales,
+                             nullptr,
+                             C,
+                             total_rows_before_expert,
+                             total_rows,
+                             gemm_n,
+                             gemm_k,
+                             num_experts,
+                             stream);
+}

csrc/gpu/moe/fused_moe/fused_moe.cu

@@ -0,0 +1,192 @@
+// 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.
+
+
+// Ignore CUTLASS warnings about type punning
+
+#pragma once
+
+#include "cutlass/numeric_conversion.h"
+#include "helper.h"
+#include "moe/fused_moe_helper.h"
+
+__global__ void compute_total_rows_before_expert_kernel(
+    int* sorted_experts,
+    const int64_t sorted_experts_len,
+    const int64_t num_experts,
+    int64_t* total_rows_before_expert) {
+  // First, compute the global tid. We only need 1 thread per expert.
+  const int expert = blockIdx.x * blockDim.x + threadIdx.x;
+  if (expert >= num_experts) return;
+
+  // This should construct the last index where each expert occurs.
+  total_rows_before_expert[expert] =
+      find_total_elts_leq_target(sorted_experts, sorted_experts_len, expert);
+}
+
+void compute_total_rows_before_expert(int* sorted_indices,
+                                      const int64_t total_indices,
+                                      const int64_t num_experts,
+                                      int64_t* total_rows_before_expert,
+                                      cudaStream_t stream) {
+  const int threads = std::min(int64_t(1024), num_experts);
+  const int blocks = (num_experts + threads - 1) / threads;
+
+  compute_total_rows_before_expert_kernel<<<blocks, threads, 0, stream>>>(
+      sorted_indices, total_indices, num_experts, total_rows_before_expert);
+}
+
+template <paddle::DataType T>
+void FusedMoeKernel(const paddle::Tensor& input,
+                    const paddle::Tensor& gate_weight,
+                    const paddle::Tensor& ffn1_weight,
+                    const paddle::optional<paddle::Tensor>& ffn1_scale,
+                    const paddle::optional<paddle::Tensor>& ffn1_bias,
+                    const paddle::Tensor& ffn2_weight,
+                    const paddle::optional<paddle::Tensor>& ffn2_scale,
+                    const paddle::optional<paddle::Tensor>& ffn2_bias,
+                    const std::string& quant_method,
+                    const int moe_topk,
+                    const bool group_moe,
+                    const bool norm_topk_prob,
+                    paddle::Tensor* output) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto* output_data = output->data<data_t>();
+
+  auto fp16_moe_gemm_runner = MoeGemmRunner<DataType_, DataType_>();
+  auto int8_moe_gemm_runner = MoeGemmRunner<DataType_, uint8_t>();
+  auto int4_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::uint4b_t>();
+
+  using NvType = typename traits_::DataType;
+  auto moe_compute = MoeHelper<data_t, NvType>(quant_method,
+                                               &fp16_moe_gemm_runner,
+                                               &int8_moe_gemm_runner,
+                                               &int4_moe_gemm_runner);
+
+  moe_compute.ComputeFFN(&input,
+                         &gate_weight,
+                         &ffn1_weight,
+                         ffn1_scale ? ffn1_scale.get_ptr() : nullptr,
+                         ffn1_bias ? ffn1_bias.get_ptr() : nullptr,
+                         &ffn2_weight,
+                         ffn2_scale ? ffn2_scale.get_ptr() : nullptr,
+                         ffn2_bias ? ffn2_bias.get_ptr() : nullptr,
+                         nullptr,
+                         moe_topk,
+                         group_moe,
+                         norm_topk_prob,
+                         1.0,  // ComputeFFN
+                         "ffn",
+                         output);
+}
+
+
+std::vector<paddle::Tensor> FusedExpertMoe(
+    const paddle::Tensor& input,
+    const paddle::Tensor& gate_weight,
+    const paddle::Tensor& ffn1_weight,
+    const paddle::Tensor& ffn2_weight,
+    const paddle::optional<paddle::Tensor>& ffn1_bias,
+    const paddle::optional<paddle::Tensor>& ffn1_scale,
+    const paddle::optional<paddle::Tensor>& ffn2_bias,
+    const paddle::optional<paddle::Tensor>& ffn2_scale,
+    const std::string& quant_method,
+    const int moe_topk,
+    const bool norm_topk_prob,
+    const bool group_moe) {
+  const auto input_type = input.dtype();
+  auto output = paddle::empty_like(input);
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      FusedMoeKernel<paddle::DataType::BFLOAT16>(input,
+                                                 gate_weight,
+                                                 ffn1_weight,
+                                                 ffn1_scale,
+                                                 ffn1_bias,
+                                                 ffn2_weight,
+                                                 ffn2_scale,
+                                                 ffn2_bias,
+                                                 quant_method,
+                                                 moe_topk,
+                                                 group_moe,
+                                                 norm_topk_prob,
+                                                 &output);
+      break;
+    case paddle::DataType::FLOAT16:
+      FusedMoeKernel<paddle::DataType::FLOAT16>(input,
+                                                gate_weight,
+                                                ffn1_weight,
+                                                ffn1_scale,
+                                                ffn1_bias,
+                                                ffn2_weight,
+                                                ffn2_scale,
+                                                ffn2_bias,
+                                                quant_method,
+                                                moe_topk,
+                                                group_moe,
+                                                norm_topk_prob,
+                                                &output);
+      break;
+    default:
+      PD_THROW("Unsupported data type for FusedMoeKernel");
+  }
+  return {output};
+}
+
+std::vector<std::vector<int64_t>> FusedExpertMoeInferShape(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<int64_t>& gate_weight_shape,
+    const std::vector<int64_t>& ffn1_weight_shape,
+    const std::vector<int64_t>& ffn2_weight_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_scale_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_scale_shape) {
+  return {input_shape};
+}
+
+std::vector<paddle::DataType> FusedExpertMoeInferDtype(
+    const paddle::DataType& input_dtype,
+    const paddle::DataType& gate_weight_dtype,
+    const paddle::DataType& ffn1_weight_dtype,
+    const paddle::DataType& ffn2_weight_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_bias_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_scale_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_bias_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_scale_dtype) {
+  return {input_dtype};
+}
+
+
+PD_BUILD_OP(fused_expert_moe)
+    .Inputs({"input",
+             "gate_weight",
+             "ffn1_weight",
+             "ffn2_weight",
+             paddle::Optional("ffn1_bias"),
+             paddle::Optional("ffn1_scale"),
+             paddle::Optional("ffn2_bias"),
+             paddle::Optional("ffn2_scale")})
+    .Outputs({"output"})
+    .Attrs({"quant_method:std::string",
+            "moe_topk:int",
+            "norm_topk_prob:bool",
+            "group_moe:bool"})
+    .SetKernelFn(PD_KERNEL(FusedExpertMoe))
+    .SetInferShapeFn(PD_INFER_SHAPE(FusedExpertMoeInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(FusedExpertMoeInferDtype));

csrc/gpu/moe/fused_moe/glog/logging.h

@@ -0,0 +1,14 @@
+#pragma once
+#include <iostream>
+
+/// @brief Defines a macro to disable log output.
+/// @details This macro returns a `NullStream` object, and all content passed via `<<` is ignored.
+///          It is designed to keep the code compilable without modifying log call statement.
+/// @param level The log level (ignored).
+class NullStream {
+public:
+    template <typename T>
+    NullStream& operator<<(const T&) { return *this; }
+};
+
+#define VLOG(level) NullStream()

csrc/gpu/moe/fused_moe/moe/fused_moe_helper.h

@@ -0,0 +1,427 @@
+
+/* 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. */
+
+#pragma once
+
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "moe/fused_moe_op.h"
+
+using namespace phi;
+
+template <typename T, int VecSize>
+__global__ void moe_token_type_ids_kernel(T *gating_output,
+                                          const int *moe_token_type_ids_out,
+                                          const int num_rows,
+                                          const int num_experts,
+                                          const int k) {
+  const int moe_token_index = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (moe_token_index >= num_rows) {
+    return;
+  }
+
+  gating_output[moe_token_index * 2] =
+      gating_output[moe_token_index * 2] +
+      (moe_token_type_ids_out[moe_token_index]) * -1e10;
+  gating_output[moe_token_index * 2 + 1] =
+      gating_output[moe_token_index * 2 + 1] +
+      (1 - moe_token_type_ids_out[moe_token_index]) * -1e10;
+}
+
+template <typename T>
+void moe_token_type_ids_kernelLauncher(T *gating_output,
+                                       const int *moe_token_type_ids_out,
+                                       const int num_rows,
+                                       const int num_experts,
+                                       const int k,
+                                       cudaStream_t stream) {
+  const int blocks = num_rows * k / 512 + 1;
+  const int threads = 512;
+  moe_token_type_ids_kernel<T, 1><<<blocks, 512, 0, stream>>>(
+      gating_output, moe_token_type_ids_out, num_rows, num_experts, k);
+}
+
+
+template <typename T, typename NvType>
+class MoeHelper {
+public:
+  MoeHelper(const std::string gemm_method,
+            MoeGemmRunner<NvType, NvType> *fp16_moe_gemm_runner,
+            MoeGemmRunner<NvType, uint8_t> *int8_moe_gemm_runner,
+            MoeGemmRunner<NvType, cutlass::uint4b_t> *int4_moe_gemm_runner,
+            int layernum = 0)
+      : gemm_method_(gemm_method),
+        fp16_moe_gemm_runner_(fp16_moe_gemm_runner),
+        int8_moe_gemm_runner_(int8_moe_gemm_runner),
+        int4_moe_gemm_runner_(int4_moe_gemm_runner),
+        layernum_(layernum) {}
+
+  // --------      getWorkspaceSize      -------- //
+  template <typename KeyT>
+  size_t getWorkspaceSize(const int64_t num_rows,
+                          const int64_t hidden_size,
+                          const int64_t inter_size,
+                          const int64_t num_experts,
+                          const int64_t k) {
+    const size_t buf_size = AlignTo16(k * num_rows * hidden_size);
+    const size_t interbuf_size = AlignTo16(k * num_rows * inter_size);
+    const size_t padded_experts = AlignTo16(num_experts);
+    const size_t num_moe_inputs = AlignTo16(k * num_rows);
+    // softmax output, permuted_rows and permuted_experts have moved to outside
+    // of moe kernel, allocate them in Encoder or Decoder before invoking
+    // FfnLayer forward.
+    size_t total_ws_bytes =
+        5 * num_moe_inputs *
+        sizeof(int);  // source_rows_, permuted_rows_, permuted_experts_
+    total_ws_bytes += buf_size * sizeof(KeyT);  // permuted_data
+    total_ws_bytes +=
+        padded_experts * sizeof(int64_t);  // Hold total_rows_before_expert_
+
+    const size_t bytes_for_fc1_result = interbuf_size * sizeof(KeyT);
+    const size_t sorter_ws_size_bytes =
+        AlignTo16(sorter_.getWorkspaceSize(num_rows));
+    sorter_.update_num_experts(num_experts);
+
+    int64_t bytes_for_intermediate_and_sorting = bytes_for_fc1_result;
+    if (sorter_ws_size_bytes > bytes_for_fc1_result) {
+      int64_t remaining_bytes =
+          AlignTo16(sorter_ws_size_bytes - bytes_for_fc1_result);
+      bytes_for_intermediate_and_sorting += remaining_bytes;
+    }
+
+    total_ws_bytes +=
+        bytes_for_intermediate_and_sorting;  // intermediate (fc1) output + cub
+                                             // sorting workspace
+
+    int64_t num_softmax_outs = 0;
+    const bool is_pow_2 =
+        (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    if (!is_pow_2 || num_experts > 256) {
+      num_softmax_outs = AlignTo16(num_rows * num_experts);
+    }
+
+    total_ws_bytes += num_softmax_outs * sizeof(float);
+
+    return total_ws_bytes;
+  }
+
+  void ComputeFFN(const paddle::Tensor *input,
+                  const paddle::Tensor *gate_weight,
+                  const paddle::Tensor *ffn1_weight,
+                  const paddle::Tensor *ffn1_scale,
+                  const paddle::Tensor *ffn1_bias,
+                  const paddle::Tensor *ffn2_weight,
+                  const paddle::Tensor *ffn2_scale,
+                  const paddle::Tensor *ffn2_bias,
+                  const paddle::Tensor *moe_token_type_ids,
+                  const int moe_topk,
+                  const bool group_moe,
+                  const bool norm_topk_prob,
+                  const float routed_scaling_factor,
+                  const std::string moe_type,
+                  paddle::Tensor *output) {
+    auto *input_activations = input->data<T>();
+    auto *gating_weights = gate_weight->data<float>();
+    const T *fc1_expert_biases = ffn1_bias ? ffn1_bias->data<T>() : nullptr;
+    const T *fc2_expert_biases = ffn2_bias ? ffn2_bias->data<T>() : nullptr;
+
+    auto *output_ = output->data<T>();
+    auto stream = input->stream();
+    auto place = input->place();
+    auto input_type = input->dtype();
+
+    auto input_dims = input->dims();
+    auto ffn1_dims = ffn1_weight->dims();
+    int64_t token_num = 0;
+    if (input_dims.size() == 3) {
+      token_num = input_dims[0] * input_dims[1];
+    } else {
+      token_num = input_dims[0];
+    }
+    const int64_t num_rows = token_num;
+
+    const int64_t hidden_size = ffn1_dims[1];
+    int64_t inter_dim = 0;
+    if (moe_type == "qkv") {
+      inter_dim = ffn1_dims[2] * ffn1_dims[3] * ffn1_dims[4];
+    } else {
+      inter_dim = ffn1_dims[2];
+    }
+
+    if (gemm_method_ == "weight_only_int4") {
+      inter_dim = inter_dim * 2;
+    }
+
+    const int64_t inter_size = inter_dim;
+    const int64_t num_experts = ffn1_dims[0];
+    const int64_t k = moe_topk;
+
+
+    int64_t bytes =
+        getWorkspaceSize<T>(num_rows, hidden_size, inter_size, num_experts, k);
+
+    // Pointers
+    int *expert_for_source_row;
+    int *source_rows_;
+    int *permuted_rows_;
+    int *permuted_experts_;
+    int *expanded_source_row_to_expanded_dest_row;
+
+    T *permuted_data_;
+    int64_t *total_rows_before_expert_;
+    T *fc1_result_;
+    float *softmax_out_;
+
+    paddle::Tensor ws_ptr_tensor =
+        GetEmptyTensor({bytes}, paddle::DataType::INT8, place);
+    int8_t *ws_ptr = ws_ptr_tensor.data<int8_t>();
+
+    const int64_t buf_size = AlignTo16(k * num_rows * hidden_size);
+    const int64_t interbuf_size = AlignTo16(k * num_rows * inter_size);
+    const int64_t padded_experts = AlignTo16(num_experts);
+    const int64_t num_moe_inputs = AlignTo16(k * num_rows);
+
+    expert_for_source_row = reinterpret_cast<int *>(ws_ptr);
+    source_rows_ = expert_for_source_row + num_moe_inputs;
+    permuted_rows_ = source_rows_ + num_moe_inputs;
+    permuted_experts_ = permuted_rows_ + num_moe_inputs;
+    expanded_source_row_to_expanded_dest_row =
+        permuted_experts_ + num_moe_inputs;
+    permuted_data_ = reinterpret_cast<T *>(
+        expanded_source_row_to_expanded_dest_row + num_moe_inputs);
+    total_rows_before_expert_ =
+        reinterpret_cast<int64_t *>(permuted_data_ + buf_size);
+    fc1_result_ =
+        reinterpret_cast<T *>(total_rows_before_expert_ + padded_experts);
+
+    const bool is_pow_2 =
+        (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    if (!is_pow_2 || num_experts > 256) {
+      softmax_out_ = reinterpret_cast<float *>(fc1_result_ + interbuf_size);
+    } else {
+      softmax_out_ = nullptr;
+    }
+
+    paddle::Tensor expert_scales_float_tensor =
+        GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+    float *expert_scales_float = expert_scales_float_tensor.data<float>();
+
+    float *softmax_max_prob = nullptr;
+    if (group_moe) {
+      paddle::Tensor softmax_max_prob_tensor = GetEmptyTensor(
+          {num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+      // (TODO: check fill success ?)
+      paddle::experimental::fill(softmax_max_prob_tensor, 0.f);
+      softmax_max_prob = softmax_max_prob_tensor.data<float>();
+    }
+
+    paddle::Tensor fc1_out_tensor =
+        GetEmptyTensor({num_rows * k, inter_size}, input_type, place);
+    T *fc1_out = fc1_out_tensor.data<T>();
+
+    auto input_cast_tensor =
+        paddle::experimental::cast(*input, paddle::DataType::FLOAT32);
+    auto gate_tensor =
+        paddle::experimental::matmul(input_cast_tensor, *gate_weight);
+    float *gating_output = gate_tensor.data<float>();
+
+    if (moe_token_type_ids) {
+      auto *moe_token_type_ids_out = moe_token_type_ids->data<int>();
+      moe_token_type_ids_kernelLauncher<float>(gating_output,
+                                               moe_token_type_ids_out,
+                                               num_rows,
+                                               num_experts,
+                                               k,
+                                               stream);
+    }
+
+    topk_gating_softmax_kernelLauncher<float>(gating_output,
+                                              expert_scales_float,
+                                              softmax_out_,
+                                              expert_for_source_row,
+                                              source_rows_,
+                                              softmax_max_prob,
+                                              num_rows,
+                                              num_experts,
+                                              k,
+                                              group_moe,
+                                              stream);
+
+    const int64_t sorter_ws_size_bytes =
+        AlignTo16(sorter_.getWorkspaceSize(int64_t(k * num_rows)));
+
+    sorter_.run(fc1_result_,
+                sorter_ws_size_bytes,
+                expert_for_source_row,
+                permuted_experts_,
+                source_rows_,
+                permuted_rows_,
+                k * num_rows,
+                false,
+                stream);
+
+    initialize_moe_routing_kernelLauncher(
+        input_activations,
+        permuted_data_,
+        permuted_rows_,
+        expanded_source_row_to_expanded_dest_row,
+        num_rows,
+        num_rows,
+        hidden_size,
+        k,
+        stream);
+
+    const int64_t expanded_active_expert_rows = k * num_rows;
+
+    compute_total_rows_before_expert(permuted_experts_,
+                                     expanded_active_expert_rows,
+                                     num_experts,
+                                     total_rows_before_expert_,
+                                     stream);
+
+    if (gemm_method_ == "weight_only_int8") {
+      int8_moe_gemm_runner_->moe_gemm_bias_act(
+          reinterpret_cast<NvType *>(permuted_data_),
+          reinterpret_cast<const uint8_t *>(ffn1_weight->data<int8_t>()),
+          reinterpret_cast<const NvType *>(ffn1_scale->data<T>()),
+          reinterpret_cast<const NvType *>(fc1_expert_biases),
+          reinterpret_cast<NvType *>(fc1_out),
+          total_rows_before_expert_,
+          expanded_active_expert_rows,
+          inter_size,
+          hidden_size,
+          num_experts,
+          "none",
+          stream);
+    } else if (gemm_method_ == "weight_only_int4") {
+      int4_moe_gemm_runner_->moe_gemm_bias_act(
+          reinterpret_cast<NvType *>(permuted_data_),
+          reinterpret_cast<const cutlass::uint4b_t *>(
+              ffn1_weight->data<int8_t>()),
+          reinterpret_cast<const NvType *>(ffn1_scale->data<T>()),
+          reinterpret_cast<const NvType *>(fc1_expert_biases),
+          reinterpret_cast<NvType *>(fc1_out),
+          total_rows_before_expert_,
+          expanded_active_expert_rows,
+          inter_size,
+          hidden_size,
+          num_experts,
+          "none",
+          stream);
+    } else {
+      fp16_moe_gemm_runner_->moe_gemm_bias_act(
+          reinterpret_cast<NvType *>(permuted_data_),
+          reinterpret_cast<const NvType *>(ffn1_weight->data<T>()),
+          nullptr,
+          reinterpret_cast<const NvType *>(fc1_expert_biases),
+          reinterpret_cast<NvType *>(fc1_out),
+          total_rows_before_expert_,
+          expanded_active_expert_rows,
+          inter_size,
+          hidden_size,
+          num_experts,
+          "none",
+          stream);
+    }
+
+    if (moe_type == "ffn") {
+      auto act_out_tensor =
+          paddle::experimental::swiglu(fc1_out_tensor, nullptr);
+      auto act_out = act_out_tensor.data<T>();
+
+      paddle::Tensor fc2_output_tensor =
+          GetEmptyTensor({k * num_rows, hidden_size}, input_type, place);
+      T *fc2_result = fc2_output_tensor.data<T>();
+
+      if (gemm_method_ == "weight_only_int8") {
+        int8_moe_gemm_runner_->moe_gemm(
+            reinterpret_cast<NvType *>(act_out),
+            reinterpret_cast<const uint8_t *>(ffn2_weight->data<int8_t>()),
+            reinterpret_cast<const NvType *>(ffn2_scale->data<T>()),
+            reinterpret_cast<NvType *>(fc2_result),
+            total_rows_before_expert_,
+            expanded_active_expert_rows,
+            hidden_size,
+            inter_size / 2,
+            num_experts,
+            stream);
+      } else if (gemm_method_ == "weight_only_int4") {
+        int4_moe_gemm_runner_->moe_gemm(
+            reinterpret_cast<NvType *>(act_out),
+            reinterpret_cast<const cutlass::uint4b_t *>(
+                ffn2_weight->data<int8_t>()),
+            reinterpret_cast<const NvType *>(ffn2_scale->data<T>()),
+            reinterpret_cast<NvType *>(fc2_result),
+            total_rows_before_expert_,
+            expanded_active_expert_rows,
+            hidden_size,
+            inter_size / 2,
+            num_experts,
+            stream);
+      } else {
+        fp16_moe_gemm_runner_->moe_gemm(
+            reinterpret_cast<NvType *>(act_out),
+            reinterpret_cast<const NvType *>(ffn2_weight->data<T>()),
+            nullptr,
+            reinterpret_cast<NvType *>(fc2_result),
+            total_rows_before_expert_,
+            expanded_active_expert_rows,
+            hidden_size,
+            inter_size / 2,
+            num_experts,
+            stream);
+      }
+
+      finalize_moe_routing_kernelLauncher(
+          fc2_result,
+          output_,
+          fc2_expert_biases,
+          reinterpret_cast<float *>(expert_scales_float),
+          expanded_source_row_to_expanded_dest_row,
+          expert_for_source_row,
+          num_rows,
+          hidden_size,
+          k,
+          static_cast<int>(1),
+          norm_topk_prob,
+          routed_scaling_factor,
+          stream);
+    } else {
+      finalize_moe_routing_kernelLauncher(
+          // fc2_result,
+          fc1_out,
+          output_,
+          fc1_expert_biases,  // fc2_expert_biases,
+          reinterpret_cast<float *>(expert_scales_float),
+          expanded_source_row_to_expanded_dest_row,
+          expert_for_source_row,
+          num_rows,
+          inter_size,
+          k,
+          static_cast<int>(0),
+          norm_topk_prob,
+          routed_scaling_factor,
+          stream);
+    }
+  }
+
+private:
+  std::string gemm_method_;
+  MoeGemmRunner<NvType, NvType> *fp16_moe_gemm_runner_;
+  MoeGemmRunner<NvType, uint8_t> *int8_moe_gemm_runner_;
+  MoeGemmRunner<NvType, cutlass::uint4b_t> *int4_moe_gemm_runner_;
+  int layernum_;
+  CubKeyValueSorter sorter_;
+};

csrc/gpu/moe/fused_moe/moe/fused_moe_imp_op.h

@@ -0,0 +1,123 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * 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.
+ */
+
+#pragma once
+#include <string>
+#include <sstream>
+#include "cub/cub.cuh"
+
+static const float HALF_FLT_MAX = 65504.F;
+static const float HALF_FLT_MIN = -65504.F;
+static inline size_t AlignTo16(const size_t& input) {
+  static constexpr int ALIGNMENT = 16;
+  return ALIGNMENT * ((input + ALIGNMENT - 1) / ALIGNMENT);
+}
+
+class CubKeyValueSorter {
+ public:
+  CubKeyValueSorter() : num_experts_(0), num_bits_(sizeof(int) * 8) {}
+
+  explicit CubKeyValueSorter(const int num_experts)
+      : num_experts_(num_experts),
+        num_bits_(static_cast<int>(log2(num_experts)) + 1) {}
+
+  void update_num_experts(const int num_experts) {
+    num_experts_ = num_experts;
+    num_bits_ = static_cast<int>(log2(num_experts)) + 1;
+  }
+
+  size_t getWorkspaceSize(const size_t num_key_value_pairs,
+                          bool descending = false) {
+    num_key_value_pairs_ = num_key_value_pairs;
+    size_t required_storage = 0;
+    int* null_int = nullptr;
+    if (descending) {
+      cub::DeviceRadixSort::SortPairsDescending(NULL,
+                                                required_storage,
+                                                null_int,
+                                                null_int,
+                                                null_int,
+                                                null_int,
+                                                num_key_value_pairs,
+                                                0,
+                                                32);
+    } else {
+      cub::DeviceRadixSort::SortPairs(NULL,
+                                      required_storage,
+                                      null_int,
+                                      null_int,
+                                      null_int,
+                                      null_int,
+                                      num_key_value_pairs,
+                                      0,
+                                      num_bits_);
+    }
+    return required_storage;
+  }
+
+  template <typename KeyT>
+  void run(void* workspace,
+           const size_t workspace_size,
+           const KeyT* keys_in,
+           KeyT* keys_out,
+           const int* values_in,
+           int* values_out,
+           const size_t num_key_value_pairs,
+           bool descending,
+           cudaStream_t stream) {
+    size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs);
+    size_t actual_ws_size = workspace_size;
+
+    if (expected_ws_size > workspace_size) {
+      std::stringstream err_ss;
+      err_ss << "[Error][CubKeyValueSorter::run]\n";
+      err_ss << "Error. The allocated workspace is too small to run this "
+                "problem.\n";
+      err_ss << "Expected workspace size of at least " << expected_ws_size
+             << " but got problem size " << workspace_size << "\n";
+      throw std::runtime_error(err_ss.str());
+    }
+    if (descending) {
+      cub::DeviceRadixSort::SortPairsDescending(workspace,
+                                                actual_ws_size,
+                                                keys_in,
+                                                keys_out,
+                                                values_in,
+                                                values_out,
+                                                num_key_value_pairs,
+                                                0,
+                                                32,
+                                                stream);
+    } else {
+      cub::DeviceRadixSort::SortPairs(workspace,
+                                      actual_ws_size,
+                                      keys_in,
+                                      keys_out,
+                                      values_in,
+                                      values_out,
+                                      num_key_value_pairs,
+                                      0,
+                                      num_bits_,
+                                      stream);
+    }
+  }
+
+ private:
+  size_t num_key_value_pairs_;
+  int num_experts_;
+  int num_bits_;
+};

csrc/gpu/moe/fused_moe/moe/fused_moe_op.h

@@ -0,0 +1,990 @@
+// /*
+//  * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+//  * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+//  *
+//  * 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.
+//  */
+
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include "moe/fused_moe_imp_op.h"
+#include "moe/fused_moe_helper.h"
+#include "cutlass/numeric_conversion.h"
+// Ignore CUTLASS warnings about type punning
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#pragma GCC diagnostic ignored "-Wunused-function"
+
+// #include "paddle/phi/backends/gpu/gpu_info.h"
+#pragma GCC diagnostic pop
+
+#include "helper.h"
+
+#define WARP_SIZE 32
+
+struct GpuLaunchConfig {
+    dim3 block_per_grid;
+    dim3 thread_per_block;
+};
+
+inline GpuLaunchConfig Get1DBlocksAnd2DGridsMoe(const int64_t cols) {
+  int blocks_x = cols;
+  int blocks_y = 1;
+  int blocks_z = 1;
+  if (blocks_x > 1024) {
+    blocks_y = 256;
+    blocks_x = (blocks_x + blocks_y - 1) / blocks_y;
+  }
+
+  GpuLaunchConfig config;
+  config.block_per_grid.x = blocks_x;
+  config.block_per_grid.y = blocks_y;
+  config.block_per_grid.z = blocks_z;
+  return config;
+}
+
+// ====================== Softmax things ===============================
+// We have our own implementation of softmax here so we can support transposing
+// the output in the softmax kernel when we extend this module to support
+// expert-choice routing.
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__
+    void group_moe_softmax(const T* input,
+                           T* output,
+                           T* softmax_max_prob,
+                           const int64_t num_cols,
+                           const int64_t softmax_num_rows) {
+  using BlockReduce = cub::BlockReduce<float, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  __shared__ float normalizing_factor;
+  __shared__ float float_max;
+  __shared__ float max_out;
+
+  int globalIdx = blockIdx.x + blockIdx.y * gridDim.x;
+  if (globalIdx >= softmax_num_rows) {
+    return;
+  }
+  const int64_t thread_row_offset = globalIdx * num_cols;
+
+  cub::Sum sum;
+  float threadData(-FLT_MAX);
+
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData = max(static_cast<float>(input[idx]), threadData);
+  }
+
+  const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+  if (threadIdx.x == 0) {
+    float_max = maxElem;
+  }
+  __syncthreads();
+
+  threadData = 0;
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData += exp((static_cast<float>(input[idx]) - float_max));
+  }
+
+  const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
+
+  if (threadIdx.x == 0) {
+    normalizing_factor = 1.f / Z;
+  }
+  __syncthreads();
+
+  threadData = 0;
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    const float val =
+        exp((static_cast<float>(input[idx]) - float_max)) * normalizing_factor;
+    output[idx] = T(val);
+    threadData = max(static_cast<float>(T(val)), threadData);
+  }
+
+  const float maxOut = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+  if (threadIdx.x == 0) {
+    // group max probs
+    max_out = 1.f / maxOut;
+    softmax_max_prob[globalIdx] = T(max_out);
+  }
+  __syncthreads();
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    // group softmax normalization
+    output[idx] = output[idx] * static_cast<T>(max_out);
+  }
+}
+
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moe_top_k(const T* inputs_after_softmax,
+                                                 T* output,
+                                                 int* indices,
+                                                 int* source_rows,
+                                                 T* softmax_max_prob,
+                                                 const int64_t num_experts,
+                                                 const int64_t k,
+                                                 const int64_t num_rows) {
+  using cub_kvp = cub::KeyValuePair<int, T>;
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x + blockIdx.y * gridDim.x;
+  if (block_row >= num_rows) {
+    return;
+  }
+
+  const bool should_process_row = true;
+  const int thread_read_offset = block_row * num_experts;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = T(-1.f);  // This is OK because inputs are probabilities
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_softmax[idx];
+
+      for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
+        const int prior_winning_expert = indices[k * block_row + prior_k];
+
+        if (prior_winning_expert == expert) {
+          inp_kvp = thread_kvp;
+        }
+      }
+
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp =
+        BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      const int idx = k * block_row + k_idx;
+      // restore normalized probes
+      output[idx] = result_kvp.value / T(softmax_max_prob[idx]);
+      indices[idx] = should_process_row ? result_kvp.key : num_experts;
+      source_rows[idx] = k_idx * num_rows + block_row;
+    }
+    __syncthreads();
+  }
+}
+
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moe_softmax(const T* input,
+                                                   T* output,
+                                                   const int64_t num_cols,
+                                                   const int64_t num_rows) {
+  using BlockReduce = cub::BlockReduce<float, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  __shared__ float normalizing_factor;
+  __shared__ float float_max;
+
+  int globalIdx = blockIdx.x + blockIdx.y * gridDim.x;
+  if (globalIdx >= num_rows) {
+    return;
+  }
+  const int64_t thread_row_offset = globalIdx * num_cols;
+
+  cub::Sum sum;
+  float threadData(-FLT_MAX);
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData = max(static_cast<float>(input[idx]), threadData);
+  }
+
+  const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+  if (threadIdx.x == 0) {
+    float_max = maxElem;
+  }
+  __syncthreads();
+
+  threadData = 0;
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData += exp((static_cast<float>(input[idx]) - float_max));
+  }
+
+  const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
+
+  if (threadIdx.x == 0) {
+    normalizing_factor = 1.f / Z;
+  }
+  __syncthreads();
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    const float val =
+        exp((static_cast<float>(input[idx]) - float_max)) * normalizing_factor;
+    output[idx] = T(val);
+  }
+}
+
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moe_top_k(const T* inputs_after_softmax,
+                                                 T* output,
+                                                 int* indices,
+                                                 int* source_rows,
+                                                 const int64_t num_experts,
+                                                 const int64_t k,
+                                                 const int64_t num_rows) {
+  using cub_kvp = cub::KeyValuePair<int, T>;
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x + blockIdx.y * gridDim.x;
+  if (block_row >= num_rows) {
+    return;
+  }
+
+  const bool should_process_row = true;
+  const int thread_read_offset = block_row * num_experts;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = T(-1.f);  // This is OK because inputs are probabilities
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_softmax[idx];
+
+      for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
+        const int prior_winning_expert = indices[k * block_row + prior_k];
+
+        if (prior_winning_expert == expert) {
+          inp_kvp = thread_kvp;
+        }
+      }
+
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp =
+        BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = result_kvp.value;
+      indices[idx] = should_process_row ? result_kvp.key : num_experts;
+      source_rows[idx] = k_idx * num_rows + block_row;
+    }
+    __syncthreads();
+  }
+}
+
+// ====================== TopK softmax things ===============================
+
+/*
+  A Top-K gating softmax written to exploit when the number of experts in the
+  MoE layers are a small power of 2. This allows us to cleanly share the rows
+  among the threads in a single warp and eliminate communication between warps
+  (so no need to use shared mem).
+
+  It fuses the softmax, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is intended for when the number of experts is a small
+  power of 2. 2) This implementation assumes k is small, but will work for any
+  k.
+*/
+
+template <typename T,
+          int VPT,
+          int NUM_EXPERTS,
+          int WARPS_PER_CTA,
+          int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+    void topk_gating_softmax(const T* input,
+                             T* output,
+                             const int64_t num_rows,
+                             int* indices,
+                             int* source_rows,
+                             const int64_t k) {
+  // We begin by enforcing compile time assertions and setting up compile time
+  // constants.
+  static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
+  static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
+                "NUM_EXPERTS must be power of 2");
+  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
+                "BYTES_PER_LDG must be power of 2");
+  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+  // Number of bytes each thread pulls in per load
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+  // Restrictions based on previous section.
+  static_assert(
+      VPT % ELTS_PER_LDG == 0,
+      "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+                "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+                "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+                "THREADS_PER_ROW can be at most warp size");
+
+  // We have NUM_EXPERTS elements per row. We specialize for small #experts
+  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+  // Restrictions for previous section.
+  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+                "The elts per row must cleanly divide the total elt per warp");
+
+  // ===================== From this point, we finally start computing run-time
+  // variables. ========================
+
+  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
+  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
+  // rows. We start by computing the start row for each block.
+  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+  // Now, using the base row per thread block, we compute the base row per warp.
+  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+  // The threads in a warp are split into sub-groups that will work on a row.
+  // We compute row offset for each thread sub-group
+  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row = warp_base_row + thread_row_in_warp;
+
+  // Threads with indices out of bounds should early exit here.
+  if (thread_row >= num_rows) return;
+  const bool should_process_row = true;
+
+  // We finally start setting up the read pointers for each thread. First, each
+  // thread jumps to the start of the row it will read.
+  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+  // Now, we compute the group each thread belong to in order to determine the
+  // first column to start loads.
+  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+  // Determine the pointer type to use to read in the data depending on the
+  // BYTES_PER_LDG template param. In theory, this can support all powers of 2
+  // up to 16.
+  using AccessType = cutlass::AlignedArray<T, ELTS_PER_LDG>;
+
+  // Finally, we pull in the data from global mem
+  cutlass::Array<T, VPT> row_chunk_input;
+  AccessType* row_chunk_vec_ptr =
+      reinterpret_cast<AccessType*>(&row_chunk_input);
+  const AccessType* vec_thread_read_ptr =
+      reinterpret_cast<const AccessType*>(thread_read_ptr);
+#pragma unroll
+  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+  }
+
+  using ComputeType = float;
+  using Converter = cutlass::NumericArrayConverter<ComputeType, T, VPT>;
+  Converter compute_type_converter;
+  cutlass::Array<ComputeType, VPT> row_chunk =
+      compute_type_converter(row_chunk_input);
+
+  // First, we perform a max reduce within the thread. We can do the max in fp16
+  // safely (I think) and just convert to float afterwards for the exp + sum
+  // reduction.
+  ComputeType thread_max = row_chunk[0];
+#pragma unroll
+  for (int ii = 1; ii < VPT; ++ii) {
+    thread_max = max(thread_max, row_chunk[ii]);
+  }
+
+// Now, we find the max within the thread group and distribute among the
+// threads. We use a butterfly reduce.
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    thread_max =
+        max(thread_max,
+            __shfl_xor_sync(0xFFFFFFFF, thread_max, mask, THREADS_PER_ROW));
+  }
+
+  // From this point, thread max in all the threads have the max within the row.
+  // Now, we subtract the max from each element in the thread and take the exp.
+  // We also compute the thread local sum.
+  float row_sum = 0;
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = expf(row_chunk[ii] - thread_max);
+    row_sum += row_chunk[ii];
+  }
+
+// Now, we perform the sum reduce within each thread group. Similar to the max
+// reduce, we use a bufferfly pattern.
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    row_sum += __shfl_xor_sync(0xFFFFFFFF, row_sum, mask, THREADS_PER_ROW);
+  }
+
+  // From this point, all threads have the max and the sum for their rows in the
+  // thread_max and thread_sum variables respectively. Finally, we can scale the
+  // rows for the softmax. Technically, for top-k gating we don't need to
+  // compute the entire softmax row. We can likely look at the maxes and only
+  // compute for the top-k values in the row. However, this kernel will likely
+  // not be a bottle neck and it seems better to closer match torch and find the
+  // argmax after computing the softmax.
+  const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+  }
+
+  // Now, softmax_res contains the softmax of the row chunk. Now, I want to find
+  // the topk elements in each row, along with the max index.​
+  int start_col = first_elt_read_by_thread;
+  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    // First, each thread does the local argmax
+    float max_val = row_chunk[0];
+    int expert = start_col;
+#pragma unroll
+    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+         ++ldg, col += COLS_PER_GROUP_LDG) {
+#pragma unroll
+      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+        // No check on the experts here since columns with the smallest index
+        // are processed first and only updated if > (not >=)
+        if (val > max_val) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads
+// reach consensus about the max. This will be useful for K > 1 so that the
+// threads can agree on "who" had the max value. That thread can then blank out
+// their max with -inf and the warp can run more iterations...
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      float other_max =
+          __shfl_xor_sync(0xFFFFFFFF, max_val, mask, THREADS_PER_ROW);
+      int other_expert =
+          __shfl_xor_sync(0xFFFFFFFF, expert, mask, THREADS_PER_ROW);
+
+      // We want lower indices to "win" in every thread so we break ties this
+      // way
+      if (other_max > max_val ||
+          (other_max == max_val && other_expert < expert)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    // Write the max for this k iteration to global memory.
+    if (thread_group_idx == 0) {
+      // The lead thread from each sub-group will write out the final results to
+      // global memory. (This will be a single) thread per row of the
+      // input/output matrices.
+      const int idx = k * thread_row + k_idx;
+      output[idx] = T(max_val);
+      indices[idx] = should_process_row ? expert : NUM_EXPERTS;
+      source_rows[idx] = k_idx * num_rows + thread_row;
+    }
+
+    // Finally, we clear the value in the thread with the current max if there
+    // is another iteration to run.
+    if (k_idx + 1 < k) {
+      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int thread_to_clear_in_group =
+          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+      // Only the thread in the group which produced the max will reset the
+      // "winning" value to -inf.
+      if (thread_group_idx == thread_to_clear_in_group) {
+        const int offset_for_expert = expert % ELTS_PER_LDG;
+        // Safe to set to any negative value since row_chunk values must be
+        // between 0 and 1.
+        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+            ComputeType(-10000.f);
+      }
+    }
+  }
+}
+
+namespace detail {
+// Constructs some constants needed to partition the work across threads at
+// compile time.
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 ||
+                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0,
+                "");
+  static constexpr int VECs_PER_THREAD =
+      std::max(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+}  // namespace detail
+
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void topk_gating_softmax_launcher_helper(const T* input,
+                                         T* output,
+                                         int* indices,
+                                         int* source_row,
+                                         const int64_t num_rows,
+                                         const int64_t num_experts,
+                                         const int64_t k,
+                                         cudaStream_t stream) {
+  static constexpr uint64_t MAX_BYTES_PER_LDG = 16;
+  static constexpr int BYTES_PER_LDG =
+      std::min(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = detail::TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int VPT = Constants::VPT;
+  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+  dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+  topk_gating_softmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+      <<<num_blocks, block_dim, 0, stream>>>(
+          input, output, num_rows, indices, source_row, k);
+}
+
+template <typename T>
+void topk_gating_softmax_kernelLauncher(const T* input,
+                                        T* output,
+                                        T* softmax,
+                                        int* indices,
+                                        int* source_row,
+                                        T* softmax_max_prob,
+                                        const int64_t num_rows,
+                                        const int64_t num_experts,
+                                        const int64_t k,
+                                        const bool group_moe,
+                                        cudaStream_t stream,
+                                        const bool topk_only_mode = false) {
+  if (topk_only_mode) {
+    static constexpr int TPB = 256;
+    const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
+    moe_top_k<T, TPB><<<config_topk.block_per_grid, TPB, 0, stream>>>(
+        input, output, indices, source_row, num_experts, k, num_rows);
+    return;
+  }
+  static constexpr int WARPS_PER_TB = 4;
+
+  #define LAUNCH_TOPK_GATING_SOFTMAX_HELPER(N)                                   \
+  case N: {                                                                    \
+    topk_gating_softmax_launcher_helper<T, N, WARPS_PER_TB>(                   \
+        input, output, indices, source_row, num_rows, num_experts, k, stream); \
+    break;                                                                     \
+  }
+  switch (num_experts) {
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(2)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(4)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(8)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(16)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(32)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(64)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(128)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(256)
+  
+    default: {
+      static constexpr int TPB = 256;
+      if (group_moe) {
+        const int group_experts = num_experts / k;
+        const int softmax_num_rows = num_rows * k;
+        const auto config_softmax = Get1DBlocksAnd2DGridsMoe(softmax_num_rows);
+        group_moe_softmax<T, TPB>
+            <<<config_softmax.block_per_grid, TPB, 0, stream>>>(
+                input,
+                softmax,
+                softmax_max_prob,
+                group_experts,
+                softmax_num_rows);
+        const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
+        moe_top_k<T, TPB>
+            <<<config_topk.block_per_grid, TPB, 0, stream>>>(softmax,
+                                                             output,
+                                                             indices,
+                                                             source_row,
+                                                             softmax_max_prob,
+                                                             num_experts,
+                                                             k,
+                                                             num_rows);
+      } else {
+        const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
+        moe_softmax<T, TPB><<<config_topk.block_per_grid, TPB, 0, stream>>>(
+            input, softmax, num_experts, num_rows);
+        moe_top_k<T, TPB>
+            <<<config_topk.block_per_grid, TPB, 0, stream>>>(softmax,
+                                                             output,
+                                                             indices,
+                                                             source_row,
+                                                             num_experts,
+                                                             k,
+                                                             num_rows);
+      }
+    }
+  }
+}
+
+// ========================== Permutation things
+// =======================================
+
+// Duplicated and permutes rows for MoE. In addition, reverse the permutation
+// map to help with finalizing routing.
+
+// "expanded_x_row" simply means that the number of values is num_rows x k. It
+// is "expanded" since we will have to duplicate some rows in the input matrix
+// to match the dimensions. Duplicates will always get routed to separate
+// experts in the end.
+
+// Note that the expanded_dest_row_to_expanded_source_row map referred to here
+// has indices in the range (0, k*rows_in_input - 1). However, it is set up so
+// that index 0, rows_in_input, 2*rows_in_input ... (k-1)*rows_in_input all map
+// to row 0 in the original matrix. Thus, to know where to read in the source
+// matrix, we simply take the modulus of the expanded index.
+
+template <typename T, int VecSize>
+__global__ void initialize_moe_routing_kernel(
+    const T* unpermuted_input,
+    T* permuted_output,
+    const int* expanded_dest_row_to_expanded_source_row,
+    int* expanded_source_row_to_expanded_dest_row,
+    const int64_t num_rows,
+    const int64_t active_rows,
+    const int64_t cols,
+    const int64_t num_rows_k) {
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec;
+
+  // Reverse permutation map.
+  // I do this so that later, we can use the source -> dest map to do the k-way
+  // reduction and unpermuting. I need the reverse map for that reduction to
+  // allow each threadblock to do 1 k-way reduce without atomics later in MoE. 1
+  // thread block will be responsible for all k summations.
+  const int expanded_dest_row = blockIdx.x + blockIdx.y * gridDim.x;
+  if (expanded_dest_row >= num_rows_k) return;
+  const int expanded_source_row =
+      expanded_dest_row_to_expanded_source_row[expanded_dest_row];
+  if (threadIdx.x == 0) {
+    expanded_source_row_to_expanded_dest_row[expanded_source_row] =
+        expanded_dest_row;
+  }
+
+  if ((blockIdx.x + blockIdx.y * gridDim.x) < active_rows) {
+    // Duplicate and permute rows
+    const int source_row = expanded_source_row % num_rows;
+
+    const T* source_row_ptr = unpermuted_input + source_row * cols;
+    T* dest_row_ptr = permuted_output + expanded_dest_row * cols;
+
+    for (int tid = threadIdx.x * VecSize; tid < cols;
+         tid += blockDim.x * VecSize) {
+      // dest_row_ptr[tid] = source_row_ptr[tid];
+      Load<T, VecSize>(&source_row_ptr[tid], &src_vec);
+      Store<T, VecSize>(src_vec, &dest_row_ptr[tid]);
+    }
+  }
+}
+
+template <typename T>
+void initialize_moe_routing_kernelLauncher(
+    const T* unpermuted_input,
+    T* permuted_output,
+    const int* expanded_dest_row_to_expanded_source_row,
+    int* expanded_source_row_to_expanded_dest_row,
+    const int64_t num_rows,
+    const int64_t active_rows,
+    const int64_t cols,
+    const int64_t k,
+    cudaStream_t stream) {
+  const int threads = std::min(cols, int64_t(1024));
+  constexpr int max_pack_size = 16 / sizeof(T);
+  const auto config_initialize = Get1DBlocksAnd2DGridsMoe(num_rows * k);
+  if (cols % max_pack_size == 0) {
+    initialize_moe_routing_kernel<T, max_pack_size>
+        <<<config_initialize.block_per_grid, threads, 0, stream>>>(
+            unpermuted_input,
+            permuted_output,
+            expanded_dest_row_to_expanded_source_row,
+            expanded_source_row_to_expanded_dest_row,
+            num_rows,
+            k * active_rows,
+            cols,
+            num_rows * k);
+  } else {
+    initialize_moe_routing_kernel<T, 1>
+        <<<config_initialize.block_per_grid, threads, 0, stream>>>(
+            unpermuted_input,
+            permuted_output,
+            expanded_dest_row_to_expanded_source_row,
+            expanded_source_row_to_expanded_dest_row,
+            num_rows,
+            k * active_rows,
+            cols,
+            num_rows * k);
+  }
+}
+
+// ============================== Infer GEMM sizes
+// =================================
+__device__ inline int find_total_elts_leq_target(int* sorted_indices,
+                                                 const int64_t arr_length,
+                                                 const int64_t target) {
+  int64_t low = 0, high = arr_length - 1, target_location = -1;
+  while (low <= high) {
+    int64_t mid = (low + high) / 2;
+
+    if (sorted_indices[mid] > target) {
+      high = mid - 1;
+    } else {
+      low = mid + 1;
+      target_location = mid;
+    }
+  }
+  return target_location + 1;
+}
+
+void compute_total_rows_before_expert(int* sorted_indices,
+                                      const int64_t total_indices,
+                                      const int64_t num_experts,
+                                      int64_t* total_rows_before_expert,
+                                      cudaStream_t stream);
+
+// Final kernel to unpermute and scale
+// This kernel unpermutes the original data, does the k-way reduction and
+// performs the final skip connection.
+template <typename T, int RESIDUAL_NUM>
+__global__ void finalize_moe_routing_kernel(
+    const T* expanded_permuted_rows,
+    T* reduced_unpermuted_output,
+    const T* bias,
+    const float* scales,
+    const int* expanded_source_row_to_expanded_dest_row,
+    const int* expert_for_source_row,
+    const int64_t cols,
+    const int64_t k,
+    const int64_t compute_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor,
+    const int64_t num_rows) {
+  const int original_row = blockIdx.x + blockIdx.y * gridDim.x;
+  // const int original_row = blockIdx.x;
+  // const int num_rows = gridDim.x;
+  if (original_row >= num_rows) return;
+  T* reduced_row_ptr = reduced_unpermuted_output + original_row * cols;
+
+  for (int tid = threadIdx.x; tid < cols; tid += blockDim.x) {
+    T thread_output{0.f};
+    float row_rescale{0.f};
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int expanded_original_row = original_row + k_idx * num_rows;
+      const int expanded_permuted_row =
+          expanded_source_row_to_expanded_dest_row[expanded_original_row];
+
+      const int64_t k_offset = original_row * k + k_idx;
+      const float row_scale = scales[k_offset];
+      row_rescale = row_rescale + row_scale;
+
+      const T* expanded_permuted_rows_row_ptr =
+          expanded_permuted_rows + expanded_permuted_row * cols;
+
+      const int expert_idx = expert_for_source_row[k_offset];
+      const T* bias_ptr = bias ? bias + expert_idx * cols : nullptr;
+      const T bias_value = bias_ptr ? bias_ptr[tid] : T{0.f};
+
+      thread_output =
+          static_cast<float>(thread_output) +
+          row_scale * static_cast<float>(
+                          expanded_permuted_rows_row_ptr[tid] +
+                          bias_value *
+                              static_cast<T>(static_cast<float>(compute_bias)));
+    }
+
+    thread_output = static_cast<float>(thread_output) /
+                    (norm_topk_prob ? row_rescale : 1.0f) *
+                    routed_scaling_factor;
+    reduced_row_ptr[tid] = thread_output;
+  }
+}
+
+template <typename T>
+void finalize_moe_routing_kernelLauncher(
+    const T* expanded_permuted_rows,
+    T* reduced_unpermuted_output,
+    const T* bias,
+    const float* scales,
+    const int* expanded_source_row_to_expanded_dest_row,
+    const int* expert_for_source_row,
+    const int64_t num_rows,
+    const int64_t cols,
+    const int64_t k,
+    const int64_t compute_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor,
+    cudaStream_t stream) {
+  const int threads = std::min(cols, int64_t(1024));
+  const auto config_final = Get1DBlocksAnd2DGridsMoe(num_rows);
+
+  finalize_moe_routing_kernel<T, 1>
+      <<<config_final.block_per_grid, threads, 0, stream>>>(
+          expanded_permuted_rows,
+          reduced_unpermuted_output,
+          bias,
+          scales,
+          expanded_source_row_to_expanded_dest_row,
+          expert_for_source_row,
+          cols,
+          k,
+          compute_bias,
+          norm_topk_prob,
+          routed_scaling_factor,
+          num_rows);
+}
+
+// ========================= TopK Softmax specializations
+// ===========================
+template void topk_gating_softmax_kernelLauncher(const float*,
+                                                 float*,
+                                                 float*,
+                                                 int*,
+                                                 int*,
+                                                 float*,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const bool,
+                                                 cudaStream_t,
+                                                 const bool);
+template void topk_gating_softmax_kernelLauncher(const half*,
+                                                 half*,
+                                                 half*,
+                                                 int*,
+                                                 int*,
+                                                 half*,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const bool,
+                                                 cudaStream_t,
+                                                 const bool);
+#ifdef PADDLE_CUDA_BF16
+template void topk_gating_softmax_kernelLauncher(const __nv_bfloat16*,
+                                                 __nv_bfloat16*,
+                                                 __nv_bfloat16*,
+                                                 int*,
+                                                 int*,
+                                                 __nv_bfloat16*,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const bool,
+                                                 cudaStream_t,
+                                                 const bool);
+#endif
+// ===================== Specializations for init routing
+// =========================
+template void initialize_moe_routing_kernelLauncher(const float*,
+                                                    float*,
+                                                    const int*,
+                                                    int*,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    cudaStream_t);
+template void initialize_moe_routing_kernelLauncher(const half*,
+                                                    half*,
+                                                    const int*,
+                                                    int*,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    cudaStream_t);
+#ifdef PADDLE_CUDA_BF16
+template void initialize_moe_routing_kernelLauncher(const __nv_bfloat16*,
+                                                    __nv_bfloat16*,
+                                                    const int*,
+                                                    int*,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    cudaStream_t);
+#endif
+// ==================== Specializations for final routing
+// ===================================
+template void finalize_moe_routing_kernelLauncher(const float*,
+                                                  float*,
+                                                  const float*,
+                                                  const float*,
+                                                  const int*,
+                                                  const int*,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const bool,
+                                                  const float,
+                                                  cudaStream_t);
+template void finalize_moe_routing_kernelLauncher(const half*,
+                                                  half*,
+                                                  const half*,
+                                                  const float*,
+                                                  const int*,
+                                                  const int*,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const bool,
+                                                  const float,
+                                                  cudaStream_t);
+#ifdef PADDLE_CUDA_BF16
+template void finalize_moe_routing_kernelLauncher(const __nv_bfloat16*,
+                                                  __nv_bfloat16*,
+                                                  const __nv_bfloat16*,
+                                                  const float*,
+                                                  const int*,
+                                                  const int*,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const bool,
+                                                  const float,
+                                                  cudaStream_t);
+#endif

csrc/gpu/moe/fused_moe/moe_dispatch.cu

@@ -0,0 +1,276 @@
+// 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.
+
+
+// Ignore CUTLASS warnings about type punning
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#pragma GCC diagnostic ignored "-Wunused-function"
+#pragma once
+
+#include "moe/fused_moe_helper.h"
+#include "moe/fused_moe_op.h"
+#pragma GCC diagnostic pop
+
+#include "helper.h"
+
+
+template <paddle::DataType T>
+void MoeDispatchKernel(const paddle::Tensor& input,
+                       const paddle::Tensor& gating_output,
+                       const int moe_topk,
+                       const bool group_moe,
+                       const bool topk_only_mode,
+                       const int num_rows,
+                       const int hidden_size,
+                       const int expert_num,
+                       paddle::Tensor* permute_input,
+                       paddle::Tensor* tokens_expert_prefix_sum,
+                       paddle::Tensor* permute_indices_per_token,
+                       paddle::Tensor* top_k_weight,
+                       paddle::Tensor* top_k_indices) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto stream = input.stream();
+  auto place = input.place();
+
+  if (group_moe) {
+    // Check if expert_num is divisible by moe_topk, else throw an error
+    PADDLE_ENFORCE_EQ(expert_num % moe_topk,
+                      0,
+                      common::errors::InvalidArgument(
+                          "The number of experts (expert_num) "
+                          "must be divisible by moe_topk. "
+                          "Got expert_num = %d and moe_topk = %d.",
+                          expert_num,
+                          moe_topk));
+  }
+
+  const int num_moe_inputs = AlignTo16(num_rows * moe_topk);
+  const int bytes = num_moe_inputs * sizeof(int);
+
+  CubKeyValueSorter sorter_;
+  sorter_.update_num_experts(expert_num);
+
+  const int sorter_ws_size_bytes =
+      AlignTo16(sorter_.getWorkspaceSize(moe_topk * num_rows));
+  const int sort_tmp_in_out_size = num_moe_inputs * 2 * sizeof(int);
+
+  paddle::Tensor ws_ptr_tensor =
+      GetEmptyTensor({bytes + sorter_ws_size_bytes + sort_tmp_in_out_size},
+                     paddle::DataType::INT8,
+                     place);
+
+  int8_t* ws_ptr = ws_ptr_tensor.data<int8_t>();
+  int* source_rows_ = reinterpret_cast<int*>(ws_ptr);
+  int8_t* sorter_ws_ptr = reinterpret_cast<int8_t*>(ws_ptr + bytes);
+  int* permuted_experts_ =
+      reinterpret_cast<int*>(sorter_ws_ptr + sorter_ws_size_bytes);
+  int* permuted_rows_ = permuted_experts_ + num_moe_inputs;
+
+  int* expert_for_source_row = top_k_indices->data<int>();
+
+  float* softmax_max_prob = nullptr;
+  if (group_moe) {
+    paddle::Tensor softmax_max_prob_tensor =
+        GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+    // (TODO: check fill success ?)
+    paddle::experimental::fill(softmax_max_prob_tensor, 0.f);
+    softmax_max_prob = softmax_max_prob_tensor.data<float>();
+  }
+
+  float* softmax_out_;
+
+  const bool is_pow_2 =
+      (expert_num != 0) && ((expert_num & (expert_num - 1)) == 0);
+
+  paddle::Tensor softmax_buffer;
+
+  if (!is_pow_2 || expert_num > 256 || group_moe) {
+    softmax_buffer = GetEmptyTensor(
+        {num_rows * expert_num}, paddle::DataType::FLOAT32, place);
+    softmax_out_ = softmax_buffer.data<float>();
+  } else {
+    softmax_out_ = nullptr;
+  }
+
+  topk_gating_softmax_kernelLauncher<float>(gating_output.data<float>(),
+                                            top_k_weight->data<float>(),
+                                            softmax_out_,
+                                            expert_for_source_row,
+                                            source_rows_,
+                                            softmax_max_prob,
+                                            num_rows,
+                                            expert_num,
+                                            moe_topk,
+                                            group_moe,
+                                            stream,
+                                            topk_only_mode);
+
+  sorter_.run(reinterpret_cast<void*>(sorter_ws_ptr),
+              sorter_ws_size_bytes,
+              expert_for_source_row,
+              permuted_experts_,
+              source_rows_,
+              permuted_rows_,
+              moe_topk * num_rows,
+              false,
+              stream);
+
+
+  initialize_moe_routing_kernelLauncher(
+      input.data<data_t>(),
+      permute_input->data<data_t>(),
+      permuted_rows_,
+      permute_indices_per_token->data<int32_t>(),
+      num_rows,
+      num_rows,
+      hidden_size,
+      moe_topk,
+      stream);
+
+
+  compute_total_rows_before_expert(
+      permuted_experts_,
+      moe_topk * num_rows,
+      expert_num,
+      tokens_expert_prefix_sum->data<int64_t>(),
+      stream);
+}
+
+
+std::vector<paddle::Tensor> MoeExpertDispatch(
+    const paddle::Tensor& input,
+    const paddle::Tensor& gating_output,
+    const int moe_topk,
+    const bool group_moe,
+    const bool topk_only_mode) {
+  const auto input_type = input.dtype();
+  auto place = input.place();
+  int token_rows = 0;
+  auto input_dims = input.dims();
+  auto gating_dims = gating_output.dims();
+  const int expert_num = gating_dims[gating_dims.size() - 1];
+
+  if (input_dims.size() == 3) {
+    token_rows = input_dims[0] * input_dims[1];
+  } else {
+    token_rows = input_dims[0];
+  }
+  const int num_rows = token_rows;
+  const int hidden_size = input.dims()[input_dims.size() - 1];
+
+  auto permute_input =
+      GetEmptyTensor({moe_topk * num_rows, hidden_size}, input_type, place);
+  // correspond to the weighted coefficients of the results from each expert.
+  auto top_k_weight =
+      GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+  auto top_k_indices =
+      GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::INT32, place);
+
+  auto tokens_expert_prefix_sum =
+      GetEmptyTensor({expert_num}, paddle::DataType::INT64, place);
+  auto permute_indices_per_token =
+      GetEmptyTensor({moe_topk, num_rows}, paddle::DataType::INT32, place);
+
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      MoeDispatchKernel<paddle::DataType::BFLOAT16>(input,
+                                                    gating_output,
+                                                    moe_topk,
+                                                    group_moe,
+                                                    topk_only_mode,
+                                                    num_rows,
+                                                    hidden_size,
+                                                    expert_num,
+                                                    &permute_input,
+                                                    &tokens_expert_prefix_sum,
+                                                    &permute_indices_per_token,
+                                                    &top_k_weight,
+                                                    &top_k_indices);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeDispatchKernel<paddle::DataType::FLOAT16>(input,
+                                                   gating_output,
+                                                   moe_topk,
+                                                   group_moe,
+                                                   topk_only_mode,
+                                                   num_rows,
+                                                   hidden_size,
+                                                   expert_num,
+                                                   &permute_input,
+                                                   &tokens_expert_prefix_sum,
+                                                   &permute_indices_per_token,
+                                                   &top_k_weight,
+                                                   &top_k_indices);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeDispatchKernel");
+  }
+  return {permute_input,
+          tokens_expert_prefix_sum,
+          permute_indices_per_token,
+          top_k_weight,
+          top_k_indices};
+}
+
+
+std::vector<std::vector<int64_t>> MoeExpertDispatchInferShape(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<int64_t>& gating_output_shape,
+    const int moe_topk) {
+  int token_rows = -1;
+
+  if (input_shape.size() == 3) {
+    token_rows = input_shape[0] * input_shape[1];
+  } else {
+    token_rows = input_shape[0];
+  }
+  const int expert_num = gating_output_shape[gating_output_shape.size() - 1];
+  const int num_rows = token_rows;
+  const int hidden_size = input_shape[input_shape.size() - 1];
+
+  return {{moe_topk * num_rows, hidden_size},
+          {expert_num},
+          {moe_topk, num_rows},
+          {num_rows, moe_topk},
+          {num_rows, moe_topk}};
+}
+
+std::vector<paddle::DataType> MoeExpertDispatchInferDtype(
+    const paddle::DataType& input_dtype,
+    const paddle::DataType& gating_output_dtype,
+    const int moe_topk) {
+  return {input_dtype,
+          paddle::DataType::INT64,
+          paddle::DataType::INT32,
+          paddle::DataType::FLOAT32,
+          paddle::DataType::INT32};
+}
+
+
+PD_BUILD_OP(moe_expert_dispatch)
+    .Inputs({"input", "gating_output"})
+    .Outputs({"permute_input",
+              "tokens_expert_prefix_sum",
+              "permute_indices_per_token",
+              "top_k_weight",
+              "top_k_indices"})
+    .Attrs({"moe_topk:int", "group_moe:bool", "topk_only_mode:bool"})
+    .SetKernelFn(PD_KERNEL(MoeExpertDispatch))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeExpertDispatchInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeExpertDispatchInferDtype));

csrc/gpu/moe/fused_moe/moe_ffn.cu

@@ -0,0 +1,232 @@
+// 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.
+
+#pragma once
+#include "cutlass/numeric_conversion.h"
+#include "helper.h"
+#include "moe/fused_moe_helper.h"
+
+template <paddle::DataType T>
+void MoeFFNKernel(const paddle::Tensor& permute_input,
+                  const paddle::Tensor& tokens_expert_prefix_sum,
+                  const paddle::Tensor& ffn1_weight,
+                  const paddle::Tensor& ffn2_weight,
+                  const paddle::optional<paddle::Tensor>& ffn1_bias,
+                  const paddle::optional<paddle::Tensor>& ffn1_scale,
+                  const paddle::optional<paddle::Tensor>& ffn2_scale,
+                  const std::string& quant_method,
+                  paddle::Tensor ffn_out) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto ffn_out_data = ffn_out.data<data_t>();
+  auto permuted_data = permute_input.data<data_t>();
+  auto place = permute_input.place();
+  auto input_type = permute_input.dtype();
+  auto stream = permute_input.stream();
+
+  auto fp16_moe_gemm_runner = MoeGemmRunner<DataType_, DataType_>();
+  auto int8_moe_gemm_runner = MoeGemmRunner<DataType_, uint8_t>();
+  auto int4_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::uint4b_t>();
+
+  const int64_t expanded_active_expert_rows = permute_input.dims()[0];
+  const int num_experts = ffn1_weight.dims()[0];
+  const int hidden_size = ffn1_weight.dims()[1];
+  int inter_dim = ffn1_weight.dims()[2];
+
+  if (quant_method == "weight_only_int4") {
+    inter_dim = inter_dim * 2;
+  }
+
+  const int64_t inter_size = inter_dim;
+
+  paddle::Tensor fc1_out_tensor = GetEmptyTensor(
+      {expanded_active_expert_rows, inter_size}, input_type, place);
+  auto fc1_out = fc1_out_tensor.data<data_t>();
+
+  using NvType = typename traits_::DataType;
+
+  auto fc1_expert_biases =
+      ffn1_bias
+          ? const_cast<paddle::Tensor*>(ffn1_bias.get_ptr())->data<data_t>()
+          : nullptr;
+
+  if (quant_method == "weight_only_int8") {
+    int8_moe_gemm_runner.moe_gemm_bias_act(
+        reinterpret_cast<const NvType*>(permuted_data),
+        reinterpret_cast<const uint8_t*>(ffn1_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn1_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<const NvType*>(fc1_expert_biases),
+        reinterpret_cast<NvType*>(fc1_out),
+        const_cast<int64_t*>(tokens_expert_prefix_sum.data<int64_t>()),
+        expanded_active_expert_rows,
+        inter_size,
+        hidden_size,
+        num_experts,
+        "none",
+        stream);
+  } else if (quant_method == "weight_only_int4") {
+    int4_moe_gemm_runner.moe_gemm_bias_act(
+        reinterpret_cast<const NvType*>(permuted_data),
+        reinterpret_cast<const cutlass::uint4b_t*>(ffn1_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn1_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<const NvType*>(fc1_expert_biases),
+        reinterpret_cast<NvType*>(fc1_out),
+        const_cast<int64_t*>(tokens_expert_prefix_sum.data<int64_t>()),
+        expanded_active_expert_rows,
+        inter_size,
+        hidden_size,
+        num_experts,
+        "none",
+        stream);
+  } else {
+    fp16_moe_gemm_runner.moe_gemm_bias_act(
+        reinterpret_cast<const NvType*>(permuted_data),
+        reinterpret_cast<const NvType*>(ffn1_weight.data<data_t>()),
+        nullptr,
+        reinterpret_cast<const NvType*>(fc1_expert_biases),
+        reinterpret_cast<NvType*>(fc1_out),
+        const_cast<int64_t*>(tokens_expert_prefix_sum.data<int64_t>()),
+        expanded_active_expert_rows,
+        inter_size,
+        hidden_size,
+        num_experts,
+        "none",
+        stream);
+  }
+
+  auto act_out_tensor = paddle::experimental::swiglu(fc1_out_tensor, nullptr);
+  auto act_out = act_out_tensor.data<data_t>();
+
+  if (quant_method == "weight_only_int8") {
+    int8_moe_gemm_runner.moe_gemm(
+        reinterpret_cast<const NvType*>(act_out),
+        reinterpret_cast<const uint8_t*>(ffn2_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn2_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<NvType*>(ffn_out_data),
+        const_cast<int64_t*>(tokens_expert_prefix_sum.data<int64_t>()),
+        expanded_active_expert_rows,
+        hidden_size,
+        inter_size / 2,
+        num_experts,
+        stream);
+
+  } else if (quant_method == "weight_only_int4") {
+    int4_moe_gemm_runner.moe_gemm(
+        reinterpret_cast<const NvType*>(act_out),
+        reinterpret_cast<const cutlass::uint4b_t*>(ffn2_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn2_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<NvType*>(ffn_out_data),
+        const_cast<int64_t*>(tokens_expert_prefix_sum.data<int64_t>()),
+        expanded_active_expert_rows,
+        hidden_size,
+        inter_size / 2,
+        num_experts,
+        stream);
+  } else {
+    fp16_moe_gemm_runner.moe_gemm(
+        reinterpret_cast<const NvType*>(act_out),
+        reinterpret_cast<const NvType*>(ffn2_weight.data<data_t>()),
+        nullptr,
+        reinterpret_cast<NvType*>(ffn_out_data),
+        const_cast<int64_t*>(tokens_expert_prefix_sum.data<int64_t>()),
+        expanded_active_expert_rows,
+        hidden_size,
+        inter_size / 2,
+        num_experts,
+        stream);
+  }
+}
+
+std::vector<paddle::Tensor> MoeExpertFFN(
+    const paddle::Tensor& permute_input,
+    const paddle::Tensor& tokens_expert_prefix_sum,
+    const paddle::Tensor& ffn1_weight,
+    const paddle::Tensor& ffn2_weight,
+    const paddle::optional<paddle::Tensor>& ffn1_bias,
+    const paddle::optional<paddle::Tensor>& ffn1_scale,
+    const paddle::optional<paddle::Tensor>& ffn2_scale,
+    const std::string& quant_method) {
+  const auto input_type = permute_input.dtype();
+  auto ffn_out = paddle::empty_like(permute_input);
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      MoeFFNKernel<paddle::DataType::BFLOAT16>(permute_input,
+                                               tokens_expert_prefix_sum,
+                                               ffn1_weight,
+                                               ffn2_weight,
+                                               ffn1_bias,
+                                               ffn1_scale,
+                                               ffn2_scale,
+                                               quant_method,
+                                               ffn_out);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeFFNKernel<paddle::DataType::FLOAT16>(permute_input,
+                                              tokens_expert_prefix_sum,
+                                              ffn1_weight,
+                                              ffn2_weight,
+                                              ffn1_bias,
+                                              ffn1_scale,
+                                              ffn2_scale,
+                                              quant_method,
+                                              ffn_out);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeExpertFFN");
+  }
+  return {ffn_out};
+}
+
+std::vector<std::vector<int64_t>> MoeExpertFFNInferShape(
+    const std::vector<int64_t>& permute_input_shape,
+    const std::vector<int64_t>& tokens_expert_prefix_sum_shape,
+    const std::vector<int64_t>& ffn1_weight_shape,
+    const std::vector<int64_t>& ffn2_weight_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_scale_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_scale_shape) {
+  return {permute_input_shape};
+}
+
+std::vector<paddle::DataType> MoeExpertFFNInferDtype(
+    const paddle::DataType& permute_input_dtype,
+    const paddle::DataType& tokens_expert_prefix_sum_dtype,
+    const paddle::DataType& ffn1_weight_dtype,
+    const paddle::DataType& ffn2_weight_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_bias_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_scale_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_scale_dtype) {
+  return {permute_input_dtype};
+}
+
+PD_BUILD_OP(moe_expert_ffn)
+    .Inputs({"permute_input",
+             "tokens_expert_prefix_sum",
+             "ffn1_weight",
+             "ffn2_weight",
+             paddle::Optional("ffn1_bias"),
+             paddle::Optional("ffn1_scale"),
+             paddle::Optional("ffn2_scale")})
+    .Outputs({"output_tensor"})
+    .Attrs({"quant_method:std::string"})
+    .SetKernelFn(PD_KERNEL(MoeExpertFFN))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeExpertFFNInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeExpertFFNInferDtype));

csrc/gpu/moe/fused_moe/moe_reduce.cu

@@ -0,0 +1,146 @@
+// 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.
+
+
+// Ignore CUTLASS warnings about type punning
+
+
+#pragma once
+
+#include "helper.h"
+#include "moe/fused_moe_helper.h"
+#include "moe/fused_moe_op.h"
+
+template <paddle::DataType T>
+void MoeReduceKernel(const paddle::Tensor& ffn_out,
+                     const paddle::Tensor& top_k_weight,
+                     const paddle::Tensor& permute_indices_per_token,
+                     const paddle::Tensor& top_k_indices,
+                     const paddle::optional<paddle::Tensor>& ffn2_bias,
+                     const bool norm_topk_prob,
+                     const float routed_scaling_factor,
+                     const int num_rows,
+                     const int hidden_size,
+                     const int topk,
+                     paddle::Tensor* output) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  auto stream = ffn_out.stream();
+
+  finalize_moe_routing_kernelLauncher(
+      ffn_out.data<data_t>(),
+      output->data<data_t>(),
+      ffn2_bias ? ffn2_bias->data<data_t>() : nullptr,
+      top_k_weight.data<float>(),
+      permute_indices_per_token.data<int32_t>(),
+      top_k_indices.data<int>(),
+      num_rows,
+      hidden_size,
+      topk,
+      static_cast<int>(1),
+      norm_topk_prob,
+      routed_scaling_factor,
+      stream);
+}
+
+
+std::vector<paddle::Tensor> MoeExpertReduce(
+    const paddle::Tensor& ffn_out,
+    const paddle::Tensor& top_k_weight,
+    const paddle::Tensor& permute_indices_per_token,
+    const paddle::Tensor& top_k_indices,
+    const paddle::optional<paddle::Tensor>& ffn2_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor) {
+  const auto input_type = ffn_out.dtype();
+  auto place = ffn_out.place();
+
+  const int topk = top_k_indices.dims()[1];
+  const int num_rows = ffn_out.dims()[0] / topk;
+  const int hidden_size = ffn_out.dims()[1];
+
+  auto output = GetEmptyTensor({num_rows, hidden_size}, input_type, place);
+
+  // Avoids ‘invalid configuration argument’ when we launch the kernel.
+  if (ffn_out.dims()[0] == 0) return {output};
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      MoeReduceKernel<paddle::DataType::BFLOAT16>(ffn_out,
+                                                  top_k_weight,
+                                                  permute_indices_per_token,
+                                                  top_k_indices,
+                                                  ffn2_bias,
+                                                  norm_topk_prob,
+                                                  routed_scaling_factor,
+                                                  num_rows,
+                                                  hidden_size,
+                                                  topk,
+                                                  &output);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeReduceKernel<paddle::DataType::BFLOAT16>(ffn_out,
+                                                  top_k_weight,
+                                                  permute_indices_per_token,
+                                                  top_k_indices,
+                                                  ffn2_bias,
+                                                  norm_topk_prob,
+                                                  routed_scaling_factor,
+                                                  num_rows,
+                                                  hidden_size,
+                                                  topk,
+                                                  &output);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeDispatchKernel");
+  }
+  return {output};
+}
+
+
+std::vector<std::vector<int64_t>> MoeExpertReduceInferShape(
+    const std::vector<int64_t>& ffn_out_shape,
+    const std::vector<int64_t>& top_k_weight_shape,
+    const std::vector<int64_t>& permute_indices_per_token_shape,
+    const std::vector<int64_t>& top_k_indices_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_bias_shape) {
+  const int topk = top_k_indices_shape[1];
+  std::vector<int64_t> fused_moe_out_shape = {ffn_out_shape[0] / topk,
+                                              ffn_out_shape[1]};
+
+  return {fused_moe_out_shape};
+}
+
+std::vector<paddle::DataType> MoeExpertReduceInferDtype(
+    const paddle::DataType& ffn_out_dtype,
+    const paddle::DataType& top_k_weight_dtype,
+    const paddle::DataType& permute_indices_per_token_dtype,
+    const paddle::DataType& top_k_indices_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_bias_dtype) {
+  return {ffn_out_dtype};
+}
+
+
+PD_BUILD_OP(moe_expert_reduce)
+    .Inputs({"ffn_out",
+             "top_k_weight",
+             "permute_indices_per_token",
+             "top_k_indices",
+             paddle::Optional("ffn2_bias")})
+    .Outputs({"output"})
+    .Attrs({"norm_topk_prob:bool", "routed_scaling_factor:float"})
+    .SetKernelFn(PD_KERNEL(MoeExpertReduce))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeExpertReduceInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeExpertReduceInferDtype));

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/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm89.cu

@@ -568,7 +568,7 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
 
   normalize_d<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(RO, m, d);
 
-  // ! here we just implement the case for fp32 acumulation
+  // ! here we just implement the case for fp32 accumulation
   if constexpr (fuse_v_scale)
   {
     float v_scale[4];

csrc/gpu/sage_attn_kernels/sageattn_utils.cuh

@@ -48,7 +48,7 @@
 
 #include "paddle/extension.h"
 
-// currently we do not support INT4, so we implement INT8 sage attention inference temperarily
+// currently we do not support INT4, so we implement INT8 sage attention inference temporarily
 
 #define FINAL_MASK 0xffffffff
 #define WARP_SIZE 32
@@ -1288,7 +1288,7 @@ using b128_t = uint4;
  * permutation in the granularity of 128bit.
  * 
  * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless,
- * meaning it does not maintain any information about the current pointer position. The offset returnd by 
+ * meaning it does not maintain any information about the current pointer position. The offset returned by 
  * the struct can be used to access the shared memory through the provided interface.
  * 
  * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j).

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;
-        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_fused_moe.py

@@ -0,0 +1,210 @@
+# 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
+from paddlenlp_ops import (
+    fused_expert_moe,
+    moe_expert_dispatch,
+    moe_expert_ffn,
+    moe_expert_reduce,
+)
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.incubate.nn.functional import swiglu
+
+paddle.seed(42)
+np.random.seed(42)
+
+class Expert(nn.Layer):
+    def __init__(self, d_model, d_feedforward):
+        super().__init__()
+        self.fc1 = nn.Linear(d_model, d_feedforward * 2)
+        self.fc2 = nn.Linear(d_feedforward, d_model)
+
+    def forward(self, x, idx):
+        x = self.fc1(x)
+        x = swiglu(x)
+        return self.fc2(x)
+
+class TestSimpleFusedMoe(unittest.TestCase):
+    def setUp(self):
+        self.set_config()
+        self.init_parameters()
+        self.prepare_data()
+
+
+    def set_config(self):
+        self.dtype = "bfloat16"
+        self.batch_size = 16
+        self.seq_len = 128
+        self.num_experts = 4
+        self.d_model = 64
+        self.d_feedforward = 128
+        self.top_k = 2
+        self.rtol = 1e-2
+        self.atol = 1e-2
+
+        paddle.set_default_dtype(self.dtype)
+
+    def init_parameters(self):
+        # 创建专家层
+        self.experts = nn.LayerList([
+            Expert(self.d_model, self.d_feedforward)
+            for _ in range(self.num_experts)
+        ])
+
+        # 初始化门控权重
+        self.gate = nn.Linear(self.d_model, self.num_experts)
+        self.gate_weight = self.gate.weight.cast("float32")
+
+    def prepare_data(self):
+        """准备输入数据"""
+        self.x = paddle.randn(
+            [self.batch_size, self.seq_len, self.d_model],
+            dtype=self.dtype
+        )
+
+        self.w0 = paddle.stack(
+            [e.fc1.weight for e in self.experts], axis=0
+        ).astype(self.dtype)
+        self.b0 = paddle.stack(
+            [e.fc1.bias for e in self.experts], axis=0
+        ).reshape([self.num_experts, 1, -1]).astype(self.dtype)
+
+        self.w1 = paddle.stack(
+            [e.fc2.weight for e in self.experts], axis=0
+        ).astype(self.dtype)
+        self.b1 = paddle.stack(
+            [e.fc2.bias for e in self.experts], axis=0
+        ).reshape([self.num_experts, 1, -1]).astype(self.dtype)
+
+    def baseline_forward(self, hidden_states):
+        """（逐个专家计算）"""
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.reshape([-1, hidden_dim])
+
+        # 路由计算
+        logits = paddle.matmul(hidden_states.cast("float32"), self.gate_weight)
+        weights = F.softmax(logits, axis=-1)
+        routing_weights, selected_experts = paddle.topk(weights, self.top_k, axis=-1)
+        # 结果累加
+        final_hidden_states = paddle.zeros_like(hidden_states)
+
+        expert_mask = paddle.transpose(
+            F.one_hot(selected_experts, num_classes=self.num_experts), [2, 1, 0]
+        )
+
+        for expert_id in range(self.num_experts):
+            expert_layer = self.experts[expert_id]
+            idx, top_x = paddle.where(expert_mask[expert_id])
+
+            current_state = paddle.index_select(
+                hidden_states, top_x, axis=0
+            ).reshape([-1, hidden_dim])
+            current_hidden_states = (
+                expert_layer(current_state, expert_id)
+                * routing_weights[top_x, idx]
+            )
+            paddle.index_add_(
+                x=final_hidden_states,
+                index=top_x.squeeze(),
+                axis=0,
+                value=current_hidden_states.to(hidden_states.dtype),
+            )
+        final_hidden_states = paddle.reshape(
+            final_hidden_states, [batch_size, seq_len, hidden_dim]
+        )
+        return final_hidden_states
+
+    def fused_forward(self, x):
+        """测试融合实现"""
+        return fused_expert_moe(
+            x,
+            self.gate_weight,
+            self.w0,
+            self.w1,
+            self.b0,
+            None,
+            self.b1,
+            None,
+            "None",
+            self.top_k,
+            False,
+            False
+        )
+
+    def split_forward(self, hidden_states):
+        '''测试拆分实现'''
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.reshape([-1, hidden_dim])
+
+        # 路由计算
+        logits = paddle.matmul(hidden_states.cast("float32"), self.gate_weight)
+        scores = F.softmax(logits, axis=-1)
+        (
+            permute_input,
+            token_nums_per_expert,
+            permute_indices_per_token,
+            top_k_weights,
+            top_k_indices,
+        ) = moe_expert_dispatch(hidden_states, scores, self.top_k, False, topk_only_mode=True)
+
+        ffn_out = moe_expert_ffn(
+            permute_input,
+            token_nums_per_expert,
+            self.w0,
+            self.w1,
+            self.b0,
+            None,
+            None,
+            "none",
+        )
+        output = moe_expert_reduce(
+                ffn_out,
+                top_k_weights,
+                permute_indices_per_token,
+                top_k_indices,
+                None,
+                norm_topk_prob=False,
+                routed_scaling_factor=1.0,
+            )
+        output = paddle.reshape(output, [batch_size, seq_len, hidden_dim])
+        return output
+
+    def test_consistency(self):
+        """测试一致性"""
+        base_out = self.baseline_forward(self.x)
+        fused_out = self.fused_forward(self.x)
+        split_out = self.split_forward(self.x)
+
+        np.testing.assert_allclose(
+            base_out.cast("float32").numpy().astype("float32"),
+            fused_out.cast("float32").numpy().astype("float32"),
+            rtol=self.rtol,
+            atol=self.atol
+        )
+        np.testing.assert_allclose(
+            base_out.cast("float32").numpy().astype("float32"),
+            split_out.cast("float32").numpy().astype("float32"),
+            rtol=self.rtol,
+            atol=self.atol
+        )
+        print("Test passed! Outputs match within tolerance.")
+
+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:
@@ -80,7 +92,7 @@ def get_gencode_flags():
 
 
 gencode_flags = get_gencode_flags()
-library_path = os.environ.get("LD_LIBRARY_PATH", "/usr/local/cuda/lib64")
+library_path = [os.environ.get("LD_LIBRARY_PATH", "/usr/local/cuda/lib64")]
 
 sources = [
     "./gpu/save_with_output.cc",
@@ -105,8 +117,7 @@ sources = [
     "./gpu/step.cu",
     "./gpu/quant_int8.cu",
     "./gpu/dequant_int8.cu",
-    "./gpu/group_quant.cu",
-    "./gpu/preprocess_for_moe.cu",
+    "./gpu/moe/preprocess_for_moe.cu",
     "./gpu/get_position_ids_and_mask_encoder_batch.cu",
     "./gpu/fused_rotary_position_encoding.cu",
     "./gpu/flash_attn_bwd.cc",
@@ -118,6 +129,10 @@ 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/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")
 
@@ -132,18 +147,22 @@ nvcc_compile_args += [
     "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
     "-U__CUDA_NO_BFLOAT162_OPERATORS__",
     "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
-    "-Igpu",
-    "-Igpu/cutlass_kernels",
-    "-Igpu/fp8_gemm_with_cutlass",
-    "-Igpu/cutlass_kernels/fp8_gemm_fused/autogen",
-    "-Ithird_party/cutlass/include",
-    "-Ithird_party/cutlass/tools/util/include",
-    "-Ithird_party/nlohmann_json/single_include",
-    "-Igpu/sample_kernels",
 ]
 
+include_dirs = [
+    "./gpu",
+    "./gpu/cutlass_kernels",
+    "./gpu/fp8_gemm_with_cutlass",
+    "./gpu/cutlass_kernels/fp8_gemm_fused/autogen",
+    "./third_party/cutlass/include",
+    "./third_party/cutlass/tools/util/include",
+    "./third_party/nlohmann_json/single_include",
+    "./gpu/sample_kernels",
+    "./gpu/moe/fused_moe",
+]
 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"]
@@ -152,12 +171,16 @@ 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"
 if os.path.isdir(fp8_auto_gen_directory):
     shutil.rmtree(fp8_auto_gen_directory)
 
+
 if cc == 89 and cuda_version >= 12.4:
     os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels.py --cuda_arch 89")
     os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels.py --cuda_arch 89")
@@ -168,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",
@@ -177,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")
@@ -209,12 +229,17 @@ if cc >= 90 and cuda_version >= 12.0:
     sources += find_end_files("./gpu/mla_attn", ".cu")
 
 ops_name = f"paddlenlp_ops_{sm_version}" if sm_version != 0 else "paddlenlp_ops"
+
 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],
+        library_dirs=library_path,
+        include_dirs=include_dirs,
     ),
 )

csrc/tools/build_wheel.sh

@@ -60,8 +60,8 @@ function generate_sm_version(){
     if [ ! -z "$SM_VERSION" ]; then
         sm_versions=($SM_VERSION )
     elif [ "$ARCHITECTURE" = "all" ]; then
-        if echo "$cuda_version >= 12.0" | awk '{if ($0) exit 0; exit 1}'; then
-          sm_versions=(70 80 80 86 89 90 )
+        if awk -v version="$cuda_version" 'BEGIN { exit !(version >= 12.0) }'; then
+          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

@@ -0,0 +1,100 @@
+// 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> AdjustBatch(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
+                                             const paddle::Tensor& cum_offsets, // [bsz, 1]
+                                             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);
+  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];  
+  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>())};
+
+  auto out = paddle::full({token_num, dim}, -2, tmp_out.type(), tmp_out.place()); 
+
+  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>()),
+          encoder_seqs_lods_vp,
+          encoder_batch_map_vp,
+          decoder_batch_map_vp,
+          dim);
+  return {out};
+}
+
+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>& 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.");
+    }
+    int64_t token_num = tmp_out_shape[0];
+    int64_t dim_embed = tmp_out_shape[1];
+    return {{token_num, dim_embed}};
+}
+
+std::vector<paddle::DataType> AdjustBatchInferDtype(const paddle::DataType& tmp_out_dtype,
+                                                         const paddle::DataType& cum_offsets_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", "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))
+    .SetInferShapeFn(PD_INFER_SHAPE(AdjustBatchInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(AdjustBatchInferDtype));

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

@@ -0,0 +1,98 @@
+// 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> GatherNextToken(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
+                                             const paddle::Tensor& cum_offsets, // [bsz, 1]
+                                             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);
+  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>())};
+
+  auto out = paddle::full({bsz, dim}, -2, tmp_out.type(), tmp_out.place()); 
+
+  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>()),
+          encoder_seqs_lods_vp,
+          encoder_batch_map_vp,
+          decoder_batch_map_vp,
+          dim);
+  return {out};
+}
+
+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>& 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.");
+    }
+    int64_t bsz = cum_offsets_shape[0];
+    int64_t dim_embed = tmp_out_shape[1];
+    return {{bsz, dim_embed}};
+}
+
+std::vector<paddle::DataType> GatherNextTokenInferDtype(const paddle::DataType& tmp_out_dtype,
+                                                         const paddle::DataType& cum_offsets_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", "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))
+    .SetInferShapeFn(PD_INFER_SHAPE(GatherNextTokenInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GatherNextTokenInferDtype));

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_padding_offset_v2.cc

@@ -16,10 +16,15 @@
 #include "paddle/extension.h"
 #include "xpu/plugin.h"
 
-std::vector<paddle::Tensor> GetPaddingOffset(const paddle::Tensor& input_ids,
+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::Tensor& seq_len,
+                                             const paddle::optional<paddle::Tensor>& draft_tokens,
+                                             const paddle::optional<paddle::Tensor>& seq_lens_encoder) {
+    if (draft_tokens) {
+        PD_THROW("speculative decoding is not supported in XPU.");
+    }
   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);
@@ -60,35 +65,29 @@ std::vector<paddle::Tensor> GetPaddingOffset(const paddle::Tensor& input_ids,
           cu_seqlens_k};
 }
 
-std::vector<std::vector<int64_t>> GetPaddingOffsetInferShape(
-    const std::vector<int64_t>& input_ids_shape,
-    const std::vector<int64_t>& cum_offsets_shape,
-    const std::vector<int64_t>& token_num_shape,
-    const std::vector<int64_t>& seq_len_shape) {
-  int64_t bsz = seq_len_shape[0];
-  int64_t seq_len = input_ids_shape[1];
-  return {{-1}, {bsz}, {-1}, {bsz + 1}, {bsz + 1}};
+std::vector<std::vector<int64_t>> GetPaddingOffsetV2InferShape(const std::vector<int64_t>& input_ids_shape,
+                                                             const std::vector<int64_t>& cum_offsets_shape,
+                                                             const std::vector<int64_t>& token_num_shape,
+                                                             const std::vector<int64_t>& seq_len_shape,
+                                                             const std::vector<int64_t>& draft_tokens_shape,
+                                                             const std::vector<int64_t>& seq_lens_encoder_shape) {
+    int64_t bsz = seq_len_shape[0];
+    int64_t seq_len = input_ids_shape[1];
+    return {{-1}, {bsz}, {-1}, {bsz + 1}, {bsz + 1}};
 }
 
-std::vector<paddle::DataType> GetPaddingOffsetInferDtype(
-    const paddle::DataType& input_ids_dtype,
-    const paddle::DataType& cum_offsets_dtype,
-    const paddle::DataType& token_num_dtype,
-    const paddle::DataType& seq_len_dtype) {
-  return {input_ids_dtype,
-          seq_len_dtype,
-          seq_len_dtype,
-          seq_len_dtype,
-          seq_len_dtype};
+std::vector<paddle::DataType> GetPaddingOffsetV2InferDtype(const paddle::DataType& input_ids_dtype,
+                                                         const paddle::DataType& cum_offsets_dtype,
+                                                         const paddle::DataType& token_num_dtype,
+                                                         const paddle::DataType& seq_len_dtype,
+                                                         const paddle::DataType& draft_tokens_dtype,
+                                                         const paddle::DataType& seq_lens_encoder_dtype) {
+    return {input_ids_dtype, seq_len_dtype, seq_len_dtype, seq_len_dtype, seq_len_dtype};
 }
 
 PD_BUILD_OP(get_padding_offset_v2)
-    .Inputs({"input_ids", "token_num", "cum_offsets", "seq_len"})
-    .Outputs({"x_remove_padding",
-              "cum_offsets_out",
-              "padding_offset",
-              "cu_seqlens_q",
-              "cu_seqlens_k"})
-    .SetKernelFn(PD_KERNEL(GetPaddingOffset))
-    .SetInferShapeFn(PD_INFER_SHAPE(GetPaddingOffsetInferShape))
-    .SetInferDtypeFn(PD_INFER_DTYPE(GetPaddingOffsetInferDtype));
+    .Inputs({"input_ids", "cum_offsets", "token_num", "seq_len", paddle::Optional("draft_tokens"), paddle::Optional("seq_lens_encoder"),})
+    .Outputs({"x_remove_padding", "cum_offsets_out", "padding_offset", "cu_seqlens_q", "cu_seqlens_k"})
+    .SetKernelFn(PD_KERNEL(GetPaddingOffsetV2))
+    .SetInferShapeFn(PD_INFER_SHAPE(GetPaddingOffsetV2InferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GetPaddingOffsetV2InferDtype));

csrc/xpu/src/get_position_ids.cc

@@ -0,0 +1,53 @@
+// 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;
+
+void GetPositionIdsKernel(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& position_ids) {
+
+  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 bs = seq_lens_encoder.shape()[0];
+
+  int ret = baidu::xpu::api::plugin::get_position_ids(
+        xpu_ctx->x_context(),
+        seq_lens_encoder.data<int32_t>(),
+        seq_lens_decoder.data<int32_t>(),
+        seq_lens_this_time.data<int32_t>(),
+        const_cast<int32_t*>(position_ids.data<int32_t>()),
+        bs
+  );
+  PD_CHECK(ret == 0, "api::plugin::get_position_ids failed");
+}
+
+PD_BUILD_OP(get_position_ids)
+    .Inputs({"seq_lens_encoder", "seq_lens_decoder", "seq_lens_this_time", "position_ids"})
+    .Outputs({"position_ids_out"})
+    .SetInplaceMap({{"position_ids", "position_ids_out"}})
+    .SetKernelFn(PD_KERNEL(GetPositionIdsKernel));

csrc/xpu/src/get_position_ids_v2.cc

@@ -0,0 +1,82 @@
+// 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> GetPositionIdsKernelV2(
+    const paddle::Tensor& seq_lens_encoder,
+    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);
+
+  const int bs = seq_lens_this_time.shape()[0];
+
+  std::vector<int> seq_lens_this_time_cpu(bs, 0);
+  int r = xpu_memcpy(seq_lens_this_time_cpu.data(),
+                 seq_lens_this_time.data<int>(),
+                 sizeof(int32_t) * bs,
+                 XPUMemcpyKind::XPU_DEVICE_TO_HOST);
+  int total_len = 0;
+  for (int i = 0; i < bs; ++i) {
+    total_len += seq_lens_this_time_cpu[i];
+  }
+
+  auto position_ids_out = paddle::full({total_len}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
+
+  int ret = baidu::xpu::api::plugin::get_position_ids(
+        xpu_ctx->x_context(),
+        seq_lens_encoder.data<int32_t>(),
+        seq_lens_decoder.data<int32_t>(),
+        seq_lens_this_time.data<int32_t>(),
+        const_cast<int32_t*>(position_ids_out.data<int32_t>()),
+        bs
+  );
+  PD_CHECK(ret == 0, "api::plugin::get_position_ids failed"); 
+  
+  baidu::xpu::api::plugin::print_times("[TIME END] GetPositionIdsKernelV2" );
+  return {position_ids_out};
+}
+
+std::vector<std::vector<int64_t>> GetPositionIdsV2InferShape(const std::vector<int64_t>& seq_lens_encoder,
+                                                             const std::vector<int64_t>& seq_lens_decoder,
+                                                             const std::vector<int64_t>& seq_lens_this_time) {
+    return {{-1}};
+}
+
+std::vector<paddle::DataType> GetPositionIdsV2InferDtype(const paddle::DataType& seq_lens_encoder_dtype,
+                                                         const paddle::DataType& seq_lens_decoder_dtype,
+                                                         const paddle::DataType& seq_lens_this_time_dtype) {
+    return {seq_lens_encoder_dtype};
+}
+
+PD_BUILD_OP(get_position_ids_v2)
+    .Inputs({"seq_lens_encoder", "seq_lens_decoder", "seq_lens_this_time"})
+    .Outputs({"position_ids_out"})
+    .SetKernelFn(PD_KERNEL(GetPositionIdsKernelV2))
+    .SetInferShapeFn(PD_INFER_SHAPE(GetPositionIdsV2InferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GetPositionIdsV2InferDtype));;

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

@@ -0,0 +1,604 @@
+// 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
+
+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> MlaAttn(
+    const paddle::Tensor& q,
+    const paddle::Tensor& k,
+    const paddle::Tensor& v,
+    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 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) {
+  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 auto& key_cache_dims = key_cache.dims();
+  const auto& value_cache_dims = value_cache.dims();
+  const int bsz = seq_lens_encoder.dims()[0];
+  const int token_num = input_dims[0];
+  const int num_head = key_cache_dims[1];
+  const int dim_qk = key_cache_dims[3];
+  const int dim_v = value_cache_dims[3];
+  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 block_size = key_cache_dims[2];  
+  const int max_seq_len = block_size * max_block_per_seq;
+  // 初始化输入：q k v
+  auto q_xft = baidu::xpu::xft::xftTensor<QType, 2>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(q.data<qdata_t>())),
+      std::array<int64_t, 2>{q.shape()[0],
+                             q.shape()[1]});
+  auto k_xft = baidu::xpu::xft::xftTensor<QType, 2>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(k.data<qdata_t>())),
+      std::array<int64_t, 2>{k.shape()[0],
+                             k.shape()[1]});
+  auto v_xft = baidu::xpu::xft::xftTensor<QType, 2>(
+      reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(v.data<qdata_t>())),
+      std::array<int64_t, 2>{v.shape()[0],
+                             v.shape()[1]});                             
+  // 初始化输入：k cache
+  auto key_cache_xft = baidu::xpu::xft::xftTensor<CacheType, 4>(
+  reinterpret_cast<CacheType*>(const_cast<paddle::bfloat16*>(key_cache.data<cache_t>())),
+  std::array<int64_t, 4>{key_cache.shape()[0],
+                          key_cache.shape()[1],
+                          key_cache.shape()[2],
+                          key_cache.shape()[3]});   
+  // 初始化输入：v cache                                    
+  auto value_cache_xft = baidu::xpu::xft::xftTensor<CacheType, 4>(
+  reinterpret_cast<CacheType*>(const_cast<paddle::bfloat16*>(value_cache.data<cache_t>())),
+  std::array<int64_t, 4>{value_cache.shape()[0],
+                          value_cache.shape()[1],
+                          value_cache.shape()[2],
+                          value_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 * 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 判断逻辑
+  std::vector<int> seq_lens_encoder_cpu(bsz, 0);
+  std::vector<int> seq_lens_decoder_cpu(bsz, 0);
+  std::vector<int> encoder_batch_map; // 
+  std::vector<int> decoder_batch_map; // 
+  std::vector<int> encoder_batch_idx; // 去除空隙的batch map
+  std::vector<int> decoder_batch_idx; // 去除空隙的batch map
+  std::vector<int> encoder_seq_lod;
+  std::vector<int> decoder_context_len;
+  std::vector<int> decoder_context_len_cache;
+  xpu_wait(xpu_ctx->x_context()->xpu_stream); // 是否需要！！！！TODO
+  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;
+  int64_t total_enc_len = 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++;
+      total_enc_len += seq_lens_encoder_cpu[i];
+      encoder_batch_map.push_back(i);
+      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_map.push_back(i);
+      decoder_batch_idx.push_back(i - batch_offset);
+      decoder_context_len.push_back(seq_lens_decoder_cpu[i] + 1);
+      decoder_context_len_cache.push_back(seq_lens_decoder_cpu[i]);
+    }
+    else{
+        batch_offset++;
+    }
+  }
+
+  // encoder
+  if(max_enc_len_this_time.data<int>()[0] > 0){
+    // q_lod
+    baidu::xpu::api::VectorParam<int32_t> context_len_vp =
+        baidu::xpu::api::VectorParam<int32_t>{encoder_seq_lod.data(), enc_batch + 1, nullptr}
+            .to_xpu(RAII_GUARD);
+    // real batch（encoder阶段 kv cache写需要）
+    baidu::xpu::api::VectorParam<int32_t> valid_batch_vp =
+        baidu::xpu::api::VectorParam<int32_t>{encoder_batch_map.data(), enc_batch, nullptr}
+            .to_xpu(RAII_GUARD);
+    // 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 =
+        baidu::xpu::api::VectorParam<int32_t>{encoder_seq_lod.data(), enc_batch + 1, nullptr}
+            .to_xpu(RAII_GUARD);  
+    // 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
+    
+    // k v cache write
+    std::vector<int> start_tokens(enc_batch, 0);
+    baidu::xpu::api::VectorParam<int32_t> start_tokens_vp =
+        baidu::xpu::api::VectorParam<int32_t>{start_tokens.data(), enc_batch, nullptr}.to_xpu(RAII_GUARD);
+    int ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
+            xpu_ctx->x_context(),
+            reinterpret_cast<float16*>(k_xft.data()),
+            reinterpret_cast<float16*>(key_cache_xft.data()),
+            block_tables_xft.data(),
+            context_len_vp,
+            start_tokens_vp,
+            valid_batch_vp,
+            enc_batch,
+            num_head,
+            dim_qk,
+            block_batch,
+            block_size,
+            max_block_per_seq,
+            "BLHD", // qkv_layout,
+            "HLD", // page_layout,
+            nullptr, // scale,
+            nullptr, // batch_max_ptrs, 
+            nullptr); // max_ptrs
+    ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
+            xpu_ctx->x_context(),
+            reinterpret_cast<float16*>(v_xft.data()),
+            reinterpret_cast<float16*>(value_cache_xft.data()),
+            block_tables_xft.data(),
+            context_len_vp,
+            start_tokens_vp,
+            valid_batch_vp,
+            enc_batch,
+            num_head,
+            dim_v,
+            block_batch,
+            block_size,
+            max_block_per_seq,
+            "BLHD", // qkv_layout,
+            "HLD", // page_layout,
+            nullptr, // scale,
+            nullptr, // batch_max_ptrs, 
+            nullptr); // max_ptrs
+    // fmha op
+    using FMHA_Type = typename baidu::xpu::xft::FMHA_QBF16_KVBF16;
+    // auto fmha_op = baidu::xpu::xft::FMHAOperation<FMHA_Type>(
+    //         enc_batch,
+    //         attn_param_.head_dim_,
+    //         attn_param_.q_head_num_,
+    //         attn_param_.kv_head_num_,
+    //         page_param_.max_context_len_,  
+    //         page_param_.max_context_len_,
+    //         attn_param_.context_len_vp_,
+    //         attn_param_.encoder_kv_lods_vp_);
+
+    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;
+  }
+
+  // decoder
+  if(max_dec_len_this_time.data<int>()[0] > 0){
+    // context_len
+    baidu::xpu::api::VectorParam<int32_t> context_len_vp =
+        baidu::xpu::api::VectorParam<int32_t>{decoder_context_len.data(), dec_batch, nullptr}
+            .to_xpu(RAII_GUARD);
+    baidu::xpu::api::VectorParam<int32_t> context_len_vp_cache =
+        baidu::xpu::api::VectorParam<int32_t>{decoder_context_len_cache.data(), dec_batch, nullptr}
+            .to_xpu(RAII_GUARD);
+    // real batch     
+    baidu::xpu::api::VectorParam<int32_t> valid_batch_vp =
+        baidu::xpu::api::VectorParam<int32_t>{decoder_batch_map.data(), dec_batch, nullptr}
+            .to_xpu(RAII_GUARD);
+    // prefix (not support)
+    baidu::xpu::api::VectorParam<int32_t> prefix_lens_vp = baidu::xpu::api::VectorParam<int32_t>();
+    // kv_lod (decoder 不需要)
+    baidu::xpu::api::VectorParam<int32_t> encoder_kv_lods_vp = baidu::xpu::api::VectorParam<int32_t>();
+    // page_param_
+    baidu::xpu::xft::PageParam page_param_(block_batch, // block_batch_
+                                            max_seq_len, // max_dec_len_this_time.data<int>()[0],max_seq_len // max_context_len_ // max_dec_len_this_time.data<int>()[0]
+                                            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
+    // k v cache write
+    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);
+    int ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
+            xpu_ctx->x_context(),
+            reinterpret_cast<float16*>(k_xft.data() + total_enc_len * dim_qk * num_head),
+            reinterpret_cast<float16*>(key_cache_xft.data()),
+            block_tables_xft.data(),
+            kv_seq_lod_vp,
+            context_len_vp_cache,
+            valid_batch_vp,
+            dec_batch,
+            num_head,
+            dim_qk,
+            block_batch,
+            block_size,
+            max_block_per_seq,
+            "BLHD", // qkv_layout,
+            "HLD", // page_layout,
+            nullptr, // scale,
+            nullptr, // batch_max_ptrs, 
+            nullptr); // max_ptrs
+    ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
+            xpu_ctx->x_context(),
+            reinterpret_cast<float16*>(v_xft.data() + total_enc_len * dim_v * num_head),
+            reinterpret_cast<float16*>(value_cache_xft.data()),
+            block_tables_xft.data(),
+            kv_seq_lod_vp,
+            context_len_vp_cache,
+            valid_batch_vp,
+            dec_batch,
+            num_head,
+            dim_v,
+            block_batch,
+            block_size,
+            max_block_per_seq,
+            "BLHD", // qkv_layout,
+            "HLD", // page_layout,
+            nullptr, // scale,
+            nullptr, // batch_max_ptrs, 
+            nullptr); // max_ptrs
+    // paged attention op
+    using PA_Type = typename baidu::xpu::xft::PA_QBF16_KVBF16;
+    using TQ = typename PA_Type::Qtype; // bfloat16
+    using TKV = typename PA_Type::KVtype; // bfloat16
+    using TO = typename PA_Type::Otype; // bfloat16
+    using TID = typename PA_Type::TIDtype; // int
+    using TL = float;
+    using TGEMM0 = float;
+    using TGEMM1 = float;
+    using TEW = float;
+
+    auto pa_func = &baidu::xpu::xfa::paged_attention_xft<TQ, TKV, TO, TL, TGEMM0, TGEMM1, TEW, TID, false>;
+
+    int ret = pa_func(
+            xpu_ctx->x_context(),
+            fmha_out_xft.data() + total_enc_len * dim_v * num_head,
+            q_xft.data() + total_enc_len * dim_qk * num_head,
+            nullptr,   /* k_cur */
+            nullptr, /* v_cur */
+            const_cast<TKV*>(key_cache_xft.data()),
+            const_cast<TKV*>(value_cache_xft.data()),
+            attn_param_.kv_head_num_,
+            softmax_scale,
+            block_tables_xft.data(),
+            attn_param_.context_len_vp_,
+            attn_param_.valid_batch_vp_,
+            page_param_.block_size_,
+            page_param_.max_context_len_,
+            nullptr, // (TBIAS*)alibi_slopes,
+            block_batch, // page_param_.block_batch_,
+            attn_param_.q_head_num_,
+            attn_param_.head_dim_,
+            page_param_.max_num_blocks_per_seq_,
+            nullptr, // shift,
+            nullptr, // smooth,
+            nullptr, // query_maxptr,
+            nullptr, // key_cache_maxptr,
+            nullptr, // value_cache_maxptr,
+            nullptr, // p_k_scales_inv,
+            nullptr, // p_k_zeros,
+            nullptr, // p_v_scales_inv,
+            nullptr, // p_v_zeros,
+            nullptr, // out_maxptr,
+            dim_v); // v_head_dim
+  }
+    return {fmha_out};   
+}
+
+std::vector<std::vector<int64_t>> MlaAttnInferShape(
+    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>& key_cache_shape,
+    const std::vector<int64_t>& value_cache_shape,
+    const std::vector<int64_t>& seq_lens_encoder_shape,
+    const std::vector<int64_t>& seq_lens_decoder_shape,
+    const std::vector<int64_t>& seq_lens_this_time_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>& encoder_batch_ids_shape,
+    const std::vector<int64_t>& encoder_tile_ids_per_batch_shape,
+    const std::vector<int64_t>& encoder_num_blocks_shape,
+    const std::vector<int64_t>& kv_batch_ids_shape,
+    const std::vector<int64_t>& kv_tile_ids_per_batch_shape,
+    const std::vector<int64_t>& kv_num_blocks_shape,
+    const std::vector<int64_t>& decoder_batch_ids_shape,
+    const std::vector<int64_t>& decoder_tile_ids_per_batch_shape,
+    const std::vector<int64_t>& decoder_num_blocks_shape,
+    const std::vector<int64_t>& max_enc_len_this_time_shape,
+    const std::vector<int64_t>& max_dec_len_this_time_shape,
+    const std::vector<int64_t>& max_len_kv_shape,
+    const paddle::optional<std::vector<int64_t>>& rotary_embs_shape,
+    const paddle::optional<std::vector<int64_t>>& attn_mask_shape,
+    const paddle::optional<std::vector<int64_t>>& qkv_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& qkv_out_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_k_quant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_v_quant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_k_dequant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_v_dequant_scales_shape,
+    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 int token_num = q_shape[0];
+  const int all_v_dim = value_cache_shape[3];
+  return {{token_num, all_v_dim}};
+}
+
+std::vector<paddle::DataType> MlaAttnInferDtype(
+    const paddle::DataType& q_dtype,
+    const paddle::DataType& k_dtype,
+    const paddle::DataType& v_dtype,
+    const paddle::DataType& key_cache_dtype,
+    const paddle::DataType& value_cache_dtype,
+    const paddle::DataType& seq_lens_encoder_dtype,
+    const paddle::DataType& seq_lens_decoder_dtype,
+    const paddle::DataType& seq_lens_this_time_dtype,
+    const paddle::DataType& padding_offsets_dtype,
+    const paddle::DataType& cum_offsets_dtype,
+    const paddle::DataType& block_tables_dtype,
+    const paddle::DataType& encoder_batch_ids_dtype,
+    const paddle::DataType& encoder_tile_ids_per_batch_dtype,
+    const paddle::DataType& encoder_num_blocks_dtype,
+    const paddle::DataType& kv_batch_ids_dtype,
+    const paddle::DataType& kv_tile_ids_per_batch_dtype,
+    const paddle::DataType& kv_num_blocks_dtype,
+    const paddle::DataType& decoder_batch_ids_dtype,
+    const paddle::DataType& decoder_tile_ids_per_batch_dtype,
+    const paddle::DataType& decoder_num_blocks_dtype,
+    const paddle::DataType& max_enc_len_this_time_dtype,
+    const paddle::DataType& max_dec_len_this_time_dtype,
+    const paddle::DataType& max_len_kv_dtype,
+    const paddle::optional<paddle::DataType>& rotary_embs_dtype,
+    const paddle::optional<paddle::DataType>& attn_mask_dtype,
+    const paddle::optional<paddle::DataType>& qkv_bias_dtype,
+    const paddle::optional<paddle::DataType>& qkv_out_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_k_quant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_v_quant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_k_dequant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_v_dequant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_k_zp_dtype,
+    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 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) {
+    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(mla_block_multihead_attention_xpu)
+    .Inputs({"q",
+             "k",
+             "v",
+             "key_cache",
+             "value_cache",
+             "seq_lens_encoder",
+             "seq_lens_decoder",
+             "seq_lens_this_time",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables",
+             "encoder_batch_ids",
+             "encoder_tile_ids_per_batch",
+             "encoder_num_blocks",
+             "kv_batch_ids",
+             "kv_tile_ids_per_batch",
+             "kv_num_blocks",
+             "decoder_batch_ids",
+             "decoder_tile_ids_per_batch",
+             "decoder_num_blocks",
+             "max_enc_len_this_time",
+             "max_dec_len_this_time",
+             "max_len_kv",
+             paddle::Optional("rotary_embs"),
+             paddle::Optional("attn_mask"),
+             paddle::Optional("qkv_bias"),
+             paddle::Optional("qkv_out_scales"),
+             paddle::Optional("cache_k_quant_scales"),
+             paddle::Optional("cache_v_quant_scales"),
+             paddle::Optional("cache_k_dequant_scales"),
+             paddle::Optional("cache_v_dequant_scales"),
+             paddle::Optional("cache_k_zp"),
+             paddle::Optional("cache_v_zp"),
+             paddle::Optional("out_linear_shifts"),
+             paddle::Optional("out_linear_smooths")})
+    .Outputs({"fmha_out"})
+    // .SetInplaceMap({{"key_cache", "key_cache_out"},
+    //                 {"value_cache", "value_cache_out"}})
+    .Attrs({"cache_quant_type: std::string",
+            "use_neox_rotary_style: bool",
+            "max_input_length: int",
+            "softmax_scale: float",
+            "quant_max_bound: float",
+            "quant_min_bound: float",
+            "out_linear_in_scale: float",
+            "speculate_max_draft_token_num: int",
+            "causal: bool",
+            "speculate_decoder: bool"})
+    .SetKernelFn(PD_KERNEL(MlaAttn))
+    .SetInferShapeFn(PD_INFER_SHAPE(MlaAttnInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MlaAttnInferDtype));
+

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_dispatch.cc

@@ -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.
+
+#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>
+namespace xftkernel = baidu::xpu::xftkernel;
+std::vector<paddle::Tensor> MoeDispatchKernel(
+        const paddle::Tensor& X,
+        const paddle::Tensor& gating_output,
+        const int moe_topk,
+        const bool group_moe,
+        const bool topk_only_mode) {
+    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);
+    int token_rows = 0;
+    auto input_dims = X.shape();
+    if (input_dims.size() == 3) {
+        token_rows = input_dims[0] * input_dims[1];
+    } else {
+        token_rows = input_dims[0];
+    }
+    const int num_rows = token_rows;
+    const int hidden_size = X.shape()[input_dims.size() - 1];
+    auto gating_dims = gating_output.shape();
+    const int expert_num = gating_dims[gating_dims.size() - 1];
+    if (group_moe) {
+        // Check if expert_num is divisible by moe_topk, else throw an error
+        PADDLE_ENFORCE_EQ(
+                expert_num % moe_topk,
+                0,
+                common::errors::InvalidArgument(
+                        "The number of experts (expert_num) "
+                        "must be divisible by moe_topk. "
+                        "Got expert_num = %d and moe_topk = %d.",
+                        expert_num,
+                        moe_topk));
+    }
+
+    auto permute_input = paddle::full({num_rows * moe_topk, hidden_size}, -1, X.type(), X.place());
+    auto token_nums_per_expert = paddle::full({expert_num}, 0, paddle::DataType::INT32, X.place());
+    auto permute_indices_per_token = paddle::full({moe_topk, num_rows}, -1, paddle::DataType::INT32, X.place());
+    auto expert_scales_float = paddle::full({num_rows, moe_topk}, 0.0f, paddle::DataType::FLOAT32, X.place());
+    auto top_k_indices = paddle::full({num_rows, moe_topk}, -1, paddle::DataType::INT32, X.place());
+
+    using XPUType = typename XPUTypeTrait<bfloat16>::Type;
+    typedef paddle::bfloat16 data_t;
+    auto index_data = permute_indices_per_token.data<int32_t>();
+//     auto index = paddle::full({num_rows, moe_topk}, -1, paddle::DataType::INT32, X.place());
+
+    int ret = xftkernel::xft_moe_group_topk_fusion<float, float, int32_t>(
+            xpu_ctx->x_context(),
+            gating_output.data<float>(),
+            expert_scales_float.data<float>(),
+            index_data,
+            nullptr,
+            num_rows,
+            expert_num,
+            0,
+            0,
+            moe_topk,  // num of shared expert
+            0);
+    PD_CHECK(ret == 0, "xftkernel::xft_moe_softmax_topk_norm_fusion failed");
+    auto sorted_tokens_num_lod = paddle::full({expert_num + 1}, -1, paddle::DataType::INT32, X.place());
+     ret = xftkernel::xft_moe_ffn_pre_sorted<XPUType, int32_t>(
+            xpu_ctx->x_context(),
+            reinterpret_cast<const XPUType*>(X.data<data_t>()),
+            index_data,
+            nullptr,
+            reinterpret_cast<XPUType*>(permute_input.data<data_t>()),
+            top_k_indices.data<int32_t>(),
+            token_nums_per_expert.data<int32_t>(),
+            sorted_tokens_num_lod.data<int32_t>(),
+            token_rows,
+            hidden_size,
+            expert_num,
+            moe_topk,
+            0);
+    PD_CHECK(ret == 0, "xftkernel::xft_moe_ffn_pre_sorted failed");
+    return {permute_input,
+            sorted_tokens_num_lod,
+            permute_indices_per_token,
+            expert_scales_float,
+            top_k_indices};
+}
+
+PD_BUILD_OP(moe_dispatch_xpu)
+        .Inputs({"X", "gating_output"})
+        .Outputs(
+                {"permute_input",
+                 "token_nums_per_expert",
+                 "permute_indices_per_token",
+                 "expert_scales_float",
+                 "expert_scales_int"})
+        .Attrs({"moe_topk: int", "group_moe: bool", "topk_only_mode: bool"})
+        .SetKernelFn(PD_KERNEL(MoeDispatchKernel));

csrc/xpu/src/moe_ffn.cc

@@ -0,0 +1,277 @@
+// 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"
+
+namespace xftkernel = baidu::xpu::xftkernel;
+
+std::vector<paddle::Tensor> MoeFFNKernel(
+    const paddle::Tensor& permute_input,
+    const paddle::Tensor& token_nums_per_expert,
+    const paddle::Tensor& ffn1_weight,
+    const paddle::Tensor& ffn2_weight,
+    const paddle::optional<paddle::Tensor>& ffn1_bias,
+    const paddle::optional<paddle::Tensor>& ffn1_scale,
+    const paddle::optional<paddle::Tensor>& ffn2_scale,
+    const std::string& quant_method) {
+  phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
+  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
+  using XPUType = typename XPUTypeTrait<bfloat16>::Type;
+  typedef paddle::bfloat16 data_t;
+  auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  auto ffn_out_paddle = permute_input.copy_to(permute_input.place(), false);
+
+
+  const int64_t expanded_active_expert_rows = permute_input.shape()[0];
+  const int num_experts = ffn1_weight.shape()[0];
+  const int hidden_size = ffn1_weight.shape()[2];
+  int64_t inter_size = ffn1_weight.shape()[1];
+  if (quant_method == "weight_only_int4") {
+    inter_size = inter_size * 2;
+  }
+
+
+  auto ffn1_input = baidu::xpu::xft::xftTensor<XPUType, 2>(
+      reinterpret_cast<XPUType*>(
+          const_cast<paddle::bfloat16*>(permute_input.data<data_t>())),
+      std::array<int64_t, 2>{permute_input.shape()[0],
+                             permute_input.shape()[1]});
+
+  auto ffn1_bias_tenosor = baidu::xpu::xft::xftVec<float>(
+      ffn1_bias ? const_cast<float*>(ffn1_bias->data<float>()) : nullptr,
+      std::array<int64_t, 1>{inter_size});
+
+  auto ffn1_out_paddle = paddle::full({expanded_active_expert_rows, inter_size},
+                                      -1,
+                                      permute_input.type(),
+                                      permute_input.place());
+  auto ffn1_out_tensor = baidu::xpu::xft::xftTensor<XPUType, 2>(
+      reinterpret_cast<XPUType*>(ffn1_out_paddle.data<data_t>()),
+      std::array<int64_t, 2>{ffn1_out_paddle.shape()[0],
+                             ffn1_out_paddle.shape()[1]});
+
+  auto token_nums_per_expert_tensor = baidu::xpu::xft::xftVec<int32_t>(
+      const_cast<int32_t*>(token_nums_per_expert.data<int>()),
+      std::array<int64_t, 1>{num_experts});
+
+  auto empty_tensor = baidu::xpu::xft::xftTensor<int32_t, 2>();
+
+  auto sorted_tokens_idx_paddle = paddle::full({expanded_active_expert_rows, 1},
+                                               -1,
+                                               paddle::DataType::INT32,
+                                               permute_input.place());
+
+  auto sorted_tokens_idx_tensor = baidu::xpu::xft::xftTensor<int32_t, 2>(
+      sorted_tokens_idx_paddle.data<int>(),
+      std::array<int64_t, 2>{sorted_tokens_idx_paddle.shape()[0],
+                             sorted_tokens_idx_paddle.shape()[1]});
+  int ret = 0;
+  if (quant_method == "weight_only_int4") {
+    auto ffn1_w_tensor = baidu::xpu::xft::xftMat<int4_t>(
+        reinterpret_cast<int4_t*>(
+            const_cast<int8_t*>(ffn1_weight.data<int8_t>())),
+        nullptr,
+        const_cast<float*>(
+            ffn1_scale->data<
+                float>()),  // 这个地方GPU是float16，我们需要的是float，需要check
+        std::array<int64_t, 2>{num_experts * inter_size, hidden_size});
+
+    ret = baidu::xpu::xft::xft_moe_sort_fc_block<XPUType,
+                                                 int4_t,
+                                                 XPUType,
+                                                 float,
+                                                 int32_t,
+                                                 int4_wo_int15>(
+        xpu_ctx->x_context(),
+        ffn1_input,
+        ffn1_w_tensor,
+        ffn1_out_tensor,
+        &ffn1_bias_tenosor,
+        token_nums_per_expert_tensor,
+        sorted_tokens_idx_tensor,
+        nullptr,
+        empty_tensor,
+        baidu::xpu::api::Activation_t::LINEAR,
+        false,
+        true,
+        1.0,
+        0.0,
+        num_experts,
+        1,  // topk
+        0,
+        2,
+        1);
+  } else {
+    auto ffn1_w_tensor = baidu::xpu::xft::xftMat<int8_t>(
+        const_cast<int8_t*>(ffn1_weight.data<int8_t>()),
+        nullptr,
+        const_cast<float*>(
+            ffn1_scale->data<
+                float>()),  // 这个地方GPU是float16，我们需要的是float，需要check
+        std::array<int64_t, 2>{num_experts * inter_size, hidden_size});
+    ret = baidu::xpu::xft::xft_moe_sort_fc_block<XPUType,
+                                                 int8_t,
+                                                 XPUType,
+                                                 float,
+                                                 int32_t,
+                                                 float>(
+        xpu_ctx->x_context(),
+        ffn1_input,
+        ffn1_w_tensor,
+        ffn1_out_tensor,
+        &ffn1_bias_tenosor,
+        token_nums_per_expert_tensor,
+        sorted_tokens_idx_tensor,
+        nullptr,
+        empty_tensor,
+        baidu::xpu::api::Activation_t::LINEAR,
+        false,
+        true,
+        1.0,
+        0.0,
+        num_experts,
+        1,  // topk
+        0,
+        2,
+        1);
+  }
+
+  auto ffn2_input_paddle =
+      paddle::full({expanded_active_expert_rows, inter_size / 2},
+                   -1,
+                   permute_input.type(),
+                   permute_input.place());
+
+
+  auto ffn2_input_tensor = baidu::xpu::xft::xftTensor<XPUType, 2>(
+      reinterpret_cast<XPUType*>(ffn2_input_paddle.data<data_t>()),
+      std::array<int64_t, 2>{ffn2_input_paddle.shape()[0],
+                             ffn2_input_paddle.shape()[1]});
+  ret = xftkernel::xft_fast_swiglu_add_mul_fusion<XPUType>(
+      xpu_ctx->x_context(),
+      ffn1_out_tensor.data(),
+      ffn2_input_tensor.data(),
+      expanded_active_expert_rows,
+      inter_size,
+      nullptr,
+      nullptr,
+      true,
+      nullptr,
+      nullptr);
+
+
+  //   auto ffn2_bias_tenosor = baidu::xpu::xft::xftVec<float>(
+  //         ffn1_bias ? const_cast<float*>(ffn2_bias->data<float>()):
+  //         nullptr, std::array<int64_t, 1>{hidden_size});
+  auto ffn2_out_tensor = baidu::xpu::xft::xftTensor<XPUType, 2>(
+      reinterpret_cast<XPUType*>(ffn_out_paddle.data<data_t>()),
+      std::array<int64_t, 2>{ffn_out_paddle.shape()[0],
+                             ffn_out_paddle.shape()[1]});
+  if (quant_method == "weight_only_int4") {
+    auto ffn2_w_tensor = baidu::xpu::xft::xftMat<int4_t>(
+        reinterpret_cast<int4_t*>(
+            const_cast<int8_t*>(ffn2_weight.data<int8_t>())),
+        nullptr,
+        const_cast<float*>(
+            ffn2_scale->data<
+                float>()),  // 这个地方GPU是float16，我们需要的是float，需要check
+        std::array<int64_t, 2>{num_experts * hidden_size, inter_size / 2});
+    ret = baidu::xpu::xft::xft_moe_sort_fc_block<XPUType,
+                                                 int4_t,
+                                                 XPUType,
+                                                 float,
+                                                 int32_t,
+                                                 int4_wo_int15>(
+        xpu_ctx->x_context(),
+        ffn2_input_tensor,
+        ffn2_w_tensor,
+        ffn2_out_tensor,
+        nullptr,
+        token_nums_per_expert_tensor,
+        sorted_tokens_idx_tensor,
+        nullptr,
+        empty_tensor,
+        baidu::xpu::api::Activation_t::LINEAR,
+        false,
+        true,
+        1.0,
+        0.0,
+        num_experts,
+        1,  // topk
+        0,
+        2,
+        1);
+  } else {
+    auto ffn2_w_tensor = baidu::xpu::xft::xftMat<int8_t>(
+        const_cast<int8_t*>(ffn2_weight.data<int8_t>()),
+        nullptr,
+        const_cast<float*>(
+            ffn2_scale->data<
+                float>()),  // 这个地方GPU是float16，我们需要的是float，需要check
+        std::array<int64_t, 2>{num_experts * hidden_size, inter_size / 2});
+    ret = baidu::xpu::xft::xft_moe_sort_fc_block<XPUType,
+                                                 int8_t,
+                                                 XPUType,
+                                                 float,
+                                                 int32_t,
+                                                 float>(
+        xpu_ctx->x_context(),
+        ffn2_input_tensor,
+        ffn2_w_tensor,
+        ffn2_out_tensor,
+        nullptr,
+        token_nums_per_expert_tensor,
+        sorted_tokens_idx_tensor,
+        nullptr,
+        empty_tensor,
+        baidu::xpu::api::Activation_t::LINEAR,
+        false,
+        true,
+        1.0,
+        0.0,
+        num_experts,
+        1,  // topk
+        0,
+        2,
+        1);
+  }
+
+
+  return {
+      ffn_out_paddle,
+  };
+}
+
+
+PD_BUILD_OP(moe_ffn_xpu)
+    .Inputs({"permute_input",
+             "token_nums_per_expert",
+             "ffn1_weight",
+             "ffn2_weight",
+             "ffn1_bias",
+             "ffn1_scale",
+             "ffn2_scale"})
+    .Outputs({"ffn_out"})
+    .Attrs({"quant_method: std::string"})
+    .SetKernelFn(PD_KERNEL(MoeFFNKernel));

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));