Update setup.py

* add rl qwen model support

* fix

* fix

* add_commit_config

* fix

.flake8

@@ -1,7 +0,0 @@
-[flake8]
-ignore = E203, E402, E501, E731, E741, W503, W605, E722, E231, W604, E702, E226, E221, E713, E271
-max-line-length = 119
-
-# E402: module level import not at top of file
-per-file-ignores =
-    __init__.py:F401,F403,E402

.github/workflows/Codestyle-Check.yml

@@ -1,50 +0,0 @@
-name: Codestyle-Check
-
-on:
-  pull_request:
-    branches:
-      - develop
-      - 'release/*'
-
-jobs:
-  pre-commit:
-    name: Pre Commit
-    if: ${{ github.repository_owner == 'PaddlePaddle' }}
-    runs-on: ubuntu-latest
-    env:
-      PR_ID: ${{ github.event.pull_request.number }}
-      BRANCH: ${{ github.event.pull_request.base.ref }}
-
-    steps:
-      - name: Cleanup
-        run: |
-          rm -rf * .[^.]*
-
-      - name: Checkout base repo
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.pull_request.base.ref }}
-          fetch-depth: 1000
-
-      - name: Merge PR to test branch
-        run: |
-          git fetch origin pull/${PR_ID}/merge
-          git checkout -b test FETCH_HEAD
-
-      - name: Setup python3.10
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-          cache: 'pip'
-
-      - name: Install dependencies
-        run: |
-          pip install pre-commit==4.2.0 cpplint==1.6.0 clang-format==13.0.0
-
-      - name: Check pre-commit
-        env:
-          SKIP_CLANG_TIDY_CHECK: "ON"
-        run: |
-          set +e
-          bash -x tools/codestyle/pre_commit.sh;EXCODE=$?
-          exit $EXCODE

.github/workflows/_accuracy_test.yml

@@ -1,187 +0,0 @@
-name: Accuracy Test
-description: "Run Accuracy Tests"
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      FASTDEPLOY_WHEEL_URL:
-        description: "URL of the FastDeploy Wheel."
-        required: true
-        type: string
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-      MODEL_CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-
-jobs:
-  accuracy_tests:
-    runs-on: [self-hosted, GPU-h20-1Cards]
-    timeout-minutes: 60
-    steps:
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-            docker pull ${docker_image}
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-
-      - name: Run FastDeploy Base Tests
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fastdeploy_wheel_url: ${{ inputs.FASTDEPLOY_WHEEL_URL }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          MODEL_CACHE_DIR: ${{ inputs.MODEL_CACHE_DIR }}
-        run: |
-          runner_name="${{ runner.name }}"
-          CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-          DEVICES=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-          DEVICE_PORT=$(echo "$DEVICES" | cut -d',' -f1)
-
-          FLASK_PORT=$((42068 + DEVICE_PORT * 100))
-          FD_API_PORT=$((42088 + DEVICE_PORT * 100))
-          FD_ENGINE_QUEUE_PORT=$((42058 + DEVICE_PORT * 100))
-          FD_METRICS_PORT=$((42078 + DEVICE_PORT * 100))
-          FD_CACHE_QUEUE_PORT=$((42098 + DEVICE_PORT * 100))
-          echo "Test ENV Parameter:"
-          echo "========================================================="
-          echo "FLASK_PORT=${FLASK_PORT}"
-          echo "FD_API_PORT=${FD_API_PORT}"
-          echo "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}"
-          echo "FD_METRICS_PORT=${FD_METRICS_PORT}"
-          echo "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}"
-          echo "DEVICES=${DEVICES}"
-          echo "========================================================="
-
-          CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-          echo "CACHE_DIR is set to ${CACHE_DIR}"
-          if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-            touch "${CACHE_DIR}/gitconfig"
-          fi
-          if [ ! -d "${MODEL_CACHE_DIR}" ]; then
-            echo "Error: MODEL_CACHE_DIR '${MODEL_CACHE_DIR}' does not exist."
-            exit 1
-          fi
-
-          PORTS=($FLASK_PORT $FD_API_PORT $FD_ENGINE_QUEUE_PORT $FD_METRICS_PORT $FD_CACHE_QUEUE_PORT)
-          LOG_FILE="./port_cleanup_$(date +%Y%m%d_%H%M%S).log"
-          echo "==== LOG_FILE is ${LOG_FILE} ===="
-
-          echo "==== PORT CLEAN BEFORE TASK RUN ====" | tee -a $LOG_FILE
-
-          for port in "${PORTS[@]}"; do
-              PIDS=$(lsof -t -i :$port || true)
-              if [ -n "$PIDS" ]; then
-                  echo "Port $port is occupied by PID(s): $PIDS" | tee -a $LOG_FILE
-                  echo "$PIDS" | xargs -r kill -9
-                  echo "Port $port cleared" | tee -a $LOG_FILE
-              else
-                  echo "Port $port is free" | tee -a $LOG_FILE
-              fi
-          done
-
-          echo "==== PORT CLEAN COMPLETE ====" | tee -a $LOG_FILE
-
-          echo "========================================================="
-          echo "Ensuring no stale container named ${runner_name} ..."
-          if [ "$(docker ps -a -q -f name=${runner_name})" ]; then
-            echo "Removing stale container: ${runner_name}"
-            docker rm -f ${runner_name} || true
-          fi
-
-          docker run --rm --ipc=host --pid=host --net=host \
-          --name ${runner_name} \
-          -v $(pwd):/workspace \
-          -w /workspace \
-          -e fastdeploy_wheel_url=${fastdeploy_wheel_url} \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          -e "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}" \
-          -e "FLASK_PORT=${FLASK_PORT}" \
-          -v "${MODEL_CACHE_DIR}:/MODELDATA" \
-          -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-          -v "${CACHE_DIR}/.cache:/root/.cache" \
-          -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-          -e TZ="Asia/Shanghai" \
-          --gpus '"device='"${DEVICES}"'"' ${docker_image} /bin/bash -xc '
-          python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-
-          pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
-
-          python -m pip install ${fastdeploy_wheel_url}
-          python -m pip install pytest
-
-          wget https://paddle-qa.bj.bcebos.com/zhengtianyu/tools/llm-deploy-linux-amd64
-          chmod +x ./llm-deploy-linux-amd64
-          ./llm-deploy-linux-amd64 -python python3.10 \
-          -model_name ERNIE-4.5-0.3B-Paddle \
-          -model_path /MODELDATA \
-          --skip install
-
-          git config --global --add safe.directory /workspace/FastDeploy
-          cd FastDeploy
-          pushd tests/ce/deploy
-          ps -ef | grep "${FD_ENGINE_QUEUE_PORT}" | grep -v grep | awk "{print \$2}" | xargs -r kill -9
-          python3.10 deploy.py > dd.log 2>&1 &
-          sleep 3
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/start \
-            -H "Content-Type: application/json" \
-            -d "{\"--model\": \"/MODELDATA/ERNIE-4.5-0.3B-Paddle\"}"
-
-          curl -X POST http://localhost:${FLASK_PORT}/wait_for_infer?timeout=90
-          popd
-
-          pushd tests/ce/accuracy_cases
-          export URL=http://localhost:${FD_API_PORT}/v1/chat/completions
-          export TEMPLATE=TOKEN_LOGPROB
-          export MODEL_SIZE=0.3B
-          TEST_EXIT_CODE=0
-          python gsm8k.py || TEST_EXIT_CODE=1
-          popd
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}" >> /workspace/FastDeploy/exit_code.env
-          '
-          if [ -f ./FastDeploy/exit_code.env ]; then
-            source ./FastDeploy/exit_code.env
-            cat ./FastDeploy/exit_code.env >> $GITHUB_ENV
-          fi
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}"
-          exit ${TEST_EXIT_CODE}

.github/workflows/_base_test.yml

@@ -1,230 +0,0 @@
-name: Base Test
-description: "Run Base Tests"
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      FASTDEPLOY_WHEEL_URL:
-        description: "URL of the FastDeploy Wheel."
-        required: true
-        type: string
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-      MODEL_CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-
-jobs:
-  base_tests:
-    runs-on: [self-hosted, GPU-h20-1Cards]
-    timeout-minutes: 60
-    steps:
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-            docker pull ${docker_image}
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-
-      - name: Run FastDeploy Base Tests
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fastdeploy_wheel_url: ${{ inputs.FASTDEPLOY_WHEEL_URL }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          MODEL_CACHE_DIR: ${{ inputs.MODEL_CACHE_DIR }}
-        run: |
-          runner_name="${{ runner.name }}"
-          CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-          DEVICES=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-          DEVICE_PORT=$(echo "$DEVICES" | cut -d',' -f1)
-
-          FLASK_PORT=$((42068 + DEVICE_PORT * 100))
-          FD_API_PORT=$((42088 + DEVICE_PORT * 100))
-          FD_ENGINE_QUEUE_PORT=$((42058 + DEVICE_PORT * 100))
-          FD_METRICS_PORT=$((42078 + DEVICE_PORT * 100))
-          FD_CACHE_QUEUE_PORT=$((42098 + DEVICE_PORT * 100))
-          echo "Test ENV Parameter:"
-          echo "========================================================="
-          echo "FLASK_PORT=${FLASK_PORT}"
-          echo "FD_API_PORT=${FD_API_PORT}"
-          echo "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}"
-          echo "FD_METRICS_PORT=${FD_METRICS_PORT}"
-          echo "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}"
-          echo "DEVICES=${DEVICES}"
-          echo "========================================================="
-
-          CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-          echo "CACHE_DIR is set to ${CACHE_DIR}"
-          if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-            touch "${CACHE_DIR}/gitconfig"
-          fi
-          if [ ! -d "${MODEL_CACHE_DIR}" ]; then
-            echo "Error: MODEL_CACHE_DIR '${MODEL_CACHE_DIR}' does not exist."
-            exit 1
-          fi
-
-          PORTS=($FLASK_PORT $FD_API_PORT $FD_ENGINE_QUEUE_PORT $FD_METRICS_PORT $FD_CACHE_QUEUE_PORT)
-          LOG_FILE="./port_cleanup_$(date +%Y%m%d_%H%M%S).log"
-          echo "==== LOG_FILE is ${LOG_FILE} ===="
-
-          echo "==== PORT CLEAN BEFORE TASK RUN ====" | tee -a $LOG_FILE
-
-          for port in "${PORTS[@]}"; do
-              PIDS=$(lsof -t -i :$port || true)
-              if [ -n "$PIDS" ]; then
-                  echo "Port $port is occupied by PID(s): $PIDS" | tee -a $LOG_FILE
-                  echo "$PIDS" | xargs -r kill -9
-                  echo "Port $port cleared" | tee -a $LOG_FILE
-              else
-                  echo "Port $port is free" | tee -a $LOG_FILE
-              fi
-          done
-
-          echo "==== PORT CLEAN COMPLETE ====" | tee -a $LOG_FILE
-
-          echo "========================================================="
-          echo "Ensuring no stale container named ${runner_name} ..."
-          if [ "$(docker ps -a -q -f name=${runner_name})" ]; then
-            echo "Removing stale container: ${runner_name}"
-            docker rm -f ${runner_name} || true
-          fi
-
-          docker run --rm --ipc=host --pid=host --net=host \
-          --name ${runner_name} \
-          -v $(pwd):/workspace \
-          -w /workspace \
-          -e fastdeploy_wheel_url=${fastdeploy_wheel_url} \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          -e "FLASK_PORT=${FLASK_PORT}" \
-          -e "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}" \
-          -v "${MODEL_CACHE_DIR}:/MODELDATA" \
-          -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-          -v "${CACHE_DIR}/.cache:/root/.cache" \
-          -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-          -e TZ="Asia/Shanghai" \
-          --gpus '"device='"${DEVICES}"'"' ${docker_image} /bin/bash -xc '
-          python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-
-          pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
-
-          python -m pip install ${fastdeploy_wheel_url}
-          python -m pip install pytest
-
-          wget https://paddle-qa.bj.bcebos.com/zhengtianyu/tools/llm-deploy-linux-amd64
-          chmod +x ./llm-deploy-linux-amd64
-          ./llm-deploy-linux-amd64 -python python3.10 \
-          -model_name ERNIE-4.5-0.3B-Paddle \
-          -model_path /MODELDATA \
-          --skip install
-
-          git config --global --add safe.directory /workspace/FastDeploy
-          cd FastDeploy
-          pushd tests/ce/deploy
-          ps -ef | grep "${FD_ENGINE_QUEUE_PORT}" | grep -v grep | awk "{print \$2}" | xargs -r kill -9
-          python3.10 deploy.py > dd.log 2>&1 &
-          sleep 3
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/start \
-            -H "Content-Type: application/json" \
-            -d "{\"--model\": \"/MODELDATA/ERNIE-4.5-0.3B-Paddle\"}"
-
-          check_service() {
-            local timeout=${1:-90}
-            local url="http://localhost:${FLASK_PORT}/wait_for_infer?timeout=${timeout}"
-            local resp
-
-            resp=$(curl -s -X POST "$url")
-
-            if echo "$resp" | grep -q "服务启动超时"; then
-              exit 8
-            fi
-          }
-
-          check_service 90
-          popd
-
-          pushd tests/ce/server
-          export URL=http://localhost:${FD_API_PORT}/v1/chat/completions
-          export TEMPLATE=TOKEN_LOGPROB
-          TEST_EXIT_CODE=0
-          python -m pytest -sv test_base_chat.py test_compare_top_logprobs.py test_logprobs.py test_params_boundary.py test_seed_usage.py test_stream.py test_evil_cases.py test_completions.py test_return_token_ids.py || TEST_EXIT_CODE=1
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/switch \
-            -H "Content-Type: application/json" \
-            -d "{\"--model\": \"/MODELDATA/ERNIE-4.5-0.3B-Paddle\", \"--early-stop-config\": \"{\\\"enable_early_stop\\\":true, \\\"window_size\\\":6, \\\"threshold\\\":0.93}\"}"
-          check_service 90
-          python -m pytest -sv test_repetition_early_stop.py || TEST_EXIT_CODE=1
-
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/switch \
-            -H "Content-Type: application/json" \
-            -d "{ \"--model\": \"/MODELDATA/ERNIE-4.5-0.3B-Paddle\", \"--max-concurrency\": 5, \"--max-waiting-time\": 1 }"
-          check_service 90
-          python -m pytest -sv test_max_concurrency.py || TEST_EXIT_CODE=1
-
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/switch \
-            -H "Content-Type: application/json" \
-            -d "{ \"--model\": \"/MODELDATA/ERNIE-4.5-0.3B-Paddle\", \"--max-concurrency\": 5000, \"--max-waiting-time\": 1 }"
-          check_service 90
-          python -m pytest -sv test_max_waiting_time.py || TEST_EXIT_CODE=1
-
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/switch \
-            -H "Content-Type: application/json" \
-            -d "{\"--model\": \"/MODELDATA/ernie-4_5-21b-a3b-bf16-paddle\", \"--config\": \"21b_mtp.yaml\", \"--enable-logprob\": \"False\"}"
-          check_service 180
-          export TEMPLATE=TOKEN_NORMAL
-          python -m pytest -sv test_seed_usage.py -k "not test_seed_stream" || TEST_EXIT_CODE=1
-
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/switch \
-            -H "Content-Type: application/json" \
-            -d "{\"--model\": \"/MODELDATA/ernie-4_5-21b-a3b-bf16-paddle\", \"--config\": \"21b_sot.yaml\", \"--enable-logprob\": \"False\"}"
-          check_service 360
-          export TEMPLATE=TOKEN_NORMAL
-          python -m pytest -sv test_seed_usage.py -k "not test_seed_stream" || TEST_EXIT_CODE=1
-
-          popd
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}" >> /workspace/FastDeploy/exit_code.env
-          '
-          if [ -f ./FastDeploy/exit_code.env ]; then
-            source ./FastDeploy/exit_code.env
-            cat ./FastDeploy/exit_code.env >> $GITHUB_ENV
-          fi
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}"
-          exit ${TEST_EXIT_CODE}

.github/workflows/_build_linux.yml

@@ -1,200 +0,0 @@
-name: FastDeploy Linux GPU Build Task
-description: "FastDeploy packages build and upload"
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      COMPILE_ARCH:
-        description: "Build GPU Archs"
-        required: true
-        type: string
-        default: "80,90"
-      WITH_NIGHTLY_BUILD:
-        description: "Enable nightly build mode (e.g. add date suffix to version)"
-        required: false
-        type: string
-        default: "OFF"
-      FD_VERSION:
-        description: "FastDeploy Package Version"
-        required: false
-        type: string
-        default: ""
-      PADDLEVERSION:
-        description: "Paddle Version Build Use"
-        required: false
-        type: string
-        default: ""
-      PADDLE_WHL_URL:
-        description: "Paddle Wheel Package URL"
-        required: false
-        type: string
-        default: ""
-      UPLOAD:
-        description: "Upload Package"
-        required: false
-        type: string
-        default: "ON"
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-    outputs:
-      wheel_path:
-        description: "Output path of the generated wheel"
-        value: ${{ jobs.fd-build.outputs.wheel_path }}
-jobs:
-  fd-build:
-    runs-on: [self-hosted, GPU-Build]
-    timeout-minutes: 360
-    outputs:
-      wheel_path: ${{ steps.set_output.outputs.wheel_path }}
-    steps:
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-          IS_PR: ${{ github.event_name == 'pull_request' }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-      - name: FastDeploy Build
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          compile_arch: ${{ inputs.COMPILE_ARCH }}
-          fd_version: ${{ inputs.FD_VERSION }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          BRANCH_REF: ${{ github.ref_name }}
-          PADDLEVERSION: ${{ inputs.PADDLEVERSION }}
-          PADDLE_WHL_URL: ${{ inputs.PADDLE_WHL_URL }}
-          WITH_NIGHTLY_BUILD: ${{ inputs.WITH_NIGHTLY_BUILD }}
-        run: |
-            set -x
-            runner_name="${{ runner.name }}"
-            CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-            gpu_id=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-
-            CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-            echo "CACHE_DIR is set to ${CACHE_DIR}"
-            if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-              touch "${CACHE_DIR}/gitconfig"
-            fi
-            PARENT_DIR=$(dirname "$WORKSPACE")
-            echo "PARENT_DIR:$PARENT_DIR"
-            docker run --rm --net=host \
-            --cap-add=SYS_PTRACE --privileged --shm-size=64G \
-            -v $(pwd):/workspace -w /workspace \
-            -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-            -v "${CACHE_DIR}/.cache:/root/.cache" \
-            -v "${CACHE_DIR}/.ccache:/root/.ccache" \
-            -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-            -e TZ="Asia/Shanghai" \
-            -e "COMPILE_ARCH=${compile_arch}" \
-            -e "FD_VERSION=${fd_version}" \
-            -e "WITH_NIGHTLY_BUILD=${WITH_NIGHTLY_BUILD}" \
-            -e "PADDLEVERSION=${PADDLEVERSION}" \
-            -e "PADDLE_WHL_URL=${PADDLE_WHL_URL}" \
-            -e "BRANCH_REF=${BRANCH_REF}" \
-            --gpus "\"device=${gpu_id}\"" ${docker_image} /bin/bash -c '
-            if [[ -n "${FD_VERSION}" ]]; then
-              export FASTDEPLOY_VERSION=${FD_VERSION}
-              echo "Custom FastDeploy version: ${FASTDEPLOY_VERSION}"
-            fi
-
-            git config --global --add safe.directory /workspace/FastDeploy
-            chown -R $(whoami) /workspace/FastDeploy
-            cd FastDeploy
-            if [[ "${WITH_NIGHTLY_BUILD}" == "ON" ]];then
-              GIT_COMMIT_TIME=$(git --no-pager show -s --format=%ci HEAD)
-              DATE_ONLY=$(echo $GIT_COMMIT_TIME | sed "s/ .*//;s/-//g")
-              echo "Git Commit Time: $GIT_COMMIT_TIME"
-              echo "Date Only: $DATE_ONLY"
-              export FASTDEPLOY_VERSION="${FASTDEPLOY_VERSION}.dev${DATE_ONLY}"
-            fi
-            # 针对不同分支和tag使用不同的PaddlePaddle安装包
-            if [[ "${PADDLE_WHL_URL}" != "" ]];then
-              python -m pip install ${PADDLE_WHL_URL}
-            elif [[ "${PADDLEVERSION}" != "" ]];then
-              python -m pip install paddlepaddle-gpu==${PADDLEVERSION} -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
-            else
-              python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-            fi
-
-            pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
-
-            python -m pip install --upgrade pip
-            python -m pip install -r requirements.txt
-            python -m pip install wheel
-            # 编译RDMA
-            export ENABLE_FD_RDMA=1
-            bash build.sh 1 python false [${COMPILE_ARCH}]
-            ls ./dist/*.whl
-            '
-      - name: Package Upload
-        id: set_output
-        env:
-          compile_arch: ${{ inputs.COMPILE_ARCH }}
-        run: |
-            set -x
-            if [[ "${{ github.event_name }}" == "pull_request" ]];then
-              commit_id=${{ github.event.pull_request.head.sha }}
-              pr_num=${{ github.event.pull_request.number }}
-              target_path=paddle-github-action/PR/FastDeploy/${pr_num}/${commit_id}/SM${compile_arch//,/_}
-            elif [[ "${{ github.ref_type }}" == "tag" ]]; then
-              commit_id=${{ github.sha }}
-              tag_name=${{ github.ref_name }}
-              target_path=paddle-github-action/TAG/FastDeploy/${tag_name}/${commit_id}/SM${compile_arch//,/_}
-            else
-              commit_id=${{ github.sha }}
-              branch_name=${{ github.ref_name }}
-              target_path=paddle-github-action/BRANCH/FastDeploy/${branch_name}/${commit_id}/SM${compile_arch//,/_}
-            fi
-            wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-            push_file=$(realpath bos_tools.py)
-            python --version
-            python -m pip install bce-python-sdk==0.9.29
-            cd FastDeploy/dist/
-            matches=($(ls fastdeploy*.whl))
-            if [ ${#matches[@]} -ne 1 ]; then
-              echo "Error: Found ${#matches[@]} matching files, expected exactly 1"
-              exit 1
-            fi
-            fd_wheel_name=${matches[0]}
-            echo "Found: $fd_wheel_name"
-            tree -L 3
-            python ${push_file} fastdeploy*.whl ${target_path}
-            target_path_stripped="${target_path#paddle-github-action/}"
-            WHEEL_PATH=https://paddle-github-action.bj.bcebos.com/${target_path_stripped}/${fd_wheel_name}
-            echo "wheel_path=${WHEEL_PATH}" >> $GITHUB_OUTPUT

.github/workflows/_ci_image_build.yml

@@ -1,73 +0,0 @@
-name: Docker Build
-description: "FastDeploy CI Image Build"
-
-on:
-  workflow_call:
-    inputs:
-      CI_DOCKER_IMAGE_NAME:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      DOCKER_IMAGE_NAME:
-        description: "Build Images"
-        required: false
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate"
-    outputs:
-      docker_name_precheck:
-        description: "Output path of the generated wheel"
-        value: ${{ jobs.docker_build.outputs.docker_name_precheck }}
-
-jobs:
-  docker_build:
-    runs-on: [self-hosted, Docker-Build]
-    outputs:
-      docker_name_precheck: ${{ steps.docker_build.outputs.docker_name_precheck }}
-    steps:
-      - name: Docker Build
-        id: docker_build
-        shell: bash
-        env:
-          docker_image_name: ${{ inputs.CI_DOCKER_IMAGE_NAME }}
-          docker_image: ${{ inputs.DOCKER_IMAGE_NAME }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-
-            # Docker Build
-            cd tools/dockerfile/
-            set -e
-            cp ../../requirements.txt ./
-            cp ../../scripts/unittest_requirement.txt ./
-            docker build -t ${docker_image_name} -f Dockerfile.ci . \
-                --network host \
-                --no-cache
-            docker push ${docker_image_name}
-            echo "docker_name_precheck=${docker_image_name}" >> $GITHUB_OUTPUT

.github/workflows/_clone_linux.yml

@@ -1,78 +0,0 @@
-name: FastDeploy Code Clone
-description: "FastDeploy clone and upload"
-
-on:
-  workflow_call:
-    inputs:
-        bos_dir:
-          type: string
-          required: false
-          default: 'FastDeploy'
-    outputs:
-      repo_archive_url:
-        description: "Compressed source code archive."
-        value: ${{ jobs.code-clone.outputs.repo_archive_url }}
-jobs:
-  code-clone:
-    runs-on:
-      group: HK-Clone
-    outputs:
-      repo_archive_url: ${{ steps.set_output.outputs.repo_archive_url }}
-    steps:
-      - name: Clone FastDeploy
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request'
-                && github.event.pull_request.base.ref
-                || github.ref_name }}
-          submodules: 'recursive'
-          fetch-depth: 1000
-
-      - name: Merge PR (if needed)
-        if: ${{ github.event_name == 'pull_request' }}
-        run: |
-          git config --global user.name "FastDeployCI"
-          git config --global user.email "fastdeploy_ci@example.com"
-          echo "Fetching and merging PR..."
-          git fetch origin pull/${{ github.event.pull_request.number }}/head:pr/${{ github.event.pull_request.number }}
-          git merge --no-ff pr/${{ github.event.pull_request.number }}
-          echo "PR Branch log "
-          git log --oneline -n 5 pr/${{ github.event.pull_request.number }}
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Code Info Show and Upload
-        id: set_output
-        env:
-          AK: paddle
-          SK: paddle
-        run: |
-          git config --unset http.https://github.com/.extraheader
-          git submodule foreach --recursive sh -c "git config --local --unset-all 'http.https://github.com/.extraheader'"
-          git submodule foreach --recursive sh -c "git config remote.origin.fetch '+refs/heads/*:refs/remotes/origin/*'"
-          echo "Current HEAD Log:"
-          git log --oneline -n 5
-          ls
-          cd ..
-          tar -zcf FastDeploy.tar.gz FastDeploy
-          if [[ "${{ github.event_name }}" == "pull_request" ]];then
-            commit_id=${{ github.event.pull_request.head.sha }}
-            pr_num=${{ github.event.pull_request.number }}
-            target_path=paddle-github-action/PR/FastDeploy/${pr_num}/${commit_id}
-          elif [[ "${{ github.ref_type }}" == "tag" ]]; then
-            commit_id=${{ github.sha }}
-            tag_name=${{ github.ref_name }}
-            target_path=paddle-github-action/TAG/FastDeploy/${tag_name}/${commit_id}
-          else
-            commit_id=${{ github.sha }}
-            branch_name=${{ github.ref_name }}
-            target_path=paddle-github-action/BRANCH/FastDeploy/${branch_name}/${commit_id}
-          fi
-          wget -O bos_tools.py -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} FastDeploy.tar.gz ${target_path}
-          target_path_stripped="${target_path#paddle-github-action/}"
-          REPO_ARCHIVE_URL=https://paddle-github-action.bj.bcebos.com/${target_path_stripped}/FastDeploy.tar.gz
-          echo "repo_archive_url=${REPO_ARCHIVE_URL}" >> $GITHUB_OUTPUT

.github/workflows/_logprob_test_linux.yml

@@ -1,185 +0,0 @@
-name: Run FastDeploy LogProb Tests
-description: "Run FastDeploy LogProb Tests"
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      PADDLETEST_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-        default: "https://xly-devops.bj.bcebos.com/PaddleTest/PaddleTest.tar.gz"
-      FASTDEPLOY_WHEEL_URL:
-        description: "URL of the FastDeploy Wheel."
-        required: true
-        type: string
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-      MODEL_CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-
-jobs:
-  run_tests_logprob:
-    runs-on: [self-hosted, GPU-h20-1Cards]
-    steps:
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          paddletest_archive_url: ${{ inputs.PADDLETEST_ARCHIVE_URL }}
-        run: |
-            docker pull ${docker_image}
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            -e "BASE_BRANCH=${BASE_BRANCH}" \
-            ${docker_image} /bin/bash -c '
-            rm -rf /workspace/*
-            '
-            wget -q ${paddletest_archive_url}
-            tar -xf PaddleTest.tar.gz
-            rm -rf PaddleTest.tar.gz
-            cd PaddleTest
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-      - name: logprob test
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fastdeploy_wheel_url: ${{ inputs.FASTDEPLOY_WHEEL_URL }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          MODEL_CACHE_DIR: ${{ inputs.MODEL_CACHE_DIR }}
-        run: |
-          runner_name="${{ runner.name }}"
-          CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-          DEVICES=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-          DEVICE_PORT=$(echo "$DEVICES" | cut -d',' -f1)
-
-          FLASK_PORT=$((42068 + DEVICE_PORT * 100))
-          FD_API_PORT=$((42088 + DEVICE_PORT * 100))
-          FD_ENGINE_QUEUE_PORT=$((42058 + DEVICE_PORT * 100))
-          FD_METRICS_PORT=$((42078 + DEVICE_PORT * 100))
-          FD_CACHE_QUEUE_PORT=$((42098 + DEVICE_PORT * 100))
-          echo "Test ENV Parameter:"
-          echo "========================================================="
-          echo "FLASK_PORT=${FLASK_PORT}"
-          echo "FD_API_PORT=${FD_API_PORT}"
-          echo "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}"
-          echo "FD_METRICS_PORT=${FD_METRICS_PORT}"
-          echo "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}"
-          echo "DEVICES=${DEVICES}"
-          echo "========================================================="
-
-          CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-          echo "CACHE_DIR is set to ${CACHE_DIR}"
-          if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-            touch "${CACHE_DIR}/gitconfig"
-          fi
-          if [ ! -d "${MODEL_CACHE_DIR}" ]; then
-            echo "Error: MODEL_CACHE_DIR '${MODEL_CACHE_DIR}' does not exist."
-            exit 1
-          fi
-
-          PORTS=($FLASK_PORT $FD_API_PORT $FD_ENGINE_QUEUE_PORT $FD_METRICS_PORT $FD_CACHE_QUEUE_PORT)
-          LOG_FILE="./port_cleanup_$(date +%Y%m%d_%H%M%S).log"
-          echo "==== LOG_FILE is ${LOG_FILE} ===="
-
-          echo "==== PORT CLEAN BEFORE TASK RUN ====" | tee -a $LOG_FILE
-
-          for port in "${PORTS[@]}"; do
-              PIDS=$(lsof -t -i :$port || true)
-              if [ -n "$PIDS" ]; then
-                  echo "Port $port is occupied by PID(s): $PIDS" | tee -a $LOG_FILE
-                  echo "$PIDS" | xargs -r kill -9
-                  echo "Port $port cleared" | tee -a $LOG_FILE
-              else
-                  echo "Port $port is free" | tee -a $LOG_FILE
-              fi
-          done
-
-          echo "==== PORT CLEAN COMPLETE ====" | tee -a $LOG_FILE
-
-          echo "========================================================="
-          echo "Ensuring no stale container named ${runner_name} ..."
-          if [ "$(docker ps -a -q -f name=${runner_name})" ]; then
-            echo "Removing stale container: ${runner_name}"
-            docker rm -f ${runner_name} || true
-          fi
-          docker run --rm --ipc=host --pid=host --net=host \
-          --name ${runner_name} \
-          -v $(pwd):/workspace \
-          -w /workspace \
-          -e fastdeploy_wheel_url=${fastdeploy_wheel_url} \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          -e "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}" \
-          -e "FLASK_PORT=${FLASK_PORT}" \
-          -v "${MODEL_CACHE_DIR}:/MODELDATA" \
-          -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-          -v "${CACHE_DIR}/.cache:/root/.cache" \
-          -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-          -e TZ="Asia/Shanghai" \
-          --gpus '"device='"${DEVICES}"'"' ${docker_image} /bin/bash -xc '
-          python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-
-          pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
-
-          python -m pip install ${fastdeploy_wheel_url}
-
-          wget https://paddle-qa.bj.bcebos.com/zhengtianyu/tools/llm-deploy-linux-amd64
-          chmod +x ./llm-deploy-linux-amd64
-          ./llm-deploy-linux-amd64 -python python3.10 \
-          -model_name ERNIE-4.5-0.3B-Paddle \
-          -model_path /MODELDATA \
-          --skip install
-
-          cd PaddleTest/framework/ServeTest
-          ps -ef | grep "${FD_ENGINE_QUEUE_PORT}" | grep -v grep | awk "{print \$2}" | xargs -r kill -9
-          python3.10 deploy.py > dd.log 2>&1 &
-          sleep 3
-          curl -X POST http://0.0.0.0:${FLASK_PORT}/start \
-              -H "Content-Type: application/json" \
-              -d "{\"--model\": \"/MODELDATA/ERNIE-4.5-0.3B-Paddle\"}"
-
-          curl -X POST http://localhost:${FLASK_PORT}/wait_for_infer?timeout=90
-          curl -s -o /dev/null -w "%{http_code}" -m 2 "http://0.0.0.0:${FD_API_PORT}/health"
-          curl -X POST "http://0.0.0.0:${FD_API_PORT}/v1/chat/completions" \
-            -H "Content-Type: application/json" \
-            -d "{\"messages\": [{\"role\": \"user\", \"content\": \"1+1=?\"}], \"logprobs\": true}"
-          set +e
-          rm -rf ./baseline_output
-          cp -r baseline/ERNIE-4.5-0.3B-Paddle ./baseline_output
-          LOGPROB_EXIT_CODE=0
-          python3.10 lanucher.py --request_template TOKEN_LOGPROB --url http://localhost:${FD_API_PORT}/v1/chat/completions  --case ./cases/demo.yaml  --concurrency 1 --name demo --exe logprob || LOGPROB_EXIT_CODE=$?
-          echo "LOGPROB_EXIT_CODE=${LOGPROB_EXIT_CODE}" > /workspace/exit_code.env
-          curl -X POST http://localhost:${FLASK_PORT}/stop
-          sleep 10s
-          cat *result.log
-          exit 0
-          '
-          if [ $? -ne 0 ];then
-            exit 1
-          fi
-
-          if [ -f exit_code.env ]; then
-            cat exit_code.env >> $GITHUB_ENV
-          fi
-      - name: logprob test result
-        if: ${{ env.LOGPROB_EXIT_CODE != 0 }}
-        shell: bash
-        run: |
-          echo "logprob test failed with exit code ${{ env.LOGPROB_EXIT_CODE }}"
-          exit 8

.github/workflows/_pre_ce_test.yml

@@ -1,148 +0,0 @@
-name: Pre-CE-Test
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddle:fastdeploy-ciuse-cuda126"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      FASTDEPLOY_WHEEL_URL:
-        description: "URL of the FastDeploy Wheel."
-        required: true
-        type: string
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-      MODEL_CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-
-jobs:
-  run_ce_cases:
-    runs-on: [self-hosted, PRE_CE_RUN_2Card]
-    timeout-minutes: 60
-    steps:
-      - name: Print current runner name
-        run: |
-          echo "Current runner name: ${{ runner.name }}"
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-            docker pull ${docker_image}
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-
-      - name: Run CI unittest
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_wheel_url: ${{ inputs.FASTDEPLOY_WHEEL_URL }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          MODEL_CACHE_DIR: ${{ inputs.MODEL_CACHE_DIR }}
-        run: |
-          runner_name="${{ runner.name }}"
-          CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-          DEVICES=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-          DEVICE_PORT=$(echo "$DEVICES" | cut -d',' -f1)
-
-          FLASK_PORT=$((42068 + DEVICE_PORT * 100))
-          FD_API_PORT=$((42088 + DEVICE_PORT * 100))
-          FD_ENGINE_QUEUE_PORT=$((42058 + DEVICE_PORT * 100))
-          FD_METRICS_PORT=$((42078 + DEVICE_PORT * 100))
-          FD_CACHE_QUEUE_PORT=$((42098 + DEVICE_PORT * 100))
-          echo "Test ENV Parameter:"
-          echo "========================================================="
-          echo "FLASK_PORT=${FLASK_PORT}"
-          echo "FD_API_PORT=${FD_API_PORT}"
-          echo "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}"
-          echo "FD_METRICS_PORT=${FD_METRICS_PORT}"
-          echo "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}"
-          echo "DEVICES=${DEVICES}"
-          echo "========================================================="
-
-          CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-          echo "CACHE_DIR is set to ${CACHE_DIR}"
-          if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-            touch "${CACHE_DIR}/gitconfig"
-          fi
-
-          PORTS=($FLASK_PORT $FD_API_PORT $FD_ENGINE_QUEUE_PORT $FD_METRICS_PORT $FD_CACHE_QUEUE_PORT)
-          LOG_FILE="./port_cleanup_$(date +%Y%m%d_%H%M%S).log"
-          echo "==== LOG_FILE is ${LOG_FILE} ===="
-
-          echo "==== PORT CLEAN BEFORE TASK RUN ====" | tee -a $LOG_FILE
-
-          for port in "${PORTS[@]}"; do
-              PIDS=$(lsof -t -i :$port || true)
-              if [ -n "$PIDS" ]; then
-                  echo "Port $port is occupied by PID(s): $PIDS" | tee -a $LOG_FILE
-                  echo "$PIDS" | xargs -r kill -9
-                  echo "Port $port cleared" | tee -a $LOG_FILE
-              else
-                  echo "Port $port is free" | tee -a $LOG_FILE
-              fi
-          done
-
-          echo "==== PORT CLEAN COMPLETE ====" | tee -a $LOG_FILE
-
-          echo "========================================================="
-          echo "Ensuring no stale container named ${runner_name} ..."
-          if [ "$(docker ps -a -q -f name=${runner_name})" ]; then
-            echo "Removing stale container: ${runner_name}"
-            docker rm -f ${runner_name} || true
-          fi
-
-          docker run --rm --net=host \
-          --name ${runner_name} \
-          -v $(pwd):/workspace \
-          -w /workspace \
-          -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-          -v "${CACHE_DIR}/.cache:/root/.cache" \
-          -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-          -v "${MODEL_CACHE_DIR}:/ModelData:ro" \
-          -e "MODEL_PATH=/ModelData" \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          -e "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}" \
-          -e "FLASK_PORT=${FLASK_PORT}" \
-          -e "fd_wheel_url=${fd_wheel_url}" \
-          --gpus "\"device=${DEVICES}\"" ${docker_image} /bin/bash -c '
-          git config --global --add safe.directory /workspace/FastDeploy
-          cd FastDeploy
-          python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-          python -m pip install ${fd_wheel_url}
-          bash scripts/run_pre_ce.sh
-          '

.github/workflows/_stable_test.yml

@@ -1,170 +0,0 @@
-name: Stable Test
-description: "Run Stable Tests"
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      FASTDEPLOY_WHEEL_URL:
-        description: "URL of the FastDeploy Wheel."
-        required: true
-        type: string
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-      MODEL_CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-
-jobs:
-  stable_tests:
-    runs-on: [self-hosted, GPU-h1z1-2Cards]
-    timeout-minutes: 60
-    steps:
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-            docker pull ${docker_image}
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-
-      - name: Run FastDeploy Stable Tests
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fastdeploy_wheel_url: ${{ inputs.FASTDEPLOY_WHEEL_URL }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          MODEL_CACHE_DIR: ${{ inputs.MODEL_CACHE_DIR }}
-        run: |
-          runner_name="${{ runner.name }}"
-          CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-          DEVICES=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-          DEVICE_PORT=$(echo "$DEVICES" | cut -d',' -f1)
-
-          FLASK_PORT=$((42068 + DEVICE_PORT * 100))
-          FD_API_PORT=$((42088 + DEVICE_PORT * 100))
-          FD_ENGINE_QUEUE_PORT=$((42058 + DEVICE_PORT * 100))
-          FD_METRICS_PORT=$((42078 + DEVICE_PORT * 100))
-          FD_CACHE_QUEUE_PORT=$((42038 + DEVICE_PORT * 100))
-          FD_INFERENCE_MSG_QUEUE_ID=$(( 42048 + DEVICE_PORT * 100))
-          echo "Test ENV Parameter:"
-          echo "========================================================="
-          echo "FLASK_PORT=${FLASK_PORT}"
-          echo "FD_API_PORT=${FD_API_PORT}"
-          echo "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}"
-          echo "FD_METRICS_PORT=${FD_METRICS_PORT}"
-          echo "FD_INFERENCE_MSG_QUEUE_ID=${FD_INFERENCE_MSG_QUEUE_ID}"
-          echo "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}"
-          echo "DEVICES=${DEVICES}"
-          echo "========================================================="
-
-          CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-          echo "CACHE_DIR is set to ${CACHE_DIR}"
-          if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-            touch "${CACHE_DIR}/gitconfig"
-          fi
-          if [ ! -d "${MODEL_CACHE_DIR}" ]; then
-            echo "Error: MODEL_CACHE_DIR '${MODEL_CACHE_DIR}' does not exist."
-            exit 1
-          fi
-
-          PORTS=($FLASK_PORT $FD_API_PORT $FD_ENGINE_QUEUE_PORT $FD_METRICS_PORT)
-          LOG_FILE="./port_cleanup_$(date +%Y%m%d_%H%M%S).log"
-          echo "==== LOG_FILE is ${LOG_FILE} ===="
-
-          echo "==== PORT CLEAN BEFORE TASK RUN ====" | tee -a $LOG_FILE
-
-          for port in "${PORTS[@]}"; do
-              PIDS=$(lsof -t -i :$port || true)
-              if [ -n "$PIDS" ]; then
-                  echo "Port $port is occupied by PID(s): $PIDS" | tee -a $LOG_FILE
-                  echo "$PIDS" | xargs -r kill -9
-                  echo "Port $port cleared" | tee -a $LOG_FILE
-              else
-                  echo "Port $port is free" | tee -a $LOG_FILE
-              fi
-          done
-
-          echo "==== PORT CLEAN COMPLETE ====" | tee -a $LOG_FILE
-
-          echo "========================================================="
-          echo "Ensuring no stale container named ${runner_name} ..."
-          if [ "$(docker ps -a -q -f name=${runner_name})" ]; then
-            echo "Removing stale container: ${runner_name}"
-            docker rm -f ${runner_name} || true
-          fi
-
-          docker run --rm --ipc=host --pid=host --net=host \
-          --name ${runner_name} \
-          -v $(pwd):/workspace \
-          -w /workspace \
-          -e fastdeploy_wheel_url=${fastdeploy_wheel_url} \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          -e "FLASK_PORT=${FLASK_PORT}" \
-          -e "FD_INFERENCE_MSG_QUEUE_ID=${FD_INFERENCE_MSG_QUEUE_ID}" \
-          -e "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}" \
-          -v "${MODEL_CACHE_DIR}:/MODELDATA" \
-          -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-          -v "${CACHE_DIR}/.cache:/root/.cache" \
-          -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-          -e TZ="Asia/Shanghai" \
-          --gpus '"device='"${DEVICES}"'"' ${docker_image} /bin/bash -xc '
-          python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-
-          pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
-
-          python -m pip install ${fastdeploy_wheel_url}
-          python -m pip install pytest
-
-          git config --global --add safe.directory /workspace/FastDeploy
-          cd FastDeploy
-          TEST_EXIT_CODE=0
-          pushd tests/ce/stable_cases
-          bash launch_model.sh /MODELDATA
-          bash run.sh || TEST_EXIT_CODE=1
-          popd
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}" >> /workspace/FastDeploy/exit_code.env
-          '
-          if [ -f ./FastDeploy/exit_code.env ]; then
-            source ./FastDeploy/exit_code.env
-            cat ./FastDeploy/exit_code.env >> $GITHUB_ENV
-          fi
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}"
-          exit ${TEST_EXIT_CODE}

.github/workflows/_unit_test_coverage.yml

@@ -1,319 +0,0 @@
-name: Coverage Check
-description: "Run FastDeploy Unit Tests and Coverage"
-
-on:
-  workflow_call:
-    inputs:
-      DOCKER_IMAGE:
-        description: "Build Images"
-        required: true
-        type: string
-        default: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:cuda126-py310"
-      FASTDEPLOY_ARCHIVE_URL:
-        description: "URL of the compressed FastDeploy code archive."
-        required: true
-        type: string
-      FASTDEPLOY_WHEEL_URL:
-        description: "URL of the FastDeploy Wheel."
-        required: true
-        type: string
-      CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-      MODEL_CACHE_DIR:
-        description: "Cache Dir Use"
-        required: false
-        type: string
-        default: ""
-    secrets:
-      github-token:
-        required: true
-
-jobs:
-  check_cov_skip:
-    uses: ./.github/workflows/check-bypass.yml
-    secrets:
-      github-token: ${{ secrets.github-token }}
-    with:
-      workflow-name: coverage
-
-  run_tests_with_coverage:
-    runs-on: [self-hosted, GPU-h1z1-2Cards]
-    timeout-minutes: 90
-    needs: check_cov_skip
-    if: needs.check_cov_skip.outputs.can-skip != 'true'
-    outputs:
-      diff_cov_file_url: ${{ steps.cov_upload.outputs.diff_cov_file_url }}
-      unittest_failed_url: ${{ steps.cov_upload.outputs.unittest_failed_url }}
-      diff_cov_result_json_url: ${{ steps.cov_upload.outputs.diff_cov_result_json_url }}
-    steps:
-      - name: Code Prepare
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-        run: |
-            set -x
-            REPO="https://github.com/${{ github.repository }}.git"
-            FULL_REPO="${{ github.repository }}"
-            REPO_NAME="${FULL_REPO##*/}"
-            BASE_BRANCH="${{ github.base_ref }}"
-            docker pull ${docker_image}
-            # Clean the repository directory before starting
-            docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
-            -e "REPO_NAME=${REPO_NAME}" \
-            ${docker_image} /bin/bash -c '
-              if [ -d ${REPO_NAME} ]; then
-                echo "Directory ${REPO_NAME} exists, removing it..."
-                rm -rf ${REPO_NAME}*
-              fi
-            '
-
-            wget -q ${fd_archive_url}
-            tar -xf FastDeploy.tar.gz
-            rm -rf FastDeploy.tar.gz
-            cd FastDeploy
-            git config --global user.name "FastDeployCI"
-            git config --global user.email "fastdeploy_ci@example.com"
-            git log -n 3 --oneline
-      - name: Run FastDeploy Unit Tests and Coverage
-        shell: bash
-        env:
-          docker_image: ${{ inputs.DOCKER_IMAGE }}
-          fd_wheel_url: ${{ inputs.FASTDEPLOY_WHEEL_URL }}
-          CACHE_DIR: ${{ inputs.CACHE_DIR }}
-          BASE_REF: ${{ github.event.pull_request.base.ref }}
-          MODEL_CACHE_DIR: ${{ inputs.MODEL_CACHE_DIR }}
-          IS_PR: ${{ github.event_name == 'pull_request' }}
-        run: |
-          if [[ "$IS_PR" == "true" ]]; then
-            echo "Running on PR"
-          else
-            echo "Not a PR"
-          fi
-          runner_name="${{ runner.name }}"
-          CARD_ID=$(echo "${runner_name}" | awk -F'-' '{print $NF}')
-          DEVICES=$(echo "$CARD_ID" | fold -w1 | paste -sd,)
-          DEVICE_PORT=$(echo "$DEVICES" | cut -d',' -f1)
-
-          FLASK_PORT=$((42068 + DEVICE_PORT * 100))
-          FD_API_PORT=$((42088 + DEVICE_PORT * 100))
-          FD_ENGINE_QUEUE_PORT=$((42058 + DEVICE_PORT * 100))
-          FD_METRICS_PORT=$((42078 + DEVICE_PORT * 100))
-          FD_CACHE_QUEUE_PORT=$((42098 + DEVICE_PORT * 100))
-          echo "Test ENV Parameter:"
-          echo "========================================================="
-          echo "FLASK_PORT=${FLASK_PORT}"
-          echo "FD_API_PORT=${FD_API_PORT}"
-          echo "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}"
-          echo "FD_METRICS_PORT=${FD_METRICS_PORT}"
-          echo "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}"
-          echo "DEVICES=${DEVICES}"
-          echo "========================================================="
-
-          CACHE_DIR="${CACHE_DIR:-$(dirname "$(dirname "${{ github.workspace }}")")}"
-          echo "CACHE_DIR is set to ${CACHE_DIR}"
-          if [ ! -f "${CACHE_DIR}/gitconfig" ]; then
-            touch "${CACHE_DIR}/gitconfig"
-          fi
-
-          PORTS=($FLASK_PORT $FD_API_PORT $FD_ENGINE_QUEUE_PORT $FD_METRICS_PORT $FD_CACHE_QUEUE_PORT)
-          LOG_FILE="./port_cleanup_$(date +%Y%m%d_%H%M%S).log"
-          echo "==== LOG_FILE is ${LOG_FILE} ===="
-
-          echo "==== PORT CLEAN BEFORE TASK RUN ====" | tee -a $LOG_FILE
-
-          for port in "${PORTS[@]}"; do
-              PIDS=$(lsof -t -i :$port || true)
-              if [ -n "$PIDS" ]; then
-                  echo "Port $port is occupied by PID(s): $PIDS" | tee -a $LOG_FILE
-                  echo "$PIDS" | xargs -r kill -9
-                  echo "Port $port cleared" | tee -a $LOG_FILE
-              else
-                  echo "Port $port is free" | tee -a $LOG_FILE
-              fi
-          done
-
-          echo "==== PORT CLEAN COMPLETE ====" | tee -a $LOG_FILE
-
-          echo "========================================================="
-          echo "Ensuring no stale container named ${runner_name} ..."
-          if [ "$(docker ps -a -q -f name=${runner_name})" ]; then
-            echo "Removing stale container: ${runner_name}"
-            docker rm -f ${runner_name} || true
-          fi
-
-          docker run --rm --net=host \
-          --name ${runner_name} \
-          --cap-add=SYS_PTRACE --shm-size=64G \
-          -v $(pwd):/workspace -w /workspace \
-          -v "${CACHE_DIR}/gitconfig:/etc/gitconfig:ro" \
-          -v "${CACHE_DIR}/.cache:/root/.cache" \
-          -v "${CACHE_DIR}/ConfigDir:/root/.config" \
-          -v "${MODEL_CACHE_DIR}:/ModelData:ro" \
-          -e "MODEL_PATH=/ModelData" \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          -e "FLASK_PORT=${FLASK_PORT}" \
-          -e "FD_CACHE_QUEUE_PORT=${FD_CACHE_QUEUE_PORT}" \
-          -e TZ="Asia/Shanghai" \
-          -e "fd_wheel_url=${fd_wheel_url}" \
-          -e "BASE_REF=${BASE_REF}" \
-          -e "IS_PR=${IS_PR}" \
-          --gpus "\"device=${DEVICES}\"" ${docker_image} /bin/bash -c '
-
-          git config --global --add safe.directory /workspace/FastDeploy
-          cd FastDeploy
-          git diff origin/${BASE_REF}..HEAD --unified=0 > diff.txt
-          python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
-          pip config set global.extra-index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
-
-          python -m pip install -r scripts/unittest_requirement.txt
-          python -m pip install ${fd_wheel_url}
-          rm -rf fastdeploy
-          # coverage subprocess use
-          python -m pip install ${fd_wheel_url} --no-deps --target=/workspace/FastDeploy
-          export PYTHONPATH=/workspace/FastDeploy/
-          if [ -d "tests/plugins" ]; then
-              cd tests/plugins
-              python setup.py install
-              cd ../..
-          else
-              echo "Warning: tests/plugins directory not found, skipping setup.py install"
-          fi
-          export COVERAGE_FILE=/workspace/FastDeploy/coveragedata/.coverage
-          export COVERAGE_RCFILE=/workspace/FastDeploy/scripts/.coveragerc
-          TEST_EXIT_CODE=0
-          bash scripts/coverage_run.sh || TEST_EXIT_CODE=8
-          echo "TEST_EXIT_CODE=${TEST_EXIT_CODE}" >> exit_code.env
-          coverage combine coveragedata/ || echo "No data to combine"
-          coverage report
-          coverage xml -o python_coverage_all.xml
-          COVERAGE_EXIT_CODE=0
-          if [[ "$IS_PR" == "true" ]]; then
-            diff-cover python_coverage_all.xml --diff-file=diff.txt --fail-under=80 --json-report diff_coverage.json || COVERAGE_EXIT_CODE=9
-            python scripts/generate_diff_coverage_xml.py diff.txt python_coverage_all.xml
-          else
-            echo "Not a PR, skipping diff-cover"
-          fi
-          echo "COVERAGE_EXIT_CODE=${COVERAGE_EXIT_CODE}" >> exit_code.env
-          '
-          if [ -f FastDeploy/exit_code.env ]; then
-            cat FastDeploy/exit_code.env >> $GITHUB_ENV
-          fi
-      - name: Upload unit resule and diff coverage to bos
-        id: cov_upload
-        shell: bash
-        run: |
-          cd FastDeploy
-          commit_id=${{ github.event.pull_request.head.sha }}
-          pr_num=${{ github.event.pull_request.number }}
-          target_path=paddle-github-action/PR/FastDeploy/${pr_num}/${commit_id}/SM${compile_arch//,/_}
-          wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py -O bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          diff_cov_file="diff_coverage.xml"
-          if [ -f ${diff_cov_file} ];then
-            python ${push_file} ${diff_cov_file} ${target_path}/CoverageData
-            target_path_stripped="${target_path#paddle-github-action/}"
-            DIFF_COV_FILE_URL=https://paddle-github-action.bj.bcebos.com/${target_path_stripped}/CoverageData/${diff_cov_file}
-            echo "diff_cov_file_url=${DIFF_COV_FILE_URL}" >> $GITHUB_OUTPUT
-            echo "diff_cov_file_url=${DIFF_COV_FILE_URL}" >> $GITHUB_ENV
-          fi
-          diff_cov_result_json="diff_coverage.json"
-          if [ -f ${diff_cov_result_json} ];then
-            python ${push_file} ${diff_cov_result_json} ${target_path}/CoverageData
-            target_path_stripped="${target_path#paddle-github-action/}"
-            DIFF_COV_JSON_URL=https://paddle-github-action.bj.bcebos.com/${target_path_stripped}/CoverageData/${diff_cov_result_json}
-            echo "diff_cov_result_json_url=${DIFF_COV_JSON_URL}" >> $GITHUB_OUTPUT
-            echo "diff_cov_result_json_url=${DIFF_COV_JSON_URL}" >> $GITHUB_ENV
-          fi
-          unittest_result="failed_tests.log"
-          if [ -s ${unittest_result} ];then
-            python ${push_file} ${unittest_result} ${target_path}/UnitTestResult
-            target_path_stripped="${target_path#paddle-github-action/}"
-            UNIT_TEST_RESULT_URL=https://paddle-github-action.bj.bcebos.com/${target_path_stripped}/UnitTestResult/${unittest_result}
-            echo "unittest_failed_url=${UNIT_TEST_RESULT_URL}" >> $GITHUB_OUTPUT
-            echo "unittest_failed_url=${UNIT_TEST_RESULT_URL}" >> $GITHUB_ENV
-          fi
-      - name: Check Unit Test Success
-        shell: bash
-        run: |
-          cd FastDeploy
-          if [ "$TEST_EXIT_CODE" -eq 8 ]; then
-            filename=$(basename "$unittest_failed_url")
-            if [ -z "${unittest_failed_url}" ]; then
-              echo "No diff unit failed file URL provided."
-            else
-              rm -rf "${filename}"
-              wget -O ${filename} ${unittest_failed_url} || echo "Download unittest file failed, but continuing..."
-            fi
-            echo "Unit tests failed (exit code 8)"
-            if [ -f "${filename}" ];then
-              echo "Failed test cases:"
-              cat "${filename}"
-            fi
-            exit "$TEST_EXIT_CODE"
-          fi
-          echo "All tests passed"
-
-      - name: Verify Code Coverage Threshold (80%)
-        if: ${{ github.event_name == 'pull_request' }}
-        shell: bash
-        run: |
-          cd FastDeploy
-          if [ "$COVERAGE_EXIT_CODE" -eq 9 ]; then
-            echo "Coverage generation failed (exit code 9)"
-            filename=$(basename "$diff_cov_result_json_url")
-            if [ -z "${diff_cov_result_json_url}" ]; then
-              echo "No diff cov result file URL provided."
-            else
-              rm -rf "${filename}"
-              wget -O ${filename} ${diff_cov_result_json_url} || echo "Download cov json file failed, but continuing..."
-            fi
-            if [ -f "${filename}" ];then
-              echo "Failed test cases:"
-              if command -v jq >/dev/null 2>&1; then
-                  jq . "${filename}"
-              else
-                  cat "${filename}"
-              fi
-            fi
-            exit "$COVERAGE_EXIT_CODE"
-          fi
-          echo "coverage passed"
-          exit 0
-
-  diff_coverage_report:
-    needs: run_tests_with_coverage
-    if: always()
-    runs-on: ubuntu-latest
-    env:
-      fd_archive_url: ${{ inputs.FASTDEPLOY_ARCHIVE_URL }}
-    steps:
-      - name: coverage diff file download
-        shell: bash
-        env:
-          diff_cov_file_url: ${{ needs.run_tests_with_coverage.outputs.diff_cov_file_url }}
-        run: |
-          wget ${fd_archive_url}
-          tar -xf FastDeploy.tar.gz
-          cd FastDeploy
-          if [ -z "${diff_cov_file_url}" ]; then
-            echo "No diff coverage file URL provided."
-            exit 0
-          fi
-          wget "${diff_cov_file_url}" -O ./diff_coverage.xml || echo "Download cov file failed, but continuing..."
-      - name: Upload diff coverage report
-        if: ${{ needs.run_tests_with_coverage.outputs.diff_cov_file_url != null && needs.run_tests_with_coverage.outputs.diff_cov_file_url != '' }}
-        uses: codecov/codecov-action@v5
-        with:
-          files: ./FastDeploy/diff_coverage.xml
-          name: python diff coverage
-          verbose: true
-          disable_search: true
-          commit_parent: false
-          flags: diff

.github/workflows/approve.yml

@@ -1,42 +0,0 @@
-name: Approval
-
-on:
-  pull_request:
-    branches:
-      - develop
-      - 'release/*'
-
-env:
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-
-jobs:
-  Approval:
-    name: Approval
-    if: ${{ github.repository_owner == 'PaddlePaddle' }}
-    runs-on: ubuntu-latest
-    env:
-      PR_ID: ${{ github.event.pull_request.number }}
-      BRANCH: ${{ github.event.pull_request.base.ref }}
-    steps:
-      - name: Checkout base repo
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.pull_request.base.ref }}
-          fetch-depth: 1000
-
-      - name: Merge PR to test branch
-        run: |
-          git fetch origin pull/${PR_ID}/merge
-          git checkout -b test FETCH_HEAD
-          git log -n 3 --oneline
-          git remote add upstream https://github.com/PaddlePaddle/FastDeploy.git
-          git fetch upstream $BRANCH
-
-      - name: Setup python3.10
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-
-      - name: Run approval check script
-        run: |
-          bash scripts/check_approval.sh

.github/workflows/ce_job.yml

@@ -1,248 +0,0 @@
-name: CE Compile Job
-
-on:
-  workflow_dispatch:
-  push:
-    branches:
-      - develop
-      - 'release/*'
-permissions: read-all
-
-concurrency:
-  group: CE-Job-${{ github.ref }}-${{ github.sha }}
-  cancel-in-progress: true
-
-jobs:
-  ce_job_pre_check:
-    runs-on: ubuntu-latest
-    env:
-      COMPILE_BRANCH: ${{ vars.COMPILE_BRANCH }}
-      CE_COMPILE_SELECTION: ${{ vars.CE_COMPILE_SELECTION }}
-      COMPILE_USE_PADDLE_WHL_URL_MAPPINGS: ${{ vars.COMPILE_USE_PADDLE_WHL_URL_MAPPINGS }}
-    outputs:
-      branch_match: ${{ steps.set_output.outputs.branch_match }}
-      compile_use_paddle_whl_url: ${{ steps.set_output.outputs.compile_use_paddle_whl_url }}
-      sm8689_match: ${{ steps.set_output.outputs.sm8689_match }}
-      sm8090_match: ${{ steps.set_output.outputs.sm8090_match }}
-
-    steps:
-      - name: Set Version
-        id: set_output
-        env:
-          COMPILE_BRANCH: ${{ env.COMPILE_BRANCH }}
-          CE_COMPILE_SELECTION: ${{ env.CE_COMPILE_SELECTION }}
-          COMPILE_USE_PADDLE_WHL_URL_MAPPINGS: ${{ env.COMPILE_USE_PADDLE_WHL_URL_MAPPINGS }}
-          GITHUB_REF_NAME: ${{ github.ref_name }}
-        run: |
-          # 选择要触发编译任务的分支 done
-          # 选择指定分支要编译的任务 8090或者8689
-          # 指定分支编译要使用的Paddle的安装包,默认使用nightly最新的
-
-          IFS=',' read -ra BRANCHES <<< "$COMPILE_BRANCH"
-          MATCH=false
-          for b in "${BRANCHES[@]}"; do
-            if [[ "$b" == "${GITHUB_REF_NAME}" ]]; then
-              MATCH=true
-              break
-            fi
-          done
-          echo "branch_match=$MATCH" >> $GITHUB_OUTPUT
-
-          # 通过变量CE_COMPILE_SELECTION中的映射关系,决定分支是编译sm8090还是sm8689
-          for pair in $(echo "$CE_COMPILE_SELECTION" | tr ';' ' '); do
-            branch=$(echo "$pair" | cut -d',' -f1)
-            compile_task_list=$(echo "$pair" | cut -d',' -f2)
-
-            if [[ "$branch" == "$GITHUB_REF_NAME" ]]; then
-
-              # 判断里面是否包含 sm8090 或 sm8689
-              if [[ "$compile_task_list" == *"sm8090"* ]]; then
-                echo "sm8090_match=true" >> $GITHUB_OUTPUT
-              fi
-              if [[ "$compile_task_list" == *"sm8689"* ]]; then
-                echo "sm8689_match=true" >> $GITHUB_OUTPUT
-              fi
-              break
-            fi
-          done
-
-          # 通过变量COMPILE_USE_PADDLE_WHL_URL_MAPPINGS中的映射关系,决定是否是安装指定版本的Paddle还是直接安装URL
-          for pair in $(echo $COMPILE_USE_PADDLE_WHL_URL_MAPPINGS | tr ';' ' '); do
-            branch=$(echo "$pair" | cut -d',' -f1)
-            paddle_whl_url=$(echo "$pair" | cut -d',' -f2)
-            if [[ "$branch" == "${{ github.ref_name }}" ]]; then
-              FOUND_PADDLE_URL="$paddle_whl_url"
-              echo "compile_use_paddle_whl_url=${FOUND_PADDLE_URL}" >> $GITHUB_OUTPUT
-              break
-            fi
-          done
-
-  print_ce_job_pre_check_outputs:
-    runs-on: ubuntu-latest
-    needs: ce_job_pre_check
-    steps:
-      - name: Print outputs as JSON
-        run: |
-          echo '${{ toJSON(needs.ce_job_pre_check.outputs) }}'
-
-
-  clone:
-    environment: CodeSync
-    name: FD-Clone-Linux
-    runs-on: ubuntu-latest
-    needs: ce_job_pre_check
-    if: ${{ needs.ce_job_pre_check.outputs.branch_match == 'true' }}
-    outputs:
-      repo_archive_url: ${{ steps.set_output.outputs.repo_archive_url }}
-    steps:
-      - name: Clone FastDeploy
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event_name == 'pull_request'
-                && github.event.pull_request.base.ref
-                || github.ref_name }}
-          submodules: 'recursive'
-          fetch-depth: 1000
-
-      - name: Python Setup
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Code Info Show and Upload
-        id: set_output
-        env:
-          AK: ${{ secrets.BOS_AK }}
-          SK: ${{ secrets.BOS_SK }}
-        run: |
-          git config --unset http.https://github.com/.extraheader
-          git submodule foreach --recursive sh -c "git config --local --unset-all 'http.https://github.com/.extraheader'"
-          git submodule foreach --recursive sh -c "git config remote.origin.fetch '+refs/heads/*:refs/remotes/origin/*'"
-          echo "Current HEAD Log:"
-          git log --oneline -n 5
-          ls
-          cd ..
-          tar -zcf FastDeploy.tar.gz FastDeploy
-          commit_id=${{ github.sha }}
-          branch_name=${{ github.ref_name }}
-          target_path=paddle-qa/BRANCH/FastDeploy/${branch_name}/${commit_id}
-          wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} FastDeploy.tar.gz ${target_path}
-          target_path_stripped="${target_path#paddle-qa/}"
-          REPO_ARCHIVE_URL=https://paddle-qa.bj.bcebos.com/${target_path_stripped}/FastDeploy.tar.gz
-          echo "repo_archive_url=${REPO_ARCHIVE_URL}" >> $GITHUB_OUTPUT
-
-  resultshow:
-    name: Show Code Archive Output
-    needs: clone
-    runs-on: ubuntu-latest
-    steps:
-      - name: Print wheel path
-        run: |
-          echo "The code archive is located at: ${{ needs.clone.outputs.repo_archive_url }}"
-
-  build_sm8090:
-    name: BUILD_SM8090
-    needs: [clone, ce_job_pre_check]
-    if: ${{ needs.ce_job_pre_check.outputs.sm8090_match == 'true' }}
-    uses: ./.github/workflows/_build_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      COMPILE_ARCH: "80,90"
-      WITH_NIGHTLY_BUILD: OFF
-      FD_VERSION: 0.0.0
-      PADDLE_WHL_URL: ${{ needs.ce_job_pre_check.outputs.compile_use_paddle_whl_url }}
-
-  build_sm8689:
-    name: BUILD_SM8689
-    needs: [clone, ce_job_pre_check]
-    if: ${{ needs.ce_job_pre_check.outputs.sm8689_match == 'true' }}
-    uses: ./.github/workflows/_build_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      COMPILE_ARCH: "86,89"
-      WITH_NIGHTLY_BUILD: OFF
-      FD_VERSION: 0.0.0
-      PADDLE_WHL_URL: ${{ needs.ce_job_pre_check.outputs.compile_use_paddle_whl_url }}
-
-  ce_upload_sm8090:
-    environment: CodeSync
-    name: CE_UPLOAD
-    needs: build_sm8090
-    runs-on: ubuntu-latest
-    env:
-      AK: ${{ secrets.BOS_AK }}
-      SK: ${{ secrets.BOS_SK }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      COMPILE_ARCH: "80,90"
-    steps:
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Wheel Info Show and Upload
-        run: |
-          echo "The wheel is located at: ${{ needs.build_sm8090.outputs.wheel_path }}"
-          wget -q --no-check-certificate ${{ needs.build_sm8090.outputs.wheel_path }}
-          filename=$(basename ${{ needs.build_sm8090.outputs.wheel_path }})
-
-          commit_id=${{ github.sha }}
-          branch_name=${{ github.ref_name }}
-          target_path=paddle-qa/paddle-pipeline/FastDeploy_ActionCE/SM${COMPILE_ARCH//,/_}/${branch_name}/${commit_id}
-
-          wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} ${filename} ${target_path}
-          target_path_stripped="${target_path#paddle-qa/}"
-          WHEEL_PATH=https://paddle-qa.bj.bcebos.com/${target_path_stripped}/${filename}
-
-          target_path_latest=paddle-qa/paddle-pipeline/FastDeploy_ActionCE/SM${COMPILE_ARCH//,/_}/${branch_name}/latest
-          python ${push_file} ${filename} ${target_path_latest}
-          target_path_stripped_latest="${target_path_latest#paddle-qa/}"
-          WHEEL_PATH_LATEST=https://paddle-qa.bj.bcebos.com/${target_path_stripped_latest}/${filename}
-          echo "commit wheel url is ${WHEEL_PATH}"
-          echo "latest wheel url is ${WHEEL_PATH_LATEST}"
-
-  ce_upload_sm8689:
-    environment: CodeSync
-    name: CE_UPLOAD
-    needs: build_sm8689
-    runs-on: ubuntu-latest
-    env:
-      AK: ${{ secrets.BOS_AK }}
-      SK: ${{ secrets.BOS_SK }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8689.outputs.wheel_path }}
-      COMPILE_ARCH: "86,89"
-    steps:
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Wheel Info Show and Upload
-        run: |
-          echo "The wheel is located at: ${{ needs.build_sm8689.outputs.wheel_path }}"
-          wget -q --no-check-certificate ${{ needs.build_sm8689.outputs.wheel_path }}
-          filename=$(basename ${{ needs.build_sm8689.outputs.wheel_path }})
-
-          commit_id=${{ github.sha }}
-          branch_name=${{ github.ref_name }}
-          target_path=paddle-qa/paddle-pipeline/FastDeploy_ActionCE/SM${COMPILE_ARCH//,/_}/${branch_name}/${commit_id}
-
-          wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} ${filename} ${target_path}
-          target_path_stripped="${target_path#paddle-qa/}"
-          WHEEL_PATH=https://paddle-qa.bj.bcebos.com/${target_path_stripped}/${filename}
-
-          target_path_latest=paddle-qa/paddle-pipeline/FastDeploy_ActionCE/SM${COMPILE_ARCH//,/_}/${branch_name}/latest
-          python ${push_file} ${filename} ${target_path_latest}
-          target_path_stripped_latest="${target_path_latest#paddle-qa/}"
-          WHEEL_PATH_LATEST=https://paddle-qa.bj.bcebos.com/${target_path_stripped_latest}/${filename}
-          echo "commit wheel url is ${WHEEL_PATH}"
-          echo "latest wheel url is ${WHEEL_PATH_LATEST}"

.github/workflows/check-bypass.yml

@@ -1,51 +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: ubuntu-latest
-    permissions:
-      contents: read
-    env:
-      CI_TEAM_MEMBERS: '["yuanlehome","YuanRisheng","Jiang-Jia-Jun","DDDivano","XieYunshen"]'
-    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/ci.yml

@@ -1,4 +1,4 @@
-name: CI_ILUVATAR
+name: CI
 
 on:
   pull_request:
@@ -6,13 +6,12 @@ on:
   workflow_dispatch:
 
 concurrency:
-  group: ${{ github.event.pull_request.number }}-iluvatar-ci
+  group: ${{ github.event.pull_request.number }}
   cancel-in-progress: true
 
 jobs:
-  CI_ILUVATAR:
-    runs-on:
-      group: IXUCA
+  build:
+    runs-on: [self-hosted, GPU-L20-4Card]
     steps:
       - name: Print current runner name
         run: |
@@ -23,7 +22,7 @@ jobs:
 
       - name: Code Checkout
         env:
-          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/device/paddle-ixuca:latest
+          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddle:fastdeploy-ciuse-cuda126
         run: |
           REPO="https://github.com/${{ github.repository }}.git"
           FULL_REPO="${{ github.repository }}"
@@ -52,7 +51,7 @@ jobs:
 
       - name: Run CI unittest
         env:
-          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/device/paddle-ixuca:latest
+          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddle:fastdeploy-ciuse-cuda126
         run: |
           runner_name="${{ runner.name }}"
           last_char="${runner_name: -1}"
@@ -60,7 +59,7 @@ jobs:
           if [[ "$last_char" =~ [0-3] ]]; then
             gpu_id="$last_char"
           else
-            gpu_id="0"
+            gpu_id="0"  
           fi
           FD_API_PORT=$((9180 + gpu_id * 100))
           FD_ENGINE_QUEUE_PORT=$((9150 + gpu_id * 100))
@@ -68,18 +67,17 @@ jobs:
 
           PARENT_DIR=$(dirname "$WORKSPACE")
           echo "PARENT_DIR:$PARENT_DIR"
-          docker run --rm --net=host --pid=host --cap-add=ALL --privileged --shm-size=64G  \
-          -v /usr/src:/usr/src -v /lib/modules:/lib/modules -v /dev:/dev \
-          -v $(pwd):/workspace -w /workspace \
-          -v "/data1/fastdeploy:/data1/fastdeploy" \
-          -e "MODEL_PATH=/ssd3/model" \
-          -e "http_proxy=$(git config --global --get http.proxy)" \
-          -e "https_proxy=$(git config --global --get https.proxy)" \
+          docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
+          -v "/ssd4/GithubActions/gitconfig:/etc/gitconfig:ro" \
+          -v "/ssd4/GithubActions/ModelData:/ModelData:ro" \
+          -v "/ssd4/GithubActions/CacheDir:/root/.cache" \
+          -v "/ssd4/GithubActions/ConfigDir:/root/.config" \
+          -e "MODEL_PATH=/ModelData" \
           -e "FD_API_PORT=${FD_API_PORT}" \
           -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
           -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-           ${docker_image} /bin/bash -c "
+          --gpus device=${gpu_id} ${docker_image} /bin/bash -c "
           git config --global --add safe.directory /workspace/FastDeploy
           cd FastDeploy
-          bash scripts/run_ci_iluvatar.sh
-          "
+          bash scripts/run_ci.sh
+          "

.github/workflows/ci_gcu.yml

@@ -1,98 +0,0 @@
-name: CI_GCU
-
-on:
-  pull_request:
-    branches:
-      - develop
-      - 'release/*'
-  workflow_dispatch:
-
-concurrency:
-  group: ${{ github.event.pull_request.number }}-gcu-ci
-  cancel-in-progress: true
-
-jobs:
-  CI_GCU:
-    runs-on:
-      group: GCU
-    steps:
-      - name: Print current runner name
-        run: |
-          echo "Current runner name: ${{ runner.name }}"
-
-      - name: Code Checkout
-        env:
-          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/device/paddle-gcu:topsrider3.5.102-ubuntu20-x86_64-gcc84
-        run: |
-          REPO="https://github.com/${{ github.repository }}.git"
-          FULL_REPO="${{ github.repository }}"
-          REPO_NAME="${FULL_REPO##*/}"
-          BASE_BRANCH="${{ github.base_ref }}"
-          # Clean the repository directory before starting
-          docker run --rm --net=host -v $(pwd):/workspace \
-          -v ${{ github.workspace }}/../../..:${{ github.workspace }}/../../.. \
-          -w /workspace \
-          -e "REPO_NAME=${REPO_NAME}" \
-          -e "BASE_BRANCH=${BASE_BRANCH}" \
-          ${docker_image} /bin/bash -c '
-            if [ -d ${REPO_NAME} ]; then
-              echo "Directory ${REPO_NAME} exists, removing it..."
-              rm -rf ${REPO_NAME}
-            fi
-          '
-          git config --global user.name "FastDeployCI"
-          git config --global user.email "fastdeploy_ci@example.com"
-          source ${{ github.workspace }}/../../../proxy
-          git clone ${REPO} ${REPO_NAME} -b ${BASE_BRANCH}
-          cd FastDeploy
-          if [ "${{ github.event_name }}" = "pull_request" ]; then
-            git fetch origin pull/${{ github.event.pull_request.number }}/head:pr/${{ github.event.pull_request.number }}
-            git merge pr/${{ github.event.pull_request.number }}
-            git log -n 3 --oneline
-          else
-            git checkout ${{ github.sha }}
-            git log -n 3 --oneline
-          fi
-          echo "Copy models..."
-          sudo mkdir -p ci_models && sudo cp -r /work/deps/ERNIE-4.5-21B-A3B-Paddle ci_models
-          echo "Copy models done."
-
-      - name: Run CI unittest
-        env:
-          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/device/paddle-gcu:topsrider3.5.102-ubuntu20-x86_64-gcc84
-        run: |
-          runner_name="${{ runner.name }}"
-          last_char="${runner_name: -1}"
-
-          if [[ "$last_char" =~ [0-3] ]]; then
-            gcu_id="$last_char"
-          else
-            gcu_id="0"
-          fi
-          FD_API_PORT=$((9180 + gcu_id * 100))
-          FD_ENGINE_QUEUE_PORT=$((9150 + gcu_id * 100))
-          FD_METRICS_PORT=$((9170 + gcu_id * 100))
-
-          PARENT_DIR=$(dirname "$WORKSPACE")
-          echo "PARENT_DIR:$PARENT_DIR"
-          echo "Install drivers..."
-          cd /work/deps
-          sudo bash TopsRider_i3x_*_deb_amd64.run --driver --no-auto-load -y
-          cd -
-          echo "Create docker..."
-          docker run --rm --network=host --ipc=host --privileged \
-          -v $(pwd):/workspace \
-          -v /home:/home \
-          -v /work:/work \
-          -w /workspace \
-          -e "MODEL_PATH=./ci_models" \
-          -e "http_proxy=$(git config --global --get http.proxy)" \
-          -e "https_proxy=$(git config --global --get https.proxy)" \
-          -e "FD_API_PORT=${FD_API_PORT}" \
-          -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
-          -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
-          ${docker_image} /bin/bash -c "
-          git config --global --add safe.directory /workspace/FastDeploy
-          cd FastDeploy
-          bash scripts/run_ci_gcu.sh
-          "

.github/workflows/ci_image_update.yml

@@ -1,174 +0,0 @@
-name: CI Images Build
-
-on:
-  workflow_dispatch:
-  schedule:
-    - cron: '0 18 * * *'   # 2:00 AM China Standard Time (UTC+8)
-
-permissions: read-all
-
-concurrency:
-  group: CI-Images-Build-${{ github.ref }}-${{ github.sha }}
-  cancel-in-progress: true
-
-
-jobs:
-  clone:
-    environment: CodeSync
-    name: FD-Clone-Linux
-    runs-on: ubuntu-latest
-    outputs:
-      repo_archive_url: ${{ steps.set_output.outputs.repo_archive_url }}
-    steps:
-      - name: Clone FastDeploy
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.ref_name }}
-          submodules: 'recursive'
-          fetch-depth: 1000
-
-      - name: Python Setup
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Code Info Show and Upload
-        id: set_output
-        env:
-          AK: ${{ secrets.BOS_AK }}
-          SK: ${{ secrets.BOS_SK }}
-        run: |
-          git config --unset http.https://github.com/.extraheader
-          git submodule foreach --recursive sh -c "git config --local --unset-all 'http.https://github.com/.extraheader'"
-          git submodule foreach --recursive sh -c "git config remote.origin.fetch '+refs/heads/*:refs/remotes/origin/*'"
-          echo "Current HEAD Log:"
-          git log --oneline -n 5
-          ls
-          cd ..
-          tar -zcf FastDeploy.tar.gz FastDeploy
-          if [[ "${{ github.ref_type }}" == "tag" ]]; then
-            commit_id=${{ github.sha }}
-            tag_name=${{ github.ref_name }}
-            target_path=paddle-qa/TAG/FastDeploy/${tag_name}/${commit_id}
-          else
-            commit_id=${{ github.sha }}
-            branch_name=${{ github.ref_name }}
-            target_path=paddle-qa/BRANCH/FastDeploy/${branch_name}/${commit_id}
-          fi
-          wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} FastDeploy.tar.gz ${target_path}
-          target_path_stripped="${target_path#paddle-qa/}"
-          REPO_ARCHIVE_URL=https://paddle-qa.bj.bcebos.com/${target_path_stripped}/FastDeploy.tar.gz
-          echo "repo_archive_url=${REPO_ARCHIVE_URL}" >> $GITHUB_OUTPUT
-
-  resultshow:
-    name: Show Code Archive Output
-    needs: clone
-    runs-on: ubuntu-latest
-    steps:
-      - name: Print wheel path
-        run: |
-          echo "The code archive is located at: ${{ needs.clone.outputs.repo_archive_url }}"
-
-  ci_image_build:
-    name: CI Images Build
-    needs: clone
-    uses: ./.github/workflows/_ci_image_build.yml
-    with:
-      CI_DOCKER_IMAGE_NAME: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate-precheck
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-
-
-  build_sm8090:
-    name: BUILD_SM8090
-    needs: [clone, ci_image_build]
-    uses: ./.github/workflows/_build_linux.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      COMPILE_ARCH: "90"
-      WITH_NIGHTLY_BUILD: ${{ needs.publish_pre_check.outputs.with_nightly_build }}
-      FD_VERSION: ${{ needs.publish_pre_check.outputs.fd_version }}
-      PADDLEVERSION: ${{ needs.publish_pre_check.outputs.compile_use_paddle_version }}
-      PADDLE_WHL_URL: ${{ needs.publish_pre_check.outputs.compile_use_paddle_whl_url }}
-
-
-  unittest_coverage:
-    name: Run FastDeploy Unit Tests and Coverage
-    needs: [clone,build_sm8090,ci_image_build]
-    uses: ./.github/workflows/_unit_test_coverage.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  logprob_test:
-    name: Run FastDeploy LogProb Tests
-    needs: [build_sm8090,ci_image_build]
-    uses: ./.github/workflows/_logprob_test_linux.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      PADDLETEST_ARCHIVE_URL: "https://xly-devops.bj.bcebos.com/PaddleTest/PaddleTest.tar.gz"
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  pre_ce_test:
-    name: Extracted partial CE model tasks to run in CI.
-    needs: [clone,build_sm8090,ci_image_build]
-    uses: ./.github/workflows/_pre_ce_test.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  base_test:
-    name: Run Base Tests
-    needs: [clone,build_sm8090,ci_image_build]
-    uses: ./.github/workflows/_base_test.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  accuracy_test:
-    name: Run Accuracy Tests
-    needs: [clone,build_sm8090,ci_image_build]
-    uses: ./.github/workflows/_accuracy_test.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  stable_test:
-    name: Run Stable Tests
-    needs: [clone,build_sm8090,ci_image_build]
-    uses: ./.github/workflows/_stable_test.yml
-    with:
-      DOCKER_IMAGE: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-
-  publish_pre_check:
-    name: Publish Docker Images Pre Check
-    needs: [ci_image_build, unittest_coverage,logprob_test,pre_ce_test,base_test,accuracy_test,stable_test]
-    runs-on: [self-hosted, Docker-Build]
-    steps:
-      - name: Images Uploading
-        env:
-          images_name: ${{ needs.ci_image_build.outputs.docker_name_precheck }}
-          ci_image_name: "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate"
-        run: |
-          echo "images_name=${images_name}"
-          docker images ${ci_image_name}
-          docker tag ${images_name} ${ci_image_name}
-          docker push ${ci_image_name}

.github/workflows/ci_xpu.yml

@@ -2,9 +2,7 @@ name: CI_XPU
 
 on:
   pull_request:
-    branches:
-      - develop
-      - 'release/*'
+    branches: [ develop ]
   workflow_dispatch:
 
 concurrency:
@@ -12,7 +10,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  CI_XPU:
+  build:
     runs-on: [self-hosted, XPU-P800-8Card]
     steps:
       - name: Print current runner name
@@ -24,16 +22,14 @@ jobs:
 
       - name: Code Checkout
         env:
-          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/fastdeploy-xpu:2.1.0
+          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/fastdeploy-xpu:2.0.0
         run: |
           REPO="https://github.com/${{ github.repository }}.git"
           FULL_REPO="${{ github.repository }}"
           REPO_NAME="${FULL_REPO##*/}"
-          BASE_BRANCH="${{ github.base_ref }}"
           # Clean the repository directory before starting
           docker run --rm --net=host -v $(pwd):/workspace -w /workspace \
           -e "REPO_NAME=${REPO_NAME}" \
-          -e "BASE_BRANCH=${BASE_BRANCH}" \
           ${docker_image} /bin/bash -c '
             if [ -d ${REPO_NAME} ]; then
               echo "Directory ${REPO_NAME} exists, removing it..."
@@ -42,7 +38,7 @@ jobs:
           '
           git config --global user.name "FastDeployCI"
           git config --global user.email "fastdeploy_ci@example.com"
-          git clone ${REPO} ${REPO_NAME} -b ${BASE_BRANCH}
+          git clone ${REPO} ${REPO_NAME}
           cd FastDeploy
           if [ "${{ github.event_name }}" = "pull_request" ]; then
             git fetch origin pull/${{ github.event.pull_request.number }}/head:pr/${{ github.event.pull_request.number }}
@@ -55,7 +51,7 @@ jobs:
 
       - name: Run CI unittest
         env:
-          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/fastdeploy-xpu:2.1.0
+          docker_image: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/fastdeploy-xpu:2.0.0
         run: |
           runner_name="${{ runner.name }}"
           last_char="${runner_name: -1}"
@@ -63,7 +59,7 @@ jobs:
           if [[ "$last_char" =~ [0-3] ]]; then
             gpu_id="$last_char"
           else
-            gpu_id="0"
+            gpu_id="0"  
           fi
           FD_API_PORT=$((9180 + gpu_id * 100))
           FD_ENGINE_QUEUE_PORT=$((9150 + gpu_id * 100))
@@ -77,7 +73,6 @@ jobs:
           -e "MODEL_PATH=/ssd3/model" \
           -e "http_proxy=$(git config --global --get http.proxy)" \
           -e "https_proxy=$(git config --global --get https.proxy)" \
-          -e "no_proxy=bcebos.com,mirrors.tuna.tsinghua.edu.cn,127.0.0.1,localhost" \
           -e "FD_API_PORT=${FD_API_PORT}" \
           -e "FD_ENGINE_QUEUE_PORT=${FD_ENGINE_QUEUE_PORT}" \
           -e "FD_METRICS_PORT=${FD_METRICS_PORT}" \
@@ -85,4 +80,4 @@ jobs:
           git config --global --add safe.directory /workspace/FastDeploy
           cd FastDeploy
           bash scripts/run_ci_xpu.sh
-          "
+          "

.github/workflows/gh-pages.yml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: 3.x
-      - run: pip install mkdocs-material mkdocs-get-deps mkdocs-material-extensions mkdocs-multilang mkdocs-static-i18n
+      - run: pip install mkdocs-material mkdocs-get-deps mkdocs-material-extensions mkdocs-multilang
       - name: Deploy to GitHub Pages
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/pr_build_and_test.yml

@@ -1,97 +0,0 @@
-name: PR Build and Test
-on:
-  pull_request:
-    types: [opened, synchronize]
-    branches: [develop, release/**]
-permissions: read-all
-
-concurrency:
-  group: ${{ github.event.pull_request.number }}-${{ github.workflow }}
-  cancel-in-progress: true
-
-jobs:
-  clone:
-    name: FD-Clone-Linux
-    uses: ./.github/workflows/_clone_linux.yml
-
-  build:
-    name: FD-Build-Linux
-    needs: clone
-    uses: ./.github/workflows/_build_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      COMPILE_ARCH: "90"
-      WITH_NIGHTLY_BUILD: "OFF"
-      FD_VERSION: "0.0.0"
-
-  resultshow:
-    name: Use Build Output
-    needs: build
-    runs-on: ubuntu-latest
-    steps:
-      - name: Print wheel path
-        run: |
-          echo "The built wheel is located at: ${{ needs.build.outputs.wheel_path }}"
-
-  unittest_coverage:
-    name: Run FastDeploy Unit Tests and Coverage
-    needs: [clone,build]
-    uses: ./.github/workflows/_unit_test_coverage.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  logprob_test:
-    name: Run FastDeploy LogProb Tests
-    needs: [build]
-    uses: ./.github/workflows/_logprob_test_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      PADDLETEST_ARCHIVE_URL: "https://xly-devops.bj.bcebos.com/PaddleTest/PaddleTest.tar.gz"
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  pre_ce_test:
-    name: Extracted partial CE model tasks to run in CI.
-    needs: [clone,build]
-    uses: ./.github/workflows/_pre_ce_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  base_test:
-    name: Run Base Tests
-    needs: [clone,build]
-    uses: ./.github/workflows/_base_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  accuracy_test:
-    name: Run Accuracy Tests
-    needs: [clone,build]
-    uses: ./.github/workflows/_accuracy_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  stable_test:
-    name: Run Stable Tests
-    needs: [clone,build]
-    uses: ./.github/workflows/_stable_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"

.github/workflows/publish_job.yml

@@ -1,331 +0,0 @@
-name: Publish Job
-
-on:
-  workflow_dispatch:
-  schedule:
-    - cron: '0 18 * * *'   # 2:00 AM China Standard Time (UTC+8)
-  push:
-    # branches:
-    #   - develop
-    tags:
-      - '*'
-
-permissions: read-all
-
-concurrency:
-  group: Publish-Job-${{ github.ref }}-${{ github.sha }}
-  cancel-in-progress: true
-
-
-jobs:
-  publish_pre_check:
-    runs-on: ubuntu-latest
-    if: |
-      github.event.repository.fork == false &&
-      (
-        (github.event_name == 'schedule' && github.ref_name == 'develop') ||
-        (github.event_name == 'push' && github.ref_type == 'tag') ||
-        ((github.event_name == 'workflow_dispatch') &&
-          (github.ref_name == 'develop' || github.ref_type == 'tag'))
-      )
-    env:
-      TAG_VERSION_MAPPINGS: ${{ vars.TAG_VERSION_MAPPINGS }}
-      FD_VERSION_DEV: ${{ vars.FD_VERSION_DEV }}
-      COMPILE_USE_PADDLE_WHL_URL_MAPPINGS: ${{ vars.COMPILE_USE_PADDLE_WHL_URL_MAPPINGS }}
-    outputs:
-      compile_use_paddle_version: ${{ steps.set_output.outputs.compile_use_paddle_version }}
-      compile_continue: ${{ steps.set_output.outputs.compile_continue }}
-      fd_version: ${{ steps.set_output.outputs.fd_version }}
-      with_nightly_build: ${{ steps.set_output.outputs.with_nightly_build }}
-      compile_use_paddle_whl_url: ${{ steps.set_output.outputs.compile_use_paddle_whl_url }}
-
-    steps:
-      - name: Get tag version
-        if: github.ref_type == 'tag'
-        run: |
-          TAG_NAME="${GITHUB_REF##*/}"       # 提取 tag 名称，比如 v2.1.0
-          TAG_VERSION="${TAG_NAME#v}"            # 去掉前缀 v
-          echo "FD_VERSION=$TAG_VERSION" >> $GITHUB_ENV
-
-      - name: Check FD version to Paddle version mapping
-        if: github.ref_type == 'tag'
-        env:
-          TARGET_FD: ${{ env.FD_VERSION }}
-        run: |
-          FOUND_PADDLE=""
-          # 遍历映射
-          for pair in $(echo $TAG_VERSION_MAPPINGS | tr ';' ' '); do
-            fd=$(echo "$pair" | cut -d',' -f1)
-            paddle=$(echo "$pair" | cut -d',' -f2)
-            if [[ "$fd" == "$TARGET_FD" ]]; then
-              FOUND_PADDLE="$paddle"
-              break
-            fi
-          done
-
-          if [[ -z "$FOUND_PADDLE" ]]; then
-            echo "No Paddle version found for FD $TARGET_FD"
-          else
-            echo "FD $TARGET_FD maps to Paddle $FOUND_PADDLE"
-            echo "PADDLE_VERSION=$FOUND_PADDLE" >> $GITHUB_ENV
-          fi
-      - name: Set Version
-        id: set_output
-        env:
-          PADDLE_VERSION: ${{ env.PADDLE_VERSION }}
-          FD_VERSION: ${{ env.FD_VERSION }}
-        run: |
-          if [[ "${{ github.ref_type }}" == "tag" ]]; then
-            if [[ -z "$PADDLE_VERSION" ]]; then
-              compile_continue=false
-            else
-              compile_use_paddle_version=$PADDLE_VERSION
-              compile_continue=true
-            fi
-            fd_version=$FD_VERSION
-          fi
-          if [[ "${{ github.ref_name }}" == "develop" ]];then
-            compile_continue=true
-            compile_use_paddle_version=""
-            fd_version=${FD_VERSION_DEV}
-            with_nightly_build=ON
-          fi
-          # Todo
-          # 通过变量COMPILE_USE_PADDLE_WHL_URL_MAPPINGS中的映射关系,决定是否是安装指定版本的Paddle还是直接安装URL
-          for pair in $(echo $COMPILE_USE_PADDLE_WHL_URL_MAPPINGS | tr ';' ' '); do
-            branch=$(echo "$pair" | cut -d',' -f1)
-            paddle_whl_url=$(echo "$pair" | cut -d',' -f2)
-            if [[ "$branch" == "${{ github.ref_name }}" ]]; then
-              FOUND_PADDLE_URL="$paddle_whl_url"
-              echo "compile_use_paddle_whl_url=${FOUND_PADDLE_URL}" >> $GITHUB_OUTPUT
-              compile_continue=true
-              break
-            fi
-          done
-          echo "compile_continue=${compile_continue}" >> $GITHUB_OUTPUT
-          echo "compile_use_paddle_version=${compile_use_paddle_version}" >> $GITHUB_OUTPUT
-          echo "fd_version=${fd_version}" >> $GITHUB_OUTPUT
-          echo "with_nightly_build=${with_nightly_build:-OFF}" >> $GITHUB_OUTPUT
-
-  print_publish_pre_check_outputs:
-    runs-on: ubuntu-latest
-    needs: publish_pre_check
-    steps:
-      - name: Print outputs as JSON
-        run: |
-          echo '${{ toJSON(needs.publish_pre_check.outputs) }}'
-
-  clone:
-    environment: CodeSync
-    name: FD-Clone-Linux
-    runs-on: ubuntu-latest
-    needs: publish_pre_check
-    if: ${{ needs.publish_pre_check.outputs.compile_continue == 'true' }}
-    outputs:
-      repo_archive_url: ${{ steps.set_output.outputs.repo_archive_url }}
-    steps:
-      - name: Clone FastDeploy
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.ref_name }}
-          submodules: 'recursive'
-          fetch-depth: 1000
-
-      - name: Python Setup
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Code Info Show and Upload
-        id: set_output
-        env:
-          AK: ${{ secrets.BOS_AK }}
-          SK: ${{ secrets.BOS_SK }}
-        run: |
-          git config --unset http.https://github.com/.extraheader
-          git submodule foreach --recursive sh -c "git config --local --unset-all 'http.https://github.com/.extraheader'"
-          git submodule foreach --recursive sh -c "git config remote.origin.fetch '+refs/heads/*:refs/remotes/origin/*'"
-          echo "Current HEAD Log:"
-          git log --oneline -n 5
-          ls
-          cd ..
-          tar -zcf FastDeploy.tar.gz FastDeploy
-          if [[ "${{ github.ref_type }}" == "tag" ]]; then
-            commit_id=${{ github.sha }}
-            tag_name=${{ github.ref_name }}
-            target_path=paddle-qa/TAG/FastDeploy/${tag_name}/${commit_id}
-          else
-            commit_id=${{ github.sha }}
-            branch_name=${{ github.ref_name }}
-            target_path=paddle-qa/BRANCH/FastDeploy/${branch_name}/${commit_id}
-          fi
-          wget  -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} FastDeploy.tar.gz ${target_path}
-          target_path_stripped="${target_path#paddle-qa/}"
-          REPO_ARCHIVE_URL=https://paddle-qa.bj.bcebos.com/${target_path_stripped}/FastDeploy.tar.gz
-          echo "repo_archive_url=${REPO_ARCHIVE_URL}" >> $GITHUB_OUTPUT
-
-  resultshow:
-    name: Show Code Archive Output
-    needs: clone
-    runs-on: ubuntu-latest
-    steps:
-      - name: Print wheel path
-        run: |
-          echo "The code archive is located at: ${{ needs.clone.outputs.repo_archive_url }}"
-
-  build_sm8090:
-    name: BUILD_SM8090
-    needs: [clone, publish_pre_check]
-    uses: ./.github/workflows/_build_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      COMPILE_ARCH: "80,90"
-      WITH_NIGHTLY_BUILD: ${{ needs.publish_pre_check.outputs.with_nightly_build }}
-      FD_VERSION: ${{ needs.publish_pre_check.outputs.fd_version }}
-      PADDLEVERSION: ${{ needs.publish_pre_check.outputs.compile_use_paddle_version }}
-      PADDLE_WHL_URL: ${{ needs.publish_pre_check.outputs.compile_use_paddle_whl_url }}
-
-  build_sm8689:
-    name: BUILD_SM8689
-    needs: [clone, publish_pre_check]
-    uses: ./.github/workflows/_build_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      COMPILE_ARCH: "86,89"
-      WITH_NIGHTLY_BUILD: ${{ needs.publish_pre_check.outputs.with_nightly_build }}
-      FD_VERSION: ${{ needs.publish_pre_check.outputs.fd_version }}
-      PADDLEVERSION: ${{ needs.publish_pre_check.outputs.compile_use_paddle_version }}
-      PADDLE_WHL_URL: ${{ needs.publish_pre_check.outputs.compile_use_paddle_whl_url }}
-
-  paddle_pypi_upload_sm8090:
-    environment: PaddleSourceUpload
-    name: PADDLE_PYPI_UPLOAD_8090
-    needs: build_sm8090
-    runs-on: ubuntu-latest
-    env:
-      AK: ${{ secrets.BOS_AK }}
-      SK: ${{ secrets.BOS_SK }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      COMPILE_ARCH: "80,90"
-    steps:
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Wheel Info Show and Upload
-        if: github.ref_name == 'develop' || github.ref_type == 'tag'
-        run: |
-          echo "The wheel is located at: ${FASTDEPLOY_WHEEL_URL}"
-          wget -q --no-check-certificate ${FASTDEPLOY_WHEEL_URL}
-          filename=$(basename ${FASTDEPLOY_WHEEL_URL})
-          if [[ "${{ github.ref_name }}" == "develop" ]];then
-            target_path=paddle-whl/nightly/fastdeploy-gpu-${COMPILE_ARCH//,/_}/fastdeploy-gpu
-          elif [[ "${{ github.ref_type }}" == "tag" ]]; then
-            target_path=paddle-whl/stable/fastdeploy-gpu-${COMPILE_ARCH//,/_}/fastdeploy-gpu
-          else
-            echo "Not develop or tag, do nothing"
-          fi
-          wget -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} ${filename} ${target_path}
-
-  paddle_pypi_upload_sm8689:
-    environment: PaddleSourceUpload
-    name: PADDLE_PYPI_UPLOAD_8689
-    needs: build_sm8689
-    runs-on: ubuntu-latest
-    env:
-      AK: ${{ secrets.BOS_AK }}
-      SK: ${{ secrets.BOS_SK }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8689.outputs.wheel_path }}
-      COMPILE_ARCH: "86,89"
-    steps:
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      - name: Wheel Info Show and Upload
-        if: github.ref_name == 'develop' || github.ref_type == 'tag'
-        run: |
-          echo "The wheel is located at: ${FASTDEPLOY_WHEEL_URL}"
-          wget -q --no-check-certificate ${FASTDEPLOY_WHEEL_URL}
-          filename=$(basename ${FASTDEPLOY_WHEEL_URL})
-          if [[ "${{ github.ref_name }}" == "develop" ]];then
-            target_path=paddle-whl/nightly/fastdeploy-gpu-${COMPILE_ARCH//,/_}/fastdeploy-gpu
-          elif [[ "${{ github.ref_type }}" == "tag" ]]; then
-            target_path=paddle-whl/stable/fastdeploy-gpu-${COMPILE_ARCH//,/_}/fastdeploy-gpu
-          else
-            echo "Not develop or tag, do nothing"
-          fi
-          wget -q --no-proxy --no-check-certificate https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddlePaddle/PaddleTest/tools/bos_tools.py
-          push_file=$(realpath bos_tools.py)
-          python -m pip install bce-python-sdk==0.9.29
-          ls
-          python ${push_file} ${filename} ${target_path}
-
-  unittest_coverage:
-    name: Run FastDeploy Unit Tests and Coverage
-    needs: [clone,build_sm8090]
-    uses: ./.github/workflows/_unit_test_coverage.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
-  logprob_test:
-    name: Run FastDeploy LogProb Tests
-    needs: [build_sm8090]
-    uses: ./.github/workflows/_logprob_test_linux.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      PADDLETEST_ARCHIVE_URL: "https://xly-devops.bj.bcebos.com/PaddleTest/PaddleTest.tar.gz"
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  pre_ce_test:
-    name: Extracted partial CE model tasks to run in CI.
-    needs: [clone,build_sm8090]
-    uses: ./.github/workflows/_pre_ce_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  base_test:
-    name: Run Base Tests
-    needs: [clone,build_sm8090]
-    uses: ./.github/workflows/_base_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  accuracy_test:
-    name: Run Accuracy Tests
-    needs: [clone,build_sm8090]
-    uses: ./.github/workflows/_accuracy_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build_sm8090.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"
-
-  stable_test:
-    name: Run Stable Tests
-    needs: [clone,build_sm8090]
-    uses: ./.github/workflows/_stable_test.yml
-    with:
-      DOCKER_IMAGE: ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlepaddle/paddleqa:fastdeploy-ciuse-cuda126-dailyupdate
-      FASTDEPLOY_ARCHIVE_URL: ${{ needs.clone.outputs.repo_archive_url }}
-      FASTDEPLOY_WHEEL_URL: ${{ needs.build.outputs.wheel_path }}
-      MODEL_CACHE_DIR: "/ssd2/actions-runner/ModelData"

.gitignore

@@ -121,7 +121,7 @@ dmypy.json
 FETCH_HEAD
 
 #log
-log/
+log*/
 
 checkpoints/
 checkpoints_origin/
@@ -156,23 +156,9 @@ nohup.out
 custom_ops/gpu_ops/fp8_deep_gemm/deep_gemm/include/cutlass
 custom_ops/gpu_ops/fp8_deep_gemm/deep_gemm/include/cute
 
-#marlin_kernel
-custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/kernel_*.cu
-
-#machete_kernel
-custom_ops/gpu_ops/machete/generated
-
 # buff
 custom_ops/tmp*
 
 build
 
 .ccls-cache
-
-third_party
-
-custom_ops/gpu_ops/w4afp8_gemm/w4afp8_gemm_*.cu
-custom_ops/gpu_ops/w4afp8_gemm/w4afp8_gemm_template.h
-
-custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8Afp8_sparse_gemm_*.cu
-custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8Afp8_sparse_gemm_template.h

.gitmodules

@@ -1,10 +0,0 @@
-[submodule "custom_ops/third_party/DeepGEMM"]
-	path = custom_ops/third_party/DeepGEMM
-	url = https://github.com/deepseek-ai/DeepGEMM.git
-	ignore = all
-[submodule "custom_ops/third_party/cutlass"]
-	path = custom_ops/third_party/cutlass
-	url = https://github.com/NVIDIA/cutlass.git
-[submodule "custom_ops/third_party/nlohmann_json"]
-	path = custom_ops/third_party/nlohmann_json
-	url = https://github.com/nlohmann/json.git

.pre-commit-config.yaml

@@ -3,30 +3,14 @@ default_install_hook_types:
   - commit-msg
 default_stages:
   - pre-commit # Run locally
-  - commit-msg
 #   - manual # Run in CI
 repos:
--   repo: https://github.com/psf/black.git
-    rev: 25.1.0
-    hooks:
-    -   id: black
-        files: \.(py|pyi)$
-        additional_dependencies: [toml]
-# 自动排序
--   repo: https://github.com/PyCQA/isort
-    rev: 5.11.5
-    hooks:
-    -   id: isort
--   repo: https://github.com/PyCQA/flake8
-    rev: 7.0.0
-    hooks:
-    -   id: flake8
 # 代码检查
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.11.7
   hooks:
   - id: ruff
-    args: [--output-format, github, --fix, --line-length=120, --config, pyproject.toml]
+    args: [--output-format, github, --fix, --line-length=120]
 # # 拼写检查
 # - repo: https://github.com/codespell-project/codespell
 #   rev: v2.4.1
@@ -34,13 +18,17 @@ repos:
 #   - id: codespell
 #     additional_dependencies: ['tomli']
 #     args: ['--toml', 'pyproject.toml']
-
+# 自动排序
+- repo: https://github.com/PyCQA/isort
+  rev: 6.0.1
+  hooks:
+  - id: isort
 # markdown
 - repo: https://github.com/jackdewinter/pymarkdown
   rev: v0.9.29
   hooks:
   - id: pymarkdown
-    args: ["-d", "MD029,MD031", fix]
+    args: [fix]
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v5.0.0
   hooks:

README.md

@@ -1,4 +1,3 @@
-English | [简体中文](README_CN.md)
 <p align="center">
   <a href="https://github.com/PaddlePaddle/FastDeploy/releases"><img src="https://github.com/user-attachments/assets/42b0039f-39e3-4279-afda-6d1865dfbffb" width="500"></a>
 </p>
@@ -9,28 +8,20 @@ English | [简体中文](README_CN.md)
     <a href="https://github.com/PaddlePaddle/FastDeploy/commits"><img src="https://img.shields.io/github/commit-activity/m/PaddlePaddle/FastDeploy?color=3af"></a>
     <a href="https://github.com/PaddlePaddle/FastDeploy/issues"><img src="https://img.shields.io/github/issues/PaddlePaddle/FastDeploy?color=9cc"></a>
     <a href="https://github.com/PaddlePaddle/FastDeploy/stargazers"><img src="https://img.shields.io/github/stars/PaddlePaddle/FastDeploy?color=ccf"></a>
-
 </p>
 
 <p align="center">
-     <a href="https://trendshift.io/repositories/4046" target="_blank"><img src="https://trendshift.io/api/badge/repositories/4046" alt="PaddlePaddle%2FFastDeploy | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a></br>
     <a href="https://paddlepaddle.github.io/FastDeploy/get_started/installation/nvidia_gpu/"><b> Installation </b></a>
     |
     <a href="https://paddlepaddle.github.io/FastDeploy/get_started/quick_start"><b> Quick Start </b></a>
     |
     <a href="https://paddlepaddle.github.io/FastDeploy/supported_models/"><b> Supported Models </b></a>
-
 </p>
 
 --------------------------------------------------------------------------------
-# FastDeploy : Inference and Deployment Toolkit for LLMs and VLMs based on PaddlePaddle
+# FastDeploy 2.0: Inference and Deployment Toolkit for LLMs and VLMs based on PaddlePaddle
 
 ## News
-**[2025-09] 🔥 FastDeploy v2.2 is newly released!** It now offers compatibility with models in the HuggingFace ecosystem, has further optimized performance, and newly adds support for [baidu/ERNIE-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking)!
-
-**[2025-08] 🔥 Released FastDeploy v2.1:** A brand-new KV Cache scheduling strategy has been introduced, and expanded support for PD separation and CUDA Graph across more models. Enhanced hardware support has been added for platforms like Kunlun and Hygon, along with comprehensive optimizations to improve the performance of both the service and inference engine.
-
-**[2025-07] The FastDeploy 2.0 Inference Deployment Challenge is now live!** Complete the inference deployment task for the ERNIE 4.5 series open-source models to win official FastDeploy 2.0 merch and generous prizes! 🎁 You're welcome to try it out and share your feedback! 📌[Sign up here](https://www.wjx.top/vm/meSsp3L.aspx#) 📌[Event details](https://github.com/PaddlePaddle/FastDeploy/discussions/2728)
 
 **[2025-06] 🔥 Released FastDeploy v2.0:** Supports inference and deployment for ERNIE 4.5. Furthermore, we open-source an industrial-grade PD disaggregation with context caching, dynamic role switching for effective resource utilization to further enhance inference performance for MoE models.
 
@@ -52,16 +43,14 @@ English | [简体中文](README_CN.md)
 
 ## Installation
 
-FastDeploy supports inference deployment on **NVIDIA GPUs**, **Kunlunxin XPUs**, **Iluvatar GPUs**, **Enflame GCUs**, **Hygon DCUs** and other hardware. For detailed installation instructions:
+FastDeploy supports inference deployment on **NVIDIA GPUs**, **Kunlunxin XPUs**, **Iluvatar GPUs**, **Enflame GCUs**, and other hardware. For detailed installation instructions:
 
 - [NVIDIA GPU](./docs/get_started/installation/nvidia_gpu.md)
 - [Kunlunxin XPU](./docs/get_started/installation/kunlunxin_xpu.md)
 - [Iluvatar GPU](./docs/get_started/installation/iluvatar_gpu.md)
 - [Enflame GCU](./docs/get_started/installation/Enflame_gcu.md)
-- [Hygon DCU](./docs/get_started/installation/hygon_dcu.md)
-- [MetaX GPU](./docs/get_started/installation/metax_gpu.md)
 
-**Note:** We are actively working on expanding hardware support. Additional hardware platforms including Ascend NPU are currently under development and testing. Stay tuned for updates!
+**Note:** We are actively working on expanding hardware support. Additional hardware platforms including Ascend NPU, Hygon DCU, and MetaX GPU are currently under development and testing. Stay tuned for updates!
 
 ## Get Started
 
@@ -71,12 +60,19 @@ Learn how to use FastDeploy through our documentation:
 - [ERNIE-4.5-VL Multimodal Model Deployment](./docs/get_started/ernie-4.5-vl.md)
 - [Offline Inference Development](./docs/offline_inference.md)
 - [Online Service Deployment](./docs/online_serving/README.md)
-- [Best Practices](./docs/best_practices/README.md)
+- [Full Supported Models List](./docs/supported_models.md)
 
 ## Supported Models
 
-Learn how to download models, enable using the torch format, and more:
-- [Full Supported Models List](./docs/supported_models.md)
+| Model | Data Type | PD Disaggregation | Chunked Prefill | Prefix Caching |  MTP | CUDA Graph | Maximum Context Length |
+|:--- | :------- | :---------- | :-------- | :-------- | :----- | :----- | :----- |
+|ERNIE-4.5-300B-A47B | BF16/WINT4/WINT8/W4A8C8/WINT2/FP8 | ✅| ✅ | ✅|✅(WINT4)| WIP |128K |
+|ERNIE-4.5-300B-A47B-Base| BF16/WINT4/WINT8 | ✅| ✅ | ✅|✅(WINT4)| WIP | 128K |
+|ERNIE-4.5-VL-424B-A47B | BF16/WINT4/WINT8 | WIP | ✅ | WIP | ❌ | WIP |128K |
+|ERNIE-4.5-VL-28B-A3B | BF16/WINT4/WINT8 | ❌ | ✅ | WIP | ❌ | WIP |128K |
+|ERNIE-4.5-21B-A3B | BF16/WINT4/WINT8/FP8  |  ❌ |  ✅ |  ✅ | WIP | ✅|128K |
+|ERNIE-4.5-21B-A3B-Base | BF16/WINT4/WINT8/FP8  |  ❌ |  ✅ |  ✅ | WIP | ✅|128K |
+|ERNIE-4.5-0.3B | BF16/WINT8/FP8  |  ❌ |  ✅ |  ✅ | ❌ | ✅| 128K |
 
 ## Advanced Usage
 

README_CN.md

@@ -1,89 +0,0 @@
-[English](README.md) | 简体中文
-<p align="center">
-  <a href="https://github.com/PaddlePaddle/FastDeploy/releases"><img src="https://github.com/user-attachments/assets/42b0039f-39e3-4279-afda-6d1865dfbffb" width="500"></a>
-</p>
-<p align="center">
-    <a href=""><img src="https://img.shields.io/badge/python-3.10-aff.svg"></a>
-    <a href=""><img src="https://img.shields.io/badge/os-linux-pink.svg"></a>
-    <a href="https://github.com/PaddlePaddle/FastDeploy/graphs/contributors"><img src="https://img.shields.io/github/contributors/PaddlePaddle/FastDeploy?color=9ea"></a>
-    <a href="https://github.com/PaddlePaddle/FastDeploy/commits"><img src="https://img.shields.io/github/commit-activity/m/PaddlePaddle/FastDeploy?color=3af"></a>
-    <a href="https://github.com/PaddlePaddle/FastDeploy/issues"><img src="https://img.shields.io/github/issues/PaddlePaddle/FastDeploy?color=9cc"></a>
-    <a href="https://github.com/PaddlePaddle/FastDeploy/stargazers"><img src="https://img.shields.io/github/stars/PaddlePaddle/FastDeploy?color=ccf"></a>
-
-</p>
-
-<p align="center">
-     <a href="https://trendshift.io/repositories/4046" target="_blank"><img src="https://trendshift.io/api/badge/repositories/4046" alt="PaddlePaddle%2FFastDeploy | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a></br>
-    <a href="https://paddlepaddle.github.io/FastDeploy/zh/get_started/installation/nvidia_gpu/"><b> 安装指导 </b></a>
-    |
-    <a href="https://paddlepaddle.github.io/FastDeploy/zh/get_started/quick_start"><b> 快速入门 </b></a>
-    |
-    <a href="https://paddlepaddle.github.io/FastDeploy/zh/supported_models/"><b> 支持模型列表 </b></a>
-
-</p>
-
---------------------------------------------------------------------------------
-# FastDeploy ：基于飞桨的大语言模型与视觉语言模型推理部署工具包
-
-## 最新活动
-**[2025-09] 🔥 FastDeploy v2.2 全新发布**: HuggingFace生态模型兼容，性能进一步优化，更新增对[baidu/ERNIE-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking)支持!
-
-**[2025-08] FastDeploy v2.1 发布**:全新的KV Cache调度策略，更多模型支持PD分离和CUDA Graph，昆仑、海光等更多硬件支持增强，全方面优化服务和推理引擎的性能。
-
-**[2025-07] 《FastDeploy2.0推理部署实测》专题活动已上线!** 完成文心4.5系列开源模型的推理部署等任务，即可获得骨瓷马克杯等FastDeploy2.0官方周边及丰富奖金！🎁 欢迎大家体验反馈～ 📌[报名地址](https://www.wjx.top/vm/meSsp3L.aspx#)   📌[活动详情](https://github.com/PaddlePaddle/FastDeploy/discussions/2728)
-
-## 关于
-
-**FastDeploy** 是基于飞桨（PaddlePaddle）的大语言模型（LLM）与视觉语言模型（VLM）推理部署工具包，提供**开箱即用的生产级部署方案**，核心技术特性包括：
-
-- 🚀 **负载均衡式PD分解**：工业级解决方案，支持上下文缓存与动态实例角色切换，在保障SLO达标和吞吐量的同时优化资源利用率
-- 🔄 **统一KV缓存传输**：轻量级高性能传输库，支持智能NVLink/RDMA选择
-- 🤝 **OpenAI API服务与vLLM兼容**：单命令部署，兼容[vLLM](https://github.com/vllm-project/vllm/)接口
-- 🧮 **全量化格式支持**：W8A16、W8A8、W4A16、W4A8、W2A16、FP8等
-- ⏩ **高级加速技术**：推测解码、多令牌预测（MTP）及分块预填充
-- 🖥️ **多硬件支持**：NVIDIA GPU、昆仑芯XPU、海光DCU、昇腾NPU、天数智芯GPU、燧原GCU、沐曦GPU等
-
-## 要求
-
-- 操作系统: Linux
-- Python: 3.10 ~ 3.12
-
-## 安装
-
-FastDeploy 支持在**英伟达（NVIDIA）GPU**、**昆仑芯（Kunlunxin）XPU**、**天数（Iluvatar）GPU**、**燧原（Enflame）GCU**、**海光（Hygon）DCU** 以及其他硬件上进行推理部署。详细安装说明如下：
-
-- [英伟达 GPU](./docs/zh/get_started/installation/nvidia_gpu.md)
-- [昆仑芯 XPU](./docs/zh/get_started/installation/kunlunxin_xpu.md)
-- [天数 CoreX](./docs/zh/get_started/installation/iluvatar_gpu.md)
-- [燧原 S60](./docs/zh/get_started/installation/Enflame_gcu.md)
-- [海光 DCU](./docs/zh/get_started/installation/hygon_dcu.md)
-- [沐曦 GPU](./docs/zh/get_started/installation/metax_gpu.md)
-
-**注意:** 我们正在积极拓展硬件支持范围。目前，包括昇腾（Ascend）NPU 等其他硬件平台正在开发测试中。敬请关注更新！
-
-## 入门指南
-
-通过我们的文档了解如何使用 FastDeploy：
-- [10分钟快速部署](./docs/zh/get_started/quick_start.md)
-- [ERNIE-4.5 部署](./docs/zh/get_started/ernie-4.5.md)
-- [ERNIE-4.5-VL 部署](./docs/zh/get_started/ernie-4.5-vl.md)
-- [离线推理](./docs/zh/offline_inference.md)
-- [在线服务](./docs/zh/online_serving/README.md)
-- [最佳实践](./docs/zh/best_practices/README.md)
-
-## 支持模型列表
-
-通过我们的文档了解如何下载模型，如何支持torch格式等：
-- [模型支持列表](./docs/zh/supported_models.md)
-
-## 进阶用法
-
-- [量化](./docs/zh/quantization/README.md)
-- [分离式部署](./docs/zh/features/disaggregated.md)
-- [投机解码](./docs/zh/features/speculative_decoding.md)
-- [前缀缓存](./docs/zh/features/prefix_caching.md)
-- [分块预填充](./docs/zh/features/chunked_prefill.md)
-
-## 致谢
-
-FastDeploy 依据 [Apache-2.0 开源许可证](./LICENSE). 进行授权。在开发过程中，我们参考并借鉴了 [vLLM](https://github.com/vllm-project/vllm) 的部分代码，以保持接口兼容性，在此表示衷心感谢。

benchmarks/README.md

@@ -41,10 +41,7 @@ python -m pip install -r requirements.txt
 --metric-percentiles 80,95,99,99.9,99.95,99.99：性能结果中展示的性能指标分位值
 --num-prompts 1：总计发送多少条请求
 --max-concurrency 1：压测并发数
---save-result：开启结果保存，结果文件会存入json，默认False不保存
---debug：开启debug模式，逐条打印payload和output内容，默认False
---shuffle：是否打乱数据集，默认False不打乱
---seed：打乱数据集时的随机种子，默认0
+--save-result：开启结果保存，结果文件会存入json
 ```
 
 ##### /v1/chat/completions接口压测单条数据调试
@@ -108,30 +105,3 @@ python benchmark_serving.py \
   --save-result > infer_log.txt 2>&1 &
 ```
 
-### 投机解码性能测试工具
-
-#### 使用方式：
-
-```bash
-python benchmarks/benchmark_mtp.py \
-  --host 127.0.0.1 --port 8000 \
-  --max-concurrency 16 32 64 96 --num-prompts 256 \
-  --acceptance-rate 0.8 --draft-token-steps 1 2 3 \
-  --s_itl-base-model 15.88 22.84 16.47 16.93 \
-  --dataset-name EBChat \
-  --dataset-path ./filtered_sharedgpt_2000_input_1136_output_200_fd.json
-```
-
-#### 参数说明
-
-```bash
---host：服务ip地址，用于组url
---port：服务HTTP端口，用于组url
---max-concurrency：测试并发数
---num-prompts：总计发送多少条请求
---acceptance-rate：投机解码的模拟接受率
---draft-token-steps：投机解码的步数
---s_itl-base-model：主模型的解码延迟，可由上述的性能压测工具获得，与batch-size一一对应
---dataset-name：指定数据集类，指定为"EBChat"可读取转存的FD格式数据集
---dataset-path：测试数据集路径
-```

benchmarks/backend_request_func.py

@@ -29,14 +29,13 @@ from typing import Optional
 import aiohttp
 from tqdm.asyncio import tqdm
 
+
 AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
 
 
 @dataclass
 class RequestFuncInput:
     """Input for requesting LLMs via API"""
-
-    no: int
     prompt: str
     history_QA: Optional[dict]
     hyper_parameters: dict
@@ -50,14 +49,11 @@ class RequestFuncInput:
     multi_modal_content: Optional[dict] = None
     ignore_eos: bool = False
     language: Optional[str] = None
-    debug: bool = False
 
 
 @dataclass
 class RequestFuncOutput:
     """Output for requesting LLMs via API"""
-
-    no: int = 0
     generated_text: str = ""
     reasoning_content: str = ""
     success: bool = False
@@ -68,7 +64,7 @@ class RequestFuncOutput:
     itl: list = field(default_factory=list)  # list of inter-token latencies
     tpot: float = 0.0  # avg next-token latencies
     prompt_len: int = 0
-    prompt_tokens: int = 0  # 推理侧返回输入token数
+    prompt_tokens: int = 0 # 推理侧返回输入token数
     error: str = ""
 
 
@@ -78,19 +74,22 @@ async def async_request_eb_openai_chat_completions(
 ) -> RequestFuncOutput:
     """Request an LLM using EB OpenAI"""
     api_url = request_func_input.api_url
-    assert api_url.endswith(("completions", "profile")), "OpenAI Chat Completions API URL must end with 'completions'."
+    assert api_url.endswith(
+        ("completions", "profile")
+    ), "OpenAI Chat Completions API URL must end with 'completions'."
 
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
         content = [{"type": "text", "text": request_func_input.prompt}]
         if request_func_input.multi_modal_content:
             content.append(request_func_input.multi_modal_content)
         payload = {
-            "model": request_func_input.model,
+            "model": "default",
             "messages": request_func_input.history_QA,
             "stream": True,
             "stream_options": {
                 "include_usage": True,
-                "continuous_usage_stats": True,
+                "continuous_usage_stats": True
             },
         }
         # 超参由yaml传入
@@ -98,10 +97,6 @@ async def async_request_eb_openai_chat_completions(
 
         if request_func_input.ignore_eos:
             payload["ignore_eos"] = request_func_input.ignore_eos
-
-        if request_func_input.debug:
-            print(f"payload:{json.dumps(payload, ensure_ascii=False)}")
-
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
@@ -109,20 +104,21 @@ async def async_request_eb_openai_chat_completions(
 
         output = RequestFuncOutput()
         output.prompt_len = 0
-        output.no = request_func_input.no
 
         ttft = 0.0
         st = time.perf_counter()
         most_recent_timestamp = st
         try:
-            async with session.post(url=api_url, json=payload, headers=headers) as response:
+            async with session.post(url=api_url, json=payload,
+                                    headers=headers) as response:
                 if response.status == 200:
                     async for chunk_bytes in response.content:
                         chunk_bytes = chunk_bytes.strip()
                         if not chunk_bytes:
                             continue
 
-                        chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")
+                        chunk = chunk_bytes.decode("utf-8").removeprefix(
+                            "data: ")
                         if chunk != "[DONE]":
                             # print("####chunk:", chunk, type(chunk))
                             timestamp = time.perf_counter()
@@ -136,20 +132,21 @@ async def async_request_eb_openai_chat_completions(
                                     ttft = timestamp - st
                                     output.ttft = ttft
                                     # cached_tokens
-                                    output.prompt_len = (
-                                        data["usage"].get("prompt_tokens_details", {}).get("cached_tokens", 0)
-                                    )
+                                    output.prompt_len = data["usage"]["prompt_tokens_details"]["cached_tokens"]
 
                                 # Decoding phase
                                 else:
-                                    output.itl.append(timestamp - most_recent_timestamp)
+                                    output.itl.append(timestamp -
+                                                      most_recent_timestamp)
 
                                 output.generated_text += content or ""
                                 output.reasoning_content += reason_content or ""
-                                output.arrival_time.append(choices[0].get("arrival_time", timestamp))
-                            elif usage := data.get("usage", {}):
-                                output.output_tokens = usage.get("completion_tokens", 0)
-                                output.prompt_tokens = usage.get("prompt_tokens", 0)
+                                output.arrival_time.append(choices[0].get("arrival_time"))
+                            elif usage := data.get("usage"):
+                                output.output_tokens = usage.get(
+                                    "completion_tokens")
+                                output.prompt_tokens = usage.get(
+                                    "prompt_tokens")
 
                             most_recent_timestamp = timestamp
 
@@ -162,12 +159,7 @@ async def async_request_eb_openai_chat_completions(
                     output.latency = most_recent_timestamp - st
                 else:
                     error_text = await response.text()
-                    print(
-                        "####error response:",
-                        error_text,
-                        "####payload:",
-                        payload,
-                    )
+                    print("####error response:", error_text, "####payload:", payload)
                     output.error = error_text or ""
                     output.success = False
         except Exception:
@@ -181,8 +173,6 @@ async def async_request_eb_openai_chat_completions(
                 f.write(str(output) + "\n")
     if pbar:
         pbar.update(1)
-    if request_func_input.debug:
-        print("#####final_output:", output)
     return output
 
 
@@ -196,14 +186,15 @@ async def async_request_eb_openai_completions(
         ("completions", "profile")
     ), "OpenAI Completions API URL must end with 'completions' or 'profile'."
 
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
         payload = {
-            "model": request_func_input.model,
+            "model": "default",
             "prompt": request_func_input.prompt,
             "stream": True,
             "stream_options": {
                 "include_usage": True,
-                "continuous_usage_stats": True,
+                "continuous_usage_stats": True
             },
         }
         # 超参由yaml传入
@@ -211,25 +202,19 @@ async def async_request_eb_openai_completions(
 
         if request_func_input.ignore_eos:
             payload["ignore_eos"] = request_func_input.ignore_eos
-
-        if request_func_input.debug:
-            print("payload:", json.dumps(payload, ensure_ascii=False))
-
         headers = {
-            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
-            "Content-Type": "application/json",
+            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
         }
 
         output = RequestFuncOutput()
         output.prompt_len = request_func_input.prompt_len
-        output.no = request_func_input.no
 
         generated_text = ""
-        ttft = 0.0
         st = time.perf_counter()
         most_recent_timestamp = st
         try:
-            async with session.post(url=api_url, json=payload, headers=headers) as response:
+            async with session.post(url=api_url, json=payload,
+                                    headers=headers) as response:
                 if response.status == 200:
                     first_chunk_received = False
                     async for chunk_bytes in response.content:
@@ -237,10 +222,10 @@ async def async_request_eb_openai_completions(
                         if not chunk_bytes:
                             continue
 
-                        chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")
+                        chunk = chunk_bytes.decode("utf-8").removeprefix(
+                            "data: ")
                         if chunk != "[DONE]":
                             # print("####chunk:", chunk, chunk.usage)
-                            timestamp = time.perf_counter()
                             data = json.loads(chunk)
 
                             # NOTE: Some completion API might have a last
@@ -250,40 +235,35 @@ async def async_request_eb_openai_completions(
                                 # Note that text could be empty here
                                 # e.g. for special tokens
                                 text = choices[0].get("text")
-
+                                timestamp = time.perf_counter()
                                 # First token
                                 if not first_chunk_received:
                                     first_chunk_received = True
-                                    ttft = timestamp - st
+                                    ttft = time.perf_counter() - st
                                     output.ttft = ttft
 
                                 # Decoding phase
                                 else:
-                                    output.itl.append(timestamp - most_recent_timestamp)
-
-                                generated_text += text or ""
+                                    output.itl.append(timestamp -
+                                                      most_recent_timestamp)
 
                                 most_recent_timestamp = timestamp
-                                output.arrival_time.append(choices[0].get("arrival_time", timestamp))
+                                output.arrival_time.append(choices[0].get("arrival_time"))
+                                generated_text += text or ""
                             elif usage := data.get("usage"):
-                                output.prompt_tokens = usage.get("prompt_tokens")
-                                output.output_tokens = usage.get("completion_tokens")
+                                output.prompt_tokens = usage.get(
+                                    "prompt_tokens")
+                                output.output_tokens = usage.get(
+                                    "completion_tokens")
                     if first_chunk_received:
                         output.success = True
                     else:
                         output.success = False
                         output.error = (
-                            "Never received a valid chunk to calculate TTFT." "This response will be marked as failed!"
-                        )
-
+                            "Never received a valid chunk to calculate TTFT."
+                            "This response will be marked as failed!")
                     output.generated_text = generated_text
                     output.latency = most_recent_timestamp - st
-
-                    if output.generated_text == "":
-                        output.success = False
-                        output.error = "No generated text found!"
-                    else:
-                        output.success = True
                 else:
                     output.error = response.reason or ""
                     output.success = False
@@ -292,9 +272,6 @@ async def async_request_eb_openai_completions(
             exc_info = sys.exc_info()
             output.error = "".join(traceback.format_exception(*exc_info))
 
-        if request_func_input.debug:
-            print(f"final_output:{output}")
-
     if pbar:
         pbar.update(1)
     return output
@@ -308,7 +285,8 @@ async def async_request_tgi(
     api_url = request_func_input.api_url
     assert api_url.endswith("generate_stream")
 
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
         params = {
             "max_new_tokens": request_func_input.output_len,
             "do_sample": True,
@@ -355,7 +333,8 @@ async def async_request_tgi(
 
                         # Decoding phase
                         else:
-                            output.itl.append(timestamp - most_recent_timestamp)
+                            output.itl.append(timestamp -
+                                              most_recent_timestamp)
 
                         most_recent_timestamp = timestamp
                         output.arrival_time.append(data["arrival_time"])
@@ -384,7 +363,8 @@ async def async_request_trt_llm(
     api_url = request_func_input.api_url
     assert api_url.endswith("generate_stream")
 
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
         payload = {
             "accumulate_tokens": True,
             "text_input": request_func_input.prompt,
@@ -409,7 +389,8 @@ async def async_request_trt_llm(
                         if not chunk_bytes:
                             continue
 
-                        chunk = chunk_bytes.decode("utf-8").removeprefix("data:")
+                        chunk = chunk_bytes.decode("utf-8").removeprefix(
+                            "data:")
 
                         data = json.loads(chunk)
                         output.generated_text += data["text_output"]
@@ -421,7 +402,8 @@ async def async_request_trt_llm(
 
                         # Decoding phase
                         else:
-                            output.itl.append(timestamp - most_recent_timestamp)
+                            output.itl.append(timestamp -
+                                              most_recent_timestamp)
 
                         most_recent_timestamp = timestamp
 
@@ -446,7 +428,8 @@ async def async_request_deepspeed_mii(
     pbar: Optional[tqdm] = None,
 ) -> RequestFuncOutput:
     """Request an LLM using Deepspeed MII"""
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
 
         payload = {
             "prompt": request_func_input.prompt,
@@ -464,16 +447,19 @@ async def async_request_deepspeed_mii(
 
         st = time.perf_counter()
         try:
-            async with session.post(url=request_func_input.api_url, json=payload) as response:
+            async with session.post(url=request_func_input.api_url,
+                                    json=payload) as response:
                 if response.status == 200:
                     parsed_resp = await response.json()
                     output.latency = time.perf_counter() - st
                     if "choices" in parsed_resp:
-                        output.generated_text = parsed_resp["choices"][0]["text"]
+                        output.generated_text = parsed_resp["choices"][0][
+                            "text"]
                     elif "text" in parsed_resp:
                         output.generated_text = parsed_resp["text"][0]
                     else:
-                        output.error = "Unexpected response format: " "neither 'choices' nor 'text' found"
+                        output.error = ("Unexpected response format: "
+                                        "neither 'choices' nor 'text' found")
                         output.success = False
                     output.success = True
                 else:
@@ -499,22 +485,26 @@ async def async_request_openai_completions(
         ("completions", "profile")
     ), "OpenAI Completions API URL must end with 'completions' or 'profile'."
 
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
         payload = {
-            "model": (request_func_input.model_name if request_func_input.model_name else request_func_input.model),
+            "model": request_func_input.model_name \
+                if request_func_input.model_name else request_func_input.model,
             "prompt": request_func_input.prompt,
             # "temperature": 0.0,
             "max_tokens": request_func_input.output_len,
             "logprobs": request_func_input.logprobs,
             "stream": True,
-            # "stream_options": {
+            #"stream_options": {
             #    "include_usage": True,
-            # },
+            #},
         }
         if request_func_input.ignore_eos:
             payload["ignore_eos"] = request_func_input.ignore_eos
 
-        headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
+        headers = {
+            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
+        }
 
         output = RequestFuncOutput()
         output.prompt_len = request_func_input.prompt_len
@@ -523,7 +513,8 @@ async def async_request_openai_completions(
         st = time.perf_counter()
         most_recent_timestamp = st
         try:
-            async with session.post(url=api_url, json=payload, headers=headers) as response:
+            async with session.post(url=api_url, json=payload,
+                                    headers=headers) as response:
                 if response.status == 200:
                     first_chunk_received = False
                     async for chunk_bytes in response.content:
@@ -531,7 +522,8 @@ async def async_request_openai_completions(
                         if not chunk_bytes:
                             continue
 
-                        chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")
+                        chunk = chunk_bytes.decode("utf-8").removeprefix(
+                            "data: ")
                         if chunk != "[DONE]":
                             # print("####chunk:", chunk, type(chunk))
                             data = json.loads(chunk)
@@ -552,19 +544,21 @@ async def async_request_openai_completions(
 
                                 # Decoding phase
                                 else:
-                                    output.itl.append(timestamp - most_recent_timestamp)
+                                    output.itl.append(timestamp -
+                                                      most_recent_timestamp)
 
                                 most_recent_timestamp = timestamp
                                 generated_text += text or ""
                             elif usage := data.get("usage"):
-                                output.output_tokens = usage.get("completion_tokens")
+                                output.output_tokens = usage.get(
+                                    "completion_tokens")
                     if first_chunk_received:
                         output.success = True
                     else:
                         output.success = False
                         output.error = (
-                            "Never received a valid chunk to calculate TTFT." "This response will be marked as failed!"
-                        )
+                            "Never received a valid chunk to calculate TTFT."
+                            "This response will be marked as failed!")
                     output.generated_text = generated_text
                     output.latency = most_recent_timestamp - st
                 else:
@@ -587,24 +581,25 @@ async def async_request_openai_audio(
     """Request an LLM using OpenAI"""
     # Lazy import without PlaceholderModule to avoid vllm dep.
     import soundfile
-
     api_url = request_func_input.api_url
     assert api_url.endswith(
-        ("transcriptions", "translations")
-    ), "OpenAI Chat Completions API URL must end with 'transcriptions' "
+        ("transcriptions", "translations"
+         )), "OpenAI Chat Completions API URL must end with 'transcriptions' "
     "or `translations`."
 
-    async with aiohttp.ClientSession(trust_env=True, timeout=AIOHTTP_TIMEOUT) as session:
+    async with aiohttp.ClientSession(trust_env=True,
+                                     timeout=AIOHTTP_TIMEOUT) as session:
         content = [{"type": "text", "text": request_func_input.prompt}]
         payload = {
-            "model": (request_func_input.model_name if request_func_input.model_name else request_func_input.model),
+            "model": request_func_input.model_name \
+                if request_func_input.model_name else request_func_input.model,
             "temperature": 0.0,
             "max_completion_tokens": request_func_input.output_len,
             "stream": True,
             "language": "en",
             # Flattened due to multipart/form-data
             "stream_include_usage": True,
-            "stream_continuous_usage_stats": True,
+            "stream_continuous_usage_stats": True
         }
         if request_func_input.extra_body:
             payload.update(request_func_input.extra_body)
@@ -619,9 +614,9 @@ async def async_request_openai_audio(
             buffer.seek(0)
             return buffer
 
-        with to_bytes(*request_func_input.multi_modal_content["audio"]) as f:
+        with to_bytes(*request_func_input.multi_modal_content['audio']) as f:
             form = aiohttp.FormData()
-            form.add_field("file", f, content_type="audio/wav")
+            form.add_field('file', f, content_type='audio/wav')
             for key, value in payload.items():
                 form.add_field(key, str(value))
 
@@ -633,20 +628,24 @@ async def async_request_openai_audio(
             st = time.perf_counter()
             most_recent_timestamp = st
             try:
-                async with session.post(url=api_url, data=form, headers=headers) as response:
+                async with session.post(url=api_url,
+                                        data=form,
+                                        headers=headers) as response:
                     if response.status == 200:
                         async for chunk_bytes in response.content:
                             chunk_bytes = chunk_bytes.strip()
                             if not chunk_bytes:
                                 continue
 
-                            chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")
+                            chunk = chunk_bytes.decode("utf-8").removeprefix(
+                                "data: ")
                             if chunk != "[DONE]":
                                 timestamp = time.perf_counter()
                                 data = json.loads(chunk)
 
                                 if choices := data.get("choices"):
-                                    content = choices[0]["delta"].get("content")
+                                    content = choices[0]["delta"].get(
+                                        "content")
                                     # First token
                                     if ttft == 0.0:
                                         ttft = timestamp - st
@@ -654,11 +653,13 @@ async def async_request_openai_audio(
 
                                     # Decoding phase
                                     else:
-                                        output.itl.append(timestamp - most_recent_timestamp)
+                                        output.itl.append(
+                                            timestamp - most_recent_timestamp)
 
                                     generated_text += content or ""
                                 elif usage := data.get("usage"):
-                                    output.output_tokens = usage.get("completion_tokens")
+                                    output.output_tokens = usage.get(
+                                        "completion_tokens")
 
                                 most_recent_timestamp = timestamp
 
@@ -692,11 +693,8 @@ ASYNC_REQUEST_FUNCS = {
 }
 
 OPENAI_COMPATIBLE_BACKENDS = [
-    k
-    for k, v in ASYNC_REQUEST_FUNCS.items()
-    if v
-    in (
-        async_request_openai_completions,
-        async_request_eb_openai_chat_completions,
-    )
+    k for k, v in ASYNC_REQUEST_FUNCS.items()
+    if v in (async_request_openai_completions,
+             async_request_eb_openai_chat_completions)
 ]
+

benchmarks/benchmark_dataset.py

@@ -26,10 +26,10 @@ from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from dataclasses import dataclass
 from io import BytesIO
-from typing import Any, Optional, Union
-
+from typing import Any, Callable, Optional, Union
 from PIL import Image
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -39,7 +39,6 @@ class SampleRequest:
     Represents a single inference request for benchmarking.
     """
 
-    no: int
     prompt: Union[str, Any]
     history_QA: Union[str, Any]
     json_data: Optional[dict]
@@ -49,7 +48,6 @@ class SampleRequest:
 
 class BenchmarkDataset(ABC):
     """BenchmarkDataset"""
-
     DEFAULT_SEED = 0
     IS_MULTIMODAL = False
 
@@ -57,7 +55,6 @@ class BenchmarkDataset(ABC):
         self,
         dataset_path: Optional[str] = None,
         random_seed: int = DEFAULT_SEED,
-        shuffle: bool = False,
         hyperparameter_path: Optional[str] = None,
     ) -> None:
         """
@@ -71,9 +68,9 @@ class BenchmarkDataset(ABC):
         self.dataset_path = dataset_path
         # Set the random seed, ensuring that a None value is replaced with the
         # default seed.
-        self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED
+        self.random_seed = (random_seed
+                            if random_seed is not None else self.DEFAULT_SEED)
         self.data = None
-        self.shuffle = shuffle
         self.hyperparameter_path = hyperparameter_path
         self.hyperparameters = {}
 
@@ -88,7 +85,8 @@ class BenchmarkDataset(ABC):
             NotImplementedError: If a subclass does not implement this method.
         """
         # TODO (jenniferzhao): add support for downloading data
-        raise NotImplementedError("load_data must be implemented in subclasses.")
+        raise NotImplementedError(
+            "load_data must be implemented in subclasses.")
 
     @abstractmethod
     def sample(self, num_requests: int) -> list[SampleRequest]:
@@ -107,7 +105,8 @@ class BenchmarkDataset(ABC):
         """
         raise NotImplementedError("sample must be implemented in subclasses.")
 
-    def maybe_oversample_requests(self, requests: list[SampleRequest], num_requests: int) -> None:
+    def maybe_oversample_requests(self, requests: list[SampleRequest],
+                                  num_requests: int) -> None:
         """
         Oversamples the list of requests if its size is less than the desired
         number.
@@ -118,9 +117,11 @@ class BenchmarkDataset(ABC):
         """
         if len(requests) < num_requests:
             random.seed(self.random_seed)
-            additional = random.choices(requests, k=num_requests - len(requests))
+            additional = random.choices(requests,
+                                        k=num_requests - len(requests))
             requests.extend(additional)
-            logger.info("Oversampled requests to reach %d total samples.", num_requests)
+            logger.info("Oversampled requests to reach %d total samples.",
+                        num_requests)
 
 
 def is_valid_sequence(
@@ -140,12 +141,14 @@ def is_valid_sequence(
     """
     # Check for invalid conditions
     prompt_too_short = prompt_len < min_len
-    output_too_short = (not skip_min_output_len_check) and (output_len < min_len)
+    output_too_short = (not skip_min_output_len_check) and (output_len
+                                                            < min_len)
     prompt_too_long = prompt_len > max_prompt_len
     combined_too_long = (prompt_len + output_len) > max_total_len
 
     # Return True if none of the invalid conditions are met
-    return not (prompt_too_short or output_too_short or prompt_too_long or combined_too_long)
+    return not (prompt_too_short or output_too_short or prompt_too_long
+                or combined_too_long)
 
 
 def process_image(image: Any) -> Mapping[str, Any]:
@@ -168,25 +171,28 @@ def process_image(image: Any) -> Mapping[str, Any]:
     Raises:
         ValueError: If the input is not a supported type.
     """
-    if isinstance(image, dict) and "bytes" in image:
-        image = Image.open(BytesIO(image["bytes"]))
+    if isinstance(image, dict) and 'bytes' in image:
+        image = Image.open(BytesIO(image['bytes']))
     if isinstance(image, Image.Image):
         image = image.convert("RGB")
         with io.BytesIO() as image_data:
             image.save(image_data, format="JPEG")
-            image_base64 = base64.b64encode(image_data.getvalue()).decode("utf-8")
+            image_base64 = base64.b64encode(
+                image_data.getvalue()).decode("utf-8")
         return {
             "type": "image_url",
-            "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"},
+            "image_url": {
+                "url": f"data:image/jpeg;base64,{image_base64}"
+            },
         }
 
     if isinstance(image, str):
-        image_url = image if image.startswith(("http://", "file://")) else f"file://{image}"
+        image_url = (image if image.startswith(
+            ("http://", "file://")) else f"file://{image}")
         return {"type": "image_url", "image_url": {"url": image_url}}
 
-    raise ValueError(
-        f"Invalid image input {image}. Must be a PIL.Image.Image" " or str or dictionary with raw image bytes."
-    )
+    raise ValueError(f"Invalid image input {image}. Must be a PIL.Image.Image"
+                     " or str or dictionary with raw image bytes.")
 
 
 class EBDataset(BenchmarkDataset):
@@ -213,10 +219,6 @@ class EBDataset(BenchmarkDataset):
         with open(self.dataset_path, encoding="utf-8") as f:
             self.data = [json.loads(i.strip()) for i in f.readlines()]
 
-        if self.shuffle:
-            random.seed(self.random_seed)
-            random.shuffle(self.data)
-
     def sample(
         self,
         num_requests: int,
@@ -227,7 +229,6 @@ class EBDataset(BenchmarkDataset):
         **kwargs,
     ) -> list:
         samples: list = []
-        cnt = 1
         for entry in self.data:
             if len(samples) >= num_requests:
                 break
@@ -241,17 +242,15 @@ class EBDataset(BenchmarkDataset):
             new_output_len = int(entry["max_dec_len"])
 
             if enable_multimodal_chat:
-                prompt = self.apply_multimodal_chat_transformation(prompt, None)
+                prompt = self.apply_multimodal_chat_transformation(
+                    prompt, None)
             samples.append(
                 SampleRequest(
-                    no=cnt,
                     prompt=prompt,
                     prompt_len=self.prompt_len,
                     history_QA=[],
                     expected_output_len=new_output_len,
-                )
-            )
-            cnt += 1
+                ))
 
         self.maybe_oversample_requests(samples, num_requests)
         return samples
@@ -262,7 +261,6 @@ class EBChatDataset(BenchmarkDataset):
     Implements the ShareGPT dataset.  Loads data from a JSON file and generates
     sample requests based on conversation turns.
     """
-
     prompt_len: int
 
     def __init__(self, **kwargs) -> None:
@@ -276,10 +274,6 @@ class EBChatDataset(BenchmarkDataset):
         with open(self.dataset_path, encoding="utf-8") as f:
             self.data = [json.loads(i.strip()) for i in f.readlines()]
 
-        if self.shuffle:
-            random.seed(self.random_seed)
-            random.shuffle(self.data)
-
     def sample(
         self,
         num_requests: int,
@@ -290,7 +284,6 @@ class EBChatDataset(BenchmarkDataset):
         **kwargs,
     ) -> list:
         samples: list = []
-        cnt = 1
         for entry in self.data:
             if len(samples) >= num_requests:
                 break
@@ -300,18 +293,17 @@ class EBChatDataset(BenchmarkDataset):
             new_output_len = int(entry.get("max_tokens", 12288))
 
             if enable_multimodal_chat:
-                prompt = self.apply_multimodal_chat_transformation(prompt, None)
+                prompt = self.apply_multimodal_chat_transformation(
+                    prompt, None)
             samples.append(
                 SampleRequest(
-                    no=cnt,
                     json_data=json_data,
                     prompt=prompt,
                     prompt_len=0,
                     history_QA=history_QA,
                     expected_output_len=new_output_len,
-                )
-            )
-            cnt += 1
+                ))
 
         self.maybe_oversample_requests(samples, num_requests)
         return samples
+

benchmarks/benchmark_mtp.py

@@ -1,178 +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 argparse
-import asyncio
-import contextlib
-import os
-from typing import Union
-
-from benchmark_dataset import EBChatDataset, EBDataset
-from benchmark_serving import benchmark
-
-
-def prepare_input_requests(num_prompts: int, dataset_name: str, dataset_path: str) -> Union[EBDataset, EBChatDataset]:
-    dataset_mapping = {
-        "EB": lambda: EBDataset(dataset_path=dataset_path).sample(num_requests=num_prompts),
-        "EBChat": lambda: EBChatDataset(dataset_path=dataset_path).sample(num_requests=num_prompts),
-    }
-
-    try:
-        input_requests = dataset_mapping[dataset_name]()
-    except KeyError as err:
-        raise ValueError(f"Unknown dataset: {dataset_name}") from err
-
-    return input_requests
-
-
-class FakeTokenizer:
-    def encode(self, text: str, add_special_tokens: bool = False):
-        return []
-
-
-def send_one_batch(base_url, max_concurrency, input_requests, disable_tqdm):
-    selected_percentile_metrics = ["s_itl"]
-    selected_percentiles = []
-    # Run benchmark
-    results = asyncio.run(
-        benchmark(
-            backend="openai-chat",
-            api_url=f"{base_url}/v1/chat/completions",
-            base_url=base_url,
-            model_id="default",
-            model_name="default",
-            input_requests=input_requests,
-            hyper_parameters={},
-            logprobs=None,
-            request_rate=float("inf"),
-            burstiness=1.0,
-            disable_tqdm=disable_tqdm,
-            profile=False,
-            selected_percentile_metrics=selected_percentile_metrics,
-            selected_percentiles=selected_percentiles,
-            ignore_eos=False,
-            goodput_config_dict=None,
-            max_concurrency=max_concurrency,
-            lora_modules=None,
-            extra_body=None,
-        )
-    )
-
-    record = {
-        "mean_s_itl_ms": results["mean_s_itl_ms"],
-    }
-
-    return record
-
-
-def calculate_speedup(acceptance_rate, draft_token_step, t_ori, t_mtp):
-
-    tmp = 0.0
-    for i in range(draft_token_step):
-        tmp += pow(acceptance_rate, i + 1)
-
-    r_ac = tmp / (1 + tmp)
-
-    return t_ori / ((1 - r_ac) * t_mtp)
-
-
-def main(args):
-    base_url = f"http://{args.host}:{args.port}"
-
-    input_requests = prepare_input_requests(args.num_prompts, args.dataset_name, args.dataset_path)
-
-    if len(args.max_concurrency) != len(args.s_itl_base_model):
-        raise ValueError("--max_concurrency should be same length as --s_itl_base_model")
-
-    for max_concurrency, s_itl in zip(args.max_concurrency, args.s_itl_base_model):
-        # Wramup
-        print("Starting warmup...")
-        with open(os.devnull, "w") as f:
-            with contextlib.redirect_stdout(f):
-                send_one_batch(
-                    base_url,
-                    max_concurrency,
-                    input_requests[0:max_concurrency],
-                    True,
-                )
-
-        # Benchmark
-        record = send_one_batch(base_url, max_concurrency, input_requests, False)
-
-        metric_header = "Speed up"
-        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
-        for draft_token_step in args.draft_token_steps:
-            speedup = calculate_speedup(
-                args.acceptance_rate,
-                draft_token_step,
-                s_itl,
-                record["mean_s_itl_ms"],
-            )
-            print("{:<40} {:<10.2f}".format(f"Speed up on {draft_token_step} steps draft", speedup))
-        print("=" * 50)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="127.0.0.1",
-    )
-    parser.add_argument(
-        "--port",
-        type=str,
-        default="8000",
-    )
-    parser.add_argument(
-        "--max-concurrency",
-        type=int,
-        nargs="+",
-        default=(1, 2, 4, 8, 16, 32),
-    )
-    parser.add_argument(
-        "--num-prompts",
-        type=int,
-        default=128,
-    )
-    parser.add_argument(
-        "--acceptance-rate",
-        type=float,
-        default=0.8,
-    )
-    parser.add_argument(
-        "--draft-token-steps",
-        type=int,
-        nargs="+",
-        default=(1, 2),
-    )
-    parser.add_argument(
-        "--s_itl-base-model",
-        type=float,
-        nargs="+",
-    )
-    parser.add_argument(
-        "--dataset-name",
-        type=str,
-        default="EBChat",
-    )
-    parser.add_argument(
-        "--dataset-path",
-        type=str,
-    )
-    args = parser.parse_args()
-
-    main(args)

benchmarks/benchmark_utils.py

@@ -24,11 +24,9 @@ import os
 from typing import Any
 
 
-def convert_to_pytorch_benchmark_format(
-    args: argparse.Namespace,
-    metrics: dict[str, list],
-    extra_info: dict[str, Any],
-) -> list:
+def convert_to_pytorch_benchmark_format(args: argparse.Namespace,
+                                        metrics: dict[str, list],
+                                        extra_info: dict[str, Any]) -> list:
     """
     Save the benchmark results in the format used by PyTorch OSS benchmark with
     on metric per record
@@ -56,10 +54,12 @@ def convert_to_pytorch_benchmark_format(
             },
         }
 
-        tp = record["benchmark"]["extra_info"]["args"].get("tensor_parallel_size")
+        tp = record["benchmark"]["extra_info"]["args"].get(
+            "tensor_parallel_size")
         # Save tensor_parallel_size parameter if it's part of the metadata
         if not tp and "tensor_parallel_size" in extra_info:
-            record["benchmark"]["extra_info"]["args"]["tensor_parallel_size"] = extra_info["tensor_parallel_size"]
+            record["benchmark"]["extra_info"]["args"][
+                "tensor_parallel_size"] = extra_info["tensor_parallel_size"]
 
         records.append(record)
 
@@ -68,7 +68,6 @@ def convert_to_pytorch_benchmark_format(
 
 class InfEncoder(json.JSONEncoder):
     """InfEncoder"""
-
     def clear_inf(self, o: Any):
         """clear_inf"""
         if isinstance(o, dict):
@@ -88,3 +87,4 @@ def write_to_json(filename: str, records: list) -> None:
     """write_to_json"""
     with open(filename, "w") as f:
         json.dump(records, f, cls=InfEncoder)
+

benchmarks/quick_benchmark.py

@@ -25,32 +25,32 @@ import os
 import random
 import time
 import warnings
-from argparse import ArgumentParser as FlexibleArgumentParser
+import yaml
+import requests
+import copy
 from collections.abc import AsyncGenerator, Iterable
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Optional
 
 import numpy as np
-import requests
-import yaml
-from backend_request_func import (
-    ASYNC_REQUEST_FUNCS,
-    OPENAI_COMPATIBLE_BACKENDS,
-    RequestFuncInput,
-    RequestFuncOutput,
-)
-from benchmark_dataset import EBChatDataset, EBDataset, SampleRequest
-from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
+from backend_request_func import (ASYNC_REQUEST_FUNCS,
+                                  OPENAI_COMPATIBLE_BACKENDS, RequestFuncInput,
+                                  RequestFuncOutput)
 from tqdm.asyncio import tqdm
 
+from argparse import ArgumentParser as FlexibleArgumentParser
+
+from benchmark_dataset import (SampleRequest, EBDataset, EBChatDataset)
+from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
+
 MILLISECONDS_TO_SECONDS_CONVERSION = 1000
 
 
+
 @dataclass
 class BenchmarkMetrics:
     """Class containing all metrics that are used in this script"""
-
     completed: int
     total_input: int
     total_output: int
@@ -133,7 +133,8 @@ async def get_request(
     input_requests: Iterable[SampleRequest] = iter(input_requests)
 
     # Calculate scale parameter theta to maintain the desired request_rate.
-    assert burstiness > 0, f"A positive burstiness factor is expected, but given {burstiness}."
+    assert burstiness > 0, (
+        f"A positive burstiness factor is expected, but given {burstiness}.")
     theta = 1.0 / (request_rate * burstiness)
 
     for request in input_requests:
@@ -159,7 +160,7 @@ def calculate_metrics(
 ) -> tuple[BenchmarkMetrics, list[int]]:
     """Calculates various performance metrics based on the inputs and outputs."""
     input_lens: list[int] = []
-    infer_input_lens: list[int] = []  # 推理侧输入token数
+    infer_input_lens: list[int] = [] # 推理侧输入token数
     actual_output_lens: list[int] = []
     total_input = 0
     completed = 0
@@ -209,9 +210,8 @@ def calculate_metrics(
             s_e2els.append(outputs[i].arrival_time[-1])
             # 解码速度去掉首token
             if len(outputs[i].arrival_time) > 2:
-                s_decodes.append(
-                    (outputs[i].output_tokens - 1) / (outputs[i].arrival_time[-1] - outputs[i].arrival_time[1])
-                )
+                s_decodes.append((outputs[i].output_tokens - 1) /
+                                 (outputs[i].arrival_time[-1] - outputs[i].arrival_time[1]))
             completed += 1
         else:
             actual_output_lens.append(0)
@@ -224,13 +224,16 @@ def calculate_metrics(
 
         if "ttft" in goodput_config_dict:
             valid_metrics.append(ttfts)
-            slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
+            slo_values.append(goodput_config_dict["ttft"] /
+                              MILLISECONDS_TO_SECONDS_CONVERSION)
         if "tpot" in goodput_config_dict:
             valid_metrics.append(all_tpots)
-            slo_values.append(goodput_config_dict["tpot"] / MILLISECONDS_TO_SECONDS_CONVERSION)
+            slo_values.append(goodput_config_dict["tpot"] /
+                              MILLISECONDS_TO_SECONDS_CONVERSION)
         if "e2el" in goodput_config_dict:
             valid_metrics.append(e2els)
-            slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
+            slo_values.append(goodput_config_dict["e2el"] /
+                              MILLISECONDS_TO_SECONDS_CONVERSION)
 
         for req_metric in zip(*valid_metrics):
             is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
@@ -239,9 +242,9 @@ def calculate_metrics(
 
     if completed == 0:
         warnings.warn(
-            "All requests failed. This is likely due to a misconfiguration " "on the benchmark arguments.",
-            stacklevel=2,
-        )
+            "All requests failed. This is likely due to a misconfiguration "
+            "on the benchmark arguments.",
+            stacklevel=2)
     metrics = BenchmarkMetrics(
         completed=completed,
         total_input=total_input,
@@ -250,50 +253,64 @@ def calculate_metrics(
         request_goodput=good_completed / dur_s,
         output_throughput=sum(actual_output_lens) / dur_s,
         total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
-        mean_s_decode=np.mean(s_decodes or 0) * 1,  # ttfts is empty if streaming is not supported by backend
+        mean_s_decode=np.mean(s_decodes or 0) *
+                       1,  # ttfts is empty if streaming is not supported by backend
         std_s_decode=np.std(s_decodes or 0) * 1,
         median_s_decode=np.median(s_decodes or 0) * 1,
-        percentiles_s_decode=[(p, np.percentile(s_decodes or 0, p) * 1) for p in selected_percentiles],
-        mean_ttft_ms=np.mean(ttfts or 0) * 1000,  # ttfts is empty if streaming is not supported by backend
+        percentiles_s_decode=[(p, np.percentile(s_decodes or 0, p) * 1)
+                               for p in selected_percentiles],
+        mean_ttft_ms=np.mean(ttfts or 0) *
+        1000,  # ttfts is empty if streaming is not supported by backend
         std_ttft_ms=np.std(ttfts or 0) * 1000,
         median_ttft_ms=np.median(ttfts or 0) * 1000,
-        percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
-        mean_s_ttft_ms=np.mean(s_ttfts or 0) * 1000,  # ttfts is empty if streaming is not supported by backend
+        percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000)
+                             for p in selected_percentiles],
+        mean_s_ttft_ms=np.mean(s_ttfts or 0) *
+                     1000,  # ttfts is empty if streaming is not supported by backend
         std_s_ttft_ms=np.std(s_ttfts or 0) * 1000,
         median_s_ttft_ms=np.median(s_ttfts or 0) * 1000,
-        percentiles_s_ttft_ms=[(p, np.percentile(s_ttfts or 0, p) * 1000) for p in selected_percentiles],
+        percentiles_s_ttft_ms=[(p, np.percentile(s_ttfts or 0, p) * 1000)
+                             for p in selected_percentiles],
         mean_tpot_ms=np.mean(tpots or 0) * 1000,
         std_tpot_ms=np.std(tpots or 0) * 1000,
         median_tpot_ms=np.median(tpots or 0) * 1000,
-        percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
+        percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000)
+                             for p in selected_percentiles],
         mean_itl_ms=np.mean(itls or 0) * 1000,
         std_itl_ms=np.std(itls or 0) * 1000,
         median_itl_ms=np.median(itls or 0) * 1000,
-        percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
+        percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000)
+                            for p in selected_percentiles],
         mean_s_itl_ms=np.mean(s_itls or 0) * 1000,
         std_s_itl_ms=np.std(s_itls or 0) * 1000,
         median_s_itl_ms=np.median(s_itls or 0) * 1000,
-        percentiles_s_itl_ms=[(p, np.percentile(s_itls or 0, p) * 1000) for p in selected_percentiles],
+        percentiles_s_itl_ms=[(p, np.percentile(s_itls or 0, p) * 1000)
+                            for p in selected_percentiles],
         mean_e2el_ms=np.mean(e2els or 0) * 1000,
         std_e2el_ms=np.std(e2els or 0) * 1000,
         median_e2el_ms=np.median(e2els or 0) * 1000,
-        percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles],
+        percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000)
+                             for p in selected_percentiles],
         mean_s_e2el_ms=np.mean(s_e2els or 0) * 1000,
         std_s_e2el_ms=np.std(s_e2els or 0) * 1000,
         median_s_e2el_ms=np.median(s_e2els or 0) * 1000,
-        percentiles_s_e2el_ms=[(p, np.percentile(s_e2els or 0, p) * 1000) for p in selected_percentiles],
+        percentiles_s_e2el_ms=[(p, np.percentile(s_e2els or 0, p) * 1000)
+                             for p in selected_percentiles],
         mean_input_len=np.mean(input_lens or 0) * 1,
         std_input_len=np.std(input_lens or 0) * 1,
         median_input_len=np.median(input_lens or 0) * 1,
-        percentiles_input_len=[(p, np.percentile(input_lens or 0, p)) for p in selected_percentiles],
+        percentiles_input_len=[(p, np.percentile(input_lens or 0, p))
+                             for p in selected_percentiles],
         mean_s_input_len=np.mean(infer_input_lens or 0) * 1,
         std_s_input_len=np.std(infer_input_lens or 0) * 1,
         median_s_input_len=np.median(infer_input_lens or 0) * 1,
-        percentiles_s_input_len=[(p, np.percentile(infer_input_lens or 0, p)) for p in selected_percentiles],
+        percentiles_s_input_len=[(p, np.percentile(infer_input_lens or 0, p))
+                               for p in selected_percentiles],
         mean_output_len=np.mean(actual_output_lens or 0) * 1,
         std_output_len=np.std(actual_output_lens or 0) * 1,
         median_output_len=np.median(actual_output_lens or 0) * 1,
-        percentiles_output_len=[(p, np.percentile(actual_output_lens or 0, p)) for p in selected_percentiles],
+        percentiles_output_len=[(p, np.percentile(actual_output_lens or 0, p))
+                               for p in selected_percentiles],
     )
 
     return metrics, actual_output_lens
@@ -334,22 +351,20 @@ async def benchmark(
 
     if lora_modules:
         # For each input request, choose a LoRA module at random.
-        lora_modules = iter([random.choice(lora_modules) for _ in range(len(input_requests))])
+        lora_modules = iter(
+            [random.choice(lora_modules) \
+                for _ in range(len(input_requests))])
 
     if profile:
         print("Starting profiler...")
-        test_prompt = None
-        test_output_len = None
-        profile_input = RequestFuncInput(
-            model=model_id,
-            model_name=model_name,
-            prompt=test_prompt,
-            api_url=base_url + "/start_profile",
-            output_len=test_output_len,
-            logprobs=logprobs,
-            ignore_eos=ignore_eos,
-            extra_body=extra_body,
-        )
+        profile_input = RequestFuncInput(model=model_id,
+                                         model_name=model_name,
+                                         prompt=test_prompt,
+                                         api_url=base_url + "/start_profile",
+                                         output_len=test_output_len,
+                                         logprobs=logprobs,
+                                         ignore_eos=ignore_eos,
+                                         extra_body=extra_body)
         profile_output = await request_func(request_func_input=profile_input)
         if profile_output.success:
             print("Profiler started")
@@ -369,16 +384,19 @@ async def benchmark(
     # and it will simplify the code in limited_request_func.
     #    semaphore = (asyncio.Semaphore(max_concurrency)
     #                 if max_concurrency else contextlib.nullcontext())
-    semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None
+    semaphore = (asyncio.Semaphore(max_concurrency)
+                 if max_concurrency else None)
 
     async def limited_request_func(request_func_input, pbar):
         if semaphore is None:
-            return await request_func(request_func_input=request_func_input, pbar=pbar)
+            return await request_func(request_func_input=request_func_input,
+                                      pbar=pbar)
         async with semaphore:
-            return await request_func(request_func_input=request_func_input, pbar=pbar)
+            return await request_func(request_func_input=request_func_input,
+                                      pbar=pbar)
 
     benchmark_start_time = time.perf_counter()
-
+    
     print(f"开始时间：{datetime.now()}")
     tasks: list[asyncio.Task] = []
     async for request in get_request(input_requests, request_rate, burstiness):
@@ -391,26 +409,25 @@ async def benchmark(
             req_lora_module = next(lora_modules)
             req_model_id, req_model_name = req_lora_module, req_lora_module
 
-        request_func_input = RequestFuncInput(
-            model=req_model_id,
-            model_name=req_model_name,
-            prompt=prompt,
-            prompt_len=0,
-            history_QA=history_QA,
-            hyper_parameters=hyper_parameters,
-            api_url=api_url,
-            output_len=output_len,
-            logprobs=logprobs,
-            ignore_eos=ignore_eos,
-            extra_body=extra_body,
-        )
-        tasks.append(asyncio.create_task(limited_request_func(request_func_input=request_func_input, pbar=pbar)))
+        request_func_input = RequestFuncInput(model=req_model_id,
+                                              model_name=req_model_name,
+                                              prompt=prompt,
+                                              prompt_len=0,
+                                              history_QA=history_QA,
+                                              hyper_parameters=hyper_parameters,
+                                              api_url=api_url,
+                                              output_len=output_len,
+                                              logprobs=logprobs,
+                                              ignore_eos=ignore_eos,
+                                              extra_body=extra_body)
+        tasks.append(
+            asyncio.create_task(
+                limited_request_func(request_func_input=request_func_input,
+                                     pbar=pbar)))
     outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)
     print(f"完成时间：{datetime.now()}")
     if profile:
         print("Stopping profiler...")
-        test_output_len = None
-        test_output_len = None
         profile_input = RequestFuncInput(
             model=model_id,
             prompt=test_prompt,
@@ -437,16 +454,22 @@ async def benchmark(
     )
     print("Benchmark complete!!!")
 
-    print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
+    print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='='))
     print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
-    print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))
+    print("{:<40} {:<10.2f}".format("Benchmark duration (s):",
+                                    benchmark_duration))
     print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
-    print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
-    print("{:<40} {:<10.3f}".format("Request throughput (req/s):", metrics.request_throughput))
+    print("{:<40} {:<10}".format("Total generated tokens:",
+                                 metrics.total_output))
+    print("{:<40} {:<10.3f}".format("Request throughput (req/s):",
+                                    metrics.request_throughput))
     if goodput_config_dict:
-        print("{:<40} {:<10.2f}".format("Request goodput (req/s):", metrics.request_goodput))
-    print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):", metrics.output_throughput))
-    print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):", metrics.total_token_throughput))
+        print("{:<40} {:<10.2f}".format("Request goodput (req/s):",
+                                        metrics.request_goodput))
+    print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):",
+                                    metrics.output_throughput))
+    print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):",
+                                    metrics.total_token_throughput))
 
     result = {
         "duration": benchmark_duration,
@@ -454,7 +477,8 @@ async def benchmark(
         "total_input_tokens": metrics.total_input,
         "total_output_tokens": metrics.total_output,
         "request_throughput": metrics.request_throughput,
-        "request_goodput:": (metrics.request_goodput if goodput_config_dict else None),
+        "request_goodput:":
+        metrics.request_goodput if goodput_config_dict else None,
         "output_throughput": metrics.output_throughput,
         "total_token_throughput": metrics.total_token_throughput,
         "input_lens": [output.prompt_len for output in outputs],
@@ -467,6 +491,7 @@ async def benchmark(
         "reasoning_contents": [output.reasoning_content for output in outputs],
         "errors": [output.error for output in outputs],
     }
+    quick_result = copy.deepcopy(result)
 
     def process_one_metric(
         # E.g., "ttft"
@@ -480,25 +505,24 @@ async def benchmark(
         # metric.
         if metric_attribute_name not in selected_percentile_metrics:
             return
-        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
-        print(
-            "{:<40} {:<10.2f}".format(
-                f"Mean {metric_name} (ms):",
-                getattr(metrics, f"mean_{metric_attribute_name}_ms"),
-            )
-        )
-        print(
-            "{:<40} {:<10.2f}".format(
-                f"Median {metric_name} (ms):",
-                getattr(metrics, f"median_{metric_attribute_name}_ms"),
-            )
-        )
-        result[f"mean_{metric_attribute_name}_ms"] = getattr(metrics, f"mean_{metric_attribute_name}_ms")
-        result[f"median_{metric_attribute_name}_ms"] = getattr(metrics, f"median_{metric_attribute_name}_ms")
-        result[f"std_{metric_attribute_name}_ms"] = getattr(metrics, f"std_{metric_attribute_name}_ms")
-        for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}_ms"):
+        print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))
+        print("{:<40} {:<10.2f}".format(
+            f"Mean {metric_name} (ms):",
+            getattr(metrics, f"mean_{metric_attribute_name}_ms")))
+        print("{:<40} {:<10.2f}".format(
+            f"Median {metric_name} (ms):",
+            getattr(metrics, f"median_{metric_attribute_name}_ms")))
+        result[f"mean_{metric_attribute_name}_ms"] = getattr(
+            metrics, f"mean_{metric_attribute_name}_ms")
+        result[f"median_{metric_attribute_name}_ms"] = getattr(
+            metrics, f"median_{metric_attribute_name}_ms")
+        result[f"std_{metric_attribute_name}_ms"] = getattr(
+            metrics, f"std_{metric_attribute_name}_ms")
+        for p, value in getattr(metrics,
+                                f"percentiles_{metric_attribute_name}_ms"):
             p_word = str(int(p)) if int(p) == p else str(p)
-            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", value))
+            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):",
+                                            value))
             result[f"p{p_word}_{metric_attribute_name}_ms"] = value
 
     def process_one_length(
@@ -513,31 +537,31 @@ async def benchmark(
         # metric.
         if metric_attribute_name not in selected_percentile_metrics:
             return
-        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
-        print(
-            "{:<40} {:<10.2f}".format(
-                f"Mean {metric_name}:",
-                getattr(metrics, f"mean_{metric_attribute_name}"),
-            )
-        )
-        print(
-            "{:<40} {:<10.2f}".format(
-                f"Median {metric_name}:",
-                getattr(metrics, f"median_{metric_attribute_name}"),
-            )
-        )
-        result[f"mean_{metric_attribute_name}"] = getattr(metrics, f"mean_{metric_attribute_name}")
-        result[f"median_{metric_attribute_name}"] = getattr(metrics, f"median_{metric_attribute_name}")
-        result[f"std_{metric_attribute_name}"] = getattr(metrics, f"std_{metric_attribute_name}")
-        for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}"):
+        print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))
+        print("{:<40} {:<10.2f}".format(
+            f"Mean {metric_name}:",
+            getattr(metrics, f"mean_{metric_attribute_name}")))
+        print("{:<40} {:<10.2f}".format(
+            f"Median {metric_name}:",
+            getattr(metrics, f"median_{metric_attribute_name}")))
+        result[f"mean_{metric_attribute_name}"] = getattr(
+            metrics, f"mean_{metric_attribute_name}")
+        result[f"median_{metric_attribute_name}"] = getattr(
+            metrics, f"median_{metric_attribute_name}")
+        result[f"std_{metric_attribute_name}"] = getattr(
+            metrics, f"std_{metric_attribute_name}")
+        for p, value in getattr(metrics,
+                                f"percentiles_{metric_attribute_name}"):
             p_word = str(int(p)) if int(p) == p else str(p)
-            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name}:", value))
+            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name}:",
+                                            value))
             result[f"p{p_word}_{metric_attribute_name}"] = value
 
     process_one_length("s_decode", "Decode", "解码速度(tok/s)")
     process_one_metric("ttft", "TTFT", "Time to First Token")
     process_one_metric("s_ttft", "S_TTFT", "Infer Time to First Token")
-    process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")
+    process_one_metric("tpot", "TPOT",
+                       "Time per Output Token (excl. 1st token)")
     process_one_metric("itl", "ITL", "Inter-token Latency")
     process_one_metric("s_itl", "S_ITL", "Infer Inter-token Latency")
     process_one_metric("e2el", "E2EL", "End-to-end Latency")
@@ -557,7 +581,6 @@ def quick_summary(quick_result, selected_percentile_metrics, metrics):
     """
     快速评估
     """
-
     def process_quick_metric(
         metric_attribute_name: str,
         metric_name: str,
@@ -565,7 +588,7 @@ def quick_summary(quick_result, selected_percentile_metrics, metrics):
     ):
         if metric_attribute_name not in selected_percentile_metrics:
             return
-        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
+        print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))
         mean_value = getattr(metrics, f"mean_{metric_attribute_name}_ms")
         print("{:<40} {:<10.2f}".format(f"Mean {metric_name} (ms):", mean_value))
         quick_result[f"mean_{metric_attribute_name}_ms"] = mean_value
@@ -577,17 +600,17 @@ def quick_summary(quick_result, selected_percentile_metrics, metrics):
     ):
         if metric_attribute_name not in selected_percentile_metrics:
             return
-        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
+        print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))
         mean_value = getattr(metrics, f"mean_{metric_attribute_name}")
         print("{:<40} {:<10.2f}".format(f"Mean {metric_name}:", mean_value))
         quick_result[f"mean_{metric_attribute_name}"] = mean_value
-
     print("\n\n\n")
-    print("{s:{c}^{n}}".format(s=" Benchmark Quick Summary ", n=50, c="="))
+    print("{s:{c}^{n}}".format(s=' Benchmark Quick Summary ', n=50, c='='))
     process_quick_length("s_decode", "Decode", "解码速度(tok/s)")
     process_quick_metric("ttft", "TTFT", "Time to First Token")
     process_quick_metric("s_ttft", "S_TTFT", "Infer Time to First Token")
-    process_quick_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")
+    process_quick_metric("tpot", "TPOT",
+                       "Time per Output Token (excl. 1st token)")
     process_quick_metric("itl", "ITL", "Inter-token Latency")
     process_quick_metric("s_itl", "S_ITL", "Infer Inter-token Latency")
     process_quick_metric("e2el", "E2EL", "End-to-end Latency")
@@ -610,14 +633,12 @@ def check_goodput_args(args):
                 raise ValueError(
                     f"Invalid metric name found, {slo_name}: {slo_val}. "
                     "The service level objective name should be one of "
-                    f"{VALID_NAMES!s}. "
-                )
+                    f"{str(VALID_NAMES)}. ")
             if slo_val < 0:
                 raise ValueError(
                     f"Invalid value found, {slo_name}: {slo_val}. "
                     "The service level objective value should be "
-                    "non-negative."
-                )
+                    "non-negative.")
     return goodput_config_dict
 
 
@@ -631,43 +652,37 @@ def parse_goodput(slo_pairs):
     except ValueError as err:
         raise argparse.ArgumentTypeError(
             "Invalid format found for service level objectives. "
-            'Specify service level objectives for goodput as "KEY:VALUE" '
+            "Specify service level objectives for goodput as \"KEY:VALUE\" "
             "pairs, where the key is a metric name, and the value is a "
-            "number in milliseconds."
-        ) from err
+            "number in milliseconds.") from err
     return goodput_config_dict
 
 
-def save_to_pytorch_benchmark_format(args: argparse.Namespace, results: dict[str, Any], file_name: str) -> None:
+def save_to_pytorch_benchmark_format(args: argparse.Namespace,
+                                     results: dict[str, Any],
+                                     file_name: str) -> None:
     """Save the benchmarking results to PyTorch Benchmark Format JSON file"""
     metrics = [
-        "median_ttft_ms",
-        "mean_ttft_ms",
-        "std_ttft_ms",
-        "p99_ttft_ms",
-        "mean_tpot_ms",
-        "median_tpot_ms",
-        "std_tpot_ms",
-        "p99_tpot_ms",
-        "median_itl_ms",
-        "mean_itl_ms",
-        "std_itl_ms",
-        "p99_itl_ms",
+        "median_ttft_ms", "mean_ttft_ms", "std_ttft_ms", "p99_ttft_ms",
+        "mean_tpot_ms", "median_tpot_ms", "std_tpot_ms", "p99_tpot_ms",
+        "median_itl_ms", "mean_itl_ms", "std_itl_ms", "p99_itl_ms"
     ]
     # These raw data might be useful, but they are rather big. They can be added
     # later if needed
     ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"]
     pt_records = convert_to_pytorch_benchmark_format(
         args=args,
-        metrics={k: [results[k]] for k in metrics},
-        extra_info={k: results[k] for k in results if k not in metrics and k not in ignored_metrics},
-    )
+        metrics={k: [results[k]]
+                 for k in metrics},
+        extra_info={
+            k: results[k]
+            for k in results if k not in metrics and k not in ignored_metrics
+        })
     if pt_records:
         # Don't use json suffix here as we don't want CI to pick it up
         pt_file = f"{os.path.splitext(file_name)[0]}.pytorch.json"
         write_to_json(pt_file, pt_records)
 
-
 def check_health(api_base_url: str) -> bool:
     health_url = api_base_url.rstrip("/") + "/health"
     try:
@@ -682,7 +697,6 @@ def check_health(api_base_url: str) -> bool:
         print(f"[HEALTH] Failed to connect to {health_url}: {e}")
         return False
 
-
 def main(args: argparse.Namespace):
     """Main entry point"""
     print(args)
@@ -693,6 +707,7 @@ def main(args: argparse.Namespace):
     model_id = args.model
     model_name = args.served_model_name
     tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
+    tokenizer_mode = args.tokenizer_mode
 
     if args.base_url is not None:
         api_url = f"{args.base_url}{args.endpoint}"
@@ -702,17 +717,23 @@ def main(args: argparse.Namespace):
         base_url = f"http://{args.host}:{args.port}"
 
     if args.dataset_name is None:
-        raise ValueError("Please specify '--dataset-name' and the corresponding " "'--dataset-path' if required.")
+        raise ValueError(
+            "Please specify '--dataset-name' and the corresponding "
+            "'--dataset-path' if required.")
 
     # For datasets that follow a similar structure, use a mapping.
     dataset_mapping = {
-        "EB": lambda: EBDataset(random_seed=args.seed, dataset_path=args.dataset_path).sample(
-            num_requests=args.num_prompts,
-            output_len=args.sharegpt_output_len,
+        "EB":
+        lambda: EBDataset(random_seed=args.seed,
+                          dataset_path=args.dataset_path).sample(
+                              num_requests=args.num_prompts,
+                              output_len=args.sharegpt_output_len,
         ),
-        "EBChat": lambda: EBChatDataset(random_seed=args.seed, dataset_path=args.dataset_path).sample(
-            num_requests=args.num_prompts,
-            output_len=args.sharegpt_output_len,
+        "EBChat":
+            lambda: EBChatDataset(random_seed=args.seed,
+                                  dataset_path=args.dataset_path).sample(
+                num_requests=args.num_prompts,
+                output_len=args.sharegpt_output_len,
         ),
     }
 
@@ -730,14 +751,15 @@ def main(args: argparse.Namespace):
             "top_p": args.top_p,
             "top_k": args.top_k,
             "min_p": args.min_p,
-            "temperature": args.temperature,
-        }.items()
-        if v is not None
+            "temperature": args.temperature
+        }.items() if v is not None
     }
 
     # Sampling parameters are only supported by openai-compatible backend.
     if sampling_params and args.backend not in OPENAI_COMPATIBLE_BACKENDS:
-        raise ValueError("Sampling parameters are only supported by openai-compatible " "backends.")
+        raise ValueError(
+            "Sampling parameters are only supported by openai-compatible "
+            "backends.")
 
     if "temperature" not in sampling_params:
         sampling_params["temperature"] = 0.0  # Default to greedy decoding.
@@ -768,14 +790,15 @@ def main(args: argparse.Namespace):
             disable_tqdm=args.disable_tqdm,
             profile=args.profile,
             selected_percentile_metrics=args.percentile_metrics.split(","),
-            selected_percentiles=[float(p) for p in args.metric_percentiles.split(",")],
+            selected_percentiles=[
+                float(p) for p in args.metric_percentiles.split(",")
+            ],
             ignore_eos=args.ignore_eos,
             goodput_config_dict=goodput_config_dict,
             max_concurrency=args.max_concurrency,
             lora_modules=args.lora_modules,
             extra_body=sampling_params,
-        )
-    )
+        ))
 
     # Save config and results to json
     if args.save_result:
@@ -796,23 +819,22 @@ def main(args: argparse.Namespace):
                     kvstring = item.split("=")
                     result_json[kvstring[0].strip()] = kvstring[1].strip()
                 else:
-                    raise ValueError("Invalid metadata format. Please use KEY=VALUE format.")
+                    raise ValueError(
+                        "Invalid metadata format. Please use KEY=VALUE format."
+                    )
 
         if not args.save_detailed:
             # Remove fields with too many data points
             for field in [
-                "input_lens",
-                "output_lens",
-                "ttfts",
-                "itls",
-                "generated_texts",
-                "errors",
+                    "input_lens", "output_lens", "ttfts", "itls",
+                    "generated_texts", "errors"
             ]:
                 if field in result_json:
                     del result_json[field]
 
         # Traffic
-        result_json["request_rate"] = args.request_rate if args.request_rate < float("inf") else "inf"
+        result_json["request_rate"] = (args.request_rate if args.request_rate
+                                       < float("inf") else "inf")
         result_json["burstiness"] = args.burstiness
         result_json["max_concurrency"] = args.max_concurrency
 
@@ -821,19 +843,21 @@ def main(args: argparse.Namespace):
 
         # Save to file
         base_model_id = model_id.split("/")[-1]
-        max_concurrency_str = f"-concurrency{args.max_concurrency}" if args.max_concurrency is not None else ""
-        file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"
+        max_concurrency_str = (f"-concurrency{args.max_concurrency}"
+                               if args.max_concurrency is not None else "")
+        file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  #noqa
         if args.result_filename:
             file_name = args.result_filename
         if args.result_dir:
             file_name = os.path.join(args.result_dir, file_name)
-        with open(file_name, "w", encoding="utf-8") as outfile:
+        with open(file_name, "w", encoding='utf-8') as outfile:
             json.dump(result_json, outfile)
         save_to_pytorch_benchmark_format(args, result_json, file_name)
 
 
 if __name__ == "__main__":
-    parser = FlexibleArgumentParser(description="Benchmark the online serving throughput.")
+    parser = FlexibleArgumentParser(
+        description="Benchmark the online serving throughput.")
     parser.add_argument(
         "--backend",
         type=str,
@@ -859,29 +883,18 @@ if __name__ == "__main__":
         "--dataset-name",
         type=str,
         default="sharegpt",
-        choices=[
-            "sharegpt",
-            "burstgpt",
-            "sonnet",
-            "random",
-            "hf",
-            "EB",
-            "EBChat",
-        ],
+        choices=["sharegpt", "burstgpt", "sonnet", "random", "hf", "EB", "EBChat"],
         help="Name of the dataset to benchmark on.",
     )
-    parser.add_argument(
-        "--dataset-path",
-        type=str,
-        default=None,
-        help="Path to the sharegpt/sonnet dataset. " "Or the huggingface dataset ID if using HF dataset.",
-    )
-    parser.add_argument(
-        "--hyperparameter-path",
-        type=str,
-        default=None,
-        help="Path to the hyperparameter. ",
-    )
+    parser.add_argument("--dataset-path",
+                        type=str,
+                        default=None,
+                        help="Path to the sharegpt/sonnet dataset. "
+                        "Or the huggingface dataset ID if using HF dataset.")
+    parser.add_argument("--hyperparameter-path",
+                        type=str,
+                        default=None,
+                        help="Path to the hyperparameter. ")
     parser.add_argument(
         "--max-concurrency",
         type=int,
@@ -893,8 +906,7 @@ if __name__ == "__main__":
         "initiated, this argument will control how many are actually allowed "
         "to execute at a time. This means that when used in combination, the "
         "actual request rate may be lower than specified with --request-rate, "
-        "if the server is not processing requests fast enough to keep up.",
-    )
+        "if the server is not processing requests fast enough to keep up.")
 
     parser.add_argument(
         "--model",
@@ -905,7 +917,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--tokenizer",
         type=str,
-        help="Name or path of the tokenizer, if not using the default tokenizer.",
+        help="Name or path of the tokenizer, if not using the default tokenizer.",  # noqa: E501
     )
     parser.add_argument("--use-beam-search", action="store_true")
     parser.add_argument(
@@ -918,13 +930,11 @@ if __name__ == "__main__":
         "--logprobs",
         type=int,
         default=None,
-        help=(
-            "Number of logprobs-per-token to compute & return as part of "
-            "the request. If unspecified, then either (1) if beam search "
-            "is disabled, no logprobs are computed & a single dummy "
-            "logprob is returned for each token; or (2) if beam search "
-            "is enabled 1 logprob per token is computed"
-        ),
+        help=("Number of logprobs-per-token to compute & return as part of "
+              "the request. If unspecified, then either (1) if beam search "
+              "is disabled, no logprobs are computed & a single dummy "
+              "logprob is returned for each token; or (2) if beam search "
+              "is enabled 1 logprob per token is computed"),
     )
     parser.add_argument(
         "--request-rate",
@@ -961,7 +971,8 @@ if __name__ == "__main__":
     parser.add_argument(
         "--profile",
         action="store_true",
-        help="Use Torch Profiler. The endpoint must be launched with " "VLLM_TORCH_PROFILER_DIR to enable profiler.",
+        help="Use Torch Profiler. The endpoint must be launched with "
+        "VLLM_TORCH_PROFILER_DIR to enable profiler.",
     )
     parser.add_argument(
         "--save-result",
@@ -1002,38 +1013,35 @@ if __name__ == "__main__":
         "--ignore-eos",
         action="store_true",
         help="Set ignore_eos flag when sending the benchmark request."
-        "Warning: ignore_eos is not supported in deepspeed_mii and tgi.",
-    )
+        "Warning: ignore_eos is not supported in deepspeed_mii and tgi.")
     parser.add_argument(
         "--percentile-metrics",
         type=str,
         default="ttft,tpot,itl",
         help="Comma-separated list of selected metrics to report percentils. "
         "This argument specifies the metrics to report percentiles. "
-        'Allowed metric names are "ttft", "tpot", "itl", "e2el". '
-        'Default value is "ttft,tpot,itl".',
-    )
+        "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". "
+        "Default value is \"ttft,tpot,itl\".")
     parser.add_argument(
         "--metric-percentiles",
         type=str,
         default="99",
         help="Comma-separated list of percentiles for selected metrics. "
-        'To report 25-th, 50-th, and 75-th percentiles, use "25,50,75". '
-        'Default value is "99". '
-        'Use "--percentile-metrics" to select metrics.',
+        "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
+        "Default value is \"99\". "
+        "Use \"--percentile-metrics\" to select metrics.",
     )
     parser.add_argument(
         "--goodput",
         nargs="+",
         required=False,
-        help='Specify service level objectives for goodput as "KEY:VALUE" '
+        help="Specify service level objectives for goodput as \"KEY:VALUE\" "
         "pairs, where the key is a metric name, and the value is in "
-        'milliseconds. Multiple "KEY:VALUE" pairs can be provided, '
+        "milliseconds. Multiple \"KEY:VALUE\" pairs can be provided, "
         "separated by spaces. Allowed request level metric names are "
-        '"ttft", "tpot", "e2el". For more context on the definition of '
+        "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of "
         "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "
-        "and the blog: https://hao-ai-lab.github.io/blogs/distserve",
-    )
+        "and the blog: https://hao-ai-lab.github.io/blogs/distserve")
 
     # group for dataset specific arguments
     sonnet_group = parser.add_argument_group("sonnet dataset options")
@@ -1061,8 +1069,8 @@ if __name__ == "__main__":
         "--sharegpt-output-len",
         type=int,
         default=None,
-        help="Output length for each request. Overrides the output length " "from the ShareGPT dataset.",
-    )
+        help="Output length for each request. Overrides the output length "
+        "from the ShareGPT dataset.")
 
     random_group = parser.add_argument_group("random dataset options")
     random_group.add_argument(
@@ -1090,24 +1098,29 @@ if __name__ == "__main__":
         "--random-prefix-len",
         type=int,
         default=0,
-        help=(
-            "Number of fixed prefix tokens before the random context "
-            "in a request. "
-            "The total input length is the sum of `random-prefix-len` and "
-            "a random "
-            "context length sampled from [input_len * (1 - range_ratio), "
-            "input_len * (1 + range_ratio)]."
-        ),
+        help=("Number of fixed prefix tokens before the random context "
+              "in a request. "
+              "The total input length is the sum of `random-prefix-len` and "
+              "a random "
+              "context length sampled from [input_len * (1 - range_ratio), "
+              "input_len * (1 + range_ratio)]."),
     )
 
     hf_group = parser.add_argument_group("hf dataset options")
-    hf_group.add_argument("--hf-subset", type=str, default=None, help="Subset of the HF dataset.")
-    hf_group.add_argument("--hf-split", type=str, default=None, help="Split of the HF dataset.")
+    hf_group.add_argument("--hf-subset",
+                          type=str,
+                          default=None,
+                          help="Subset of the HF dataset.")
+    hf_group.add_argument("--hf-split",
+                          type=str,
+                          default=None,
+                          help="Split of the HF dataset.")
     hf_group.add_argument(
         "--hf-output-len",
         type=int,
         default=None,
-        help="Output length for each request. Overrides the output lengths " "from the sampled HF dataset.",
+        help="Output length for each request. Overrides the output lengths "
+        "from the sampled HF dataset.",
     )
 
     sampling_group = parser.add_argument_group("sampling parameters")
@@ -1115,58 +1128,52 @@ if __name__ == "__main__":
         "--top-p",
         type=float,
         default=None,
-        help="Top-p sampling parameter. Only has effect on openai-compatible " "backends.",
-    )
+        help="Top-p sampling parameter. Only has effect on openai-compatible "
+        "backends.")
     sampling_group.add_argument(
         "--top-k",
         type=int,
         default=None,
-        help="Top-k sampling parameter. Only has effect on openai-compatible " "backends.",
-    )
+        help="Top-k sampling parameter. Only has effect on openai-compatible "
+        "backends.")
     sampling_group.add_argument(
         "--min-p",
         type=float,
         default=None,
-        help="Min-p sampling parameter. Only has effect on openai-compatible " "backends.",
-    )
+        help="Min-p sampling parameter. Only has effect on openai-compatible "
+        "backends.")
     sampling_group.add_argument(
         "--temperature",
         type=float,
         default=None,
         help="Temperature sampling parameter. Only has effect on "
         "openai-compatible backends. If not specified, default to greedy "
-        "decoding (i.e. temperature==0.0).",
-    )
+        "decoding (i.e. temperature==0.0).")
 
     parser.add_argument(
-        "--tokenizer-mode",
+        '--tokenizer-mode',
         type=str,
         default="auto",
-        choices=["auto", "slow", "mistral", "custom"],
+        choices=['auto', 'slow', 'mistral', 'custom'],
         help='The tokenizer mode.\n\n* "auto" will use the '
         'fast tokenizer if available.\n* "slow" will '
-        "always use the slow tokenizer. \n* "
+        'always use the slow tokenizer. \n* '
         '"mistral" will always use the `mistral_common` tokenizer. \n*'
-        '"custom" will use --tokenizer to select the preregistered tokenizer.',
-    )
+        '"custom" will use --tokenizer to select the preregistered tokenizer.')
 
-    parser.add_argument(
-        "--served-model-name",
-        type=str,
-        default=None,
-        help="The model name used in the API. "
-        "If not specified, the model name will be the "
-        "same as the ``--model`` argument. ",
-    )
+    parser.add_argument("--served-model-name",
+                        type=str,
+                        default=None,
+                        help="The model name used in the API. "
+                        "If not specified, the model name will be the "
+                        "same as the ``--model`` argument. ")
 
-    parser.add_argument(
-        "--lora-modules",
-        nargs="+",
-        default=None,
-        help="A subset of LoRA module names passed in when "
-        "launching the server. For each request, the "
-        "script chooses a LoRA module at random.",
-    )
+    parser.add_argument("--lora-modules",
+                        nargs='+',
+                        default=None,
+                        help="A subset of LoRA module names passed in when "
+                        "launching the server. For each request, the "
+                        "script chooses a LoRA module at random.")
 
     args = parser.parse_args()
 

benchmarks/yaml/eb45-128k-wint4-a800-tp8.yaml

@@ -6,4 +6,3 @@ tensor_parallel_size: 8
 max_num_batched_tokens: 4096
 max_num_partial_prefills: 3
 max_long_partial_prefills: 3
-quantization: wint4

benchmarks/yaml/eb45-128k-wint4-tp1-plas.yaml

@@ -1,6 +0,0 @@
-tensor_parallel_size: 1
-max_model_len: 131072
-max_num_seqs: 32
-quantization: wint4
-max_num_batched_tokens: 8192
-plas_attention_config: '{"plas_encoder_top_k_left": 50, "plas_encoder_top_k_right": 60, "plas_decoder_top_k_left": 100, "plas_decoder_top_k_right": 120}'

benchmarks/yaml/eb45-128k-wint8-a800-tp8.yaml

@@ -6,4 +6,3 @@ tensor_parallel_size: 8
 max_num_batched_tokens: 4096
 max_num_partial_prefills: 3
 max_long_partial_prefills: 3
-quantization: wint8

benchmarks/yaml/eb45-21B-vl-128k-wint4-h800-tp1.yaml

@@ -7,4 +7,4 @@ tensor_parallel_size: 1
 enable_chunked_prefill: True
 max_num_batched_tokens: 384
 quantization: wint4
-reasoning_parser: ernie-45-vl
+reasoning_parser: ernie-45-vl

benchmarks/yaml/eb45-32k-w4a8c8-tp4_decode.yaml

@@ -12,4 +12,4 @@ rdma_comm_ports: "7671,7672,7673,7674"
 pd_comm_port: "2334"
 max_num_batched_tokens: 384
 max_num_partial_prefills: 3
-max_long_partial_prefills: 3
+max_long_partial_prefills: 3

benchmarks/yaml/eb45-32k-w4a8c8-tp4_prefill.yaml

@@ -9,4 +9,4 @@ cache_queue_port: 55664
 engine_worker_queue_port: 6677
 cache_transfer_protocol: "rdma,ipc"
 rdma_comm_ports: "7675,7676,7677,7678"
-pd_comm_port: "2333"
+pd_comm_port: "2333"

benchmarks/yaml/eb45-32k-wint2-tp4.yaml

@@ -1,5 +0,0 @@
-max_model_len: 32768
-max_num_seqs: 256
-kv_cache_ratio: 0.75
-tensor_parallel_size: 4
-gpu_memory_utilization: 0.9

benchmarks/yaml/eb45-32k-wint4-a800-tp4.yaml

@@ -3,4 +3,3 @@ max_num_seqs: 96
 gpu_memory_utilization: 0.9
 kv_cache_ratio: 0.71
 tensor_parallel_size: 4
-quantization: wint4

benchmarks/yaml/eb45-32k-wint4-h800-dp8_prefill.yaml

@@ -10,4 +10,4 @@ engine_worker_queue_port: 6677
 num_gpu_blocks_override: 1024
 cache_transfer_protocol: "rdma"
 rdma_comm_ports: "7671,7672,7673,7674,7675,7676,7677,7678"
-pd_comm_port: "2334"
+pd_comm_port: "2334"

benchmarks/yaml/eb45-32k-wint4-mtp-tp4-decode.yaml

@@ -10,4 +10,4 @@ splitwise_role: decode
 engine_worker_queue_port: 6678
 cache_transfer_protocol: "rdma,ipc"
 rdma_comm_ports: "7671,7672,7673,7674"
-pd_comm_port: "2334"
+pd_comm_port: "2334"

benchmarks/yaml/eb45-32k-wint4-mtp-tp4-prefill.yaml

@@ -9,4 +9,4 @@ cache_queue_port: 55664
 engine_worker_queue_port: 6677
 cache_transfer_protocol: "rdma,ipc"
 rdma_comm_ports: "7675,7676,7677,7678"
-pd_comm_port: "2333"
+pd_comm_port: "2333"

benchmarks/yaml/eb45-32k-wint4-tp4_decode.yaml

@@ -12,5 +12,4 @@ rdma_comm_ports: "7671,7672,7673,7674"
 pd_comm_port: "2334"
 max_num_batched_tokens: 384
 max_num_partial_prefills: 3
-max_long_partial_prefills: 3
-quantization: wint4
+max_long_partial_prefills: 3

benchmarks/yaml/eb45-32k-wint4-tp4_prefill.yaml

@@ -9,5 +9,4 @@ cache_queue_port: 55664
 engine_worker_queue_port: 6677
 cache_transfer_protocol: "rdma,ipc"
 rdma_comm_ports: "7675,7676,7677,7678"
-pd_comm_port: "2333"
-quantization: wint4
+pd_comm_port: "2333"

benchmarks/yaml/eb45-32k-wint8-a800-tp8.yaml

@@ -3,4 +3,3 @@ max_num_seqs: 96
 gpu_memory_utilization: 0.9
 kv_cache_ratio: 0.71
 tensor_parallel_size: 8
-quantization: wint8

benchmarks/yaml/eb45-8k-fp8-tp1-dp8_ep.yaml

@@ -1,6 +0,0 @@
-num_gpu_blocks_override: 1024
-max_model_len: 8192
-max_num_seqs: 64
-data_parallel_size: 8
-tensor_parallel_size: 1
-enable_expert_parallel: True

benchmarks/yaml/eb45-vl-128k-wint4-h800-tp8.yaml

@@ -1,11 +0,0 @@
-enable_mm: True
-max_model_len: 131072
-max_num_seqs: 56
-gpu_memory_utilization: 0.8
-kv_cache_ratio: 0.8
-tensor_parallel_size: 8
-quantization: wint4
-limit_mm_per_prompt: '{"image": 100, "video": 100}'
-enable_chunked_prefill: True
-max_num_batched_tokens: 384
-reasoning_parser: ernie-45-vl

benchmarks/yaml/eb45-vl-32k-wint8-a800-tp8.yaml

@@ -1,7 +1,7 @@
 enable_mm: True
 max_model_len: 32768
 max_num_seqs: 36
-gpu_memory_utilization: 0.9
+gpu_memory_utilization: 0.95
 kv_cache_ratio: 0.8
 tensor_parallel_size: 8
 quantization: wint8

benchmarks/yaml/eb45-vl-32k-wint8-h800-tp8.yaml

@@ -1,7 +1,7 @@
 enable_mm: True
 max_model_len: 32768
 max_num_seqs: 36
-gpu_memory_utilization: 0.85
+gpu_memory_utilization: 0.8
 kv_cache_ratio: 0.8
 tensor_parallel_size: 8
 quantization: wint8

benchmarks/yaml/eb45-vl-lite-32k-bf16-a800-tp1.yaml

@@ -1,9 +0,0 @@
-enable_mm: True
-max_model_len: 32768
-max_num_seqs: 128
-gpu_memory_utilization: 0.9
-kv_cache_ratio: 0.71
-tensor_parallel_size: 1
-enable_chunked_prefill: True
-max_num_batched_tokens: 384
-reasoning_parser: ernie-45-vl

benchmarks/yaml/eb45-vl-lite-32k-wint4-a800-tp1.yaml

@@ -1,10 +0,0 @@
-enable_mm: True
-max_model_len: 32768
-max_num_seqs: 128
-gpu_memory_utilization: 0.9
-kv_cache_ratio: 0.71
-tensor_parallel_size: 1
-enable_chunked_prefill: True
-max_num_batched_tokens: 384
-quantization: wint4
-reasoning_parser: ernie-45-vl

benchmarks/yaml/eb45-vl-lite-32k-wint8-a800-tp1.yaml

@@ -1,10 +0,0 @@
-enable_mm: True
-max_model_len: 32768
-max_num_seqs: 128
-gpu_memory_utilization: 0.9
-kv_cache_ratio: 0.71
-tensor_parallel_size: 1
-enable_chunked_prefill: True
-max_num_batched_tokens: 384
-quantization: wint8
-reasoning_parser: ernie-45-vl

benchmarks/yaml/eb45t_0dot3b-32k-bf16-a30-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/eb45t_0dot3b-32k-bf16-h800-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/eb45t_0dot3b-32k-wint8-a30-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wint8
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/eb45t_0dot3b-32k-wint8-h800-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wint8
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/eb45t_21b-32k-bf16-h800-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/eb45t_21b-32k-wint4-h800-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wint4
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/eb45t_300b-32k-wint4-h800-tp4-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 96
 gpu_memory_utilization: 0.9
 kv_cache_ratio: 0.71
 tensor_parallel_size: 4
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen2_7b-32k-bf16-a30-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen2_7b-32k-bf16-h800-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen2_7b-32k-fp8-h800-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wfp8afp8
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3_0dot6b-32k-bf16-a30-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3_0dot6b-32k-bf16-h800-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3_0dot6b-32k-wint8-a30-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wint8
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3_0dot6b-32k-wint8-h800-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wint8
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3_30b-32k-bf16-h800-tp1-static.yaml

@@ -2,5 +2,4 @@ max_model_len: 32768
 max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3_30b-32k-wint4-h800-tp1-static.yaml

@@ -3,5 +3,4 @@ max_num_seqs: 128
 kv_cache_ratio: 0.75
 tensor_parallel_size: 1
 quantization: wint4
-graph_optimization_config:
-  graph_opt_level: 1
+enable_static_graph_inference: True

benchmarks/yaml/qwen3moe235b-32k-wint4-h800-tp4.yaml

@@ -3,4 +3,4 @@ max_num_seqs: 75
 gpu_memory_utilization: 0.85
 kv_cache_ratio: 0.75
 quantization: wint4
-tensor_parallel_size: 4
+tensor_parallel_size: 4

benchmarks/yaml/qwen3moe235b-32k-wint8-h800-tp4.yaml

@@ -3,4 +3,4 @@ max_num_seqs: 25
 gpu_memory_utilization: 0.9
 kv_cache_ratio: 0.75
 quantization: wint8
-tensor_parallel_size: 4
+tensor_parallel_size: 4

benchmarks/yaml/request_yaml/eb45-vl-128k.yaml

@@ -1 +0,0 @@
-max_tokens: 131071

benchmarks/yaml/request_yaml/eb45-vl-32k.yaml

@@ -1 +0,0 @@
-max_tokens: 12288

benchmarks/yaml/request_yaml/quick_benchmark.yaml

@@ -1,3 +1,3 @@
 metadata:
   min_tokens: 32
-max_tokens: 33
+max_tokens: 33

benchmarks/yaml/request_yaml/qwen2-32k.yaml

@@ -5,4 +5,4 @@ metadata:
 max_tokens: 12288
 repetition_penalty: 1.05
 frequency_penalty: 0
-presence_penalty: 0
+presence_penalty: 0

benchmarks/yaml/request_yaml/qwen3-32k.yaml

@@ -5,4 +5,4 @@ metadata:
 max_tokens: 12288
 repetition_penalty: 1.0
 frequency_penalty: 0
-presence_penalty: 1.5
+presence_penalty: 1.5

benchmarks/yaml/request_yaml/vLLM_default.yaml

@@ -1,11 +0,0 @@
-top_p: 1.0
-temperature: 1.0
-metadata:
-  min_tokens: 1
-max_tokens: 30721
-repetition_penalty: 1.0
-frequency_penalty: 0
-presence_penalty: 0
-skip_special_tokens: false
-chat_template_kwargs:
-  enable_thinking: true

benchmarks/yaml/request_yaml/x1-128k.yaml

@@ -1,8 +0,0 @@
-top_p: 0.95
-temperature: 0.6
-metadata:
-  min_tokens: 1
-max_tokens: 131071
-repetition_penalty: 1.0
-frequency_penalty: 0
-presence_penalty: 0

benchmarks/yaml/x1-32k-wint8-p800-tp8.yaml

@@ -3,4 +3,4 @@ max_num_seqs: 64
 gpu_memory_utilization: 0.9
 tensor_parallel_size: 8
 quantization: wint8
-reasoning_parser: ernie-x1
+reasoning_parser: ernie-x1

benchmarks/yaml/x1-64k-w4a8c8-tp4.yaml

@@ -1,10 +0,0 @@
-reasoning-parser: ernie_x1
-tool_call_parser: ernie_x1
-tensor_parallel_size: 4
-max_model_len: 65536
-max_num_seqs: 128
-enable_prefix_caching: True
-enable_chunked_prefill: True
-gpu_memory_utilization: 0.85
-use_cudagraph: True
-enable_custom_all_reduce: True

benchmarks/yaml/x1-a3b-128k-wint8-h800-tp1.yaml

@@ -1,6 +0,0 @@
-tensor_parallel_size: 1
-max_model_len: 131072
-max_num_seqs: 32
-reasoning_parser: ernie_x1
-tool_call_parser: ernie_x1
-load_choices: "default_v1"

build.sh

@@ -18,9 +18,6 @@ BUILD_WHEEL=${1:-1}
 PYTHON_VERSION=${2:-"python"}
 export python=$PYTHON_VERSION
 FD_CPU_USE_BF16=${3:-"false"}
-# FD_BUILDING_ARCS: Specify target CUDA architectures for custom ops, e.g., "[80, 90, 100]".
-# For SM90 (Hopper), use 90. For SM100 (Blackwell), use 100.
-# These will be translated to 90a / 100a in setup_ops.py for specific features.
 FD_BUILDING_ARCS=${4:-""}
 
 
@@ -34,6 +31,7 @@ EGG_DIR="fastdeploy.egg-info"
 
 # custom_ops directory config
 OPS_SRC_DIR="custom_ops"
+OPS_TMP_DIR_BASE="tmp_base"
 OPS_TMP_DIR="tmp"
 
 # command line log config
@@ -70,6 +68,7 @@ function copy_ops(){
     PY_VERSION="py${PY_MAIN_VERSION}.${PY_SUB_VERSION}"
     SYSTEM_VERSION=`${python} -c "import platform; print(platform.system().lower())"`
     PROCESSOR_VERSION=`${python} -c "import platform; print(platform.processor())"`
+    WHEEL_BASE_NAME="fastdeploy_base_ops-${OPS_VERSION}-${PY_VERSION}-${SYSTEM_VERSION}-${PROCESSOR_VERSION}.egg"
     WHEEL_NAME="fastdeploy_ops-${OPS_VERSION}-${PY_VERSION}-${SYSTEM_VERSION}-${PROCESSOR_VERSION}.egg"
     WHEEL_CPU_NAME="fastdeploy_cpu_ops-${OPS_VERSION}-${PY_VERSION}-${SYSTEM_VERSION}-${PROCESSOR_VERSION}.egg"
     is_rocm=`$python -c "import paddle; print(paddle.is_compiled_with_rocm())"`
@@ -79,11 +78,13 @@ function copy_ops(){
       echo -e "ROCM ops have been copy to fastdeploy"
       return
     fi
+    mkdir -p ../fastdeploy/model_executor/ops/base
     is_cuda=`$python -c "import paddle; print(paddle.is_compiled_with_cuda())"`
     if [ "$is_cuda" = "True" ]; then
       DEVICE_TYPE="gpu"
+      cp -r ./${OPS_TMP_DIR_BASE}/${WHEEL_BASE_NAME}/* ../fastdeploy/model_executor/ops/base
       cp -r ./${OPS_TMP_DIR}/${WHEEL_NAME}/* ../fastdeploy/model_executor/ops/gpu
-      echo -e "CUDA ops have been copy to fastdeploy"
+      echo -e "BASE and CUDA ops have been copy to fastdeploy"
       return
     fi
 
@@ -103,49 +104,27 @@ function copy_ops(){
       return
     fi
 
-    if_corex=`$python -c "import paddle; print(paddle.is_compiled_with_custom_device(\"iluvatar_gpu\"))"`
-    if [ "$if_corex" = "True" ]; then
-      DEVICE_TYPE="iluvatar-gpu"
-      cp -r ./${OPS_TMP_DIR}/${WHEEL_NAME}/* ../fastdeploy/model_executor/ops/iluvatar
-      echo -e "Iluvatar ops have been copy to fastdeploy"
-      return
-    fi
-
-    is_gcu=`$python -c "import paddle; print(paddle.is_compiled_with_custom_device('gcu'))"`
-    if [ "$is_gcu" = "True" ]; then
-      DEVICE_TYPE="gcu"
-      cp -r ${OPS_TMP_DIR}/${WHEEL_NAME}/* ../fastdeploy/model_executor/ops/gcu
-      echo -e "gcu ops have been copy to fastdeploy"
-      return
-    fi
-
-    is_maca=`$python -c "import paddle; print(paddle.device.is_compiled_with_custom_device('metax_gpu'))"`
-    if [ "$is_maca" = "True" ]; then
-      DEVICE_TYPE="metax_gpu"
-      mkdir -p ../fastdeploy/model_executor/ops/base
-      cp -r ./${OPS_TMP_DIR_BASE}/${WHEEL_BASE_NAME}/* ../fastdeploy/model_executor/ops/base
-      cp -r ./${OPS_TMP_DIR}/${WHEEL_NAME}/* ../fastdeploy/model_executor/ops/gpu
-      echo -e "MACA ops have been copy to fastdeploy"
-      return
-    fi
-
     DEVICE_TYPE="cpu"
+    cp -r ./${OPS_TMP_DIR_BASE}/${WHEEL_BASE_NAME}/* ../fastdeploy/model_executor/ops/base
     cd ../../../../
     cp -r ${OPS_TMP_DIR}/${WHEEL_CPU_NAME}/* ../fastdeploy/model_executor/ops/cpu
-    echo -e "CPU ops have been copy to fastdeploy"
+    echo -e "BASE and CPU ops have been copy to fastdeploy"
     return
 }
 
 function build_and_install_ops() {
   cd $OPS_SRC_DIR
   export no_proxy=bcebos.com,paddlepaddle.org.cn,${no_proxy}
+  echo -e "${BLUE}[build]${NONE} build and install fastdeploy_base_ops..."
+  ${python} setup_ops_base.py install --install-lib ${OPS_TMP_DIR_BASE}
+  find ${OPS_TMP_DIR_BASE} -type f -name "*.o" -exec rm -f {} \;
   echo -e "${BLUE}[build]${NONE} build and install fastdeploy_ops..."
   TMP_DIR_REAL_PATH=`readlink -f ${OPS_TMP_DIR}`
   is_xpu=`$python -c "import paddle; print(paddle.is_compiled_with_xpu())"`
   if [ "$is_xpu" = "True" ]; then
-    cd xpu_ops
+    cd xpu_ops/src
     bash build.sh ${TMP_DIR_REAL_PATH}
-    cd ..
+    cd ../..
   elif [ "$FD_CPU_USE_BF16" == "true" ]; then
     if [ "$FD_BUILDING_ARCS" == "" ]; then
       FD_CPU_USE_BF16=True ${python} setup_ops.py install --install-lib ${OPS_TMP_DIR}
@@ -184,6 +163,17 @@ function build_and_install() {
     exit 1
   fi
   echo -e "${BLUE}[build]${NONE} ${GREEN}build fastdeploy wheel success${NONE}\n"
+
+  echo -e "${BLUE}[install]${NONE} installing fastdeploy..."
+  cd $DIST_DIR
+  find . -name "fastdeploy*.whl" | xargs ${python} -m pip install --force-reinstall --no-cache-dir
+  if [ $? -ne 0 ]; then
+    cd ..
+    echo -e "${RED}[FAIL]${NONE} install fastdeploy wheel failed"
+    exit 1
+  fi
+  echo -e "${BLUE}[install]${NONE} ${GREEN}fastdeploy install success${NONE}\n"
+  cd ..
 }
 
 function version_info() {
@@ -191,10 +181,7 @@ function version_info() {
   fastdeploy_git_commit_id=$(git rev-parse HEAD)
   paddle_version=$(${python} -c "import paddle; print(paddle.__version__)")
   paddle_git_commit_id=$(${python} -c "import paddle; print(paddle.__git_commit__)")
-  cuda_version="nvcc-not-installed"
-  if command -v nvcc &> /dev/null; then
-    cuda_version=$(nvcc -V | grep -Po "(?<=release )[\d.]+(?=, V)")
-  fi
+  cuda_version=$(nvcc -V | grep -Po "(?<=release )[\d.]+(?=, V)")
   cxx_version=$(g++ --version | head -n 1 | grep -Po "(?<=\) )[\d.]+")
 
   echo "fastdeploy GIT COMMIT ID: $fastdeploy_git_commit_id" > $output_file
@@ -212,6 +199,7 @@ function cleanup() {
   fi
 
   rm -rf $OPS_SRC_DIR/$BUILD_DIR $OPS_SRC_DIR/$EGG_DIR
+  rm -rf $OPS_SRC_DIR/$OPS_TMP_DIR_BASE
   rm -rf $OPS_SRC_DIR/$OPS_TMP_DIR
 }
 

custom_ops/cpu_ops/get_padding_offset.cc

@@ -84,6 +84,7 @@ std::vector<paddle::Tensor> GetPaddingOffset(const paddle::Tensor &input_ids,
                    seq_length,
                    bsz);
     return {x_remove_padding,
+            cum_offsets_out,
             padding_offset,
             cu_seqlens_q,
             cu_seqlens_k};
@@ -96,7 +97,7 @@ std::vector<std::vector<int64_t>> GetPaddingOffsetInferShape(
     const std::vector<int64_t> &seq_len_shape) {
     int64_t bsz = seq_len_shape[0];
     int64_t seq_len = input_ids_shape[1];
-    return {{-1}, {-1}, {bsz + 1}, {bsz + 1}};
+    return {{-1}, {bsz}, {-1}, {bsz + 1}, {bsz + 1}};
 }
 
 std::vector<paddle::DataType> GetPaddingOffsetInferDtype(
@@ -105,6 +106,7 @@ std::vector<paddle::DataType> GetPaddingOffsetInferDtype(
     const paddle::DataType &token_num_dtype,
     const paddle::DataType &seq_len_dtype) {
     return {input_ids_dtype,
+            seq_len_dtype,
             seq_len_dtype,
             seq_len_dtype,
             seq_len_dtype};
@@ -113,6 +115,7 @@ std::vector<paddle::DataType> GetPaddingOffsetInferDtype(
 PD_BUILD_STATIC_OP(get_padding_offset_cpu)
     .Inputs({"input_ids", "cum_offsets", "token_num", "seq_len"})
     .Outputs({"x_remove_padding",
+              "cum_offsets_out",
               "padding_offset",
               "cu_seqlens_q",
               "cu_seqlens_k"})

custom_ops/cpu_ops/rebuild_padding.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -19,11 +19,10 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-
 template <typename T>
 void RebuildPaddingCPUImpl(T *output_data,
                            const T *input_data,
-                           const int *cu_seqlens_q_data,
+                           const int *cum_offsets_data,
                            const int *seq_len_this_time_data,
                            const int *seq_lens_decoder_data,
                            const int *seq_lens_encoder_data,
@@ -41,12 +40,11 @@ void RebuildPaddingCPUImpl(T *output_data,
         if (seq_lens_decoder_data[bi] == 0 && seq_lens_encoder_data[bi] == 0) {
             continue;
         }
-
         if (seq_lens_encoder_data[bi] > 0) {
             seq_id = seq_lens_encoder_data[bi] - 1;
         }
-
-        const int ori_token_idx = cu_seqlens_q_data[bi] + seq_id;
+        const int ori_token_idx =
+            bi * max_input_length - cum_offsets_data[bi] + seq_id;
         const int src_offset = ori_token_idx * dim_embed + bias_idx;
 
         output_data[i] = input_data[src_offset];
@@ -56,7 +54,7 @@ void RebuildPaddingCPUImpl(T *output_data,
 template <typename T>
 void RebuildAppendPaddingCPUImpl(T *output_data,
                                  const T *input_data,
-                                 const int *cu_seqlens_q_data,
+                                 const int *cum_offsets_data,
                                  const int *seq_len_this_time_data,
                                  const int *seq_lens_decoder_data,
                                  const int *seq_lens_encoder_data,
@@ -71,32 +69,30 @@ void RebuildAppendPaddingCPUImpl(T *output_data,
         int bi = ori_token_id / max_input_length;
         if (seq_len_this_time_data[bi] == 0 ||
             (seq_lens_decoder_data[bi] == 0 &&
-            seq_lens_encoder_data[bi] == 0)) {
-                continue;
-            }
+             seq_lens_encoder_data[bi] == 0)) {
+            continue;
+        }
         int seq_id = 0;
-
         if (seq_lens_encoder_data[bi] > 0) {
             seq_id = seq_lens_encoder_data[bi] - 1;
         }
-        int input_token_id = cu_seqlens_q_data[bi] + seq_id;
+        int input_token_id = ori_token_id - cum_offsets_data[bi] + seq_id;
         int bias_idx = i % dim_embed;
         int src_offset = input_token_id * dim_embed + bias_idx;
-
         output_data[i] = input_data[src_offset];
     }
 }
 
 std::vector<paddle::Tensor> RebuildPaddingCPU(
     const paddle::Tensor &tmp_out,
-    const paddle::Tensor &cu_seqlens_q,
+    const paddle::Tensor &cum_offsets,
     const paddle::Tensor &seq_len_this_time,
     const paddle::Tensor &seq_lens_decoder,
     const paddle::Tensor &seq_lens_encoder,
     const paddle::optional<paddle::Tensor> &output_padding_offset,
     int max_input_length) {
     auto tmp_out_cpu = tmp_out.copy_to(paddle::CPUPlace(), true);
-    auto cu_seqlens_q_cpu = cu_seqlens_q.copy_to(paddle::CPUPlace(), true);
+    auto cum_offsets_cpu = cum_offsets.copy_to(paddle::CPUPlace(), true);
     auto seq_len_this_time_cpu =
         seq_len_this_time.copy_to(paddle::CPUPlace(), true);
     auto seq_lens_decoder_cpu =
@@ -111,7 +107,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
 
     int token_num = tmp_out_cpu.shape()[0];
     int dim_embed = tmp_out_cpu.shape()[1];
-    int bsz = cu_seqlens_q_cpu.shape()[0] - 1;
+    int bsz = cum_offsets_cpu.shape()[0];
 
     paddle::Tensor out;
     if (output_padding_offset_cpu) {
@@ -132,7 +128,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
             {bsz, dim_embed}, 0, tmp_out_cpu.dtype(), paddle::CPUPlace());
     }
 
-    const int *cu_seqlens_q_data = cu_seqlens_q_cpu.data<int>();
+    const int *cum_offsets_data = cum_offsets_cpu.data<int>();
     const int *seq_len_this_time_data = seq_len_this_time_cpu.data<int>();
     const int *seq_lens_decoder_data = seq_lens_decoder_cpu.data<int>();
     const int *seq_lens_encoder_data = seq_lens_encoder_cpu.data<int>();
@@ -145,7 +141,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
             case paddle::DataType::FLOAT32:
                 RebuildAppendPaddingCPUImpl<float>(out.data<float>(),
                                                    tmp_out_cpu.data<float>(),
-                                                   cu_seqlens_q_data,
+                                                   cum_offsets_data,
                                                    seq_len_this_time_data,
                                                    seq_lens_decoder_data,
                                                    seq_lens_encoder_data,
@@ -158,7 +154,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
                 RebuildAppendPaddingCPUImpl<paddle::float16>(
                     out.data<paddle::float16>(),
                     tmp_out_cpu.data<paddle::float16>(),
-                    cu_seqlens_q_data,
+                    cum_offsets_data,
                     seq_len_this_time_data,
                     seq_lens_decoder_data,
                     seq_lens_encoder_data,
@@ -171,7 +167,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
                 RebuildAppendPaddingCPUImpl<paddle::bfloat16>(
                     out.data<paddle::bfloat16>(),
                     tmp_out_cpu.data<paddle::bfloat16>(),
-                    cu_seqlens_q_data,
+                    cum_offsets_data,
                     seq_len_this_time_data,
                     seq_lens_decoder_data,
                     seq_lens_encoder_data,
@@ -190,7 +186,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
             case paddle::DataType::FLOAT32:
                 RebuildPaddingCPUImpl<float>(out.data<float>(),
                                              tmp_out_cpu.data<float>(),
-                                             cu_seqlens_q_data,
+                                             cum_offsets_data,
                                              seq_len_this_time_data,
                                              seq_lens_decoder_data,
                                              seq_lens_encoder_data,
@@ -202,7 +198,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
                 RebuildPaddingCPUImpl<paddle::float16>(
                     out.data<paddle::float16>(),
                     tmp_out_cpu.data<paddle::float16>(),
-                    cu_seqlens_q_data,
+                    cum_offsets_data,
                     seq_len_this_time_data,
                     seq_lens_decoder_data,
                     seq_lens_encoder_data,
@@ -211,10 +207,11 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
                     elem_nums);
                 break;
             case paddle::DataType::BFLOAT16:
+
                 RebuildPaddingCPUImpl<paddle::bfloat16>(
                     out.data<paddle::bfloat16>(),
                     tmp_out_cpu.data<paddle::bfloat16>(),
-                    cu_seqlens_q_data,
+                    cum_offsets_data,
                     seq_len_this_time_data,
                     seq_lens_decoder_data,
                     seq_lens_encoder_data,
@@ -233,7 +230,7 @@ std::vector<paddle::Tensor> RebuildPaddingCPU(
 
 std::vector<std::vector<int64_t>> RebuildPaddingInferShape(
     const std::vector<int64_t> &tmp_out_shape,
-    const std::vector<int64_t> &cu_seqlens_q_shape,
+    const std::vector<int64_t> &cum_offsets_shape,
     const std::vector<int64_t> &seq_len_this_time_shape,
     const std::vector<int64_t> &seq_lens_decoder_shape,
     const std::vector<int64_t> &seq_lens_encoder_shape,
@@ -242,14 +239,14 @@ std::vector<std::vector<int64_t>> RebuildPaddingInferShape(
     if (output_padding_offset_shape) {
         return {{-1, dim_embed}};
     } else {
-        int64_t bsz = cu_seqlens_q_shape[0] - 1;
+        int64_t bsz = cum_offsets_shape[0];
         return {{bsz, dim_embed}};
     }
 }
 
 std::vector<paddle::DataType> RebuildPaddingInferDtype(
     const paddle::DataType &tmp_out_dtype,
-    const paddle::DataType &cu_seqlens_q_dtype,
+    const paddle::DataType &cum_offsets_dtype,
     const paddle::DataType &seq_len_this_time_dtype,
     const paddle::DataType &seq_lens_decoder_dtype,
     const paddle::DataType &seq_lens_encoder_dtype,
@@ -259,7 +256,7 @@ std::vector<paddle::DataType> RebuildPaddingInferDtype(
 
 PD_BUILD_STATIC_OP(rebuild_padding_cpu)
     .Inputs({"tmp_out",
-             "cu_seqlens_q",
+             "cum_offsets",
              "seq_len_this_time",
              "seq_lens_decoder",
              "seq_lens_encoder",

custom_ops/cpu_ops/set_value_by_flags.cc

@@ -14,7 +14,7 @@
 
 #include "paddle/extension.h"
 
-void set_value_by_flags_and_idx(const bool *stop_flags,
+void set_value_by_flag_and_id(const bool *stop_flags,
                               int64_t *pre_ids_all,
                               const int64_t *input_ids,
                               const int *seq_lens_encoder,
@@ -50,7 +50,7 @@ void SetValueByFlagsAndIdx(const paddle::Tensor &pre_ids_all,
     int length = pre_ids_all_shape[1];
     int length_input_ids = input_ids.shape()[1];
 
-    set_value_by_flags_and_idx(stop_flags.data<bool>(),
+    set_value_by_flag_and_id(stop_flags.data<bool>(),
                              const_cast<int64_t *>(pre_ids_all.data<int64_t>()),
                              input_ids.data<int64_t>(),
                              seq_lens_encoder.data<int>(),

custom_ops/cpu_ops/update_inputs.cc

@@ -46,7 +46,7 @@ void update_inputs_kernel(bool *not_need_stop,
     not_need_stop[0] = stop_sum < stop_nums[0];
 }
 
-void UpdateInputs(const paddle::Tensor &stop_flags,
+void UpdateInputes(const paddle::Tensor &stop_flags,
                    const paddle::Tensor &not_need_stop,
                    const paddle::Tensor &seq_lens_this_time,
                    const paddle::Tensor &seq_lens_encoder,
@@ -90,4 +90,4 @@ PD_BUILD_STATIC_OP(update_inputs_cpu)
                     {"seq_lens_encoder", "seq_lens_encoder_out"},
                     {"seq_lens_decoder", "seq_lens_decoder_out"},
                     {"input_ids", "input_ids_out"}})
-    .SetKernelFn(PD_KERNEL(UpdateInputs));
+    .SetKernelFn(PD_KERNEL(UpdateInputes));

custom_ops/gpu_ops/append_attention.cu

@@ -38,7 +38,7 @@ class type2value<phi::dtype::float16> {
 
 
 template <paddle::DataType D>
-void AppendAttentionKernel(
+std::vector<paddle::Tensor> AppendAttentionKernel(
     const AppendAttnMetaData& meta_data,
     const paddle::Tensor& qkv,
     const paddle::Tensor& key_cache,
@@ -46,8 +46,8 @@ void AppendAttentionKernel(
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder,
     const paddle::Tensor& seq_lens_this_time,
-    const paddle::Tensor& batch_id_per_token,
-    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,
@@ -60,7 +60,6 @@ void AppendAttentionKernel(
     const paddle::Tensor& decoder_num_blocks,
     const paddle::Tensor& set_max_lengths,
     const paddle::Tensor& max_len_kv,
-    paddle::Tensor& fmha_out,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& attn_mask,
     const paddle::optional<paddle::Tensor>& qkv_bias,
@@ -73,11 +72,7 @@ void 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>& mask_offset,
     const paddle::optional<paddle::Tensor>& kv_signal_data,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
     const bool rope_3d,
@@ -123,6 +118,27 @@ void AppendAttentionKernel(
   } else {
     qkv_out = qkv;
   }
+  paddle::Tensor fmha_out;
+  if (out_linear_in_scale > 0.0) {
+    if (fabs(quant_max_bound - 127.0f) < 0.000001) {
+      fmha_out = GetEmptyTensor(
+        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims},
+        paddle::DataType::INT8,
+        qkv.place());
+    } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
+      fmha_out = GetEmptyTensor(
+        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims},
+        paddle::DataType::FLOAT8_E4M3FN,
+        qkv.place());
+    }else{
+      PD_THROW("Only supported attr of quant_max_bound in ['127', '448'].");
+    }
+  } else {
+    fmha_out = GetEmptyTensor(
+        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims},
+        D,
+        qkv.place());
+  }
 
   auto dispatch_CascadeAppendAttentionKernel = [&](auto temp_args,
         const paddle::Tensor& lambda_batch_ids,
@@ -140,8 +156,8 @@ void AppendAttentionKernel(
           key_cache,
           value_cache,
           attn_mask,
-          cache_quant_type_str == "block_wise_fp8" ? cache_k_quant_scales : cache_k_dequant_scales,
-          cache_quant_type_str == "block_wise_fp8" ? cache_v_quant_scales : cache_v_dequant_scales,
+          cache_k_dequant_scales,
+          cache_v_dequant_scales,
           cache_k_zp,
           cache_v_zp,
           out_linear_shifts,
@@ -149,8 +165,8 @@ void AppendAttentionKernel(
           seq_lens_this_time,
           seq_lens_decoder,
           seq_lens_encoder,
-          batch_id_per_token,
-          cu_seqlens_q,
+          padding_offsets,
+          cum_offsets,
           block_tables,
           lambda_batch_ids,
           lambda_tile_ids_per_batch,
@@ -186,8 +202,8 @@ void AppendAttentionKernel(
           seq_lens_this_time,
           seq_lens_encoder,
           seq_lens_decoder,
-          batch_id_per_token,
-          cu_seqlens_q,
+          padding_offsets,
+          cum_offsets,
           block_tables,
           kv_batch_ids,
           kv_tile_ids_per_batch,
@@ -207,10 +223,7 @@ void AppendAttentionKernel(
           main_stream,
           &qkv_out,
           const_cast<paddle::Tensor*>(&key_cache),
-          const_cast<paddle::Tensor*>(&value_cache),
-          q_norm_weight,
-          k_norm_weight,
-          rms_norm_eps);
+          const_cast<paddle::Tensor*>(&value_cache));
     };
 
     if (qkv_out_scales) {
@@ -261,8 +274,8 @@ void AppendAttentionKernel(
             qkv,  // [token_num, num_heads, head_dim]
             seq_lens_decoder,
             seq_lens_encoder,
-            batch_id_per_token,
-            cu_seqlens_q,
+            padding_offsets,
+            cum_offsets,
             block_tables,
             rotary_embs,
             qkv_out_scales,
@@ -273,23 +286,19 @@ void AppendAttentionKernel(
             cache_v_zp,
             cache_quant_type_str,
             use_neox_rotary_style,
-            rope_3d,
             max_input_length,
             exec_stream,
             &qkv_out,
             const_cast<paddle::Tensor*>(&key_cache),
-            const_cast<paddle::Tensor*>(&value_cache),
-            q_norm_weight,
-            k_norm_weight,
-            rms_norm_eps);
+            const_cast<paddle::Tensor*>(&value_cache));
       } else {
         SpeculateWriteCacheWithRoPEKernel<data_t, data_t>(
             meta_data,
             qkv_out,  // [token_num, num_heads, head_dim]
             seq_lens_decoder,
             seq_lens_encoder,
-            batch_id_per_token,
-            cu_seqlens_q,
+            padding_offsets,
+            cum_offsets,
             block_tables,
             rotary_embs,
             qkv_out_scales,
@@ -300,15 +309,11 @@ void AppendAttentionKernel(
             cache_v_zp,
             cache_quant_type_str,
             use_neox_rotary_style,
-            rope_3d,
             max_input_length,
             exec_stream,
             &qkv_out,
             const_cast<paddle::Tensor*>(&key_cache),
-            const_cast<paddle::Tensor*>(&value_cache),
-            q_norm_weight,
-            k_norm_weight,
-            rms_norm_eps);
+            const_cast<paddle::Tensor*>(&value_cache));
       }
     } else {
       if (qkv_out_scales) {
@@ -317,7 +322,8 @@ void AppendAttentionKernel(
             qkv,  // [token_num, num_heads, head_dim]
             seq_lens_decoder,
             seq_lens_encoder,
-            cu_seqlens_q,
+            padding_offsets,
+            cum_offsets,
             block_tables,
             rotary_embs,
             qkv_out_scales,
@@ -333,17 +339,15 @@ void AppendAttentionKernel(
             exec_stream,
             &qkv_out,
             const_cast<paddle::Tensor*>(&key_cache),
-            const_cast<paddle::Tensor*>(&value_cache),
-            q_norm_weight,
-            k_norm_weight,
-            rms_norm_eps);
+            const_cast<paddle::Tensor*>(&value_cache));
       } else {
         DecoderWriteCacheWithRoPEKernel<data_t, data_t>(
             meta_data,
             qkv_out,  // [token_num, num_heads, head_dim]
             seq_lens_decoder,
             seq_lens_encoder,
-            cu_seqlens_q,
+            padding_offsets,
+            cum_offsets,
             block_tables,
             rotary_embs,
             qkv_out_scales,
@@ -359,10 +363,7 @@ void AppendAttentionKernel(
             exec_stream,
             &qkv_out,
             const_cast<paddle::Tensor*>(&key_cache),
-            const_cast<paddle::Tensor*>(&value_cache),
-            q_norm_weight,
-            k_norm_weight,
-            rms_norm_eps);
+            const_cast<paddle::Tensor*>(&value_cache));
       }
     }
 
@@ -391,6 +392,8 @@ void AppendAttentionKernel(
       cudaStreamWaitEvent(main_stream, decoder_event);
     }
   }
+
+  return {fmha_out, qkv_out};
 }
 
 std::vector<paddle::Tensor> AppendAttention(
@@ -400,8 +403,8 @@ std::vector<paddle::Tensor> AppendAttention(
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder,
     const paddle::Tensor& seq_lens_this_time,
-    const paddle::Tensor& batch_id_per_token,
-    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,
@@ -426,11 +429,7 @@ std::vector<paddle::Tensor> AppendAttention(
     const paddle::optional<paddle::Tensor>& cache_v_zp,
     const paddle::optional<paddle::Tensor>& out_linear_shifts,
     const paddle::optional<paddle::Tensor>& out_linear_smooths,
-    const paddle::optional<paddle::Tensor>& mask_offset,
     const paddle::optional<paddle::Tensor>& kv_signal_data,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps,
     const std::string& compute_dtype,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
@@ -463,62 +462,10 @@ std::vector<paddle::Tensor> AppendAttention(
 
   meta_data.max_blocks_per_seq = block_tables.dims()[1];
   meta_data.block_size = key_cache.dims()[2];
-  meta_data.batch_size = seq_lens_this_time.dims()[0];
+  meta_data.batch_size = cum_offsets.dims()[0];
 
-  // template dtype generation
-  phi::DataType dtype_id;
-  switch (qkv.dtype()) {
-    case paddle::DataType::FLOAT16:  {dtype_id = phi::DataType::FLOAT16;  break;}
-    case paddle::DataType::BFLOAT16: {dtype_id = phi::DataType::BFLOAT16; break;}
-    case paddle::DataType::INT32: {
-      if (compute_dtype == "bf16") {
-        dtype_id = phi::DataType::BFLOAT16;
-        break;
-      } else if (compute_dtype == "fp16") {
-        dtype_id = phi::DataType::FLOAT16;
-        break;
-      } else {
-        PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16'].");
-        break;
-      }
-    }
-    default: {
-      PD_THROW(
-          "NOT supported data type. "
-          "Only float16 and bfloat16 are supported. ");
-      break;
-    }
-  }
-
-  // fmha_out generation, rewrite from AppendAttentionKernel
-  paddle::Tensor fmha_out;
-  if (out_linear_in_scale > 0.0) {
-    if (fabs(quant_max_bound - 127.0f) < 0.000001) {
-      fmha_out = GetEmptyTensor(
-        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims},
-        paddle::DataType::INT8,
-        qkv.place());
-    } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
-      fmha_out = GetEmptyTensor(
-        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims},
-        paddle::DataType::FLOAT8_E4M3FN,
-        qkv.place());
-    } else{
-      PD_THROW("Only supported attr of quant_max_bound in ['127', '448'].");
-    }
-  } else {
-    fmha_out = GetEmptyTensor(
-        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims},
-        dtype_id,
-        qkv.place());
-  }
-
-  if (mask_offset) {
-    meta_data.mask_offset = mask_offset.get().data<int>();
-  }
-
-  auto dispatch_by_template = [&](auto temp_args) -> void {
-      AppendAttentionKernel<type2value<decltype(temp_args)>::value>(
+  auto dispatch_by_template = [&](auto temp_args) -> std::vector<paddle::Tensor> {
+      return AppendAttentionKernel<type2value<decltype(temp_args)>::value>(
           meta_data,
           qkv,
           key_cache,
@@ -526,8 +473,8 @@ std::vector<paddle::Tensor> AppendAttention(
           seq_lens_encoder,
           seq_lens_decoder,
           seq_lens_this_time,
-          batch_id_per_token,
-          cu_seqlens_q,
+          padding_offsets,
+          cum_offsets,
           block_tables,
           encoder_batch_ids,
           encoder_tile_ids_per_batch,
@@ -540,7 +487,6 @@ std::vector<paddle::Tensor> AppendAttention(
           decoder_num_blocks,
           set_max_lengths,
           max_len_kv,
-          fmha_out,
           rotary_embs,
           attn_mask,
           qkv_bias,
@@ -553,11 +499,7 @@ std::vector<paddle::Tensor> AppendAttention(
           cache_v_zp,
           out_linear_shifts,
           out_linear_smooths,
-          mask_offset,
           kv_signal_data,
-          q_norm_weight,
-          k_norm_weight,
-          rms_norm_eps,
           cache_quant_type_str,
           use_neox_rotary_style,
           rope_3d,
@@ -572,198 +514,35 @@ std::vector<paddle::Tensor> AppendAttention(
           speculate_max_draft_token_num,
           causal,
           speculate_decoder);
-  };
 
+  };
 
   phi::dtype::float16 fp16_dtype;
   phi::dtype::bfloat16 bp16_dtype;
-  switch (dtype_id){
-      case phi::DataType::FLOAT16: {
-        dispatch_by_template(fp16_dtype);
-        return {fmha_out};
-        }
-      case phi::DataType::BFLOAT16: {
-        dispatch_by_template(bp16_dtype);
-        return {fmha_out};
-        }
-      default:
-        PD_THROW(
+
+  switch (qkv.dtype()) {
+    case paddle::DataType::FLOAT16: return dispatch_by_template(fp16_dtype);
+    case paddle::DataType::BFLOAT16: return dispatch_by_template(bp16_dtype);
+    case paddle::DataType::INT32: {
+      if (compute_dtype == "bf16") {
+        return dispatch_by_template(bp16_dtype);
+      } else if (compute_dtype == "fp16") {
+        return dispatch_by_template(fp16_dtype);
+      } else {
+        PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16'].");
+        break;
+      }
+    }
+    default: {
+      PD_THROW(
           "NOT supported data type. "
           "Only float16 and bfloat16 are supported. ");
-        break;
+      break;
+    }
   }
-
   return {paddle::Tensor{}};
 }
 
-void AppendAttentionWithOutput(
-    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& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
-    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& set_max_lengths,
-    const paddle::Tensor& max_len_kv,
-    paddle::Tensor& fmha_out,
-    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>& mask_offset,
-    const paddle::optional<paddle::Tensor>& kv_signal_data,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps,
-    const std::string& compute_dtype,
-    const std::string& cache_quant_type_str,
-    const bool use_neox_rotary_style,
-    const bool rope_3d,
-    const int max_input_length,
-    const float quant_max_bound,
-    const float quant_min_bound,
-    const float out_linear_in_scale,
-    const int encoder_block_shape_q,
-    const int decoder_block_shape_q,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
-    const int speculate_max_draft_token_num,
-    const bool causal,
-    const bool speculate_decoder) {
-  AppendAttnMetaData meta_data;
-
-  const auto& qkv_dims = qkv.dims();
-  const auto& key_cache_dims = key_cache.dims();
-  meta_data.token_nums = qkv_dims[0];
-  meta_data.kv_num_heads = key_cache_dims[1];
-  meta_data.head_dims = key_cache_dims[3];
-  // TODO: trick method support c4, add attr head_dims in the future
-  if (cache_quant_type_str == "cache_int4_zp") {
-    meta_data.head_dims *= 2;
-  }
-  const int total_num_head =
-      qkv_dims[qkv_dims.size() - 1] / meta_data.head_dims;
-  meta_data.q_num_heads = total_num_head - 2 * meta_data.kv_num_heads;
-
-  meta_data.max_blocks_per_seq = block_tables.dims()[1];
-  meta_data.block_size = key_cache.dims()[2];
-  meta_data.batch_size = seq_lens_this_time.dims()[0];
-
-  if (mask_offset) {
-    meta_data.mask_offset = mask_offset.get().data<int>();
-  }
-
-  auto dispatch_by_template = [&](auto temp_args) -> void {
-      AppendAttentionKernel<type2value<decltype(temp_args)>::value>(
-          meta_data,
-          qkv,
-          key_cache,
-          value_cache,
-          seq_lens_encoder,
-          seq_lens_decoder,
-          seq_lens_this_time,
-          batch_id_per_token,
-          cu_seqlens_q,
-          block_tables,
-          encoder_batch_ids,
-          encoder_tile_ids_per_batch,
-          encoder_num_blocks,
-          kv_batch_ids,
-          kv_tile_ids_per_batch,
-          kv_num_blocks,
-          decoder_batch_ids,
-          decoder_tile_ids_per_batch,
-          decoder_num_blocks,
-          set_max_lengths,
-          max_len_kv,
-          fmha_out,
-          rotary_embs,
-          attn_mask,
-          qkv_bias,
-          qkv_out_scales,
-          cache_k_quant_scales,
-          cache_v_quant_scales,
-          cache_k_dequant_scales,
-          cache_v_dequant_scales,
-          cache_k_zp,
-          cache_v_zp,
-          out_linear_shifts,
-          out_linear_smooths,
-          mask_offset,
-          kv_signal_data,
-          q_norm_weight,
-          k_norm_weight,
-          rms_norm_eps,
-          cache_quant_type_str,
-          use_neox_rotary_style,
-          rope_3d,
-          max_input_length,
-          quant_max_bound,
-          quant_min_bound,
-          out_linear_in_scale,
-          encoder_block_shape_q,
-          decoder_block_shape_q,
-          max_partition_size,
-          encoder_max_partition_size,
-          speculate_max_draft_token_num,
-          causal,
-          speculate_decoder);
-  };
-
-  phi::dtype::float16 fp16_dtype;
-  phi::dtype::bfloat16 bp16_dtype;
-
-  switch (qkv.dtype()) {
-    case paddle::DataType::FLOAT16: {
-      dispatch_by_template(fp16_dtype);
-      break;
-    }
-    case paddle::DataType::BFLOAT16: {
-      dispatch_by_template(bp16_dtype);
-      break;
-    }
-    case paddle::DataType::INT32: {
-      if (compute_dtype == "bf16") {
-        dispatch_by_template(bp16_dtype);
-        break;
-      } else if (compute_dtype == "fp16") {
-        dispatch_by_template(fp16_dtype);
-        break;
-      } else {
-        PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16'].");
-        break;
-      }
-    }
-    default: {
-      PD_THROW(
-          "NOT supported data type. "
-          "Only float16 and bfloat16 are supported. ");
-      break;
-    }
-  }
-}
-
-
 std::vector<std::vector<int64_t>> AppendAttentionInferShape(
     const std::vector<int64_t>& qkv_shape,
     const std::vector<int64_t>& key_cache_shape,
@@ -771,8 +550,8 @@ std::vector<std::vector<int64_t>> AppendAttentionInferShape(
     const std::vector<int64_t>& seq_lens_encoder_shape,
     const std::vector<int64_t>& seq_lens_decoder_shape,
     const std::vector<int64_t>& seq_lens_this_time_shape,
-    const std::vector<int64_t>& batch_id_per_token_shape,
-    const std::vector<int64_t>& cu_seqlens_q_shape,
+    const std::vector<int64_t>& padding_offsets_shape,
+    const std::vector<int64_t>& cum_offsets_shape,
     const std::vector<int64_t>& block_tables_shape,
     const std::vector<int64_t>& encoder_batch_ids_shape,
     const std::vector<int64_t>& encoder_tile_ids_per_batch_shape,
@@ -797,11 +576,7 @@ std::vector<std::vector<int64_t>> AppendAttentionInferShape(
     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>>& mask_offset_shape,
     const paddle::optional<std::vector<int64_t>>& kv_signal_data_shape,
-    const paddle::optional<std::vector<int64_t>>& q_norm_weight_shape,
-    const paddle::optional<std::vector<int64_t>>& k_norm_weight_shape,
-    const float rms_norm_eps,
     const std::string& compute_dtype,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
@@ -825,7 +600,7 @@ std::vector<std::vector<int64_t>> AppendAttentionInferShape(
   }
   const int total_num_head = qkv_shape[qkv_shape.size() - 1] / head_dim;
   const int num_heads = total_num_head - 2 * kv_num_heads;
-  return {{token_num, num_heads * head_dim}};
+  return {{token_num, num_heads * head_dim}, qkv_shape};
 }
 
 std::vector<paddle::DataType> AppendAttentionInferDtype(
@@ -835,8 +610,8 @@ std::vector<paddle::DataType> AppendAttentionInferDtype(
     const paddle::DataType& seq_lens_encoder_dtype,
     const paddle::DataType& seq_lens_decoder_dtype,
     const paddle::DataType& seq_lens_this_time_dtype,
-    const paddle::DataType& batch_id_per_token_dtype,
-    const paddle::DataType& cu_seqlens_q_dtype,
+    const paddle::DataType& padding_offsets_dtype,
+    const paddle::DataType& cum_offsets_dtype,
     const paddle::DataType& block_tables_dtype,
     const paddle::DataType& encoder_batch_ids_dtype,
     const paddle::DataType& encoder_tile_ids_per_batch_dtype,
@@ -861,11 +636,7 @@ std::vector<paddle::DataType> AppendAttentionInferDtype(
     const paddle::optional<paddle::DataType>& cache_v_zp_dtype,
     const paddle::optional<paddle::DataType>& out_linear_shifts_dtype,
     const paddle::optional<paddle::DataType>& out_linear_smooths_dtype,
-    const paddle::optional<paddle::DataType>& mask_offset_dtype,
     const paddle::optional<paddle::DataType>& kv_signal_data_dtype,
-    const paddle::optional<paddle::DataType>& q_norm_weight_dtype,
-    const paddle::optional<paddle::DataType>& k_norm_weight_dtype,
-    const float rms_norm_eps,
     const std::string& compute_dtype,
     const std::string& cache_quant_type_str,
     const bool use_neox_rotary_style,
@@ -884,148 +655,32 @@ std::vector<paddle::DataType> AppendAttentionInferDtype(
   if (compute_dtype == "bf16") {
     if (out_linear_in_scale > 0.0) {
       if (fabs(quant_max_bound - 127.0f) < 0.000001) {
-        return {paddle::DataType::INT8};
+        return {paddle::DataType::INT8, paddle::DataType::BFLOAT16};
       } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
-        return {paddle::DataType::FLOAT8_E4M3FN};
+        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::BFLOAT16};
       }else{
         PD_THROW("Only supported attr of quant_max_bound in ['127.0', '448.0'].");
       }
     } else {
-      return {paddle::DataType::BFLOAT16};
+      return {paddle::DataType::BFLOAT16, paddle::DataType::BFLOAT16};
     }
   } else if (compute_dtype == "fp16") {
     if (out_linear_in_scale > 0.0) {
       if (fabs(quant_max_bound - 127.0f) < 0.000001) {
-        return {paddle::DataType::INT8};
+        return {paddle::DataType::INT8, paddle::DataType::FLOAT16};
       } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
-        return {paddle::DataType::FLOAT8_E4M3FN};
+        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT16};
       }else{
         PD_THROW("Only supported attr of quant_max_bound in ['127.0', '448.0'].");
       }
     } else {
-      return {paddle::DataType::FLOAT16};
+      return {paddle::DataType::FLOAT16, paddle::DataType::FLOAT16};
     }
   } else {
     PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16'].");
   }
 }
 
-std::vector<std::vector<int64_t>> AppendAttentionWithOutputInferShape(
-    const std::vector<int64_t>& qkv_shape,
-    const std::vector<int64_t>& key_cache_shape,
-    const std::vector<int64_t>& value_cache_shape,
-    const std::vector<int64_t>& seq_lens_encoder_shape,
-    const std::vector<int64_t>& seq_lens_decoder_shape,
-    const std::vector<int64_t>& seq_lens_this_time_shape,
-    const std::vector<int64_t>& batch_id_per_token_shape,
-    const std::vector<int64_t>& cu_seqlens_q_shape,
-    const std::vector<int64_t>& block_tables_shape,
-    const std::vector<int64_t>& encoder_batch_ids_shape,
-    const std::vector<int64_t>& encoder_tile_ids_per_batch_shape,
-    const std::vector<int64_t>& encoder_num_blocks_shape,
-    const std::vector<int64_t>& kv_batch_ids_shape,
-    const std::vector<int64_t>& kv_tile_ids_per_batch_shape,
-    const std::vector<int64_t>& kv_num_blocks_shape,
-    const std::vector<int64_t>& decoder_batch_ids_shape,
-    const std::vector<int64_t>& decoder_tile_ids_per_batch_shape,
-    const std::vector<int64_t>& decoder_num_blocks_shape,
-    const std::vector<int64_t>& set_max_lengths_shape,
-    const std::vector<int64_t>& max_len_kv_shape,
-    const std::vector<int64_t>& fmha_out_shape,
-    const paddle::optional<std::vector<int64_t>>& rotary_embs_shape,
-    const paddle::optional<std::vector<int64_t>>& attn_mask_shape,
-    const paddle::optional<std::vector<int64_t>>& qkv_bias_shape,
-    const paddle::optional<std::vector<int64_t>>& qkv_out_scales_shape,
-    const paddle::optional<std::vector<int64_t>>& cache_k_quant_scales_shape,
-    const paddle::optional<std::vector<int64_t>>& cache_v_quant_scales_shape,
-    const paddle::optional<std::vector<int64_t>>& cache_k_dequant_scales_shape,
-    const paddle::optional<std::vector<int64_t>>& cache_v_dequant_scales_shape,
-    const paddle::optional<std::vector<int64_t>>& cache_k_zp_shape,
-    const paddle::optional<std::vector<int64_t>>& cache_v_zp_shape,
-    const paddle::optional<std::vector<int64_t>>& out_linear_shifts_shape,
-    const paddle::optional<std::vector<int64_t>>& out_linear_smooths_shape,
-    const paddle::optional<std::vector<int64_t>>& mask_offset_shape,
-    const paddle::optional<std::vector<int64_t>>& kv_signal_data_shape,
-    const paddle::optional<std::vector<int64_t>>& q_norm_weight_shape,
-    const paddle::optional<std::vector<int64_t>>& k_norm_weight_shape,
-    const float rms_norm_eps,
-    const std::string& compute_dtype,
-    const std::string& cache_quant_type_str,
-    const bool use_neox_rotary_style,
-    const bool rope_3d,
-    const int max_input_length,
-    const float quant_max_bound,
-    const float quant_min_bound,
-    const float out_linear_in_scale,
-    const int encoder_block_shape_q,
-    const int decoder_block_shape_q,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
-    const int speculate_max_draft_token_num,
-    const bool causal,
-    const bool speculate_decoder) {
-  return {fmha_out_shape};
-}
-
-std::vector<paddle::DataType> AppendAttentionWithOutputInferDtype(
-    const paddle::DataType& qkv_dtype,
-    const paddle::DataType& key_cache_dtype,
-    const paddle::DataType& value_cache_dtype,
-    const paddle::DataType& seq_lens_encoder_dtype,
-    const paddle::DataType& seq_lens_decoder_dtype,
-    const paddle::DataType& seq_lens_this_time_dtype,
-    const paddle::DataType& batch_id_per_token_dtype,
-    const paddle::DataType& cu_seqlens_q_dtype,
-    const paddle::DataType& block_tables_dtype,
-    const paddle::DataType& encoder_batch_ids_dtype,
-    const paddle::DataType& encoder_tile_ids_per_batch_dtype,
-    const paddle::DataType& encoder_num_blocks_dtype,
-    const paddle::DataType& kv_batch_ids_dtype,
-    const paddle::DataType& kv_tile_ids_per_batch_dtype,
-    const paddle::DataType& kv_num_blocks_dtype,
-    const paddle::DataType& decoder_batch_ids_dtype,
-    const paddle::DataType& decoder_tile_ids_per_batch_dtype,
-    const paddle::DataType& decoder_num_blocks_dtype,
-    const paddle::DataType& set_max_lengths_dtype,
-    const paddle::DataType& max_len_kv_dtype,
-    const paddle::DataType& fmha_out_dtype,
-    const paddle::optional<paddle::DataType>& rotary_embs_dtype,
-    const paddle::optional<paddle::DataType>& attn_mask_dtype,
-    const paddle::optional<paddle::DataType>& qkv_bias_dtype,
-    const paddle::optional<paddle::DataType>& qkv_out_scales_dtype,
-    const paddle::optional<paddle::DataType>& cache_k_quant_scales_dtype,
-    const paddle::optional<paddle::DataType>& cache_v_quant_scales_dtype,
-    const paddle::optional<paddle::DataType>& cache_k_dequant_scales_dtype,
-    const paddle::optional<paddle::DataType>& cache_v_dequant_scales_dtype,
-    const paddle::optional<paddle::DataType>& cache_k_zp_dtype,
-    const paddle::optional<paddle::DataType>& cache_v_zp_dtype,
-    const paddle::optional<paddle::DataType>& out_linear_shifts_dtype,
-    const paddle::optional<paddle::DataType>& out_linear_smooths_dtype,
-    const paddle::optional<paddle::DataType>& mask_offset_dtype,
-    const paddle::optional<paddle::DataType>& kv_signal_data_dtype,
-    const paddle::optional<paddle::DataType>& q_norm_weight_dtype,
-    const paddle::optional<paddle::DataType>& k_norm_weight_dtype,
-    const float rms_norm_eps,
-    const std::string& compute_dtype,
-    const std::string& cache_quant_type_str,
-    const bool use_neox_rotary_style,
-    const bool rope_3d,
-    const int max_input_length,
-    const float quant_max_bound,
-    const float quant_min_bound,
-    const float out_linear_in_scale,
-    const int encoder_block_shape_q,
-    const int decoder_block_shape_q,
-    const int max_partition_size,
-    const int encoder_max_partition_size,
-    const int speculate_max_draft_token_num,
-    const bool causal,
-    const bool speculate_decoder) {
-  return {fmha_out_dtype};
-}
-
-
-
 PD_BUILD_STATIC_OP(append_attention)
     .Inputs({"qkv",
              "key_cache",
@@ -1033,8 +688,8 @@ PD_BUILD_STATIC_OP(append_attention)
              "seq_lens_encoder",
              "seq_lens_decoder",
              "seq_lens_this_time",
-             "batch_id_per_token",
-             "cu_seqlens_q",
+             "padding_offsets",
+             "cum_offsets",
              "block_tables",
              "encoder_batch_ids",
              "encoder_tile_ids_per_batch",
@@ -1059,15 +714,11 @@ PD_BUILD_STATIC_OP(append_attention)
              paddle::Optional("cache_v_zp"),
              paddle::Optional("out_linear_shifts"),
              paddle::Optional("out_linear_smooths"),
-             paddle::Optional("mask_offset"),
-             paddle::Optional("kv_signal_data"),
-             paddle::Optional("q_norm_weight"),
-             paddle::Optional("k_norm_weight")})
-    .Outputs({"fmha_out", "key_cache_out", "value_cache_out"})
+             paddle::Optional("kv_signal_data")})
+    .Outputs({"fmha_out", "qkv_out", "key_cache_out", "value_cache_out"})
     .SetInplaceMap({{"key_cache", "key_cache_out"},
                     {"value_cache", "value_cache_out"}})
-    .Attrs({"rms_norm_eps: float",
-            "compute_type: std::string",
+    .Attrs({"compute_type: std::string",
             "cache_quant_type: std::string",
             "use_neox_rotary_style: bool",
             "rope_3d: bool",
@@ -1081,71 +732,7 @@ PD_BUILD_STATIC_OP(append_attention)
             "encoder_max_partition_size: int",
             "speculate_max_draft_token_num: int",
             "causal: bool",
-            "speculate_decoder: bool",
-            })
+            "speculate_decoder: bool"})
     .SetKernelFn(PD_KERNEL(AppendAttention))
     .SetInferShapeFn(PD_INFER_SHAPE(AppendAttentionInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(AppendAttentionInferDtype));
-
-PD_BUILD_STATIC_OP(append_attention_with_output)
-    .Inputs({"qkv",
-             "key_cache",
-             "value_cache",
-             "seq_lens_encoder",
-             "seq_lens_decoder",
-             "seq_lens_this_time",
-             "batch_id_per_token",
-             "cu_seqlens_q",
-             "block_tables",
-             "encoder_batch_ids",
-             "encoder_tile_ids_per_batch",
-             "encoder_num_blocks",
-             "kv_batch_ids",
-             "kv_tile_ids_per_batch",
-             "kv_num_blocks",
-             "decoder_batch_ids",
-             "decoder_tile_ids_per_batch",
-             "decoder_num_blocks",
-             "set_max_lengths",
-             "max_len_kv",
-             "fmha_out",
-             paddle::Optional("rotary_embs"),
-             paddle::Optional("attn_mask"),
-             paddle::Optional("qkv_bias"),
-             paddle::Optional("qkv_out_scales"),
-             paddle::Optional("cache_k_quant_scales"),
-             paddle::Optional("cache_v_quant_scales"),
-             paddle::Optional("cache_k_dequant_scales"),
-             paddle::Optional("cache_v_dequant_scales"),
-             paddle::Optional("cache_k_zp"),
-             paddle::Optional("cache_v_zp"),
-             paddle::Optional("out_linear_shifts"),
-             paddle::Optional("out_linear_smooths"),
-             paddle::Optional("mask_offset"),
-             paddle::Optional("kv_signal_data"),
-             paddle::Optional("q_norm_weight"),
-             paddle::Optional("k_norm_weight")})
-    .Outputs({"fmha_out_out", "qkv_out", "key_cache_out", "value_cache_out"})
-    .SetInplaceMap({{"fmha_out", "fmha_out_out"},
-                    {"key_cache", "key_cache_out"},
-                    {"value_cache", "value_cache_out"}})
-    .Attrs({"rms_norm_eps: float",
-            "compute_type: std::string",
-            "cache_quant_type: std::string",
-            "use_neox_rotary_style: bool",
-            "rope_3d: bool",
-            "max_input_length: int",
-            "quant_max_bound: float",
-            "quant_min_bound: float",
-            "out_linear_in_scale: float",
-            "encoder_block_shape_q: int",
-            "decoder_block_shape_q: int",
-            "max_partition_size: int",
-            "encoder_max_partition_size: int",
-            "speculate_max_draft_token_num: int",
-            "causal: bool",
-            "speculate_decoder: bool",
-            })
-    .SetKernelFn(PD_KERNEL(AppendAttentionWithOutput))
-    .SetInferShapeFn(PD_INFER_SHAPE(AppendAttentionWithOutputInferShape))
-    .SetInferDtypeFn(PD_INFER_DTYPE(AppendAttentionWithOutputInferDtype));

custom_ops/gpu_ops/append_attn/append_attention_c16_impl.cuh

@@ -41,9 +41,8 @@ __global__ void multi_query_append_attention_kernel(
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ batch_ids,
     const int *__restrict__ tile_ids_per_batch,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
-    const int *__restrict__ mask_offset,
     const int max_seq_len,
     const int max_dec_len,
     const int max_block_num_per_seq,
@@ -52,7 +51,6 @@ __global__ void multi_query_append_attention_kernel(
     const float quant_min_bound,
     const float in_scale,
     const uint32_t chunk_size,
-    const int num_blocks_x_cpu,
     T *__restrict__ tmp_workspace,  // split kv [token_num, num_chunks,
                                     // num_heads, head_dim]
     float *__restrict__ tmp_m,      // [token_num, num_chunks, num_heads]
@@ -75,11 +73,6 @@ __global__ void multi_query_append_attention_kernel(
 
   block_table_now = block_table + batch_id * max_block_num_per_seq;
 
-  //When cudagraph capture prefill, may launch more gridDim.x
-  if(btid >= static_cast<uint32_t>(num_blocks_x_cpu)){
-    return;
-  }
-
   const uint32_t q_len = seq_lens[batch_id];
   if (q_len <= 0) {
     return;
@@ -121,7 +114,8 @@ __global__ void multi_query_append_attention_kernel(
   const uint32_t kv_n_stride = kv_num_heads * BLOCK_SIZE * HEAD_DIM;
   const uint32_t kv_h_stride = BLOCK_SIZE * HEAD_DIM;
   const uint32_t kv_b_stride = HEAD_DIM;
-  const uint32_t q_start_seq_id = cu_seqlens_q[batch_id];
+  const uint32_t q_start_seq_id =
+      batch_id * max_seq_len - __ldg(&cum_offsets[batch_id]);
   const uint32_t q_base_seq_id_this_block =
       (tile_id * NUM_WARPS + wid) * num_frags_x * 16;
   const uint32_t q_offset = q_start_seq_id * q_ori_n_stride +
@@ -148,7 +142,6 @@ __global__ void multi_query_append_attention_kernel(
   } else {
     o_base_ptr_int8 = out + o_offset;
   }
-  const int *mask_offset_this_seq = mask_offset ? mask_offset + q_start_seq_id * 2 : nullptr;
   smem_t qo_smem(smem);
 
   uint32_t q_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -187,7 +180,7 @@ __global__ void multi_query_append_attention_kernel(
                          kv_len - q_len +
                              tile_id * num_rows_per_block / GROUP_SIZE,
                          chunk_start)))
-              : mask_offset ? 0 : chunk_len) /
+              : chunk_len) /
       (num_frags_z * 16);
   uint32_t k_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
       8 * (tid / 16) + tid % 8, (tid % 16) / 8);
@@ -253,16 +246,12 @@ __global__ void multi_query_append_attention_kernel(
              NUM_WARPS,
              num_frags_x,
              num_frags_y,
-             num_frags_z>(nullptr,
-                          q_base_seq_id_this_block,
+             num_frags_z>(q_base_seq_id_this_block,
                           kv_idx_base,
                           q_len,
                           kv_len,
                           chunk_end,
-                          -1,
-                          s_frag,
-                          mask_offset_this_seq);
-
+                          s_frag);
     }
 
     // update m,d
@@ -416,10 +405,8 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ batch_ids,
     const int *__restrict__ tile_ids_per_batch,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
-    const int *__restrict__ mask_offset,
-    const bool *__restrict__ attn_mask,    // [bsz, max_q, max_q] for tree-mask
     const int max_seq_len,
     const int max_dec_len,
     const int max_block_num_per_seq,
@@ -428,14 +415,12 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
     const float quant_min_bound,
     const float in_scale,
     const uint32_t chunk_size,
-    const int num_blocks_x_cpu,
     T *__restrict__ tmp_workspace,  // split kv [token_num, num_chunks,
                                     // num_heads, head_dim]
     float *__restrict__ tmp_m,      // [token_num, num_chunks, num_heads]
     float *__restrict__ tmp_d,      // [token_num, num_chunks, num_heads]
     OutT *__restrict__ out,
-    const int speculate_max_draft_token_num = 5,
-    const int32_t attn_mask_len = -1) {
+    const int speculate_max_draft_token_num = 5) {
   constexpr uint32_t num_vecs_per_head = HEAD_DIM / num_elems_per_128b<T>();
   static_assert(NUM_WARP_Q == 1, "NUM_WARP_Q must be 1");
   static_assert(NUM_WARP_KV == 4, "NUM_WARP_KV must be 4");
@@ -452,11 +437,6 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
   const uint32_t num_rows_per_block = num_frags_x * 16;
   const int *block_table_now = block_table + batch_id * max_block_num_per_seq;
 
-  //When cudagraph capture prefill, may launch more gridDim.x
-  if(btid >= static_cast<uint32_t>(num_blocks_x_cpu)){
-    return;
-  }
-
   const uint32_t q_len = seq_lens[batch_id];
   if (q_len <= 0) {
     return;
@@ -497,7 +477,8 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
   const uint32_t kv_n_stride = kv_num_heads * BLOCK_SIZE * HEAD_DIM;
   const uint32_t kv_h_stride = BLOCK_SIZE * HEAD_DIM;
   const uint32_t kv_b_stride = HEAD_DIM;
-  const uint32_t q_start_seq_id = cu_seqlens_q[batch_id];
+  const uint32_t q_start_seq_id =
+      batch_id * max_seq_len - __ldg(&cum_offsets[batch_id]);
   const uint32_t q_base_seq_id_this_block = tile_id * num_frags_x * 16;
   const uint32_t q_offset = q_start_seq_id * q_ori_n_stride +
                             q_head_idx * HEAD_DIM +
@@ -523,7 +504,7 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
           tid % 8 * num_elems_per_128b<T>();
     }
   }
-  const int *mask_offset_this_seq = mask_offset ? mask_offset + q_start_seq_id * 2 : nullptr;
+
   smem_t qo_smem(smem);
 
   uint32_t q_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -561,9 +542,10 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
   const uint32_t mask_check_iteration =
       (CAUSAL ? (min(chunk_len,
                      sub_if_greater_or_zero(
-                         kv_len - q_len,
+                         kv_len - q_len +
+                             tile_id * num_rows_per_block / GROUP_SIZE,
                          chunk_start)))
-              : mask_offset ? 0 : chunk_len) /
+              : chunk_len) /
       (NUM_WARP_KV * num_frags_z * 16);
 
   uint32_t k_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -631,15 +613,12 @@ __global__ void multi_query_append_attention_warp1_4_kernel(
              NUM_WARPS,
              num_frags_x,
              num_frags_y,
-             num_frags_z>(attn_mask ? attn_mask + batch_id * attn_mask_len *attn_mask_len : nullptr,
-                          q_base_seq_id_this_block,
+             num_frags_z>(q_base_seq_id_this_block,
                           kv_idx_base + wid * num_frags_z * 16,
                           q_len,
                           kv_len,
                           chunk_end,
-                          attn_mask_len,
-                          s_frag,
-                          mask_offset_this_seq);
+                          s_frag);
     }
 
     // update m,d
@@ -796,8 +775,8 @@ void MultiQueryAppendAttention(
     const paddle::Tensor &seq_lens_q,
     const paddle::Tensor &seq_lens_kv,
     const paddle::Tensor &seq_lens_encoder,
-    const paddle::Tensor &batch_id_per_token,
-    const paddle::Tensor &cu_seqlens_q,
+    const paddle::Tensor &padding_offsets,
+    const paddle::Tensor &cum_offsets,
     const paddle::Tensor &block_table,
     const paddle::Tensor &batch_ids,
     const paddle::Tensor &tile_ids_per_batch,
@@ -903,9 +882,8 @@ void MultiQueryAppendAttention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -914,7 +892,6 @@ void MultiQueryAppendAttention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           nullptr,
           nullptr,
           nullptr,
@@ -962,9 +939,8 @@ void MultiQueryAppendAttention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -973,7 +949,6 @@ void MultiQueryAppendAttention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
           static_cast<float *>(tmp_m->ptr()),
           static_cast<float *>(tmp_d->ptr()),
@@ -999,7 +974,7 @@ void MultiQueryAppendAttention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                cu_seqlens_q.data<int>(),
+                cum_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
                                  const_cast<T *>(shift_bias.get().data<T>()))
                            : nullptr,
@@ -1034,8 +1009,7 @@ void MultiQueryAppendAttention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                batch_id_per_token.data<int>(),
-                cu_seqlens_q.data<int>(),
+                padding_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
                                  const_cast<T *>(shift_bias.get().data<T>()))
                            : nullptr,
@@ -1088,18 +1062,12 @@ void MultiQueryAppendAttention(
     if (!is_decoder) {
       chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
     }
+    const int num_chunks = div_up(max_dec_len, chunk_size);
 
-    int32_t attn_mask_len;
-    if (attn_mask) {
-        attn_mask_len = attn_mask.get().shape()[1];
-    } else {
-        attn_mask_len = -1;
-    }
-
-    const int num_chunks = div_up(max_seq_len, chunk_size);
     dim3 grids(num_blocks_x_cpu, num_chunks, kv_num_heads);
     dim3 blocks(32, num_warps);
-    if (num_chunks <= 0) {
+
+    if (num_chunks <= 1) {
       auto nosplit_kv_kernel =
           multi_query_append_attention_warp1_4_kernel<NV_TYPE,
                                                       false,
@@ -1135,11 +1103,8 @@ void MultiQueryAppendAttention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
-          attn_mask ? const_cast<bool *>(attn_mask.get().data<bool>())
-                        : nullptr,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1148,13 +1113,11 @@ void MultiQueryAppendAttention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           nullptr,
           nullptr,
           nullptr,
           reinterpret_cast<OUT_NV_TYPE *>(out->data<OutT>()),
-          speculate_max_draft_token_num,
-          attn_mask_len);
+          speculate_max_draft_token_num);
     } else {
       phi::Allocator::AllocationPtr tmp_workspace, tmp_m, tmp_d;
       if (is_decoder) {
@@ -1199,8 +1162,8 @@ void MultiQueryAppendAttention(
           reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_k.data<T>())),
           reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_v.data<T>())),
           shift_bias ? reinterpret_cast<NV_TYPE *>(
-                            const_cast<T *>(shift_bias.get().data<T>()))
-                      : nullptr,
+                           const_cast<T *>(shift_bias.get().data<T>()))
+                     : nullptr,
           smooth_weight ? reinterpret_cast<NV_TYPE *>(
                               const_cast<T *>(smooth_weight.get().data<T>()))
                         : nullptr,
@@ -1208,11 +1171,8 @@ void MultiQueryAppendAttention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
-          attn_mask ? const_cast<bool *>(attn_mask.get().data<bool>())
-                        : nullptr,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1221,13 +1181,11 @@ void MultiQueryAppendAttention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
           static_cast<float *>(tmp_m->ptr()),
           static_cast<float *>(tmp_d->ptr()),
           reinterpret_cast<OUT_NV_TYPE *>(out->data<OutT>()),
-          speculate_max_draft_token_num,
-          attn_mask_len);
+          speculate_max_draft_token_num);
 
       // merge
       constexpr int vec_size = num_elems_per_128b<NV_TYPE>();
@@ -1249,10 +1207,10 @@ void MultiQueryAppendAttention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                cu_seqlens_q.data<int>(),
+                cum_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
-                                  const_cast<T *>(shift_bias.get().data<T>()))
-                            : nullptr,
+                                 const_cast<T *>(shift_bias.get().data<T>()))
+                           : nullptr,
                 smooth_weight ? reinterpret_cast<NV_TYPE *>(const_cast<T *>(
                                     smooth_weight.get().data<T>()))
                               : nullptr,
@@ -1269,14 +1227,14 @@ void MultiQueryAppendAttention(
         constexpr int blockx = HEAD_DIM / vec_size;
         constexpr int blocky = (128 + blockx - 1) / blockx;
         dim3 grids_merge(min(sm_count * 4, token_num),
-                          num_heads);
+                         num_heads);
         dim3 blocks_merge(blockx, blocky);
         merge_multi_chunks_v2_kernel<NV_TYPE,
-                                      vec_size,
-                                      blocky,
-                                      HEAD_DIM,
-                                      OUT_NV_TYPE,
-                                      ENABLE_PREFILL>
+                                     vec_size,
+                                     blocky,
+                                     HEAD_DIM,
+                                     OUT_NV_TYPE,
+                                     ENABLE_PREFILL>
             <<<grids_merge, blocks_merge, 0, stream>>>(
                 reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
                 static_cast<float *>(tmp_m->ptr()),
@@ -1284,11 +1242,10 @@ void MultiQueryAppendAttention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                batch_id_per_token.data<int>(),
-                cu_seqlens_q.data<int>(),
+                padding_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
-                                  const_cast<T *>(shift_bias.get().data<T>()))
-                            : nullptr,
+                                 const_cast<T *>(shift_bias.get().data<T>()))
+                           : nullptr,
                 smooth_weight ? reinterpret_cast<NV_TYPE *>(const_cast<T *>(
                                     smooth_weight.get().data<T>()))
                               : nullptr,
@@ -1332,8 +1289,8 @@ void CascadeAppendAttentionC16Kernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -1395,8 +1352,8 @@ void CascadeAppendAttentionC16Kernel(
                                 seq_lens_q,
                                 seq_lens_kv,
                                 seq_lens_encoder,
-                                batch_id_per_token,
-                                cu_seqlens_q,
+                                padding_offsets,
+                                cum_offsets,
                                 block_table,
                                 batch_ids,
                                 tile_ids_per_batch,

custom_ops/gpu_ops/append_attn/append_attention_c4_impl.cuh

@@ -46,9 +46,8 @@ __global__ void multi_query_append_attention_c4_kernel(
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ batch_ids,
     const int *__restrict__ tile_ids_per_batch,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
-    const int *__restrict__ mask_offset,
     const int max_seq_len,
     const int max_dec_len,
     const int max_block_num_per_seq,
@@ -57,7 +56,6 @@ __global__ void multi_query_append_attention_c4_kernel(
     const float quant_min_bound,
     const float in_scale,
     const uint32_t chunk_size,
-    const int num_blocks_x_cpu,
     T *__restrict__ tmp_workspace,  // split kv [token_num, num_chunks,
                                     // num_heads, head_dim]
     float *__restrict__ tmp_m,      // [token_num, num_chunks, num_heads]
@@ -86,11 +84,6 @@ __global__ void multi_query_append_attention_c4_kernel(
 
   block_table_now = block_table + batch_id * max_block_num_per_seq;
 
-  //When cudagraph capture prefill, may launch more gridDim.x
-  if(btid >= static_cast<uint32_t>(num_blocks_x_cpu)){
-    return;
-  }
-
   const uint32_t q_len = seq_lens[batch_id];
   if (q_len <= 0) {
     return;
@@ -151,7 +144,8 @@ __global__ void multi_query_append_attention_c4_kernel(
   const uint32_t kv_h_stride = BLOCK_SIZE * HEAD_DIM / 2;
   const uint32_t kv_b_stride = HEAD_DIM / 2;
   const uint32_t kv_d_stride = BLOCK_SIZE / 2;
-  const uint32_t q_start_seq_id = cu_seqlens_q[batch_id];
+  const uint32_t q_start_seq_id =
+      batch_id * max_seq_len - __ldg(&cum_offsets[batch_id]);
   const uint32_t q_base_seq_id_this_block =
       (tile_id * NUM_WARPS + wid) * num_frags_x * 16;
   const uint32_t q_offset = q_start_seq_id * q_ori_n_stride +
@@ -179,7 +173,6 @@ __global__ void multi_query_append_attention_c4_kernel(
   } else {
     o_base_ptr_int8 = out + o_offset;
   }
-  const int *mask_offset_this_seq = mask_offset ? mask_offset + q_start_seq_id * 2 : nullptr;
   smem_t qo_smem(smem);
 
   uint32_t q_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -256,7 +249,7 @@ __global__ void multi_query_append_attention_c4_kernel(
                          kv_len - q_len +
                              tile_id * num_rows_per_block / GROUP_SIZE,
                          chunk_start)))
-              : mask_offset ? 0 : chunk_len) /
+              : chunk_len) /
       (num_frags_z * 16);
 
   uint32_t k_smem_offset_r =
@@ -341,15 +334,12 @@ __global__ void multi_query_append_attention_c4_kernel(
              NUM_WARPS,
              num_frags_x,
              num_frags_y,
-             num_frags_z>(nullptr,
-                          q_base_seq_id_this_block,
+             num_frags_z>(q_base_seq_id_this_block,
                           kv_idx_base,
                           q_len,
                           kv_len,
                           chunk_end,
-                          -1,
-                          s_frag,
-                          mask_offset_this_seq);
+                          s_frag);
     }
 
     update_mdo_states<num_frags_x, num_frags_y, num_frags_z>(
@@ -514,10 +504,8 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ batch_ids,
     const int *__restrict__ tile_ids_per_batch,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
-    const int *__restrict__ mask_offset,
-    const bool *__restrict__ attn_mask,    // [bsz, max_q, max_q] for tree-mask
     const int max_seq_len,
     const int max_dec_len,
     const int max_block_num_per_seq,
@@ -526,14 +514,12 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
     const float quant_min_bound,
     const float in_scale,
     const uint32_t chunk_size,
-    const int num_blocks_x_cpu,
     T *__restrict__ tmp_workspace,  // split kv [token_num, num_chunks,
                                     // num_heads, head_dim]
     float *__restrict__ tmp_m,      // [token_num, num_chunks, num_heads]
     float *__restrict__ tmp_d,      // [token_num, num_chunks, num_heads]
     OutT *__restrict__ out,
-    const int speculate_max_draft_token_num = 5,
-    const int32_t attn_mask_len = -1) {
+    const int speculate_max_draft_token_num = 5) {
   constexpr uint32_t num_vecs_per_head = HEAD_DIM / num_elems_per_128b<T>();
   constexpr uint32_t num_vecs_per_head_k =
       HEAD_DIM / 2 / num_elems_per_128b<CacheT>();
@@ -556,11 +542,6 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
   const uint32_t num_rows_per_block = num_frags_x * 16;
   const int *block_table_now = block_table + batch_id * max_block_num_per_seq;
 
-  //When cudagraph capture prefill, may launch more gridDim.x
-  if(btid >= static_cast<uint32_t>(num_blocks_x_cpu)){
-    return;
-  }
-
   const uint32_t q_len = seq_lens[batch_id];
   if (q_len <= 0) {
     return;
@@ -620,7 +601,8 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
   const uint32_t kv_h_stride = BLOCK_SIZE * HEAD_DIM / 2;
   const uint32_t kv_b_stride = HEAD_DIM / 2;
   const uint32_t kv_d_stride = BLOCK_SIZE / 2;
-  const uint32_t q_start_seq_id = cu_seqlens_q[batch_id];
+  const uint32_t q_start_seq_id =
+      batch_id * max_seq_len - __ldg(&cum_offsets[batch_id]);
   const uint32_t q_base_seq_id_this_block = tile_id * num_frags_x * 16;
   const uint32_t q_offset = q_start_seq_id * q_ori_n_stride +
                             q_head_idx * HEAD_DIM +
@@ -647,7 +629,7 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
           tid % 8 * num_elems_per_128b<T>();
     }
   }
-  const int *mask_offset_this_seq = mask_offset ? mask_offset + q_start_seq_id * 2 : nullptr;
+
   smem_t qo_smem(smem);
 
   uint32_t q_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -723,9 +705,10 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
   const uint32_t mask_check_iteration =
       (CAUSAL ? (min(chunk_len,
                      sub_if_greater_or_zero(
-                         kv_len - q_len,
+                         kv_len - q_len +
+                             tile_id * num_rows_per_block / GROUP_SIZE,
                          chunk_start)))
-              : mask_offset ? 0 : chunk_len) /
+              : chunk_len) /
       (NUM_WARP_KV * num_frags_z * 16);
 
   uint32_t k_smem_offset_r =
@@ -807,15 +790,12 @@ __global__ void multi_query_append_attention_c4_warp1_4_kernel(
              NUM_WARPS,
              num_frags_x,
              num_frags_y,
-             num_frags_z>(attn_mask ? attn_mask + batch_id * attn_mask_len *attn_mask_len : nullptr,
-                          q_base_seq_id_this_block,
+             num_frags_z>(q_base_seq_id_this_block,
                           kv_idx_base + wid * num_frags_z * 16,
                           q_len,
                           kv_len,
                           chunk_end,
-                          attn_mask_len,
-                          s_frag,
-                          mask_offset_this_seq);
+                          s_frag);
     }
 
     update_mdo_states<num_frags_x, num_frags_y, num_frags_z>(
@@ -982,8 +962,8 @@ void MultiQueryAppendC4Attention(
     const paddle::Tensor &seq_lens_q,
     const paddle::Tensor &seq_lens_kv,
     const paddle::Tensor &seq_lens_encoder,
-    const paddle::Tensor &batch_id_per_token,
-    const paddle::Tensor &cu_seqlens_q,
+    const paddle::Tensor &padding_offsets,
+    const paddle::Tensor &cum_offsets,
     const paddle::Tensor &block_table,
     const paddle::Tensor &batch_ids,
     const paddle::Tensor &tile_ids_per_batch,
@@ -1108,9 +1088,8 @@ void MultiQueryAppendC4Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1119,7 +1098,6 @@ void MultiQueryAppendC4Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           nullptr,
           nullptr,
           nullptr,
@@ -1173,9 +1151,8 @@ void MultiQueryAppendC4Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1184,7 +1161,6 @@ void MultiQueryAppendC4Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
           static_cast<float *>(tmp_m->ptr()),
           static_cast<float *>(tmp_d->ptr()),
@@ -1210,7 +1186,7 @@ void MultiQueryAppendC4Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                cu_seqlens_q.data<int>(),
+                cum_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
                                  const_cast<T *>(shift_bias.get().data<T>()))
                            : nullptr,
@@ -1245,8 +1221,7 @@ void MultiQueryAppendC4Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                batch_id_per_token.data<int>(),
-                cu_seqlens_q.data<int>(),
+                padding_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
                                  const_cast<T *>(shift_bias.get().data<T>()))
                            : nullptr,
@@ -1311,18 +1286,10 @@ void MultiQueryAppendC4Attention(
     if (!is_decoder) {
       chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
     }
-
-    const int num_chunks = div_up(max_seq_len, chunk_size);
-    int32_t attn_mask_len;
-    if (attn_mask) {
-        attn_mask_len = attn_mask.get().shape()[1];
-    } else {
-        attn_mask_len = -1;
-    }
-
+    const int num_chunks = div_up(max_dec_len, chunk_size);
     dim3 grids(num_blocks_x_cpu, num_chunks, kv_num_heads);
     dim3 blocks(32, num_warps);
-    if (num_chunks <= 0) {
+    if (num_chunks <= 1) {
       auto nosplit_kv_kernel =
           multi_query_append_attention_c4_warp1_4_kernel<NV_TYPE,
                                                          uint8_t,
@@ -1366,11 +1333,8 @@ void MultiQueryAppendC4Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
-          attn_mask ? const_cast<bool *>(attn_mask.get().data<bool>())
-                        : nullptr,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1379,13 +1343,11 @@ void MultiQueryAppendC4Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           nullptr,
           nullptr,
           nullptr,
           reinterpret_cast<OUT_NV_TYPE *>(out->data<OutT>()),
-          speculate_max_draft_token_num,
-          attn_mask_len);
+          speculate_max_draft_token_num);
     } else {
       phi::Allocator::AllocationPtr tmp_workspace, tmp_m, tmp_d;
       if (is_decoder) {
@@ -1431,15 +1393,15 @@ void MultiQueryAppendC4Attention(
           const_cast<uint8_t *>(cache_v.data<uint8_t>()),
           reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_k_scale.data<T>())),
           cache_k_zp ? reinterpret_cast<NV_TYPE *>(
-                            const_cast<T *>(cache_k_zp.get().data<T>()))
-                      : nullptr,
+                           const_cast<T *>(cache_k_zp.get().data<T>()))
+                     : nullptr,
           reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_v_scale.data<T>())),
           cache_v_zp ? reinterpret_cast<NV_TYPE *>(
-                            const_cast<T *>(cache_v_zp.get().data<T>()))
-                      : nullptr,
+                           const_cast<T *>(cache_v_zp.get().data<T>()))
+                     : nullptr,
           shift_bias ? reinterpret_cast<NV_TYPE *>(
-                            const_cast<T *>(shift_bias.get().data<T>()))
-                      : nullptr,
+                           const_cast<T *>(shift_bias.get().data<T>()))
+                     : nullptr,
           smooth_weight ? reinterpret_cast<NV_TYPE *>(
                               const_cast<T *>(smooth_weight.get().data<T>()))
                         : nullptr,
@@ -1447,11 +1409,8 @@ void MultiQueryAppendC4Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
-          attn_mask ? const_cast<bool *>(attn_mask.get().data<bool>())
-                        : nullptr,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1460,13 +1419,11 @@ void MultiQueryAppendC4Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
           static_cast<float *>(tmp_m->ptr()),
           static_cast<float *>(tmp_d->ptr()),
           reinterpret_cast<OUT_NV_TYPE *>(out->data<OutT>()),
-          speculate_max_draft_token_num,
-          attn_mask_len);
+          speculate_max_draft_token_num);
       // merge
       constexpr int vec_size = num_elems_per_128b<NV_TYPE>();
       if (is_decoder) {
@@ -1487,10 +1444,10 @@ void MultiQueryAppendC4Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                cu_seqlens_q.data<int>(),
+                cum_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
-                                  const_cast<T *>(shift_bias.get().data<T>()))
-                            : nullptr,
+                                 const_cast<T *>(shift_bias.get().data<T>()))
+                           : nullptr,
                 smooth_weight ? reinterpret_cast<NV_TYPE *>(const_cast<T *>(
                                     smooth_weight.get().data<T>()))
                               : nullptr,
@@ -1507,14 +1464,14 @@ void MultiQueryAppendC4Attention(
         constexpr int blockx = HEAD_DIM / vec_size;
         constexpr int blocky = (128 + blockx - 1) / blockx;
         dim3 grids_merge(min(sm_count * 4, token_num),
-                          num_heads);
+                         num_heads);
         dim3 blocks_merge(blockx, blocky);
         merge_multi_chunks_v2_kernel<NV_TYPE,
-                                      vec_size,
-                                      blocky,
-                                      HEAD_DIM,
-                                      OUT_NV_TYPE,
-                                      ENABLE_PREFILL>
+                                     vec_size,
+                                     blocky,
+                                     HEAD_DIM,
+                                     OUT_NV_TYPE,
+                                     ENABLE_PREFILL>
             <<<grids_merge, blocks_merge, 0, stream>>>(
                 reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
                 static_cast<float *>(tmp_m->ptr()),
@@ -1522,11 +1479,10 @@ void MultiQueryAppendC4Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                batch_id_per_token.data<int>(),
-                cu_seqlens_q.data<int>(),
+                padding_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
-                                  const_cast<T *>(shift_bias.get().data<T>()))
-                            : nullptr,
+                                 const_cast<T *>(shift_bias.get().data<T>()))
+                           : nullptr,
                 smooth_weight ? reinterpret_cast<NV_TYPE *>(const_cast<T *>(
                                     smooth_weight.get().data<T>()))
                               : nullptr,
@@ -1570,8 +1526,8 @@ void CascadeAppendAttentionC4Kernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -1637,8 +1593,8 @@ void CascadeAppendAttentionC4Kernel(
                                 seq_lens_q,
                                 seq_lens_kv,
                                 seq_lens_encoder,
-                                batch_id_per_token,
-                                cu_seqlens_q,
+                                padding_offsets,
+                                cum_offsets,
                                 block_table,
                                 batch_ids,
                                 tile_ids_per_batch,

custom_ops/gpu_ops/append_attn/append_attention_c8_impl.cuh

@@ -32,24 +32,22 @@ template <typename T,
           typename OutT = T,
           bool ENABLE_PREFILL = true,
           bool is_scale_channel_wise = false,
-          bool IsFP8 = false,
-          bool IsDynamicC8 = false>
+          bool IsFP8=false>
 __global__ void multi_query_append_attention_c8_kernel(
     T *__restrict__ q,             // [token_num, (num_heads + 2* kv_num_head) * head_dim]
     CacheT *__restrict__ cache_k,  // [max_block_num, num_heads, block_size,
                                    // head_dim]
     CacheT *__restrict__ cache_v,
-    const T *__restrict__ cache_k_scale,  // [num_kv_heads] or [max_block_num, num_heads, block_size]
-    const T *__restrict__ cache_v_scale,  // [num_kv_heads] or [max_block_num, num_heads, block_size]
+    const T *__restrict__ cache_k_scale,  // [num_kv_heads]
+    const T *__restrict__ cache_v_scale,  // [num_kv_heads]
     const T *__restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
     const T *__restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
     const int *__restrict__ seq_lens,
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ batch_ids,
     const int *__restrict__ tile_ids_per_batch,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
-    const int *__restrict__ mask_offset,
     const int max_seq_len,
     const int max_dec_len,
     const int max_block_num_per_seq,
@@ -58,7 +56,6 @@ __global__ void multi_query_append_attention_c8_kernel(
     const float quant_min_bound,
     const float in_scale,
     const uint32_t chunk_size,
-    const int num_blocks_x_cpu,
     T *__restrict__ tmp_workspace,  // split kv [token_num, num_chunks,
                                     // num_heads, head_dim]
     float *__restrict__ tmp_m,      // [token_num, num_chunks, num_heads]
@@ -88,40 +85,33 @@ __global__ void multi_query_append_attention_c8_kernel(
 
   block_table_now = block_table + batch_id * max_block_num_per_seq;
 
-  //When cudagraph capture prefill, may launch more gridDim.x
-  if(btid >= static_cast<uint32_t>(num_blocks_x_cpu)){
-    return;
-  }
-
   const uint32_t q_len = seq_lens[batch_id];
   if (q_len <= 0) {
     return;
   }
 
-  T cache_k_scale_reg[IsDynamicC8 ? num_frags_z * 2 : num_frags_y * 4];
-  T cache_v_scale_reg[IsDynamicC8 ? num_frags_z * 4 : num_frags_y * 2];
-  if constexpr (!IsDynamicC8) {
-    if constexpr (is_scale_channel_wise) {
-      int scale_col_base = threadIdx.x % 4 * 2 + kv_head_idx * HEAD_DIM;
-      const T *cache_k_scale_cur_head = cache_k_scale + scale_col_base;
-      for (int i = 0; i < num_frags_y; ++i) {
-        const int scale_idx = i * 16;
-        cache_k_scale_reg[i * 4] = cache_k_scale_cur_head[scale_idx];
-        cache_k_scale_reg[i * 4 + 1] = cache_k_scale_cur_head[scale_idx + 1];
-        cache_k_scale_reg[i * 4 + 2] = cache_k_scale_cur_head[scale_idx + 8];
-        cache_k_scale_reg[i * 4 + 3] = cache_k_scale_cur_head[scale_idx + 9];
-      }
-      scale_col_base = threadIdx.x / 4 + kv_head_idx * HEAD_DIM;
-      const T *cache_v_scale_cur_head = cache_v_scale + scale_col_base;
-      for (int i = 0; i < num_frags_y; ++i) {
-        const int scale_idx = i * 16;
-        cache_v_scale_reg[i * 2] = cache_v_scale_cur_head[scale_idx];
-        cache_v_scale_reg[i * 2 + 1] = cache_v_scale_cur_head[scale_idx + 8];
-      }
-    } else {
-      cache_k_scale_reg[0] = cache_k_scale[kv_head_idx];
-      cache_v_scale_reg[0] = cache_v_scale[kv_head_idx];
+  T cache_k_scale_reg[num_frags_y * 4];
+  T cache_v_scale_reg[num_frags_y * 2];
+  if (is_scale_channel_wise) {
+    int scale_col_base = threadIdx.x % 4 * 2 + kv_head_idx * HEAD_DIM;
+    const T *cache_k_scale_cur_head = cache_k_scale + scale_col_base;
+    for (int i = 0; i < num_frags_y; ++i) {
+      const int scale_idx = i * 16;
+      cache_k_scale_reg[i * 4] = cache_k_scale_cur_head[scale_idx];
+      cache_k_scale_reg[i * 4 + 1] = cache_k_scale_cur_head[scale_idx + 1];
+      cache_k_scale_reg[i * 4 + 2] = cache_k_scale_cur_head[scale_idx + 8];
+      cache_k_scale_reg[i * 4 + 3] = cache_k_scale_cur_head[scale_idx + 9];
     }
+    scale_col_base = threadIdx.x / 4 + kv_head_idx * HEAD_DIM;
+    const T *cache_v_scale_cur_head = cache_v_scale + scale_col_base;
+    for (int i = 0; i < num_frags_y; ++i) {
+      const int scale_idx = i * 16;
+      cache_v_scale_reg[i * 2] = cache_v_scale_cur_head[scale_idx];
+      cache_v_scale_reg[i * 2 + 1] = cache_v_scale_cur_head[scale_idx + 8];
+    }
+  } else {
+    cache_k_scale_reg[0] = cache_k_scale[kv_head_idx];
+    cache_v_scale_reg[0] = cache_v_scale[kv_head_idx];
   }
 
   const uint32_t q_end =
@@ -161,7 +151,8 @@ __global__ void multi_query_append_attention_c8_kernel(
   const uint32_t kv_h_stride = BLOCK_SIZE * HEAD_DIM;
   const uint32_t kv_b_stride = HEAD_DIM;
   const uint32_t kv_d_stride = BLOCK_SIZE;
-  const uint32_t q_start_seq_id = cu_seqlens_q[batch_id];
+  const uint32_t q_start_seq_id =
+      batch_id * max_seq_len - __ldg(&cum_offsets[batch_id]);
   const uint32_t q_base_seq_id_this_block =
       (tile_id * NUM_WARPS + wid) * num_frags_x * 16;
   const uint32_t q_offset = q_start_seq_id * q_ori_n_stride +
@@ -189,7 +180,6 @@ __global__ void multi_query_append_attention_c8_kernel(
   } else {
     o_base_ptr_int8 = out + o_offset;
   }
-  const int *mask_offset_this_seq = mask_offset ? mask_offset + q_start_seq_id * 2 : nullptr;
   smem_t qo_smem(smem);
 
   uint32_t q_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -210,13 +200,6 @@ __global__ void multi_query_append_attention_c8_kernel(
   smem_t k_smem(smem + NUM_WARPS * num_frags_x * 16 * HEAD_DIM * sizeof(T)),
       v_smem(smem + NUM_WARPS * num_frags_x * 16 * HEAD_DIM * sizeof(T) +
              num_frags_z * 16 * HEAD_DIM * sizeof(CacheT));
-  T* k_smem_scale = nullptr;
-  T* v_smem_scale = nullptr;
-  if constexpr (IsDynamicC8) {
-    k_smem_scale = reinterpret_cast<T*>(smem + NUM_WARPS * num_frags_x * 16 * HEAD_DIM * sizeof(T) +
-                                         num_frags_z * 16 * HEAD_DIM * sizeof(CacheT) * 2);
-    v_smem_scale = k_smem_scale + num_frags_z * 16;
-  }
 
 
   const uint32_t num_iterations = div_up(
@@ -234,7 +217,7 @@ __global__ void multi_query_append_attention_c8_kernel(
                          kv_len - q_len +
                              tile_id * num_rows_per_block / GROUP_SIZE,
                          chunk_start)))
-              : mask_offset ? 0 : chunk_len) /
+              : chunk_len) /
       (num_frags_z * 16);
 
   uint32_t k_smem_offset_r =
@@ -298,22 +281,10 @@ __global__ void multi_query_append_attention_c8_kernel(
 
 #pragma unroll 1
   for (uint32_t iter = 0; iter < num_iterations; ++iter) {
-    if constexpr (IsDynamicC8) {
-      produce_k_dynamic_scale<BLOCK_SIZE, num_frags_z, NUM_WARP_Q, T>(
-        k_smem_scale,
-        cache_k_scale_reg,
-        block_table_now,
-        cache_k_scale,
-        kv_idx_base,
-        kv_num_heads,
-        kv_head_idx,
-        chunk_end
-      );
-    }
     wait_group<1>();
     __syncthreads();
     // s = qk
-    compute_qk_c8<num_frags_x, num_frags_y, num_frags_z, T, CacheT, is_scale_channel_wise, IsFP8, IsDynamicC8>(
+    compute_qk_c8<num_frags_x, num_frags_y, num_frags_z, T, CacheT, is_scale_channel_wise, IsFP8>(
         &qo_smem,
         &q_smem_offset_r,
         &k_smem,
@@ -330,15 +301,12 @@ __global__ void multi_query_append_attention_c8_kernel(
              NUM_WARPS,
              num_frags_x,
              num_frags_y,
-             num_frags_z>(nullptr,
-                          q_base_seq_id_this_block,
+             num_frags_z>(q_base_seq_id_this_block,
                           kv_idx_base,
                           q_len,
                           kv_len,
                           chunk_end,
-                          -1,
-                          s_frag,
-                          mask_offset_this_seq);
+                          s_frag);
     }
 
     // update m,d
@@ -346,7 +314,6 @@ __global__ void multi_query_append_attention_c8_kernel(
         s_frag, o_frag, m_frag, d_frag);
     __syncthreads();
 
-    const int ori_kv_idx_base = kv_idx_base;
     kv_idx_base += num_frags_z * 16;
     produce_k_blockwise_c8<SharedMemFillMode::kNoFill,
                            NUM_WARPS,
@@ -365,18 +332,6 @@ __global__ void multi_query_append_attention_c8_kernel(
                                        chunk_end,
                                        const_k_offset);
     commit_group();
-    if constexpr (IsDynamicC8) {
-      produce_v_dynamic_scale<BLOCK_SIZE, num_frags_z, NUM_WARP_Q, T>(
-        v_smem_scale,
-        cache_v_scale_reg,
-        block_table_now,
-        cache_v_scale,
-        ori_kv_idx_base,
-        kv_num_heads,
-        kv_head_idx,
-        chunk_end
-      );
-    }
     wait_group<1>();
     __syncthreads();
 
@@ -387,9 +342,7 @@ __global__ void multi_query_append_attention_c8_kernel(
                      BLOCK_SIZE,
                      T,
                      CacheT,
-                     is_scale_channel_wise,
-                     IsFP8,
-                     IsDynamicC8>(
+                     is_scale_channel_wise, IsFP8>(
         &v_smem, &v_smem_offset_r, s_frag, o_frag, d_frag, cache_v_scale_reg);
     __syncthreads();
 
@@ -506,25 +459,22 @@ template <typename T,
           typename OutT = T,
           bool ENABLE_PREFILL = true,
           bool is_scale_channel_wise=false,
-          bool IsFP8 = false,
-          bool IsDynamicC8 = false>
+          bool IsFP8=false>
 __global__ void multi_query_append_attention_c8_warp1_4_kernel(
     T *__restrict__ q,             // [token_num, (num_heads + 2* kv_num_head) * head_dim]
     CacheT *__restrict__ cache_k,  // [max_block_num, num_heads, block_size,
                                    // head_dim]
     CacheT *__restrict__ cache_v,
-    const T *__restrict__ cache_k_scale,  // [num_kv_heads] or [max_block_num, num_heads, block_size]
-    const T *__restrict__ cache_v_scale,  // [num_kv_heads] or [max_block_num, num_heads, block_size]
+    const T *__restrict__ cache_k_scale,  // [num_kv_heads, head_dim]
+    const T *__restrict__ cache_v_scale,  // [num_kv_heads, head_dim]
     const T *__restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
     const T *__restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
     const int *__restrict__ seq_lens,
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ batch_ids,
     const int *__restrict__ tile_ids_per_batch,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
-    const int *__restrict__ mask_offset,
-    const bool *__restrict__ attn_mask,   // [bsz, max_q, max_q] for tree-mask
     const int max_seq_len,
     const int max_dec_len,
     const int max_block_num_per_seq,
@@ -533,14 +483,12 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
     const float quant_min_bound,
     const float in_scale,
     const uint32_t chunk_size,
-    const int num_blocks_x_cpu,
     T *__restrict__ tmp_workspace,  // split kv [token_num, num_chunks,
                                     // num_heads, head_dim]
     float *__restrict__ tmp_m,      // [token_num, num_chunks, num_heads]
     float *__restrict__ tmp_d,      // [token_num, num_chunks, num_heads]
     OutT *__restrict__ out,
-    const int speculate_max_draft_token_num = 5,
-    const int32_t attn_mask_len = -1) {
+    const int speculate_max_draft_token_num = 5) {
   constexpr uint32_t num_vecs_per_head = HEAD_DIM / num_elems_per_128b<T>();
   constexpr uint32_t num_vecs_per_head_k =
       HEAD_DIM / num_elems_per_128b<CacheT>();
@@ -563,39 +511,32 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
   const uint32_t num_rows_per_block = num_frags_x * 16;
   const int *block_table_now = block_table + batch_id * max_block_num_per_seq;
 
-  //When cudagraph capture prefill, may launch more gridDim.x
-  if(btid >= static_cast<uint32_t>(num_blocks_x_cpu)){
-    return;
-  }
-
   const uint32_t q_len = seq_lens[batch_id];
   if (q_len <= 0) {
     return;
   }
-  T cache_k_scale_reg[IsDynamicC8 ? num_frags_z * 2 : num_frags_y * 4];
-  T cache_v_scale_reg[IsDynamicC8 ? num_frags_z * 4 : num_frags_y * 2];
-  if constexpr (!IsDynamicC8) {
-    if constexpr (is_scale_channel_wise) {
-      int scale_col_base = threadIdx.x % 4 * 2 + kv_head_idx * HEAD_DIM;
-      const T *cache_k_scale_cur_head = cache_k_scale + scale_col_base;
-      for (int i = 0; i < num_frags_y; ++i) {
-        const int scale_idx = i * 16;
-        cache_k_scale_reg[i * 4] = cache_k_scale_cur_head[scale_idx];
-        cache_k_scale_reg[i * 4 + 1] = cache_k_scale_cur_head[scale_idx + 1];
-        cache_k_scale_reg[i * 4 + 2] = cache_k_scale_cur_head[scale_idx + 8];
-        cache_k_scale_reg[i * 4 + 3] = cache_k_scale_cur_head[scale_idx + 9];
-      }
-      scale_col_base = threadIdx.x / 4 + kv_head_idx * HEAD_DIM;
-      const T *cache_v_scale_cur_head = cache_v_scale + scale_col_base;
-      for (int i = 0; i < num_frags_y; ++i) {
-        const int scale_idx = i * 16;
-        cache_v_scale_reg[i * 2] = cache_v_scale_cur_head[scale_idx];
-        cache_v_scale_reg[i * 2 + 1] = cache_v_scale_cur_head[scale_idx + 8];
-      }
-    } else {
-      cache_k_scale_reg[0] = cache_k_scale[kv_head_idx];
-      cache_v_scale_reg[0] = cache_v_scale[kv_head_idx];
+  T cache_k_scale_reg[num_frags_y * 4];
+  T cache_v_scale_reg[num_frags_y * 2];
+  if (is_scale_channel_wise) {
+    int scale_col_base = threadIdx.x % 4 * 2 + kv_head_idx * HEAD_DIM;
+    const T *cache_k_scale_cur_head = cache_k_scale + scale_col_base;
+    for (int i = 0; i < num_frags_y; ++i) {
+      const int scale_idx = i * 16;
+      cache_k_scale_reg[i * 4] = cache_k_scale_cur_head[scale_idx];
+      cache_k_scale_reg[i * 4 + 1] = cache_k_scale_cur_head[scale_idx + 1];
+      cache_k_scale_reg[i * 4 + 2] = cache_k_scale_cur_head[scale_idx + 8];
+      cache_k_scale_reg[i * 4 + 3] = cache_k_scale_cur_head[scale_idx + 9];
     }
+    scale_col_base = threadIdx.x / 4 + kv_head_idx * HEAD_DIM;
+    const T *cache_v_scale_cur_head = cache_v_scale + scale_col_base;
+    for (int i = 0; i < num_frags_y; ++i) {
+      const int scale_idx = i * 16;
+      cache_v_scale_reg[i * 2] = cache_v_scale_cur_head[scale_idx];
+      cache_v_scale_reg[i * 2 + 1] = cache_v_scale_cur_head[scale_idx + 8];
+    }
+  } else {
+    cache_k_scale_reg[0] = cache_k_scale[kv_head_idx];
+    cache_v_scale_reg[0] = cache_v_scale[kv_head_idx];
   }
   const uint32_t q_end =
       min(q_len, div_up((tile_id + 1) * num_rows_per_block, GROUP_SIZE));
@@ -634,7 +575,8 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
   const uint32_t kv_h_stride = BLOCK_SIZE * HEAD_DIM;
   const uint32_t kv_b_stride = HEAD_DIM;
   const uint32_t kv_d_stride = BLOCK_SIZE;
-  const uint32_t q_start_seq_id = cu_seqlens_q[batch_id];
+  const uint32_t q_start_seq_id =
+      batch_id * max_seq_len - __ldg(&cum_offsets[batch_id]);
   const uint32_t q_base_seq_id_this_block = tile_id * num_frags_x * 16;
   const uint32_t q_offset = q_start_seq_id * q_ori_n_stride +
                             q_head_idx * HEAD_DIM +
@@ -661,7 +603,7 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
           tid % 8 * num_elems_per_128b<T>();
     }
   }
-  const int *mask_offset_this_seq = mask_offset ? mask_offset + q_start_seq_id * 2 : nullptr;
+
   smem_t qo_smem(smem);
 
   uint32_t q_smem_offset_r = smem_t::get_permuted_offset<num_vecs_per_head>(
@@ -686,13 +628,6 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
   smem_t k_smem(smem + num_frags_x * 16 * HEAD_DIM * sizeof(T)),
       v_smem(smem + num_frags_x * 16 * HEAD_DIM * sizeof(T) +
              NUM_WARP_KV * num_frags_z * 16 * HEAD_DIM * sizeof(CacheT));
-  T* k_smem_scale = nullptr;
-  T* v_smem_scale = nullptr;
-  if constexpr (IsDynamicC8) {
-    k_smem_scale = reinterpret_cast<T*>(smem + num_frags_x * 16 * HEAD_DIM * sizeof(T) +
-                                        NUM_WARP_KV * num_frags_z * 16 * HEAD_DIM * sizeof(CacheT) * 2);
-    v_smem_scale = k_smem_scale + NUM_WARP_KV * num_frags_z * 16;
-  }
 
   const uint32_t num_iterations = div_up(
       CAUSAL
@@ -709,7 +644,7 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
                          kv_len - q_len +
                              tile_id * num_rows_per_block / GROUP_SIZE,
                          chunk_start)))
-              : mask_offset ? 0 : chunk_len) /
+              : chunk_len) /
       (NUM_WARP_KV * num_frags_z * 16);
 
   uint32_t k_smem_offset_r =
@@ -775,23 +710,11 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
   commit_group();
 #pragma unroll 1
   for (uint32_t iter = 0; iter < num_iterations; ++iter) {
-    if constexpr (IsDynamicC8) {
-      produce_k_dynamic_scale<BLOCK_SIZE, num_frags_z, NUM_WARP_Q, T>(
-        k_smem_scale,
-        cache_k_scale_reg,
-        block_table_now,
-        cache_k_scale,
-        kv_idx_base,
-        kv_num_heads,
-        kv_head_idx,
-        chunk_end
-      );
-    }
     wait_group<1>();
     __syncthreads();
 
     // s = qk
-    compute_qk_c8<num_frags_x, num_frags_y, num_frags_z, T, CacheT, is_scale_channel_wise, IsFP8, IsDynamicC8>(
+    compute_qk_c8<num_frags_x, num_frags_y, num_frags_z, T, CacheT, is_scale_channel_wise, IsFP8>(
         &qo_smem,
         &q_smem_offset_r,
         &k_smem,
@@ -807,16 +730,12 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
              NUM_WARPS,
              num_frags_x,
              num_frags_y,
-             num_frags_z>(attn_mask ? attn_mask + batch_id * attn_mask_len *attn_mask_len : nullptr,
-                          q_base_seq_id_this_block,
+             num_frags_z>(q_base_seq_id_this_block,
                           kv_idx_base + wid * num_frags_z * 16,
                           q_len,
                           kv_len,
                           chunk_end,
-                          attn_mask_len,
-                          s_frag,
-                          mask_offset_this_seq);
-
+                          s_frag);
     }
 
     // update m,d
@@ -824,7 +743,6 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
         s_frag, o_frag, m_frag, d_frag);
     __syncthreads();
 
-    const uint32_t ori_kv_idx_base = kv_idx_base;
     kv_idx_base += NUM_WARP_KV * num_frags_z * 16;
     produce_k_blockwise_c8<SharedMemFillMode::kNoFill,
                            NUM_WARPS,
@@ -843,18 +761,6 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
                                        chunk_end,
                                        const_k_offset);
     commit_group();
-    if constexpr (IsDynamicC8) {
-      produce_v_dynamic_scale<BLOCK_SIZE, num_frags_z, NUM_WARP_Q, T>(
-        v_smem_scale,
-        cache_v_scale_reg,
-        block_table_now,
-        cache_v_scale,
-        ori_kv_idx_base,
-        kv_num_heads,
-        kv_head_idx,
-        chunk_end
-      );
-    }
     wait_group<1>();
     __syncthreads();
 
@@ -865,9 +771,7 @@ __global__ void multi_query_append_attention_c8_warp1_4_kernel(
                                   BLOCK_SIZE,
                                   T,
                                   CacheT,
-                                  is_scale_channel_wise,
-                                  IsFP8,
-                                  IsDynamicC8>(
+                                  is_scale_channel_wise, IsFP8>(
         &v_smem, &v_smem_offset_r, s_frag, o_frag, d_frag, cache_v_scale_reg);
     __syncthreads();
 
@@ -981,8 +885,7 @@ template <typename T,
           uint32_t NUM_WARP_Q,
           typename OutT = T,
           bool ENABLE_PREFILL = true,
-          bool IsFP8 = false,
-          bool IsDynamicC8 = false>
+          bool IsFP8=false>
 void MultiQueryAppendC8Attention(
     const AppendAttnMetaData &meta_data,
     const paddle::Tensor &qkv,
@@ -996,8 +899,8 @@ void MultiQueryAppendC8Attention(
     const paddle::Tensor &seq_lens_q,
     const paddle::Tensor &seq_lens_kv,
     const paddle::Tensor &seq_lens_encoder,
-    const paddle::Tensor &batch_id_per_token,
-    const paddle::Tensor &cu_seqlens_q,
+    const paddle::Tensor &padding_offsets,
+    const paddle::Tensor &cum_offsets,
     const paddle::Tensor &block_table,
     const paddle::Tensor &batch_ids,
     const paddle::Tensor &tile_ids_per_batch,
@@ -1040,8 +943,7 @@ void MultiQueryAppendC8Attention(
     constexpr uint32_t num_frags_z = BLOCK_SIZE / 16;
     constexpr uint32_t smem_size =
         num_warps * num_frags_x * 16 * HEAD_DIM * sizeof(T) +
-        num_frags_z * 16 * HEAD_DIM * sizeof(uint8_t) * 2 +
-        num_frags_z * 16 * sizeof(T) * 2;
+        num_frags_z * 16 * HEAD_DIM * sizeof(uint8_t) * 2;
     auto split_kv_kernel =
         multi_query_append_attention_c8_kernel<NV_TYPE,
                                                uint8_t,
@@ -1058,9 +960,7 @@ void MultiQueryAppendC8Attention(
                                                num_frags_y,
                                                OUT_NV_TYPE,
                                                ENABLE_PREFILL,
-                                               false,
-                                               IsFP8,
-                                               IsDynamicC8>;
+                                               false, IsFP8>;
     if (is_scale_channel_wise) {
       split_kv_kernel =
         multi_query_append_attention_c8_kernel<NV_TYPE,
@@ -1078,9 +978,7 @@ void MultiQueryAppendC8Attention(
                                                num_frags_y,
                                                OUT_NV_TYPE,
                                                ENABLE_PREFILL,
-                                               true,
-                                               IsFP8,
-                                               IsDynamicC8>;
+                                               true, IsFP8>;
     }
     if (smem_size >= 48 * 1024) {
       cudaFuncSetAttribute(split_kv_kernel,
@@ -1114,9 +1012,7 @@ void MultiQueryAppendC8Attention(
                                                  num_frags_y,
                                                  OUT_NV_TYPE,
                                                  ENABLE_PREFILL,
-                                                 false,
-                                                 IsFP8,
-                                                 IsDynamicC8>;
+                                                 false, IsFP8>;
       if (is_scale_channel_wise) {
         nosplit_kv_kernel =
           multi_query_append_attention_c8_kernel<NV_TYPE,
@@ -1134,9 +1030,7 @@ void MultiQueryAppendC8Attention(
                                                  num_frags_y,
                                                  OUT_NV_TYPE,
                                                  ENABLE_PREFILL,
-                                                 true,
-                                                 IsFP8,
-                                                 IsDynamicC8>;
+                                                 true, IsFP8>;
       }
       if (smem_size >= 48 * 1024) {
         cudaFuncSetAttribute(nosplit_kv_kernel,
@@ -1160,9 +1054,8 @@ void MultiQueryAppendC8Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1171,7 +1064,6 @@ void MultiQueryAppendC8Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           nullptr,
           nullptr,
           nullptr,
@@ -1219,9 +1111,8 @@ void MultiQueryAppendC8Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1230,7 +1121,6 @@ void MultiQueryAppendC8Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
           static_cast<float *>(tmp_m->ptr()),
           static_cast<float *>(tmp_d->ptr()),
@@ -1256,7 +1146,7 @@ void MultiQueryAppendC8Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                cu_seqlens_q.data<int>(),
+                cum_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
                                  const_cast<T *>(shift_bias.get().data<T>()))
                            : nullptr,
@@ -1291,8 +1181,7 @@ void MultiQueryAppendC8Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                batch_id_per_token.data<int>(),
-                cu_seqlens_q.data<int>(),
+                padding_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
                                  const_cast<T *>(shift_bias.get().data<T>()))
                            : nullptr,
@@ -1316,8 +1205,7 @@ void MultiQueryAppendC8Attention(
     constexpr uint32_t num_frags_z = BLOCK_SIZE / 16 / NUM_WARP_KV * 2;
     constexpr uint32_t smem_size =
         num_frags_x * 16 * HEAD_DIM * sizeof(T) +
-        NUM_WARP_KV * num_frags_z * 16 * HEAD_DIM * sizeof(uint8_t) * 2 +
-        NUM_WARP_KV * num_frags_z * 16 * sizeof(T) * 2;
+        NUM_WARP_KV * num_frags_z * 16 * HEAD_DIM * sizeof(uint8_t) * 2;
     auto split_kv_kernel =
         multi_query_append_attention_c8_warp1_4_kernel<NV_TYPE,
                                                        uint8_t,
@@ -1334,9 +1222,7 @@ void MultiQueryAppendC8Attention(
                                                        num_frags_y,
                                                        OUT_NV_TYPE,
                                                        ENABLE_PREFILL,
-                                                       false,
-                                                       IsFP8,
-                                                       IsDynamicC8>;
+                                                       false, IsFP8>;
     if (is_scale_channel_wise) {
       split_kv_kernel =
         multi_query_append_attention_c8_warp1_4_kernel<NV_TYPE,
@@ -1354,9 +1240,7 @@ void MultiQueryAppendC8Attention(
                                                        num_frags_y,
                                                        OUT_NV_TYPE,
                                                        ENABLE_PREFILL,
-                                                       true,
-                                                       IsFP8,
-                                                       IsDynamicC8>;
+                                                       true, IsFP8>;
     }
     if (smem_size >= 48 * 1024) {
       cudaFuncSetAttribute(split_kv_kernel,
@@ -1371,17 +1255,10 @@ void MultiQueryAppendC8Attention(
       chunk_size = static_cast<uint32_t>(encoder_max_partition_size);
     }
 
-    const int num_chunks = div_up(max_seq_len, chunk_size);
-    int32_t attn_mask_len;
-    if (attn_mask) {
-        attn_mask_len = attn_mask.get().shape()[1];
-    } else {
-        attn_mask_len = -1;
-    }
-
+    const int num_chunks = div_up(max_dec_len, chunk_size);
     dim3 grids(num_blocks_x_cpu, num_chunks, kv_num_heads);
     dim3 blocks(32, num_warps);
-    if (num_chunks <= 0) {
+    if (num_chunks <= 1) {
       auto nosplit_kv_kernel =
           multi_query_append_attention_c8_warp1_4_kernel<NV_TYPE,
                                                          uint8_t,
@@ -1398,9 +1275,7 @@ void MultiQueryAppendC8Attention(
                                                          num_frags_y,
                                                          OUT_NV_TYPE,
                                                          ENABLE_PREFILL,
-                                                         false,
-                                                         IsFP8,
-                                                         IsDynamicC8>;
+                                                         false, IsFP8>;
       if (is_scale_channel_wise) {
         nosplit_kv_kernel =
           multi_query_append_attention_c8_warp1_4_kernel<NV_TYPE,
@@ -1418,9 +1293,7 @@ void MultiQueryAppendC8Attention(
                                                          num_frags_y,
                                                          OUT_NV_TYPE,
                                                          ENABLE_PREFILL,
-                                                         true,
-                                                         IsFP8,
-                                                         IsDynamicC8>;
+                                                         true, IsFP8>;
       }
       if (smem_size >= 48 * 1024) {
         cudaFuncSetAttribute(nosplit_kv_kernel,
@@ -1444,11 +1317,8 @@ void MultiQueryAppendC8Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
-          attn_mask ? const_cast<bool *>(attn_mask.get().data<bool>())
-                        : nullptr,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1457,13 +1327,11 @@ void MultiQueryAppendC8Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           nullptr,
           nullptr,
           nullptr,
           reinterpret_cast<OUT_NV_TYPE *>(out->data<OutT>()),
-          speculate_max_draft_token_num,
-          attn_mask_len);
+          speculate_max_draft_token_num);
     } else {
       phi::Allocator::AllocationPtr tmp_workspace, tmp_m, tmp_d;
       if (is_decoder) {
@@ -1510,8 +1378,8 @@ void MultiQueryAppendC8Attention(
           reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_k_scale.data<T>())),
           reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_v_scale.data<T>())),
           shift_bias ? reinterpret_cast<NV_TYPE *>(
-                            const_cast<T *>(shift_bias.get().data<T>()))
-                      : nullptr,
+                           const_cast<T *>(shift_bias.get().data<T>()))
+                     : nullptr,
           smooth_weight ? reinterpret_cast<NV_TYPE *>(
                               const_cast<T *>(smooth_weight.get().data<T>()))
                         : nullptr,
@@ -1519,11 +1387,8 @@ void MultiQueryAppendC8Attention(
           seq_lens_kv.data<int>(),
           batch_ids.data<int>(),
           tile_ids_per_batch.data<int>(),
-          cu_seqlens_q.data<int>(),
+          cum_offsets.data<int>(),
           block_table.data<int>(),
-          meta_data.mask_offset,
-          attn_mask ? const_cast<bool *>(attn_mask.get().data<bool>())
-                        : nullptr,
           max_seq_len,
           max_dec_len,
           max_block_num_per_seq,
@@ -1532,13 +1397,11 @@ void MultiQueryAppendC8Attention(
           quant_min_bound,
           in_scale,
           chunk_size,
-          num_blocks_x_cpu,
           reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
           static_cast<float *>(tmp_m->ptr()),
           static_cast<float *>(tmp_d->ptr()),
           reinterpret_cast<OUT_NV_TYPE *>(out->data<OutT>()),
-          speculate_max_draft_token_num,
-          attn_mask_len);
+          speculate_max_draft_token_num);
       // merge
       constexpr int vec_size = num_elems_per_128b<NV_TYPE>();
       if (is_decoder) {
@@ -1554,10 +1417,10 @@ void MultiQueryAppendC8Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                cu_seqlens_q.data<int>(),
+                cum_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
-                                  const_cast<T *>(shift_bias.get().data<T>()))
-                            : nullptr,
+                                 const_cast<T *>(shift_bias.get().data<T>()))
+                           : nullptr,
                 smooth_weight ? reinterpret_cast<NV_TYPE *>(const_cast<T *>(
                                     smooth_weight.get().data<T>()))
                               : nullptr,
@@ -1574,14 +1437,14 @@ void MultiQueryAppendC8Attention(
         constexpr int blockx = HEAD_DIM / vec_size;
         constexpr int blocky = (128 + blockx - 1) / blockx;
         dim3 grids_merge(min(sm_count * 4, token_num),
-                          num_heads);
+                         num_heads);
         dim3 blocks_merge(blockx, blocky);
         merge_multi_chunks_v2_kernel<NV_TYPE,
-                                      vec_size,
-                                      blocky,
-                                      HEAD_DIM,
-                                      OUT_NV_TYPE,
-                                      ENABLE_PREFILL>
+                                     vec_size,
+                                     blocky,
+                                     HEAD_DIM,
+                                     OUT_NV_TYPE,
+                                     ENABLE_PREFILL>
             <<<grids_merge, blocks_merge, 0, stream>>>(
                 reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
                 static_cast<float *>(tmp_m->ptr()),
@@ -1589,11 +1452,10 @@ void MultiQueryAppendC8Attention(
                 seq_lens_q.data<int>(),
                 seq_lens_kv.data<int>(),
                 seq_lens_encoder.data<int>(),
-                batch_id_per_token.data<int>(),
-                cu_seqlens_q.data<int>(),
+                padding_offsets.data<int>(),
                 shift_bias ? reinterpret_cast<NV_TYPE *>(
-                                  const_cast<T *>(shift_bias.get().data<T>()))
-                            : nullptr,
+                                 const_cast<T *>(shift_bias.get().data<T>()))
+                           : nullptr,
                 smooth_weight ? reinterpret_cast<NV_TYPE *>(const_cast<T *>(
                                     smooth_weight.get().data<T>()))
                               : nullptr,
@@ -1637,8 +1499,8 @@ void CascadeAppendAttentionC8Kernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -1655,7 +1517,6 @@ void CascadeAppendAttentionC8Kernel(
     const bool causal,
     const bool is_decoder,
     const bool enable_prefill,
-    const std::string& cache_quant_type_str,
     cudaStream_t& stream,
     paddle::Tensor* out) {
   const auto token_num = meta_data.token_nums;
@@ -1664,7 +1525,6 @@ void CascadeAppendAttentionC8Kernel(
   const auto num_heads = meta_data.q_num_heads;
   const auto group_size = meta_data.q_num_heads / meta_data.kv_num_heads;
   const auto head_dim = meta_data.head_dims;
-  bool is_dynamic_cfp8 = cache_quant_type_str == "block_wise_fp8";
 
   DISPATCH_CAUSAL(
       causal,
@@ -1683,46 +1543,43 @@ void CascadeAppendAttentionC8Kernel(
                       BLOCK_SIZE,
                       {DISPATCH_BLOCKSHAPE_Q(
                           block_shape_q, BLOCK_SHAPE_Q, NUM_WARP_Q, {
-                            DISPATCH_DyCfp8(is_dynamic_cfp8, IsDynamicC8, {
-                              MultiQueryAppendC8Attention<T,
-                                                          GROUP_SIZE,
-                                                          HEAD_DIM,
-                                                          BLOCK_SIZE,
-                                                          CAUSAL,
-                                                          BLOCK_SHAPE_Q,
-                                                          NUM_WARP_Q,
-                                                          OutT,
-                                                          ENABLE_PREFILL,
-                                                          IsFP8,
-                                                          IsDynamicC8>(
-                                  meta_data,
-                                  qkv,
-                                  cache_k,
-                                  cache_v,
-                                  attn_mask,
-                                  cache_k_scale.get(),
-                                  cache_v_scale.get(),
-                                  shift_bias,
-                                  smooth_weight,
-                                  seq_lens_q,
-                                  seq_lens_kv,
-                                  seq_lens_encoder,
-                                  batch_id_per_token,
-                                  cu_seqlens_q,
-                                  block_table,
-                                  batch_ids,
-                                  tile_ids_per_batch,
-                                  num_blocks,
-                                  max_seq_len,
-                                  max_dec_len,
-                                  quant_max_bound,
-                                  quant_min_bound,
-                                  in_scale,
-                                  max_partition_size,
-                                  encoder_max_partition_size,
-                                  speculate_max_draft_token_num,
-                                  is_decoder,
-                                  stream,
-                                  out);
-      })})})})})})})
+                            MultiQueryAppendC8Attention<T,
+                                                        GROUP_SIZE,
+                                                        HEAD_DIM,
+                                                        BLOCK_SIZE,
+                                                        CAUSAL,
+                                                        BLOCK_SHAPE_Q,
+                                                        NUM_WARP_Q,
+                                                        OutT,
+                                                        ENABLE_PREFILL, IsFP8>(
+                                meta_data,
+                                qkv,
+                                cache_k,
+                                cache_v,
+                                attn_mask,
+                                cache_k_scale.get(),
+                                cache_v_scale.get(),
+                                shift_bias,
+                                smooth_weight,
+                                seq_lens_q,
+                                seq_lens_kv,
+                                seq_lens_encoder,
+                                padding_offsets,
+                                cum_offsets,
+                                block_table,
+                                batch_ids,
+                                tile_ids_per_batch,
+                                num_blocks,
+                                max_seq_len,
+                                max_dec_len,
+                                quant_max_bound,
+                                quant_min_bound,
+                                in_scale,
+                                max_partition_size,
+                                encoder_max_partition_size,
+                                speculate_max_draft_token_num,
+                                is_decoder,
+                                stream,
+                                out);
+                          })})})})})})
 }

custom_ops/gpu_ops/append_attn/append_attention_func.cuh

@@ -384,113 +384,6 @@ __device__ __forceinline__ void produce_v_blockwise_c8(
   }
 }
 
-template<uint32_t block_size,
-         uint32_t num_frags_z,
-         uint32_t NUM_WARP_Q,
-         typename T>
-__device__ __forceinline__ void produce_k_dynamic_scale(
-  T* k_smem_scale,
-  T* cache_k_reg,
-  const int* block_table_now,
-  const T* cache_k_scale,
-  const uint32_t kv_idx,
-  const uint32_t kv_num_heads,
-  const uint32_t kv_head_idx,
-  const uint32_t chunk_end
-) {
-  const uint32_t tx = threadIdx.x, ty = threadIdx.y;
-  if constexpr (NUM_WARP_Q == 4) {
-    // 4 warps shared block_size
-    const uint32_t tid = ty * 32 + tx;
-    int block_id = __ldg(&block_table_now[kv_idx / block_size]);
-    if (block_id < 0) block_id = 0;
-    const T* cache_k_scale_now = cache_k_scale + block_id * kv_num_heads * block_size + kv_head_idx * block_size;
-    if (tid < block_size) {
-      k_smem_scale[tid] = cache_k_scale_now[tid];
-    }
-    __syncthreads();
-    const uint32_t row_id = tx / 4;
-    for (uint32_t fz = 0; fz < num_frags_z; fz++) {
-      cache_k_reg[fz * 2] = k_smem_scale[fz * 16 + row_id];
-      cache_k_reg[fz * 2 + 1] = k_smem_scale[fz * 16 + row_id + 8];
-    }
-  } else {
-    // 1 warp 32 tokens
-    const uint32_t kv_idx_now = kv_idx + block_size * ty / 2;
-    int block_id = __ldg(&block_table_now[kv_idx_now / block_size]);
-    if (block_id < 0) block_id = 0;
-    const T* cache_k_scale_now = cache_k_scale + block_id * kv_num_heads * block_size + kv_head_idx * block_size;
-    const int kv_idx_this_thread = kv_idx + ty * 32 + tx;
-    if (kv_idx_this_thread < chunk_end) {
-      k_smem_scale[ty * 32 + tx] = cache_k_scale_now[(ty  % 2) * 32 + tx];
-    } else {
-      k_smem_scale[ty * 32 + tx] = 0;
-    }
-    __syncwarp();
-    const uint32_t row_id = tx / 4;
-    for (uint32_t fz = 0; fz < num_frags_z; fz++) {
-      cache_k_reg[fz * 2] = k_smem_scale[ty * 32 + fz * 16 + row_id];
-      cache_k_reg[fz * 2 + 1] = k_smem_scale[ty * 32 + fz * 16 + row_id + 8];
-    }
-  }
-}
-
-template<uint32_t block_size,
-         uint32_t num_frags_z,
-         uint32_t NUM_WARP_Q,
-         typename T>
-__device__ __forceinline__ void produce_v_dynamic_scale(
-  T* v_smem_scale,
-  T* cache_v_reg,
-  const int* block_table_now,
-  const T* cache_v_scale,
-  const uint32_t kv_idx,
-  const uint32_t kv_num_heads,
-  const uint32_t kv_head_idx,
-  const uint32_t chunk_end
-) {
-  const uint32_t tx = threadIdx.x, ty = threadIdx.y;
-
-  if constexpr (NUM_WARP_Q == 4) {
-    // 4 warps shared block_size
-    const uint32_t tid = ty * 32 + tx;
-    int block_id = __ldg(&block_table_now[kv_idx / block_size]);
-    if (block_id < 0) block_id = 0;
-    const T* cache_v_scale_now = cache_v_scale + block_id * kv_num_heads * block_size + kv_head_idx * block_size;
-    if (tid < block_size) {
-      v_smem_scale[tid] = cache_v_scale_now[tid];
-    }
-    __syncthreads();
-    const uint32_t row_id = tx % 4 * 2;
-    for (uint32_t fz = 0; fz < num_frags_z; fz++) {
-      cache_v_reg[fz * 4] = v_smem_scale[fz * 16 + row_id];
-      cache_v_reg[fz * 4 + 1] = v_smem_scale[fz * 16 + row_id + 1];
-      cache_v_reg[fz * 4 + 2] = v_smem_scale[fz * 16 + row_id + 8];
-      cache_v_reg[fz * 4 + 3] = v_smem_scale[fz * 16 + row_id + 9];
-    }
-  } else {
-    // 1 warp 32 tokens
-    const uint32_t kv_idx_now = kv_idx + block_size * ty / 2;
-    int block_id = __ldg(&block_table_now[kv_idx_now / block_size]);
-    if (block_id < 0) block_id = 0;
-    const T* cache_v_scale_now = cache_v_scale + block_id * kv_num_heads * block_size + kv_head_idx * block_size;
-    const int kv_idx_this_thread = kv_idx + ty * 32 + tx;
-    if (kv_idx_this_thread < chunk_end) {
-      v_smem_scale[ty * 32 + tx] = cache_v_scale_now[(ty  % 2) * 32 + tx];
-    } else {
-      v_smem_scale[ty * 32 + tx] = 0;
-    }
-    __syncwarp();
-    const uint32_t row_id = tx % 4 * 2;
-    for (uint32_t fz = 0; fz < num_frags_z; fz++) {
-      cache_v_reg[fz * 4] = v_smem_scale[ty * 32 + fz * 16 + row_id];
-      cache_v_reg[fz * 4 + 1] = v_smem_scale[ty * 32 + fz * 16 + row_id + 1];
-      cache_v_reg[fz * 4 + 2] = v_smem_scale[ty * 32 + fz * 16 + row_id + 8];
-      cache_v_reg[fz * 4 + 3] = v_smem_scale[ty * 32 + fz * 16 + row_id + 9];
-    }
-  }
-}
-
 template <SharedMemFillMode fill_mode,
           uint32_t num_warps,
           uint32_t block_size,
@@ -923,8 +816,7 @@ template <uint32_t num_frags_x,
           typename T,
           typename CacheT,
           bool is_scale_channel_wise = false,
-          bool IsFP8 = false,
-          bool IsDynamicC8 = false>
+          bool IsFP8=false>
 __device__ __forceinline__ void compute_qk_c8(smem_t* q_smem,
                                               uint32_t* q_smem_offset_r,
                                               smem_t* k_smem,
@@ -968,27 +860,20 @@ __device__ __forceinline__ void compute_qk_c8(smem_t* q_smem,
         convert_c8<T,IsFP8>(b_frag_dq_T, b_frag[fy * 2]);
         convert_c8<T,IsFP8>(b_frag_dq_T + 4, b_frag[fy * 2 + 1]);
         // scale zp
-        if constexpr (!IsDynamicC8) {
-          if constexpr (is_scale_channel_wise) {
-            const int scale_col = (ky * 2 + fy) * 4;
-            b_frag_dq_T[0] *= cache_k_scale[scale_col];
-            b_frag_dq_T[1] *= cache_k_scale[scale_col + 1];
-            b_frag_dq_T[2] *= cache_k_scale[scale_col + 2];
-            b_frag_dq_T[3] *= cache_k_scale[scale_col + 3];
-            b_frag_dq_T[4] *= cache_k_scale[scale_col];
-            b_frag_dq_T[5] *= cache_k_scale[scale_col + 1];
-            b_frag_dq_T[6] *= cache_k_scale[scale_col + 2];
-            b_frag_dq_T[7] *= cache_k_scale[scale_col + 3];
-          } else {
-#pragma unroll
-            for (uint32_t b_i = 0; b_i < 8; ++b_i) {
-              b_frag_dq_T[b_i] *= cache_k_scale[0];
-            }
-          }
+        if constexpr (is_scale_channel_wise) {
+          const int scale_col = (ky * 2 + fy) * 4;
+          b_frag_dq_T[0] *= cache_k_scale[scale_col];
+          b_frag_dq_T[1] *= cache_k_scale[scale_col + 1];
+          b_frag_dq_T[2] *= cache_k_scale[scale_col + 2];
+          b_frag_dq_T[3] *= cache_k_scale[scale_col + 3];
+          b_frag_dq_T[4] *= cache_k_scale[scale_col];
+          b_frag_dq_T[5] *= cache_k_scale[scale_col + 1];
+          b_frag_dq_T[6] *= cache_k_scale[scale_col + 2];
+          b_frag_dq_T[7] *= cache_k_scale[scale_col + 3];
         } else {
 #pragma unroll
           for (uint32_t b_i = 0; b_i < 8; ++b_i) {
-            b_frag_dq_T[b_i] *= cache_k_scale[fz * 2 + b_i / 4];
+            b_frag_dq_T[b_i] *= cache_k_scale[0];
           }
         }
 #pragma unroll
@@ -1020,15 +905,12 @@ template <typename T,
           uint32_t num_frags_y,
           uint32_t num_frags_z,
           bool IS_SYSTEM = false>
-__device__ __forceinline__ void mask_s(const bool* attn_mask,
-                                       const uint32_t qo_idx_base,
+__device__ __forceinline__ void mask_s(const uint32_t qo_idx_base,
                                        const uint32_t kv_idx_base,
                                        const uint32_t qo_len,
                                        const uint32_t kv_len,
                                        const uint32_t chunk_end,
-                                       const int32_t attn_mask_len,
-                                       float (*s_frag)[num_frags_z][8],
-                                       const int *mask_offset = nullptr) {
+                                       float (*s_frag)[num_frags_z][8]) {
   const uint32_t tx = threadIdx.x;
 #pragma unroll
   for (uint32_t fx = 0; fx < num_frags_x; ++fx) {
@@ -1042,21 +924,10 @@ __device__ __forceinline__ void mask_s(const bool* attn_mask,
                                  group_size,
                          kv_idx = kv_idx_base + fz * 16 + 2 * (tx % 4) +
                                   8 * (reg_id / 4) + reg_id % 2;
-          bool out_of_boundary;
-          if (mask_offset) {
-            out_of_boundary = q_idx < qo_len ? (kv_idx >= mask_offset[q_idx * 2 + 1] || kv_idx < mask_offset[q_idx * 2]) : true;
-          } else {
-            out_of_boundary =
-                (causal
-                    ? (kv_idx > kv_len + q_idx - qo_len || (kv_idx >= chunk_end))
-                    : kv_idx >= chunk_end);
-            if (attn_mask != nullptr && kv_idx > kv_len - qo_len && kv_idx < chunk_end && attn_mask_len > 0 && q_idx < static_cast<uint32_t>(attn_mask_len)) {
-              const int32_t mask_idx = q_idx * attn_mask_len + kv_idx - kv_len + qo_len;
-              bool mask = attn_mask[mask_idx];
-              out_of_boundary |= mask;
-            }
-          }
-
+          const bool out_of_boundary =
+              (causal
+                   ? (kv_idx > kv_len + q_idx - qo_len || (kv_idx >= chunk_end))
+                   : kv_idx >= chunk_end);
           if constexpr (std::is_same<T, half>::value) {
             s_frag[fx][fz][reg_id] =
                 out_of_boundary ? -5e4f : s_frag[fx][fz][reg_id];
@@ -1064,7 +935,6 @@ __device__ __forceinline__ void mask_s(const bool* attn_mask,
             s_frag[fx][fz][reg_id] =
                 out_of_boundary ? -3.0e+30f : s_frag[fx][fz][reg_id];
           }
-          // printf("tid: %d. qk[%u,%u] = %f, mask: %d \n ", threadIdx.x, kv_idx, q_idx, static_cast<float>(s_frag[fx][fz][reg_id]), int(out_of_boundary));
         } else {
           const uint32_t q_idx = qo_idx_base,
                          kv_idx = kv_idx_base + fz * 16 + 2 * (tx % 4) +
@@ -1208,9 +1078,7 @@ template <uint32_t num_frags_x,
           uint32_t block_size,
           typename T,
           typename CacheT,
-          bool is_scale_channel_wise = false,
-          bool IsFP8 = false,
-          bool IsDynamicC8 = false>
+          bool is_scale_channel_wise = false, bool IsFP8=false>
 __device__ __forceinline__ void compute_sfm_v_c8(
     smem_t* v_smem,
     uint32_t* v_smem_offset_r,
@@ -1252,28 +1120,16 @@ __device__ __forceinline__ void compute_sfm_v_c8(
         convert_c8<T,IsFP8>(b_frag_dq_T, b_frag[fz * 2]);
         convert_c8<T,IsFP8>(b_frag_dq_T + 4, b_frag[fz * 2 + 1]);
         // scale zp
-        if constexpr (!IsDynamicC8) {
-          if constexpr (is_scale_channel_wise) {
+        if constexpr (is_scale_channel_wise) {
 #pragma unroll
-            for (uint32_t b_i = 0; b_i < 8; ++b_i) {
-              b_frag_dq_T[b_i] *= cache_v_scale[b_i / 4 + fy * 2];
-            }
-          } else {
-#pragma unroll
-            for (uint32_t b_i = 0; b_i < 8; ++b_i) {
-              b_frag_dq_T[b_i] *= cache_v_scale[0];
-            }
+          for (uint32_t b_i = 0; b_i < 8; ++b_i) {
+            b_frag_dq_T[b_i] *= cache_v_scale[b_i / 4 + fy * 2];
           }
         } else {
-          const int scale_col = (kz * 2 + fz) * 4;
-          b_frag_dq_T[0] *= cache_v_scale[scale_col];
-          b_frag_dq_T[1] *= cache_v_scale[scale_col + 1];
-          b_frag_dq_T[2] *= cache_v_scale[scale_col + 2];
-          b_frag_dq_T[3] *= cache_v_scale[scale_col + 3];
-          b_frag_dq_T[4] *= cache_v_scale[scale_col];
-          b_frag_dq_T[5] *= cache_v_scale[scale_col + 1];
-          b_frag_dq_T[6] *= cache_v_scale[scale_col + 2];
-          b_frag_dq_T[7] *= cache_v_scale[scale_col + 3];
+#pragma unroll
+          for (uint32_t b_i = 0; b_i < 8; ++b_i) {
+            b_frag_dq_T[b_i] *= cache_v_scale[0];
+          }
         }
 #pragma unroll
         for (uint32_t fx = 0; fx < num_frags_x; ++fx) {  // m: num_frags_x * 16
@@ -1300,9 +1156,7 @@ template <uint32_t num_frags_x,
           uint32_t block_size,
           typename T,
           typename CacheT,
-          bool is_scale_channel_wise = false,
-          bool IsFP8 = false,
-          bool IsDynamicC8 = false>
+          bool is_scale_channel_wise = false, bool IsFP8=false>
 __device__ __forceinline__ void compute_sfm_v_c8_iter_sq_bvec(
     smem_t* v_smem,
     uint32_t* v_smem_offset_r,
@@ -1346,28 +1200,16 @@ __device__ __forceinline__ void compute_sfm_v_c8_iter_sq_bvec(
         convert_c8<T,IsFP8>(b_frag_dq_T, b_frag[fz * 2]);
         convert_c8<T,IsFP8>(b_frag_dq_T + 4, b_frag[fz * 2 + 1]);
         // scale zp
-        if constexpr (!IsDynamicC8) {
-          if constexpr (is_scale_channel_wise) {
+        if constexpr (is_scale_channel_wise) {
 #pragma unroll
-            for (uint32_t b_i = 0; b_i < 8; ++b_i) {
-              b_frag_dq_T[b_i] *= cache_v_scale[b_i / 4 + fy * 2];
-            }
-          } else {
-            #pragma unroll
-            for (uint32_t b_i = 0; b_i < 8; ++b_i) {
-              b_frag_dq_T[b_i] *= cache_v_scale[0];
-            }
+          for (uint32_t b_i = 0; b_i < 8; ++b_i) {
+            b_frag_dq_T[b_i] *= cache_v_scale[b_i / 4 + fy * 2];
           }
         } else {
-          const int scale_col = (kz * 2 + fz) * 4;
-          b_frag_dq_T[0] *= cache_v_scale[scale_col];
-          b_frag_dq_T[1] *= cache_v_scale[scale_col + 1];
-          b_frag_dq_T[2] *= cache_v_scale[scale_col + 2];
-          b_frag_dq_T[3] *= cache_v_scale[scale_col + 3];
-          b_frag_dq_T[4] *= cache_v_scale[scale_col];
-          b_frag_dq_T[5] *= cache_v_scale[scale_col + 1];
-          b_frag_dq_T[6] *= cache_v_scale[scale_col + 2];
-          b_frag_dq_T[7] *= cache_v_scale[scale_col + 3];
+          #pragma unroll
+          for (uint32_t b_i = 0; b_i < 8; ++b_i) {
+            b_frag_dq_T[b_i] *= cache_v_scale[0];
+          }
         }
 #pragma unroll
         for (uint32_t fx = 0; fx < num_frags_x; ++fx) {  // m: num_frags_x * 16
@@ -2010,7 +1852,7 @@ __global__ void merge_multi_chunks_kernel(
     const float* __restrict__ multi_d,  // [token_num, num_chunks, num_heads]
     const int* __restrict__ seq_lens_q,
     const int* __restrict__ seq_lens_kv,
-    const int* __restrict__ batch_id_per_token,
+    const int* __restrict__ padding_offsets,
     const T* __restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
     const T* __restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
     T* __restrict__ out,
@@ -2024,7 +1866,8 @@ __global__ void merge_multi_chunks_kernel(
     const int head_dim) {
   const int vid = threadIdx.x, hid = threadIdx.y;
   const int qid = blockIdx.x;
-  const uint32_t bid = batch_id_per_token[qid];
+  const uint32_t ori_token_id = qid + padding_offsets[qid];
+  const uint32_t bid = ori_token_id / max_seq_len;
   if (seq_lens_q[bid] <= 0 || seq_lens_kv[bid] <= 0) {
     return;
   }
@@ -2268,7 +2111,7 @@ __global__ void merge_multi_chunks_decoder_kernel(
     const int *__restrict__ seq_lens_q,
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ seq_lens_encoder,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ cum_offsets,
     const T *__restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
     const T *__restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
     OutT *__restrict__ out,
@@ -2284,7 +2127,7 @@ __global__ void merge_multi_chunks_decoder_kernel(
   const int bid = blockIdx.x, hid = blockIdx.y;
   __shared__ T smem[bdy * HEAD_DIM];
   __shared__ float md_smem[bdy * 2];
-  const int start_token_idx = cu_seqlens_q[bid];
+  const int start_token_idx = bid * max_seq_len - cum_offsets[bid];
   const int seq_len_q = seq_lens_q[bid];
   if (seq_len_q == 0) return;
   int seq_len_kv = seq_lens_kv[bid];
@@ -2397,8 +2240,7 @@ __global__ void merge_multi_chunks_v2_kernel(
     const int *__restrict__ seq_lens_q,
     const int *__restrict__ seq_lens_kv,
     const int *__restrict__ seq_lens_encoder,
-    const int *__restrict__ batch_id_per_token,
-    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ padding_offsets,
     const T *__restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
     const T *__restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
     OutT *__restrict__ out,
@@ -2417,8 +2259,9 @@ __global__ void merge_multi_chunks_v2_kernel(
   __shared__ T smem[bdy * HEAD_DIM];
   __shared__ float md_smem[bdy * 2];
   for (int qid = blockIdx.x; qid < token_num; qid += gridDim.x) {
-    const uint32_t bid = batch_id_per_token[qid];
-    const uint32_t local_seq_id = qid - cu_seqlens_q[bid];
+    const uint32_t ori_token_id = qid + padding_offsets[qid];
+    const uint32_t bid = ori_token_id / max_seq_len;
+    const uint32_t local_seq_id = ori_token_id % max_seq_len;
     const int seq_len_q = seq_lens_q[bid];
     if (seq_len_q == 0) continue;
     int seq_len_kv = seq_lens_kv[bid];

custom_ops/gpu_ops/append_attn/append_attention_kernel.h

@@ -40,8 +40,8 @@ void CascadeAppendAttentionC16Kernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -85,8 +85,8 @@ void CascadeAppendAttentionC8Kernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -103,7 +103,6 @@ void CascadeAppendAttentionC8Kernel(
     const bool causal,
     const bool is_decoder,
     const bool enable_prefill,
-    const std::string& cache_quant_type_str,
     cudaStream_t& stream,
     paddle::Tensor* out);
 
@@ -131,8 +130,8 @@ void CascadeAppendAttentionC4Kernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -176,8 +175,8 @@ void CascadeAppendAttentionKernel(
     const paddle::Tensor& seq_lens_q,
     const paddle::Tensor& seq_lens_kv,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_table,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids_per_batch,
@@ -212,8 +211,8 @@ void CascadeAppendAttentionKernel(
                                                 seq_lens_q,
                                                 seq_lens_kv,
                                                 seq_lens_encoder,
-                                                batch_id_per_token,
-                                                cu_seqlens_q,
+                                                padding_offsets,
+                                                cum_offsets,
                                                 block_table,
                                                 batch_ids,
                                                 tile_ids_per_batch,
@@ -247,8 +246,8 @@ void CascadeAppendAttentionKernel(
                                                 seq_lens_q,
                                                 seq_lens_kv,
                                                 seq_lens_encoder,
-                                                batch_id_per_token,
-                                                cu_seqlens_q,
+                                                padding_offsets,
+                                                cum_offsets,
                                                 block_table,
                                                 batch_ids,
                                                 tile_ids_per_batch,
@@ -265,10 +264,9 @@ void CascadeAppendAttentionKernel(
                                                 causal,
                                                 is_decoder,
                                                 enable_prefill,
-                                                cache_quant_type_str,
                                                 stream,
                                                 out);
-    } else if (cache_quant_type_str == "cache_fp8" or cache_quant_type_str == "block_wise_fp8") {
+    } else if (cache_quant_type_str == "cache_fp8") {
         CascadeAppendAttentionC8Kernel<T, OutT, true>(meta_data,
                                                 qkv,
                                                 cache_k,
@@ -283,8 +281,8 @@ void CascadeAppendAttentionKernel(
                                                 seq_lens_q,
                                                 seq_lens_kv,
                                                 seq_lens_encoder,
-                                                batch_id_per_token,
-                                                cu_seqlens_q,
+                                                padding_offsets,
+                                                cum_offsets,
                                                 block_table,
                                                 batch_ids,
                                                 tile_ids_per_batch,
@@ -301,7 +299,6 @@ void CascadeAppendAttentionKernel(
                                                 causal,
                                                 is_decoder,
                                                 enable_prefill,
-                                                cache_quant_type_str,
                                                 stream,
                                                 out);
     } else if (cache_quant_type_str == "cache_int4_zp") {
@@ -319,8 +316,8 @@ void CascadeAppendAttentionKernel(
                                                 seq_lens_q,
                                                 seq_lens_kv,
                                                 seq_lens_encoder,
-                                                batch_id_per_token,
-                                                cu_seqlens_q,
+                                                padding_offsets,
+                                                cum_offsets,
                                                 block_table,
                                                 batch_ids,
                                                 tile_ids_per_batch,

custom_ops/gpu_ops/append_attn/decode_attention_kernel.cu

@@ -35,7 +35,7 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
                                                     const T * __restrict__ multi_d, // [bsz, num_chunks, num_heads]
                                                     const int * __restrict__ seq_lens_q,
                                                     const int * __restrict__ seq_lens_kv,
-                                                    const int * __restrict__ cu_seqlens_q,
+                                                    const int * __restrict__ cum_offsets,
                                                     const T * __restrict__ shift_bias, // [q_num_heads * HEAD_DIM]
                                                     const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
                                                     OutT * __restrict__ out, // [token_num, num_heads, head_dim]
@@ -59,7 +59,7 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
   __shared__ T smem[bdy * HEAD_DIM];
   __shared__ T md_smem[bdy * 2];
 
-  const int start_token_ids = cu_seqlens_q[qid];
+  const int start_token_ids = qid * max_seq_len - __ldg(&cum_offsets[qid]);
   using LoadT = AlignedVector<T, vec_size>;
   LoadT load_vec;
   LoadT res_vec;
@@ -134,7 +134,7 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
                                                     const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
                                                     const int * __restrict__ seq_lens_q,
                                                     const int * __restrict__ seq_lens_kv,
-                                                    const int * __restrict__ cu_seqlens_q,
+                                                    const int * __restrict__ cum_offsets,
                                                     const int * __restrict__ block_table, // [bsz, block_num_per_seq]
                                                     const int max_seq_len,
                                                     const int max_dec_len,
@@ -171,8 +171,8 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   }
   kv_len += q_len;
   const uint32_t num_chunk_this_seq = div_up(kv_len, chunk_size);
-  const uint32_t q_start_idx = cu_seqlens_q[bid];
-  const uint32_t q_write_idx = cu_seqlens_q[bid];
+  const uint32_t q_start_idx = bid * max_seq_len - __ldg(&cum_offsets[bid]);
+  const uint32_t q_write_idx = bid * max_seq_len - __ldg(&cum_offsets[bid]);
   if (chunk_id >= num_chunk_this_seq) {
     return;
   }
@@ -317,8 +317,8 @@ void MultiQueryDecoderAttention(
   const paddle::optional<paddle::Tensor>& smooth_weight,
   const paddle::Tensor &seq_lens_q,
   const paddle::Tensor &seq_lens_kv,
-  const paddle::Tensor &batch_id_per_token,
-  const paddle::Tensor &cu_seqlens_q,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
   const paddle::Tensor &block_table,
   const int max_seq_len,
   const int max_dec_len,
@@ -393,7 +393,7 @@ void MultiQueryDecoderAttention(
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
       seq_lens_q.data<int>(),
       seq_lens_kv.data<int>(),
-      cu_seqlens_q.data<int>(),
+      cum_offsets.data<int>(),
       block_table.data<int>(),
       max_seq_len,
       max_dec_len,
@@ -430,7 +430,7 @@ void MultiQueryDecoderAttention(
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
       seq_lens_q.data<int>(),
       seq_lens_kv.data<int>(),
-      cu_seqlens_q.data<int>(),
+      cum_offsets.data<int>(),
       block_table.data<int>(),
       max_seq_len,
       max_dec_len,
@@ -456,7 +456,7 @@ void MultiQueryDecoderAttention(
       reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
       seq_lens_q.data<int>(),
       seq_lens_kv.data<int>(),
-      cu_seqlens_q.data<int>(),
+      cum_offsets.data<int>(),
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>())),
@@ -483,8 +483,8 @@ void DecodeMLAAttentionKernel(
   const paddle::optional<paddle::Tensor>& smooth_weight,
   const paddle::Tensor &seq_lens_q, // q_seq_len is 1
   const paddle::Tensor &seq_lens_kv,
-  const paddle::Tensor &batch_id_per_token,
-  const paddle::Tensor &cu_seqlens_q,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
   const paddle::Tensor &block_table,
   int max_seq_len,
   int max_dec_len,
@@ -513,7 +513,7 @@ void DecodeMLAAttentionKernel(
           {DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE,
               {DISPATCH_DEAL_EACH_TIME(deal_each_time, DEAL_EACH_TIME,
                   {MultiQueryDecoderAttention<T, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V, BLOCK_SIZE, CAUSAL, 2, 16, DEAL_EACH_TIME>(
-                  meta_data, stream, q, cache_k, cache_v, attn_mask, shift_bias, smooth_weight, seq_lens_q, seq_lens_kv, batch_id_per_token, cu_seqlens_q,
+                  meta_data, stream, q, cache_k, cache_v, attn_mask, shift_bias, smooth_weight, seq_lens_q, seq_lens_kv, padding_offsets, cum_offsets,
                   block_table, max_seq_len, max_dec_len, rope_scale, rope_theta, softmax_scale, in_scale, out);})})})})})});
 }
 
@@ -527,8 +527,8 @@ template void DecodeMLAAttentionKernel<paddle::bfloat16>(
   const paddle::optional<paddle::Tensor>& smooth_weight,
   const paddle::Tensor &seq_lens_q, // q_seq_len is 1
   const paddle::Tensor &seq_lens_kv,
-  const paddle::Tensor &batch_id_per_token,
-  const paddle::Tensor &cu_seqlens_q,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
   const paddle::Tensor &block_table,
   int max_seq_len,
   int max_dec_len,
@@ -548,8 +548,8 @@ template void DecodeMLAAttentionKernel<paddle::float16>(
   const paddle::optional<paddle::Tensor>& smooth_weight,
   const paddle::Tensor &seq_lens_q, // q_seq_len is 1
   const paddle::Tensor &seq_lens_kv,
-  const paddle::Tensor &batch_id_per_token,
-  const paddle::Tensor &cu_seqlens_q,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
   const paddle::Tensor &block_table,
   int max_seq_len,
   int max_dec_len,

custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_kernel.cu

@@ -15,74 +15,14 @@
 #include "decoder_write_cache_with_rope_kernel.h"
 #include "utils.cuh"
 
-template <typename T, typename QKV_TYPE>
-void append_decode_cache_rope_qk_norm(const QKV_TYPE* qkv,
-                              T* key_cache,
-                              T* value_cache,
-                              T* qkv_out,
-                              const int* block_tables,
-                              const int* cu_seqlens_q,
-                              const int* seq_lens,
-                              const int* seq_lens_encoder,
-                              const float* cos_emb,
-                              const float* sin_emb,
-                              const float* qkv_out_scales,
-                              const T* qkv_biases,
-                              const int max_seq_len,
-                              const int max_blocks_per_seq,
-                              const int num_heads,
-                              const int kv_num_heads,
-                              const int dim_head,
-                              const int block_size,
-                              const int bsz,
-                              const cudaStream_t& stream,
-                              const bool use_neox_style,
-                              const bool rope_3d,
-                              const float* q_norm_weight,
-                              const float* k_norm_weight,
-                              const float rms_norm_eps) {
-  const uint32_t elem_nums =
-      use_neox_style ? bsz * (num_heads + 2 * kv_num_heads) * dim_head / 2
-                     : bsz * (num_heads + 2 * kv_num_heads) * dim_head;
-  constexpr int HEAD_DIM = 128;
-
-  constexpr int PackSize = HEAD_DIM / kWarpSize;
-  const int pack_num = elem_nums / PackSize;
-  const int blocksize = 128;
-  int grid_size = 1;
-  GetNumBlocks<128>(pack_num, &grid_size);
-  dim3 block_dim(kWarpSize, blocksize / kWarpSize, 1);
-  append_decode_cache_T_rope_qk_norm_kernel<T, PackSize>
-      <<<grid_size, block_dim, 0, stream>>>(reinterpret_cast<const T*>(qkv),
-                                            key_cache,
-                                            value_cache,
-                                            qkv_out,
-                                            block_tables,
-                                            cu_seqlens_q,
-                                            seq_lens,
-                                            seq_lens_encoder,
-                                            cos_emb,
-                                            sin_emb,
-                                            max_seq_len,
-                                            max_blocks_per_seq,
-                                            num_heads,
-                                            dim_head,
-                                            block_size,
-                                            elem_nums,
-                                            kv_num_heads,
-                                            rope_3d,
-                                            q_norm_weight,
-                                            k_norm_weight,
-                                            rms_norm_eps);
-}
-
 template <typename T, typename QKV_TYPE>
 void append_decode_cache_rope(const QKV_TYPE* qkv,
                               T* key_cache,
                               T* value_cache,
                               T* qkv_out,
                               const int* block_tables,
-                              const int* cu_seqlens_q,
+                              const int* padding_offsets,
+                              const int* cum_offsets,
                               const int* seq_lens,
                               const int* seq_lens_encoder,
                               const float* cos_emb,
@@ -94,7 +34,6 @@ void append_decode_cache_rope(const QKV_TYPE* qkv,
                               const int num_heads,
                               const int kv_num_heads,
                               const int dim_head,
-                              const int rotary_dim,
                               const int block_size,
                               const int bsz,
                               const cudaStream_t& stream,
@@ -118,7 +57,8 @@ void append_decode_cache_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -131,38 +71,16 @@ void append_decode_cache_rope(const QKV_TYPE* qkv,
               dim_head,
               block_size,
               elem_nums,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     } else {
-      if (rotary_dim < dim_head){
-        append_decode_cache_T_neox_partial_rope_kernel<T, PackSize>
+      append_decode_cache_T_neox_rope_kernel<T, PackSize>
           <<<grid_size, blocksize, 0, stream>>>(reinterpret_cast<const T*>(qkv),
                                                 key_cache,
                                                 value_cache,
                                                 qkv_out,
                                                 block_tables,
-                                                cu_seqlens_q,
-                                                seq_lens,
-                                                seq_lens_encoder,
-                                                cos_emb,
-                                                sin_emb,
-                                                max_seq_len,
-                                                max_blocks_per_seq,
-                                                num_heads,
-                                                dim_head,
-                                                rotary_dim,
-                                                block_size,
-                                                elem_nums,
-                                                kv_num_heads,
-                                                rope_3d);
-      }else{
-        append_decode_cache_T_neox_rope_kernel<T, PackSize>
-          <<<grid_size, blocksize, 0, stream>>>(reinterpret_cast<const T*>(qkv),
-                                                key_cache,
-                                                value_cache,
-                                                qkv_out,
-                                                block_tables,
-                                                cu_seqlens_q,
+                                                padding_offsets,
+                                                cum_offsets,
                                                 seq_lens,
                                                 seq_lens_encoder,
                                                 cos_emb,
@@ -173,9 +91,7 @@ void append_decode_cache_rope(const QKV_TYPE* qkv,
                                                 dim_head,
                                                 block_size,
                                                 elem_nums,
-                                                kv_num_heads,
-                                                rope_3d);
-      }
+                                                kv_num_heads);
     }
   } else {
     if (qkv_out_scales) {
@@ -186,7 +102,8 @@ void append_decode_cache_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -208,7 +125,8 @@ void append_decode_cache_rope(const QKV_TYPE* qkv,
                                                 value_cache,
                                                 qkv_out,
                                                 block_tables,
-                                                cu_seqlens_q,
+                                                padding_offsets,
+                                                cum_offsets,
                                                 seq_lens,
                                                 seq_lens_encoder,
                                                 cos_emb,
@@ -231,7 +149,8 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
                                    uint8_t* value_cache,
                                    T* qkv_out,
                                    const int* block_tables,
-                                   const int* cu_seqlens_q,
+                                   const int* padding_offsets,
+                                   const int* cum_offsets,
                                    const int* seq_lens,
                                    const int* seq_lens_encoder,
                                    const float* cos_emb,
@@ -263,7 +182,8 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -278,8 +198,7 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               block_size,
               127.0f,
               -127.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     } else {
       append_decode_cache_int8_neox_rope_kernel<T, 4>
           <<<grids, num_warps * 32, 0, stream>>>(
@@ -288,7 +207,8 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -301,8 +221,7 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               block_size,
               127.0f,
               -127.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     }
   } else {
     if (qkv_out_scales) {
@@ -313,7 +232,8 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -328,8 +248,7 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               block_size,
               127.0f,
               -127.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     } else {
       append_decode_cache_int8_rope_kernel<T, 4, 0, 128, is_scale_channel_wise, IsFP8>
           <<<grids, num_warps * 32, 0, stream>>>(
@@ -338,7 +257,8 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -351,8 +271,7 @@ void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
               block_size,
               127.0f,
               -127.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     }
   }
 }
@@ -363,7 +282,8 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
                                    uint8_t* value_cache,
                                    T* qkv_out,
                                    const int* block_tables,
-                                   const int* cu_seqlens_q,
+                                   const int* padding_offsets,
+                                   const int* cum_offsets,
                                    const int* seq_lens,
                                    const int* seq_lens_encoder,
                                    const float* cos_emb,
@@ -397,7 +317,8 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -414,8 +335,7 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               block_size,
               7.0f,
               -8.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     } else {
       append_decode_cache_int4_neox_rope_kernel<T, 4>
           <<<grids, num_warps * 32, 0, stream>>>(
@@ -424,7 +344,8 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -439,8 +360,7 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               block_size,
               7.0f,
               -8.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     }
   } else {
     if (qkv_out_scales) {
@@ -451,7 +371,8 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -468,8 +389,7 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               block_size,
               7.0f,
               -8.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     } else {
       append_decode_cache_int4_rope_kernel<T, 4>
           <<<grids, num_warps * 32, 0, stream>>>(
@@ -478,7 +398,8 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               value_cache,
               qkv_out,
               block_tables,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               seq_lens,
               seq_lens_encoder,
               cos_emb,
@@ -493,8 +414,7 @@ void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
               block_size,
               7.0f,
               -8.0f,
-              kv_num_heads,
-              rope_3d);
+              kv_num_heads);
     }
   }
 }
@@ -504,7 +424,8 @@ void DecoderWriteCacheWithRoPEKernel(
     const paddle::Tensor& qkv,
     const paddle::Tensor& seq_lens,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& qkv_out_scales,
@@ -520,10 +441,7 @@ void DecoderWriteCacheWithRoPEKernel(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps) {
+    paddle::Tensor* value_cache_out) {
   typedef cascade_attn_type_traits<T> traits_;
   typedef cascade_attn_type_traits<QKV_TYPE> qkt_nv_type_;
   typedef typename traits_::type DataType_;
@@ -540,31 +458,86 @@ void DecoderWriteCacheWithRoPEKernel(
   const float* cos_emb =
       rotary_embs ? rotary_embs.get().data<float>() : nullptr;
   const float* sin_emb;
-  int rotary_dim = dim_head;
   if (rotary_embs) {
     sin_emb =
         use_neox_rotary_style
             ? rotary_embs.get().data<float>() + max_seq_len * dim_head
             : rotary_embs.get().data<float>() + max_seq_len * dim_head / 2;
-    rotary_dim = rotary_embs.get().dims()[rotary_embs.get().dims().size()-1] * 2;
-    if(rotary_dim < dim_head){
-      if (!use_neox_rotary_style || qkv_out_scales || q_norm_weight || k_norm_weight|| cache_quant_type_str != "none"){
-        PADDLE_THROW(phi::errors::Fatal(
-          "partial_rotary_factor < 1.0 only supports neox_rotary_style=True, qkv_out_scales is None, q_norm_weight/k_norm_weight) is None, and cache_quant_type_str is 'none'."));
-      }
-      sin_emb = rotary_embs.get().data<float>() + max_seq_len * rotary_dim / 2;
-    }
   }
-
-  if (q_norm_weight && k_norm_weight) {
-    if (cache_quant_type_str == "none") {
-      append_decode_cache_rope_qk_norm(
+  if (cache_quant_type_str == "none") {
+    append_decode_cache_rope(
+        reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+        reinterpret_cast<DataType_*>(key_cache_out->data<T>()),
+        reinterpret_cast<DataType_*>(value_cache_out->data<T>()),
+        reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+        block_tables.data<int>(),
+        padding_offsets.data<int>(),
+        cum_offsets.data<int>(),
+        seq_lens.data<int>(),
+        seq_lens_encoder.data<int>(),
+        cos_emb,
+        sin_emb,
+        qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+        qkv_biases ? reinterpret_cast<DataType_*>(
+                         const_cast<T*>(qkv_biases.get().data<T>()))
+                   : nullptr,
+        max_seq_len,
+        max_blocks_per_seq,
+        num_heads,
+        kv_num_heads,
+        dim_head,
+        block_size,
+        bsz,
+        stream,
+        use_neox_rotary_style,
+        rope_3d);
+  } else if (cache_quant_type_str == "cache_int8") {
+    bool is_scale_channel_wise = false;
+    if (cache_k_scale && cache_k_scale.get().dims()[0] == dim_head * kv_num_heads) {
+      is_scale_channel_wise = true;
+    }
+    if (is_scale_channel_wise) {
+      append_decode_cache_int8_rope<DataType_, QKV_TYPE, true>(
+        reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+        key_cache_out->data<uint8_t>(),
+        value_cache_out->data<uint8_t>(),
+        reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+        block_tables.data<int>(),
+        padding_offsets.data<int>(),
+        cum_offsets.data<int>(),
+        seq_lens.data<int>(),
+        seq_lens_encoder.data<int>(),
+        cos_emb,
+        sin_emb,
+        qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+        qkv_biases ? reinterpret_cast<DataType_*>(
+                         const_cast<T*>(qkv_biases.get().data<T>()))
+                   : nullptr,
+        cache_k_scale ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(cache_k_scale.get().data<T>()))
+                      : nullptr,
+        cache_v_scale ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(cache_v_scale.get().data<T>()))
+                      : nullptr,
+        max_seq_len,
+        max_blocks_per_seq,
+        num_heads,
+        kv_num_heads,
+        dim_head,
+        block_size,
+        bsz,
+        stream,
+        use_neox_rotary_style,
+        rope_3d);
+    } else {
+      append_decode_cache_int8_rope<DataType_, QKV_TYPE, false>(
           reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
-          reinterpret_cast<DataType_*>(key_cache_out->data<T>()),
-          reinterpret_cast<DataType_*>(value_cache_out->data<T>()),
+          key_cache_out->data<uint8_t>(),
+          value_cache_out->data<uint8_t>(),
           reinterpret_cast<DataType_*>(qkv_out->data<T>()),
           block_tables.data<int>(),
-          cu_seqlens_q.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
           seq_lens.data<int>(),
           seq_lens_encoder.data<int>(),
           cos_emb,
@@ -573,6 +546,12 @@ void DecoderWriteCacheWithRoPEKernel(
           qkv_biases ? reinterpret_cast<DataType_*>(
                           const_cast<T*>(qkv_biases.get().data<T>()))
                     : nullptr,
+          cache_k_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_k_scale.get().data<T>()))
+                        : nullptr,
+          cache_v_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_v_scale.get().data<T>()))
+                        : nullptr,
           max_seq_len,
           max_blocks_per_seq,
           num_heads,
@@ -582,247 +561,84 @@ void DecoderWriteCacheWithRoPEKernel(
           bsz,
           stream,
           use_neox_rotary_style,
-          rope_3d,
-          q_norm_weight ? q_norm_weight.get().data<float>() : nullptr,
-          k_norm_weight ? k_norm_weight.get().data<float>() : nullptr,
-          rms_norm_eps);
-    } else if (cache_quant_type_str == "block_wise_fp8") {
-      constexpr int num_warps = 4;
-      const int all_warps =
-          ((num_heads + 2 * kv_num_heads) + num_warps - 1) / num_warps * num_warps;
-      dim3 grids(bsz, all_warps / num_warps);
-      append_decode_cache_int8_rope_qk_norm_kernel<DataType_, 4, 0, 128, false, true>
-          <<<grids, num_warps * 32, 0, stream>>>(
-              reinterpret_cast<const DataType_*>(qkv_ptr),
-              key_cache_out->data<uint8_t>(),
-              value_cache_out->data<uint8_t>(),
-              reinterpret_cast<DataType_*>(qkv_out->data<T>()),
-              block_tables.data<int>(),
-              cu_seqlens_q.data<int>(),
-              seq_lens.data<int>(),
-              seq_lens_encoder.data<int>(),
-              cos_emb,
-              sin_emb,
-              const_cast<DataType_*>(reinterpret_cast<const DataType_*>(cache_k_scale.get().data<T>())),
-              const_cast<DataType_*>(reinterpret_cast<const DataType_*>((cache_v_scale.get().data<T>()))),
-              q_norm_weight.get().data<float>(),
-              k_norm_weight.get().data<float>(),
-              max_seq_len,
-              max_blocks_per_seq,
-              num_heads,
-              block_size,
-              127.0f,
-              -127.0f,
-              kv_num_heads,
-              rope_3d,
-              rms_norm_eps);
-    } else {
-      PD_THROW(
-          "append_decode_cache_rope_qk_norm just supports cache_quant_type none/block_wise_fp8");
+          rope_3d);
     }
+  } else if (cache_quant_type_str == "cache_fp8")  {
+      append_decode_cache_int8_rope<DataType_, QKV_TYPE, false, true>(
+          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+          key_cache_out->data<uint8_t>(),
+          value_cache_out->data<uint8_t>(),
+          reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          seq_lens_encoder.data<int>(),
+          cos_emb,
+          sin_emb,
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? reinterpret_cast<DataType_*>(
+                          const_cast<T*>(qkv_biases.get().data<T>()))
+                    : nullptr,
+          cache_k_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_k_scale.get().data<T>()))
+                        : nullptr,
+          cache_v_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_v_scale.get().data<T>()))
+                        : nullptr,
+          max_seq_len,
+          max_blocks_per_seq,
+          num_heads,
+          kv_num_heads,
+          dim_head,
+          block_size,
+          bsz,
+          stream,
+          use_neox_rotary_style,
+          rope_3d);
+  } else if (cache_quant_type_str == "cache_int4_zp") {
+    append_decode_cache_int4_rope(
+        reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+        key_cache_out->data<uint8_t>(),
+        value_cache_out->data<uint8_t>(),
+        reinterpret_cast<DataType_*>(const_cast<T*>(qkv_out->data<T>())),
+        block_tables.data<int>(),
+        padding_offsets.data<int>(),
+        cum_offsets.data<int>(),
+        seq_lens.data<int>(),
+        seq_lens_encoder.data<int>(),
+        cos_emb,
+        sin_emb,
+        qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+        qkv_biases ? reinterpret_cast<DataType_*>(
+                         const_cast<T*>(qkv_biases.get().data<T>()))
+                   : nullptr,
+        cache_k_scale ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(cache_k_scale.get().data<T>()))
+                      : nullptr,
+        cache_v_scale ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(cache_v_scale.get().data<T>()))
+                      : nullptr,
+        cache_k_zp ? reinterpret_cast<DataType_*>(
+                         const_cast<T*>(cache_k_zp.get().data<T>()))
+                   : nullptr,
+        cache_v_zp ? reinterpret_cast<DataType_*>(
+                         const_cast<T*>(cache_v_zp.get().data<T>()))
+                   : nullptr,
+        max_seq_len,
+        max_blocks_per_seq,
+        num_heads,
+        kv_num_heads,
+        dim_head,
+        block_size,
+        bsz,
+        stream,
+        use_neox_rotary_style,
+        rope_3d);
   } else {
-    if (cache_quant_type_str == "none") {
-      append_decode_cache_rope(
-          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
-          reinterpret_cast<DataType_*>(key_cache_out->data<T>()),
-          reinterpret_cast<DataType_*>(value_cache_out->data<T>()),
-          reinterpret_cast<DataType_*>(qkv_out->data<T>()),
-          block_tables.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens.data<int>(),
-          seq_lens_encoder.data<int>(),
-          cos_emb,
-          sin_emb,
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? reinterpret_cast<DataType_*>(
-                          const_cast<T*>(qkv_biases.get().data<T>()))
-                    : nullptr,
-          max_seq_len,
-          max_blocks_per_seq,
-          num_heads,
-          kv_num_heads,
-          dim_head,
-          rotary_dim,
-          block_size,
-          bsz,
-          stream,
-          use_neox_rotary_style,
-          rope_3d);
-    } else if (cache_quant_type_str == "cache_int8") {
-      bool is_scale_channel_wise = false;
-      if (cache_k_scale && cache_k_scale.get().dims()[0] == dim_head * kv_num_heads) {
-        is_scale_channel_wise = true;
-      }
-      if (is_scale_channel_wise) {
-        append_decode_cache_int8_rope<DataType_, QKV_TYPE, true>(
-          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
-          key_cache_out->data<uint8_t>(),
-          value_cache_out->data<uint8_t>(),
-          reinterpret_cast<DataType_*>(qkv_out->data<T>()),
-          block_tables.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens.data<int>(),
-          seq_lens_encoder.data<int>(),
-          cos_emb,
-          sin_emb,
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? reinterpret_cast<DataType_*>(
-                          const_cast<T*>(qkv_biases.get().data<T>()))
-                    : nullptr,
-          cache_k_scale ? reinterpret_cast<DataType_*>(
-                              const_cast<T*>(cache_k_scale.get().data<T>()))
-                        : nullptr,
-          cache_v_scale ? reinterpret_cast<DataType_*>(
-                              const_cast<T*>(cache_v_scale.get().data<T>()))
-                        : nullptr,
-          max_seq_len,
-          max_blocks_per_seq,
-          num_heads,
-          kv_num_heads,
-          dim_head,
-          block_size,
-          bsz,
-          stream,
-          use_neox_rotary_style,
-          rope_3d);
-      } else {
-        append_decode_cache_int8_rope<DataType_, QKV_TYPE, false>(
-            reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
-            key_cache_out->data<uint8_t>(),
-            value_cache_out->data<uint8_t>(),
-            reinterpret_cast<DataType_*>(qkv_out->data<T>()),
-            block_tables.data<int>(),
-            cu_seqlens_q.data<int>(),
-            seq_lens.data<int>(),
-            seq_lens_encoder.data<int>(),
-            cos_emb,
-            sin_emb,
-            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-            qkv_biases ? reinterpret_cast<DataType_*>(
-                            const_cast<T*>(qkv_biases.get().data<T>()))
-                      : nullptr,
-            cache_k_scale ? reinterpret_cast<DataType_*>(
-                                const_cast<T*>(cache_k_scale.get().data<T>()))
-                          : nullptr,
-            cache_v_scale ? reinterpret_cast<DataType_*>(
-                                const_cast<T*>(cache_v_scale.get().data<T>()))
-                          : nullptr,
-            max_seq_len,
-            max_blocks_per_seq,
-            num_heads,
-            kv_num_heads,
-            dim_head,
-            block_size,
-            bsz,
-            stream,
-            use_neox_rotary_style,
-            rope_3d);
-      }
-    } else if (cache_quant_type_str == "cache_fp8")  {
-        append_decode_cache_int8_rope<DataType_, QKV_TYPE, false, true>(
-            reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
-            key_cache_out->data<uint8_t>(),
-            value_cache_out->data<uint8_t>(),
-            reinterpret_cast<DataType_*>(qkv_out->data<T>()),
-            block_tables.data<int>(),
-            cu_seqlens_q.data<int>(),
-            seq_lens.data<int>(),
-            seq_lens_encoder.data<int>(),
-            cos_emb,
-            sin_emb,
-            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-            qkv_biases ? reinterpret_cast<DataType_*>(
-                            const_cast<T*>(qkv_biases.get().data<T>()))
-                      : nullptr,
-            cache_k_scale ? reinterpret_cast<DataType_*>(
-                                const_cast<T*>(cache_k_scale.get().data<T>()))
-                          : nullptr,
-            cache_v_scale ? reinterpret_cast<DataType_*>(
-                                const_cast<T*>(cache_v_scale.get().data<T>()))
-                          : nullptr,
-            max_seq_len,
-            max_blocks_per_seq,
-            num_heads,
-            kv_num_heads,
-            dim_head,
-            block_size,
-            bsz,
-            stream,
-            use_neox_rotary_style,
-            rope_3d);
-    } else if (cache_quant_type_str == "block_wise_fp8") {
-      constexpr int num_warps = 4;
-      const int all_warps =
-          ((num_heads + 2 * kv_num_heads) + num_warps - 1) / num_warps * num_warps;
-      dim3 grids(bsz, all_warps / num_warps);
-      append_decode_cache_int8_rope_qk_norm_kernel<DataType_, 4, 0, 128, false, true>
-          <<<grids, num_warps * 32, 0, stream>>>(
-              reinterpret_cast<const DataType_*>(qkv_ptr),
-              key_cache_out->data<uint8_t>(),
-              value_cache_out->data<uint8_t>(),
-              reinterpret_cast<DataType_*>(qkv_out->data<T>()),
-              block_tables.data<int>(),
-              cu_seqlens_q.data<int>(),
-              seq_lens.data<int>(),
-              seq_lens_encoder.data<int>(),
-              cos_emb,
-              sin_emb,
-              const_cast<DataType_*>(reinterpret_cast<const DataType_*>(cache_k_scale.get().data<T>())),
-              const_cast<DataType_*>(reinterpret_cast<const DataType_*>((cache_v_scale.get().data<T>()))),
-              nullptr,
-              nullptr,
-              max_seq_len,
-              max_blocks_per_seq,
-              num_heads,
-              block_size,
-              127.0f,
-              -127.0f,
-              kv_num_heads,
-              rope_3d,
-              rms_norm_eps);
-    } else if (cache_quant_type_str == "cache_int4_zp") {
-      append_decode_cache_int4_rope(
-          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
-          key_cache_out->data<uint8_t>(),
-          value_cache_out->data<uint8_t>(),
-          reinterpret_cast<DataType_*>(const_cast<T*>(qkv_out->data<T>())),
-          block_tables.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens.data<int>(),
-          seq_lens_encoder.data<int>(),
-          cos_emb,
-          sin_emb,
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? reinterpret_cast<DataType_*>(
-                          const_cast<T*>(qkv_biases.get().data<T>()))
-                    : nullptr,
-          cache_k_scale ? reinterpret_cast<DataType_*>(
-                              const_cast<T*>(cache_k_scale.get().data<T>()))
-                        : nullptr,
-          cache_v_scale ? reinterpret_cast<DataType_*>(
-                              const_cast<T*>(cache_v_scale.get().data<T>()))
-                        : nullptr,
-          cache_k_zp ? reinterpret_cast<DataType_*>(
-                          const_cast<T*>(cache_k_zp.get().data<T>()))
-                    : nullptr,
-          cache_v_zp ? reinterpret_cast<DataType_*>(
-                          const_cast<T*>(cache_v_zp.get().data<T>()))
-                    : nullptr,
-          max_seq_len,
-          max_blocks_per_seq,
-          num_heads,
-          kv_num_heads,
-          dim_head,
-          block_size,
-          bsz,
-          stream,
-          use_neox_rotary_style,
-          rope_3d);
-    } else {
-      PD_THROW(
-          "cache_quant_type_str should be one of [none, cache_int8, cache_fp8 "
-          "cache_int4_zp]");
-    }
+    PD_THROW(
+        "cache_quant_type_str should be one of [none, cache_int8, cache_fp8 "
+        "cache_int4_zp]");
   }
 }
 
@@ -834,7 +650,8 @@ template void DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, int>(
               // kv_num_heads, head_dim] if GQA)
     const paddle::Tensor& seq_lens,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& qkv_out_scales,
@@ -850,10 +667,7 @@ template void DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, int>(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps);
+    paddle::Tensor* value_cache_out);
 
 template void
 DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, paddle::bfloat16>(
@@ -863,7 +677,8 @@ DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, paddle::bfloat16>(
               // kv_num_heads, head_dim] if GQA)
     const paddle::Tensor& seq_lens,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& qkv_out_scales,
@@ -879,10 +694,7 @@ DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, paddle::bfloat16>(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps);
+    paddle::Tensor* value_cache_out);
 
 template void DecoderWriteCacheWithRoPEKernel<paddle::float16, int>(
     const AppendAttnMetaData& meta_data,
@@ -891,7 +703,8 @@ template void DecoderWriteCacheWithRoPEKernel<paddle::float16, int>(
               // kv_num_heads, head_dim] if GQA)
     const paddle::Tensor& seq_lens,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& qkv_out_scales,
@@ -907,10 +720,7 @@ template void DecoderWriteCacheWithRoPEKernel<paddle::float16, int>(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps);
+    paddle::Tensor* value_cache_out);
 
 template void DecoderWriteCacheWithRoPEKernel<paddle::float16, paddle::float16>(
     const AppendAttnMetaData& meta_data,
@@ -919,7 +729,8 @@ template void DecoderWriteCacheWithRoPEKernel<paddle::float16, paddle::float16>(
               // kv_num_heads, head_dim] if GQA)
     const paddle::Tensor& seq_lens,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& qkv_out_scales,
@@ -935,7 +746,4 @@ template void DecoderWriteCacheWithRoPEKernel<paddle::float16, paddle::float16>(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps);
+    paddle::Tensor* value_cache_out);

custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_kernel.h

@@ -23,7 +23,8 @@ void DecoderWriteCacheWithRoPEKernel(
               // kv_num_heads, head_dim] if GQA)
     const paddle::Tensor& seq_lens,
     const paddle::Tensor& seq_lens_encoder,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::optional<paddle::Tensor>& rotary_embs,
     const paddle::optional<paddle::Tensor>& qkv_out_scales,
@@ -39,6 +40,4 @@ void DecoderWriteCacheWithRoPEKernel(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,    const float rms_norm_eps);
+    paddle::Tensor* value_cache_out);

custom_ops/gpu_ops/append_attn/encoder_write_cache_with_rope_kernel.h

@@ -25,8 +25,8 @@ void EncoderWriteCacheWithRopeKernel(
     const paddle::Tensor& seq_lens_this_time,
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder,
-    const paddle::Tensor& batch_id_per_token,
-    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const paddle::Tensor& batch_ids,
     const paddle::Tensor& tile_ids,
@@ -46,39 +46,43 @@ void EncoderWriteCacheWithRopeKernel(
     cudaStream_t& stream,
     paddle::Tensor* qkv_out,
     paddle::Tensor* key_cache_out,
-    paddle::Tensor* value_cache_out,
-    const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,
-    const float rms_norm_eps) {
+    paddle::Tensor* value_cache_out) {
   auto token_num = meta_data.token_nums;
   auto num_heads = meta_data.q_num_heads;
   auto kv_num_heads = meta_data.kv_num_heads;
   auto head_dim = meta_data.head_dims;
   bool is_scale_channel_wise = false;
-  int rotary_dim = head_dim;
   if (cache_k_scale && cache_k_scale.get().dims()[0] == head_dim * kv_num_heads) {
     is_scale_channel_wise = true;
   }
-  if (rotary_embs){
-    rotary_dim = rotary_embs.get().dims()[rotary_embs.get().dims().size()-1] * 2;
-    if(rotary_dim < head_dim){
-      if (!use_neox_style || q_norm_weight || k_norm_weight || num_heads == kv_num_heads || is_scale_channel_wise){
-        PADDLE_THROW(phi::errors::Fatal(
-          "partial_rotary_factor < 1.0 only supports use_neox_rotary_style=True, q_norm_weight/k_norm_weight) is None, GQA and is_scale_channel_wise=false."));
-      }
-    }
-  }
 
-  if (q_norm_weight && k_norm_weight) {
-    if (num_heads != kv_num_heads && !is_scale_channel_wise && !use_neox_style) {
-      gqa_rotary_qk_norm_variable(
+  if (num_heads == kv_num_heads) {
+    rotary_qk_variable(
         qkv_out->data<T>(),
         qkv.data<QKV_TYPE>(),
         qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
         qkv_biases ? qkv_biases.get().data<T>() : nullptr,
         rotary_embs.get().data<float>(),
-        batch_id_per_token.data<int>(),
-        cu_seqlens_q.data<int>(),
+        padding_offsets.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        token_num,
+        num_heads,
+        max_seq_len,
+        rotary_embs.get().dims()[2],
+        head_dim,
+        stream,
+        use_neox_style,
+        rope_3d);
+  } else {
+    if (!is_scale_channel_wise) {
+      gqa_rotary_qk_variable(
+        qkv_out->data<T>(),
+        qkv.data<QKV_TYPE>(),
+        qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+        qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+        rotary_embs.get().data<float>(),
+        padding_offsets.data<int>(),
         seq_lens_encoder.data<int>(),
         seq_lens_decoder.data<int>(),
         token_num,
@@ -89,96 +93,44 @@ void EncoderWriteCacheWithRopeKernel(
         head_dim,
         stream,
         use_neox_style,
-        rope_3d,
-        q_norm_weight ? q_norm_weight.get().data<float>() : nullptr,
-        k_norm_weight ? k_norm_weight.get().data<float>() : nullptr,
-        rms_norm_eps);
+        rope_3d);
     } else {
-      PD_THROW(
-          "gqa_rotary_qk_norm_variable only support gqa mode. channel wise scale and neox style are not supported");
+      gqa_rotary_qk_quant_variable(
+        qkv_out->data<T>(),
+        qkv.data<QKV_TYPE>(),
+        qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+        qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+        cache_k_scale ? cache_k_scale.get().data<T>() : nullptr,
+        cache_v_scale ? cache_v_scale.get().data<T>() : nullptr,
+        rotary_embs.get().data<float>(),
+        padding_offsets.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        token_num,
+        num_heads,
+        kv_num_heads,
+        max_seq_len,
+        rotary_embs.get().dims()[2],
+        head_dim,
+        stream,
+        use_neox_style,
+        rope_3d);
     }
-  } else {
-    if (num_heads == kv_num_heads) {
-      rotary_qk_variable(
-          qkv_out->data<T>(),
-          qkv.data<QKV_TYPE>(),
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
-          rotary_embs.get().data<float>(),
-          batch_id_per_token.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens_encoder.data<int>(),
-          seq_lens_decoder.data<int>(),
-          token_num,
-          num_heads,
-          max_seq_len,
-          rotary_embs.get().dims()[2],
-          head_dim,
-          stream,
-          use_neox_style,
-          rope_3d);
-    } else {
-      if (!is_scale_channel_wise) {
-        gqa_rotary_qk_variable(
-          qkv_out->data<T>(),
-          qkv.data<QKV_TYPE>(),
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
-          rotary_embs.get().data<float>(),
-          batch_id_per_token.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens_encoder.data<int>(),
-          seq_lens_decoder.data<int>(),
-          token_num,
-          num_heads,
-          kv_num_heads,
-          max_seq_len,
-          rope_3d ? rotary_embs.get().dims()[3] : rotary_embs.get().dims()[2],
-          head_dim,
-          rotary_dim,
-          stream,
-          use_neox_style,
-          rope_3d);
-      } else {
-        gqa_rotary_qk_quant_variable(
-          qkv_out->data<T>(),
-          qkv.data<QKV_TYPE>(),
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
-          cache_k_scale ? cache_k_scale.get().data<T>() : nullptr,
-          cache_v_scale ? cache_v_scale.get().data<T>() : nullptr,
-          rotary_embs.get().data<float>(),
-          batch_id_per_token.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens_encoder.data<int>(),
-          seq_lens_decoder.data<int>(),
-          token_num,
-          num_heads,
-          kv_num_heads,
-          max_seq_len,
-          rotary_embs.get().dims()[2],
-          head_dim,
-          stream,
-          use_neox_style,
-          rope_3d);
-      }
 
-    }
   }
   const uint32_t block_size = meta_data.block_size;
   if (cache_quant_type_str == "none") {
     CascadeAppendWriteCacheKVQKV<T>(meta_data,
                                     *qkv_out,
                                     block_tables,
-                                    batch_id_per_token,
-                                    cu_seqlens_q,
+                                    padding_offsets,
                                     seq_lens_encoder,
                                     seq_lens_decoder,
                                     max_seq_len,
                                     stream,
                                     key_cache_out,
                                     value_cache_out);
-  } else if (cache_quant_type_str == "cache_int8" or cache_quant_type_str == "cache_fp8" or cache_quant_type_str == "block_wise_fp8") {
+  } else if (cache_quant_type_str == "cache_int8" or cache_quant_type_str == "cache_fp8") {
     DISPATCH_HEAD_DIM(
         head_dim, HEAD_DIM, {DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {
           CascadeAppendWriteCacheKVC8QKV<T, HEAD_DIM, BLOCK_SIZE>(
@@ -190,15 +142,15 @@ void EncoderWriteCacheWithRopeKernel(
               cache_v_scale.get(),
               seq_lens_this_time,
               seq_lens_decoder,
-              batch_id_per_token,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               block_tables,
               batch_ids,
               tile_ids,
               num_blocks,
               max_seq_len,
               is_scale_channel_wise,
-              cache_quant_type_str,
+              cache_quant_type_str == "cache_fp8",
               stream,
               key_cache_out,
               value_cache_out);
@@ -217,8 +169,8 @@ void EncoderWriteCacheWithRopeKernel(
               cache_v_zp.get(),
               seq_lens_this_time,
               seq_lens_decoder,
-              batch_id_per_token,
-              cu_seqlens_q,
+              padding_offsets,
+              cum_offsets,
               block_tables,
               batch_ids,
               tile_ids,

custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu

@@ -11,11 +11,10 @@
 // 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 "cute/tensor.hpp"
+
 #include "helper.h"
 #include "paddle/extension.h"
 #include "paddle/phi/core/memory/memcpy.h"
-#include "utils.cuh"
 
 template <int THREADBLOCK_SIZE>
 __global__ void
@@ -117,93 +116,6 @@ void GetMaxLen(const paddle::Tensor &seq_lens_tensor,
       max_len_tensor.data<int>(), batch_size);
 }
 
-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) {
-  if (threadIdx.x == 0) {
-    int index = 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) continue;
-
-      int loop_times;
-      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;
-      }
-    }
-  }
-}
-
 __global__ void split_q_block(const int *__restrict__ seq_lens_q,
                               const int *__restrict__ seq_lens_encoder,
                               int *__restrict__ batch_ids,
@@ -279,38 +191,26 @@ get_max_len_kv_ernel(int *max_seq_lens_out, const int *seq_lens_this_time,
   }
 }
 
-void GetBlockShapeAndSplitKVBlock(
+std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
     const paddle::Tensor &seq_lens_encoder,
     const paddle::Tensor &seq_lens_decoder,
-    const paddle::Tensor &seq_lens_this_time,
-    paddle::Tensor &decoder_batch_ids,          // Inplace
-    paddle::Tensor &decoder_tile_ids_per_batch, // Inplace
-    paddle::Tensor &decoder_num_blocks_cpu,     // Inplace, Pinned Memory
-    paddle::Tensor &decoder_num_blocks_device,  // Inplace
-    paddle::Tensor &decoder_chunk_size_device,  // Inplace
-    paddle::Tensor &max_len_tensor_cpu,         // Inplace, CPU
-    paddle::Tensor &encoder_batch_ids,          // Inplace
-    paddle::Tensor &encoder_tile_ids_per_batch, // Inplace
-    paddle::Tensor &encoder_num_blocks_x_cpu,   // Inplace, CPU
-    paddle::Tensor &kv_batch_ids,               // Inplace
-    paddle::Tensor &kv_tile_ids_per_batch,      // Inplace
-    paddle::Tensor &kv_num_blocks_x_cpu,        // Inplace, CPU
-    paddle::Tensor &max_len_kv_cpu,             // Inplace, CPU
-    const int encoder_block_shape_q,
-    const int decoder_block_shape_q,
-    const int group_size,
-    const int block_size,
-    const int decoder_step_token_num)
-{
+    const paddle::Tensor &seq_lens_this_time, const paddle::Tensor &cum_offsets,
+    const int encoder_block_shape_q, const int decoder_block_shape_q,
+    const int group_size, const int block_size,
+    const int decoder_step_token_num) {
   auto stream = seq_lens_encoder.stream();
-  int bsz = seq_lens_this_time.shape()[0];
-
-  paddle::Tensor max_len_tensor_gpu = GetEmptyTensor({max_len_tensor_cpu.shape()[0]}, paddle::DataType::INT32, seq_lens_this_time.place());
+  int bsz = cum_offsets.shape()[0];
+  auto max_len_tensor =
+      GetEmptyTensor({8}, paddle::DataType::INT32, seq_lens_encoder.place());
   GetMaxLen(seq_lens_decoder, seq_lens_this_time, seq_lens_encoder,
-            max_len_tensor_gpu, bsz);
-  max_len_tensor_cpu.copy_(max_len_tensor_gpu, max_len_tensor_cpu.place(), false);
+            max_len_tensor, bsz);
 
-  auto max_len_cpu_ptr = max_len_tensor_cpu.data<int>();
+  // max_len_this_time, max_enc_len_this_time, max_dec_len_this_time,
+  // max_enc_dec_len_this_time, max_just_dec_len_this_time,
+  // max_just_dec_merged_len_this_time, max_system_len,
+  // max_just_dec_len_without_system
+  auto max_len_cpu = max_len_tensor.copy_to(paddle::CPUPlace(), false);
+  auto max_len_cpu_ptr = max_len_cpu.data<int>();
   int max_len_this_time = max_len_cpu_ptr[0];
   int max_enc_len_this_time = max_len_cpu_ptr[1];
   int max_dec_len_this_time = max_len_cpu_ptr[2];
@@ -320,120 +220,34 @@ void GetBlockShapeAndSplitKVBlock(
   int max_system_len = max_len_cpu_ptr[6];
   int max_just_dec_len_without_system = max_len_cpu_ptr[7];
 
+  paddle::Tensor encoder_batch_ids;
+  paddle::Tensor encoder_tile_ids_per_batch;
+  paddle::Tensor encoder_num_blocks_x_cpu; /*cpu*/
+  paddle::Tensor kv_batch_ids;
+  paddle::Tensor kv_tile_ids_per_batch;
+  paddle::Tensor kv_num_blocks_x_cpu; /*cpu*/
+  paddle::Tensor decoder_batch_ids;
+  paddle::Tensor decoder_tile_ids_per_batch;
+  paddle::Tensor decoder_num_blocks_x_cpu; /*cpu*/
+  paddle::Tensor max_len_kv_cpu;           /*cpu*/
+
   auto max_len_kv =
       GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_decoder.place());
   get_max_len_kv_ernel<128><<<1, 128, 0, stream>>>(
       max_len_kv.data<int>(), seq_lens_this_time.data<int>(),
       seq_lens_decoder.data<int>(), bsz);
 
+  max_len_kv_cpu = max_len_kv.copy_to(paddle::CPUPlace(), false);
 
-  max_len_kv_cpu.copy_(max_len_kv, max_len_kv_cpu.place(), false);
-
-  // decoder
-  if (max_dec_len_this_time > 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);
-
-      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
-          decoder_chunk_size_device.data<int>(), 64, sizeof(int32_t), stream));
-
-      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
-          decoder_num_blocks_device.data<int>(), 0, sizeof(int32_t), stream));
-
-      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_device.data<int>(),
-                                 decoder_chunk_size_device.data<int>(),
-                                 bsz,
-                                 set_chunk_size,
-                                 block_size,
-                                 sm_cout);
-
-      decoder_num_blocks_cpu.copy_(
-          decoder_num_blocks_device, decoder_num_blocks_cpu.place(), false);
-      auto decoder_chunk_size_cpu =
-          decoder_chunk_size_device.copy_to(paddle::CPUPlace(), false);
-      const int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
-
-      //  NOTE: (changwenbin) When using auto_chunk,
-      // decode_max_tile_size must take into account the maximum case, where * 1024 can cover 128K.
-      // const uint32_t decoder_batch_shape = seq_lens_decoder.dims()[0] * 1024;
-
-      const uint32_t decoder_max_tile_size_per_bs_q =
-          div_up((decoder_step_token_num * group_size), decoder_block_shape_q);
-      const uint32_t decoder_batch_shape =
-          bsz * 1024 * decoder_max_tile_size_per_bs_q;
-
-      PADDLE_ENFORCE_GPU_SUCCESS(
-          cudaMemsetAsync(decoder_batch_ids.data<int>(),
-                          0,
-                          decoder_batch_shape * sizeof(int32_t),
-                          stream));
-      PADDLE_ENFORCE_GPU_SUCCESS(
-          cudaMemsetAsync(decoder_tile_ids_per_batch.data<int>(),
-                          0,
-                          decoder_batch_shape * sizeof(int32_t),
-                          stream));
-
-
-      split_block_for_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>(),
-          bsz,
-          chunk_size);
-
-    } else {
-        // Note:(changwenbin)In order to adapt to cudagraph, the maximum value should be taken here
-        const uint32_t decoder_max_tile_size_per_bs_q = div_up((decoder_step_token_num * group_size), decoder_block_shape_q);
-        const uint32_t decoder_batch_shape = bsz * 1024 * decoder_max_tile_size_per_bs_q;
-
-        PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(decoder_batch_ids.data<int>(), 0, decoder_batch_shape * sizeof(int32_t), stream));
-        PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(decoder_tile_ids_per_batch.data<int>(), 0, decoder_batch_shape * sizeof(int32_t), stream));
-        PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(decoder_num_blocks_device.data<int>(), 0, sizeof(int32_t), stream));
-
-        split_q_block<<<1, 32, 0, stream>>>(
-            seq_lens_this_time.data<int>(),
-            seq_lens_encoder.data<int>(),
-            decoder_batch_ids.data<int>(),
-            decoder_tile_ids_per_batch.data<int>(),
-            decoder_num_blocks_device.data<int>(),
-            bsz,
-            decoder_block_shape_q,
-            group_size);
-
-        decoder_num_blocks_cpu.copy_(
-            decoder_num_blocks_device, decoder_num_blocks_cpu.place(), false);
-        PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
-            decoder_chunk_size_device.data<int>(), 64, sizeof(int32_t), stream));
-    }
-  } else {
-      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
-          decoder_chunk_size_device.data<int>(), 64, sizeof(int32_t), stream));
-      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
-          decoder_num_blocks_device.data<int>(), 0, sizeof(int32_t), stream));
-      decoder_num_blocks_cpu.copy_(
-          decoder_num_blocks_device, decoder_num_blocks_cpu.place(), false);
-  }
-
-  // encoder
   if (max_enc_len_this_time > 0) {
-    const uint32_t max_tile_size_per_bs_kv = div_up(max_enc_dec_len_this_time, block_size);
-    const uint32_t kv_batch_shape = bsz * max_tile_size_per_bs_kv;
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(kv_batch_ids.data<int>(), 0, kv_batch_shape * sizeof(int32_t), stream));
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(kv_tile_ids_per_batch.data<int>(), 0, kv_batch_shape * sizeof(int32_t), stream));
+    const uint32_t max_tile_size_per_bs_kv =
+        div_up(max_enc_dec_len_this_time, block_size);
+    kv_batch_ids =
+        GetEmptyTensor({bsz * max_tile_size_per_bs_kv}, paddle::DataType::INT32,
+                       seq_lens_encoder.place());
+    kv_tile_ids_per_batch =
+        GetEmptyTensor({bsz * max_tile_size_per_bs_kv}, paddle::DataType::INT32,
+                       seq_lens_encoder.place());
     auto kv_num_blocks_x =
         GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
 
@@ -444,12 +258,16 @@ void GetBlockShapeAndSplitKVBlock(
         kv_tile_ids_per_batch.data<int>(), kv_num_blocks_x.data<int>(), bsz,
         block_size, block_size);
 
-    kv_num_blocks_x_cpu.copy_(kv_num_blocks_x, kv_num_blocks_x_cpu.place(), false);
-    // Clear buffer
-    const uint32_t encoder_max_tile_size_per_bs_q = div_up((max_enc_dec_len_this_time * group_size), encoder_block_shape_q);
-    const uint32_t encoder_batch_shape = bsz * encoder_max_tile_size_per_bs_q;
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(encoder_batch_ids.data<int>(), 0, encoder_batch_shape * sizeof(int32_t), stream));
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(encoder_tile_ids_per_batch.data<int>(), 0, encoder_batch_shape * sizeof(int32_t), stream));
+    kv_num_blocks_x_cpu = kv_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+
+    const uint32_t encoder_max_tile_size_per_bs_q =
+        div_up((max_enc_dec_len_this_time * group_size), encoder_block_shape_q);
+    encoder_batch_ids =
+        GetEmptyTensor({bsz * encoder_max_tile_size_per_bs_q},
+                       paddle::DataType::INT32, seq_lens_encoder.place());
+    encoder_tile_ids_per_batch =
+        GetEmptyTensor({bsz * encoder_max_tile_size_per_bs_q},
+                       paddle::DataType::INT32, seq_lens_encoder.place());
     auto encoder_num_blocks_x =
         GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
     split_q_block<<<1, 32, 0, stream>>>(seq_lens_encoder.data<int>(), nullptr,
@@ -457,38 +275,111 @@ void GetBlockShapeAndSplitKVBlock(
                                         encoder_tile_ids_per_batch.data<int>(),
                                         encoder_num_blocks_x.data<int>(), bsz,
                                         encoder_block_shape_q, group_size);
-    encoder_num_blocks_x_cpu.copy_(encoder_num_blocks_x, encoder_num_blocks_x_cpu.place(), false);
+    encoder_num_blocks_x_cpu =
+        encoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+  } else {
+    encoder_batch_ids =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+    encoder_tile_ids_per_batch =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+    encoder_num_blocks_x_cpu =
+        GetEmptyTensor({0}, paddle::DataType::INT32, paddle::CPUPlace());
+    kv_batch_ids =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+    kv_tile_ids_per_batch =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+    kv_num_blocks_x_cpu =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+  }
+  if (max_just_dec_len_this_time > 0) {
+    const uint32_t decoder_max_tile_size_per_bs_q =
+        div_up((decoder_step_token_num * group_size), decoder_block_shape_q);
+
+    decoder_batch_ids =
+        GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
+                       paddle::DataType::INT32, seq_lens_encoder.place());
+    decoder_tile_ids_per_batch =
+        GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
+                       paddle::DataType::INT32, seq_lens_encoder.place());
+    auto decoder_num_blocks_x =
+        GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+    split_q_block<<<1, 32, 0, stream>>>(
+        seq_lens_this_time.data<int>(), seq_lens_encoder.data<int>(),
+        decoder_batch_ids.data<int>(), decoder_tile_ids_per_batch.data<int>(),
+        decoder_num_blocks_x.data<int>(), bsz, decoder_block_shape_q,
+        group_size);
+    decoder_num_blocks_x_cpu =
+        decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+  } else {
+    decoder_batch_ids =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+    decoder_tile_ids_per_batch =
+        GetEmptyTensor({0}, paddle::DataType::INT32, seq_lens_encoder.place());
+    decoder_num_blocks_x_cpu =
+        GetEmptyTensor({0}, paddle::DataType::INT32, paddle::CPUPlace());
   }
 
+  return {encoder_batch_ids,
+          encoder_tile_ids_per_batch,
+          encoder_num_blocks_x_cpu, /*cpu*/
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          kv_num_blocks_x_cpu, /*cpu*/
+          decoder_batch_ids,
+          decoder_tile_ids_per_batch,
+          decoder_num_blocks_x_cpu, /*cpu*/
+          max_len_kv_cpu /*cpu*/,
+          max_len_cpu};
+}
+
+std::vector<paddle::DataType> GetBlockShapeAndSplitKVBlockInferDtype(
+    const paddle::DataType &seq_lens_encoder_dtype,
+    const paddle::DataType &seq_lens_decoder_dtype,
+    const paddle::DataType &seq_lens_this_time_dtype,
+    const paddle::DataType &cum_offsets_dtype) {
+  return {
+      paddle::DataType::INT32, paddle::DataType::INT32, paddle::DataType::INT32,
+      paddle::DataType::INT32, paddle::DataType::INT32, paddle::DataType::INT32,
+      paddle::DataType::INT32, paddle::DataType::INT32, paddle::DataType::INT32,
+      paddle::DataType::INT32, paddle::DataType::INT32};
+}
+
+std::vector<std::vector<int64_t>> GetBlockShapeAndSplitKVBlockInferShape(
+    const std::vector<int64_t> &seq_lens_encoder_shape,
+    const std::vector<int64_t> &seq_lens_decoder_shape,
+    const std::vector<int64_t> &seq_lens_this_time_shape,
+    const std::vector<int64_t> &cum_offsets_shape) {
+  std::vector<int64_t> dynamic_shape = {-1};
+
+  return {dynamic_shape,
+          dynamic_shape,
+          {1},
+          dynamic_shape,
+          dynamic_shape,
+          {1},
+          dynamic_shape,
+          dynamic_shape,
+          {1},
+          {1},
+          {8}};
 }
 
 PD_BUILD_STATIC_OP(get_block_shape_and_split_kv_block)
-    .Inputs({
-      "seq_lens_encoder",
-      "seq_lens_decoder",
-      "seq_lens_this_time",
-      "decoder_batch_ids",
-      "decoder_tile_ids_per_batch",
-      "decoder_num_blocks_cpu",
-      "decoder_num_blocks_device",
-      "decoder_chunk_size_device",
-      "max_len_tensor_cpu",
-      "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",
-      "max_len_kv_cpu"
-    })
-    .Outputs({
-
-    })
-    .Attrs({
-      "encoder_block_shape_q: int",
-      "decoder_block_shape_q: int",
-      "group_size: int",
-      "block_size: int",
-      "decoder_step_token_num: int"
-    })
-    .SetKernelFn(PD_KERNEL(GetBlockShapeAndSplitKVBlock));
+    .Inputs({"seq_lens_encoder", "seq_lens_decoder", "seq_lens_this_time",
+             "cum_offsets"})
+    .Outputs({paddle::Optional("encoder_batch_ids"),
+              paddle::Optional("encoder_tile_ids_per_batch"),
+              paddle::Optional("encoder_num_blocks"),
+              paddle::Optional("kv_batch_ids"),
+              paddle::Optional("kv_tile_ids_per_batch"),
+              paddle::Optional("kv_num_blocks"),
+              paddle::Optional("decoder_batch_ids"),
+              paddle::Optional("decoder_tile_ids_per_batch"),
+              paddle::Optional("decoder_num_blocks"),
+              paddle::Optional("max_len_kv"), "set_max_lengths"})
+    .Attrs({"encoder_block_shape_q: int", "decoder_block_shape_q: int",
+            "group_size: int", "block_size: int",
+            "decoder_step_token_num: int"})
+    .SetKernelFn(PD_KERNEL(GetBlockShapeAndSplitKVBlock))
+    .SetInferShapeFn(PD_INFER_SHAPE(GetBlockShapeAndSplitKVBlockInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GetBlockShapeAndSplitKVBlockInferDtype));