Update kernel.py

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

@@ -1,41 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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

@@ -1,77 +0,0 @@
-# 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

@@ -1,47 +0,0 @@
-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

@@ -1,47 +0,0 @@
-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

@@ -1,52 +0,0 @@
-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,80 +1,17 @@
-name: Codestyle Check
+name: Lint
 
 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
-    needs: check-bypass
-    if: ${{ needs.check-bypass.outputs.can-skip != 'true' }}
-    runs-on: [self-hosted, ernie-cpu]
+    runs-on: ubuntu-20.04
     steps:
-      - name: Run Container
-        env:
-          work_dir: ${{ github.workspace }}
-          python_version: "3.10"
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: checkout develop
         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
@@ -82,32 +19,14 @@ jobs:
           else
             echo "local develop branch exist, skipping"
           fi
-          '
-          
-      - name: Setup Environment
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip' # caching pip dependencies
+      - name: Install dependencies
         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: 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
+      - name: run the command
+        run: make lint
+    

.github/workflows/llm.yml

@@ -1,207 +0,0 @@
-# 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-24.04
+    runs-on: ubuntu-20.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-24.04
+    runs-on: ubuntu-20.04
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v1

.github/workflows/rerun.yml

@@ -1,73 +0,0 @@
-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-24.04
+    runs-on: ubuntu-20.04
 
     permissions:
       issues: write

.github/workflows/tests.yml

@@ -0,0 +1,29 @@
+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

@@ -1,120 +0,0 @@
-# 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

@@ -1,204 +0,0 @@
-# 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

@@ -134,10 +134,4 @@ output/
 autogen/
 
 # cutlass kernel
-!csrc/gpu/cutlass_kernels/gemm/collective/builders
-
-
-#fp8
-ops/csrc/fp8/deep_gemm/include/cutlass
-ops/csrc/fp8/deep_gemm/include/cute
-.ccls-cache
+!csrc/gpu/cutlass_kernels/gemm/collective/builders

.pre-commit-config.yaml

@@ -62,12 +62,3 @@ 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/zh/conf.py
+   configuration: docs/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 "committing my cool feature"
+    ➜  (my-virtual-env) git commit -m "commiting 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 "committing my cool feature"
+    ➜  (my-virtual-env) git commit -m "commiting my cool feature"
     black....................................................................Passed
     isort....................................................................Failed
     - hook id: isort

README.md

@@ -33,26 +33,16 @@
 
 ## News 📢
 
-* **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.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
 
 * **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
-
 * **2024.12.16 [PaddleNLP v3.0 Beta3](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta3)**：大模型功能全新升级，新增了 Llama-3.2、DeepSeekV2模型，升级了 TokenizerFast，快速分词，重构了 SFTTrainer，一键开启 SFT 训练。此外，PaddleNLP 还支持了优化器状态的卸载和重载功能，实现了精细化的重新计算，训练性能提升7%。在 Unified Checkpoint 方面，进一步优化了异步保存逻辑，新增 Checkpoint 压缩功能，可节省78.5%存储空间。
 最后，在大模型推理方面，升级 Append Attention，支持了 FP8量化，支持投机解码。
 
+<details><summary> <b>点击展开</b> </summary><div>
+
 * **2024.12.13 📚《飞桨大模型套件 Unified Checkpoint 技术》**，加速模型存储95%，节省空间78%。支持全分布式策略调整自适应转换，提升模型训练的灵活性与可扩展性。训练-压缩-推理统一存储协议，无需手动转换提升全流程体验。Checkpoint 无损压缩结合异步保存，实现秒级存储并降低模型存储成本。适用于智能制造、指挥交通、医疗健康、金融服务等产业实际场景。12月24日（周二）19：00直播为您详细解读该技术如何优化大模型训练流程。报名链接：https://www.wjx.top/vm/huZkHn9.aspx?udsid=787976
 
 * **2024.11.28 📚《FlashRAG-Paddle | 基于 PaddleNLP 的高效开发与评测 RAG 框架》**，为文本更快更好构建准确嵌入表示、加速推理生成速度。PaddleNLP 支持超大 Batch 嵌入表示学习与多硬件高性能推理，涵盖 INT8/INT4量化技术及多种高效注意力机制优化与 TensorCore 深度优化。内置全环节算子融合技术，使得 FlashRAG 推理性能相比 transformers 动态图提升70%以上，结合检索增强知识输出结果更加准确，带来敏捷高效的使用体验。直播时间：12月3日（周二）19：00。报名链接：https://www.wjx.top/vm/eaBa1vA.aspx?udsid=682361
@@ -69,7 +59,7 @@
 ### <a href=#多硬件训推一体> 🔧 多硬件训推一体 </a>
 
 支持英伟达 GPU、昆仑 XPU、昇腾 NPU、燧原 GCU 和海光 DCU 等多个硬件的大模型和自然语言理解模型训练和推理，套件接口支持硬件快速切换，大幅降低硬件切换研发成本。
-当前支持的自然语言理解模型：[多硬件自然语言理解模型列表](./docs/zh/model_zoo/model_list_multy_device.md)
+当前支持的自然语言理解模型：[多硬件自然语言理解模型列表](./docs/model_zoo/model_list_multy_device.md)
 
 ### <a href=#高效易用的预训练> 🚀 高效易用的预训练 </a>
 
@@ -119,15 +109,14 @@
 |           [Qwen1.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen1.5-0.5B, Qwen/Qwen1.5-0.5B-Chat, Qwen/Qwen1.5-1.8B, Qwen/Qwen1.5-1.8B-Chat, Qwen/Qwen1.5-4B, Qwen/Qwen1.5-4B-Chat, Qwen/Qwen1.5-7B, Qwen/Qwen1.5-7B-Chat, Qwen/Qwen1.5-14B, Qwen/Qwen1.5-14B-Chat, Qwen/Qwen1.5-32B, Qwen/Qwen1.5-32B-Chat, Qwen/Qwen1.5-72B, Qwen/Qwen1.5-72B-Chat, Qwen/Qwen1.5-110B, Qwen/Qwen1.5-110B-Chat, Qwen/Qwen1.5-MoE-A2.7B, Qwen/Qwen1.5-MoE-A2.7B-Chat |
 |            [Qwen2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)             | Qwen/Qwen2-0.5B, Qwen/Qwen2-0.5B-Instruct, Qwen/Qwen2-1.5B, Qwen/Qwen2-1.5B-Instruct, Qwen/Qwen2-7B, Qwen/Qwen2-7B-Instruct, Qwen/Qwen2-72B, Qwen/Qwen2-72B-Instruct, Qwen/Qwen2-57B-A14B, Qwen/Qwen2-57B-A14B-Instruct                                                                                                                                                                       |
 |          [Qwen2-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)          | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct, Qwen/Qwen2-Math-72B, Qwen/Qwen2-Math-72B-Instruct, Qwen/Qwen2-Math-RM-72B                                                                                                                                                                                                               |
-|           [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](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-14B, Qwen/Qwen2.5-14B-Instruct, 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                                                                                                                                                                                                                                                                                                                                    |
 
 * 4D 并行和算子优化已支持 LLaMA 系列、Baichuan 系列、Bloom 系列、ChatGLM 系列、Gemma 系列、Mistral 系列、OPT 系列和 Qwen 系列，【LLM】模型4D 并行和算子支持列表如下：
 
+
 | 模型名称/并行能力支持 | 数据并行 | 张量模型并行 |          | 参数分片并行 |        |        | 流水线并行 |
 |:---------------------:|:--------:|:------------:|:--------:|:------------:|:------:|:------:|:----------:|
 |                       |          |   基础能力   | 序列并行 |    stage1    | stage2 | stage3 |            |
@@ -165,32 +154,31 @@
 | [OPT](./llm/config/opt)                    |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         🚧         |  🚧  |    ✅     |      🚧      |
 | [Gemma](./llm/config/gemma)                |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
 | [Yuan](./llm/config/yuan)                  |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
-* [大模型推理](./llm/docs/predict/inference.md)已支持 LLaMA 系列、Qwen 系列、DeepSeek 系列、Mistral 系列、ChatGLM 系列、Bloom 系列和 Baichuan 系列，支持 Weight Only INT8及 INT4推理，支持 WAC（权重、激活、Cache KV）进行 INT8、FP8量化的推理，【LLM】模型推理支持列表如下：
+* [大模型推理](./llm/docs/predict/inference.md)已支持 LLaMA 系列、Qwen 系列、Mistral 系列、ChatGLM 系列、Bloom 系列和 Baichuan 系列，支持 Weight Only INT8及 INT4推理，支持 WAC（权重、激活、Cache KV）进行 INT8、FP8量化的推理，【LLM】模型推理支持列表如下：
 
-|           模型名称/量化类型支持            | FP16/BF16 | WINT8 | WINT4 | INT8-A8W8 | FP8-A8W8 | INT8-A8W8C8 |
-|:------------------------------------------:|:---------:|:-----:|:-----:|:---------:|:--------:|:-----------:|
-|    [LLaMA](./llm/docs/predict/llama.md)    |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
-|     [Qwen](./llm/docs/predict/qwen.md)     |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
-| [DeepSeek](./llm/docs/predict/deepseek.md) |     ✅     |   ✅   |   ✅   |    🚧     |    ✅     |     🚧      |
-|   [Qwen-Moe](./llm/docs/predict/qwen.md)   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|  [Mixtral](./llm/docs/predict/mixtral.md)  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|                  ChatGLM                   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|                   Bloom                    |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|                  BaiChuan                  |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |     🚧      |
+|          模型名称/量化类型支持           | FP16/BF16 | WINT8 | WINT4 | INT8-A8W8 | FP8-A8W8 | INT8-A8W8C8 |
+|:----------------------------------------:|:---------:|:-----:|:-----:|:---------:|:--------:|:-----------:|
+|   [LLaMA](./llm/docs/predict/llama.md)   |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
+|    [Qwen](./llm/docs/predict/qwen.md)    |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
+|  [Qwen-Moe](./llm/docs/predict/qwen.md)  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+| [Mixtral](./llm/docs/predict/mixtral.md) |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                 ChatGLM                  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                  Bloom                   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                 BaiChuan                 |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |     🚧      |
 
 ## 安装
 
 ### 环境依赖
 
 * python >= 3.8
-* paddlepaddle >= 3.0.0rc1
+* paddlepaddle >= 3.0.0rc0
 
 如果您尚未安装 PaddlePaddle，请参考 [飞桨官网](https://www.paddlepaddle.org.cn/) 进行安装。
 
 ### pip 安装
 
 ```shell
-pip install --upgrade paddlenlp==3.0.0b4
+pip install --upgrade paddlenlp==3.0.0b3
 ```
 
 或者可通过以下命令安装最新 develop 分支代码：
@@ -199,7 +187,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/zh/get_started/installation.rst)。
+更多关于 PaddlePaddle 和 PaddleNLP 安装的详细教程请查看[Installation](./docs/get_started/installation.rst)。
 
 ------------------------------------------------------------------------------------------
 
@@ -210,14 +198,13 @@ 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")
-# 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语言模型，我可以回答各种问题，包括但不限于：天气、新闻、历史、文化、科学、教育、娱乐等。请问您有什么需要了解的吗？']
+>>> 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语言模型，我可以回答各种问题，包括但不限于：天气、新闻、历史、文化、科学、教育、娱乐等。请问您有什么需要了解的吗？']
 ```
 
 ### 大模型预训练

csrc/cpu/setup.sh

@@ -47,11 +47,9 @@ fi
 cd xFasterTransformer
 git apply paddle.patch
 
-#4. build xFasterTransformer
-cd 3rdparty
-bash prepare_oneccl.sh
-source ./oneccl/build/_install/env/setvars.sh
-cd ..
+# #4. build xFasterTransformer
+sh ./3rdparty/prepare_oneccl.sh
+source ./3rdparty/oneccl/build/_install/env/setvars.sh
 
 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"
-template <typename T>
+git adtemplate <typename T>
 void AvxCompute(const paddle::Tensor &x,
                 const paddle::Tensor &weight,
                 bool trans,
@@ -137,12 +137,12 @@ std::vector<paddle::Tensor> InvokeAvxWeightOnly(const paddle::Tensor &x,
 
 std::vector<std::vector<int64_t>> AvxWeightOnlyInferShape(
     std::vector<int64_t> x_shape,
-    std::vector<int64_t> weight_shape) {
+    std::vector<int64_t> weigh_shape) {
   int m = 1;
   for (int i = 0; i < x_shape.size() - 1; i++) {
     m = m * x_shape[i];
   }
-  return {std::vector<int64_t>{m, weight_shape[1]}};
+  return {std::vector<int64_t>{m, weigh_shape[1]}};
 }
 
 std::vector<paddle::DataType> AvxWeightOnlyInferDtype(

csrc/cpu/src/stop_generation_multi_ends.cc

@@ -14,20 +14,9 @@
 
 #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

@@ -1,125 +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 "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

@@ -1,528 +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.
-#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

@@ -1,30 +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"
-
-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,7 +53,6 @@ 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,
@@ -141,7 +140,6 @@ 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,
@@ -169,7 +167,6 @@ 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,
@@ -571,7 +568,6 @@ 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,
@@ -636,7 +632,6 @@ 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,
@@ -684,7 +679,6 @@ 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,
@@ -733,7 +727,6 @@ 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,
@@ -780,7 +773,6 @@ 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,
@@ -839,8 +831,7 @@ 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>>& excess_blocks_shape) {
+    const paddle::optional<std::vector<int64_t>>& out_linear_smooths_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];
@@ -885,7 +876,6 @@ 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,
@@ -959,8 +949,7 @@ 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("excess_blocks")})
+             paddle::Optional("out_linear_smooths")})
     .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,11 +146,8 @@ __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 / 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 bid = bidx, gid = 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;
@@ -187,10 +184,11 @@ __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 + bidy * HEAD_DIM_QK;
+  T *cu_q_smem = q_smem + gid * 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>;
@@ -205,14 +203,14 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   }
 
 
-  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + bdy * HEAD_DIM_QK * sizeof(CacheT)); 
+  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + GROUP_SIZE * 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 + bidy;
+      const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + gid;
       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,
@@ -269,7 +267,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 + bidy;
+      const uint32_t k_seq_offset = j * bdy + gid;
       produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
         kv_smem,
         cache_k,
@@ -344,10 +342,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 > 16 ? 16 : GROUP_SIZE;
-  const int gridx = bsz * (GROUP_SIZE / blocky);
+  constexpr int blocky = GROUP_SIZE;
+  const int gridx = bsz;
   
-  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>;
@@ -359,7 +357,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 + blocky * HEAD_DIM_QK * sizeof(T);
+  size_t smem_size = cache_smem_bytes + GROUP_SIZE * HEAD_DIM_QK * sizeof(T);
 
   if (smem_size >= 48 * 1024) {
     cudaFuncSetAttribute(

csrc/gpu/append_attn/encoder_write_cache_with_rope_impl.cuh

@@ -691,99 +691,6 @@ __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,
@@ -1556,12 +1463,9 @@ 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) {
@@ -1573,58 +1477,29 @@ void CascadeAppendWriteCacheKVQKV(
   auto head_dim_v = meta_data.head_dims_v;
   auto block_size = meta_data.block_size;
 
-  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 个
-
+  const uint32_t elem_nums =
+      num_tokens * kv_num_heads * (head_dim_qk + head_dim_v);
   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);
-  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);
-  }
+  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,7 +36,6 @@ 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,
@@ -49,7 +48,6 @@ 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(
@@ -95,12 +93,9 @@ 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

@@ -50,75 +50,19 @@ __global__ void split_q_block(const int* __restrict__ seq_lens_q,
   }
 }
 
-template <uint32_t config_size>
-__global__ void search_chunk_size_for_mla(
-  const int * __restrict__ seq_lens_q,
-  const int * __restrict__ seq_lens_encoder,
-  const int * __restrict__ seq_lens_decoder,
-  int * __restrict__ num_blocks_x,
-  int * __restrict__ res_chunk_size,
-  const int bsz,
-  const int set_chunk_size,
-  const int block_size,
-  const int sm_cout) {
-  const uint32_t conf_id = threadIdx.x;
-  int gridx = 0;
-  if (set_chunk_size > 0 && conf_id == 0) {
-    for (uint32_t bid = 0; bid < bsz; bid++) {
-      int seq_len = seq_lens_q[bid];
-      int seq_len_encoder = seq_lens_encoder[bid];
-      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
-      if (seq_len == 0 || seq_len_encoder > 0) continue;
-
-      int loop_times;
-      loop_times = cute::ceil_div(seq_len_decoder, set_chunk_size);
-      gridx += loop_times;
-    }
-    *num_blocks_x = gridx;
-    *res_chunk_size = set_chunk_size;
-  } else if (conf_id < config_size) {
-    __shared__ int gridx_shared[config_size];
-    // chunk_size is a multiple of 64
-    const int chunk_size = block_size << conf_id;
-    for (uint32_t bid = 0; bid < bsz; bid++) {
-      int seq_len = seq_lens_q[bid];
-      int seq_len_encoder = seq_lens_encoder[bid];
-      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
-      if (seq_len == 0 || seq_len_encoder > 0) continue;
-
-      int loop_times;
-      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
-      gridx += loop_times;
-    }
-    gridx_shared[conf_id] = gridx;
-    __syncthreads();
-    if (threadIdx.x == 0) {
-      uint32_t res_id = 0;
-      uint32_t max_last_wave_block = 0;
-      for (uint32_t i = 1; i < config_size; i++) {
-        uint32_t last_wave_block = gridx_shared[i] % sm_cout;
-        if (last_wave_block >= max_last_wave_block) {
-          res_id = i;
-          max_last_wave_block = last_wave_block;
-        }
-      }
-      *num_blocks_x = gridx_shared[res_id];
-      *res_chunk_size = block_size << res_id;
-    }
-  }
-}
-
-
-
-__global__ void split_block_for_mla(
-  const int * __restrict__ seq_lens_q,
-  const int * __restrict__ seq_lens_encoder,
-  const int * __restrict__ seq_lens_decoder,
-  int * __restrict__ batch_ids,
-  int * __restrict__ tile_ids_per_batch,
-  const int bsz,
-  const int chunk_size) {
+__global__ void split_q_block_mla(const int * __restrict__ seq_lens_q,
+                              const int * __restrict__ seq_lens_encoder,
+                              const int * __restrict__ seq_lens_decoder,
+                              int * __restrict__ batch_ids,
+                              int * __restrict__ tile_ids_per_batch,
+                              int * __restrict__ num_blocks_x,
+                              const int bsz,
+                              const int num_rows_per_block,
+                              const int chunk_size,
+                              const int GROUP_SIZE,
+                              const bool is_encoder) {
   if (threadIdx.x == 0) {
+    int gridx = 0;
     int index = 0;
     for (uint32_t bid = 0; bid < bsz; bid++) {
       int seq_len = seq_lens_q[bid];
@@ -128,15 +72,24 @@ __global__ void split_block_for_mla(
       if (seq_len == 0) continue;
 
       int loop_times;
-      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
-      if (seq_len_encoder > 0) {
-        loop_times = 0;
+      if (is_encoder) {
+        loop_times = cute::ceil_div(seq_len * GROUP_SIZE, num_rows_per_block);
+        if (seq_len_decoder > 0) {
+          loop_times = 0;
+        }
+      } else {
+        loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+        if (seq_len_encoder > 0) {
+          loop_times = 0;
+        }
       }
       for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) {
         batch_ids[index] = bid;
         tile_ids_per_batch[index++] = tile_id;
       }
+      gridx += loop_times;
     }
+    *num_blocks_x = gridx;
   }
 }
 
@@ -202,7 +155,7 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
     const int decoder_step_token_num) {
   paddle::Tensor encoder_batch_ids, encoder_tile_ids_per_batch, encoder_num_blocks_x_cpu,
     kv_batch_ids, kv_tile_ids_per_batch, kv_num_blocks_x_cpu, decoder_batch_ids,
-    decoder_tile_ids_per_batch, decoder_num_blocks_x, decoder_num_blocks_x_cpu, decoder_chunk_size_device, decoder_chunk_size_cpu;
+    decoder_tile_ids_per_batch, decoder_num_blocks_x, decoder_num_blocks_x_cpu;
   auto stream = seq_lens_this_time.stream();
   int bsz = cum_offsets.shape()[0];
   const int encoder_block_shape_q = get_encoder_block_shape_q();
@@ -223,55 +176,32 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
   // decoder
   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 && group_size <=64) {
-      const int set_chunk_size = get_mla_dec_chunk_size(bsz);
-      decoder_chunk_size_device =
-          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
-      decoder_num_blocks_x =
-          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
-
-      int device;
-      cudaGetDevice(&device);
-      int sm_cout;
-      cudaDeviceGetAttribute(&sm_cout, cudaDevAttrMultiProcessorCount, device);
-      constexpr int config_size = 12; // search space for chunk size:[64, 128, 256, ... 131072]
-
-      search_chunk_size_for_mla<config_size><<<1, 32, 0, stream>>>(
-        seq_lens_this_time.data<int>(),
-        seq_lens_encoder.data<int>(),
-        seq_lens_decoder.data<int>(),
-        decoder_num_blocks_x.data<int>(),
-        decoder_chunk_size_device.data<int>(),
-        bsz,
-        set_chunk_size,
-        block_size,
-        sm_cout);
-
-      decoder_chunk_size_cpu =
-        decoder_chunk_size_device.copy_to(paddle::CPUPlace(), false);
-      decoder_num_blocks_x_cpu =
-        decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
-
-      const int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
-      const int num_blocks = decoder_num_blocks_x_cpu.data<int>()[0];
+    const bool mla_use_tensorcore = get_mla_use_tensorcore();
+    if (mla_use_tensorcore) {
+      const int chunk_size = get_max_partition_size(bsz);
+      const int decoder_max_tile_size_per_bs = div_up(max_len_kv_cpu.data<int>()[0], chunk_size);
       decoder_batch_ids = 
-          GetEmptyTensor({num_blocks},
+          GetEmptyTensor({bsz * decoder_max_tile_size_per_bs},
                         paddle::DataType::INT32,
                         seq_lens_encoder.place());
       decoder_tile_ids_per_batch =
-          GetEmptyTensor({num_blocks},
+          GetEmptyTensor({bsz * decoder_max_tile_size_per_bs},
                         paddle::DataType::INT32,
                         seq_lens_encoder.place());
-
-      split_block_for_mla<<<1, 32, 0, stream>>>(
+      decoder_num_blocks_x =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());                  
+      split_q_block_mla<<<1, 32, 0, stream>>>(
         seq_lens_this_time.data<int>(),
         seq_lens_encoder.data<int>(),
         seq_lens_decoder.data<int>(),
         decoder_batch_ids.data<int>(),
         decoder_tile_ids_per_batch.data<int>(),
+        decoder_num_blocks_x.data<int>(),
         bsz,
-        chunk_size
+        block_size,
+        chunk_size,
+        group_size,
+        false // is_encoder
       );
     } else {
       const uint32_t decoder_max_tile_size_per_bs_q =
@@ -294,18 +224,15 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
                                           bsz,
                                           decoder_block_shape_q,
                                           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());
     }
+
+    decoder_num_blocks_x_cpu =
+        decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
   } else {
     decoder_batch_ids =
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
     decoder_tile_ids_per_batch =
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
-    decoder_chunk_size_cpu =
-        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 =
@@ -382,7 +309,6 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
           decoder_tile_ids_per_batch,
           decoder_num_blocks_x,
           decoder_num_blocks_x_cpu, /*cpu*/
-          decoder_chunk_size_cpu, /*cpu*/
           max_len_kv_cpu /*cpu*/};
 }
 
@@ -403,7 +329,6 @@ std::vector<paddle::DataType> GetBlockShapeAndSplitKVBlockInferDtype(
           paddle::DataType::INT32,
           paddle::DataType::INT32,
           paddle::DataType::INT32,
-          paddle::DataType::INT32,
           paddle::DataType::INT32};
 }
 
@@ -426,7 +351,6 @@ std::vector<std::vector<int64_t>> GetBlockShapeAndSplitKVBlockInferShape(
           dynamic_shape,
           {1},
           {1},
-          {1},
           {1}};
 }
 
@@ -447,7 +371,6 @@ PD_BUILD_OP(get_block_shape_and_split_kv_block)
               "decoder_tile_ids_per_batch",
               "decoder_num_blocks",
               "decoder_num_blocks_cpu",
-              "decoder_chunk_size_cpu",
               "max_len_kv"})
     .Attrs({"group_size: int",
             "block_size: int",

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_bfloat16_bfloat16_kernel.cu

@@ -34,7 +34,6 @@ 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,7 +33,6 @@ 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,7 +33,6 @@ 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,7 +33,6 @@ 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

@@ -1,365 +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 "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/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h

@@ -117,7 +117,7 @@ class LeftGELUAndMul {
   CUTLASS_HOST_DEVICE
   FragmentOutput operator()(FragmentAccumulator const &lhs,
                             FragmentAccumulator const &rhs) const {
-    // Convert source to internal compute numeric type
+    // Convert source to interal compute numeric type
     NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
         accumulator_to_compute;
 

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h

@@ -117,7 +117,7 @@ class LeftSiLUAndMul {
   CUTLASS_HOST_DEVICE
   FragmentOutput operator()(FragmentAccumulator const &lhs,
                             FragmentAccumulator const &rhs) const {
-    // Convert source to internal compute numeric type
+    // Convert source to interal compute numeric type
     NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
         accumulator_to_compute;
 

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/threadblock/dual_mma_base.h

@@ -92,7 +92,7 @@ class DualMmaBase {
                               Shape::kN / WarpGemm::kN,
                               Shape::kK / WarpGemm::kK>;
 
-  /// Number of warp-level GEMM operations
+  /// Number of warp-level GEMM oeprations
   static int const kWarpGemmIterations =
       (WarpGemm::kK / Operator0::Policy::MmaShape::kK);
 

csrc/gpu/dequant_int8.cu

@@ -158,7 +158,7 @@ std::vector<paddle::DataType> DequantInt8Dtype(const paddle::DataType& input_dty
 }
 
 PD_BUILD_OP(dequant_int8)
-    .Inputs({"input","out_scale"})
+    .Inputs({"intput","out_scale"})
     .Outputs({"output"})
     .Attrs({"dtype: std::string"})
     .SetKernelFn(PD_KERNEL(DequantInt8))

csrc/gpu/env.h

@@ -31,8 +31,8 @@ inline uint32_t get_encoder_block_shape_q() {
 inline uint32_t get_max_partition_size(int bsz) {
     static const char* max_partition_size_env = std::getenv("FLAGS_cascade_attention_max_partition_size");
     static const uint32_t max_partition_size =
-            max_partition_size_env == nullptr ? 32768 : std::stoul(std::string(max_partition_size_env));
-    return max_partition_size;
+            max_partition_size_env == nullptr ? 0 : std::stoul(std::string(max_partition_size_env));
+    return (max_partition_size != 0 ? max_partition_size : (bsz == 1 ? 128 : 512));
 }
 
 inline uint32_t get_cascade_attention_deal_each_time() {
@@ -56,30 +56,16 @@ inline uint32_t get_cascade_attention_num_threads() {
     return cascade_attention_num_threads != 0 ? cascade_attention_num_threads : 128;
 }
 
-inline bool get_flags_mla_use_tensorcore() {
+inline bool get_mla_use_tensorcore() {
     static const char* mla_use_tensorcore_env = std::getenv("FLAGS_mla_use_tensorcore");
     static const uint32_t mla_use_tensorcore =
             mla_use_tensorcore_env == nullptr ? 1 : std::stoul(std::string(mla_use_tensorcore_env));
     return mla_use_tensorcore != 0 ? true : false;
 }
 
-inline bool get_mla_use_wg4() {
-    static const char* mla_use_wg4_env = std::getenv("FLAGS_mla_use_wg4");
-    static const uint32_t mla_use_wg4 =
-            mla_use_wg4_env == nullptr ? 1 : std::stoul(std::string(mla_use_wg4_env));
-    return mla_use_wg4 != 0 ? true : false;
-}
-
 inline bool enable_cuda_core_fp8_gemm() {
     static const char* enable_cuda_core_fp8_env = std::getenv("FLAGS_cuda_core_fp8_gemm");
     static const bool enable_cuda_core_fp8_gemm =
             enable_cuda_core_fp8_env != nullptr && std::string(enable_cuda_core_fp8_env) == "1";
     return enable_cuda_core_fp8_gemm;
 }
-
-inline int get_mla_dec_chunk_size(int bsz) {
-    static const char* mla_dec_chunk_size_env = std::getenv("FLAGS_mla_dec_chunk_size");
-    static const int mla_dec_chunk_size =
-            mla_dec_chunk_size_env == nullptr ? -1 : std::stoi(std::string(mla_dec_chunk_size_env));
-    return bsz > 1 ? mla_dec_chunk_size : 64;
-}

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"
 
-paddle::Tensor cutlass_fp8_fp8_half_block_gemm_fused_func(
+std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
     const paddle::Tensor& x,
     const paddle::Tensor& y,
     const paddle::Tensor& x_scale,
@@ -145,20 +145,29 @@ paddle::Tensor cutlass_fp8_fp8_half_block_gemm_fused_func(
                                     x_scale_ptr,
                                     y_scale_ptr};
     fp8_fp8_block_gemm_scale_bias_act(params);
-  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)};
+  // 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};
 }
 
 std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfBlockGemmFusedInferShape(

csrc/gpu/fused_rotary_position_encoding.cu

@@ -101,11 +101,6 @@ void FusedRotaryPositionEncoding(
   int64_t query_stride = num_heads * head_size;
   int64_t key_stride = num_kv_heads * head_size;
 
-  if (num_tokens > 65535) {
-    PD_THROW(
-        "apply_rotary_embedding_kernel launch failed when num_tokens > 65535.");
-  }
-
   dim3 grid(num_tokens);
   dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
   PD_DISPATCH_FLOATING_AND_HALF_TYPES(

csrc/gpu/get_output.cc

@@ -20,7 +20,6 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
-// #define GET_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -33,20 +32,8 @@ void GetOutput(const paddle::Tensor& x,
   if (rank_id > 0) return;
 
   static struct msgdata msg_rcv;
-  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 key_t key = ftok("./", 1);
 
   static int msgid = msgget(key, IPC_CREAT | 0666);
 

csrc/gpu/get_position_ids_and_mask_encoder_batch.cu

@@ -63,7 +63,7 @@ void GetPositionIdsAndMaskEncoderBatch(
     const paddle::Tensor& seq_lens_this_time,
     const paddle::Tensor& position_ids,
     const paddle::Tensor& mask_encoder_batch) {
-  const int bsz = seq_lens_this_time.shape()[0];
+  const int bsz = seq_lens_encoder.shape()[0];
 
   GetPositionIdsAndMaskEncoderBatchKernel<<<1, bsz, 0, position_ids.stream()>>>(
       seq_lens_encoder.data<int>(),

csrc/gpu/group_quant.cu

@@ -0,0 +1,201 @@
+// 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,7 +45,6 @@ namespace cub = hipcub;
 #include "paddle/phi/backends/gpu/gpu_info.h"
 #include "nlohmann/json.hpp"
 
-
 using json = nlohmann::json;
 
 #define CUDA_CHECK(call)                           \
@@ -162,13 +161,6 @@ public:
   typedef paddle::float8_e4m3fn data_t;
 };
 
-template <>
-class PDTraits<paddle::DataType::INT8> {
-public:
-  typedef int8_t DataType;
-  typedef int8_t data_t;
-};
-
 template <typename T, int Size>
 struct alignas(sizeof(T) * Size) AlignedVector {
   T val[Size];
@@ -245,42 +237,3 @@ inline int GetSMVersion() {
       phi::backends::gpu::GetCurrentDeviceId());
   return sm_version;
 }
-
-inline bool GetMlaUseTensorcore() {
-  static const bool flags_mla_use_tensorcore = get_flags_mla_use_tensorcore();
-  static const bool enable_mla_tensorcore = GetSMVersion() >= 90 ? true : false;
-  const bool mla_use_tensorcore = flags_mla_use_tensorcore && enable_mla_tensorcore;
-  return mla_use_tensorcore;
-}
-
-__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/int8_gemm_with_cutlass/epilogue_tensor_op_int32.h

@@ -176,7 +176,7 @@ struct DefaultIteratorsTensorOp<cutlass::bfloat16_t,
 ///
 /// Satisfies: ReadableTileIterator
 ///
-template <typename ThreadMap_  ///< Thread map (concept: OutputTileThreadMap)
+template <typename ThreadMap_  ///< Thread map (conept: OutputTileThreadMap)
           >
 class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 16, 8, 8> {
  public:

csrc/gpu/mla_attn/attention_updater.cuh

@@ -19,8 +19,6 @@
  * Modified by the FlashInfer team.
  */
 
-#pragma once
-
 #include <cute/tensor.hpp>
 #include <cutlass/detail/helper_macros.hpp>
 

csrc/gpu/mla_attn/batch_mla_with_paged_kv_cache.cu

@@ -24,7 +24,6 @@
 #include <vector>
 #include "cute/tensor.hpp"
 #include "mla_hopper.cuh"
-#include "mla_hopper_wg4.cuh"
 #include <iostream>
 #include <string>
 #include <sstream>
@@ -81,14 +80,13 @@ void BatchMLAWithPagedKVCacheKernel(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
-    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int draft_total_token_num,
+    const int draft_token_num,
     const bool causal,
     cudaStream_t& stream,
     paddle::Tensor* out) {
@@ -100,12 +98,9 @@ void BatchMLAWithPagedKVCacheKernel(
   const auto q_head_num = meta_data.q_num_heads;
   const auto max_block_num_per_seq = meta_data.max_blocks_per_seq;
   const auto max_block_num = bsz * max_block_num_per_seq;
-  const bool mla_use_wg4 = get_mla_use_wg4();
-#if CUDA_VERSION >= 12080
-  constexpr bool USE_REG_EALLOC = true;
-#else
-  constexpr bool USE_REG_EALLOC = false;
-#endif
+  const uint32_t chunk_size = get_max_partition_size(bsz);
+
+
   int q_head_dim = meta_data.head_dims;
   int k_head_dim = meta_data.head_dims;
   int v_head_dim = meta_data.head_dims_v;
@@ -116,116 +111,59 @@ void BatchMLAWithPagedKVCacheKernel(
   phi::Allocator::AllocationPtr O_tmp, m_tmp, d_tmp;
   O_tmp = allocator->Allocate(
       phi::SizeOf(q.dtype()) *
-      static_cast<size_t>(num_chunks * bsz * draft_total_token_num * q_head_num * v_head_dim));
+      static_cast<size_t>(num_chunks * bsz * draft_token_num * q_head_num * v_head_dim));
   m_tmp = allocator->Allocate(
       sizeof(float) *
-      static_cast<size_t>(num_chunks * bsz * draft_total_token_num * q_head_num));
+      static_cast<size_t>(num_chunks * bsz * draft_token_num * q_head_num));
   d_tmp = allocator->Allocate(
       sizeof(float) *
-      static_cast<size_t>(num_chunks * bsz * draft_total_token_num * q_head_num));
+      static_cast<size_t>(num_chunks * bsz * draft_token_num * q_head_num));
+
+  Params<CUTLASS_TYPE, CUTLASS_TYPE, CUTLASS_TYPE, int> params = {};
+  params.Q = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(q.data<T>()));
+  params.KV = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(latent_cache.data<T>()));
+  params.O = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(out->data<T>()));
+  params.O_tmp = reinterpret_cast<CUTLASS_TYPE*>(O_tmp->ptr());
+  params.m = reinterpret_cast<float*>(m_tmp->ptr());
+  params.d = reinterpret_cast<float*>(d_tmp->ptr());
+  params.block_tables = const_cast<int*>(block_tables.data<int>());
+  params.seq_lens_this_time = const_cast<int*>(seq_lens_this_time.data<int>());
+  params.seq_lens_encoder = const_cast<int*>(seq_lens_encoder.data<int>());
+  params.seq_lens_decoder = const_cast<int*>(seq_lens_decoder.data<int>());
+  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>());
+  params.num_blocks_x = const_cast<int*>(num_blocks_x_device.data<int>());
+  params.q_stride_bsz = q_head_num * q_head_dim;
+  params.q_stride_head_num = q_head_dim;
+  params.kv_stride_block_num = block_size * k_head_dim;
+  params.kv_stride_block_size = k_head_dim;
+  params.o_stride_bsz = q_head_num * v_head_dim;
+  params.o_stride_head_num = v_head_dim;
+  params.bsz = bsz;
+  params.token_num = token_num;
+  params.max_seq_len = max_seq_len;
+  params.max_block_num = max_block_num;
+  params.max_block_num_per_seq = max_block_num_per_seq;
+  params.q_num_head = q_head_num;
+  params.qk_head_dim = q_head_dim;
+  params.vo_head_dim = v_head_dim;
+  params.block_size = block_size;
+  params.max_draft_token_num = draft_token_num;
+  params.sm_scale = softmax_scale;
+  params.chunk_size = chunk_size;
+  params.chunk_num = num_chunks;
 
   if (q_head_dim == 576) {
-      if (mla_use_wg4 && block_size == 32) {
-        if constexpr (!std::is_same<CUTLASS_TYPE, cutlass::half_t>::value) {
-          PD_THROW("error!!! mla_wg4 not supports bf16 now!!!\n");
-        } else {
-          using ParamsType = Params<cutlass::half_t, cutlass::half_t, cutlass::half_t, int> ;
-          ParamsType params = {};
-          params.Q = reinterpret_cast<cutlass::half_t*>(const_cast<T*>(q.data<T>()));
-          params.KV = reinterpret_cast<cutlass::half_t*>(const_cast<T*>(latent_cache.data<T>()));
-          params.O = reinterpret_cast<cutlass::half_t*>(const_cast<T*>(out->data<T>()));
-          params.O_tmp = reinterpret_cast<cutlass::half_t*>(O_tmp->ptr());
-          params.m = reinterpret_cast<float*>(m_tmp->ptr());
-          params.d = reinterpret_cast<float*>(d_tmp->ptr());
-          params.block_tables = const_cast<int*>(block_tables.data<int>());
-          params.seq_lens_this_time = const_cast<int*>(seq_lens_this_time.data<int>());
-          params.seq_lens_encoder = const_cast<int*>(seq_lens_encoder.data<int>());
-          params.seq_lens_decoder = const_cast<int*>(seq_lens_decoder.data<int>());
-          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>());
-          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;
-          params.q_stride_head_num = q_head_dim;
-          params.kv_stride_block_num = block_size * k_head_dim;
-          params.kv_stride_block_size = k_head_dim;
-          params.o_stride_bsz = q_head_num * v_head_dim;
-          params.o_stride_head_num = v_head_dim;
-          params.bsz = bsz;
-          params.token_num = token_num;
-          params.max_seq_len = max_seq_len;
-          params.max_block_num = max_block_num;
-          params.max_block_num_per_seq = max_block_num_per_seq;
-          params.q_num_head = q_head_num;
-          params.qk_head_dim = q_head_dim;
-          params.vo_head_dim = v_head_dim;
-          params.block_size = block_size;
-          params.draft_total_token_num = draft_total_token_num;
-          params.sm_scale = softmax_scale;
-          params.chunk_size = chunk_size;
-          params.chunk_num = num_chunks;
-          BatchMLAWithPagedKVCacheWG4Dispatched<576, 512, half, ParamsType, USE_REG_EALLOC>(
-              params, stream
-          );
-        }
-      } else {
-        using ParamsType = Params<CUTLASS_TYPE, CUTLASS_TYPE, CUTLASS_TYPE, int>;
-        ParamsType params = {};
-        params.Q = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(q.data<T>()));
-        params.KV = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(latent_cache.data<T>()));
-        params.O = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(out->data<T>()));
-        params.O_tmp = reinterpret_cast<CUTLASS_TYPE*>(O_tmp->ptr());
-        params.m = reinterpret_cast<float*>(m_tmp->ptr());
-        params.d = reinterpret_cast<float*>(d_tmp->ptr());
-        params.block_tables = const_cast<int*>(block_tables.data<int>());
-        params.seq_lens_this_time = const_cast<int*>(seq_lens_this_time.data<int>());
-        params.seq_lens_encoder = const_cast<int*>(seq_lens_encoder.data<int>());
-        params.seq_lens_decoder = const_cast<int*>(seq_lens_decoder.data<int>());
-        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>());
-        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;
-        params.q_stride_head_num = q_head_dim;
-        params.kv_stride_block_num = block_size * k_head_dim;
-        params.kv_stride_block_size = k_head_dim;
-        params.o_stride_bsz = q_head_num * v_head_dim;
-        params.o_stride_head_num = v_head_dim;
-        params.bsz = bsz;
-        params.token_num = token_num;
-        params.max_seq_len = max_seq_len;
-        params.max_block_num = max_block_num;
-        params.max_block_num_per_seq = max_block_num_per_seq;
-        params.q_num_head = q_head_num;
-        params.qk_head_dim = q_head_dim;
-        params.vo_head_dim = v_head_dim;
-        params.block_size = block_size;
-        params.draft_total_token_num = draft_total_token_num;
-        params.sm_scale = softmax_scale;
-        params.chunk_size = chunk_size;
-        params.chunk_num = num_chunks;
-        BatchMLAWithPagedKVCacheDispatched<576, 512, NV_TYPE, ParamsType, USE_REG_EALLOC>(
-            params, stream
-        );
-      }
+      BatchMLAWithPagedKVCacheDispatched<576, 512, NV_TYPE>(
+          params, stream
+      );
   } else {
       PD_THROW("error!!! q_head_dim must be 576 !!!\n");
   }
 }
 
-
 template void BatchMLAWithPagedKVCacheKernel<paddle::bfloat16>(
     const AppendAttnMetaData& meta_data,
     const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]
@@ -249,14 +187,13 @@ template void BatchMLAWithPagedKVCacheKernel<paddle::bfloat16>(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
-    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int draft_total_token_num,
+    const int draft_token_num,
     const bool causal,
     cudaStream_t& stream,
     paddle::Tensor* out);
@@ -285,14 +222,13 @@ template void BatchMLAWithPagedKVCacheKernel<paddle::float16>(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
-    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int draft_total_token_num,
+    const int draft_token_num,
     const bool causal,
     cudaStream_t& stream,
     paddle::Tensor* out);

csrc/gpu/mla_attn/batch_mla_with_paged_kv_cache.h

@@ -57,14 +57,13 @@ void BatchMLAWithPagedKVCacheKernel(
     const paddle::Tensor& num_blocks_x_device,
     const std::string& cache_quant_type_str,
     const int num_blocks_x,
-    const int chunk_size,
     const int max_seq_len,
     const int max_dec_len,
     const float softmax_scale,
     const float quant_max_bound,
     const float quant_min_bound,
     const float in_scale,
-    const int draft_total_token_num,
+    const int draft_token_num,
     const bool causal,
     cudaStream_t& stream,
     paddle::Tensor* out);

csrc/gpu/mla_attn/epilogue.cuh

@@ -19,7 +19,6 @@
  * Modified by the FlashInfer team.
  */
 
-#pragma once
 
 #ifndef ATTENTION_HOPPER_EPILOGUE_CUH_
 #define ATTENTION_HOPPER_EPILOGUE_CUH_
@@ -129,12 +128,10 @@ struct CollectiveEpilogue {
                             const int seq_len_now,
                             const int start_token_idx,
                             const int tile_idx,
-                            const int q_tile_idx,
                             const int kv_len,
                             const int chunk_size,
-                            const int draft_total_token_num,
-                            const int o_stride_head) {
-    const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+                            const int max_draft_token_num,
+                            const int o_stride_bsz) {
     const int num_chunks = cute::ceil_div(kv_len, chunk_size);
     Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
     auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
@@ -152,7 +149,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 * 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;
+    auto O_ptr = num_chunks == 1 ? epilogue_params.O_ptr + start_token_idx * o_stride_bsz : epilogue_params.O_ptr_tmp + (tile_idx * bsz + bid) * max_draft_token_num * o_stride_bsz;
     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)
@@ -167,7 +164,7 @@ struct CollectiveEpilogue {
     // copy if not out of bound
     auto predicate_fn = [&](auto coords) {
       auto s_coords = tOcOGroup(_0{}, coords);
-      return elem_less((get<0>(s_coords) + q_group_offset) / Ktraits::GROUP_SIZE, seq_len_now);
+      return elem_less(get<0>(s_coords) / Ktraits::GROUP_SIZE, seq_len_now);
     };
     copy_if(gmem_tiled_copy_O, predicate_fn, tOsOGroup, tOgOGroup);
   }

csrc/gpu/mla_attn/kernel_traits.cuh

@@ -19,8 +19,6 @@
  * Modified by the FlashInfer team.
  */
 
-#pragma once
-
 #ifndef ATTENTION_HOPPER_KERNEL_TRAITS_CUH_
 #define ATTENTION_HOPPER_KERNEL_TRAITS_CUH_
 

csrc/gpu/mla_attn/kernel_traits_wg4.cuh

@@ -1,160 +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.
-
-#pragma once
-
-#ifndef ATTENTION_HOPPER_KERNEL_TRAITS_WG4_CUH_
-#define ATTENTION_HOPPER_KERNEL_TRAITS_WG4_CUH_
-
-#include <type_traits>
-#include "cute/algorithm/copy.hpp"
-#include "cute/atom/mma_atom.hpp"
-#include "cutlass/cutlass.h"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-#include "cutlass/layout/layout.h"
-#include "cutlass/numeric_types.h"
-#include "cutlass/pipeline/pipeline.hpp"
-
-namespace mla_attn {
-
-using namespace cute;
-
-template <typename MainloopPipeline, typename MainloopPipelineQ, typename MainloopPipelineQK, class DTypeQ, class DTypeKV, class DTypeQKAccum, class DTypeOut, class IdType,
-          int BLOCK_SHAPE_KV, class SmemLayoutQ, class SmemLayoutK, class SmemLayoutP, class SmemLayoutRow, class SmemLayoutO>
-struct alignas(16) SharedStorageQKVWG4O {
-  alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutQ>> smem_q;
-  alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutP>> smem_p;
-  alignas(16) cute::array_aligned<DTypeQKAccum, cute::cosize_v<SmemLayoutRow>> smem_scale;
-  union {
-    alignas(16) cute::array_aligned<DTypeKV, cute::cosize_v<SmemLayoutK>> smem_kv;
-    alignas(16) cute::array_aligned<DTypeOut, cute::cosize_v<SmemLayoutO>> smem_o;
-  };
-  struct {
-    alignas(16) typename MainloopPipelineQ::SharedStorage pipeline_q;
-    alignas(16) typename MainloopPipelineQK::SharedStorage pipeline_qk;
-    alignas(16) typename MainloopPipeline::SharedStorage pipeline_kv;
-  };
-};
-
-template <bool USE_TMA_LOAD_KV_, int HEAD_DIM_QK_, int HEAD_DIM_VO_, int GROUP_SIZE_, int BLOCK_SHAPE_Q_, int BLOCK_SHAPE_KV_,
-          int NUM_STAGES_, typename DTypeQ_, typename DTypeKV_, typename DTypeO_, typename IdType_, typename NV_TYPE_>
-struct AttentionKernelWG4Traits {
-
-  using DTypeQ = DTypeQ_;
-  using DTypeKV = DTypeKV_;
-  using DTypeO = DTypeO_;
-  using IdType = IdType_;
-//   using DTypeQKAccum = float;
-//   using DTypePVAccum = float;
-  using DTypeQKAccum = DTypeQ;
-  using DTypePVAccum = DTypeO;
-  using NV_TYPE = NV_TYPE_;
-  
-  
-  static constexpr bool USE_TMA_LOAD_KV = USE_TMA_LOAD_KV_;
-  static constexpr int GROUP_SIZE = GROUP_SIZE_;
-  static constexpr int BLOCK_SHAPE_Q = BLOCK_SHAPE_Q_;
-  static_assert(BLOCK_SHAPE_Q % 64 == 0);
-  static constexpr int BLOCK_SHAPE_KV = BLOCK_SHAPE_KV_;
-  static constexpr int HEAD_DIM_QK = HEAD_DIM_QK_;
-  static constexpr int HEAD_DIM_VO = HEAD_DIM_VO_;
-  static constexpr int NUM_PER_STAGE = BLOCK_SHAPE_KV * HEAD_DIM_QK;
-  static_assert(HEAD_DIM_QK % 32 == 0);
-  static_assert(HEAD_DIM_VO % 32 == 0);
-
-
-  using TileShape_QKD = Shape<Int<BLOCK_SHAPE_Q>, Int<BLOCK_SHAPE_KV>, Int<HEAD_DIM_QK>>;
-  using TileShape_PDV = Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
-
-  static constexpr int NUM_STAGES = NUM_STAGES_;
-
-  using AtomLayoutQKD = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _1, _1>>;
-  using AtomLayoutPV = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _2, _1>>;
-  using TiledMmaQK = decltype(cute::make_tiled_mma(
-      cute::GMMA::ss_op_selector<DTypeQ, DTypeKV, DTypeQKAccum, TileShape_QKD>(), AtomLayoutQKD{}));
-  using TiledMmaPV = decltype(cute::make_tiled_mma(
-      cute::GMMA::rs_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
-                                 GMMA::Major::K, GMMA::Major::MN>(),
-      AtomLayoutPV{}));
-  using TiledMmaPVSS = decltype(cute::make_tiled_mma(
-      cute::GMMA::ss_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
-                                 GMMA::Major::K, GMMA::Major::MN>(),
-      AtomLayoutPV{}));
-
-  static constexpr int NUM_MMA_THREADS_QK = size(TiledMmaQK{});
-  static constexpr int NUM_MMA_THREADS_PV = size(TiledMmaPVSS{});
-  static constexpr int NUM_MMA_THREADS = NUM_MMA_THREADS_PV;
-  static constexpr int NUM_PRODUCER_THREADS = cutlass::NumThreadsPerWarpGroup;
-  static constexpr int NUM_THREADS = NUM_MMA_THREADS_QK + NUM_MMA_THREADS_PV + NUM_PRODUCER_THREADS;
-  static constexpr int NUM_WARPS = NUM_THREADS / cutlass::NumThreadsPerWarp;
-
-  using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
-                                   decltype(cute::get<2>(TileShape_QKD{}))>());
-  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{})));
-
-  using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeKV, decltype(cute::get<1>(TileShape_QKD{})),
-                                   decltype(cute::get<2>(TileShape_QKD{}))>());
-  using SmemLayoutK = decltype(tile_to_shape(
-      SmemLayoutAtomK{},
-      make_shape(shape<1>(TileShape_QKD{}), shape<2>(TileShape_QKD{}), Int<NUM_STAGES>{})));
-  using SmemLayoutVt = decltype(composition(
-      SmemLayoutK{}, make_ordered_layout(make_shape(get<2>(TileShape_QKD{}),
-                                                    get<1>(TileShape_QKD{}), Int<NUM_STAGES>{}),
-                                         Step<_2, _1, _3>{})));
-  using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeKV, decltype(cute::get<2>(TileShape_PDV{})),
-                                   decltype(cute::get<1>(TileShape_PDV{}))>());
-  using SmemLayoutV = decltype(tile_to_shape(
-      SmemLayoutAtomV{},
-      make_shape(get<2>(TileShape_PDV{}), get<1>(TileShape_PDV{}), Int<1>{})));
-
-  // Note this is the transpose in terms of the view, not in terms of memory.
-  using SmemLayoutVtOneStage = decltype(composition(
-      SmemLayoutV{}, make_ordered_layout(make_shape(get<1>(TileShape_PDV{}),
-                                                    get<2>(TileShape_PDV{}), Int<1>{}),
-                                         Step<_2, _1, _3>{})));
-
-  using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})),
-                                   decltype(cute::get<1>(TileShape_PDV{}))>());
-  using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
-
-  using SmemCopyAtom = Copy_Atom<cute::SM90_U32x4_STSM_N, DTypeQ>;
-  // todo: permute layout
-  using SmemLayoutRow = Layout<Shape<Int<2>, Int<128>, Int<2>>, Stride<Int<1>, Int<2>, Int<256>>>;
-
-  using SmemLayoutAtomP = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
-                                   decltype(cute::get<1>(TileShape_QKD{}))>());
-  static constexpr int QK_STAGE = 2;
-  using SmemLayoutPSS = decltype(tile_to_shape(SmemLayoutAtomP{}, make_shape(shape<0>(TileShape_QKD{}), shape<1>(TileShape_QKD{}), Int<QK_STAGE>{})));
-
-  using MainloopPipelineQ = typename cutlass::PipelineAsync<1>;
-  using PipelineStateQ = typename cutlass::PipelineState<1>;
-  using MainloopPipelineQK = typename cutlass::PipelineAsync<QK_STAGE>;
-  using PipelineStateQK = typename cutlass::PipelineState<QK_STAGE>;
-  using MainloopPipeline =
-      std::conditional_t<USE_TMA_LOAD_KV, typename cutlass::PipelineTmaAsync<NUM_STAGES>,
-                         typename cutlass::PipelineAsync<NUM_STAGES>>;
-  using PipelineState = typename cutlass::PipelineState<NUM_STAGES>;
-
-  using SharedStorage = SharedStorageQKVWG4O<MainloopPipeline, MainloopPipelineQ, MainloopPipelineQK, DTypeQ, DTypeKV, DTypeQKAccum, DTypeO, IdType, BLOCK_SHAPE_KV,
-                                             SmemLayoutQ, SmemLayoutK, SmemLayoutPSS, SmemLayoutRow, SmemLayoutO>;
-};
-
-}  // namespace mla_attn
-
-#endif // ATTENTION_HOPPER_KERNEL_TRAITS_WG4_CUH_

csrc/gpu/mla_attn/mainloop_load.cuh

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#pragma once
-
 #ifndef ATTENTION_HOPPER_MAINLOOP_LOAD_CUH_
 #define ATTENTION_HOPPER_MAINLOOP_LOAD_CUH_
 
@@ -135,7 +133,6 @@ struct CollectiveMainloop {
     IdType const* cumsum_q_seqlens;
     IdType const* batch_ids;
     IdType const* tile_ids_per_batch;
-    IdType const* q_tile_ids_per_batch;
     IdType const* num_blocks_x;
     float sm_scale;
     int bsz;
@@ -149,7 +146,7 @@ struct CollectiveMainloop {
     int o_stride_head_num;
     int chunk_size;
     int chunk_num;
-    int draft_total_token_num;
+    int max_draft_token_num;
   };
 
   // Device side kernel params
@@ -168,7 +165,6 @@ struct CollectiveMainloop {
     IdType* cumsum_q_seqlens;
     IdType* batch_ids;
     IdType* tile_ids_per_batch;
-    IdType* q_tile_ids_per_batch;
     IdType* num_blocks_x;
     float sm_scale;
     int bsz;
@@ -182,7 +178,7 @@ struct CollectiveMainloop {
     int o_stride_head_num;
     int chunk_size;
     int chunk_num;
-    int draft_total_token_num;
+    int max_draft_token_num;
     TMA_KV tma_load_KV;
   };
 
@@ -207,7 +203,6 @@ struct CollectiveMainloop {
             const_cast<IdType*>(args.cumsum_q_seqlens),
             const_cast<IdType*>(args.batch_ids),
             const_cast<IdType*>(args.tile_ids_per_batch),
-            const_cast<IdType*>(args.q_tile_ids_per_batch),
             const_cast<IdType*>(args.num_blocks_x),
             args.sm_scale,
             args.bsz,
@@ -221,7 +216,7 @@ struct CollectiveMainloop {
             args.o_stride_head_num,
             args.chunk_size,
             args.chunk_num,
-            args.draft_total_token_num,
+            args.max_draft_token_num,
             tma_load_KV
             };
   }
@@ -239,11 +234,9 @@ struct CollectiveMainloop {
                              PipelineStateQ& smem_pipe_write_q,
                              SharedStorage& shared_storage,
                              const int thread_idx,
-                             const int bid,
-                             const int q_tile_idx) {
-    const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
-    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;
+                             const int bid) {
+    int start_q_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+    int offset_Q = mainloop_params.q_stride_bsz * start_q_token_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{});
@@ -260,7 +253,7 @@ struct CollectiveMainloop {
     int valid_q_size = mainloop_params.seq_lens_this_time[bid];
     auto q_predicate_fn = [&](auto coords) {
       auto s_coords = tQcQGroup(_0{}, coords);
-      return elem_less((get<0>(s_coords) + q_group_offset) / Ktraits::GROUP_SIZE, valid_q_size);
+      return elem_less(get<0>(s_coords) / Ktraits::GROUP_SIZE, valid_q_size);
     };
     Tensor tQgQiGroup = flatten_1(tQgQ);
     Tensor tQsQiGroup = flatten_1(tQsQ);

csrc/gpu/mla_attn/mainloop_mma.cuh

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#pragma once
-
 #ifndef ATTENTION_HOPPER_MAINLOOP_MMA_CUH_
 #define ATTENTION_HOPPER_MAINLOOP_MMA_CUH_
 
@@ -41,7 +39,6 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
                             const int kv_len,
                             const int qo_len,
                             const int tile_idx,
-                            const int q_tile_idx,
                             SharedStorage& shared_storage) {
   using DTypeQ = typename Ktraits::DTypeQ;
   using DTypeKV = typename Ktraits::DTypeKV;
@@ -65,11 +62,8 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
   static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
   static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
 
-  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
-
   Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
   Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-  // todo: split rope/norope for sVt, not use s1~s2.
   Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
   Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
   Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
@@ -125,7 +119,7 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
         Tensor tScS = threadMmaQK.partition_C(cS);
 #pragma unroll
         for (int i = 0; i < size(tSrS); ++i) {
-          int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+          int qo_idx = get<0>(tScS(i)) / Ktraits::GROUP_SIZE;
           int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
           if constexpr (!CAUSAL) {  // Just masking based on col
             if (kv_idx >= kv_len) {
@@ -189,13 +183,13 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
       const int warp_idx = thread_idx / 32;
 #pragma unroll
       for (int w_i = 0; w_i < 2; ++w_i) {
-        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
+        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i;
         const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
 
         if (token_idx < qo_len) {
-          // const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
-          const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
-          const int write_idx = bid * bid_offset + token_group_idx;
+          const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
+          const int bid_offset = mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE;
+          const int write_idx = bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx;
           mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
           mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
         }
@@ -251,11 +245,10 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
                                        const int kv_len,
                                        const int qo_len,
                                        const int tile_idx,
-                                       const int q_tile_idx,
                                        SharedStorage& shared_storage) {
   using DTypeQ = typename Ktraits::DTypeQ;
   using DTypeKV = typename Ktraits::DTypeKV;
-  using DTypeMD = typename Ktraits::DTypeO; 
+  using DTypeMD = typename Ktraits::DTypeO; // !!! bf16
   using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
   using IdType = typename Ktraits::IdType;
   using TileShape_QKD = typename Ktraits::TileShape_QKD;
@@ -275,14 +268,8 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
   static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
   static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
 
-  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
-
-  cutlass::FastDivmod stage_div(2);
-  int quotient, remainder;
-
   Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
   Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-  // TODO: split rope/norope for sVt, not use s1~s4.
   Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
   Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
   Tensor sVt_s3 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 2 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
@@ -338,7 +325,7 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       Tensor tScS = threadMmaQK.partition_C(cS);
 #pragma unroll
       for (int i = 0; i < size(tSrS); ++i) {
-        int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+        int qo_idx = get<0>(tScS(i)) / Ktraits::GROUP_SIZE;
         int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
         if constexpr (!CAUSAL) {  // Just masking based on col
           if (kv_idx >= kv_len) {
@@ -355,9 +342,8 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     Tensor scale_o = attention_updater.update</*init=*/true>(tSrS);
     Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
     // gather qk gemm res
-    stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
-    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, remainder));
-    cute::copy(scale_o, tScalesScale(_, remainder));
+    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_read_kv.index() % 2));
+    cute::copy(scale_o, tScalesScale(_, smem_pipe_read_kv.index() % 2));
     // r2s fence wgmma
     cutlass::arch::fence_view_async_shared();
     cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
@@ -375,19 +361,18 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
                                           tSrS);
       attention_updater.rescale_o(tOrO);
-      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv_cur.index());
       // last pv gemm
       if (smem_pipe_read_kv_cur.index() == 0) {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
                                              tOrV1(_, _, _, _0{}), tOrO);
       } else if (smem_pipe_read_kv_cur.index() == 1) {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
                                              tOrV2(_, _, _, _0{}), tOrO);
       } else if (smem_pipe_read_kv_cur.index() == 2) {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
                                              tOrV3(_, _, _, _0{}), tOrO);
       } else {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
                                              tOrV4(_, _, _, _0{}), tOrO);
       }
       // wait cur qk gemm
@@ -398,7 +383,7 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
         Tensor tScS = threadMmaQK.partition_C(cS);
 #pragma unroll
         for (int i = 0; i < size(tSrS); ++i) {
-          int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+          int qo_idx = get<0>(tScS(i)) / Ktraits::GROUP_SIZE;
           int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
           if constexpr (!CAUSAL) {  // Just masking based on col
             if (kv_idx >= kv_len) {
@@ -414,10 +399,10 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       // update s (exp(s - m))
       Tensor scale_o = attention_updater.update</*init=*/false>(tSrS);
       Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
+
       // gather qk gemm res
-      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
-      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, remainder));
-      cute::copy(scale_o, tScalesScale(_, remainder));
+      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_read_kv.index() % 2));
+      cute::copy(scale_o, tScalesScale(_, smem_pipe_read_kv.index() % 2));
       // r2s fence wgmma
       cutlass::arch::fence_view_async_shared();
       // make sure tSrS r2s done
@@ -432,18 +417,17 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     ++smem_pipe_read_q;
     // compute last pv
     attention_updater.rescale_o(tOrO);
-    stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
     if (smem_pipe_read_kv.index() == 0) {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                            tOrV1(_, _, _, _0{}), tOrO);
     } else if (smem_pipe_read_kv.index() == 1) {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                            tOrV2(_, _, _, _0{}), tOrO);
     } else if (smem_pipe_read_kv.index() == 2) {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                            tOrV3(_, _, _, _0{}), tOrO);
     } else {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                            tOrV4(_, _, _, _0{}), tOrO);
     }
     scale_o = attention_updater.finalize(tSrS);
@@ -453,8 +437,8 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     ++smem_pipe_read_kv;
     if (chunk_num_this_seq == 1) {
       // norm
-      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
-      cute::copy(scale_o, tScalesScale(_, remainder));
+      cute::copy(scale_o, tScalesScale(_, smem_pipe_read_kv.index() % 2));
+
       cutlass::arch::NamedBarrier::arrive(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2LastSync));
       attention_updater.rescale_o(tOrO);
     }
@@ -463,13 +447,13 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       const int warp_idx = thread_idx / 32;
 #pragma unroll
       for (int w_i = 0; w_i < 2; ++w_i) {
-        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
+        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i;
         const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
 
         if (token_idx < qo_len) {
-          // const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
-          const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
-          const int write_idx = bid * bid_offset + token_group_idx;
+          const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
+          const int bid_offset = mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE;
+          const int write_idx = bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx;
           mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
           mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
         }
@@ -482,21 +466,20 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       consumer_wait(pipeline_kv, smem_pipe_read_kv);
       cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
       // A: tPsP
-      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
-      cute::copy(tScalesScale(_, remainder), scale_o);
+      cute::copy(tScalesScale(_, smem_pipe_read_kv.index() % 2), scale_o);
       // rescale
       attention_updater.rescale_o(tOrO, scale_o);
       if (smem_pipe_read_kv.index() == 0) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                             tOrV1(_, _, _, _0{}), tOrO);
       } else if (smem_pipe_read_kv.index() == 1) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                             tOrV2(_, _, _, _0{}), tOrO);
       } else if (smem_pipe_read_kv.index() == 2) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                             tOrV3(_, _, _, _0{}), tOrO);
       } else {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
                                             tOrV4(_, _, _, _0{}), tOrO);
       }
       pipeline_kv.consumer_release(smem_pipe_read_kv);
@@ -504,459 +487,14 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     }
     if (chunk_num_this_seq == 1) {
       // norm
-      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
       cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2LastSync));
-      cute::copy(tScalesScale(_, remainder), scale_o);
+      cute::copy(tScalesScale(_, smem_pipe_read_kv.index() % 2), scale_o);
       attention_updater.rescale_o(tOrO, scale_o);
     }
   }
   return;
 }
 
-template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ, typename MainloopPipelineQK,
-          typename PipelineState, typename PipelineStateQ, typename PipelineStateQK, typename SharedStorage, typename AttentionUpdater>
-CUTLASS_DEVICE void mma_qk_one_stages(const Params& mainloop_params,
-                                      MainloopPipelineQ pipeline_q,
-                                      PipelineStateQ& smem_pipe_read_q,
-                                      MainloopPipelineQK pipeline_qk,
-                                      PipelineStateQK& smem_pipe_write_qk,
-                                      MainloopPipeline pipeline_kv,
-                                      PipelineState& smem_pipe_read_kv,
-                                      AttentionUpdater& attention_updater,
-                                      const int& thread_idx, 
-                                      const int& bid,
-                                      const int kv_len,
-                                      const int qo_len,
-                                      const int tile_idx,
-                                      const int q_tile_idx,
-                                      SharedStorage& shared_storage) {
-  using DTypeQ = typename Ktraits::DTypeQ;
-  using DTypeKV = typename Ktraits::DTypeKV;
-  using DTypeMD = typename Ktraits::DTypeO; 
-  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
-  using IdType = typename Ktraits::IdType;
-  using TileShape_QKD = typename Ktraits::TileShape_QKD;
-  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
-  using SmemLayoutK = typename Ktraits::SmemLayoutK;
-  using SmemLayoutV = typename Ktraits::SmemLayoutV;
-  using SmemLayoutPSS = typename Ktraits::SmemLayoutPSS;
-  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
-  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
-  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
-
-  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
-  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
-
-  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
-
-  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
-
-  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
-  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-
-  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutPSS{});
-  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
-
-  typename Ktraits::TiledMmaQK tiled_mma_qk;
-  auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
-  Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
-  Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
-  auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
-  auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
-  Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
-  Tensor tScalesScale = s_scale(_, thread_idx, _);
-  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _); // (bsz * draft_token_num * num_head)
-  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
-
-  const int start_len = tile_idx * mainloop_params.chunk_size;
-  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-  const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
-  int kv_tile_idx = end_tile_idx;
-
-  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
-    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
-    pipeline.consumer_wait(smem_pipe_read, barrier_token);
-  };
-  Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
-
-  // DO QK gemm
-  auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
-  // wait q
-  consumer_wait(pipeline_q, smem_pipe_read_q);
-  constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
-  int masking_step = n_masking_steps;
-  for (; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
-    // Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
-    consumer_wait(pipeline_kv, smem_pipe_read_kv);
-    // gemm next qk
-    gemm</*init=*/true, /*wg_wait=*/0>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                        tSrS);
-    // release last kv
-    pipeline_kv.consumer_release(smem_pipe_read_kv);
-    ++smem_pipe_read_kv;
-    // mask
-    if (masking_step > 0) {
-      Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
-      Tensor tScS = threadMmaQK.partition_C(cS);
-#pragma unroll
-      for (int i = 0; i < size(tSrS); ++i) {
-        int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
-        int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
-        if constexpr (!CAUSAL) {  // Just masking based on col
-          if (kv_idx >= kv_len) {
-            tSrS(i) = AttentionUpdater::fill_value;
-          }
-        } else {
-          if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
-            tSrS(i) = AttentionUpdater::fill_value;
-          }
-        }
-      }
-    }
-    // update s (exp(s - m))
-    Tensor scale_o = (kv_tile_idx == end_tile_idx ? attention_updater.update<true>(tSrS) : attention_updater.update<false>(tSrS));
-    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
-
-    // r2s
-    pipeline_qk.producer_acquire(smem_pipe_write_qk);
-    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
-    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-    // r2s fence wgmma
-    cutlass::arch::fence_view_async_shared();
-    pipeline_qk.producer_commit(smem_pipe_write_qk);
-    ++smem_pipe_write_qk;
-  }
-  // wg1 release q
-  pipeline_q.consumer_release(smem_pipe_read_q);
-  ++smem_pipe_read_q;
-
-  // update last softmax scales
-  Tensor scale_o = attention_updater.finalize(tSrS);
-  if (chunk_num_this_seq == 1) {
-    // norm
-    pipeline_qk.producer_acquire(smem_pipe_write_qk);
-    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-    pipeline_qk.producer_commit(smem_pipe_write_qk);
-    ++smem_pipe_write_qk;
-  }
-
-  // WG1 write m,d back to gmem
-  if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
-    const int warp_idx = thread_idx / 32;
-#pragma unroll
-    for (int w_i = 0; w_i < 2; ++w_i) {
-      const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
-      const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
-      if (token_idx < qo_len) {
-        // const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
-        const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
-        const int write_idx = bid * bid_offset + token_group_idx;
-        mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
-        mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
-      }
-    }
-  }
-  return;
-}
-
-template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ, typename MainloopPipelineQK,
-          typename PipelineState, typename PipelineStateQ, typename PipelineStateQK, typename SharedStorage, typename AttentionUpdater>
-CUTLASS_DEVICE void mma_qk_two_stages(const Params& mainloop_params,
-                                      MainloopPipelineQ pipeline_q,
-                                      PipelineStateQ& smem_pipe_read_q,
-                                      MainloopPipelineQK pipeline_qk,
-                                      PipelineStateQK& smem_pipe_write_qk,
-                                      MainloopPipeline pipeline_kv,
-                                      PipelineState& smem_pipe_read_kv,
-                                      AttentionUpdater& attention_updater,
-                                      const int& thread_idx, 
-                                      const int& bid,
-                                      const int kv_len,
-                                      const int qo_len,
-                                      const int tile_idx,
-                                      const int q_tile_idx,
-                                      SharedStorage& shared_storage) {
-  using DTypeQ = typename Ktraits::DTypeQ;
-  using DTypeKV = typename Ktraits::DTypeKV;
-  using DTypeMD = typename Ktraits::DTypeO; 
-  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
-  using IdType = typename Ktraits::IdType;
-  using TileShape_QKD = typename Ktraits::TileShape_QKD;
-  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
-  using SmemLayoutK = typename Ktraits::SmemLayoutK;
-  using SmemLayoutV = typename Ktraits::SmemLayoutV;
-  using SmemLayoutPSS = typename Ktraits::SmemLayoutPSS;
-  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
-  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
-  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
-
-  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
-  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
-
-  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
-
-  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
-
-  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
-  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-
-  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutPSS{});
-  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
-
-  typename Ktraits::TiledMmaQK tiled_mma_qk;
-  auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
-  Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
-  Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
-  auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
-  auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
-  Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
-  Tensor tScalesScale = s_scale(_, thread_idx, _);
-  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _); // (bsz * draft_token_num * num_head)
-  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
-
-  const int start_len = tile_idx * mainloop_params.chunk_size;
-  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-  int kv_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
-
-  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
-    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
-    pipeline.consumer_wait(smem_pipe_read, barrier_token);
-  };
-
-  Tensor tSrS1 = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
-  Tensor tSrS2 = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
-
-  // DO QK gemm
-  auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
-  // wait q
-  consumer_wait(pipeline_q, smem_pipe_read_q);
-  // wait k
-  consumer_wait(pipeline_kv, smem_pipe_read_kv);
-  // first qk gemm
-  if (smem_pipe_write_qk.index() == 0) {
-    gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                        tSrS1);
-  } else {
-    gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                        tSrS2);
-  }
-  constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
-  int masking_step = n_masking_steps;
-  --kv_tile_idx;
-  for (; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
-    Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
-    PipelineState smem_pipe_read_kv_cur = smem_pipe_read_kv;
-    ++smem_pipe_read_kv;
-    // wait next kv
-    consumer_wait(pipeline_kv, smem_pipe_read_kv);
-    // gemm next qk
-    if (smem_pipe_write_qk.index() == 1) {
-      gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                          tSrS1);
-    } else {
-      gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                          tSrS2);
-    }
-    // wait last qk done
-    warpgroup_wait<1>();
-    // release last kv
-    pipeline_kv.consumer_release(smem_pipe_read_kv_cur);
-    // mask now
-    if (masking_step > 0) {
-      Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
-      Tensor tScS = threadMmaQK.partition_C(cS);
-#pragma unroll
-      for (int i = 0; i < size(tSrS1); ++i) {
-        int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
-        int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
-        if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
-          if (smem_pipe_write_qk.index() == 0) {
-            tSrS1(i) = AttentionUpdater::fill_value;
-          } else {
-            tSrS2(i) = AttentionUpdater::fill_value;
-          }
-        }
-      }
-    }
-    // update s (exp(s - m))
-    if (smem_pipe_write_qk.index() == 0) {
-      Tensor scale_o = attention_updater.update</*init=*/false>(tSrS1);
-      Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS1));
-      // r2s
-      pipeline_qk.producer_acquire(smem_pipe_write_qk);
-      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
-      cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-    } else {
-      Tensor scale_o = attention_updater.update</*init=*/false>(tSrS2);
-      Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS2));
-      // r2s
-      pipeline_qk.producer_acquire(smem_pipe_write_qk);
-      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
-      cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-    }
-    // r2s fence wgmma
-    cutlass::arch::fence_view_async_shared();
-    pipeline_qk.producer_commit(smem_pipe_write_qk);
-    ++smem_pipe_write_qk;
-  }
-  // wg1 release q
-  pipeline_q.consumer_release(smem_pipe_read_q);
-  ++smem_pipe_read_q;
-  // wait last qk
-  warpgroup_wait<0>();
-  // wg1 release kv
-  pipeline_kv.consumer_release(smem_pipe_read_kv);
-  ++smem_pipe_read_kv;
-  // update s (exp(s - m)) last
-  if (smem_pipe_write_qk.index() == 0) {
-    Tensor scale_o = attention_updater.update</*init=*/false>(tSrS1);
-    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS1));
-    // r2s
-    pipeline_qk.producer_acquire(smem_pipe_write_qk);
-    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
-    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-  } else {
-    Tensor scale_o = attention_updater.update</*init=*/false>(tSrS2);
-    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS2));
-    // r2s
-    pipeline_qk.producer_acquire(smem_pipe_write_qk);
-    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
-    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-  }
-  // r2s fence wgmma
-  cutlass::arch::fence_view_async_shared();
-  pipeline_qk.producer_commit(smem_pipe_write_qk);
-  ++smem_pipe_write_qk;
-
-  // update last softmax scales
-  Tensor scale_o = attention_updater.finalize(tSrS1);
-  if (chunk_num_this_seq == 1) {
-    // norm
-    pipeline_qk.producer_acquire(smem_pipe_write_qk);
-    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
-    pipeline_qk.producer_commit(smem_pipe_write_qk);
-    ++smem_pipe_write_qk;
-  }
-
-  // WG1 write m,d back to gmem
-  if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
-    const int warp_idx = thread_idx / 32;
-#pragma unroll
-    for (int w_i = 0; w_i < 2; ++w_i) {
-      const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
-      const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
-      if (token_idx < qo_len) {
-        const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
-        const int write_idx = bid * bid_offset + token_group_idx;
-        mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
-        mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
-      }
-    }
-  }
-  return;
-}
-
-template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQK,
-          typename PipelineState, typename PipelineStateQK, typename SharedStorage, typename FrgTensorO, typename AttentionUpdater>
-CUTLASS_DEVICE void mma_pv_one_stages(const Params& mainloop_params,
-                                      MainloopPipelineQK pipeline_qk,
-                                      PipelineStateQK& smem_pipe_read_qk,
-                                      MainloopPipeline pipeline_kv,
-                                      PipelineState& smem_pipe_read_kv,
-                                      AttentionUpdater& attention_updater,
-                                      FrgTensorO& tOrO, 
-                                      const int thread_idx, 
-                                      const int bid,
-                                      const int kv_len,
-                                      const int qo_len,
-                                      const int tile_idx,
-                                      SharedStorage& shared_storage) {
-  using DTypeQ = typename Ktraits::DTypeQ;
-  using DTypeKV = typename Ktraits::DTypeKV;
-  using DTypeMD = typename Ktraits::DTypeO;
-  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
-  using IdType = typename Ktraits::IdType;
-  using TileShape_QKD = typename Ktraits::TileShape_QKD;
-  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
-  using SmemLayoutK = typename Ktraits::SmemLayoutK;
-  using SmemLayoutV = typename Ktraits::SmemLayoutV;
-  using SmemLayoutPSS = typename Ktraits::SmemLayoutPSS;
-  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
-  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
-  using SmemLayoutVtOneStage = typename Ktraits::SmemLayoutVtOneStage;
-  static_assert(is_rmem<FrgTensorO>::value, "O tensor must be rmem resident.");
-
-  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
-  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
-
-  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
-  // todo: split rope/norope for sVt, not use s1~s4.
-  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
-  Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-  Tensor sVt_s3 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 2 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-  Tensor sVt_s4 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 3 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-
-  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutPSS{});
-  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
-  Tensor tScalesScale = s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup, _);
-
-  typename Ktraits::TiledMmaPVSS tiled_mma_pv_ss;
-  auto threadMmaPVSS = tiled_mma_pv_ss.get_thread_slice(thread_idx);
-  Tensor tOrV1 = threadMmaPVSS.partition_fragment_B(sVt_s1);
-  Tensor tOrV2 = threadMmaPVSS.partition_fragment_B(sVt_s2);
-  Tensor tOrV3 = threadMmaPVSS.partition_fragment_B(sVt_s3);
-  Tensor tOrV4 = threadMmaPVSS.partition_fragment_B(sVt_s4);
-  Tensor tOrP_S = threadMmaPVSS.partition_fragment_A(sPSS);
-  const int start_len = tile_idx * mainloop_params.chunk_size;
-  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-  const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
-  int kv_tile_idx = end_tile_idx;
-
-  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
-    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
-    pipeline.consumer_wait(smem_pipe_read, barrier_token);
-  };
-
-  Tensor scale_o = make_tensor<DTypeQKAccum>(Shape<_2>{});
-  for (; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
-    // wait v
-    consumer_wait(pipeline_kv, smem_pipe_read_kv);
-    // wait p
-    consumer_wait(pipeline_qk, smem_pipe_read_qk);
-    // A: tPsP
-    cute::copy(tScalesScale(_, smem_pipe_read_qk.index()), scale_o);
-    // rescale
-    attention_updater.rescale_o(tOrO, scale_o);
-    // gemm pv
-    if (smem_pipe_read_kv.index() == 0) {
-      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
-                                          tOrV1(_, _, _, _0{}), tOrO);
-    } else if (smem_pipe_read_kv.index() == 1) {
-      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
-                                          tOrV2(_, _, _, _0{}), tOrO);
-    } else if (smem_pipe_read_kv.index() == 2) {
-      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
-                                          tOrV3(_, _, _, _0{}), tOrO);
-    } else if (smem_pipe_read_kv.index() == 3) {
-      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
-                                          tOrV4(_, _, _, _0{}), tOrO);
-    }
-    pipeline_qk.consumer_release(smem_pipe_read_qk);
-    ++smem_pipe_read_qk;
-    pipeline_kv.consumer_release(smem_pipe_read_kv);
-    ++smem_pipe_read_kv;
-  }
-  if (chunk_num_this_seq == 1) {
-    // norm
-    consumer_wait(pipeline_qk, smem_pipe_read_qk);
-    cute::copy(tScalesScale(_, smem_pipe_read_qk.index()), scale_o);
-    pipeline_qk.consumer_release(smem_pipe_read_qk);
-    ++smem_pipe_read_qk;
-    attention_updater.rescale_o(tOrO, scale_o);
-  }
-  return;
-}
-
 }  // namespace mla_attn
 
 #endif  // ATTENTION_HOPPER_MAINLOOP_MMA_CUH_

csrc/gpu/mla_attn/mla_hopper.cuh

@@ -19,8 +19,6 @@
  * Modified by the FlashInfer team.
  */
 
-#pragma once
-
 #ifndef ATTENTION_HOPPER_PREFILL_SM90_CUH_
 #define ATTENTION_HOPPER_PREFILL_SM90_CUH_
 
@@ -65,8 +63,8 @@ struct Params {
     alignas(16) DTypeKV *KV; // [max_block_num, block_size, dim_head]
     alignas(16) DTypeO *O; // [token_num, head_num, dim_head]
     alignas(16) DTypeO *O_tmp; // [num_chunks, bsz, head_num, dim_head]
-    alignas(16) float *m; // [num_chunks, bsz * draft_total_token_num * head_num]
-    alignas(16) float *d; // [num_chunks, bsz * draft_total_token_num * head_num]
+    alignas(16) float *m; // [num_chunks, bsz * max_draft_token_num * head_num]
+    alignas(16) float *d; // [num_chunks, bsz * max_draft_token_num * head_num]
 
     alignas(16) IdType *block_tables;
     alignas(16) IdType *seq_lens_this_time;
@@ -76,8 +74,7 @@ struct Params {
     alignas(16) IdType *padding_offsets;
 
     alignas(16) IdType *batch_ids;
-    alignas(16) IdType *q_tile_ids_per_batch = nullptr;
-    alignas(16) IdType *tile_ids_per_batch = nullptr;
+    alignas(16) IdType *tile_ids_per_batch;
     alignas(16) IdType *num_blocks_x;
 
 
@@ -99,10 +96,9 @@ struct Params {
     int qk_head_dim;
     int vo_head_dim;
     int block_size;
-    int draft_total_token_num;
+    int max_draft_token_num;
     int chunk_size;
     int chunk_num;
-    int num_blocks_x_int;
 
     float sm_scale;
 };
@@ -117,14 +113,12 @@ 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;                            \
   }
-template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, bool USE_REG_EALLOC = false>
+
+template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, int SM_COUNT = 132>
 __global__ void __launch_bounds__(Ktraits::NUM_WARPS * cutlass::NumThreadsPerWarp, 1)
 MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
                      typename CollectiveMainloop::Params const mainloop_params,
@@ -204,47 +198,33 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
   
   if (warp_group_idx == 0) {
     // producer
-    if (USE_REG_EALLOC) {
-      cutlass::arch::warpgroup_reg_dealloc<88>();
-    }
+    cutlass::arch::warpgroup_reg_dealloc<72>();
     const uint32_t warp_idx_in_warpgroup = __shfl_sync(0xffffffff, warp_idx % 4, 0);
     
     PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state<MainloopPipelineQ>();
     PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
-    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 ? 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];
-    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
-    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+    for (int i = blockIdx.x; i < num_blocks_x; i += SM_COUNT) {
+      const int bid = mainloop_params.batch_ids[i];
+      const int tile_id = mainloop_params.tile_ids_per_batch[i];
+      const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+      const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
+      const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+      const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_THREADS,
+                                        /*id=*/static_cast<int>(NamedBarriers::kWG0WG1WG2Sync));
 
-    // load Q
-    collective_mainloop.load_q(
-        mainloop_params,
-        pipeline_q,
-        smem_pipe_write_q,
-        shared_storage,
-        threadIdx.x,
-        bid,
-        q_tile_id);
-
-    if constexpr (!use_tma_load_kv) {
-      // load kv
-      collective_mainloop.load_kv(
+      // load Q
+      collective_mainloop.load_q(
           mainloop_params,
-          pipeline_kv,
-          smem_pipe_write_kv,
+          pipeline_q,
+          smem_pipe_write_q,
           shared_storage,
-          bid,
-          seq_len_decoder_now,
-          tile_id
-      );
-    } else {
-      if (warp_idx_in_warpgroup == 0) {
-        // load kv tma
-        collective_mainloop.load_kv_tma(
+          threadIdx.x,
+          bid);
+
+      if constexpr (!use_tma_load_kv) {
+        // load kv
+        collective_mainloop.load_kv(
             mainloop_params,
             pipeline_kv,
             smem_pipe_write_kv,
@@ -253,13 +233,24 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
             seq_len_decoder_now,
             tile_id
         );
+      } else {
+        if (warp_idx_in_warpgroup == 0) {
+          // load kv tma
+          collective_mainloop.load_kv_tma(
+              mainloop_params,
+              pipeline_kv,
+              smem_pipe_write_kv,
+              shared_storage,
+              bid,
+              seq_len_decoder_now,
+              tile_id
+          );
+        }
       }
     }
   } else {
     // consumer
-    if (USE_REG_EALLOC) {
-      cutlass::arch::warpgroup_reg_alloc<208>();
-    }
+    cutlass::arch::warpgroup_reg_alloc<216>(); 
     PipelineStateQ smem_pipe_read_q;
     PipelineState smem_pipe_read_kv;
 
@@ -267,73 +258,72 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
     Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_PDV{}));
 
     auto attention_updater = OnlineSoftmax<2 * size<1>(tOrO), /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
-    const int block_id = blockIdx.x;
-    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 ? 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];
-    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
-    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+    for (int i = blockIdx.x; i < num_blocks_x; i += SM_COUNT) {
+      clear(tOrO);
+      clear(attention_updater.scores_scale);
+      const int bid = mainloop_params.batch_ids[i];
+      const int tile_id = mainloop_params.tile_ids_per_batch[i];
+      const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+      const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
+      const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+      const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_THREADS,
+                                        /*id=*/static_cast<int>(NamedBarriers::kWG0WG1WG2Sync));
 
-    if constexpr (BLOCK_SHAPE_KV == 64) {
-      mma_f16<Ktraits, CAUSAL>(
-        mainloop_params, 
-        pipeline_q, 
-        smem_pipe_read_q,
-        pipeline_kv, 
-        smem_pipe_read_kv,
-        tOrO, 
-        attention_updater,
-        threadIdx.x - NUM_COPY_THREADS,
-        bid,
-        seq_len_decoder_now,
-        seq_len_now,
-        tile_id,
-        q_tile_id,
-        shared_storage);
-    } else if (BLOCK_SHAPE_KV == 32) {
-      mma_f16_two_stages<Ktraits, CAUSAL>(
-        mainloop_params, 
-        pipeline_q, 
-        smem_pipe_read_q,
-        pipeline_kv, 
-        smem_pipe_read_kv,
-        tOrO, 
-        attention_updater,
-        threadIdx.x - NUM_COPY_THREADS,
-        bid,
-        seq_len_decoder_now,
-        seq_len_now,
-        tile_id,
-        q_tile_id,
-        shared_storage);
+      if constexpr (BLOCK_SHAPE_KV == 64) {
+        mma_f16<Ktraits, CAUSAL>(
+          mainloop_params, 
+          pipeline_q, 
+          smem_pipe_read_q,
+          pipeline_kv, 
+          smem_pipe_read_kv,
+          tOrO, 
+          attention_updater,
+          threadIdx.x - NUM_COPY_THREADS,
+          bid,
+          seq_len_decoder_now,
+          seq_len_now,
+          tile_id,
+          shared_storage);
+      } else if (BLOCK_SHAPE_KV == 32) {
+        mma_f16_two_stages<Ktraits, CAUSAL>(
+          mainloop_params, 
+          pipeline_q, 
+          smem_pipe_read_q,
+          pipeline_kv, 
+          smem_pipe_read_kv,
+          tOrO, 
+          attention_updater,
+          threadIdx.x - NUM_COPY_THREADS,
+          bid,
+          seq_len_decoder_now,
+          seq_len_now,
+          tile_id,
+          shared_storage);
+      }
+
+      collective_epilogue.store(
+          epilogue_params, 
+          tOrO, 
+          attention_updater.get_lse(),
+          shared_storage,
+          tiled_mma_pv, 
+          threadIdx.x - NUM_COPY_THREADS,
+          bid,
+          mainloop_params.bsz,
+          seq_len_now,
+          start_token_idx,
+          tile_id,
+          seq_len_decoder_now,
+          mainloop_params.chunk_size,
+          mainloop_params.max_draft_token_num,
+          mainloop_params.o_stride_bsz);
     }
-
-    collective_epilogue.store(
-        epilogue_params, 
-        tOrO, 
-        attention_updater.get_lse(),
-        shared_storage,
-        tiled_mma_pv, 
-        threadIdx.x - NUM_COPY_THREADS,
-        bid,
-        mainloop_params.bsz,
-        seq_len_now,
-        start_token_idx,
-        tile_id,
-        q_tile_id,
-        seq_len_decoder_now,
-        mainloop_params.chunk_size,
-        mainloop_params.draft_total_token_num,
-        mainloop_params.o_stride_head_num);
   }
 }
 
 
-template <typename KernelTraits, bool CAUSAL, typename Params, bool USE_REG_EALLOC=false>
+template <typename KernelTraits, bool CAUSAL, typename Params>
 cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
                                                            cudaStream_t stream) {
   using DTypeQ = typename KernelTraits::DTypeQ;
@@ -349,7 +339,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
   typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments({
       make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.qk_head_dim), make_stride(params.qk_head_dim, _1{})), // layout q
       make_layout(make_shape(params.block_size, params.qk_head_dim, params.max_block_num), make_stride(params.qk_head_dim, _1{}, params.block_size * params.qk_head_dim)),
-      make_layout(make_shape(params.chunk_num, params.bsz * params.draft_total_token_num * params.q_num_head), make_stride(params.bsz * params.draft_total_token_num * params.q_num_head, _1{})),
+      make_layout(make_shape(params.chunk_num, params.bsz * params.max_draft_token_num * params.q_num_head), make_stride(params.bsz * params.max_draft_token_num * params.q_num_head, _1{})),
       params.Q,
       params.KV,
       params.m,
@@ -361,7 +351,6 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
       params.cumsum_q_seqlens,
       params.batch_ids,
       params.tile_ids_per_batch,
-      params.q_tile_ids_per_batch,
       params.num_blocks_x,
       params.sm_scale,
       params.bsz,
@@ -375,7 +364,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
       params.o_stride_head_num,
       params.chunk_size,
       params.chunk_num,
-      params.draft_total_token_num
+      params.max_draft_token_num
   });
   typename CollectiveEpilogue::Params epilogue_params = CollectiveEpilogue::to_underlying_arguments_ntma({
       params.O,
@@ -386,7 +375,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
 
   // Get the ptr to kernel function.
   auto kernel =
-      MLAWithKVCacheKernel<CollectiveMainloop, CollectiveEpilogue, KernelTraits, CAUSAL, USE_REG_EALLOC>;
+      MLAWithKVCacheKernel<CollectiveMainloop, CollectiveEpilogue, KernelTraits, CAUSAL, 132>;
   int smem_size = sizeof(typename KernelTraits::SharedStorage);
   cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
   int device;
@@ -396,8 +385,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
   int act_blocks_per_sm;
   cudaOccupancyMaxActiveBlocksPerMultiprocessor(
       &act_blocks_per_sm, kernel, KernelTraits::NUM_WARPS * 32, smem_size);
-
-  dim3 grid_dims = {params.num_blocks_x_int, 1, 1}; // todo: split kv
+  dim3 grid_dims = {multiprocessor_count, 1, 1}; // todo: split kv
   static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32;
   dim3 block_dims(ctaSize, 1, 1);
   kernel<<<grid_dims, block_dims, smem_size, stream>>>(
@@ -426,53 +414,31 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
       params.vo_head_dim,
       params.token_num,
       params.bsz,
-      params.draft_total_token_num
+      params.max_draft_token_num
     );
   }
   return cudaSuccess;
 }
 
-template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params, bool USE_REG_EALLOC=false>
+template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params>
 cudaError_t BatchMLAWithPagedKVCacheDispatched(Params& params, cudaStream_t stream) {
   constexpr bool CAUSAL = true;
   if constexpr (HEAD_DIM_QK == 576) {
-    if (params.block_size == 32) {
-      DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
-        BatchMLAWithPagedKVCacheKernelTraitsDispatched<
-            AttentionKernelTraits</*USE_TMA_LOAD_KV=*/true, 
-                                  HEAD_DIM_QK, 
-                                  HEAD_DIM_VO, 
-                                  GROUP_SIZE,
-                                  /*BLOCK_SHAPE_Q_=*/64,
-                                  /*BLOCK_SHAPE_KV_=*/32,
-                                  /*NUM_STAGES_=*/4, 
-                                  typename Params::DTypeQ,
-                                  typename Params::DTypeKV, 
-                                  typename Params::DTypeO,
-                                  typename Params::IdType,
-                                  NV_TYPE>,
-            CAUSAL,
-            Params,
-            USE_REG_EALLOC>(params, stream);)
-    } else if (params.block_size == 64) {
-      DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
-        BatchMLAWithPagedKVCacheKernelTraitsDispatched<
-            AttentionKernelTraits</*USE_TMA_LOAD_KV=*/true, 
-                                  HEAD_DIM_QK, 
-                                  HEAD_DIM_VO, 
-                                  GROUP_SIZE,
-                                  /*BLOCK_SHAPE_Q_=*/64,
-                                  /*BLOCK_SHAPE_KV_=*/64,
-                                  /*NUM_STAGES_=*/2, 
-                                  typename Params::DTypeQ,
-                                  typename Params::DTypeKV, 
-                                  typename Params::DTypeO,
-                                  typename Params::IdType,
-                                  NV_TYPE>,
-            CAUSAL,
-            Params,
-            USE_REG_EALLOC>(params, stream);)
-    }
+    DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
+      BatchMLAWithPagedKVCacheKernelTraitsDispatched<
+          AttentionKernelTraits</*USE_TMA_LOAD_KV=*/false, 
+                                HEAD_DIM_QK, 
+                                HEAD_DIM_VO, 
+                                GROUP_SIZE,
+                                /*BLOCK_SHAPE_Q_=*/64,
+                                /*BLOCK_SHAPE_KV_=*/64,
+                                /*NUM_STAGES_=*/2, 
+                                typename Params::DTypeQ,
+                                typename Params::DTypeKV, 
+                                typename Params::DTypeO,
+                                typename Params::IdType, 
+                                NV_TYPE>,
+          CAUSAL>(params, stream);)
   } else {
     return cudaErrorNotSupported;
   }

csrc/gpu/mla_attn/mla_hopper_wg4.cuh

@@ -1,426 +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.
-
-#pragma once
-
-#ifndef ATTENTION_HOPPER_PREFILL_WG4_SM90_CUH_
-#define ATTENTION_HOPPER_PREFILL_WG4_SM90_CUH_
-
-#include <cuda.h>
-#include <cuda_device_runtime_api.h>
-#include <cutlass/arch/reg_reconfig.h>
-#include <cutlass/array.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/numeric_conversion.h>
-#include <cutlass/numeric_types.h>
-
-#include <type_traits>
-#include <vector>
-
-#include "attention_updater.cuh"
-#include "cute/tensor.hpp"
-#include "cutlass/pipeline/pipeline.hpp"
-#include "epilogue.cuh"
-#include "helper.h"
-#include "kernel_traits_wg4.cuh"
-#include "mla_hopper.cuh"
-#include "mainloop_mma.cuh"
-#include "mainloop_load.cuh"
-#include "utils.cuh"
-
-#ifdef DEBUG_MLA
-#undef DEBUG_MLA
-#endif
-// #define DEBUG_MLA
-
-namespace mla_attn {
-
-using namespace cute;
-
-template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, bool USE_REG_EALLOC=false, bool USE_QK_TWO_STAGE=false>
-__global__ void __launch_bounds__(Ktraits::NUM_WARPS * cutlass::NumThreadsPerWarp, 1)
-MLAWithKVCacheWG4Kernel(CUTE_GRID_CONSTANT
-                        typename CollectiveMainloop::Params const mainloop_params,
-                        CUTE_GRID_CONSTANT
-                        typename CollectiveEpilogue::Params const epilogue_params) {
-  int warp_group_idx = cutlass::canonical_warp_group_idx();
-  using DTypeQ = typename Ktraits::DTypeQ;
-  using DTypeKV = typename Ktraits::DTypeKV;
-  using DTypeO = typename Ktraits::DTypeO;
-  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
-  using TileShape_QKD = typename Ktraits::TileShape_QKD;
-  using TileShape_PDV = typename Ktraits::TileShape_PDV;
-
-  static constexpr int NUM_MMA_THREADS_QK = Ktraits::NUM_MMA_THREADS_QK;
-  static constexpr int NUM_MMA_THREADS_PV = Ktraits::NUM_MMA_THREADS_PV;
-  static constexpr int NUM_COPY_THREADS = Ktraits::NUM_PRODUCER_THREADS;
-  static constexpr int BLOCK_SHAPE_Q = Ktraits::BLOCK_SHAPE_Q;
-  static constexpr int BLOCK_SHAPE_KV = Ktraits::BLOCK_SHAPE_KV;
-  const int num_blocks_x = mainloop_params.num_blocks_x[0];
-
-  static constexpr bool use_tma_load_kv = CollectiveMainloop::USE_TMA_LOAD_KV;
-
-  using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
-  using PipelineParams = typename MainloopPipeline::Params;
-  using PipelineState = typename MainloopPipeline::PipelineState;
-
-  using MainloopPipelineQ = typename CollectiveMainloop::MainloopPipelineQ;
-  using PipelineParamsQ = typename MainloopPipelineQ::Params;
-  using PipelineStateQ = typename MainloopPipelineQ::PipelineState;
-
-  using MainloopPipelineQK = typename Ktraits::MainloopPipelineQK;
-  using PipelineParamsQK = typename MainloopPipelineQK::Params;
-  using PipelineStateQK = typename MainloopPipelineQK::PipelineState;
-
-  extern __shared__ char shared_memory[];
-  auto& shared_storage = *reinterpret_cast<typename Ktraits::SharedStorage*>(shared_memory);
-
-  PipelineParams pipeline_params;
-  pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::Producer
-                                             : MainloopPipeline::ThreadCategory::Consumer;
-  if constexpr (use_tma_load_kv) {
-    int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
-    pipeline_params.is_leader = warp_group_thread_idx == 0;
-    pipeline_params.num_consumers = NUM_MMA_THREADS_QK + NUM_MMA_THREADS_PV;
-  } else {
-    pipeline_params.producer_arv_count = NUM_COPY_THREADS;
-    pipeline_params.consumer_arv_count = NUM_MMA_THREADS_QK + NUM_MMA_THREADS_PV;
-  }
-
-  PipelineParamsQ pipeline_params_q;
-  if (warp_group_idx == 0) {
-    pipeline_params_q.role = MainloopPipelineQ::ThreadCategory::Producer;
-  } else if (warp_group_idx == 1) {
-    pipeline_params_q.role = MainloopPipelineQ::ThreadCategory::Consumer;
-  } else {
-    pipeline_params_q.role = MainloopPipelineQ::ThreadCategory::NonParticipant;
-  }
-  pipeline_params_q.producer_arv_count = NUM_COPY_THREADS;
-  pipeline_params_q.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; // just one wg qk
-  
-
-  MainloopPipelineQ pipeline_q(shared_storage.pipeline_q, pipeline_params_q);
-
-  PipelineParamsQK pipeline_params_qk;
-  if (warp_group_idx == 0) {
-    pipeline_params_qk.role = MainloopPipelineQK::ThreadCategory::NonParticipant;
-  } else if (warp_group_idx == 1) {
-    pipeline_params_qk.role = MainloopPipelineQK::ThreadCategory::Producer;
-  } else {
-    pipeline_params_qk.role = MainloopPipelineQK::ThreadCategory::Consumer;
-  }
-  pipeline_params_qk.producer_arv_count = NUM_MMA_THREADS_QK;
-  pipeline_params_qk.consumer_arv_count = NUM_MMA_THREADS_PV;
-  MainloopPipelineQK pipeline_qk(shared_storage.pipeline_qk, pipeline_params_qk);
-
-  MainloopPipeline pipeline_kv = [&] {
-    if constexpr (use_tma_load_kv) {
-      pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesKV;
-      return MainloopPipeline(shared_storage.pipeline_kv, pipeline_params,
-                              /*cluster_shape=*/Shape<_1, _1, _1>{});
-    } else {
-      return MainloopPipeline(shared_storage.pipeline_kv, pipeline_params);
-    }
-  }();
-  __syncthreads();
-
-  CollectiveMainloop collective_mainloop;
-  CollectiveEpilogue collective_epilogue;
-  
-  if (warp_group_idx == 0) {
-    // producer
-    if constexpr (USE_REG_EALLOC) {
-        cutlass::arch::warpgroup_reg_dealloc<64>();
-    }
-    int const warp_idx = cutlass::canonical_warp_idx_sync();
-    const uint32_t warp_idx_in_warpgroup = __shfl_sync(0xffffffff, warp_idx % 4, 0);
-    PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state<MainloopPipelineQ>();
-    PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
-    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 ? 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;
-    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
-    // load Q
-    collective_mainloop.load_q(
-        mainloop_params,
-        pipeline_q,
-        smem_pipe_write_q,
-        shared_storage,
-        threadIdx.x,
-        bid,
-        q_tile_id);
-
-    if constexpr (!use_tma_load_kv) {
-      // load kv
-      collective_mainloop.load_kv(
-          mainloop_params,
-          pipeline_kv,
-          smem_pipe_write_kv,
-          shared_storage,
-          bid,
-          seq_len_decoder_now,
-          tile_id
-      );
-    } else {
-      if (warp_idx_in_warpgroup == 0) {
-        // load kv tma
-        collective_mainloop.load_kv_tma(
-            mainloop_params,
-            pipeline_kv,
-            smem_pipe_write_kv,
-            shared_storage,
-            bid,
-            seq_len_decoder_now,
-            tile_id
-        );
-      }
-    }
-  } else if (warp_group_idx == 1) {
-    // mma qk
-    if constexpr (USE_REG_EALLOC) {
-        cutlass::arch::warpgroup_reg_dealloc<64>(); 
-    }
-
-    auto attention_updater = OnlineSoftmax<2 * BLOCK_SHAPE_Q / 64, /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
-    PipelineStateQ smem_pipe_read_q;
-    PipelineState smem_pipe_read_kv;
-    PipelineStateQK smem_pipe_write_qk = cutlass::make_producer_start_state<MainloopPipelineQK>();
-    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 ? 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;
-    if constexpr (USE_QK_TWO_STAGE) {
-      mma_qk_two_stages<Ktraits, CAUSAL>(
-          mainloop_params, 
-          pipeline_q, 
-          smem_pipe_read_q,
-          pipeline_qk,
-          smem_pipe_write_qk,
-          pipeline_kv, 
-          smem_pipe_read_kv,
-          attention_updater, 
-          threadIdx.x - NUM_COPY_THREADS,
-          bid,
-          seq_len_decoder_now,
-          seq_len_now,
-          tile_id,
-          q_tile_id,
-          shared_storage);
-    } else {
-      mma_qk_one_stages<Ktraits, CAUSAL>(
-          mainloop_params, 
-          pipeline_q, 
-          smem_pipe_read_q,
-          pipeline_qk,
-          smem_pipe_write_qk,
-          pipeline_kv, 
-          smem_pipe_read_kv,
-          attention_updater, 
-          threadIdx.x - NUM_COPY_THREADS,
-          bid,
-          seq_len_decoder_now,
-          seq_len_now,
-          tile_id,
-          q_tile_id,
-          shared_storage);
-    }
-  } else {
-    // mm pv and store 160 reg
-    if constexpr (USE_REG_EALLOC) {
-        cutlass::arch::warpgroup_reg_alloc<184>();
-    }
-    typename Ktraits::TiledMmaPVSS tiled_mma_pv;
-    Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_PDV{}));
-    auto attention_updater = OnlineSoftmax<2 * BLOCK_SHAPE_Q / 64, /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
-    PipelineState smem_pipe_read_kv;
-    PipelineStateQK smem_pipe_read_qk;
-    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 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;
-    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
-    clear(tOrO);
-    clear(attention_updater.scores_scale);
-    mma_pv_one_stages<Ktraits, CAUSAL>(
-        mainloop_params, 
-        pipeline_qk, 
-        smem_pipe_read_qk,
-        pipeline_kv, 
-        smem_pipe_read_kv,
-        attention_updater,
-        tOrO,
-        threadIdx.x - NUM_COPY_THREADS - NUM_MMA_THREADS_QK,
-        bid,
-        seq_len_decoder_now,
-        seq_len_now,
-        tile_id,
-        shared_storage);
-    collective_epilogue.store(
-      epilogue_params, 
-      tOrO, 
-      attention_updater.get_lse(),
-      shared_storage,
-      tiled_mma_pv, 
-      threadIdx.x - NUM_COPY_THREADS - NUM_MMA_THREADS_QK,
-      bid,
-      mainloop_params.bsz,
-      seq_len_now,
-      start_token_idx,
-      tile_id,
-      q_tile_id,
-      seq_len_decoder_now,
-      mainloop_params.chunk_size,
-      mainloop_params.draft_total_token_num,
-      mainloop_params.o_stride_head_num);
-  }
-}
-
-template <typename KernelTraits, bool CAUSAL, typename Params, bool USE_REG_EALLOC=false, bool USE_QK_TWO_STAGE=false>
-cudaError_t BatchMLAWithPagedKVCacheWG4KernelTraitsDispatched(Params& params,
-                                                              cudaStream_t stream) {
-  using DTypeQ = typename KernelTraits::DTypeQ;
-  using DTypeKV = typename KernelTraits::DTypeKV;
-  using DTypeO = typename KernelTraits::DTypeO;
-  using IdType = typename KernelTraits::IdType;
-  using NV_TYPE = typename KernelTraits::NV_TYPE;
-
-  using CollectiveMainloop = CollectiveMainloop<KernelTraits, CAUSAL>;
-  using CollectiveEpilogue = CollectiveEpilogue<KernelTraits>;
-
-  typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments({
-      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.qk_head_dim), make_stride(params.qk_head_dim, _1{})), // layout q
-      make_layout(make_shape(params.block_size, params.qk_head_dim, params.max_block_num), make_stride(params.qk_head_dim, _1{}, params.block_size * params.qk_head_dim)),
-      make_layout(make_shape(params.chunk_num, params.bsz * params.draft_total_token_num * params.q_num_head), make_stride(params.bsz * params.draft_total_token_num * params.q_num_head, _1{})),
-      params.Q,
-      params.KV,
-      params.m,
-      params.d,
-      params.block_tables,
-      params.seq_lens_this_time,
-      params.seq_lens_encoder,
-      params.seq_lens_decoder,
-      params.cumsum_q_seqlens,
-      params.batch_ids,
-      params.tile_ids_per_batch,
-      params.q_tile_ids_per_batch,
-      params.num_blocks_x,
-      params.sm_scale,
-      params.bsz,
-      params.max_block_num,
-      params.max_block_num_per_seq,
-      params.q_stride_bsz,
-      params.q_stride_head_num,
-      params.kv_stride_block_num,
-      params.kv_stride_block_size,
-      params.o_stride_bsz,
-      params.o_stride_head_num,
-      params.chunk_size,
-      params.chunk_num,
-      params.draft_total_token_num
-  });
-  typename CollectiveEpilogue::Params epilogue_params = CollectiveEpilogue::to_underlying_arguments_ntma({
-      params.O,
-      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.vo_head_dim), make_stride(params.vo_head_dim, _1{})), // layout O
-      params.O_tmp,
-      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.vo_head_dim), make_stride(params.vo_head_dim, _1{})) // layout O_tmp
-  });
-
-  // Get the ptr to kernel function.
-  auto kernel =
-      MLAWithKVCacheWG4Kernel<CollectiveMainloop, CollectiveEpilogue, KernelTraits, CAUSAL, USE_REG_EALLOC, USE_QK_TWO_STAGE>;
-  int smem_size = sizeof(typename KernelTraits::SharedStorage);
-  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-  int device;
-  cudaGetDevice(&device);
-  int multiprocessor_count;
-  cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device);
-  int act_blocks_per_sm;
-  cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-      &act_blocks_per_sm, kernel, KernelTraits::NUM_WARPS * 32, smem_size);
-  dim3 grid_dims = {params.num_blocks_x_int, 1, 1};
-  static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32;
-  dim3 block_dims(ctaSize, 1, 1);
-  kernel<<<grid_dims, block_dims, smem_size, stream>>>(
-    mainloop_params, epilogue_params
-  );
-  if (params.chunk_num > 1) {
-    constexpr int vec_size = 16 / sizeof(DTypeO);
-    constexpr int merge_block_size = 256;
-    constexpr int blockx = KernelTraits::HEAD_DIM_VO / vec_size;
-    constexpr int blocky = (merge_block_size + blockx - 1) / blockx;
-    dim3 grids_merge(min(multiprocessor_count, params.token_num), params.q_num_head); // 128k is too large
-    dim3 blocks_merge(blockx, blocky);
-    merge_multi_chunks_kernel<NV_TYPE, vec_size, blocky, KernelTraits::HEAD_DIM_VO><<<grids_merge, blocks_merge, 0, stream>>>(
-      reinterpret_cast<NV_TYPE*>(params.O_tmp),
-      params.m,
-      params.d,
-      params.seq_lens_this_time,
-      params.seq_lens_decoder,
-      params.seq_lens_encoder,
-      params.padding_offsets,
-      reinterpret_cast<NV_TYPE*>(params.O),
-      params.max_seq_len,
-      params.chunk_num,
-      params.q_num_head,
-      params.chunk_size,
-      params.vo_head_dim,
-      params.token_num,
-      params.bsz,
-      params.draft_total_token_num
-    );
-  }
-  return cudaSuccess;
-}
-
-template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params, bool USE_REG_EALLOC=false, bool USE_QK_TWO_STAGE=false>
-cudaError_t BatchMLAWithPagedKVCacheWG4Dispatched(Params& params, cudaStream_t stream) {
-  constexpr bool CAUSAL = true;
-  if constexpr (HEAD_DIM_QK == 576) {
-    if (params.block_size == 32) {
-      DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
-        BatchMLAWithPagedKVCacheWG4KernelTraitsDispatched<
-            AttentionKernelWG4Traits</*USE_TMA_LOAD_KV=*/false, 
-                                     HEAD_DIM_QK, 
-                                     HEAD_DIM_VO, 
-                                     GROUP_SIZE,
-                                     /*BLOCK_SHAPE_Q_=*/64,
-                                     /*BLOCK_SHAPE_KV_=*/32,
-                                     /*NUM_STAGES_=*/4, 
-                                     typename Params::DTypeQ,
-                                     typename Params::DTypeKV, 
-                                     typename Params::DTypeO,
-                                     typename Params::IdType,
-                                     NV_TYPE>,
-            CAUSAL,
-            Params,
-            USE_REG_EALLOC,
-            USE_QK_TWO_STAGE>(params, stream);)
-    }
-  } else {
-    return cudaErrorNotSupported;
-  }
-  return cudaSuccess;
-};
-
-}  // namespace mla_attn
-
-#endif  // ATTENTION_HOPPER_PREFILL_WG4_SM90_CUH_

csrc/gpu/mla_attn/named_barrier.cuh

@@ -19,8 +19,6 @@
  * Modified by the FlashInfer team.
  */
 
-#pragma once
-
 #ifndef ATTENTION_HOPPER_NAMED_BARRIERS_CUH_
 #define ATTENTION_HOPPER_NAMED_BARRIERS_CUH_
 
@@ -42,7 +40,6 @@ enum class NamedBarriers {
   kWG1WG2Sync = 7,
   kWG0WG1WG2Sync = 8,
   kWG1WG2LastSync = 9,
-  kAllWGSync = 10,
 };
 
 }  // namespace mla_attn

csrc/gpu/mla_attn/utils.cuh

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#pragma once
-
 #ifndef ATTENTION_HOPPER_UTILS_CUH_
 #define ATTENTION_HOPPER_UTILS_CUH_
 
@@ -266,7 +264,7 @@ __global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [
                                           const int head_dim,
                                           const int token_num,
                                           const int bsz,
-                                          const int draft_total_token_num) {
+                                          const int max_draft_token_num) {
   const int vid = threadIdx.x, ty = threadIdx.y;
   const int hid = blockIdx.y;
   __shared__ T smem[bdy * HEAD_DIM];
@@ -310,13 +308,13 @@ __global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [
 
     for (int i = ty; i < num_chunks_this_seq; i += bdy) {
       uint32_t offset;
-      offset = ((i * bsz + bid) * draft_total_token_num + local_seq_id) * num_heads + hid;
+      offset = ((i * bsz + bid) * max_draft_token_num + local_seq_id) * num_heads + hid;
       float m_prev = m;
       float d_prev = d;
       const float m_now = multi_m[offset];
       const float d_now = multi_d[offset];
       m = max(m_prev, m_now);
-      offset = (((i * bsz + bid) * draft_total_token_num + local_seq_id) * num_heads + hid) * head_dim + vid * vec_size;
+      offset = (((i * bsz + bid) * max_draft_token_num + local_seq_id) * num_heads + hid) * head_dim + vid * vec_size;
       Load<T, vec_size>(&multi_out[offset], &load_vec);
       const float scale1 = __expf(m_prev - m), scale2 = __expf(m_now - m);
       const T scale1_T = static_cast<T>(scale1), scale2_T = static_cast<T>(scale2);

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

@@ -1,93 +0,0 @@
-/***************************************************************************************************
- * 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

@@ -1,346 +0,0 @@
-/***************************************************************************************************
- * 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

@@ -1,611 +0,0 @@
-/***************************************************************************************************
- * 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

@@ -1,85 +0,0 @@
-/*
- * 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

@@ -1,24 +0,0 @@
-/*
- * 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

@@ -1,25 +0,0 @@
-/*
- * 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

@@ -1,26 +0,0 @@
-/*
- * 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

@@ -1,23 +0,0 @@
-/*
- * 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

@@ -1,23 +0,0 @@
-/*
- * 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

@@ -1,23 +0,0 @@
-/*
- * 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

@@ -1,819 +0,0 @@
-/*
- * 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

@@ -1,192 +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.
-
-
-// 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

@@ -1,14 +0,0 @@
-#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

@@ -1,427 +0,0 @@
-
-/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
-
-#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

@@ -1,123 +0,0 @@
-/*
- * 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

@@ -1,990 +0,0 @@
-// /*
-//  * 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

@@ -1,276 +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.
-
-
-// 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

@@ -1,232 +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.
-
-#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

@@ -1,146 +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.
-
-
-// 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/multi_head_latent_attention.cu

@@ -38,7 +38,6 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
     const paddle::Tensor& decoder_num_blocks_cpu,
-    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
@@ -59,22 +58,24 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_draft_total_token_num,
+    const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
   typedef PDTraits<D> traits_;
   typedef typename traits_::data_t data_t;
+
   int decoder_num_blocks_data = decoder_num_blocks_cpu.data<int>()[0];
   int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
-  int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
   int max_len_kv_data = max_len_kv.data<int>()[0];
-  auto mla_use_tensorcore = GetMlaUseTensorcore();
 
-  if (speculate_decoder && !mla_use_tensorcore) {
+  const bool mla_use_tensorcore = get_mla_use_tensorcore();
+  auto sm_version = GetSMVersion();
+  if ((speculate_decoder || mla_use_tensorcore) && sm_version < 90) {
     PD_THROW("Please use speculate_decoder=0 and FLAGS_mla_use_tensorcore=0 when sm < 90.");
   }
 
   auto main_stream = query.stream();
+
   paddle::Tensor fmha_out = paddle::full(
       {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v},
       0,
@@ -105,14 +106,13 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
                                              decoder_num_blocks,
                                              cache_quant_type_str,
                                              decoder_num_blocks_data,
-                                             chunk_size,
                                              max_input_length,
                                              max_len_kv_data,
                                              softmax_scale,
                                              quant_max_bound,
                                              quant_min_bound,
                                              out_linear_in_scale,
-                                             speculate_draft_total_token_num,
+                                             speculate_max_draft_token_num,
                                              causal,
                                              main_stream,
                                              &fmha_out);
@@ -163,7 +163,6 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
     const paddle::Tensor& decoder_num_blocks_cpu,
-    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
@@ -186,7 +185,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_draft_total_token_num,
+    const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
   AppendAttnMetaData meta_data;
@@ -205,7 +204,6 @@ 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,
@@ -229,7 +227,6 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           decoder_tile_ids_per_batch,
           decoder_num_blocks,
           decoder_num_blocks_cpu,
-          decoder_chunk_size_cpu,
           max_enc_len_this_time,
           max_dec_len_this_time,
           max_len_kv,
@@ -250,11 +247,10 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           quant_max_bound,
           quant_min_bound,
           out_linear_in_scale,
-          speculate_draft_total_token_num,
+          speculate_max_draft_token_num,
           causal,
           speculate_decoder);
     }
-#endif
     case paddle::DataType::FLOAT16: {
       return MultiHeadLatentAttentionKernel<paddle::DataType::FLOAT16>(
           meta_data,
@@ -278,7 +274,6 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           decoder_tile_ids_per_batch,
           decoder_num_blocks,
           decoder_num_blocks_cpu,
-          decoder_chunk_size_cpu,
           max_enc_len_this_time,
           max_dec_len_this_time,
           max_len_kv,
@@ -299,7 +294,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           quant_max_bound,
           quant_min_bound,
           out_linear_in_scale,
-          speculate_draft_total_token_num,
+          speculate_max_draft_token_num,
           causal,
           speculate_decoder);
     }
@@ -333,7 +328,6 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
     const std::vector<int64_t>& decoder_tile_ids_per_batch_shape,
     const std::vector<int64_t>& decoder_num_blocks_shape,
     const std::vector<int64_t>& decoder_num_blocks_cpu_shape,
-    const std::vector<int64_t>& decoder_chunk_size_cpu_shape,
     const std::vector<int64_t>& max_enc_len_this_time_shape,
     const std::vector<int64_t>& max_dec_len_this_time_shape,
     const std::vector<int64_t>& max_len_kv_shape,
@@ -356,7 +350,7 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_draft_total_token_num,
+    const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
   const int token_num = query_shape[0];
@@ -368,8 +362,6 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
   return {{token_num, num_heads * head_dim_v}};
 }
 
-
-
 std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const paddle::DataType& query_dtype,
     const paddle::DataType& key_cache_dtype,
@@ -391,7 +383,6 @@ std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const paddle::DataType& decoder_tile_ids_per_batch_dtype,
     const paddle::DataType& decoder_num_blocks_dtype,
     const paddle::DataType& decoder_num_blocks_cpu_dtype,
-    const paddle::DataType& decoder_chunk_size_cpu_dtype,
     const paddle::DataType& max_enc_len_this_time_dtype,
     const paddle::DataType& max_dec_len_this_time_dtype,
     const paddle::DataType& max_len_kv_dtype,
@@ -414,7 +405,7 @@ std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_draft_total_token_num,
+    const int speculate_max_draft_token_num,
     const bool causal,
     const bool speculate_decoder) {
   if (compute_dtype == "bf16") {
@@ -447,7 +438,6 @@ PD_BUILD_OP(multi_head_latent_attention)
              "decoder_tile_ids_per_batch",
              "decoder_num_blocks",
              "decoder_num_blocks_cpu",
-             "decoder_chunk_size_cpu",
              "max_enc_len_this_time",
              "max_dec_len_this_time",
              "max_len_kv",
@@ -471,7 +461,7 @@ PD_BUILD_OP(multi_head_latent_attention)
             "quant_max_bound: float",
             "quant_min_bound: float",
             "out_linear_in_scale: float",
-            "speculate_draft_total_token_num: int",
+            "speculate_max_draft_token_num: int",
             "causal: bool",
             "speculate_decoder: bool"})
     .SetKernelFn(PD_KERNEL(MultiHeadLatentAttention))

csrc/gpu/preprocess_for_moe.cu

@@ -140,14 +140,12 @@ 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("【ERROR】, please add the current number of experts in the case");
+      printf("errors");
     }
 
     const int block_threads = 256;

csrc/gpu/quant_int8.cu

@@ -271,7 +271,7 @@ std::vector<paddle::DataType> QuantInt8Dtype(const paddle::DataType& input_dtype
 }
 
 PD_BUILD_OP(quant_int8)
-    .Inputs({"input", paddle::Optional("shift"),paddle::Optional("smooth") })
+    .Inputs({"intput", paddle::Optional("shift"),paddle::Optional("smooth") })
     .Outputs({"output"})
     .Attrs({"scale: float","round_type: int","max_bound: float", "min_bound: float"})
     .SetKernelFn(PD_KERNEL(QuantInt8))

csrc/gpu/quantization/per_tensor_quant_fp8.cu

@@ -1,169 +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 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

@@ -1,251 +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>
-
-
-__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};
 }
 
-paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
+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,
@@ -153,7 +153,7 @@ paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num
                 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 @@ paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num
                 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 @@ paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num
                 seq_lens_encoder,
                 output_padding_offset,
                 max_input_length
-            )[0];
+            );
         }
         default: {
             PD_THROW(
@@ -184,22 +184,6 @@ paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num
     }
 }
 
-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_fused.cuh

@@ -1,19 +0,0 @@
-#pragma once
-#include "paddle/extension.h"
-
-std::vector<paddle::Tensor> per_warp_int8_cuda(paddle::Tensor& q,
-                                            paddle::Tensor& k,
-                                            paddle::Tensor& km,
-                                            int BLKQ,
-                                            int WARPQ,
-                                            int BLKK,
-                                            int tensor_layout);
-                    
-std::vector<paddle::Tensor> per_channel_fp8(paddle::Tensor& v,
-                                            int tensor_layout,
-                                            float scale_max,
-                                            bool smooth_v);
-
-std::vector<paddle::Tensor> sub_mean(paddle::Tensor& v,
-                                    paddle::Tensor& vm,
-                                    int tensor_layout);

csrc/gpu/save_output_dygraph.cu

@@ -1,76 +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"
-
-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,7 +20,6 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
-// #define SAVE_WITH_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -36,25 +35,7 @@ void SaveOutMmsg(const paddle::Tensor& x,
     auto not_need_stop_data = not_need_stop.data<bool>()[0];
 
     static struct msgdata msg_sed;
-    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 key_t key = ftok("./", 1);
     static int msgid = msgget(key, IPC_CREAT | 0666);
 
     msg_sed.mtype = 1;

csrc/gpu/set_preids_token_penalty_multi_scores.cu

@@ -41,18 +41,18 @@ __global__ void set_preids_token_penalty_multi_scores_kernel(const bool *stop_fl
     T *logits_now = logits + bi * length;
     int tid = threadIdx.x;
 
-    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
+    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 (seq_len_dec == 0 && seq_len_enc == 0) return; // stoped
 
         const int step_idx_now = step_idx[bi];
         if (tid == 0 && step_idx_now >= 0) {
             if (seq_len_enc > 0) { // encoder, get last token accord to seq_lens_encoder
                 pre_ids_now[step_idx_now] = input_ids_now[seq_len_enc - 1];
-            } else { // decoder, get first token
+            } else { // decoedr, get first token
                 pre_ids_now[step_idx_now] = input_ids_now[0];
             }
         }

csrc/gpu/set_value_by_flags_v2.cu

@@ -29,11 +29,11 @@ __global__ void set_value_by_flag_and_id_v2(const bool *stop_flags,
         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 (seq_len_dec == 0 && seq_len_enc == 0) return; // stopped
+        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stoped
         if (step_idx[tid] >= 0) {
             if (seq_len_dec == 0) { // encoder, get last token accord to seq_lens_encoder
                 pre_ids_all_now[step_idx[tid]] = input_ids_now[seq_len_enc - 1];
-            } else { // decoder, get first token
+            } else { // decoedr, get first token
                 pre_ids_all_now[step_idx[tid]] = input_ids_now[0];
             }
         }

csrc/gpu/speculate_decoding_kernels/speculate_set_value_by_flags.cu

@@ -30,7 +30,7 @@ __global__ void speculate_set_value_by_flag_and_id(int64_t *pre_ids_all,
         const int64_t *accept_tokens_now = accept_tokens + tid * max_draft_tokens;
         const int seq_len_dec = seq_lens_decoder[tid];
         const int seq_len_enc = seq_lens_encoder[tid];
-        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stopped
+        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stoped
         if (step_idx[tid] >= 0) {
             for (int i = 0; i < accept_num[tid]; i++) {
                 pre_ids_all_now[step_idx[tid] - i] = accept_tokens_now[accept_num[tid] - 1 - i];

csrc/gpu/unittest/run_mla_benchmark.sh

@@ -1,23 +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.
-
-#请使用cuda-12.8编译和运行benchnark
-export FLAGS_mla_use_tensorcore=1
-export FLAGS_mla_use_wg4=0
-export FLAGS_mla_dec_chunk_size=-1
-
-# python test_absorb_mla.py
-
-ncu=`which ncu`
-${ncu} --section regex:'^(?!Nvlink)' -o profile_mla_kernel -f --import-source on --cache-control=all --clock-control=base -k regex:"MLAWithKVCacheKernel|merge_multi_chunks_kernel" --print-source=cuda,sass --page source python test_absorb_mla.py

csrc/gpu/unittest/test_absorb_mla.py

@@ -62,28 +62,28 @@ MAX_DEC_LEN = 1024
 NUM_Q_HEAD = 8
 NUM_KV_HEAD = 1
 dtype = "bfloat16"
-DRAFT_TOTAL_TOKEN_NUM = 1
+MAX_DRAFT_TOKEN_NUM = 1
 CAUSAL = True
 
 SPECULATE_DECODER = False
 
 
 def ref_attention(q_all, p_compressed_kv, compressed_kv, p_key_pe, key_pe, bsz, cache_length, softmax_scale):
-    q = q_all.reshape([bsz, DRAFT_TOTAL_TOKEN_NUM, NUM_Q_HEAD, HEAD_DIM_QK])
+    q = q_all.reshape([bsz, MAX_DRAFT_TOKEN_NUM, NUM_Q_HEAD, HEAD_DIM_QK])
     p_compressed_kv = p_compressed_kv.reshape((bsz, cache_length, HEAD_DIM_V))
-    compressed_kv = compressed_kv.reshape((bsz, DRAFT_TOTAL_TOKEN_NUM, HEAD_DIM_V))
+    compressed_kv = compressed_kv.reshape((bsz, MAX_DRAFT_TOKEN_NUM, HEAD_DIM_V))
     p_key_pe = p_key_pe.reshape((bsz, cache_length, PE_SIZE))
-    key_pe = key_pe.reshape((bsz, DRAFT_TOTAL_TOKEN_NUM, PE_SIZE))
+    key_pe = key_pe.reshape((bsz, MAX_DRAFT_TOKEN_NUM, PE_SIZE))
 
     p_k = paddle.concat([p_compressed_kv, p_key_pe], axis=-1).reshape((bsz, cache_length, HEAD_DIM_QK))
-    k = paddle.concat([compressed_kv, key_pe], axis=-1).reshape((bsz, DRAFT_TOTAL_TOKEN_NUM, HEAD_DIM_QK))
+    k = paddle.concat([compressed_kv, key_pe], axis=-1).reshape((bsz, MAX_DRAFT_TOKEN_NUM, HEAD_DIM_QK))
     k = paddle.concat([p_k, k], axis=1).transpose([0, 2, 1])
     v = paddle.concat([p_compressed_kv, compressed_kv], axis=1).reshape(
-        (bsz, cache_length + DRAFT_TOTAL_TOKEN_NUM, HEAD_DIM_V)
+        (bsz, cache_length + MAX_DRAFT_TOKEN_NUM, HEAD_DIM_V)
     )
 
-    out = paddle.zeros(shape=[bsz, DRAFT_TOTAL_TOKEN_NUM, NUM_Q_HEAD, HEAD_DIM_V], dtype=dtype)
-    for i in range(DRAFT_TOTAL_TOKEN_NUM):
+    out = paddle.zeros(shape=[bsz, MAX_DRAFT_TOKEN_NUM, NUM_Q_HEAD, HEAD_DIM_V], dtype=dtype)
+    for i in range(MAX_DRAFT_TOKEN_NUM):
         for j in range(NUM_Q_HEAD):
 
             query = q[:, i, j, :].reshape([-1, 1, HEAD_DIM_QK])
@@ -147,7 +147,7 @@ def test_append_c16_attention(cache_length, bsz):
         cache_length,
     ] * bsz
     seq_lens_this_time = [
-        DRAFT_TOTAL_TOKEN_NUM,
+        MAX_DRAFT_TOKEN_NUM,
     ] * bsz
     max_enc_len_this_time = max(seq_lens_enc)
     max_dec_len_this_time = max(seq_lens_dec)
@@ -186,7 +186,6 @@ def test_append_c16_attention(cache_length, bsz):
         decoder_tile_ids_per_batch,
         decoder_num_blocks_device,
         decoder_num_blocks,
-        decoder_chunk_size,
         max_len_kv,
     ) = paddlenlp_ops.get_block_shape_and_split_kv_block(
         seq_lens_encoder,
@@ -197,7 +196,7 @@ def test_append_c16_attention(cache_length, bsz):
         cum_offsets,
         NUM_Q_HEAD // NUM_KV_HEAD,
         BLOCK_SIZE,
-        DRAFT_TOTAL_TOKEN_NUM,
+        MAX_DRAFT_TOKEN_NUM,
     )
     softmax_scale = HEAD_DIM_QK ** (-0.5)
 
@@ -230,6 +229,7 @@ def test_append_c16_attention(cache_length, bsz):
 
     ref_out = ref_attention(query, p_compressed_kv, compressed_kv, p_key_pe, key_pe, bsz, cache_length, softmax_scale)
     paddle.device.synchronize()
+
     s_time = 0
     for i in range(RUN_TIME + WARM_UP):
         if i == WARM_UP:
@@ -255,7 +255,6 @@ def test_append_c16_attention(cache_length, bsz):
             decoder_tile_ids_per_batch,
             decoder_num_blocks_device,
             decoder_num_blocks,
-            decoder_chunk_size,
             max_enc_len_this_time,
             max_dec_len_this_time,
             max_len_kv,
@@ -278,7 +277,7 @@ def test_append_c16_attention(cache_length, bsz):
             0.0,
             0.0,
             0.0,  # out_linear_in_scale
-            DRAFT_TOTAL_TOKEN_NUM,  # DRAFT_TOTAL_TOKEN_NUM
+            MAX_DRAFT_TOKEN_NUM,  # max_draft_token_num
             CAUSAL,
             SPECULATE_DECODER,
         )
@@ -295,7 +294,6 @@ def test_append_c16_attention(cache_length, bsz):
 
 
 if __name__ == "__main__":
-    print("CUDA_VERSION >=12.8 is required to compile and run mla benchmark")
     for cache_length in [1023, 2047]:
         for bsz in [1, 16, 32, 64, 128, 256]:
             test_append_c16_attention(cache_length, bsz)

csrc/gpu/unittest/test_dynamic_per_tensor_quant_fp8.py

@@ -1,88 +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.
-
-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()