diff --git a/.devcontainer/cuda/devcontainer.json b/.devcontainer/cu126/devcontainer.json similarity index 64% rename from .devcontainer/cuda/devcontainer.json rename to .devcontainer/cu126/devcontainer.json index 7d88c73d6d..8c0e885fbb 100644 --- a/.devcontainer/cuda/devcontainer.json +++ b/.devcontainer/cu126/devcontainer.json @@ -1,8 +1,8 @@ { "name": "CUDA Development Container", "build": { - "dockerfile": "Dockerfile", - "context": "." + "dockerfile": "../../docker/Dockerfile.cu126.dev", + "context": "../../" }, "runArgs": [ "--gpus=all" @@ -12,8 +12,10 @@ "extensions": [ "llvm-vs-code-extensions.vscode-clangd", "ms-python.python", - "ms-python.black-formatter", - "nvidia.nsight-vscode-edition" + "nvidia.nsight-vscode-edition", + "xaver.clang-format", + "charliermarsh.ruff", + "matangover.mypy" ] } }, diff --git a/.devcontainer/cu128/devcontainer.json b/.devcontainer/cu128/devcontainer.json new file mode 100644 index 0000000000..b1afbf30a5 --- /dev/null +++ b/.devcontainer/cu128/devcontainer.json @@ -0,0 +1,26 @@ +{ + "name": "CUDA Development Container", + "build": { + "dockerfile": "../../docker/Dockerfile.cu128.dev", + "context": "../../" + }, + "runArgs": [ + "--gpus=all" + ], + "customizations": { + "vscode": { + "extensions": [ + "llvm-vs-code-extensions.vscode-clangd", + "ms-python.python", + "nvidia.nsight-vscode-edition", + "xaver.clang-format", + "charliermarsh.ruff", + "matangover.mypy" + ] + } + }, + "mounts": [ + "type=bind,source=${localEnv:HOME}/.ssh,target=/home/devuser/.ssh,readonly" + ], + "remoteUser": "devuser" +} diff --git a/.devcontainer/cu129/devcontainer.json b/.devcontainer/cu129/devcontainer.json new file mode 100644 index 0000000000..5829d9eae6 --- /dev/null +++ b/.devcontainer/cu129/devcontainer.json @@ -0,0 +1,26 @@ +{ + "name": "CUDA Development Container", + "build": { + "dockerfile": "../../docker/Dockerfile.cu129.dev", + "context": "../../" + }, + "runArgs": [ + "--gpus=all" + ], + "customizations": { + "vscode": { + "extensions": [ + "llvm-vs-code-extensions.vscode-clangd", + "ms-python.python", + "nvidia.nsight-vscode-edition", + "xaver.clang-format", + "charliermarsh.ruff", + "matangover.mypy" + ] + } + }, + "mounts": [ + "type=bind,source=${localEnv:HOME}/.ssh,target=/home/devuser/.ssh,readonly" + ], + "remoteUser": "devuser" +} diff --git a/.devcontainer/cu130/devcontainer.json b/.devcontainer/cu130/devcontainer.json new file mode 100644 index 0000000000..08b8d763fd --- /dev/null +++ b/.devcontainer/cu130/devcontainer.json @@ -0,0 +1,26 @@ +{ + "name": "CUDA Development Container", + "build": { + "dockerfile": "../../docker/Dockerfile.cu130.dev", + "context": "../../" + }, + "runArgs": [ + "--gpus=all" + ], + "customizations": { + "vscode": { + "extensions": [ + "llvm-vs-code-extensions.vscode-clangd", + "ms-python.python", + "nvidia.nsight-vscode-edition", + "xaver.clang-format", + "charliermarsh.ruff", + "matangover.mypy" + ] + } + }, + "mounts": [ + "type=bind,source=${localEnv:HOME}/.ssh,target=/home/devuser/.ssh,readonly" + ], + "remoteUser": "devuser" +} diff --git a/.devcontainer/cuda/install/install_python.sh b/.devcontainer/cuda/install/install_python.sh deleted file mode 100644 index 598d7fd126..0000000000 --- a/.devcontainer/cuda/install/install_python.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -set -e -set -u -set -o pipefail - - -# Install python and pip. Don't modify this to add Python package dependencies, -wget -O Miniforge3.sh "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh" -bash Miniforge3.sh -b -p /home/devuser/conda - -/home/devuser/conda/bin/conda create -n $1 python=3.12 diff --git a/.devcontainer/cuda/install/install_python_packages.sh b/.devcontainer/cuda/install/install_python_packages.sh deleted file mode 100644 index 49151a224f..0000000000 --- a/.devcontainer/cuda/install/install_python_packages.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/bin/bash -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -set -e -set -u - -pip3 install ninja pytest numpy scipy build cuda-python pytest -pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128 -pip3 install pre-commit diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000000..fc3b20c491 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,45 @@ +# Code Owners File +# Generated automatically from git history analysis +# Analysis period: 180 days +# Minimum commits threshold: 1 + +benchmarks/ @bkryu @jiahanc @cyx-6 @yzh119 @nv-yunzheq +benchmarks/routines/ @bkryu @nv-yunzheq @jiahanc @cyx-6 @nvmbreughe +ci/ @cyx-6 @yzh119 @nvmbreughe +ci/scripts/ @cyx-6 +ci/scripts/jenkins/ @cyx-6 +csrc/ @wenscarl @yzh119 @cyx-6 @djmmoss @nv-yunzheq +csrc/fused_moe/ @nv-yunzheq @yzh119 @yongwww @djmmoss @cyx-6 +csrc/fused_moe/cutlass_backend/ @nv-yunzheq @yzh119 @yongwww @djmmoss @cyx-6 +csrc/nv_internal/ @wenscarl @djmmoss @nv-yunzheq @yongwww @cyx-6 +csrc/nv_internal/cpp/ @wenscarl @bkryu @yongwww @djmmoss @joker-eph +csrc/nv_internal/include/ @wenscarl @nv-yunzheq +csrc/nv_internal/tensorrt_llm/ @wenscarl @djmmoss @nv-yunzheq @yongwww @cyx-6 +csrc/xqa/ @cyx-6 @yzh119 +docs/ @yzh119 @cyx-6 @wenscarl @nv-yunzheq @aleozlx +flashinfer/ @yzh119 @cyx-6 @nvmbreughe @aleozlx @wenscarl +flashinfer-cubin/ @yzh119 @cyx-6 +flashinfer-cubin/flashinfer_cubin/ @yzh119 +flashinfer-jit-cache/ @yzh119 @cyx-6 +flashinfer-jit-cache/flashinfer_jit_cache/ @yzh119 +flashinfer/comm/ @yzh119 @cyx-6 @nvmbreughe @wenscarl @djmmoss +flashinfer/cudnn/ @Anerudhan @yzh119 @cyx-6 @Anerudhan +flashinfer/cute_dsl/ @yzh119 @kaixih @Amir-19 @aleozlx +flashinfer/dsv3_ops/ @nvmbreughe +flashinfer/fused_moe/ @djmmoss @jiahanc @yzh119 @cyx-6 @aleozlx +flashinfer/gemm/ @nvmbreughe +flashinfer/jit/ @yzh119 @cyx-6 @aleozlx @jiahanc @nvmbreughe +flashinfer/jit/attention/ @yzh119 @cyx-6 @Anerudhan @joker-eph +flashinfer/jit/gemm/ @yzh119 @nv-yunzheq @jiahanc +flashinfer/logits_processor/ @cyx-6 @yzh119 +flashinfer/profiler/ @cyx-6 +flashinfer/triton/ @nvmbreughe @cyx-6 +flashinfer/tuning_configs/ @kaixih +include/ @yzh119 @jiahanc @nvmbreughe @IwakuraRein @bkryu +include/flashinfer/ @yzh119 @jiahanc @nvmbreughe @IwakuraRein @bkryu +include/flashinfer/attention/ @yzh119 @kahyunnam @joker-eph +include/flashinfer/comm/ @yongwww @nvmbreughe @djmmoss @yzh119 @cyx-6 +include/flashinfer/gemm/ @ttyio @yongwww @yzh119 @nvmbreughe @aleozlx +include/flashinfer/trtllm/ @jiahanc @joker-eph @aleozlx @yzh119 @wenscarl +profiler/ @cyx-6 +scripts/ @yzh119 @nvmbreughe @dierksen @yongwww @bkryu diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000000..f5c54c66e8 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,30 @@ + + +## 📌 Description + + + +## 🔍 Related Issues + + + +## 🚀 Pull Request Checklist + +Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. + +### ✅ Pre-commit Checks + +- [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). +- [ ] I have installed the hooks with `pre-commit install`. +- [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. + +> If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). + +## 🧪 Tests + +- [ ] Tests have been added or updated as needed. +- [ ] All tests are passing (`unittest`, etc.). + +## Reviewer Notes + + diff --git a/.github/workflows/build-doc.yml b/.github/workflows/build-doc.yml index 647e63c93b..bf13f3e586 100644 --- a/.github/workflows/build-doc.yml +++ b/.github/workflows/build-doc.yml @@ -5,6 +5,10 @@ on: branches: - main + pull_request: + branches: + - main + # Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages permissions: contents: read @@ -27,6 +31,11 @@ jobs: with: submodules: recursive + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + - name: Configuring build Environment run: | sudo apt-get update @@ -37,17 +46,18 @@ jobs: python -m pip install -r docs/requirements.txt - name: Bulid Documentation - if: github.ref == 'refs/heads/main' run: | cd docs - make html + ./build_docs.sh - name: Upload artifact + if: github.event_name == 'push' && github.ref == 'refs/heads/main' uses: actions/upload-pages-artifact@v3 with: # Upload entire repository path: 'docs/_build/html' - name: Deploy to GitHub Pages + if: github.event_name == 'push' && github.ref == 'refs/heads/main' id: deployment uses: actions/deploy-pages@v4 diff --git a/.github/workflows/new-issue.yml b/.github/workflows/new-issue.yml new file mode 100644 index 0000000000..a00ace949f --- /dev/null +++ b/.github/workflows/new-issue.yml @@ -0,0 +1,20 @@ +name: Triage new issue + +on: + issues: + types: + - opened + +jobs: + label_issues: + runs-on: ubuntu-latest + permissions: + issues: write + steps: + - name: Add needs-triage label + run: gh issue edit "$NUMBER" --add-label "$LABELS" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GH_REPO: ${{ github.repository }} + NUMBER: ${{ github.event.issue.number }} + LABELS: needs-triage diff --git a/.github/workflows/nightly-release.yml b/.github/workflows/nightly-release.yml new file mode 100644 index 0000000000..7c57d4bd7a --- /dev/null +++ b/.github/workflows/nightly-release.yml @@ -0,0 +1,421 @@ +name: Nightly Release + +on: + schedule: + # Run at 00:00 UTC every day + - cron: '0 0 * * *' + workflow_dispatch: + inputs: + date_suffix: + description: 'Date suffix for dev version (YYYYMMDD, leave empty for today)' + required: false + type: string + +jobs: + setup: + runs-on: ubuntu-latest + outputs: + dev_suffix: ${{ steps.set-suffix.outputs.dev_suffix }} + release_tag: ${{ steps.set-suffix.outputs.release_tag }} + version: ${{ steps.set-suffix.outputs.version }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set date suffix and release tag + id: set-suffix + run: | + # Read version from version.txt + VERSION=$(cat version.txt | tr -d '[:space:]') + + # Set date suffix + if [ -n "${{ inputs.date_suffix }}" ]; then + DEV_SUFFIX="${{ inputs.date_suffix }}" + else + DEV_SUFFIX=$(date -u +%Y%m%d) + fi + + # Create release tag with version + RELEASE_TAG="nightly-v${VERSION}-${DEV_SUFFIX}" + + echo "version=${VERSION}" >> $GITHUB_OUTPUT + echo "dev_suffix=${DEV_SUFFIX}" >> $GITHUB_OUTPUT + echo "release_tag=${RELEASE_TAG}" >> $GITHUB_OUTPUT + echo "Base version: ${VERSION}" + echo "Using dev suffix: ${DEV_SUFFIX}" + echo "Release tag: ${RELEASE_TAG}" + + build-flashinfer-python: + needs: setup + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build wheel + + - name: Build flashinfer-python wheel and sdist + env: + FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} + run: | + echo "Building flashinfer-python with dev suffix: ${FLASHINFER_DEV_RELEASE_SUFFIX}" + echo "Git commit: $(git rev-parse HEAD)" + python -m build + ls -lh dist/ + + - name: Upload flashinfer-python artifact + uses: actions/upload-artifact@v4 + with: + name: flashinfer-python-dist + path: dist/* + retention-days: 7 + + build-flashinfer-cubin: + needs: setup + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine wheel + + - name: Build flashinfer-cubin wheel + env: + FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} + run: | + echo "Building flashinfer-cubin with dev suffix: ${FLASHINFER_DEV_RELEASE_SUFFIX}" + echo "Git commit: $(git rev-parse HEAD)" + cd flashinfer-cubin + rm -rf dist build *.egg-info + python -m build --wheel + ls -lh dist/ + mkdir -p ../dist + cp dist/*.whl ../dist/ + + - name: Upload flashinfer-cubin artifact + uses: actions/upload-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist/*.whl + retention-days: 7 + + build-flashinfer-jit-cache: + needs: setup + strategy: + fail-fast: false + matrix: + cuda: ["12.8", "12.9", "13.0"] + arch: ['x86_64', 'aarch64'] + + runs-on: [self-hosted, "${{ matrix.arch == 'aarch64' && 'arm64' || matrix.arch }}"] + + steps: + - name: Display Machine Information + run: | + echo "CPU: $(nproc) cores, $(lscpu | grep 'Model name' | cut -d':' -f2 | xargs)" + echo "RAM: $(free -h | awk '/^Mem:/ {print $7 " available out of " $2}')" + echo "Disk: $(df -h / | awk 'NR==2 {print $4 " available out of " $2}')" + echo "Architecture: $(uname -m)" + + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Build wheel in container + env: + DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }} + FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f' }} + FLASHINFER_DEV_RELEASE_SUFFIX: ${{ needs.setup.outputs.dev_suffix }} + run: | + # Extract CUDA major and minor versions + CUDA_MAJOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f1) + CUDA_MINOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f2) + export CUDA_MAJOR + export CUDA_MINOR + export FLASHINFER_LOCAL_VERSION="cu${CUDA_MAJOR}${CUDA_MINOR}" + + chown -R $(id -u):$(id -g) ${{ github.workspace }} + mkdir -p ${{ github.workspace }}/ci-cache + chown -R $(id -u):$(id -g) ${{ github.workspace }}/ci-cache + + # Run the build script inside the container with proper mounts + docker run --rm \ + -v ${{ github.workspace }}:/workspace \ + -v ${{ github.workspace }}/ci-cache:/ci-cache \ + -e FLASHINFER_CI_CACHE=/ci-cache \ + -e CUDA_VERSION="${{ matrix.cuda }}" \ + -e CUDA_MAJOR="$CUDA_MAJOR" \ + -e CUDA_MINOR="$CUDA_MINOR" \ + -e FLASHINFER_LOCAL_VERSION="$FLASHINFER_LOCAL_VERSION" \ + -e FLASHINFER_DEV_RELEASE_SUFFIX="${FLASHINFER_DEV_RELEASE_SUFFIX}" \ + -e ARCH="${{ matrix.arch }}" \ + -e FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST}" \ + --user $(id -u):$(id -g) \ + -w /workspace \ + ${{ env.DOCKER_IMAGE }} \ + bash /workspace/scripts/build_flashinfer_jit_cache_whl.sh + timeout-minutes: 180 + + - name: Display wheel size + run: du -h flashinfer-jit-cache/dist/* + + - name: Create artifact name + id: artifact-name + run: | + CUDA_NO_DOT=$(echo "${{ matrix.cuda }}" | tr -d '.') + echo "name=jit-cache-cu${CUDA_NO_DOT}-${{ matrix.arch }}" >> $GITHUB_OUTPUT + + - name: Upload flashinfer-jit-cache artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ steps.artifact-name.outputs.name }} + path: flashinfer-jit-cache/dist/*.whl + retention-days: 7 + + create-release: + needs: [setup, build-flashinfer-python, build-flashinfer-cubin, build-flashinfer-jit-cache] + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Create GitHub Release (empty first) + env: + GH_TOKEN: ${{ github.token }} + run: | + TAG="${{ needs.setup.outputs.release_tag }}" + + # Delete existing release and tag if they exist + if gh release view "$TAG" &>/dev/null; then + echo "Deleting existing release: $TAG" + gh release delete "$TAG" --yes --cleanup-tag + fi + + # Create new release without assets first + gh release create "$TAG" \ + --title "Nightly Release v${{ needs.setup.outputs.version }}-${{ needs.setup.outputs.dev_suffix }}" \ + --notes "Automated nightly build for version ${{ needs.setup.outputs.version }} (dev${{ needs.setup.outputs.dev_suffix }})" \ + --prerelease + + - name: Download flashinfer-python artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-python-dist + path: dist-python/ + + - name: Upload flashinfer-python to release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-python/* --clobber + + - name: Download flashinfer-cubin artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist-cubin/ + + - name: Upload flashinfer-cubin to release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-cubin/* --clobber + + - name: Upload flashinfer-jit-cache wheels to release (one at a time to avoid OOM) + env: + GH_TOKEN: ${{ github.token }} + run: | + # Upload jit-cache wheels one at a time to avoid OOM + # Each wheel can be several GB, so we download, upload, delete, repeat + mkdir -p dist-jit-cache + + for cuda in 128 129 130; do + for arch in x86_64 aarch64; do + ARTIFACT_NAME="jit-cache-cu${cuda}-${arch}" + echo "Processing ${ARTIFACT_NAME}..." + + # Download this specific artifact + gh run download ${{ github.run_id }} -n "${ARTIFACT_NAME}" -D dist-jit-cache/ || { + echo "Warning: Failed to download ${ARTIFACT_NAME}, skipping..." + continue + } + + # Upload to release + if [ -n "$(ls -A dist-jit-cache/)" ]; then + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-jit-cache/* --clobber + echo "✅ Uploaded ${ARTIFACT_NAME}" + fi + + # Clean up to save disk space before next iteration + rm -rf dist-jit-cache/* + done + done + + test-nightly-build: + needs: [setup, build-flashinfer-python, build-flashinfer-cubin, build-flashinfer-jit-cache] + strategy: + fail-fast: false + matrix: + cuda: ["12.9", "13.0"] + test-shard: [1, 2, 3, 4, 5] + runs-on: [self-hosted, G5, X64] + + steps: + - name: Display Machine Information + run: | + echo "CPU: $(nproc) cores, $(lscpu | grep 'Model name' | cut -d':' -f2 | xargs)" + echo "RAM: $(free -h | awk '/^Mem:/ {print $7 " available out of " $2}')" + echo "Disk: $(df -h / | awk 'NR==2 {print $4 " available out of " $2}')" + echo "Architecture: $(uname -m)" + + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: true + + - name: Download flashinfer-python artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-python-dist + path: dist-python/ + + - name: Download flashinfer-cubin artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist-cubin/ + + - name: Download flashinfer-jit-cache artifact + uses: actions/download-artifact@v4 + with: + name: jit-cache-cu${{ matrix.cuda == '12.9' && '129' || '130' }}-x86_64 + path: dist-jit-cache/ + + - name: Get Docker image tag + id: docker-tag + run: | + CUDA_VERSION="cu${{ matrix.cuda == '12.9' && '129' || '130' }}" + DOCKER_TAG=$(grep "flashinfer/flashinfer-ci-${CUDA_VERSION}" ci/docker-tags.yml | cut -d':' -f2 | tr -d ' ') + echo "cuda_version=${CUDA_VERSION}" >> $GITHUB_OUTPUT + echo "tag=${DOCKER_TAG}" >> $GITHUB_OUTPUT + + - name: Run nightly build tests in Docker (shard ${{ matrix.test-shard }}) + env: + CUDA_VISIBLE_DEVICES: 0 + run: | + DOCKER_IMAGE="flashinfer/flashinfer-ci-${{ steps.docker-tag.outputs.cuda_version }}:${{ steps.docker-tag.outputs.tag }}" + bash ci/bash.sh ${DOCKER_IMAGE} \ + -e TEST_SHARD ${{ matrix.test-shard }} \ + -e FLASHINFER_JIT_CACHE_REPORT_FILE /workspace/jit_cache_report_shard${{ matrix.test-shard }}_cuda${{ matrix.cuda }}.json \ + ./scripts/task_test_nightly_build.sh + + - name: Upload JIT cache report + if: always() + uses: actions/upload-artifact@v4 + with: + name: jit-cache-report-shard${{ matrix.test-shard }}-cuda${{ matrix.cuda }} + path: jit_cache_report_shard${{ matrix.test-shard }}_cuda${{ matrix.cuda }}.json + if-no-files-found: ignore + retention-days: 7 + + jit-cache-summary: + needs: test-nightly-build + if: always() + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Download all JIT cache reports + uses: actions/download-artifact@v4 + with: + pattern: jit-cache-report-* + path: jit-reports/ + merge-multiple: true + + - name: Merge and print JIT cache summary + run: | + # Merge all report files into one + mkdir -p merged-reports + cat jit-reports/*.json > merged-reports/all_reports.json 2>/dev/null || echo "No JIT cache reports found" + + # Print summary + if [ -f merged-reports/all_reports.json ] && [ -s merged-reports/all_reports.json ]; then + python scripts/print_jit_cache_summary.py merged-reports/all_reports.json + else + echo "✅ No missing JIT cache modules - all tests passed!" + fi + + update-wheel-index: + needs: [setup, create-release, test-nightly-build] + runs-on: ubuntu-latest + steps: + - name: Checkout flashinfer repo + uses: actions/checkout@v4 + + - name: Download all artifacts + uses: actions/download-artifact@v4 + with: + path: artifacts/ + + - name: Collect wheels and sdist + run: | + mkdir -p dist + find artifacts/ -name "*.whl" -exec cp {} dist/ \; + find artifacts/ -name "*.tar.gz" -exec cp {} dist/ \; + ls -lh dist/ + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/flashinfer-ai/whl.git flashinfer-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Update wheel index + run: | + python3 scripts/update_whl_index.py \ + --dist-dir dist \ + --output-dir flashinfer-whl \ + --release-tag "${{ needs.setup.outputs.release_tag }}" \ + --nightly + + - name: Push wheel index + run: | + cd flashinfer-whl + git config --local user.name "github-actions[bot]" + git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com" + git add -A + git commit -m "update whl for nightly ${{ needs.setup.outputs.dev_suffix }}" + git push diff --git a/.github/workflows/release-ci-docker.yml b/.github/workflows/release-ci-docker.yml index 2f45ae302b..6df0c14066 100644 --- a/.github/workflows/release-ci-docker.yml +++ b/.github/workflows/release-ci-docker.yml @@ -2,21 +2,135 @@ name: Release CI Docker on: workflow_dispatch: - + push: + branches: + - main + paths: + - 'docker/**' + - '.github/workflows/release-ci-docker.yml' + - 'requirements.txt' + pull_request: + branches: + - main + paths: + - 'docker/**' + - '.github/workflows/release-ci-docker.yml' + - 'requirements.txt' jobs: - build: + generate-tag: runs-on: ubuntu-latest + outputs: + date_sha: ${{ steps.generate_tag.outputs.date_sha }} + steps: + - uses: actions/checkout@v4 + + - name: Generate version tag + id: generate_tag + run: | + DATE_SHA=$(date +'%Y%m%d')-$(git rev-parse --short HEAD) + echo "date_sha=${DATE_SHA}" >> $GITHUB_OUTPUT + echo "Generated version tag: ${DATE_SHA}" + + build: + runs-on: [self-hosted, x86_64] + needs: generate-tag + strategy: + matrix: + cuda: [cu126, cu128, cu129, cu130] + arch: [amd64, arm64] steps: - uses: actions/checkout@v4 + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + - name: Login to Docker Hub + if: github.event_name != 'pull_request' uses: docker/login-action@v3 with: username: flashinfer password: ${{ secrets.DOCKERHUB_TOKEN }} - - uses: docker/build-push-action@v4 + - name: Build and push ${{ matrix.cuda }} ${{ matrix.arch }} image + uses: docker/build-push-action@v5 with: - context: docker - file: docker/Dockerfile.ci_gpu - push: true - tags: flashinfer/flashinfer-ci:latest + context: . + file: docker/Dockerfile.${{ matrix.cuda }} + platforms: linux/${{ matrix.arch }} + push: ${{ github.event_name != 'pull_request' }} + tags: | + flashinfer/flashinfer-ci-${{ matrix.cuda }}:${{ matrix.arch }}-${{ needs.generate-tag.outputs.date_sha }} + cache-from: type=registry,ref=flashinfer/flashinfer-ci-${{ matrix.cuda }}:buildcache-${{ matrix.arch }} + cache-to: ${{ github.event_name != 'pull_request' && format('type=registry,ref=flashinfer/flashinfer-ci-{0}:buildcache-{1},mode=max', matrix.cuda, matrix.arch) || '' }} + provenance: false + sbom: false + + create-manifests: + if: github.event_name != 'pull_request' + runs-on: ubuntu-latest + needs: [generate-tag, build] + strategy: + matrix: + cuda: [cu126, cu128, cu129, cu130] + steps: + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: flashinfer + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Create and push multi-arch manifest for ${{ matrix.cuda }} + env: + DATE_SHA: ${{ needs.generate-tag.outputs.date_sha }} + run: | + docker buildx imagetools create \ + -t flashinfer/flashinfer-ci-${{ matrix.cuda }}:${DATE_SHA} \ + -t flashinfer/flashinfer-ci-${{ matrix.cuda }}:latest \ + flashinfer/flashinfer-ci-${{ matrix.cuda }}:amd64-${DATE_SHA} \ + flashinfer/flashinfer-ci-${{ matrix.cuda }}:arm64-${DATE_SHA} + + update-docker-tags: + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + needs: [generate-tag, create-manifests] + steps: + - uses: actions/checkout@v4 + + - name: Update docker-tags.yml + env: + DATE_SHA: ${{ needs.generate-tag.outputs.date_sha }} + run: | + cat > ci/docker-tags.yml << EOF + flashinfer/flashinfer-ci-cu126: ${DATE_SHA} + flashinfer/flashinfer-ci-cu128: ${DATE_SHA} + flashinfer/flashinfer-ci-cu129: ${DATE_SHA} + flashinfer/flashinfer-ci-cu130: ${DATE_SHA} + EOF + + - name: Create Pull Request + uses: peter-evans/create-pull-request@v6 + with: + token: ${{ secrets.FLASHINFER_BOT_TOKEN }} + commit-message: "Update Docker tags to ${{ needs.generate-tag.outputs.date_sha }}" + title: "Update Docker CI tags to ${{ needs.generate-tag.outputs.date_sha }}" + body: | + This PR updates the Docker CI image tags to the latest version: `${{ needs.generate-tag.outputs.date_sha }}` + + Updated images: + - flashinfer/flashinfer-ci-cu126:${{ needs.generate-tag.outputs.date_sha }} + - flashinfer/flashinfer-ci-cu128:${{ needs.generate-tag.outputs.date_sha }} + - flashinfer/flashinfer-ci-cu129:${{ needs.generate-tag.outputs.date_sha }} + - flashinfer/flashinfer-ci-cu130:${{ needs.generate-tag.outputs.date_sha }} + + Auto-generated by [release-ci-docker workflow](https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}) + branch: update-docker-tags-${{ needs.generate-tag.outputs.date_sha }} + delete-branch: true + labels: | + docker + automated diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000000..b11e72e1f7 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,444 @@ +name: Release + +on: + workflow_dispatch: + inputs: + tag: + description: 'Tag (e.g., v1.2.3) to build and publish' + required: true + type: string + pull_request: + paths: + - '.github/workflows/release.yml' + +jobs: + setup: + runs-on: ubuntu-latest + outputs: + version: ${{ steps.set-version.outputs.version }} + release_tag: ${{ steps.set-version.outputs.release_tag }} + is_pr: ${{ steps.set-version.outputs.is_pr }} + steps: + - name: Determine if running in PR mode + id: check-mode + run: | + if [ "${{ github.event_name }}" == "pull_request" ]; then + echo "is_pr=true" >> $GITHUB_OUTPUT + echo "Running in PR mode (dry-run)" + else + echo "is_pr=false" >> $GITHUB_OUTPUT + echo "Running in release mode" + fi + + - name: Validate tag format + if: github.event_name == 'workflow_dispatch' + run: | + if [[ ! "${{ inputs.tag }}" =~ ^v[0-9]+\.[0-9]+\.[0-9]+([a-z0-9]+)?$ ]]; then + echo "Error: Tag '${{ inputs.tag }}' does not match the expected format (e.g., v1.2.3 or v1.2.3.post1 or v1.2.3rc1)" + exit 1 + fi + echo "✓ Tag format is valid: ${{ inputs.tag }}" + + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.head_ref || inputs.tag }} + submodules: true + + - name: Verify tag matches version.txt + id: set-version + run: | + # Check version.txt - this is the source of truth + if [ ! -f "version.txt" ]; then + echo "Error: version.txt file not found!" + exit 1 + fi + + VERSION_TXT=$(cat version.txt | tr -d '[:space:]') + + # In PR mode, use version from version.txt with a test suffix + if [ "${{ github.event_name }}" == "pull_request" ]; then + echo "is_pr=true" >> $GITHUB_OUTPUT + echo "version=${VERSION_TXT}" >> $GITHUB_OUTPUT + echo "release_tag=test-v${VERSION_TXT}-pr${{ github.event.pull_request.number }}" >> $GITHUB_OUTPUT + echo "✓ PR mode: using test tag test-v${VERSION_TXT}-pr${{ github.event.pull_request.number }}" + else + # Extract version from tag (remove 'v' prefix) + TAG_VERSION="${{ inputs.tag }}" + TAG_VERSION="${TAG_VERSION#v}" + + if [ "$TAG_VERSION" != "$VERSION_TXT" ]; then + echo "❌ CRITICAL ERROR: version.txt does not match tag!" + echo " Tag version: $TAG_VERSION" + echo " version.txt: $VERSION_TXT" + echo "" + echo "Please update version.txt to match the release version before creating a release." + echo "The tag should be 'v$VERSION_TXT' (e.g., if version.txt contains '1.2.3', tag should be 'v1.2.3')" + exit 1 + fi + + echo "is_pr=false" >> $GITHUB_OUTPUT + echo "✓ version.txt matches tag version: $VERSION_TXT" + echo "version=${VERSION_TXT}" >> $GITHUB_OUTPUT + echo "release_tag=${{ inputs.tag }}" >> $GITHUB_OUTPUT + fi + + build-flashinfer-python: + needs: setup + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.head_ref || inputs.tag }} + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build wheel + + - name: Build flashinfer-python wheel and sdist + run: | + echo "Building flashinfer-python..." + echo "Git commit: $(git rev-parse HEAD)" + python -m build + ls -lh dist/ + + - name: Upload flashinfer-python artifact + uses: actions/upload-artifact@v4 + with: + name: flashinfer-python-dist + path: dist/* + retention-days: 7 + + build-flashinfer-cubin: + needs: setup + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.head_ref || inputs.tag }} + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine wheel + + - name: Build flashinfer-cubin wheel + run: | + echo "Building flashinfer-cubin..." + echo "Git commit: $(git rev-parse HEAD)" + cd flashinfer-cubin + rm -rf dist build *.egg-info + python -m build --wheel + ls -lh dist/ + mkdir -p ../dist + cp dist/*.whl ../dist/ + + - name: Upload flashinfer-cubin artifact + uses: actions/upload-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist/*.whl + retention-days: 7 + + build-flashinfer-jit-cache: + needs: setup + strategy: + fail-fast: false + matrix: + cuda: ["12.8", "12.9", "13.0"] + arch: ['x86_64', 'aarch64'] + + runs-on: [self-hosted, "${{ matrix.arch == 'aarch64' && 'arm64' || matrix.arch }}"] + + steps: + - name: Display Machine Information + run: | + echo "CPU: $(nproc) cores, $(lscpu | grep 'Model name' | cut -d':' -f2 | xargs)" + echo "RAM: $(free -h | awk '/^Mem:/ {print $7 " available out of " $2}')" + echo "Disk: $(df -h / | awk 'NR==2 {print $4 " available out of " $2}')" + echo "Architecture: $(uname -m)" + + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.head_ref || inputs.tag }} + submodules: true + + - name: Build wheel in container + env: + DOCKER_IMAGE: ${{ matrix.arch == 'aarch64' && format('pytorch/manylinuxaarch64-builder:cuda{0}', matrix.cuda) || format('pytorch/manylinux2_28-builder:cuda{0}', matrix.cuda) }} + FLASHINFER_CUDA_ARCH_LIST: ${{ matrix.cuda < '13.0' && '7.5 8.0 8.9 9.0a 10.0a 12.0a' || '7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f' }} + run: | + # Extract CUDA major and minor versions + CUDA_MAJOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f1) + CUDA_MINOR=$(echo "${{ matrix.cuda }}" | cut -d'.' -f2) + export CUDA_MAJOR + export CUDA_MINOR + export FLASHINFER_LOCAL_VERSION="cu${CUDA_MAJOR}${CUDA_MINOR}" + + chown -R $(id -u):$(id -g) ${{ github.workspace }} + mkdir -p ${{ github.workspace }}/ci-cache + chown -R $(id -u):$(id -g) ${{ github.workspace }}/ci-cache + + # Run the build script inside the container with proper mounts + docker run --rm \ + -v ${{ github.workspace }}:/workspace \ + -v ${{ github.workspace }}/ci-cache:/ci-cache \ + -e FLASHINFER_CI_CACHE=/ci-cache \ + -e CUDA_VERSION="${{ matrix.cuda }}" \ + -e CUDA_MAJOR="$CUDA_MAJOR" \ + -e CUDA_MINOR="$CUDA_MINOR" \ + -e FLASHINFER_LOCAL_VERSION="$FLASHINFER_LOCAL_VERSION" \ + -e ARCH="${{ matrix.arch }}" \ + -e FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST}" \ + --user $(id -u):$(id -g) \ + -w /workspace \ + ${{ env.DOCKER_IMAGE }} \ + bash /workspace/scripts/build_flashinfer_jit_cache_whl.sh + timeout-minutes: 180 + + - name: Display wheel size + run: du -h flashinfer-jit-cache/dist/* + + - name: Create artifact name + id: artifact-name + run: | + CUDA_NO_DOT=$(echo "${{ matrix.cuda }}" | tr -d '.') + echo "name=jit-cache-cu${CUDA_NO_DOT}-${{ matrix.arch }}" >> $GITHUB_OUTPUT + + - name: Upload flashinfer-jit-cache artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ steps.artifact-name.outputs.name }} + path: flashinfer-jit-cache/dist/*.whl + retention-days: 7 + + create-release: + needs: [setup, build-flashinfer-python, build-flashinfer-cubin, build-flashinfer-jit-cache] + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.tag }} + fetch-depth: 0 + + - name: Generate release notes + id: release-notes + env: + GH_TOKEN: ${{ github.token }} + run: | + TAG="${{ needs.setup.outputs.release_tag }}" + + # Get the previous release tag + PREV_TAG=$(git tag --sort=-version:refname | grep -A1 "^${TAG}$" | tail -n1 || echo "") + + if [ -z "$PREV_TAG" ]; then + echo "No previous tag found, generating notes from all commits" + NOTES=$(gh api /repos/${{ github.repository }}/releases/generate-notes \ + -f tag_name="$TAG" \ + -f target_commitish="${{ github.sha }}" \ + --jq '.body') + else + echo "Generating notes from ${PREV_TAG} to ${TAG}" + NOTES=$(gh api /repos/${{ github.repository }}/releases/generate-notes \ + -f tag_name="$TAG" \ + -f previous_tag_name="$PREV_TAG" \ + -f target_commitish="${{ github.sha }}" \ + --jq '.body') + fi + + # Save to file for multi-line output + echo "$NOTES" > /tmp/release_notes.md + echo "notes_file=/tmp/release_notes.md" >> $GITHUB_OUTPUT + + - name: Create GitHub Release (empty first) + env: + GH_TOKEN: ${{ github.token }} + run: | + TAG="${{ needs.setup.outputs.release_tag }}" + + # Delete existing release and tag if they exist + if gh release view "$TAG" &>/dev/null; then + echo "Deleting existing release: $TAG" + gh release delete "$TAG" --yes --cleanup-tag + fi + + # Read release notes from file + NOTES=$(cat ${{ steps.release-notes.outputs.notes_file }}) + + # Create new release without assets first + gh release create "$TAG" \ + --title "Release v${{ needs.setup.outputs.version }}" \ + --notes "$NOTES" + + - name: Download flashinfer-python artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-python-dist + path: dist-python/ + + - name: Upload flashinfer-python to release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-python/* --clobber + + - name: Download flashinfer-cubin artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist-cubin/ + + - name: Upload flashinfer-cubin to release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-cubin/* --clobber + + - name: Upload flashinfer-jit-cache wheels to release (one at a time to avoid OOM) + env: + GH_TOKEN: ${{ github.token }} + run: | + # Upload jit-cache wheels one at a time to avoid OOM + # Each wheel can be several GB, so we download, upload, delete, repeat + mkdir -p dist-jit-cache + + for cuda in 128 129 130; do + for arch in x86_64 aarch64; do + ARTIFACT_NAME="jit-cache-cu${cuda}-${arch}" + echo "Processing ${ARTIFACT_NAME}..." + + # Download this specific artifact + gh run download ${{ github.run_id }} -n "${ARTIFACT_NAME}" -D dist-jit-cache/ || { + echo "Warning: Failed to download ${ARTIFACT_NAME}, skipping..." + continue + } + + # Upload to release + if [ -n "$(ls -A dist-jit-cache/)" ]; then + gh release upload "${{ needs.setup.outputs.release_tag }}" dist-jit-cache/* --clobber + echo "✅ Uploaded ${ARTIFACT_NAME}" + fi + + # Clean up to save disk space before next iteration + rm -rf dist-jit-cache/* + done + done + + publish-to-pypi: + needs: [setup, create-release] + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + steps: + - name: Download flashinfer-python artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-python-dist + path: dist-python/ + + - name: Download flashinfer-cubin artifact + uses: actions/download-artifact@v4 + with: + name: flashinfer-cubin-wheel + path: dist-cubin/ + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install twine + run: | + python -m pip install --upgrade pip + pip install twine + + - name: Check packages with twine + run: | + echo "Running twine check..." + twine check dist-python/* dist-cubin/* + echo "✓ Package validation passed" + + - name: Upload flashinfer-python to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + echo "Uploading flashinfer (sdist and wheel) to PyPI..." + twine upload --verbose --non-interactive dist-python/* + echo "✓ Successfully uploaded flashinfer to PyPI" + + - name: Upload flashinfer-cubin to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + echo "Uploading flashinfer-cubin wheel to PyPI..." + twine upload --verbose --non-interactive dist-cubin/*.whl + echo "✓ Successfully uploaded flashinfer-cubin to PyPI" + + update-wheel-index: + needs: [setup, create-release] + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + steps: + - name: Checkout flashinfer repo + uses: actions/checkout@v4 + with: + ref: ${{ inputs.tag }} + + - name: Download all artifacts + uses: actions/download-artifact@v4 + with: + path: artifacts/ + + - name: Collect wheels and sdist + run: | + mkdir -p dist + find artifacts/ -name "*.whl" -exec cp {} dist/ \; + find artifacts/ -name "*.tar.gz" -exec cp {} dist/ \; + ls -lh dist/ + + - name: Clone wheel index + run: git clone https://oauth2:${WHL_TOKEN}@github.com/flashinfer-ai/whl.git flashinfer-whl + env: + WHL_TOKEN: ${{ secrets.WHL_TOKEN }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Update wheel index + run: | + python3 scripts/update_whl_index.py \ + --dist-dir dist \ + --output-dir flashinfer-whl \ + --release-tag "${{ needs.setup.outputs.release_tag }}" + + - name: Push wheel index + run: | + cd flashinfer-whl + git config --local user.name "github-actions[bot]" + git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com" + git add -A + git commit -m "update whl for release ${{ needs.setup.outputs.version }}" + git push diff --git a/.github/workflows/release_wheel.yml b/.github/workflows/release_wheel.yml deleted file mode 100644 index c0b8f2cd5b..0000000000 --- a/.github/workflows/release_wheel.yml +++ /dev/null @@ -1,135 +0,0 @@ -# Adapted from https://github.com/punica-ai/punica/blob/591b59899f0a20760821785d06b331c8a2e5cb86/.github/workflows/release_wheel.yml -name: Release -on: - workflow_dispatch: - inputs: - tag_name: - required: true - type: string - workflow_call: - inputs: - tag_name: - required: true - type: string - secrets: - WHL_TOKEN: - required: true - # PYPI_TEST_TOKEN: - # required: true - -env: - TORCH_CUDA_ARCH_LIST: "7.5 8.0 8.9 9.0+PTX" - -jobs: - build: - strategy: - fail-fast: false - matrix: - cuda: ["11.8", "12.1", "12.4", "12.6"] - torch: ["2.4", "2.5", "2.6"] - exclude: # We use release_wheel_sglang.yml for faster release and verification. If everything is okay, then we trigger release_wheel.yml. This combination (cuda 12.4 or 11.8 + torch 2.5) is already handled in release_wheel_sglang.yml - - cuda: "12.4" - torch: "2.5" - - cuda: "11.8" - torch: "2.5" - - cuda: "12.1" - torch: "2.6" - - cuda: "12.6" - torch: "2.4" - - cuda: "12.6" - torch: "2.5" - - - runs-on: [self-hosted] - steps: - - uses: actions/checkout@v4 - with: - submodules: true - - - name: Set torch and cuda version value - id: set_torch_and_cuda_version - run: | - IFS='.' read -r major minor <<< "${{ matrix.torch }}" - version_value=$((major * 100 + minor)) - echo "TORCH_VERSION=$version_value" >> $GITHUB_OUTPUT - IFS='.' read -r major minor <<< "${{ matrix.cuda }}" - version_value=$((major * 100 + minor)) - echo "CUDA_VERSION=$version_value" >> $GITHUB_OUTPUT - - - name: Build wheel - env: - TORCH_VERSION: ${{ steps.set_torch_and_cuda_version.outputs.TORCH_VERSION }} - CUDA_VERSION: ${{ steps.set_torch_and_cuda_version.outputs.CUDA_VERSION }} - BUILDER_IMAGE: ${{ ((steps.set_torch_and_cuda_version.outputs.TORCH_VERSION == 206 && steps.set_torch_and_cuda_version.outputs.CUDA_VERSION >= 1206) || steps.set_torch_and_cuda_version.outputs.TORCH_VERSION > 206) && 'pytorch/manylinux2_28-builder:cuda' || 'pytorch/manylinux-builder:cuda' }}${{ matrix.cuda }} - run: | - echo "TORCH_VERSION: $TORCH_VERSION" - echo "CUDA_VERSION: $CUDA_VERSION" - echo "BUILDER_IMAGE: $BUILDER_IMAGE" - chown -R $CI_UID:$CI_GID "$GITHUB_WORKSPACE" - docker run --rm -t \ - -v "$CI_RUNNER_CACHE_DIR":/ci-cache \ - -v "$GITHUB_WORKSPACE":/app \ - -e FLASHINFER_CI_CACHE=/ci-cache \ - -e FLASHINFER_CI_CUDA_VERSION=${{ matrix.cuda }} \ - -e FLASHINFER_CI_TORCH_VERSION=${{ matrix.torch }} \ - -e FLASHINFER_CI_PYTHON_VERSION=3.10 \ - -e TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" \ - -e MAX_JOBS=128 \ - --user $CI_UID:$CI_GID \ - $BUILDER_IMAGE \ - bash /app/scripts/run-ci-build-wheel.sh - timeout-minutes: 120 - - run: du -h dist/* - - - uses: actions/upload-artifact@v4 - with: - name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }} - path: dist/* - - release: - needs: build - runs-on: [self-hosted] - steps: - - uses: actions/download-artifact@v4 - with: - path: dist/ - merge-multiple: true - pattern: wheel-* - - - run: ls -lah dist/ - - - uses: softprops/action-gh-release@v1 - with: - tag_name: ${{ inputs.tag_name }} - files: | - dist/flashinfer*.whl - - - uses: softprops/action-gh-release@v1 - with: - tag_name: ${{ inputs.tag_name }} - files: | - dist/flashinfer-*.tar.gz - - - name: Clone wheel index - run: git clone https://oauth2:${WHL_TOKEN}@github.com/flashinfer-ai/whl.git flashinfer-whl - env: - WHL_TOKEN: ${{ secrets.WHL_TOKEN }} - - - name: Update wheel index - run: python3 scripts/update_whl_index.py - - - name: Push wheel index - run: | - cd flashinfer-whl - git config --local user.name "github-actions[bot]" - git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com" - git add -A - git commit -m "update whl" - git push - - # - name: Upload sdist to pypi - # run: | - # pip install twine - # python -m twine upload --repository testpypi --username=__token__ dist/*.tar.gz - # env: - # TWINE_PASSWORD: ${{ secrets.PYPI_TEST_TOKEN }} diff --git a/.github/workflows/release_wheel_sglang.yml b/.github/workflows/release_wheel_sglang.yml deleted file mode 100644 index 2d7a2ef352..0000000000 --- a/.github/workflows/release_wheel_sglang.yml +++ /dev/null @@ -1,114 +0,0 @@ -name: Release Wheel -on: - workflow_dispatch: - inputs: - tag_name: - required: true - type: string - workflow_call: - inputs: - tag_name: - required: true - type: string - secrets: - WHL_TOKEN: - required: true - -env: - TORCH_CUDA_ARCH_LIST: "7.5 8.0 8.9 9.0+PTX" - -jobs: - build: - strategy: - fail-fast: false - matrix: - cuda: ["11.8", "12.4"] - torch: ["2.5"] - - runs-on: [self-hosted] - steps: - - uses: actions/checkout@v4 - with: - submodules: true - - - name: Set torch and cuda version value - id: set_torch_and_cuda_version - run: | - IFS='.' read -r major minor <<< "${{ matrix.torch }}" - version_value=$((major * 100 + minor)) - echo "TORCH_VERSION=$version_value" >> $GITHUB_OUTPUT - IFS='.' read -r major minor <<< "${{ matrix.cuda }}" - version_value=$((major * 100 + minor)) - echo "CUDA_VERSION=$version_value" >> $GITHUB_OUTPUT - - - name: Build wheel - env: - TORCH_VERSION: ${{ steps.set_torch_and_cuda_version.outputs.TORCH_VERSION }} - CUDA_VERSION: ${{ steps.set_torch_and_cuda_version.outputs.CUDA_VERSION }} - BUILDER_IMAGE: ${{ ((steps.set_torch_and_cuda_version.outputs.TORCH_VERSION == 206 && steps.set_torch_and_cuda_version.outputs.CUDA_VERSION >= 1206) || steps.set_torch_and_cuda_version.outputs.TORCH_VERSION > 206) && 'pytorch/manylinux2_28-builder:cuda' || 'pytorch/manylinux-builder:cuda' }}${{ matrix.cuda }} - run: | - echo "TORCH_VERSION: $TORCH_VERSION" - echo "CUDA_VERSION: $CUDA_VERSION" - echo "BUILDER_IMAGE: $BUILDER_IMAGE" - chown -R $CI_UID:$CI_GID "$GITHUB_WORKSPACE" - docker run --rm -t \ - -v "$CI_RUNNER_CACHE_DIR":/ci-cache \ - -v "$GITHUB_WORKSPACE":/app \ - -e FLASHINFER_CI_CACHE=/ci-cache \ - -e FLASHINFER_CI_CUDA_VERSION=${{ matrix.cuda }} \ - -e FLASHINFER_CI_TORCH_VERSION=${{ matrix.torch }} \ - -e FLASHINFER_CI_PYTHON_VERSION=3.10 \ - -e FLASHINFER_HEAD_DIMS="64,128,256" \ - -e TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" \ - -e MAX_JOBS=128 \ - --user $CI_UID:$CI_GID \ - $BUILDER_IMAGE \ - bash /app/scripts/run-ci-build-wheel.sh - timeout-minutes: 120 - - run: du -h dist/* - - - uses: actions/upload-artifact@v4 - with: - name: wheel-cuda${{ matrix.cuda }}-torch${{ matrix.torch }} - path: dist/* - - release: - needs: build - runs-on: [self-hosted] - steps: - - uses: actions/download-artifact@v4 - with: - path: dist/ - merge-multiple: true - pattern: wheel-* - - - run: ls -lah dist/ - - - uses: softprops/action-gh-release@v1 - with: - tag_name: ${{ inputs.tag_name }} - files: | - dist/flashinfer*.whl - - - uses: softprops/action-gh-release@v1 - with: - tag_name: ${{ inputs.tag_name }} - files: | - dist/flashinfer-*.tar.gz - - - name: Clone wheel index - run: git clone https://oauth2:${WHL_TOKEN}@github.com/flashinfer-ai/whl.git flashinfer-whl - env: - WHL_TOKEN: ${{ secrets.WHL_TOKEN }} - - - name: Update wheel index - run: python3 scripts/update_whl_index.py - - - name: Push wheel index - run: | - cd flashinfer-whl - git config --local user.name "github-actions[bot]" - git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com" - git add -A - git commit -m "update whl" - git push diff --git a/.github/workflows/update-codeowners.yml b/.github/workflows/update-codeowners.yml new file mode 100644 index 0000000000..a1beed6339 --- /dev/null +++ b/.github/workflows/update-codeowners.yml @@ -0,0 +1,118 @@ +name: Update CODEOWNERS + +on: + schedule: + # Run weekly on Monday at 00:00 UTC + - cron: '0 0 * * 1' + workflow_dispatch: # Allow manual triggering + +permissions: + contents: write + pull-requests: write + +jobs: + update-codeowners: + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + DEPTH: 3 + MIN_COMMITS: 1 + DAYS_BACK: 180 + TOP_N: 5 + steps: + - name: Checkout repository + uses: actions/checkout@v4.2.2 + with: + fetch-depth: 0 # Fetch full history for accurate analysis + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Run CODEOWNERS analyzer + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + python scripts/codeowner_analyzer.py \ + --output .github/CODEOWNERS \ + --depth ${{ env.DEPTH }} \ + --min-commits ${{ env.MIN_COMMITS }} \ + --days-back ${{ env.DAYS_BACK }} \ + --top-n ${{ env.TOP_N }} \ + --allowed-users-file scripts/authorized_codeowner.txt + + - name: Check for changes + id: check_changes + run: | + # Check if CODEOWNERS file is new (unstaged) or has changes + if git ls-files --error-unmatch .github/CODEOWNERS >/dev/null 2>&1; then + # File is tracked, check for changes + if git diff --quiet .github/CODEOWNERS; then + echo "changed=false" >> $GITHUB_OUTPUT + echo "No changes detected in CODEOWNERS" + else + echo "changed=true" >> $GITHUB_OUTPUT + echo "Changes detected in CODEOWNERS" + fi + else + # File is untracked (newly created) + echo "changed=true" >> $GITHUB_OUTPUT + echo "CODEOWNERS file is new" + fi + + - name: Create Pull Request + if: steps.check_changes.outputs.changed == 'true' + uses: peter-evans/create-pull-request@v7 + with: + token: ${{ secrets.FLASHINFER_BOT_TOKEN }} + committer: flashinfer-bot + author: flashinfer-bot + commit-message: | + chore: update CODEOWNERS based on git history + + Auto-generated CODEOWNERS update based on commit activity over the last ${{ env.DAYS_BACK }} days. + + 🤖 Generated with [Claude Code](https://claude.com/claude-code) + + Co-Authored-By: Claude + branch: auto-update-codeowners + base: main + delete-branch: true + title: 'chore: Update CODEOWNERS' + body: | + ## Summary + + This PR updates the CODEOWNERS file based on git commit history analysis from the last ${{ env.DAYS_BACK }} days. + + ## Changes + + - Updated `.github/CODEOWNERS` with current code ownership based on: + - Commit frequency + - File coverage + - Commit recency + + ## How to Review + + 1. Review the changes to `.github/CODEOWNERS` + 2. Verify that the assigned owners are appropriate for each module + 3. Make manual adjustments if needed before merging + + ## Notes + + - This is an automated PR generated weekly + - Minimum commits threshold: ${{ env.MIN_COMMITS }} + - Analysis period: ${{ env.DAYS_BACK }} days + - Directory depth: ${{ env.DEPTH }} levels + - Top N owners per module: ${{ env.TOP_N }} + + --- + + 🤖 This PR was automatically generated by the [update-codeowners workflow](.github/workflows/update-codeowners.yml) + labels: | + automated + maintenance + assignees: | + + reviewers: | diff --git a/.gitignore b/.gitignore index 2075e5321f..53964e1be8 100644 --- a/.gitignore +++ b/.gitignore @@ -7,13 +7,13 @@ compile_commands.json # Generated files csrc/generated/ +csrc/nv_internal/tensorrt_llm/cutlass_instantiations/ docs/generated/ flashinfer/_build_meta.py flashinfer/data/ flashinfer/_version.py flashinfer/__config__.py flashinfer/jit/aot_config.py -src/generated/ csrc/aot_default_additional_params.h *.out @@ -80,6 +80,7 @@ coverage.xml .hypothesis/ .pytest_cache/ cover/ +junit/ # Translations *.mo diff --git a/.gitmodules b/.gitmodules index 741dda47ff..a45558d527 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ [submodule "3rdparty/cutlass"] path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git +[submodule "3rdparty/spdlog"] + path = 3rdparty/spdlog + url = https://github.com/gabime/spdlog.git diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2f5a784ab0..d1ac2c37a7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -35,31 +35,38 @@ repos: - id: remove-crlf # Formatters - - repo: https://github.com/psf/black-pre-commit-mirror - rev: 24.8.0 - hooks: - - id: black - - - repo: https://github.com/pycqa/isort - rev: 5.13.2 - hooks: - - id: isort - args: ["--profile=black"] - exclude: ^flashinfer/(cascade\.py|_version\.py)$ - - repo: https://github.com/pre-commit/mirrors-clang-format rev: v19.1.1 hooks: - id: clang-format types_or: [c++, c, cuda] exclude: | - (?x)^(3rdparty/.* src/generated/.* flashinfer/jit/aot_config.py)$ + (?x)^(3rdparty/.* flashinfer/jit/aot_config.py)$ + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: 'v1.17.1' # Use the sha / tag you want to point at + hooks: + - id: mypy + args: ["--config-file", "pyproject.toml", "--exclude", "flashinfer-cubin/"] + files: ^flashinfer/ + exclude: ^(flashinfer-cubin/|3rdparty/|build/) + + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.12.8 + hooks: + # Run the linter. + - id: ruff-check + # Run the formatter. + - id: ruff-format + types_or: [ python, pyi ] - repo: https://github.com/cheshirekow/cmake-format-precommit rev: v0.6.13 hooks: - id: cmake-format additional_dependencies: [pyyaml>=5.1] + # REUSE compliance check for specific directories only - repo: local hooks: diff --git a/3rdparty/cutlass b/3rdparty/cutlass index df8a550d39..f3fde58372 160000 --- a/3rdparty/cutlass +++ b/3rdparty/cutlass @@ -1 +1 @@ -Subproject commit df8a550d3917b0e97f416b2ed8c2d786f7f686a3 +Subproject commit f3fde58372d33e9a5650ba7b80fc48b3b49d40c8 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000000..32ca4e6d12 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,46 @@ +# Installation + +For development, the easiest way to install flashinfer is through editable installation: + +``` +git clone git@github.com:flashinfer-ai/flashinfer.git --recursive +pip install --no-build-isolation -e . -v +``` + +We recommend using the `--no-build-isolation` flag to ensure compatibility with your existing environment. Without it, `pip` may attempt to resolve dependencies (e.g., `torch`) from PyPI, which could pull in packages built with older CUDA versions and lead to incompatibility issues. + +# Code Structure + +``` +flashinfer/ +| --include/ # kernel definitions and common utilities functions +| --csrc/ # op registration to frameworks (pytorch), and binding codes +| --python/ # python interface exposed to users +| --docs/ # documentation (using sphinx) +| --tests/ # unittests in python (using pytest) +| --benchmarks/ # kernel benchmarks in python +| --3rdparty/ # 3rdparty dependencies such as cutlass +``` + +Kernel definitions (framework-agnostic cuda code, accepting raw pointer as input) should be placed under the `include` directory. Whenever possible, reuse existing FlashInfer infrastructure such as logging, exception handling, and utility functions. +The operator registration code (i.e., framework-specific components, accepting torch tensors as input) should reside in the `csrc` directory. This is where Torch headers may be included and operators can be bound to PyTorch. Note that Torch headers must not be included in any files under the `include` directory. + +Code Contribution Procedure +* Write kernel definitions in `include/` +* Write kernel registration and pytorch interface under `csrc/` +* Write python interface under `python/` +* Write unit tests in `tests/` +* (Optional) Add benchmark suites under `benchmark/` +* Update (python) documentation index under `docs/` +* Update `pyproject.toml` if you created new module in flashinfer + +# Release Versioning + +When incrementing a version and creating a release, follow a "right-shifted" versioning scheme similar to [vLLM Release Versioning](https://github.com/vllm-project/vllm/blob/main/RELEASE.md) (`major.minor.patch[.post1]`) [^1]. In particular: + +* _major_ increment signals architectural milestone and/or when incompatible API changes are made, similar to PyTorch 2.0. +* _minor_ increment signals significant backwards-compatible new features +* _patch_ increment signals small backwards-compatible features (e.g. new kernels, new SM support, etc) and backwards-compatible bug fixes +* _post1_ is an optional suffix for a quick follow up release with just backwards-compatible bug fixes + +[^1]: We have not followed this strictly through v0.4.0. But after v0.4.0, the versioning should follow this "right-shifted" versioning scheme. diff --git a/Jenkinsfile b/Jenkinsfile index a832264056..62c28f171c 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -37,9 +37,11 @@ // import org.jenkinsci.plugins.pipeline.modeldefinition.Utils -// These are set at runtime from data in ci/jenkins/docker-images.yml, update -// image tags in that file -docker_run = "bash ci/bash.sh flashinfer/flashinfer-ci:latest" + +def getDockerRun(cuda_version, dockerTags) { + def image_name = "flashinfer/flashinfer-ci-${cuda_version}" + return "bash ci/bash.sh ${image_name}:${dockerTags[image_name]}" +} def per_exec_ws(folder) { return "workspace/exec_${env.EXECUTOR_NUMBER}/" + folder @@ -61,6 +63,56 @@ def unpack_lib(name, libs) { """ } +def should_skip_build() { + // Skip build if changes are only in documentation/config directories + def skip_patterns = [ + 'README.md', + '.github/', + 'docs/', + 'docker/', + 'licenses/', + 'LICENSE', + 'NOTICE', + 'version.txt' + ] + + if (env.CHANGE_ID) { + // This is a PR build, check changed files + def changedFiles = [] + try { + changedFiles = sh( + script: 'git diff --name-only origin/${CHANGE_TARGET}...HEAD', + returnStdout: true + ).trim().split('\n') + } catch (Exception e) { + echo "Could not determine changed files: ${e.toString()}" + return false + } + + if (changedFiles.size() == 0) { + return false + } + + // Check if all changed files match skip patterns + def allSkippable = changedFiles.every { file -> + skip_patterns.any { pattern -> + if (pattern.endsWith('/')) { + file.startsWith(pattern) + } else { + file == pattern + } + } + } + + if (allSkippable) { + echo "Skipping build - all changes are in documentation/config files: ${changedFiles}" + return true + } + } + + return false +} + def cancel_previous_build() { // cancel previous build if it is not on main. if (env.BRANCH_NAME != 'main') { @@ -73,10 +125,13 @@ def cancel_previous_build() { } def is_last_build() { - // whether it is last build - def job = Jenkins.instance.getItem(env.JOB_NAME) - def lastBuild = job.getLastBuild() - return lastBuild.getNumber() == env.BUILD_NUMBER + // check whether it is last build + try { + return currentBuild.number == currentBuild.rawBuild.project.getLastBuild().number + } catch (Throwable ex) { + echo 'Error during check is_last_build ' + ex.toString() + return false + } } def init_git(submodule = false) { @@ -94,6 +149,25 @@ def init_git(submodule = false) { } } +def run_with_spot_retry(spot_node_type, on_demand_node_type, test_name, test_closure) { + try { + test_closure(spot_node_type) + } catch (hudson.AbortException abortEx) { + echo "Received normal AbortException, exit now: " + abortEx.toString() + throw abortEx + } catch (Throwable ex) { + echo "Exception during SPOT run for ${test_name}: " + ex.toString() + if (is_last_build()) { + echo "Exception during SPOT run for ${test_name}: " + ex.toString() + " retry on-demand" + currentBuild.result = 'SUCCESS' + test_closure(on_demand_node_type) + } else { + echo 'Exit since it is not last build' + throw ex + } + } +} + // stage('Lint') { // node('CPU-SPOT') { // ws(per_exec_ws('flashinfer-lint')) { @@ -102,118 +176,176 @@ def init_git(submodule = false) { // } // } -def run_unittest_CPU_AOT_COMPILE(node_type) { - echo "Running CPU AOT Compile Unittest" - node(node_type) { - ws(per_exec_ws('flashinfer-aot')) { - init_git(true) - sh(script: "ls -alh", label: 'Show work directory') - sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info') - sh(script: "${docker_run} --no-gpu ./scripts/task_test_aot_build_import.sh", label: 'Test AOT Build and Import') +def run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, cuda_version) { + echo "Running CPU JIT Cache Package Build and Import Unittest with CUDA ${cuda_version}" + + if (node_type.contains('SPOT')) { + // Add timeout only for spot instances - node allocation only + def node_allocated = false + + try { + timeout(time: 15, unit: 'MINUTES') { + // Only timeout the node allocation, not the test execution + node(node_type) { + node_allocated = true + // Just mark that we got the node, don't run tests here + } + } + + // If we reach here, node allocation was successful + // Now run the tests without any timeout + node(node_type) { + ws(per_exec_ws('flashinfer-jit-cache')) { + init_git(true) + def dockerTags = readYaml file: 'ci/docker-tags.yml' + def docker_run = getDockerRun(cuda_version, dockerTags) + sh(script: "ls -alh", label: 'Show work directory') + sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info') + sh(script: "${docker_run} --no-gpu ./scripts/task_test_jit_cache_package_build_import.sh", label: 'Test JIT Cache Package Build and Import') + } + } + } catch (Exception e) { + if (!node_allocated) { + echo "Node allocation timeout or failure after 15 minutes for ${node_type}: ${e.toString()}" + } + throw e + } + } else { + // No timeout for non-spot instances + node(node_type) { + ws(per_exec_ws('flashinfer-jit-cache')) { + init_git(true) + def dockerTags = readYaml file: 'ci/docker-tags.yml' + def docker_run = getDockerRun(cuda_version, dockerTags) + sh(script: "ls -alh", label: 'Show work directory') + sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info') + sh(script: "${docker_run} --no-gpu ./scripts/task_test_jit_cache_package_build_import.sh", label: 'Test JIT Cache Package Build and Import') + } } } } -def shard_run_unittest_GPU(node_type, shard_id) { - echo "Running unittest on ${node_type}, shard ${shard_id}" - node(node_type) { - ws(per_exec_ws('flashinfer-unittest')) { - init_git(true) // we need cutlass submodule - sh(script: "ls -alh", label: 'Show work directory') - sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info') - sh(script: "${docker_run} ./scripts/task_jit_run_tests_part${shard_id}.sh", label: 'JIT Unittest Part ${shard_id}') +def shard_run_unittest_GPU(node_type, shard_id, cuda_version) { + echo "Running unittest on ${node_type}, shard ${shard_id}, CUDA ${cuda_version}" + + if (node_type.contains('SPOT')) { + // Add timeout only for spot instances - node allocation only + def node_allocated = false + + try { + timeout(time: 15, unit: 'MINUTES') { + // Only timeout the node allocation, not the test execution + node(node_type) { + node_allocated = true + // Just mark that we got the node, don't run tests here + } + } + + // If we reach here, node allocation was successful + // Now run the tests without any timeout + node(node_type) { + ws(per_exec_ws('flashinfer-unittest')) { + init_git(true) // we need cutlass submodule + def dockerTags = readYaml file: 'ci/docker-tags.yml' + def docker_run = getDockerRun(cuda_version, dockerTags) + sh(script: "ls -alh", label: 'Show work directory') + sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info') + sh(script: "${docker_run} ./scripts/task_jit_run_tests_part${shard_id}.sh", label: 'JIT Unittest Part ${shard_id}') + } + } + } catch (Exception e) { + if (!node_allocated) { + echo "Node allocation timeout or failure after 15 minutes for ${node_type}: ${e.toString()}" + } + throw e + } + } else { + // No timeout for non-spot instances + node(node_type) { + ws(per_exec_ws('flashinfer-unittest')) { + init_git(true) // we need cutlass submodule + def dockerTags = readYaml file: 'ci/docker-tags.yml' + def docker_run = getDockerRun(cuda_version, dockerTags) + sh(script: "ls -alh", label: 'Show work directory') + sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info') + sh(script: "${docker_run} ./scripts/task_jit_run_tests_part${shard_id}.sh", label: 'JIT Unittest Part ${shard_id}') + } } } } stage('Unittest') { + if (should_skip_build()) { + echo "Skipping tests - only documentation/config files changed" + Utils.markStageSkippedForConditional('Unittest') + return + } + cancel_previous_build() parallel( failFast: true, - 'AOT-Build-Import': { - try { - run_unittest_CPU_AOT_COMPILE('CPU-LARGE-SPOT') - } catch (Throwable ex) { - if (is_last_build()) { - // retry if we are currently at last build - // mark the current stage as success - // and try again via on demand node - echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand' - currentBuild.result = 'SUCCESS' - run_unittest_CPU_AOT_COMPILE('CPU-LARGE') - } else { - echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build' - throw ex - } - } + // CUDA 12.6 AOT Tests + 'AOT-Build-Import-x86-64-cu126': { + run_with_spot_retry('CPU-LARGE-SPOT', 'CPU-LARGE', 'AOT-Build-Import-x86-64-cu126', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu126') }) }, - 'JIT-Unittest-1': { - try { - shard_run_unittest_GPU('GPU-G5-SPOT', 1) - } catch (Throwable ex) { - if (is_last_build()) { - // retry if we are currently at last build - // mark the current stage as success - // and try again via on demand node - echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand' - currentBuild.result = 'SUCCESS' - shard_run_unittest_GPU('GPU-G5', 1) - } else { - echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build' - throw ex - } - } + 'AOT-Build-Import-aarch64-cu126': { + run_with_spot_retry('ARM-LARGE-SPOT', 'ARM-LARGE', 'AOT-Build-Import-aarch64-cu126', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu126') }) }, - 'JIT-Unittest-2': { - try { - shard_run_unittest_GPU('GPU-G5-SPOT', 2) - } catch (Throwable ex) { - if (is_last_build()) { - // retry if we are currently at last build - // mark the current stage as success - // and try again via on demand node - echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand' - currentBuild.result = 'SUCCESS' - shard_run_unittest_GPU('GPU-G5', 2) - } else { - echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build' - throw ex - } - } + // CUDA 12.8 AOT Tests + 'AOT-Build-Import-x86-64-cu128': { + run_with_spot_retry('CPU-LARGE-SPOT', 'CPU-LARGE', 'AOT-Build-Import-x86-64-cu128', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu128') }) }, - 'JIT-Unittest-3': { - try { - shard_run_unittest_GPU('GPU-G5-SPOT', 3) - } catch (Throwable ex) { - if (is_last_build()) { - // retry if we are currently at last build - // mark the current stage as success - // and try again via on demand node - echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand' - currentBuild.result = 'SUCCESS' - shard_run_unittest_GPU('GPU-G5', 3) - } else { - echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build' - throw ex - } - } + 'AOT-Build-Import-aarch64-cu128': { + run_with_spot_retry('ARM-LARGE-SPOT', 'ARM-LARGE', 'AOT-Build-Import-aarch64-cu128', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu128') }) + }, + // CUDA 12.9 AOT Tests + 'AOT-Build-Import-x86-64-cu129': { + run_with_spot_retry('CPU-LARGE-SPOT', 'CPU-LARGE', 'AOT-Build-Import-x86-64-cu129', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu129') }) + }, + 'AOT-Build-Import-aarch64-cu129': { + run_with_spot_retry('ARM-LARGE-SPOT', 'ARM-LARGE', 'AOT-Build-Import-aarch64-cu129', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu129') }) + }, + // CUDA 13.0 AOT Tests + 'AOT-Build-Import-x86-64-cu130': { + run_with_spot_retry('CPU-LARGE-SPOT', 'CPU-LARGE', 'AOT-Build-Import-x86-64-cu130', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu130') }) + }, + 'AOT-Build-Import-aarch64-cu130': { + run_with_spot_retry('ARM-LARGE-SPOT', 'ARM-LARGE', 'AOT-Build-Import-aarch64-cu130', + { node_type -> run_unittest_CPU_JIT_CACHE_PACKAGE_BUILD_IMPORT(node_type, 'cu130') }) + }, + // JIT unittest only for cu129 + 'JIT-Unittest-1-cu129': { + run_with_spot_retry('GPU-G5-SPOT', 'GPU-G5', 'JIT-Unittest-1-cu129', + { node_type -> shard_run_unittest_GPU(node_type, 1, 'cu129') }) + }, + 'JIT-Unittest-2-cu129': { + run_with_spot_retry('GPU-G5-SPOT', 'GPU-G5', 'JIT-Unittest-2-cu129', + { node_type -> shard_run_unittest_GPU(node_type, 2, 'cu129') }) + }, + 'JIT-Unittest-3-cu129': { + run_with_spot_retry('GPU-G5-SPOT', 'GPU-G5', 'JIT-Unittest-3-cu129', + { node_type -> shard_run_unittest_GPU(node_type, 3, 'cu129') }) + }, + 'JIT-Unittest-4-cu129': { + run_with_spot_retry('GPU-G5-SPOT', 'GPU-G5', 'JIT-Unittest-4-cu129', + { node_type -> shard_run_unittest_GPU(node_type, 4, 'cu129') }) + }, + 'JIT-Unittest-5-cu129': { + run_with_spot_retry('GPU-G5-SPOT', 'GPU-G5', 'JIT-Unittest-5-cu129', + { node_type -> shard_run_unittest_GPU(node_type, 5, 'cu129') }) + }, + // JIT unit test for CUDA 12.9 with SM75 (AWS G4) + // For now, we only enable sampling test for SM75 + 'JIT-Unittest-G4-cu129': { + run_with_spot_retry('GPU-G4-SPOT', 'GPU-G4', 'JIT-Unittest-G4-cu129', + { node_type -> shard_run_unittest_GPU(node_type, 3, 'cu129') }) }, - 'JIT-Unittest-4': { - try { - shard_run_unittest_GPU('GPU-G5-SPOT', 4) - } catch (Throwable ex) { - if (is_last_build()) { - // retry if we are currently at last build - // mark the current stage as success - // and try again via on demand node - echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand' - currentBuild.result = 'SUCCESS' - shard_run_unittest_GPU('GPU-G5', 4) - } else { - echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build' - throw ex - } - } - } ) } diff --git a/LICENSE b/LICENSE index 7c8f7e1403..ba557f2bb9 100644 --- a/LICENSE +++ b/LICENSE @@ -221,3 +221,9 @@ BSD 3-Clause "New" License 3rdparty/cutlass include/flashinfer/attention/hopper/block_sparse_gather.cuh + +MIT License +----------- + +3rdparty/spdlog +3rdparty/spdlog/include/spdlog/fmt/bundled (fmt library) diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000000..c0fea60adb --- /dev/null +++ b/NOTICE @@ -0,0 +1,8 @@ +FlashInfer +Copyright 2025 NVIDIA +Copyright 2023-2025 FlashInfer community (https://flashinfer.ai/) + +------------------------------------------------------------------------------------------------- +Some of the code in this project are adapted from other open-source projects with different +licenses. This product also bundles some third-party components under other open source licenses. +See licenses/ for text of these licenses. diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000000..e7e17156a4 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,241 @@ +# FlashInfer Perf Benchmarking Framework -- `flashinfer_benchmark.py` + +The aim of `flashinfer_benchmark.py` is to provide a single framework for benchmarking any FlashInfer kernel and replace standalone benchmarking scripts. + +## Overview + +This framework provides tools to: +- Benchmark FlashInfer's Attention, GEMM, and MOE API performance from different kernel backends such as FlashAttention2/3, cuDNN, cuBLAS, CUTLASS, and TensorRT-LLM +- Compare performance across different configurations +- Batch performance test multiple attention test cases + +Currently supports testing most attention, gemm, and fused MOE APIs: +- Attention: + - `BatchDecodeWithPagedKVCacheWrapper` - Decode attention with paged KV cache. + - Also supports computationally similar `cudnn_batch_decode_with_kv_cache` and `trtllm_batch_decode_with_kv_cache`. + - `BatchPrefillWithPagedKVCacheWrapper` - Prefill attention with paged KV cache. + - Also supports computationally similar `cudnn_batch_prefill_with_kv_cache` and `trtllm_batch_context_with_kv_cache`. + - `BatchPrefillWithRaggedKVCacheWrapper` - Prefill attention with ragged KV cache. + - Also supports computationally similar `cudnn_batch_prefill_with_kv_cache` and `trtllm_ragged_attention_deepseek`. + - `BatchMLAPagedAttentionWrapper` - MLA attention proposed in DeepSeek series of models. + - Also supports computationally similar `trtllm_batch_decode_with_kv_cache_mla`. +- GEMM: + - `gemm_fp8_nt_groupwise` - GEMM with FP8 data types using groupwise scaling. + - `group_gemm_fp8_nt_groupwise` - Group GEMM with FP8 data types using groupwise scaling. + - `bmm_fp8` - Batched matrix multiplication with FP8 inputs. + - `mm_fp4` - Matrix multiplication with NVFP4 inputs. +- MOE: + - `trtllm_fp4_block_scale_moe` - MOE with FP4 quantized weights and block-wise scaling. + - `trtllm_fp8_block_scale_moe` - MOE with FP8 quantized weights and block-wise scaling. + - `trtllm_fp8_per_tensor_scale_moe` - MOE with FP8 quantized weights and per-tensor scaling. + - `cutlass_fused_moe` - CUTLASS fused MoE (base/fp8/nvfp4 variants with optional TP/EP) + +## Quick Start +### Single Test Run +A test case is generally invoked as `python3 flashinfer_benchmark.py --routine `. + +*See samples in samples/sample_testlist.txt for various example test flags.* +Example commands and outputs areas follows + +```bash +# bmm_fp8 +$ python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +[INFO] args = Namespace(routine='bmm_fp8', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=256, m=1, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cublas', 'cutlass'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testBmmFp8 +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +[VVERBOSE] input_fp8.shape = torch.Size([256, 1, 7168]) +[VVERBOSE] input_fp8.dtype = torch.float8_e4m3fn +[VVERBOSE] mat2_fp8.shape = torch.Size([256, 7168, 1024]) +[VVERBOSE] mat2_fp8.dtype = torch.float8_e4m3fn +[VVERBOSE] input_inv_s = tensor(0.0109, device='cuda:0') +[VVERBOSE] input_inv_s.dtype = torch.float32 +[VVERBOSE] mat2_inv_s = tensor(0.0135, device='cuda:0') +[VVERBOSE] mat2_inv_s.dtype = torch.float32 +[PERF] cudnn :: median time 0.285 ms; std 0.000 ms; achieved tflops 13.180 TFLOPs/sec; achieved tb_per_sec 0.026 TB/sec +[PERF] cublas :: median time 0.286 ms; std 0.000 ms; achieved tflops 13.159 TFLOPs/sec; achieved tb_per_sec 0.026 TB/sec +[PERF] cutlass :: median time 0.266 ms; std 0.001 ms; achieved tflops 14.137 TFLOPs/sec; achieved tb_per_sec 0.028 TB/sec + +# non-paged (ragged) prefill +$ python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "DeepSeek-R1" +[INFO] args = Namespace(routine='BatchPrefillWithRaggedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='DeepSeek-R1', generate_repro_command=True, repro_command='', backends=['fa2', 'fa3', 'cutlass', 'cudnn'], page_size=0, batch_size=16, s_qo=1024, s_kv=1024, num_qo_heads=128, num_kv_heads=128, head_dim_qk=192, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=True, random_actual_seq_len=True) +[INFO] Running testBatchPrefillWithRaggedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +[WARNING] fa3 for routine BatchPrefillWithRaggedKVCacheWrapper is not supported on compute capability 10.0. Skipping. +[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q +[VERBOSE] Average actual qo seq len: 327 +[VERBOSE] Average actual kv seq len: 327 +[VVERBOSE] actual_seq_lens_q.flatten() = tensor([103, 436, 861, 271, 107, 72, 701, 21, 615, 122, 467, 215, 331, 459, + 88, 373], dtype=torch.int32) +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([103, 436, 861, 271, 107, 72, 701, 21, 615, 122, 467, 215, 331, 459, + 88, 373], dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([5242, 128, 192]) +[VVERBOSE] k.shape = torch.Size([5242, 128, 192]) +[VVERBOSE] v.shape = torch.Size([5242, 128, 128]) +[VVERBOSE] qo_indptr.shape = torch.Size([17]) +[VVERBOSE] kv_indptr.shape = torch.Size([17]) +[VVERBOSE] scale = 0.07216878364870323 +[PERF] fa2 :: median time 0.495 ms; std 0.006 ms; achieved tflops 219.336 TFLOPs/sec; achieved tb_per_sec 1.736 TB/sec +[PERF] cutlass :: median time 0.530 ms; std 0.002 ms; achieved tflops 204.674 TFLOPs/sec; achieved tb_per_sec 1.620 TB/sec +[PERF] cudnn :: median time 0.313 ms; std 0.000 ms; achieved tflops 346.715 TFLOPs/sec; achieved tb_per_sec 2.745 TB/sec +``` + +### Batch Testing + +Run multiple tests from a file and save results: +```bash +python3 flashinfer_benchmark.py --testlist samples/sample_testlist.txt --output_path samples/sample_testlist_output.csv +``` + +See `samples/sample_testlist.txt` for an example stdout output from the above command; `samples/sample_testlist_output.csv` for csv output from the same run. + +The output CSV will contain detailed metrics including: +- Median execution time +- Standard deviation +- TFLOPS/sec +- Memory throughput (TB/sec) +- Input flags +- Reproducer commands if `--generate_repro_command` is provided + +## Command Line Arguments +### General Flags +| Flag | Description | +|--------------------------|-------------------------------------------------------------------------------------------------------------| +| `--routine` | Test routine to run: `BatchDecodeWithPagedKVCacheWrapper`, `BatchPrefillWithPagedKVCacheWrapper`, `BatchPrefillWithRaggedKVCacheWrapper`, `BatchMLAPagedAttentionWrapper`, `gemm_fp8_nt_groupwise`, `group_gemm_fp8_nt_groupwise`, `bmm_fp8`, `mm_fp4`, `trtllm_fp4_block_scale_moe`, `trtllm_fp8_block_scale_moe`, `trtllm_fp8_per_tensor_scale_moe`, `cutlass_fused_moe` | +| `--num_iters` | Number of iterations for performance measurement | +| `--dry_run_iters` | Number of warmup iterations | +| `--no_cuda_graph` | Disable CUDA graph to execute kernels outside of the graph. | +| `--use_cupti` | Use CUPTI for timing GPU kernels when available. | +| `--refcheck` | Verify outputs match between different backends | +| `--allow_output_mismatch`| Continue testing even if outputs don't pass refcheck | +| `--random_seed` | Random seed for reproducibility | +| `--output_path` | Path to save CSV results | +| `--testlist` | Path to a file containing a list of test cases to run in batch mode | +| `--verbose`, `-v` | Print additional information (can be used multiple times for more verbosity, e.g. `-vv`) | +| `--case_tag` | Optional tag for the test case, useful for annotating or filtering results in the output CSV. | +| `--generate_repro_command`| If set, prints a reproducer command for the test case and stores it in the output CSV. | +| `--backends` | Space-separated list of backends to test, e.g. fa2, fa2_tc, fa3, cudnn, cutlass, trtllm, trtllm-gen, trtllm-native, cublas| + +### Attention Flags +| Flag | Description | +|--------------------------|-------------------------------------------------------------------------------------------------------------| +| `--page_size` | Page size for paged attention. Required for paged attention tests. | +| `--batch_size` | Number of sequences to process in parallel | +| `--s_qo` | Query/output sequence length. Should be 1 for decode tests. | +| `--s_kv` | Key/value sequence length (context length) | +| `--num_qo_heads` | Number of query/output attention heads | +| `--num_kv_heads` | Number of key/value attention heads | +| `--head_dim_qk` | Head dimension for Q/K. Must be 128 or 192. | +| `--head_dim_vo` | Head dimension for V/O. Usually equals head_dim_qk. | +| `--head_dim_ckv` | Head dimension for C/K/V (MLA attention). | +| `--head_dim_kpe` | Head dimension for KPE (MLA attention). | +| `--q_dtype` | Data type for the query tensor. Default: bfloat16. Currently only bfloat16 is supported. | +| `--kv_dtype` | Data type for the key and value tensors. Default: bfloat16. Currently only bfloat16 is supported. | +| `--causal` | Use causal attention masking (prefill only) | +| `--random_actual_seq_len`| Use random sequence lengths up to max length. If False, use max length. | + +### GEMM Flags +| Flag | Description | +|--------------------------|-------------------------------------------------------------------------------------------------------------| +| `--m` | Number of rows of matrix A and output matrix (GEMM M dimension) | +| `--n` | Number of columns of matrix B and output matrix (GEMM N dimension) | +| `--k` | Number of columns of matrix A / rows of matrix B (GEMM K dimension) | +| `--tile_size` | Tile size for the GEMM operation (affects performance and scaling) | +| `--group_size` | Number of groups for group GEMM (batching multiple GEMMs together) | +| `--scale_major_mode` | Layout for FP8 scaling: `MN` (per output tile) or `K` (per input tile) | +| `--out_dtype` | Output data type: `bfloat16` or `float16` | +| `--mma_sm` | Number of SMs to use for the MMA operation (1 or 2) | +| `--input_dtype` | Data type for input matrix (for FP8 GEMM, e.g. `fp8_e4m3`) | +| `--mat2_dtype` | Data type for second matrix (for FP8 GEMM, e.g. `fp8_e4m3`) | +| `--use_128x4_sf_layout` | Use 128x4 scale/format layout for FP4 GEMM (for `mm_fp4` routine) | +| `--use_nvfp4` | Whether to use nvfp4 quantization or mxfp4 quantization, defaults to False.(for `mm_fp4` routine) | +| `--autotune` | Enable autotune for supported operation (`trtllm` and `cutlass` backends for `mm_fp4` and `bmm_fp8` routines)| + +### MOE Flags +| Flag | Description | +|--------------------------|-------------------------------------------------------------------------------------------------------------| +| `--num_tokens` | Number of input tokens | +| `--hidden_size` | Hidden dimension size | +| `--intermediate_size` | Intermediate dimension size (FF layer dimension) | +| `--num_experts` | Total number of experts | +| `--top_k` | Number of experts to route to per token | +| `--n_group` | Number of expert groups (for DeepSeek routing). Default: 1 | +| `--topk_group` | Number of groups to consider for top-k routing. Default: 1 | +| `--routed_scaling_factor`| Scaling factor for routing. Default: 2.5 | +| `--local_expert_offset` | Offset of local experts in global expert space. Default: 0 | +| `--local_num_experts` | Number of experts handled by this device. Default: equals num_experts | +| `--tile_tokens_dim` | Tile dimension for tokens. Default: 8 | +| `--routing_method` | Routing method: `renormalize`, `deepseek_v3`, `llama4`, `renormalize_naive`. Default: `deepseek_v3`. | +| `--use_shuffled_weight` | Whether to use shuffled weight layout | +| `--weight_layout` | Weight layout: 0=MajorK, 1=MajorMn, 2=BlockMajorK. Default: 0 | +| `--use_routing_bias` | Whether to use routing bias | +| `--use_routing_scales_on_input` | Whether to use routing scales on input (for Llama4 routing) | +| `--input_dtype` | Data type of the input hidden states. Default: bfloat16 | +| `--weight_dtype` | Data type of the weights (before quantization). Default: bfloat16 | +| `--cutlass_variant` | CUTLASS MoE variant: `base` (no quant), `fp8` (per-tensor FP8), `nvfp4` (FP4 block-scale) | +| `--quantized_input` | For `nvfp4` only: quantize input activations to FP4 | +| `--tp_size` | Tensor-parallel world size | +| `--tp_rank` | Tensor-parallel rank | +| `--ep_size` | Expert-parallel world size | +| `--ep_rank` | Expert-parallel rank | +| `--gated_act` | Gated activation function: `swiglu` (default) or `geglu` | +| `--autotune` | Enable autotune for supported operation | + +### MOE Routing Method Compatibility + +| Routing Method | Requirements | Compatible MOE Types | +|------------------------|--------------|---------------------| +| **deepseek_v3** | `top_k <= 8`, `topk_group <= 4`, requires `--n_group`, `--topk_group`, `--routed_scaling_factor`, `--use_routing_bias` | FP4, FP8 Block Scale | +| **renormalize** | `top_k == 1` for FP8 Block Scale, `top_k <= 8` for FP4. Do NOT use `--n_group` or `--topk_group` | All MOE types | +| **llama4** | `top_k == 1`, requires `--routed_scaling_factor`, `--use_routing_bias`, `--use_routing_scales_on_input`. Do NOT use `--n_group` or `--topk_group` | FP8 Per-Tensor | +| **renormalize_naive** | `top_k == 1` for FP8 Block Scale, `top_k <= 8` for FP4. Do NOT use `--n_group` or `--topk_group` | FP4 primarily | + +Notes: +- Group parameters (`--n_group`, `--topk_group`) are ONLY used with DeepSeekV3 routing method. Using them with other routing methods will cause the error: "Routing kernel with groups implies DeepSeekV3 routing method." +- Different MOE kernel implementations have different `top_k` constraints. FP8 MOE kernels (both Block Scale and Per-Tensor) have stricter limits than FP4 for non-DeepSeekV3 routing methods. +- FP8 MOE kernels require integer values for group parameters, while FP4 MOE kernels accept optional values. +- CUTLASS fused MoE (`cutlass_fused_moe`) ignores `--routing_method`, `--n_group`, and `--topk_group`; it computes routing via softmax+top-k internally from the provided logits. + +## `flashinfer_benchmark.py` Routine & Backend Support Matrix +The following table summarizes the support surface of each routine & backend's on various [CUDA Compute Capabilities](https://developer.nvidia.com/cuda-gpus). + +Each column represents a compute capability. Backends inside cells represent supported backends. A blank cell means no backend is supported for that routine at that compute capability. + + +| Routine | 7.5 | 8.0 | 8.6 | 8.9 | 9.0 | 10.0 | 10.3 | 12.0 | +|---------|-----|-----|-----|-----|-----|-------|-------|-------| +| **BatchDecodeWithPagedKVCacheWrapper** | fa2 | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native | fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native | fa2, fa2_tc, cudnn | +| **BatchPrefillWithPagedKVCacheWrapper** | | fa2, cudnn | fa2, cudnn | fa2, cudnn | fa2, fa3, cudnn | fa2, cudnn, trtllm-gen, trtllm-native | fa2, cudnn, trtllm-gen, trtllm-native | fa2, cudnn | +| **BatchPrefillWithRaggedKVCacheWrapper** | | fa2, cudnn | fa2, cudnn | fa2, cudnn | fa2, fa3, cudnn | fa2, cudnn, cutlass, trtllm-native | fa2, cudnn, cutlass, trtllm-native | fa2, cudnn | +| **BatchMLAPagedAttentionWrapper** | | fa2 | fa2 | fa2 | fa2, fa3 | fa2, cutlass, trtllm-native | fa2, cutlass, trtllm-native | fa2 | +| **gemm_fp8_nt_groupwise** | | | | | | cutlass | cutlass | | +| **group_gemm_fp8_nt_groupwise** | | | | | | cutlass | cutlass | | +| **bmm_fp8** | | | | cudnn, cublas | cudnn, cublas | cudnn, cublas, cutlass | cudnn, cublas, cutlass | cudnn, cublas | +| **mm_fp4** | | | | | | cudnn, trtllm, cutlass | cudnn, trtllm, cutlass | cudnn | +| **trtllm_fp4_block_scale_moe** | | | | | | trtllm | trtllm | | +| **trtllm_fp8_block_scale_moe** | | | | | | trtllm | trtllm | | +| **trtllm_fp8_per_tensor_scale_moe** | | | | | | trtllm | trtllm | | +| **cutlass_fused_moe** | | | | | | cutlass | cutlass | | + +Backend Legend: +- fa2: FlashAttention2 +- fa2_tc: FlashAttention2 (with Tensor Cores for `BatchDecodeWithPagedKVCacheWrapper`) +- fa3: FlashAttention-3 +- cudnn: cuDNN +- cutlass: CUTLASS +- trtllm: TensorRT-LLM +- trtllm-gen: TensorRT-LLM +- trtllm-native: TensorRT-LLM (out-of-wrapper) diff --git a/benchmarks/bench_append_paged_kv_cache.py b/benchmarks/bench_append_paged_kv_cache.py index 46bbbbff62..a7ba8b6b6f 100644 --- a/benchmarks/bench_append_paged_kv_cache.py +++ b/benchmarks/bench_append_paged_kv_cache.py @@ -1,11 +1,12 @@ import argparse import dataclasses -from typing import Tuple, cast +from typing import Tuple +import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time @dataclasses.dataclass(kw_only=True) @@ -108,7 +109,8 @@ def fn_convert() -> Tuple[torch.Tensor, torch.Tensor]: ) batch_indices, positions = fn_convert() - convert_latency_ms = cast(float, do_bench(fn_convert)) + convert_latencies = bench_gpu_time(fn_convert) + convert_latency_ms = np.median(convert_latencies) @torch.cuda.nvtx.range(f"append model={model_name}, seqlens={seqlens}") def fn() -> None: @@ -124,7 +126,8 @@ def fn() -> None: "NHD", ) - latency_ms = cast(float, do_bench(fn)) + latencies = bench_gpu_time(fn) + latency_ms = np.median(latencies) all_layers_latency_ms = convert_latency_ms + latency_ms * model.num_layers throughput = ( k.numel() @@ -136,10 +139,10 @@ def fn() -> None: print( f"model: {model_name:8}", f"seqlens: {seqlens!r:{seqlen_strlen}}", - f"convert: {convert_latency_ms*1e3:2.0f}us", - f"1layer: {latency_ms*1e3:2.0f}us", - f"{model.num_layers}layers: {all_layers_latency_ms*1e3:3.0f}us", - f"throughput: {throughput*1e-9:8.3f}GB/s", + f"convert: {convert_latency_ms * 1e3:2.0f}us", + f"1layer: {latency_ms * 1e3:2.0f}us", + f"{model.num_layers}layers: {all_layers_latency_ms * 1e3:3.0f}us", + f"throughput: {throughput * 1e-9:8.3f}GB/s", ) print("---") diff --git a/benchmarks/bench_append_paged_mla_kv_cache.py b/benchmarks/bench_append_paged_mla_kv_cache.py index 3716c0bf4e..f1355213d7 100644 --- a/benchmarks/bench_append_paged_mla_kv_cache.py +++ b/benchmarks/bench_append_paged_mla_kv_cache.py @@ -1,11 +1,12 @@ import argparse import dataclasses -from typing import Tuple, cast +from typing import Tuple +import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time @dataclasses.dataclass(kw_only=True) @@ -92,7 +93,8 @@ def fn_convert() -> Tuple[torch.Tensor, torch.Tensor]: ) batch_indices, positions = fn_convert() - convert_latency_ms = cast(float, do_bench(fn_convert)) + convert_latencies = bench_gpu_time(fn_convert) + convert_latency_ms = np.median(convert_latencies) @torch.cuda.nvtx.range(f"append model={model_name}, seqlens={seqlens}") def fn() -> None: @@ -108,7 +110,8 @@ def fn() -> None: kv_last_page_len, ) - latency_ms = cast(float, do_bench(fn)) + latencies = bench_gpu_time(fn) + latency_ms = np.median(latencies) all_layers_latency_ms = convert_latency_ms + latency_ms * model.num_layers throughput = ( (ckv.numel() + kpe.numel()) @@ -119,10 +122,10 @@ def fn() -> None: print( f"model: {model_name:8}", f"seqlens: {seqlens!r:{seqlen_strlen}}", - f"convert: {convert_latency_ms*1e3:2.0f}us", - f"1layer: {latency_ms*1e3:2.0f}us", - f"{model.num_layers}layers: {all_layers_latency_ms*1e3:3.0f}us", - f"throughput: {throughput*1e-9:8.3f}GB/s", + f"convert: {convert_latency_ms * 1e3:2.0f}us", + f"1layer: {latency_ms * 1e3:2.0f}us", + f"{model.num_layers}layers: {all_layers_latency_ms * 1e3:3.0f}us", + f"throughput: {throughput * 1e-9:8.3f}GB/s", ) print("---") diff --git a/benchmarks/bench_attention_sink_triton_sgl_context.py b/benchmarks/bench_attention_sink_triton_sgl_context.py new file mode 100644 index 0000000000..e7d7457852 --- /dev/null +++ b/benchmarks/bench_attention_sink_triton_sgl_context.py @@ -0,0 +1,634 @@ +# bench: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/triton_ops/extend_attention.py + +""" +Memory-efficient attention for prefill. +It supports page size = 1 and prefill with KV cache (i.e. extend). +""" + +import torch +import triton +import triton.language as tl +import numpy as np + +from flashinfer.testing.utils import bench_gpu_time + +_is_cuda = True +if _is_cuda: + CUDA_CAPABILITY = torch.cuda.get_device_capability() + +_is_hip = False + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _fwd_kernel( + Q_Extend, + K_Extend, + V_Extend, + O_Extend, + K_Buffer, + V_Buffer, + qo_indptr, + kv_indptr, + kv_indices, + mask_ptr, + mask_indptr, + sink_ptr, + sm_scale, + kv_group_num, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + SLIDING_WINDOW_SIZE: tl.constexpr, + logit_cap: tl.constexpr, + Lq: tl.constexpr, + Lv: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + USE_CUSTOM_MASK: tl.constexpr, + IS_CAUSAL: tl.constexpr, + SKIP_PREFIX_CUSTOM_MASK: tl.constexpr, + STORE_TRANSPOSE: tl.constexpr, + HAS_SINK: tl.constexpr, +): + cur_seq = tl.program_id(0) + cur_head = tl.program_id(1) + cur_block_m = tl.program_id(2) + cur_kv_head = cur_head // kv_group_num + + cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq) + cur_seq_len_extend = tl.load(qo_indptr + cur_seq + 1) - cur_seq_extend_start_idx + cur_seq_kv_start_idx = tl.load(kv_indptr + cur_seq) + cur_seq_len_prefix = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx + cur_seq_len = cur_seq_len_prefix + cur_seq_len_extend + + if USE_CUSTOM_MASK: + cur_seq_mask_start_idx = tl.load(mask_indptr + cur_seq) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + offs_m = tl.arange(0, BLOCK_M) + mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_len_extend + + mask_d = offs_d < Lq + mask_dv = offs_dv < Lv + + offs_q = ( + (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] + ) + q = tl.load( + Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0 + ) + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + offs_qpe = ( + (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_qbs + + cur_head * stride_qh + + offs_dpe[None, :] + ) + qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0) + + # stage 1: compute scores with prefix + offs_n = tl.arange(0, BLOCK_N) + + acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32) + deno = tl.zeros([BLOCK_M], dtype=tl.float32) + e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + + for start_n in range(0, cur_seq_len_prefix, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + mask_n = (start_n + offs_n) < cur_seq_len_prefix + + final_mask = mask_m[:, None] & mask_n[None, :] + if USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK: + custom_mask = tl.load( + mask_ptr + + cur_seq_mask_start_idx + + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len + + start_n + + offs_n[None, :], + mask=(mask_m[:, None] & mask_n[None, :]), + other=0, + ) + final_mask &= custom_mask + if SLIDING_WINDOW_SIZE > 0: + # Add mask where q_id <= kv_id + sliding_window_size + # q_id = prefix_len + cur_m, kv_id = cur_n + window_mask = ( + cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m[:, None] + ) <= (start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE) + final_mask &= window_mask + + SKIP_TILE = False + if (USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK) or SLIDING_WINDOW_SIZE > 0: + SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 + + if not SKIP_TILE: + offs_kv_loc = tl.load( + kv_indices + cur_seq_kv_start_idx + start_n + offs_n, + mask=mask_n, + other=0, + ) + + # load k in transposed way + offs_buf_k = ( + offs_kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[:, None] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(mask_n[None, :]) & (mask_d[:, None]), + other=0.0, + ) + + qk = tl.dot(q.to(k.dtype), k) + if BLOCK_DPE > 0: + offs_kpe = ( + offs_kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Buffer + offs_kpe, + mask=mask_n[None, :], + other=0.0, + ) + qk += tl.dot(qpe.to(kpe.dtype), kpe) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(final_mask, qk, float("-inf")) + + row_max = tl.max(qk, 1) + row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max) + n_e_max = tl.maximum(row_max_fixed, e_max) + + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + deno = deno * re_scale + tl.sum(p, 1) + + offs_buf_v = ( + offs_kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=mask_n[:, None] & mask_dv[None, :], + other=0.0, + ) + p = p.to(v.dtype) + acc = acc * re_scale[:, None] + tl.dot(p, v) + + e_max = n_e_max + + # stage 2: compute the triangle part + + cur_block_m_end = ( + cur_seq_len_extend + if not IS_CAUSAL + else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M) + ) + for start_n in range(0, cur_block_m_end, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + mask_n = (start_n + offs_n) < cur_block_m_end + + final_mask = mask_m[:, None] & mask_n[None, :] + if USE_CUSTOM_MASK: + custom_mask = tl.load( + mask_ptr + + cur_seq_mask_start_idx + + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len + + cur_seq_len_prefix + + start_n + + offs_n[None, :], + mask=(mask_m[:, None] & mask_n[None, :]), + other=0, + ) + custom_mask &= mask_m[:, None] & mask_n[None, :] + final_mask &= custom_mask + elif IS_CAUSAL: + mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= ( + start_n + offs_n[None, :] + ) + mask_causual &= mask_m[:, None] & mask_n[None, :] + final_mask &= mask_causual + else: + mask_non_causal = mask_m[:, None] & mask_n[None, :] + final_mask &= mask_non_causal + + if SLIDING_WINDOW_SIZE > 0: + # Add mask where q_id <= kv_id + sliding_window_size + window_mask = (cur_block_m * BLOCK_M + offs_m[:, None]) <= ( + start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE + ) + final_mask &= window_mask + + SKIP_TILE = False + if USE_CUSTOM_MASK or SLIDING_WINDOW_SIZE > 0: + SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 + + if not SKIP_TILE: + # load k in transposed way + offs_k = ( + (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] + ) + k = tl.load( + K_Extend + offs_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0 + ) + + qk = tl.dot(q, k, out_dtype=tl.float32) + if BLOCK_DPE > 0: + offs_kpe = ( + (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs + + cur_kv_head * stride_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Extend + offs_kpe, + mask=mask_n[None, :], + other=0.0, + ) + qk += tl.dot(qpe, kpe) + + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(final_mask, qk, float("-inf")) + + row_max = tl.max(qk, 1) + row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max) + n_e_max = tl.maximum(row_max_fixed, e_max) + + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + deno = deno * re_scale + tl.sum(p, 1) + + offs_v = ( + (cur_seq_extend_start_idx + start_n + offs_n[:, None]) * stride_vbs + + cur_kv_head * stride_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0 + ) + p = p.to(v.dtype) + acc = acc * re_scale[:, None] + tl.dot(p, v) + + e_max = n_e_max + + if HAS_SINK: + cur_sink = tl.load(sink_ptr + cur_head) + deno += tl.exp(cur_sink - e_max) + + offs_o = ( + (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_obs + + cur_head * stride_oh + + offs_dv[None, :] + ) + if STORE_TRANSPOSE: + tl.store( + O_Extend + offs_o.T, + (acc / deno[:, None]).T, + mask=(mask_m[:, None] & mask_dv[None, :]).T, + ) + else: + tl.store( + O_Extend + offs_o, + acc / deno[:, None], + mask=mask_m[:, None] & mask_dv[None, :], + ) + + +def extend_attention_fwd( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + custom_mask, + is_causal, + mask_indptr, + max_len_extend, + sm_scale=None, + logit_cap=0.0, + skip_prefix_custom_mask=True, + sliding_window_size=-1, + sinks=None, +): + """ + q_extend, k_extend, v_extend, o_extend: contiguous tensors + + k_buffer, v_buffer: (prefix + extend) tensors in mem_manager + """ + Lq, Lk, Lv = ( + q_extend.shape[-1], + k_extend.shape[-1], + v_extend.shape[-1], + ) + + if Lq == 576: + BLOCK_DMODEL = 512 + BLOCK_DPE = 64 + elif Lq == 288: + BLOCK_DMODEL = 256 + BLOCK_DPE = 32 + elif Lq == 192: + BLOCK_DMODEL = 128 + BLOCK_DPE = 64 + else: + BLOCK_DMODEL = triton.next_power_of_2(Lq) + BLOCK_DPE = 0 + BLOCK_DV = triton.next_power_of_2(Lv) + + if _is_hip: + BLOCK_M, BLOCK_N = (64, 64) + num_warps = 4 + + else: + if _is_cuda and CUDA_CAPABILITY[0] >= 9: + if Lq <= 256: + BLOCK_M, BLOCK_N = (128, 64) + else: + BLOCK_M, BLOCK_N = (32, 64) + elif _is_cuda and CUDA_CAPABILITY[0] >= 8: + # sm86/sm89 has a much smaller shared memory size (100K) than sm80 (160K) + if CUDA_CAPABILITY[1] == 9 or CUDA_CAPABILITY[1] == 6: + if Lq <= 128: + BLOCK_M, BLOCK_N = (64, 128) + elif Lq <= 256: + BLOCK_M, BLOCK_N = (64, 64) + else: + BLOCK_M, BLOCK_N = (32, 32) + else: + if Lq <= 128: + BLOCK_M, BLOCK_N = (128, 128) + elif Lq <= 256: + BLOCK_M, BLOCK_N = (64, 64) + else: + BLOCK_M, BLOCK_N = (32, 64) + else: + BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32) + + num_warps = 4 if Lk <= 64 else 8 + + sm_scale = sm_scale or 1.0 / (Lq**0.5) + batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1] + kv_group_num = q_extend.shape[1] // k_extend.shape[1] + + USE_CUSTOM_MASK = custom_mask is not None + # Skip custom mask for prefix part + SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask + + HAS_SINK = sinks is not None + + grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M)) + num_stages = 1 + + extra_kargs = {} + if _is_hip: + extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2} + + _fwd_kernel[grid]( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + custom_mask, + mask_indptr, + sinks, + sm_scale, + kv_group_num, + q_extend.stride(0), + q_extend.stride(1), + k_extend.stride(0), + k_extend.stride(1), + v_extend.stride(0), + v_extend.stride(1), + o_extend.stride(0), + o_extend.stride(1), + k_buffer.stride(0), + k_buffer.stride(1), + v_buffer.stride(0), + v_buffer.stride(1), + SLIDING_WINDOW_SIZE=sliding_window_size, + logit_cap=logit_cap, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DPE=BLOCK_DPE, + BLOCK_DV=BLOCK_DV, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + Lq=Lq, + Lv=Lv, + USE_CUSTOM_MASK=USE_CUSTOM_MASK, + IS_CAUSAL=is_causal, + SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK, + HAS_SINK=HAS_SINK, + STORE_TRANSPOSE=_is_hip, + num_warps=num_warps, + num_stages=num_stages, + **extra_kargs, + ) + + +def bench_extend_attention_sink_triton_sgl( + batch_size, seq_len, head_qo_num, head_kv_num, head_dim, bench_with_sink +): + torch.manual_seed(42) + dtype = torch.bfloat16 + device = "cuda:0" + + # Split S into prefix and extend lengths + prefill_len = seq_len // 2 # Similar to test's N_CTX // 2 + extend_len = seq_len // 4 # Make extend length smaller than prefix + + # Calculate total tokens and extend tokens + total_extend_tokens = batch_size * extend_len + total_prefix_tokens = batch_size * prefill_len + + # Create query, key, value tensors for extension + q_extend = torch.randn( + total_extend_tokens, head_qo_num, head_dim, dtype=dtype, device=device + ) + k_extend = torch.randn( + total_extend_tokens, head_kv_num, head_dim, dtype=dtype, device=device + ) + v_extend = torch.randn( + total_extend_tokens, head_kv_num, head_dim, dtype=dtype, device=device + ) + o_extend = torch.empty_like(q_extend) + + # Create key-value buffers for prefix + k_buffer = torch.randn( + total_prefix_tokens, head_kv_num, head_dim, dtype=dtype, device=device + ) + v_buffer = torch.randn( + total_prefix_tokens, head_kv_num, head_dim, dtype=dtype, device=device + ) + + # Create index pointers + qo_indptr = torch.arange( + 0, (batch_size + 1) * extend_len, extend_len, device=device + ).to(torch.int32) + kv_indptr = torch.arange( + 0, (batch_size + 1) * prefill_len, prefill_len, device=device + ).to(torch.int32) + kv_indices = torch.arange(0, total_prefix_tokens, device=device).to(torch.int32) + + sm_scale = 1.0 / (head_dim**0.5) + # sliding_window = 128 # From GPT-OSS config, skip for now + sliding_window = -1 + + sink = ( + torch.randn(head_qo_num, device=device, dtype=torch.float32) + if bench_with_sink + else None + ) + + # warmup + for _ in range(5): + extend_attention_fwd( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + custom_mask=None, + is_causal=True, + mask_indptr=None, + max_len_extend=extend_len, + sm_scale=sm_scale, + sliding_window_size=sliding_window, + sinks=sink, + ) + + # benchmark + torch.cuda.synchronize() + measurements = bench_gpu_time( + lambda: extend_attention_fwd( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + custom_mask=None, + is_causal=True, + mask_indptr=None, + max_len_extend=extend_len, + sm_scale=sm_scale, + sliding_window_size=sliding_window, + sinks=sink, + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + kv_cache_numel = k_buffer.numel() + v_buffer.numel() + io = ( + q_extend.numel() * q_extend.element_size() + + kv_cache_numel * k_buffer.element_size() + ) + print( + f"batch_size={batch_size}, seq_len={seq_len}, num_qo_heads={head_qo_num}, num_kv_heads={head_kv_num}, head_dim={head_dim}" + ) + print(f"execution time: {ms}ms") + print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s") + + +# gpt oss +# head_num = 64 +# head_dim = 64 +# head_kv_num = 8 +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Benchmark SGLANG Extend Attention with Sink" + ) + parser.add_argument( + "--head_dim", type=int, default=64, help="Dimension of each head" + ) + parser.add_argument( + "--head_kv_num", type=int, default=8, help="Number of key/value heads" + ) + parser.add_argument( + "--head_qo_num", + type=int, + default=64, + help="Number of query heads", + ) + parser.add_argument("--sink", action="store_true", help="Whether to test with sink") + parser.add_argument( + "--batch_sizes", + type=int, + nargs="+", + default=[4, 128, 256], + help="List of batch sizes to test", + ) + parser.add_argument( + "--seq_lens", + type=int, + nargs="+", + default=[1024, 4096, 8192], + help="List of sequence lengths to test", + ) + + args = parser.parse_args() + + for batch_size in args.batch_sizes: + for seq_len in args.seq_lens: + bench_extend_attention_sink_triton_sgl( + batch_size=batch_size, + seq_len=seq_len, + head_qo_num=args.head_qo_num, + head_kv_num=args.head_kv_num, + head_dim=args.head_dim, + bench_with_sink=args.sink, + ) diff --git a/benchmarks/bench_attention_sink_triton_sgl_decode.py b/benchmarks/bench_attention_sink_triton_sgl_decode.py new file mode 100644 index 0000000000..08ada939b2 --- /dev/null +++ b/benchmarks/bench_attention_sink_triton_sgl_decode.py @@ -0,0 +1,877 @@ +# bench: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/triton_ops/decode_attention.py +# mypy: disable-error-code="no-redef" + +""" +Memory-efficient attention for decoding. +It supports page size = 1. +""" + +import torch +import triton +import triton.language as tl +import numpy as np + +from flashinfer.testing.utils import bench_gpu_time + +_is_hip = False + +_MIN_BLOCK_KV = 32 + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _fwd_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + kv_indptr, + kv_indices, + Att_Out, + Att_Lse, + num_kv_splits, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + kv_group_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + split_kv_id = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + + cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch) + cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx + kv_splits = tl.load(num_kv_splits + cur_batch) + + off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d + + kv_len_per_split = ( + tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + ) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + e_max = -float("inf") + e_sum = 0.0 + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + q = tl.load(Q + off_q, mask=mask_d, other=0.0) + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + kv_loc = tl.load( + kv_indices + cur_batch_kv_start_idx + offs_n, + mask=offs_n < split_kv_end, + other=0, + ) + offs_buf_k = ( + kv_loc[:, None] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[None, :] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]), + other=0.0, + ) + qk = tl.sum(q[None, :] * k, 1) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(offs_n < split_kv_end, qk, float("-inf")) + + offs_buf_v = ( + kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 0), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max) + acc *= re_scale + acc += tl.sum(p[:, None] * v, 0) + + e_sum = e_sum * re_scale + tl.sum(p, 0) + e_max = n_e_max + + offs_mid_o = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_dv + ) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum, + mask=(mask_dv), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + ) // Lv + + tl.store( + Att_Lse + offs_mid_o_1, + e_max + tl.log(e_sum), + ) + + +def _decode_att_m_fwd( + q, + k_buffer, + v_buffer, + att_out, + att_lse, + kv_indptr, + kv_indices, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, +): + BLOCK = 64 + # [TODO] work around SGPR limit on MI3xx + if _is_hip: + BLOCK = 8 + MAX_KV_SPLITS = max_kv_splits + Lk = k_buffer.shape[-1] + Lv = v_buffer.shape[-1] + + batch, head_num = kv_indptr.shape[0] - 1, q.shape[1] + + grid = (batch, head_num, MAX_KV_SPLITS) + kv_group_num = q.shape[1] // k_buffer.shape[1] + + if kv_group_num == 1: + num_warps = 4 + else: + num_warps = 2 + if _is_hip: + num_warps = 1 + + BLOCK_DMODEL = triton.next_power_of_2(Lk) + BLOCK_DV = triton.next_power_of_2(Lv) + + _fwd_kernel_stage1[grid]( + q, + k_buffer, + v_buffer, + sm_scale, + kv_indptr, + kv_indices, + att_out, + att_lse, + num_kv_splits, + q.stride(0), + q.stride(1), + k_buffer.stride(0), + k_buffer.stride(1), + v_buffer.stride(0), + v_buffer.stride(1), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + kv_group_num=kv_group_num, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DV=BLOCK_DV, + BLOCK_N=BLOCK, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + logit_cap=logit_cap, + num_warps=num_warps, + num_stages=2, + Lk=Lk, + Lv=Lv, + ) + + +@triton.jit +def _fwd_grouped_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + kv_indptr, + kv_indices, + Att_Out, + Att_Lse, + num_kv_splits, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + split_kv_id = tl.program_id(2) + + # ruff: noqa: SIM300 + if BLOCK_H < kv_group_num: + VALID_BLOCK_H: tl.constexpr = BLOCK_H + else: + VALID_BLOCK_H: tl.constexpr = kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H + mask_h = mask_h & (cur_head < q_head_num) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + + cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch) + cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx + kv_splits = tl.load(num_kv_splits + cur_batch) + + offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :] + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + mask_dpe = offs_dpe < Lk + off_qpe = ( + cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :] + ) + + kv_len_per_split = ( + tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + ) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0) + if BLOCK_DPE > 0: + qpe = tl.load( + Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0 + ) + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + kv_loc = tl.load( + kv_indices + cur_batch_kv_start_idx + offs_n, + mask=offs_n < split_kv_end, + other=0, + ) + offs_buf_k = ( + kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[:, None] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]), + other=0.0, + ) + qk = tl.dot(q, k.to(q.dtype)) + if BLOCK_DPE > 0: + offs_buf_kpe = ( + kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Buffer + offs_buf_kpe, + mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]), + other=0.0, + ) + qk += tl.dot(qpe, kpe.to(qpe.dtype)) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where( + mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf") + ) + + offs_buf_v = ( + kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 1), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + acc *= re_scale[:, None] + acc += tl.dot(p.to(v.dtype), v) + + e_sum = e_sum * re_scale + tl.sum(p, 1) + e_max = n_e_max + + offs_mid_o = ( + cur_batch * stride_mid_ob + + cur_head[:, None] * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_dv[None, :] + ) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_dv[None, :]), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + ) // Lv + + tl.store( + Att_Lse + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +def _decode_grouped_att_m_fwd( + q, + k_buffer, + v_buffer, + att_out, + att_lse, + kv_indptr, + kv_indices, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, +): + BLOCK = 32 + Lk = k_buffer.shape[-1] + Lv = v_buffer.shape[-1] + + # [TODO] work around shmem limit on MI3xx + if _is_hip and Lk >= 576: + BLOCK = 16 + + if Lk == 576: + BLOCK_DMODEL = 512 + BLOCK_DPE = 64 + elif Lk == 288: + BLOCK_DMODEL = 256 + BLOCK_DPE = 32 + else: + BLOCK_DMODEL = triton.next_power_of_2(Lk) + BLOCK_DPE = 0 + BLOCK_DV = triton.next_power_of_2(Lv) + + batch, head_num = kv_indptr.shape[0] - 1, q.shape[1] + kv_group_num = q.shape[1] // k_buffer.shape[1] + + BLOCK_H = 16 + MAX_KV_SPLITS = max_kv_splits + grid = ( + batch, + triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), + MAX_KV_SPLITS, + ) + + extra_kargs = {} + num_stages = 2 + if _is_hip: + # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html + # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py + extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2} + num_stages = 1 + + _fwd_grouped_kernel_stage1[grid]( + q, + k_buffer, + v_buffer, + sm_scale, + kv_indptr, + kv_indices, + att_out, + att_lse, + num_kv_splits, + q.stride(0), + q.stride(1), + k_buffer.stride(0), + k_buffer.stride(1), + v_buffer.stride(0), + v_buffer.stride(1), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + kv_group_num=kv_group_num, + q_head_num=head_num, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DPE=BLOCK_DPE, + BLOCK_DV=BLOCK_DV, + BLOCK_N=BLOCK, + BLOCK_H=BLOCK_H, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + logit_cap=logit_cap, + num_warps=4, + num_stages=num_stages, + Lk=Lk, + Lv=Lv, + **extra_kargs, + ) + + +@triton.jit +def _fwd_kernel_stage2( + Mid_O, + Mid_O_1, + O, + kv_indptr, + num_kv_splits, + sink_ptr, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_obs, + stride_oh, + MAX_KV_SPLITS: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + BLOCK_DV: tl.constexpr, + Lv: tl.constexpr, + HAS_SINK: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + + cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load( + kv_indptr + cur_batch + ) + kv_splits = tl.load(num_kv_splits + cur_batch) + + offs_d = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lv + + e_sum = 0.0 + e_max = -float("inf") + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh) // Lv + kv_len_per_split = ( + tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + ) + + for split_kv_id in range(0, MAX_KV_SPLITS): + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + if split_kv_end > split_kv_start: + tv = tl.load( + Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0 + ) + tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv) + n_e_max = tl.maximum(tlogic, e_max) + + old_scale = tl.exp(e_max - n_e_max) + acc *= old_scale + exp_logic = tl.exp(tlogic - n_e_max) + acc += exp_logic * tv + + e_sum = e_sum * old_scale + exp_logic + e_max = n_e_max + + if HAS_SINK: + cur_sink = tl.load(sink_ptr + cur_head) + e_sum += tl.exp(cur_sink - e_max) + + tl.store( + O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, + acc / e_sum, + mask=mask_d, + ) + + +def _decode_softmax_reducev_fwd( + logits, + lse, + q, + o, + v_buffer, + kv_indptr, + num_kv_splits, + max_kv_splits, + sinks=None, +): + batch, head_num = q.shape[0], q.shape[1] + Lv = v_buffer.shape[-1] + BLOCK_DV = triton.next_power_of_2(Lv) + + MAX_KV_SPLITS = max_kv_splits + HAS_SINK = sinks is not None + + extra_kargs = {} + if _is_hip: + # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html + # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py + extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2} + + grid = (batch, head_num) + _fwd_kernel_stage2[grid]( + logits, + lse, + o, + kv_indptr, + num_kv_splits, + sinks, + logits.stride(0), + logits.stride(1), + logits.stride(2), + o.stride(0), + o.stride(1), + MAX_KV_SPLITS=MAX_KV_SPLITS, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + BLOCK_DV=BLOCK_DV, + Lv=Lv, + HAS_SINK=HAS_SINK, + num_warps=4, + num_stages=2, + **extra_kargs, + ) + + +def decode_attention_fwd_normal( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=None, +): + _decode_att_m_fwd( + q, + k_buffer, + v_buffer, + attn_logits, + attn_lse, + kv_indptr, + kv_indices, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, + ) + _decode_softmax_reducev_fwd( + attn_logits, + attn_lse, + q, + o, + v_buffer, + kv_indptr, + num_kv_splits, + max_kv_splits, + sinks, + ) + + +def decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=None, +): + _decode_grouped_att_m_fwd( + q, + k_buffer, + v_buffer, + attn_logits, + attn_lse, + kv_indptr, + kv_indices, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap, + ) + _decode_softmax_reducev_fwd( + attn_logits, + attn_lse, + q, + o, + v_buffer, + kv_indptr, + num_kv_splits, + max_kv_splits, + sinks, + ) + + +def decode_attention_fwd( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=None, +): + assert max_kv_splits == attn_logits.shape[2] + assert q.shape[0] <= kv_indptr.shape[0] - 1 + assert q.shape[0] <= attn_logits.shape[0] + + kv_group_num = q.shape[1] // v_buffer.shape[1] + + if kv_group_num == 1: + # MHA + decode_attention_fwd_normal( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=logit_cap, + sinks=sinks, + ) + else: + # GQA/MQA/MLA + decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=logit_cap, + sinks=sinks, + ) + + +def bench_decode_attention_sink_triton_sgl( + batch_size, seq_len, head_qo_num, head_kv_num, head_dim, bench_with_sink +): + torch.manual_seed(42) + device = "cuda:0" + + dtype = torch.bfloat16 + total_tokens = batch_size * seq_len + device = torch.device("cuda") + sm_scale = 1.0 / (head_dim**0.5) + max_kv_splits = 8 + num_kv_splits = torch.full((batch_size,), 4, dtype=torch.int32, device="cuda") + + # q represents the new token being generated, one per batch + q = torch.randn(batch_size, head_qo_num, head_dim, dtype=dtype, device="cuda") + + # k_buffer and v_buffer represent all previous tokens + k_buffer = torch.randn( + total_tokens, head_kv_num, head_dim, dtype=dtype, device="cuda" + ) + v_buffer = torch.randn( + total_tokens, head_kv_num, head_dim, dtype=dtype, device="cuda" + ) + + o = torch.zeros(batch_size, head_qo_num, head_dim, dtype=dtype, device="cuda") + + b_seq_len = torch.full((batch_size,), seq_len, device="cuda") + + kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda") + kv_indptr[1 : batch_size + 1] = torch.cumsum(b_seq_len, dim=0) + kv_indices = torch.arange(total_tokens, device="cuda") + + attn_logits1 = torch.empty( + (batch_size, head_qo_num, max_kv_splits, head_dim), + dtype=torch.float32, + device="cuda", + ) + attn_lse1 = torch.empty( + (batch_size, head_qo_num, max_kv_splits, head_dim), + dtype=torch.float32, + device="cuda", + ) + sink = ( + torch.randn(head_qo_num, device=device, dtype=torch.float32) + if bench_with_sink + else None + ) + + # warmup + for _ in range(5): + decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits1, + attn_lse1, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=sink, + ) + + # benchmark + measurements = bench_gpu_time( + lambda: decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + attn_logits1, + attn_lse1, + num_kv_splits, + max_kv_splits, + sm_scale, + logit_cap=0.0, + sinks=sink, + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + kv_cache_numel = k_buffer.numel() + v_buffer.numel() + io = q.numel() * q.element_size() + kv_cache_numel * k_buffer.element_size() + print( + f"batch_size={batch_size}, seq_len={seq_len}, num_qo_heads={head_qo_num}, num_kv_heads={head_kv_num}, head_dim={head_dim}" + ) + print(f"execution time: {ms}ms") + print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s") + + +# gpt oss +# head_num = 64 +# head_dim = 64 +# head_kv_num = 8 +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Benchmark SGLANG Decode Attention with Sink" + ) + parser.add_argument( + "--head_dim", type=int, default=64, help="Dimension of each head" + ) + parser.add_argument( + "--head_kv_num", type=int, default=8, help="Number of key/value heads" + ) + parser.add_argument( + "--head_qo_num", + type=int, + default=64, + help="Number of query heads", + ) + parser.add_argument("--sink", action="store_true", help="Whether to test with sink") + parser.add_argument( + "--batch_sizes", + type=int, + nargs="+", + default=[4, 128, 256], + help="List of batch sizes to test", + ) + parser.add_argument( + "--seq_lens", + type=int, + nargs="+", + default=[1024, 4096, 8192, 16384], + help="List of sequence lengths to test", + ) + + args = parser.parse_args() + + for batch_size in args.batch_sizes: + for seq_len in args.seq_lens: + bench_decode_attention_sink_triton_sgl( + batch_size=batch_size, + seq_len=seq_len, + head_qo_num=args.head_qo_num, + head_kv_num=args.head_kv_num, + head_dim=args.head_dim, + bench_with_sink=args.sink, + ) diff --git a/benchmarks/bench_batch_attention.py b/benchmarks/bench_batch_attention.py new file mode 100644 index 0000000000..c94a86eacc --- /dev/null +++ b/benchmarks/bench_batch_attention.py @@ -0,0 +1,599 @@ +from __future__ import annotations + +import argparse +import os +from typing import List, Sequence, Tuple + +import numpy as np +import pandas as pd +import torch +import time +import flashinfer +from flashinfer.testing.utils import bench_gpu_time +import matplotlib.pyplot as plt + +NUM_LAYERS = 36 # QWen3 8b + +save_dir = "bench_plots/persistent" + + +def plot_original_comparison(df_all): + """Generate the original comparison plots.""" + cases = {"1": "chunked"} + + # Clean up potential duplicated header rows after appends and ensure numeric types + if "scheduler" in df_all.columns: + df_all = df_all[df_all["scheduler"] != "scheduler"] + df_all["seq_cfg_id"] = pd.to_numeric(df_all["seq_cfg_id"], errors="coerce") + df_all["bandwidth_GB_s"] = pd.to_numeric(df_all["bandwidth_GB_s"], errors="coerce") + df_all = df_all.dropna( + subset=["seq_cfg_id", "bandwidth_GB_s"] + ) # keep only valid rows + df_all["seq_cfg_id"] = df_all["seq_cfg_id"].astype(int) + repeats = df_all["num_repeats"].unique() + + # Pick an available case (prefer 4: hybrid-chunked). Avoid empty plots when requested case is missing + available_cases = sorted(df_all["seq_cfg_id"].dropna().unique().tolist()) + if not available_cases: + raise ValueError("No rows found in bench_batch_attention.csv") + preferred_case = 1 if 1 in available_cases else available_cases[0] + selected_case_label = cases.get(str(preferred_case), f"case {preferred_case}") + + for repeat in repeats: + # Collect data for all cases + case_names = [] + batch_prefill_values = [] + persistent_original_values = [] + decode_prefill_values = [] + decode_len = df_all["decode_len"].unique()[0] // 1024 + prefill_chunk_size = df_all["prefill_chunk_size"].unique()[0] // 1024 + for case in [str(preferred_case)]: + # Compute averages for this specific repeat count + batch_prefill_avg = df_all[ + (df_all["scheduler"] == "BatchPrefillWithPagedKVCacheWrapper") + & (df_all["seq_cfg_id"] == int(case)) + & (df_all["num_repeats"] == repeat) + ]["bandwidth_GB_s"].mean() + decode_prefill_avg = df_all[ + (df_all["scheduler"] == "Decode + Prefill") + & (df_all["seq_cfg_id"] == int(case)) + & (df_all["num_repeats"] == repeat) + ]["bandwidth_GB_s"].mean() + persistent_original_avg = df_all[ + (df_all["scheduler"] == "BatchAttentionWrapper") + & (df_all["seq_cfg_id"] == int(case)) + & (df_all["num_repeats"] == repeat) + ]["bandwidth_GB_s"].mean() + + case_names.append(selected_case_label) + # Replace NaNs with zeros to ensure bars render even if some schedulers are missing + batch_prefill_values.append(np.nan_to_num(batch_prefill_avg, nan=0.0)) + decode_prefill_values.append(np.nan_to_num(decode_prefill_avg, nan=0.0)) + persistent_original_values.append( + np.nan_to_num(persistent_original_avg, nan=0.0) + ) + + # Create grouped bar plot + x = np.arange(len(case_names)) + width = 0.2 + # Group persistent schedules together on the left, others on the right + offsets = np.array([0.5, 1.5, 2.5]) * width + + plt.figure(figsize=(12, 8)) # Increased height for better spacing + bars1 = plt.bar( + x + offsets[0], + persistent_original_values, + width, + label="Persistent", + color="#2ca02c", + ) + bars2 = plt.bar( + x + offsets[1], + batch_prefill_values, + width, + label="Batch Prefill", + color="#ff7f0e", + ) + bars3 = plt.bar( + x + offsets[2], + decode_prefill_values, + width, + label="Decode + Prefill", + color="#9467bd", + ) + + plt.ylabel("Average Bandwidth (GB/s)") + plt.title( + f"Average Bandwidth ({selected_case_label}, {repeat} repeats, {prefill_chunk_size}k prefill, {decode_len}k decode)" + ) + + # Add more space above the highest bar for legend and value labels + max_value = max( + max(batch_prefill_values), + max(decode_prefill_values), + max(persistent_original_values), + ) + plt.ylim(0, max_value * 1.08) # 8% more space above the highest bar + + plt.legend(fontsize=8, loc="upper right") + + # Add value labels on bars + def add_value_labels(bars): + for bar in bars: + height = bar.get_height() + if height == 0: + continue + # Use proportional offset based on data range + offset = max_value * 0.02 # 2% of max value as offset + plt.text( + bar.get_x() + bar.get_width() / 2.0, + height + offset, + f"{height:.1f}", + ha="center", + va="bottom", + fontsize=8, + ) + + add_value_labels(bars1) # Persistent (original) + add_value_labels(bars2) # Batch Prefill + add_value_labels(bars3) # Decode + Prefill + + plt.tight_layout() + plt.savefig( + f"{save_dir}/persistent_comparison_{repeat}_repeats_{prefill_chunk_size}k_prefill_{decode_len}k_decode.png", + dpi=300, + bbox_inches="tight", + ) + plt.show() + + +def run_bench( + decode_kv_lens: Sequence[int], + decode_qo_lens: Sequence[int], + prefill_kv_lens: Sequence[int], + prefill_qo_lens: Sequence[int], + *, + page_block_size: int, + num_kv_heads: int, + num_qo_heads: int, + head_dim: int, + device: int = 0, + causal: bool = True, + repeats: int = 50, +) -> Tuple[float, float, float, float, float, float, float]: + kv_lens = list(decode_kv_lens) + list(prefill_kv_lens) + seq_lens = torch.tensor(kv_lens, dtype=torch.int32) + q_lens = torch.tensor( + list(decode_qo_lens) + list(prefill_qo_lens), dtype=torch.int32 + ) + seq_lens_blocks = torch.ceil(seq_lens / page_block_size).int() + + q_indptr = torch.cat([torch.tensor([0]), torch.cumsum(q_lens, 0)], dim=0).int() + kv_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(seq_lens_blocks, 0)], dim=0 + ).int() + num_blocks = kv_indptr[-1].item() + + q = torch.rand( + q_indptr[-1].item(), num_qo_heads, head_dim, dtype=torch.bfloat16, device=device + ) + kv_data = torch.randn( + num_blocks, + 2, + page_block_size, + num_kv_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + + # old + wrapper_old = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device), + kv_layout="NHD", + backend="fa2", + ) + last_page_len = (seq_lens - 1) % page_block_size + 1 + + def old_plan(): + wrapper_old.plan( + q_indptr.to(device), + kv_indptr.to(device), + torch.arange(num_blocks, dtype=torch.int32, device=device), + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal=causal, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + + old_plan() # warmup module loading + start_time = time.perf_counter() + old_plan() + end_time = time.perf_counter() + measurements_old = bench_gpu_time( + lambda: wrapper_old.run(q, kv_data), repeat_iters=repeats + ) + ms_old = np.mean(measurements_old) + (end_time - start_time) * 1000 / NUM_LAYERS + + # Fused kernel + wrapper = flashinfer.BatchAttention(kv_layout="NHD") + + def persistent_plan(): + wrapper.plan( + q_indptr.to(device), + kv_indptr.to(device), + torch.arange(num_blocks, dtype=torch.int32, device=device), + seq_lens.to(device), + num_qo_heads, + num_kv_heads, + head_dim, + head_dim, + page_block_size, + causal=causal, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + + persistent_plan() # warmup module loading + start_time = time.perf_counter() + persistent_plan() + end_time = time.perf_counter() + measurements_new_normal = bench_gpu_time( + lambda: wrapper.run(q, kv_data), + repeat_iters=repeats, + ) + ms_new_normal = ( + np.mean(measurements_new_normal) + (end_time - start_time) * 1000 / NUM_LAYERS + ) + o, _ = wrapper.run(q, kv_data) + + # Separate prefill and decode wrappers + q_lens_d = torch.tensor(decode_qo_lens, dtype=torch.int32, device=device) + q_indptr_d = torch.cat( + [torch.tensor([0], device=device), torch.cumsum(q_lens_d, 0)], dim=0 + ).int() + seq_lens_d = torch.tensor(decode_kv_lens, dtype=torch.int32) + seq_lens_blocks_d = torch.ceil(seq_lens_d / page_block_size).int() + kv_indptr_d = torch.cat( + [torch.tensor([0]), torch.cumsum(seq_lens_blocks_d, 0)], dim=0 + ).int() + num_blocks_d = kv_indptr_d[-1].item() + if len(decode_qo_lens) > 0: + wrapper_decode = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device), + kv_layout="NHD", + backend="fa2", + use_tensor_cores=True, + ) + + last_page_len_d = (seq_lens_d - 1) % page_block_size + 1 + + q_d = q[: q_indptr_d[-1].item()] + kv_data_d = kv_data[:num_blocks_d] + start_time_d = time.perf_counter() + wrapper_decode.plan( + kv_indptr_d.to(device), + torch.arange(num_blocks_d, dtype=torch.int32, device=device), + last_page_len_d.to(device), + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + end_time_d = time.perf_counter() + measurements_decode = bench_gpu_time( + lambda: wrapper_decode.run(q_d, kv_data_d), repeat_iters=repeats + ) + ms_decode = ( + np.mean(measurements_decode) + + (end_time_d - start_time_d) * 1000 / NUM_LAYERS + ) + else: + ms_decode = 0 + + wrapper_prefill = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device), + kv_layout="NHD", + backend="fa2", + ) + q_lens_p = torch.tensor(prefill_qo_lens, dtype=torch.int32, device=device) + q_indptr_p = torch.cat( + [torch.tensor([0], device=device), torch.cumsum(q_lens_p, 0)], dim=0 + ).int() + seq_lens_p = torch.tensor(prefill_kv_lens, dtype=torch.int32, device=device) + seq_lens_blocks_p = torch.ceil(seq_lens_p / page_block_size).int() + kv_indptr_p = torch.cat( + [torch.tensor([0], device=device), torch.cumsum(seq_lens_blocks_p, 0)], + dim=0, + ).int() + num_blocks_p = kv_indptr_p[-1].item() + q_p = q[q_indptr_d[-1].item() :] + kv_data_p = kv_data[num_blocks_d:] + last_page_len_p = (seq_lens_p - 1) % page_block_size + 1 + + if len(prefill_qo_lens) > 0: + start_time_p = time.perf_counter() + wrapper_prefill.plan( + q_indptr_p.to(device), + kv_indptr_p.to(device), + torch.arange(num_blocks_p, dtype=torch.int32, device=device), + last_page_len_p.to(device), + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal=causal, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + end_time_p = time.perf_counter() + measurements_prefill = bench_gpu_time( + lambda: wrapper_prefill.run(q_p, kv_data_p), repeat_iters=repeats + ) + ms_prefill = ( + np.mean(measurements_prefill) + + (end_time_p - start_time_p) * 1000 / NUM_LAYERS + ) + else: + ms_prefill = 0 + + ms_separate = ms_prefill + ms_decode + + total_bytes = ( + q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + ) + mem_MB = total_bytes / 1024**2 + bw_old = total_bytes / (ms_old * 1e-3) / 1024**3 + bw_new_normal = total_bytes / (ms_new_normal * 1e-3) / 1024**3 + bw_separate = total_bytes / (ms_separate * 1e-3) / 1024**3 + + return ( + ms_old, + ms_new_normal, + ms_separate, + mem_MB, + bw_old, + bw_new_normal, + bw_separate, + ) # type: ignore + + +def synthesize_seq_len_configs( + decode_len, prefill_len, prefill_chunk_size, num_prefill_reqs, num_decode_reqs +) -> Tuple[List[List[Tuple[int, int]]], List[List[Tuple[int, int]]]]: + decode_lens: List[List[Tuple[int, int]]] = [ + [(decode_len, 1)] * num_decode_reqs, + ] + prefill_lens: List[List[Tuple[int, int]]] = [ + [(prefill_len, prefill_chunk_size)] * num_prefill_reqs, + ] + + return decode_lens, prefill_lens + + +def main(args: argparse.Namespace) -> None: + # If plotting mode, load existing data and plot + if args.plot: + if not os.path.exists("bench_batch_attention.csv"): + print("Error: bench_batch_attention.csv not found. Run benchmark first.") + return + + df_all = pd.read_csv("bench_batch_attention.csv") + os.makedirs(save_dir, exist_ok=True) + + plot_original_comparison(df_all) + + return + + # Benchmark mode + np.random.seed(42) + torch.random.manual_seed(42) + decode_len = args.decode_len + prefill_len = args.prefill_len + prefill_chunk_size = args.prefill_chunk_size + num_prefill_reqs = args.num_prefill_reqs + num_decode_reqs = args.num_decode_reqs + + decode_lens, prefill_lens = synthesize_seq_len_configs( + decode_len, prefill_len, prefill_chunk_size, num_prefill_reqs, num_decode_reqs + ) + if num_prefill_reqs == 0: + prefill_chunk_size = 0 + prefill_len = 0 + if num_decode_reqs == 0: + decode_len = 0 + + combinations = [ + { + "page_block_size": 1, + "head_dim": 128, + "num_kv_heads": 8, + "num_qo_heads": 32, + "model_name": "Qwen-8B", + }, + { + "page_block_size": 1, + "head_dim": 128, + "num_kv_heads": 8, + "num_qo_heads": 64, + "model_name": "Llama-3.1-70B", + }, + { + "page_block_size": 1, + "head_dim": 64, + "num_kv_heads": 4, + "num_qo_heads": 64, + "model_name": "Qwen-MoE-235B", + }, + ] + records_old = [] + records_new = [] + records_separate = [] + for cfg_id, (decode_case, prefill_case) in enumerate( + zip(decode_lens, prefill_lens, strict=True), start=1 + ): + prefill_kv_lens = [p[0] for p in prefill_case] + prefill_qo_lens = [p[1] for p in prefill_case] + decode_kv_lens = [p[0] for p in decode_case] + decode_qo_lens = [p[1] for p in decode_case] + for param in combinations: + pbs, hd, n_kv, n_qo, model_name = ( + param["page_block_size"], # type: ignore + param["head_dim"], # type: ignore + param["num_kv_heads"], # type: ignore + param["num_qo_heads"], # type: ignore + param["model_name"], + ) + ( + ms_old, + ms_new_normal, + ms_separate, + mem_MB, + bw_old, + bw_new_normal, + bw_separate, + ) = run_bench( + decode_kv_lens, + decode_qo_lens, + prefill_kv_lens, + prefill_qo_lens, + page_block_size=pbs, # type: ignore + num_kv_heads=n_kv, # type: ignore + num_qo_heads=n_qo, # type: ignore + head_dim=hd, # type: ignore + device=0, + causal=True, + repeats=args.repeats, + ) + records_old.extend( + [ + { + "scheduler": "BatchPrefillWithPagedKVCacheWrapper", + "seq_cfg_id": cfg_id, + "page_size": pbs, + "head_dim": hd, + "num_kv_heads": n_kv, + "num_qo_heads": n_qo, + "model_name": model_name, + "time_ms": ms_old, + "memory_MB": mem_MB, + "bandwidth_GB_s": bw_old, + "num_repeats": args.repeats, + "decode_len": decode_len, + "prefill_len": prefill_len, + "prefill_chunk_size": prefill_chunk_size, + "num_decode_reqs": num_decode_reqs, + }, + ] + ) + records_new.extend( + [ + { + "scheduler": "BatchAttentionWrapper", + "seq_cfg_id": cfg_id, + "page_size": pbs, + "head_dim": hd, + "num_kv_heads": n_kv, + "num_qo_heads": n_qo, + "model_name": model_name, + "time_ms": ms_new_normal, + "memory_MB": mem_MB, + "bandwidth_GB_s": bw_new_normal, + "num_repeats": args.repeats, + "decode_len": decode_len, + "prefill_len": prefill_len, + "prefill_chunk_size": prefill_chunk_size, + "num_decode_reqs": num_decode_reqs, + }, + ] + ) + records_separate.extend( + [ + { + "scheduler": "Decode + Prefill", + "seq_cfg_id": cfg_id, + "page_size": pbs, + "head_dim": hd, + "num_kv_heads": n_kv, + "num_qo_heads": n_qo, + "model_name": model_name, + "time_ms": ms_separate, + "memory_MB": mem_MB, + "bandwidth_GB_s": bw_separate, + "num_repeats": args.repeats, + "decode_len": decode_len, + "prefill_len": prefill_len, + "prefill_chunk_size": prefill_chunk_size, + "num_decode_reqs": num_decode_reqs, + }, + ] + ) + df = pd.DataFrame( + records_old + records_new + records_separate, + columns=[ + "scheduler", + "seq_cfg_id", + "page_size", + "head_dim", + "num_kv_heads", + "num_qo_heads", + "model_name", + "time_ms", + "memory_MB", + "bandwidth_GB_s", + "num_repeats", + "decode_len", + "prefill_len", + "prefill_chunk_size", + "num_decode_reqs", + ], + ) + file_name = "bench_batch_attention.csv" + if os.path.exists(file_name) and args.overwrite: + os.remove(file_name) + + # Append if file exists; write header only on first write + append_mode = "a" if os.path.exists(file_name) else "w" + write_header = append_mode == "w" + df.to_csv(file_name, index=False, mode=append_mode, header=write_header) + + # Only drop columns for printing, not for CSV + df_print = df.drop( + columns=[ + "page_size", + "num_repeats", + "decode_len", + "prefill_len", + "prefill_chunk_size", + "num_decode_reqs", + ] + ) + print(df_print.to_markdown(index=False, floatfmt=".2f")) + + +if __name__ == "__main__": + # Benchmark different attention schedulers and optionally plot results + + parser = argparse.ArgumentParser() + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--prefill_chunk_size", type=int, default=4096) + parser.add_argument("--num_prefill_reqs", type=int, default=1) + parser.add_argument("--num_decode_reqs", type=int, default=128) + parser.add_argument("--decode_len", type=int, default=8192) + parser.add_argument("--prefill_len", type=int, default=8192) + parser.add_argument( + "--plot", + action="store_true", + help="Plot existing benchmark results instead of running benchmark", + ) + parser.add_argument( + "--overwrite", action="store_true", help="Overwrite existing CSV file" + ) + args = parser.parse_args() + main(args) diff --git a/benchmarks/bench_batch_decode.py b/benchmarks/bench_batch_decode.py index 5a7bbdc23c..81413e1a26 100644 --- a/benchmarks/bench_batch_decode.py +++ b/benchmarks/bench_batch_decode.py @@ -16,9 +16,9 @@ import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time page_block_size = 16 num_kv_heads = 4 @@ -67,14 +67,15 @@ def bench_batch_decode( q_data_type=q_dtype, ) - ms = do_bench(lambda: wrapper.run(q, kv_data)) + measurements = bench_gpu_time(lambda: wrapper.run(q, kv_data)) + ms = np.median(measurements) io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() print( f"batch_size={batch_size}, seq_len={seq_len}, num_qo_heads={num_qo_heads}, num_kv_heads={num_kv_heads}, head_dim={head_dim}, page_block_size={page_block_size}, q_dtype={q_dtype}, kv_dtype={kv_dtype}" ) print(f"execution time: {ms}ms") - print(f"memory bandwidth: {io / ms / 1024 / 1024 :.2f} GB/s") + print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s") if __name__ == "__main__": diff --git a/benchmarks/bench_blackwell_attention.py b/benchmarks/bench_blackwell_attention.py new file mode 100644 index 0000000000..73b0cd0b3c --- /dev/null +++ b/benchmarks/bench_blackwell_attention.py @@ -0,0 +1,126 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import ( + bench_gpu_time, + attention_tflops_per_sec_with_actual_seq_lens, +) + + +def bench_fmha_blackwell( + batch_size, + qkv_len, + num_heads, + head_dim, + causal, + dtype, +): + q = torch.randn( + batch_size * qkv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + k = torch.randn( + batch_size * qkv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + v = torch.randn( + batch_size * qkv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + + qo_segment_offsets = ( + torch.arange(0, batch_size + 1, device="cuda", dtype=torch.int32) * qkv_len + ) + kv_segment_offsets = ( + torch.arange(0, batch_size + 1, device="cuda", dtype=torch.int32) * qkv_len + ) + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + torch.empty(128 * 1024 * 1024, dtype=dtype, device="cuda"), + kv_layout="NHD", + backend="cutlass", + ) + wrapper.plan( + qo_segment_offsets, + kv_segment_offsets, + num_heads, + num_heads, + head_dim, + head_dim_vo=head_dim, + causal=causal, + q_data_type=dtype, + kv_data_type=dtype, + ) + _o = wrapper.run(q, k, v) + measurements = bench_gpu_time( + lambda: wrapper.run(q, k, v), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + + TFLOPS = attention_tflops_per_sec_with_actual_seq_lens( + torch.full((batch_size,), qkv_len), + torch.full((batch_size,), qkv_len), + head_dim, + head_dim, + num_heads, + causal, + ms, + ) + print( + f"bench_fmha_blackwell (batch_size={batch_size}, qkv_len={qkv_len}, num_heads={num_heads}, head_dim={head_dim}, causal={causal}), flops: {TFLOPS:.3f} TFLOPs/s" + ) + + +if __name__ == "__main__": + print("\n === head_dim=128 ===") + bench_fmha_blackwell(128, 512, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(64, 1024, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(32, 2048, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(16, 4096, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(8, 8192, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(4, 16384, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(2, 32768, 32, 128, False, torch.bfloat16) + bench_fmha_blackwell(1, 65536, 32, 128, False, torch.bfloat16) + + bench_fmha_blackwell(128, 512, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(64, 1024, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(32, 2048, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(16, 4096, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(8, 8192, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(4, 16384, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(2, 32768, 32, 128, True, torch.bfloat16) + bench_fmha_blackwell(1, 65536, 32, 128, True, torch.bfloat16) + + print("\n === head_dim=64 ===") + bench_fmha_blackwell(128, 512, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(64, 1024, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(32, 2048, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(16, 4096, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(8, 8192, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(4, 16384, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(2, 32768, 32, 64, False, torch.bfloat16) + bench_fmha_blackwell(1, 65536, 32, 64, False, torch.bfloat16) + + bench_fmha_blackwell(128, 512, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(64, 1024, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(32, 2048, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(16, 4096, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(8, 8192, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(4, 16384, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(2, 32768, 32, 64, True, torch.bfloat16) + bench_fmha_blackwell(1, 65536, 32, 64, True, torch.bfloat16) diff --git a/benchmarks/bench_block_sparse_attention.py b/benchmarks/bench_block_sparse_attention.py new file mode 100644 index 0000000000..2da2478a6f --- /dev/null +++ b/benchmarks/bench_block_sparse_attention.py @@ -0,0 +1,157 @@ +""" +Copyright (c) 2024 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import ( + bench_gpu_time, + attention_tflops_per_sec_with_actual_seq_lens, +) + + +def bench_variable_block_sparse_attention( + num_qo_heads, + num_kv_heads, + head_dim, + seq_len, + num_blocks_row, + num_blocks_col, + block_density, +): + if num_qo_heads % num_kv_heads != 0: + return + if seq_len // num_blocks_row < 1: + return + if seq_len // num_blocks_col < 1: + return + + # synthesize uniform block sz + block_row_sz = torch.ones(num_blocks_row, dtype=torch.int32) * ( + seq_len // num_blocks_row + ) + block_row_sz[-1] = seq_len - (seq_len // num_blocks_row) * (num_blocks_row - 1) + block_row_sz = block_row_sz.unsqueeze(0).repeat(num_kv_heads, 1) + + block_col_sz = torch.ones(num_blocks_col, dtype=torch.int32) * ( + seq_len // num_blocks_col + ) + block_col_sz[-1] = seq_len - (seq_len // num_blocks_col) * (num_blocks_col - 1) + block_col_sz = block_col_sz.unsqueeze(0).repeat(num_kv_heads, 1) + + block_mask_map = ( + torch.rand(num_kv_heads, num_blocks_row, num_blocks_col) < block_density + ) + + q = torch.randn(num_qo_heads, seq_len, head_dim, dtype=torch.half, device="cuda") + k = torch.randn(num_kv_heads, seq_len, head_dim, dtype=torch.half, device="cuda") + v = torch.randn(num_kv_heads, seq_len, head_dim, dtype=torch.half, device="cuda") + + float_workspace_buffer = torch.empty( + 128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0" + ) + sparse_wrapper_fa2 = flashinfer.sparse.VariableBlockSparseAttentionWrapper( + float_workspace_buffer, backend="fa2" + ) + sparse_wrapper_fa3 = flashinfer.sparse.VariableBlockSparseAttentionWrapper( + float_workspace_buffer, backend="fa3" + ) + + sparse_wrapper_fa2.plan( + block_mask_map=block_mask_map, + block_row_sz=block_row_sz, + block_col_sz=block_col_sz, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + q_data_type=torch.half, + ) + sparse_wrapper_fa3.plan( + block_mask_map=block_mask_map, + block_row_sz=block_row_sz, + block_col_sz=block_col_sz, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + q_data_type=torch.half, + ) + + # Benchmark sparse attention with FA2 + measurements_fa2 = bench_gpu_time( + lambda: sparse_wrapper_fa2.run(q, k, v), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + sparse_ms_fa2 = np.median(measurements_fa2) + + # Benchmark sparse attention with FA3 + measurements_fa3 = bench_gpu_time( + lambda: sparse_wrapper_fa3.run(q, k, v), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + sparse_ms_fa3 = np.median(measurements_fa3) + + q = torch.randn(seq_len, num_qo_heads, head_dim, dtype=torch.half, device="cuda") + k = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda") + v = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda") + dense_sm80_ms, dense_sm90_ms = ( + np.median( + bench_gpu_time( + lambda: flashinfer.single_prefill_with_kv_cache_return_lse( + q, k, v, causal=False, backend=backend + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ) + for backend in ["fa2", "fa3"] + ) + + def flops(ms): + return attention_tflops_per_sec_with_actual_seq_lens( + torch.tensor([seq_len]), + torch.tensor([seq_len]), + head_dim, + head_dim, + num_qo_heads, + False, + ms, + ) + + print( + f"bench_variable_block_sparse_attention (num_qo_heads={num_qo_heads}, num_kv_heads={num_kv_heads}, head_dim={head_dim}, seq_len={seq_len}, num_blocks_row={num_blocks_row}, num_blocks_col={num_blocks_col}, block_density={block_density}), sparse fa2-template: {flops(sparse_ms_fa2):.3f} TFLOPs/s, sparse fa3-template: {flops(sparse_ms_fa3):.3f} TFLOPs/s, dense fa2-template: {flops(dense_sm80_ms):.3f} TFLOPs/s, dense fa3-template: {flops(dense_sm90_ms):.3f} TFLOPs/s" + ) + + +if __name__ == "__main__": + for num_qo_heads in [32]: + for num_kv_heads in [32]: + for head_dim in [128]: + for seq_len in [8192, 16384, 32768]: + for num_blocks_row in [20]: + for num_blocks_col in [50]: + for block_density in [0.1, 0.3, 0.5, 0.7, 0.9]: + bench_variable_block_sparse_attention( + num_qo_heads, + num_kv_heads, + head_dim, + seq_len, + num_blocks_row, + num_blocks_col, + block_density, + ) diff --git a/benchmarks/bench_cute_dsl_blockscaled_gemm.py b/benchmarks/bench_cute_dsl_blockscaled_gemm.py new file mode 100644 index 0000000000..fb444b019d --- /dev/null +++ b/benchmarks/bench_cute_dsl_blockscaled_gemm.py @@ -0,0 +1,202 @@ +import json +import random +import cutlass +from flashinfer.cute_dsl.blockscaled_gemm import ( + create_scale_factor_tensor, + grouped_gemm_nt_masked, # deepgemm-like python interface for DLFW integration +) +import torch +import cutlass.torch as cutlass_torch +from flashinfer.cute_dsl.utils import get_cutlass_dtype +from flashinfer.testing.utils import bench_kineto, count_bytes + + +ab_dtype = "float4_e2m1fn" +sf_dtype = "float8_e4m3fn" +c_dtype = "bfloat16" +sf_vec_size = 16 + +# DeepGEMM case +a_major = "k" +b_major = "k" +c_major = "n" + + +def bench_one(num_groups, max_m, expected_m_per_group, n, k): + data = create_data( + num_groups=num_groups, + max_m=max_m, + expected_m_per_group=expected_m_per_group, + n=n, + k=k, + ) + + def test_func(): + grouped_gemm_nt_masked( + lhs=data["a"], + rhs=data["b"], + out=data["c"], + masked_m=data["masked_m"], + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + alpha_dtype="float32", + ) + + t = bench_kineto( + test_func, + "Sm100BlockScaledPersistentDenseGemmKernel", + suppress_kineto_output=True, + ) + + valid_m = data["masked_m"].sum().item() + t_calibrated = t / valid_m * (expected_m_per_group * num_groups) + + tflops = 2 * valid_m * n * k / t / 1e12 + gb_per_s = ( + ( + count_bytes(data["a"], data["c"]) * valid_m / (max_m * num_groups) + + count_bytes(data["b"]) + ) + / 1e9 + / t + ) + + print( + f" > Perf ({num_groups=}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}): " + f"{t * 1e6:4.0f} us | {tflops:4.0f} TFLOPS | {gb_per_s:4.0f} GB/s" + ) + + metrics = dict( + num_groups=num_groups, + m_per_group=expected_m_per_group, + valid_m=valid_m, + n=n, + k=k, + t_us_raw=t * 1e6, + t_us_calibrated=t_calibrated * 1e6, + tflops=tflops, + gb_per_s=gb_per_s, + ) + print(f"MAIN_OUTPUT={json.dumps(metrics)}") + + +# ref: DeepGEMM +def enumerate_m_grouped_masked(): + max_m = 4096 + + cases = [ + # GB200 cases + (6, 1024), + (6, 512), + # DeepGEMM default cases + (1, 1024), + (2, 512), + (4, 256), + ] + # more GB200 cases + num_experts = 288 + num_experts_per_token = 8 + for num_ranks in [4, 8, 16, 32, 36, 48, 72]: + for num_tokens in [64, 128, 256, 384, 512, 768, 1024]: + num_groups = num_experts // num_ranks + expected_m_per_group = num_tokens * num_experts_per_token // num_groups + cases.append((num_groups, expected_m_per_group)) + + for num_groups, expected_m_per_group in cases: + for n, k in ( + (4096, 7168), + (7168, 2048), + ): + yield dict( + num_groups=num_groups, + max_m=max_m, + expected_m_per_group=expected_m_per_group, + n=n, + k=k, + ) + + +# Copy and modified from test_cute_dsl_blockscaled_gemm.py, may extract common logic later if needed +def create_data(num_groups, max_m, expected_m_per_group, n, k, device="cuda:0"): + device = torch.device(device) + l = num_groups + m = max_m + + a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like( + a_ref, + get_cutlass_dtype(ab_dtype), + is_dynamic_layout=True, + assumed_align=16, + ) + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_ref, + get_cutlass_dtype(ab_dtype), + is_dynamic_layout=True, + assumed_align=16, + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, + get_cutlass_dtype(c_dtype), + is_dynamic_layout=True, + assumed_align=16, + ) + + # for deepgemm-like python interface + if ab_dtype == "float4_e2m1fn": + m, k, l = a_torch.shape + n, k, l = b_torch.shape + # slice into half after flatten + half_len_a = a_torch.numel() // 2 + half_len_b = b_torch.numel() // 2 + a_torch = ( + a_torch.permute(2, 0, 1) + .flatten()[:half_len_a] + .reshape(l, m, k // 2) + .permute(1, 2, 0) + ) + b_torch = ( + b_torch.permute(2, 0, 1) + .flatten()[:half_len_b] + .reshape(l, n, k // 2) + .permute(1, 2, 0) + ) + + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( + l, m, k, sf_vec_size, get_cutlass_dtype(sf_dtype), device + ) + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( + l, n, k, sf_vec_size, get_cutlass_dtype(sf_dtype), device + ) + + masked_m_tensor = create_masked_m( + num_groups=num_groups, expected_m_per_group=expected_m_per_group, max_m=max_m + ) + + return dict( + a=(a_torch, sfa_torch), + b=(b_torch, sfb_torch), + c=c_torch, + masked_m=masked_m_tensor, + ) + + +def create_masked_m(num_groups, expected_m_per_group, max_m): + """Align with DeepGEMM :: generate_m_grouped_masked""" + masked_m = torch.empty((num_groups,), device="cuda", dtype=torch.int) + for j in range(num_groups): + masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3)) + assert masked_m.amax().item() <= max_m + return masked_m + + +if __name__ == "__main__": + torch.manual_seed(42) + random.seed(42) + for config in enumerate_m_grouped_masked(): + bench_one(**config) diff --git a/benchmarks/bench_cutlass_fused_moe.py b/benchmarks/bench_cutlass_fused_moe.py new file mode 100644 index 0000000000..e0dff8e215 --- /dev/null +++ b/benchmarks/bench_cutlass_fused_moe.py @@ -0,0 +1,229 @@ +""" +Copyright (c) 2024 by FlashInfer team. + +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 pprint + +import numpy as np +import torch +from torch.nn import functional as F + +import flashinfer.fused_moe as fused_moe +from flashinfer import fp4_quantize +from flashinfer.autotuner import AutoTuner, autotune, get_config_path +from flashinfer.testing.utils import bench_gpu_time + +FLOAT4_E2M1_MAX = 6.0 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + + +test_configs = [ + { + "hidden_size": 7168, + "num_experts": 256, + "top_k": 8, + "intermediate_size": 256, + }, + { + "hidden_size": 7168, + "num_experts": 32, + "top_k": 8, + "intermediate_size": 2048, + }, +] + + +def compute_routing( + router_logits: torch.Tensor, top_k: int +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute routing weights and selected experts from router logits. + + Args: + router_logits (torch.Tensor): Router logits of shape [batch_size, num_experts] + top_k (int): Number of experts to route to per token + + Returns: + tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - routing_weights: Expert weights of shape [batch_size, top_k] + - selected_experts: Expert indices of shape [batch_size, top_k] + """ + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.float() + return routing_weights, selected_experts + + +def bench_cutlass_fused_moe( + batch_size, + hidden_size, + num_experts, + top_k, + intermediate_size, + skip_autotune, +): + torch.manual_seed(42) + quant_blocksize = 16 + round_up = lambda x, y: (x + y - 1) // y * y + e = num_experts + m = batch_size + n = intermediate_size + k = hidden_size + otype = torch.bfloat16 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=otype) / 10 + w1_cutlass = torch.cat((w1[:, n:, :], w1[:, :n, :]), dim=1).contiguous() + + sf_w1_2n = round_up(2 * n, 128) + sf_w1_k = round_up(k // quant_blocksize, 4) + w1_blockscale = torch.empty( + (e, sf_w1_2n, sf_w1_k), device="cuda", dtype=torch.float8_e4m3fn + ) + w1_blockscale_cutlass = torch.empty( + (e, sf_w1_2n, sf_w1_k), device="cuda", dtype=torch.float8_e4m3fn + ) + + w2 = torch.randn((e, k, n), device="cuda", dtype=otype) / 10 + sf_w2_k = round_up(k, 128) + sf_w2_n = round_up(n // quant_blocksize, 4) + w2_blockscale = torch.empty( + (e, sf_w2_k, sf_w2_n), device="cuda", dtype=torch.float8_e4m3fn + ) + w1_q = torch.empty((e, 2 * n, k // 2), device="cuda", dtype=torch.uint8) + w1_q_cutlass = torch.empty((e, 2 * n, k // 2), device="cuda", dtype=torch.uint8) + w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8) + w1_gs = torch.empty((e,), device="cuda", dtype=torch.float32) + w2_gs = torch.empty((e,), device="cuda", dtype=torch.float32) + + for expert in range(e): + w1_amax = torch.abs(w1).max().to(torch.float32) + w2_amax = torch.abs(w2).max().to(torch.float32) + w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax + w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax + + w1_q[expert], w1_blockscale[expert] = fp4_quantize(w1[expert], w1_gs[expert]) + + w1_q_cutlass[expert], w1_blockscale_cutlass[expert] = fp4_quantize( + w1_cutlass[expert], w1_gs[expert] + ) + + w2_q[expert], w2_blockscale[expert] = fp4_quantize(w2[expert], w2_gs[expert]) + + x = torch.randn(m, k, dtype=otype).cuda() + a1_gs = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(x).max().to( + torch.float32 + ).cuda() + a1_gs = torch.tensor(1.0, device="cuda", dtype=torch.float32) + a2_gs = torch.tensor(1.0, device="cuda", dtype=torch.float32) + router_logits = torch.randn(m, e, dtype=otype).cuda() + routing_weights, selected_experts = compute_routing(router_logits, top_k) + + flash_output = torch.zeros_like(x) + + quant_scales = [ + a1_gs, + w1_blockscale.view(torch.int32), + 1.0 / (a1_gs * w1_gs), + a2_gs, + w2_blockscale.view(torch.int32), + 1.0 / (a2_gs * w2_gs), + ] + hidden_states = x + hidden_states, input_sf = fp4_quantize(x, a1_gs) + + # Warmup + for _ in range(3): + _ = fused_moe.cutlass_fused_moe( + hidden_states, + selected_experts.to(torch.int), + routing_weights, + w1_q.contiguous().view(torch.long), + w2_q.contiguous().view(torch.long), + otype, + quant_scales=quant_scales, + input_sf=input_sf, + output=flash_output, + tune_max_num_tokens=16384, + ) + + if not skip_autotune: + with torch.inference_mode(), autotune(True): + _ = fused_moe.cutlass_fused_moe( + hidden_states, + selected_experts.to(torch.int), + routing_weights, + w1_q.contiguous().view(torch.long), + w2_q.contiguous().view(torch.long), + otype, + quant_scales=quant_scales, + input_sf=input_sf, + output=flash_output, + tune_max_num_tokens=16384, + ) + ms_list = bench_gpu_time( + lambda: fused_moe.cutlass_fused_moe( + hidden_states, + selected_experts.to(torch.int), + routing_weights, + w1_q.contiguous().view(torch.long), + w2_q.contiguous().view(torch.long), + otype, + quant_scales=quant_scales, + input_sf=input_sf, + output=flash_output, + ), + ) + median_ms = np.median(ms_list) + print(f"{'input':<15} {'weight1':<20} {'weight2':<20} {'time(ms)'}") + print( + f"{str(tuple(hidden_states.shape)):<15} {str(tuple(w1.shape)):<20} {str(tuple(w2.shape)):<20} {median_ms:.3f}" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--update-config", + action="store_true", + help="Update the config file with the new profiling results", + ) + parser.add_argument( + "--num-tokens", type=int, default=32, help="Number of tokens to profile" + ) + parser.add_argument("--skip-autotune", action="store_true", help="Skip autotuning") + args = parser.parse_args() + AutoTuner.get().clear_cache() + + for config in test_configs: + bench_cutlass_fused_moe( + args.num_tokens, + config["hidden_size"], + config["num_experts"], + config["top_k"], + config["intermediate_size"], + args.skip_autotune, + ) + + configs = AutoTuner.get().profiling_cache + if args.update_config and configs: + # The original key contains a runner's hash in k[2] which might be different across machines. + # So, we remove it for now. v[0] and v[1] are the runner id and the tactic. + converted = {str((k[0], k[1], k[3])): (v[0], v[1]) for k, v in configs.items()} + config_path = get_config_path(is_module=False) + with open(config_path, "w") as f: + f.write("best_configs = ") + pprint.pprint(converted, stream=f) + print(f"Saved the cache to {config_path}") diff --git a/benchmarks/bench_deepgemm_blackwell.py b/benchmarks/bench_deepgemm_blackwell.py new file mode 100644 index 0000000000..ec66f22341 --- /dev/null +++ b/benchmarks/bench_deepgemm_blackwell.py @@ -0,0 +1,149 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import numpy as np +import torch + +from flashinfer.gemm import ( + batch_deepgemm_fp8_nt_groupwise, + group_deepgemm_fp8_nt_groupwise, +) +from flashinfer.testing.utils import bench_gpu_time, quantize_fp8 + + +def bench_deepgemm_grouped_fp8_blackwell(batch_size, m, n, k, in_dtype, out_dtype): + """Benchmark DeepGEMM-based grouped GEMM with FP8 quantization.""" + + # Create float32 input tensors + a_f32 = torch.randn(batch_size * m, k, device="cuda", dtype=torch.float32) + b_f32 = torch.randn(batch_size, n, k, device="cuda", dtype=torch.float32) + + # Quantize tensor A using per-token quantization + a_fp8, a_scale = quantize_fp8(a_f32, (batch_size * m, k // 128), (1, 128), "K") + + # Quantize tensor B using per-block quantization + b_fp8, b_scale = quantize_fp8( + b_f32, (batch_size, n // 128, k // 128), (1, 128, 128), "K" + ) + + # Create group assignment indices + m_indices = torch.arange( + batch_size, device="cuda", dtype=torch.int32 + ).repeat_interleave(m) + + # Pre-allocate output tensor + out = torch.empty(batch_size * m, n, device="cuda", dtype=out_dtype) + + # Benchmark the DeepGEMM function + measurements = bench_gpu_time( + lambda: group_deepgemm_fp8_nt_groupwise( + a_fp8, b_fp8, a_scale, b_scale, m_indices, out=out, out_dtype=out_dtype + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + tflops_per_second = 2 * batch_size * m * n * k * 1e-9 / ms + memory_bandwidth_per_second = ( + sum( + [ + _.numel() * _.element_size() + for _ in [a_fp8, b_fp8, a_scale, b_scale, m_indices, out] + ] + ) + * 1e-9 + / ms + ) + print( + f"group_deepgemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} " + f"in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s" + f"memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s" + ) + + return tflops_per_second + + +def bench_deepgemm_batch_fp8_blackwell(batch_size, m, n, k, in_dtype, out_dtype): + """Benchmark DeepGEMM-based batch GEMM with FP8 quantization.""" + + a = torch.randn((batch_size, m, k), device="cuda", dtype=torch.float32) + b = torch.randn((batch_size, n, k), device="cuda", dtype=torch.float32) + masked_m = torch.randint(0, m, (batch_size,), device="cuda", dtype=torch.int32) + a_fp8, a_scale = quantize_fp8(a, (batch_size, m, k // 128), (1, 1, 128), "K") + b_fp8, b_scale = quantize_fp8( + b, (batch_size, n // 128, k // 128), (1, 128, 128), "K" + ) + expected_m = min(int(masked_m.float().mean()) + 1, m) + + out = torch.empty((batch_size, m, n), device="cuda", dtype=out_dtype) + + # Benchmark the DeepGEMM function + measurements = bench_gpu_time( + lambda: batch_deepgemm_fp8_nt_groupwise( + a_fp8, + b_fp8, + a_scale, + b_scale, + masked_m, + expected_m, + out=out, + out_dtype=out_dtype, + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + + tflops_per_second = 2 * batch_size * m * n * k * 1e-9 / ms + memory_bandwidth_per_second = ( + sum( + [ + _.numel() * _.element_size() + for _ in [a_fp8, b_fp8, a_scale, b_scale, masked_m, out] + ] + ) + * 1e-9 + / ms + ) + print( + f"group_deepgemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} " + f"in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s" + f"memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s" + ) + + return tflops_per_second + + +if __name__ == "__main__": + print("=== DeepGEMM Grouped FP8 GEMM Benchmark ===\n") + + for batch_size in [1, 4, 8, 64, 128, 256]: + for m in [128, 256, 1024, 8192, 16384]: + for n, k in [(128, 512), (512, 128), (4096, 7168), (7168, 2048)]: + if m // batch_size < 128: + continue + if m * batch_size <= 16384: # Limit total problem size + bench_deepgemm_grouped_fp8_blackwell( + batch_size, m, n, k, torch.float8_e4m3fn, torch.bfloat16 + ) + + for batch_size in [1, 4, 8, 64, 128, 256]: + for m in [128, 256, 1024, 8192, 16384]: + for n, k in [(128, 512), (512, 128), (4096, 7168), (7168, 2048)]: + if m * batch_size <= 16384: # Limit total problem size + bench_deepgemm_batch_fp8_blackwell( + batch_size, m, n, k, torch.float8_e4m3fn, torch.bfloat16 + ) diff --git a/benchmarks/bench_deepseek_mla.py b/benchmarks/bench_deepseek_mla.py index eb147cb1c1..b13fc6c2fd 100644 --- a/benchmarks/bench_deepseek_mla.py +++ b/benchmarks/bench_deepseek_mla.py @@ -14,10 +14,11 @@ limitations under the License. """ +import numpy as np import torch -import triton import flashinfer +from flashinfer.testing.utils import bench_gpu_time def bench_deepseek_mla_decode(batch_size, seq_len, num_heads, backend): @@ -61,11 +62,12 @@ def bench_deepseek_mla_decode(batch_size, seq_len, num_heads, backend): ) o = wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False) - ms = triton.testing.do_bench( + measurements = bench_gpu_time( lambda: wrapper.run(q_nope, q_pe, ckv, kpe), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = sum([_.numel() * _.element_size() for _ in [q_nope, q_pe, ckv, kpe, o]]) flops = 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len diff --git a/benchmarks/bench_fused_add_rmsnorm.py b/benchmarks/bench_fused_add_rmsnorm.py index e03449d6e1..d03c7605ac 100644 --- a/benchmarks/bench_fused_add_rmsnorm.py +++ b/benchmarks/bench_fused_add_rmsnorm.py @@ -1,10 +1,10 @@ import argparse -from typing import cast +import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time @torch.inference_mode() @@ -42,7 +42,8 @@ def fn() -> None: flashinfer.fused_add_rmsnorm(x, residual, weight, eps) # Run benchmarking - latency_ms = cast(float, do_bench(fn)) + measurements = bench_gpu_time(fn) + latency_ms = np.median(measurements) throughput = ( x.numel() * x.element_size() * 2 + residual.numel() * residual.element_size() * 2 @@ -52,8 +53,8 @@ def fn() -> None: f"batch_size: {batch_size:3},", f"hidden_size: {hidden_size:5},", f"dtype: {dtype_str:8},", - f"latency: {latency_ms*1e3:2.0f}us,", - f"throughput: {throughput*1e-9:7.3f}GB/s", + f"latency: {latency_ms * 1e3:2.0f}us,", + f"throughput: {throughput * 1e-9:7.3f}GB/s", ) print("---") diff --git a/benchmarks/bench_groupwise_gemm_fp8_blackwell.py b/benchmarks/bench_groupwise_gemm_fp8_blackwell.py new file mode 100644 index 0000000000..451bb902e6 --- /dev/null +++ b/benchmarks/bench_groupwise_gemm_fp8_blackwell.py @@ -0,0 +1,202 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import numpy as np +import torch +import triton +import triton.language as tl + +from flashinfer.gemm import gemm_fp8_nt_groupwise +from flashinfer.testing.utils import bench_gpu_time + + +@triton.jit +def _w8a8_block_fp8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and store the result in output + tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def triton_w8a8_block_fp8_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + out: torch.Tensor, +) -> torch.Tensor: + M = A.shape[0] + N, K = B.shape + block_n, block_k = 128, 128 + + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + + def grid(META): + return ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + _w8a8_block_fp8_matmul[grid]( + A, + B, + out, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + out.stride(-2), + out.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return out + + +def bench_groupwise_gemm_fp8_blackwell(m, n, k, in_dtype, out_dtype): + a = torch.randn((m, k), device="cuda").to(in_dtype) + b = torch.randn((n, k), device="cuda").to(in_dtype) + a_scale = torch.rand((k // 128, m), dtype=torch.float32, device="cuda") + b_scale = torch.rand((k // 128, n // 128), dtype=torch.float32, device="cuda") + + out = torch.empty((m, n), dtype=out_dtype, device="cuda") + gemm_fp8_nt_groupwise(a, b, a_scale, b_scale, out=out, scale_major_mode="MN") + + measurements = bench_gpu_time( + lambda: gemm_fp8_nt_groupwise( + a, b, a_scale, b_scale, out=out, scale_major_mode="MN" + ) + ) + ms = np.median(measurements) + tflops_per_second = 2 * m * n * k * 1e-9 / ms + print( + f"gemm_fp8_nt_groupwise {m} {n} {k} {in_dtype} {out_dtype}: {tflops_per_second:.2f} TFLOPs/s" + ) + + tl_out = torch.empty((m, n), dtype=out_dtype, device="cuda") + a_scale = a_scale.transpose(0, 1).contiguous() + b_scale = b_scale.transpose(0, 1).contiguous() + measurements = bench_gpu_time( + lambda: triton_w8a8_block_fp8_matmul(a, b, a_scale, b_scale, tl_out) + ) + ms = np.median(measurements) + tflops_per_second = 2 * m * n * k * 1e-9 / ms + print( + f"triton_gemm_fp8_nt_groupwise {m} {n} {k} {in_dtype} {out_dtype}: {tflops_per_second:.2f} TFLOPs/s" + ) + + +if __name__ == "__main__": + for m in [1024, 2048, 4096, 8192]: + for n in [1024, 2048, 4096, 8192]: + for k in [1024, 2048, 4096, 8192]: + bench_groupwise_gemm_fp8_blackwell( + m, n, k, torch.float8_e5m2, torch.bfloat16 + ) diff --git a/benchmarks/bench_groupwise_grouped_gemm_fp8_blackwell.py b/benchmarks/bench_groupwise_grouped_gemm_fp8_blackwell.py new file mode 100644 index 0000000000..340c41a220 --- /dev/null +++ b/benchmarks/bench_groupwise_grouped_gemm_fp8_blackwell.py @@ -0,0 +1,64 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import bench_gpu_time + + +def bench_groupwise_grouped_gemm_fp8_blackwell( + batch_size, m, n, k, in_dtype, out_dtype +): + torch.random.manual_seed(0) + a = torch.randn(batch_size * m, k, device="cuda:0").to(in_dtype) + b = torch.randn(batch_size, n, k, device="cuda:0").to(in_dtype) + out = torch.empty(batch_size * m, n, device="cuda:0", dtype=out_dtype) + + a_scale = torch.randn( + (k // 128, batch_size * m), dtype=torch.float32, device="cuda:0" + ) + b_scale = torch.randn( + (batch_size, k // 128, n // 128), dtype=torch.float32, device="cuda:0" + ) + + segment_offsets = torch.arange( + 0, (batch_size + 1) * m, m, device="cuda:0", dtype=torch.int32 + ) + + measurements = bench_gpu_time( + lambda: flashinfer.gemm.group_gemm_fp8_nt_groupwise( + a, b, a_scale, b_scale, segment_offsets, out=out, mma_sm=2 + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + tflops_per_second = 2 * batch_size * m * n * k * 1e-9 / ms + print( + f"group_gemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s" + ) + + +if __name__ == "__main__": + for batch_size in [1, 3, 8, 16]: + for m in [128, 512, 1024, 2048, 4096, 8192]: + for n in [1024, 2048, 4096, 8192]: + for k in [1024, 2048, 4096, 8192]: + bench_groupwise_grouped_gemm_fp8_blackwell( + batch_size, m, n, k, torch.float8_e5m2, torch.bfloat16 + ) diff --git a/benchmarks/bench_groupwise_grouped_gemm_mxfp4_blackwell.py b/benchmarks/bench_groupwise_grouped_gemm_mxfp4_blackwell.py new file mode 100644 index 0000000000..c274e65592 --- /dev/null +++ b/benchmarks/bench_groupwise_grouped_gemm_mxfp4_blackwell.py @@ -0,0 +1,125 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +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. +""" + +from itertools import product + +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import bench_gpu_time + + +def bench_groupwise_grouped_gemm_mxfp4_blackwell( + group_size, m, n, k, in_dtype, out_dtype +): + torch.random.manual_seed(0) + assert n % 8 == 0 + assert k % 128 == 0 + tile_size = 32 + alignment_sf = 128 + fp8_info = torch.finfo(in_dtype) + a = ( + torch.empty(group_size * m, k, dtype=torch.float32, device="cuda:0") + .uniform_(-fp8_info.max, fp8_info.max) + .to(in_dtype) + ) + b = torch.randint( + 0, 256, (group_size, n, k // 2), dtype=torch.uint8, device="cuda:0" + ) + out = torch.empty(group_size * m, n, dtype=out_dtype, device="cuda:0") + + a_scale = torch.randint( + 0, + 256, + ( + (group_size * m + (alignment_sf - 1) * group_size) + // alignment_sf + * alignment_sf, + k // tile_size, + ), + dtype=torch.uint8, + device="cuda:0", + ) + b_scale = torch.randint( + 0, + 256, + ( + group_size, + (n + alignment_sf - 1) // alignment_sf * alignment_sf, + k // tile_size, + ), + dtype=torch.uint8, + device="cuda:0", + ) + + segment_offsets = torch.arange( + 0, (group_size + 1) * m, m, device="cuda:0", dtype=torch.int32 + ) + + ms_best = float("inf") + config_best = None + mma_sm_list = [1, 2] + tile_m_list = [128] + tile_n_list = [64, 128, 192, 256] + tile_k_list = [128, 256] + swap_ab_list = [True, False] + for mma_sm, tile_m, tile_n, tile_k, swap_ab in product( + mma_sm_list, tile_m_list, tile_n_list, tile_k_list, swap_ab_list + ): + measurements = bench_gpu_time( + lambda: flashinfer.gemm.group_gemm_mxfp4_nt_groupwise( + a, + b, + a_scale, + b_scale, + segment_offsets, + out=out, + mma_sm=mma_sm, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + swap_ab=swap_ab, + ), + dry_run_time_ms=10, + repeat_time_ms=100, + ) + ms = np.median(measurements) + if ms < ms_best: + ms_best = ms + config_best = { + "mma_sm": mma_sm, + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "swap_ab": swap_ab, + } + tflops_per_second = 2 * group_size * m * n * k * 1e-9 / ms_best + print( + f"group_gemm_mxfp4_nt_groupwise group_size={group_size} m={m} n={n} k={k} in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s" + ) + print(f"best config: {config_best}") + print() + + +if __name__ == "__main__": + for group_size in [1, 3, 8, 16]: + for m in [128, 512, 1024, 2048, 4096, 8192]: + for n in [1024, 2048, 4096, 8192]: + for k in [1024, 2048, 4096, 8192]: + bench_groupwise_grouped_gemm_mxfp4_blackwell( + group_size, m, n, k, torch.float8_e4m3fn, torch.bfloat16 + ) diff --git a/benchmarks/bench_hopper_attention.py b/benchmarks/bench_hopper_attention.py index f5bcc19eac..c1e56e6225 100644 --- a/benchmarks/bench_hopper_attention.py +++ b/benchmarks/bench_hopper_attention.py @@ -14,10 +14,14 @@ limitations under the License. """ +import numpy as np import torch -import triton import flashinfer +from flashinfer.testing.utils import ( + bench_gpu_time, + attention_tflops_per_sec_with_actual_seq_lens, +) def bench_single_prefill(seq_len, num_heads, causal, head_dim): @@ -27,21 +31,28 @@ def bench_single_prefill(seq_len, num_heads, causal, head_dim): v = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda") sm80_ms, sm90_ms = ( - triton.testing.do_bench( - lambda: flashinfer.single_prefill_with_kv_cache_return_lse( - q, k, v, causal=causal, backend=backend - ), - warmup=100, - rep=1000, + np.median( + bench_gpu_time( + lambda: flashinfer.single_prefill_with_kv_cache_return_lse( + q, k, v, causal=causal, backend=backend + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) ) for backend in ["fa2", "fa3"] ) def flops(ms): - if causal: - return seq_len * seq_len * num_qo_heads * head_dim * 2 / ms / 1e9 - else: - return seq_len * seq_len * num_qo_heads * head_dim * 4 / ms / 1e9 + return attention_tflops_per_sec_with_actual_seq_lens( + torch.tensor([seq_len]), + torch.tensor([seq_len]), + head_dim, + head_dim, + num_qo_heads, + causal, + ms, + ) print( f"bench_single_prefill (seq_len={seq_len}, num_heads={num_heads}, causal={causal}, head_dim={head_dim}), fa2-template: {flops(sm80_ms):.3f} TFLOPs/s, fa3-template: {flops(sm90_ms):.3f} TFLOPs/s" @@ -83,23 +94,26 @@ def bench_batch_ragged_prefill(batch_size, num_heads, seq_len, causal, head_dim) ) sm80_ms, sm90_ms = ( - triton.testing.do_bench( - lambda: wrapper.run(q, k, v), - warmup=100, - rep=1000, + np.median( + bench_gpu_time( + lambda: wrapper.run(q, k, v), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) ) for wrapper in [sm80_wrapper, sm90_wrapper] ) def flops(ms): - if causal: - return ( - batch_size * seq_len * seq_len * num_qo_heads * head_dim * 2 / ms / 1e9 - ) - else: - return ( - batch_size * seq_len * seq_len * num_qo_heads * head_dim * 4 / ms / 1e9 - ) + return attention_tflops_per_sec_with_actual_seq_lens( + torch.full((batch_size,), seq_len), + torch.full((batch_size,), seq_len), + head_dim, + head_dim, + num_qo_heads, + causal, + ms, + ) print( f"bench_batch_ragged_prefill (batch_size={batch_size}, num_heads={num_heads}, seq_len={seq_len}, causal={causal}, head_dim={head_dim}), fa2-template: {flops(sm80_ms):.3f} TFLOPs/s, fa3-template: {flops(sm90_ms):.3f} TFLOPs/s" @@ -160,23 +174,26 @@ def bench_batch_paged_prefill( ) sm80_ms, sm90_ms = ( - triton.testing.do_bench( - lambda: wrapper.run(q, (k, v)), - warmup=100, - rep=1000, + np.median( + bench_gpu_time( + lambda: wrapper.run(q, (k, v)), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) ) for wrapper in [sm80_wrapper, sm90_wrapper] ) def flops(ms): - if causal: - return ( - batch_size * seq_len * seq_len * num_qo_heads * head_dim * 2 / ms / 1e9 - ) - else: - return ( - batch_size * seq_len * seq_len * num_qo_heads * head_dim * 4 / ms / 1e9 - ) + return attention_tflops_per_sec_with_actual_seq_lens( + torch.full((batch_size,), seq_len), + torch.full((batch_size,), seq_len), + head_dim, + head_dim, + num_qo_heads, + causal, + ms, + ) print( f"bench_batch_paged_prefill (page_size={page_size} batch_size={batch_size}, num_heads={num_heads}, seq_len={seq_len}, causal={causal}, head_dim={head_dim}), fa2-template: {flops(sm80_ms):.3f} TFLOPs/s, fa3-template: {flops(sm90_ms):.3f} TFLOPs/s" @@ -184,6 +201,12 @@ def flops(ms): if __name__ == "__main__": + device_capability = torch.cuda.get_device_capability() + if device_capability[0] != 9: + print(f"Current device capability: {device_capability}.") + print("Current benchmark targets capability (9, 0). Returning...") + exit() + bench_batch_paged_prefill(1, 128, 32, 1024, True, 128) bench_batch_paged_prefill(1, 64, 32, 2048, True, 128) bench_batch_paged_prefill(1, 32, 32, 4096, True, 128) diff --git a/benchmarks/bench_hopper_fp8_attention.py b/benchmarks/bench_hopper_fp8_attention.py new file mode 100644 index 0000000000..89224af622 --- /dev/null +++ b/benchmarks/bench_hopper_fp8_attention.py @@ -0,0 +1,77 @@ +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import ( + bench_gpu_time, + attention_tflops_per_sec_with_actual_seq_lens, +) + + +def bench_single_prefill(seq_len, num_heads, causal, head_dim): + num_qo_heads = num_kv_heads = num_heads + q = torch.randn(seq_len, num_qo_heads, head_dim, dtype=torch.half, device="cuda") + k = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda") + v = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda") + + sm80_ms, sm90_ms = ( + np.median( + bench_gpu_time( + lambda: flashinfer.single_prefill_with_kv_cache_return_lse( + q, k, v, causal=causal, backend=backend + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ) + for backend in ["fa2", "fa3"] + ) + + q = torch.randn( + seq_len, num_qo_heads, head_dim, dtype=torch.half, device="cuda" + ).to(dtype=torch.float8_e4m3fn) + k = torch.randn( + seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda" + ).to(dtype=torch.float8_e4m3fn) + v = torch.randn( + seq_len, num_kv_heads, head_dim, dtype=torch.half, device="cuda" + ).to(dtype=torch.float8_e4m3fn) + + fp8_sm90_ms = np.median( + bench_gpu_time( + lambda: flashinfer.single_prefill_with_kv_cache_return_lse( + q, k, v, causal=causal, backend="fa3", o_dtype=torch.half + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ) + + def flops(ms): + return attention_tflops_per_sec_with_actual_seq_lens( + torch.tensor([seq_len]), + torch.tensor([seq_len]), + head_dim, + head_dim, + num_qo_heads, + causal, + ms, + ) + + print( + f"bench_single_prefill (seq_len={seq_len}, num_heads={num_heads}, causal={causal}, head_dim={head_dim}), fa2-template: {flops(sm80_ms):.3f} TFLOPs/s, fa3-template: {flops(sm90_ms):.3f} TFLOPs/s, fa3-fp8: {flops(fp8_sm90_ms):.3f} TFLOPs/s" + ) + + +if __name__ == "__main__": + device_capability = torch.cuda.get_device_capability() + if device_capability[0] != 9: + print(f"Current device capability: {device_capability}.") + print("Current benchmark targets capability (9, 0). Returning...") + exit() + + for seq_len in [4096, 8192, 16384]: + for num_heads in [24, 32]: + for causal in [True, False]: + for head_dim in [64, 128, 256]: + bench_single_prefill(seq_len, num_heads, causal, head_dim) diff --git a/benchmarks/bench_grouped_gemm.py b/benchmarks/bench_hopper_grouped_gemm.py similarity index 86% rename from benchmarks/bench_grouped_gemm.py rename to benchmarks/bench_hopper_grouped_gemm.py index 9ac3dfcc64..d4f314cfd3 100644 --- a/benchmarks/bench_grouped_gemm.py +++ b/benchmarks/bench_hopper_grouped_gemm.py @@ -16,9 +16,9 @@ import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time def bench_grouped_gemm( @@ -41,9 +41,10 @@ def bench_grouped_gemm( device="cuda:0", ) - ms = do_bench( + measurements = bench_gpu_time( lambda: segment_gemm.run(X, W, batch_size, True, out=Y, seg_indptr=seg_indptr) ) + ms = np.median(measurements) flops = 2 * batch_size * num_tokens_per_group * d_in * d_out print( @@ -53,6 +54,12 @@ def bench_grouped_gemm( if __name__ == "__main__": + device_capability = torch.cuda.get_device_capability() + if device_capability[0] != 9: + print(f"Current device capability: {device_capability}.") + print("Current benchmark targets capability (9, 0). Returning...") + exit() + for dtype_in in [torch.float8_e4m3fn, torch.bfloat16]: for dtype_out in [torch.bfloat16]: for batch_size in [1, 3, 8, 16]: diff --git a/benchmarks/bench_mixed_attention.py b/benchmarks/bench_mixed_attention.py new file mode 100644 index 0000000000..7414a58af0 --- /dev/null +++ b/benchmarks/bench_mixed_attention.py @@ -0,0 +1,323 @@ +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import bench_gpu_time + + +def run_bench( + p_qo_lens, + p_kv_lens, + d_qo_lens, + d_kv_lens, + # page_block_size=1, + num_kv_heads=4, + num_qo_heads=28, + head_dim=128, + device=0, + causal=True, +): + # POD Attention only supports page size = 1 due to use of single prefill kernel + page_block_size = 1 + seq_lens = torch.tensor(d_kv_lens + p_kv_lens, dtype=torch.int32) + q_lens = torch.tensor(d_qo_lens + p_qo_lens, dtype=torch.int32) + + seq_lens_blocks = torch.ceil(seq_lens / page_block_size).int() + p_seq_lens_blocks = torch.ceil( + torch.tensor(p_kv_lens, dtype=torch.int32) / page_block_size + ).int() + d_seq_lens_blocks = torch.ceil( + torch.tensor(d_kv_lens, dtype=torch.int32) / page_block_size + ).int() + + q_indptr = torch.cat([torch.tensor([0]), torch.cumsum(q_lens, 0)], dim=0).int() + kv_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(seq_lens_blocks, 0)], dim=0 + ).int() + + p_q_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(torch.tensor(p_qo_lens), 0)], dim=0 + ).int() + p_kv_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(p_seq_lens_blocks, 0)], dim=0 + ).int() + + d_q_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(torch.tensor(d_qo_lens), 0)], dim=0 + ).int() + d_kv_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(d_seq_lens_blocks, 0)], dim=0 + ).int() + num_blocks = kv_indptr[-1].item() + + q = torch.rand(q_indptr[-1].item(), num_qo_heads, head_dim).to( + device, dtype=torch.bfloat16 + ) + kv_data = torch.randn(num_blocks, 2, page_block_size, num_kv_heads, head_dim).to( + device, dtype=torch.bfloat16 + ) + + workspace_buffer = torch.empty(156 * 1024 * 1024, dtype=torch.uint8, device=device) + kv_layout = "NHD" + + wrapper_old = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout=kv_layout, + backend="fa2", + ) + last_page_len = (seq_lens - 1) % page_block_size + 1 + wrapper_old.plan( + q_indptr.to(device), + kv_indptr.to(device), + torch.arange(num_blocks).int().to(device), + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal=causal, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + o = wrapper_old.run(q, kv_data) + measurements = bench_gpu_time(lambda: wrapper_old.run(q, kv_data)) + ms_old = np.median(measurements) + + wrapper_persistent = flashinfer.BatchAttention(kv_layout="NHD") + wrapper_persistent.plan( + q_indptr.to(device), + kv_indptr.to(device), + torch.arange(num_blocks, dtype=torch.int32, device=device), + seq_lens.to(device), + num_qo_heads, + num_kv_heads, + head_dim, + head_dim, + page_block_size, + causal=causal, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + o_persistent, _ = wrapper_persistent.run(q, kv_data) + measurements_persistent = bench_gpu_time(lambda: wrapper_persistent.run(q, kv_data)) + ms_persistent = np.mean(measurements_persistent) + + # Batched POD Attention + q_d = q[: d_q_indptr[-1]] + kv_d = kv_data[: d_kv_indptr[-1]].unbind(1) + q_p = q[d_q_indptr[-1] :] + kv_p = kv_data[d_kv_indptr[-1] :].unbind(1) + kv_indices_d = torch.arange(0, d_kv_indptr[-1], device=device, dtype=torch.int32) + kv_indices_p = torch.arange(0, p_kv_indptr[-1], device=device, dtype=torch.int32) + + last_page_len_d = (d_seq_lens_blocks - 1) % page_block_size + 1 + last_page_len_p = (p_seq_lens_blocks - 1) % page_block_size + 1 + wrapper_pod = flashinfer.BatchPODWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout=kv_layout, + ) + + wrapper_pod.plan( + # Prefill params + p_q_indptr.to(device), + p_kv_indptr.to(device), + kv_indices_p.to(device), + last_page_len_p, + # Decode params + d_q_indptr.to(device), + d_kv_indptr.to(device), + kv_indices_d.to(device), + last_page_len_d, + # Common params + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + page_size=page_block_size, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + o_p_batch, o_d_batch = wrapper_pod.run( + q_p, + kv_p, + q_d, + kv_d, + causal_p=causal, + ) + o_batch_pod = torch.cat([o_d_batch, o_p_batch], dim=0) + + # Verify output matches + torch.testing.assert_close( + o_batch_pod, o, rtol=4e-3, atol=4e-3, msg="Batch POD-Attention decode mismatch!" + ) + measurements = bench_gpu_time( + lambda: wrapper_pod.run( + q_p, + kv_p, + q_d, + kv_d, + causal_p=causal, + ) + ) + ms_batch_pod = np.median(measurements) + + if len(p_kv_lens) == 1: + # Single POD attention + q_d = q[: d_q_indptr[-1]] + kv_d = kv_data[: d_kv_indptr[-1]].unbind(1) + q_p = q[d_q_indptr[-1] :] + k_p, v_p = kv_data[d_kv_indptr[-1] :].unbind(1) + k_p, v_p = k_p.squeeze(1), v_p.squeeze(1) + kv_indices_d = torch.arange( + 0, d_kv_indptr[-1], device=device, dtype=torch.int32 + ) + + last_page_len_d = (d_seq_lens_blocks - 1) % page_block_size + 1 + wrapper_pod = flashinfer.PODWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout=kv_layout, + ) + wrapper_pod.plan( + d_kv_indptr.to(device), + kv_indices_d.to(device), + last_page_len=last_page_len_d, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + page_size=page_block_size, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + o_p, o_d = wrapper_pod.run( + q_p, + k_p, + v_p, + q_d, + kv_data, + causal_p=causal, + ) + o_pod = torch.cat([o_d, o_p], dim=0) + # Verify output matches + torch.testing.assert_close( + o, o_pod, rtol=4e-3, atol=4e-3, msg="POD-Attention output mismatch!" + ) + measurements = bench_gpu_time( + lambda: wrapper_pod.run( + q_p, + k_p, + v_p, + q_d, + kv_d, + causal_p=causal, + causal_d=causal, + ) + ) + ms_pod = np.median(measurements) + + # Sequential two kernels: single prefill + batch decode (tensor cores) + # Prefill using single_prefill_with_kv_cache + def _run_single_prefill(): + return flashinfer.prefill.single_prefill_with_kv_cache( + q_p, + k_p, + v_p, + causal=causal, + pos_encoding_mode="NONE", + backend="fa2", + ) + + measurements_prefill = bench_gpu_time(lambda: _run_single_prefill()) + ms_prefill = np.median(measurements_prefill) + + # Batch decode using tensor cores + wrapper_decode = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, kv_layout=kv_layout, use_tensor_cores=True + ) + wrapper_decode.plan( + d_kv_indptr.to(device), + kv_indices_d.to(device), + last_page_len_d, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + data_type=torch.bfloat16, + q_data_type=torch.bfloat16, + ) + measurements_decode = bench_gpu_time(lambda: wrapper_decode.run(q_d, kv_d)) + ms_decode = np.median(measurements_decode) + ms_seq_two_kernels = ms_prefill + ms_decode + + print(f"Elapsed time (Batched Prefill): {ms_old:.2f} ms") + print(f"Elapsed time (Batched POD Attention): {ms_batch_pod:.2f} ms") + if len(p_kv_lens) == 1: + print(f"Elapsed time (POD Attention): {ms_pod:.2f} ms") + print(f"Elapsed time (Sequential two kernels): {ms_seq_two_kernels:.2f} ms") + print(f"Elapsed time (Persistent BatchAttention): {ms_persistent:.2f} ms") + print( + f"Batch POD speedup over Persistent BatchAttention: {ms_persistent / ms_batch_pod:.2f}x" + ) + + total_bytes = ( + q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + ) + print(f"Loading memory size (MB): {total_bytes / (1024**2):.2f} MB") + + bandwidth_old_gb_s = total_bytes / (ms_old * 1e-3) / (1024**3) + + print(f"Memory bandwidth (Batched Prefill): {bandwidth_old_gb_s:.2f} GB/s") + bandwidth_batch_pod_gb_s = total_bytes / (ms_batch_pod * 1e-3) / (1024**3) + print( + f"Memory bandwidth (Batched POD Attention): {bandwidth_batch_pod_gb_s:.2f} GB/s" + ) + if len(p_kv_lens) == 1: + bandwidth_pod_gb_s = total_bytes / (ms_pod * 1e-3) / (1024**3) + print(f"Memory bandwidth (POD Attention): {bandwidth_pod_gb_s:.2f} GB/s") + bandwidth_seq_gb_s = total_bytes / (ms_seq_two_kernels * 1e-3) / (1024**3) + print( + f"Memory bandwidth (Sequential two kernels): {bandwidth_seq_gb_s:.2f} GB/s" + ) + bandwidth_persistent_gb_s = total_bytes / (ms_persistent * 1e-3) / (1024**3) + print( + f"Memory bandwidth (Persistent BatchAttention): {bandwidth_persistent_gb_s:.2f} GB/s" + ) + + +if __name__ == "__main__": + np.random.seed(42) + torch.random.manual_seed(42) + + # Irregular sequence lengths for prefill and decode + d_q_len_configs = [[1] * 128] * 7 + d_kv_len_configs = [ + [2048] * 128, + [2048] * 128, + [2048] * 128, + [2048] * 128, + [4096] * 128, + [8192] * 128, + [8192] * 128, + ] + p_q_configs = [[512], [1536], [2048] * 2, [2048], [4096], [4096], [6000]] + p_kv_configs = [[512], [1536], [2048] * 2, [2048], [4096], [4096], [7000]] + + page_block_size = 1 + num_kv_heads = 8 + num_qo_heads = 32 + head_dim = 128 + + for idx, (p_q_lens, p_kv_lens, d_q_len, d_kv_len) in enumerate( + zip(p_q_configs, p_kv_configs, d_q_len_configs, d_kv_len_configs, strict=True) + ): + print(f"===== Benchmark {idx + 1}: (kv_len, qo_len) set =====") + run_bench( + p_q_lens, + p_kv_lens, + d_q_len, + d_kv_len, + # page_block_size=page_block_size, + num_kv_heads=num_kv_heads, + num_qo_heads=num_qo_heads, + head_dim=head_dim, + device=0, + causal=True, + ) diff --git a/benchmarks/bench_mm_fp8.py b/benchmarks/bench_mm_fp8.py new file mode 100644 index 0000000000..a4df76ebd9 --- /dev/null +++ b/benchmarks/bench_mm_fp8.py @@ -0,0 +1,98 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +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. +""" + +from typing import Dict +from flashinfer.autotuner import autotune +from flashinfer.trtllm_low_latency_gemm import prepare_low_latency_gemm_weights +import numpy as np +import torch + +from flashinfer import mm_fp8 +from flashinfer.testing.utils import bench_gpu_time + +_cache_permute_indices: Dict[torch.Size, torch.Tensor] = {} + + +def to_float8( + x: torch.Tensor, dtype=torch.float8_e4m3fn +) -> tuple[torch.Tensor, torch.Tensor]: + finfo = torch.finfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + scale = finfo.max / amax + x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) + return x_scl_sat.to(dtype), scale.float().reciprocal() + + +def bench_mm_fp8(m, n, k, in_dtype, out_dtype): + torch.manual_seed(123) + input_tensor = torch.randn([m, k], device="cuda", dtype=torch.bfloat16) + input_fp8, input_inv_s = to_float8(input_tensor, dtype=in_dtype) + + # mat2 row major -> column major + mat2 = torch.randn([n, k], device="cuda", dtype=torch.bfloat16) + mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=in_dtype) + + res = torch.zeros([m, n], device="cuda", dtype=out_dtype) + global_scale = input_inv_s * mat2_inv_s + + # Do row shuffling. + prepared_weights = prepare_low_latency_gemm_weights( + mat2_fp8, _cache_permute_indices + ) + + with autotune(True): + mm_fp8( + input_fp8, + prepared_weights, + global_scale, + out=res, + ) + + measurements = bench_gpu_time( + lambda: mm_fp8( + input_fp8, + prepared_weights, + global_scale, + res, + ), + dry_run_time_ms=500, + repeat_time_ms=2500, + use_cuda_graph=True, + ) + ms = np.median(measurements) + tflops_per_second = 2 * m * n * k * 1e-9 / ms + + bandwidth = ( + ( + input_fp8.numel() * input_fp8.element_size() + + prepared_weights.numel() * prepared_weights.element_size() + + res.numel() * res.element_size() + ) + / ms + / 1e9 + ) + + print( + f"mm_fp8 m={m} n={n} k={k} in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s over {ms:.6f} ms, {bandwidth:.2f} TB/s" + ) + + +if __name__ == "__main__": + for m in [1, 2, 4, 8, 16, 32, 64]: + for n in [2560, 5120, 8192]: + for k in [16384, 32768]: + bench_mm_fp8(m, n, k, torch.float8_e4m3fn, torch.bfloat16) diff --git a/benchmarks/bench_persistent_gemm.py b/benchmarks/bench_persistent_gemm.py index 07ae8170b8..98a5fb8ccf 100644 --- a/benchmarks/bench_persistent_gemm.py +++ b/benchmarks/bench_persistent_gemm.py @@ -1,10 +1,10 @@ -import pytest +import numpy as np import torch import triton -from triton.testing import do_bench import flashinfer import flashinfer.triton +from flashinfer.testing.utils import bench_gpu_time def is_cuda(): @@ -16,7 +16,7 @@ def supports_tma(): def bench_gemm_persistent(num_sms, dtype, M, N, K, reps=1000, warmup_reps=10000): - ms = do_bench( + measurements = bench_gpu_time( lambda: flashinfer.triton.sm_constraint_gemm.gemm_persistent( a=torch.randn((M, K), device="cuda", dtype=torch.float16).to(dtype), b=torch.randn((N, K), device="cuda", dtype=torch.float16).to(dtype), @@ -24,9 +24,10 @@ def bench_gemm_persistent(num_sms, dtype, M, N, K, reps=1000, warmup_reps=10000) beta=0.0, num_sms=num_sms, ), - warmup=warmup_reps, - rep=reps, + dry_run_time_ms=warmup_reps, + repeat_time_ms=reps, ) + ms = np.median(measurements) # matmul: 2 * M * N * K # scale and add: 3 * M * N @@ -41,7 +42,7 @@ def bench_gemm_descriptor_persistent( ): if dtype == torch.float32: return - ms = do_bench( + measurements = bench_gpu_time( lambda: flashinfer.triton.sm_constraint_gemm.gemm_descriptor_persistent( a=torch.randn((M, K), device="cuda", dtype=torch.float16).to(dtype), b=torch.randn((N, K), device="cuda", dtype=torch.float16).to(dtype), @@ -49,9 +50,10 @@ def bench_gemm_descriptor_persistent( beta=0.0, num_sms=num_sms, ), - warmup=warmup_reps, - rep=reps, + dry_run_time_ms=warmup_reps, + repeat_time_ms=reps, ) + ms = np.median(measurements) # matmul: 2 * M * N * K # scale and add: 3 * M * N diff --git a/benchmarks/bench_renorm.py b/benchmarks/bench_renorm.py index 30bf256d1b..aca54c3318 100644 --- a/benchmarks/bench_renorm.py +++ b/benchmarks/bench_renorm.py @@ -1,7 +1,8 @@ +import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time def normal_distribution(std): @@ -35,19 +36,20 @@ def main(): gumbel_distribution(0.1), gumbel_distribution(1), ]: - for p in [0.1, 0.5, 0.9]: + for p in [0.1, 0.5, 0.9, 1.0]: logits = distrib((batch_size, vocab_size), device="cuda") probs = torch.softmax(logits, dim=-1) - ms = do_bench( + measurements = bench_gpu_time( lambda: flashinfer.sampling.top_p_renorm_probs(probs, p), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = (probs.numel() * probs.element_size()) * 2 bandwidth = io * 1e-6 / ms print( - f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, p: {p}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, p: {p}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" ) print("---") @@ -63,16 +65,17 @@ def main(): for k in [10, 100, 1000, 5000]: logits = distrib((batch_size, vocab_size), device="cuda") probs = torch.softmax(logits, dim=-1) - ms = do_bench( + measurements = bench_gpu_time( lambda: flashinfer.sampling.top_k_renorm_probs(probs, k), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = (probs.numel() * probs.element_size()) * 2 bandwidth = io * 1e-6 / ms print( - f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" ) print("---") @@ -87,16 +90,17 @@ def main(): ]: for k in [10, 100, 1000, 5000]: logits = distrib((batch_size, vocab_size), device="cuda") - ms = do_bench( + measurements = bench_gpu_time( lambda: flashinfer.sampling.top_k_mask_logits(logits, k), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = (logits.numel() * logits.element_size()) * 2 bandwidth = io * 1e-6 / ms print( - f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" ) diff --git a/benchmarks/bench_rope.py b/benchmarks/bench_rope.py index 65d46a9c27..235788423a 100644 --- a/benchmarks/bench_rope.py +++ b/benchmarks/bench_rope.py @@ -6,9 +6,9 @@ $ python bench_rope.py """ -import math -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Optional, Tuple, Union +import numpy as np import torch import torch.nn as nn import triton @@ -17,10 +17,10 @@ ) from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace +from flashinfer.testing.utils import bench_gpu_time class FlashInferRotaryEmbedding(nn.Module): - def __init__( self, head_size: int, @@ -194,10 +194,12 @@ def benchmark( batch_size * seq_len, num_kv_heads * head_size, dtype=dtype, device=device ) - quantiles = [0.5, 0.2, 0.8] - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: rope_forward(pos_ids, query, key), quantiles=quantiles - ) + # Get raw measurements + measurements = bench_gpu_time(lambda: rope_forward(pos_ids, query, key)) + # Calculate statistics to match original return values + ms = np.median(measurements) + min_ms = np.percentile(measurements, 20) + max_ms = np.percentile(measurements, 80) return ms, min_ms, max_ms diff --git a/benchmarks/bench_rope_quantize_fp8.py b/benchmarks/bench_rope_quantize_fp8.py new file mode 100644 index 0000000000..952b1b4b41 --- /dev/null +++ b/benchmarks/bench_rope_quantize_fp8.py @@ -0,0 +1,343 @@ +import os +import sys +from typing import Union + +from torch import nn +import flashinfer +import numpy as np +import torch +import triton +from flashinfer.testing.utils import bench_gpu_time, bench_gpu_time_with_cudagraph + +# Add the project root to Python path to import test helpers +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) +from tests.test_helpers.rope_reference import RotaryEmbedding + +mode_ncu = bool(int(os.environ.get("FLASHINFER_MODE_NCU", "0"))) + + +class FlashInferRotaryEmbedding(nn.Module): + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.head_size = head_size + self.rotary_dim = rotary_dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.is_neox_style = is_neox_style + self.dtype = dtype + + cache = self._compute_cos_sin_cache() + self.cos_sin_cache: torch.Tensor + self.register_buffer("cos_sin_cache", cache, persistent=False) + + def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor: + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim + ) + ) + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + """Compute the cos and sin cache.""" + inv_freq = self._compute_inv_freq(self.base) + t = torch.arange(self.max_position_embeddings, dtype=torch.float) + + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + cache = torch.cat((cos, sin), dim=-1) + return cache + + def _apply_rotary_emb( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + is_neox_style: bool, + ) -> torch.Tensor: + """ + Args: + x: [num_tokens, num_heads, head_size] + cos: [num_tokens, head_size // 2] + sin: [num_tokens, head_size // 2] + is_neox_style: Whether to use the Neox-style or GPT-J-style rotary + positional embeddings. + """ + cos = cos.unsqueeze(-2).to(x.dtype) + sin = sin.unsqueeze(-2).to(x.dtype) + if is_neox_style: + x1, x2 = torch.chunk(x, 2, dim=-1) + else: + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + if is_neox_style: + return torch.cat((o1, o2), dim=-1) + else: + return torch.stack((o1, o2), dim=-1).flatten(-2) + + +def benchmark_config(config_name, num_tokens, provider, enable_pdl=False): + """Benchmark a specific attention configuration.""" + input_dtype = torch.bfloat16 + device = "cuda" + quant_dtype = torch.float8_e4m3fn + + # Configuration-specific parameters + if config_name == "mla": + # MLA: Original configuration for regression testing + num_qo_heads, num_kv_heads = 128, 1 + rope_dim, no_rope_dim = 64, 512 + elif config_name == "gqa": + # GQA: Realistic grouped-query attention + num_qo_heads, num_kv_heads = 32, 8 + rope_dim, no_rope_dim = 64, 64 + elif config_name == "mha": + # MHA: Standard multi-head attention + num_qo_heads, num_kv_heads = 32, 32 + rope_dim, no_rope_dim = 64, 64 + else: + raise ValueError(f"Unknown config: {config_name}") + + total_dim = rope_dim + no_rope_dim + + # Create input tensors for both implementations + if config_name == "mla": + # MLA: 2D K tensors (shared) + q_in = torch.randn( + num_tokens, num_qo_heads, total_dim, dtype=input_dtype, device=device + ) + k_in = torch.randn(num_tokens, total_dim, dtype=input_dtype, device=device) + else: + # GQA/MHA: 3D K tensors (multiple heads) + q_in = torch.randn( + num_tokens, num_qo_heads, total_dim, dtype=input_dtype, device=device + ) + k_in = torch.randn( + num_tokens, num_kv_heads, total_dim, dtype=input_dtype, device=device + ) + + pos_ids = torch.arange(num_tokens, device=device) + + # Create reference implementation + rope_ref = RotaryEmbedding( + head_size=total_dim, + rotary_dim=rope_dim, + max_position_embeddings=4096, + base=10000, + is_neox_style=False, + dtype=input_dtype, + device=device, + ) + + run_idx = 0 + + if provider == "flashinfer": + # Split tensors for FlashInfer + q_rope = q_in[..., :rope_dim] + q_nope = q_in[..., rope_dim:] + k_rope = k_in[..., :rope_dim] + k_nope = k_in[..., rope_dim:] + + # Create output tensors + q_rope_out = torch.empty_like(q_rope, dtype=quant_dtype) + q_nope_out = torch.empty_like(q_nope, dtype=quant_dtype) + k_rope_out = torch.empty_like(k_rope, dtype=quant_dtype) + k_nope_out = torch.empty_like(k_nope, dtype=quant_dtype) + + def execute(): + nonlocal run_idx + run_idx += 1 + + if mode_ncu and run_idx == 20: + torch.cuda.cudart().cudaProfilerStart() + + flashinfer.rope.rope_quantize_fp8( + q_rope=q_rope, + k_rope=k_rope, + q_nope=q_nope, + k_nope=k_nope, + cos_sin_cache=rope_ref.cos_sin_cache, + pos_ids=pos_ids, + is_neox=False, + q_rope_out=q_rope_out, + k_rope_out=k_rope_out, + q_nope_out=q_nope_out, + k_nope_out=k_nope_out, + quant_scale_q=1.0, + quant_scale_kv=1.0, + enable_pdl=enable_pdl, + ) + + if mode_ncu and run_idx == 20: + torch.cuda.cudart().cudaProfilerStop() + + elif provider == "torch": + # Create compiled version for better performance + @torch.compile + def torch_rope_quantize(q_in, k_in, pos_ids): + # Apply RoPE using reference implementation + q_out_f16, k_out_f16 = rope_ref.forward_native(pos_ids, q_in, k_in) + + # Quantize to FP8 (PyTorch native) + q_out_f8 = q_out_f16.to(quant_dtype) + k_out_f8 = k_out_f16.to(quant_dtype) + return q_out_f8, k_out_f8 + + # Warmup the compiled function + _ = torch_rope_quantize(q_in, k_in, pos_ids) + torch.cuda.synchronize() + + def execute(): + nonlocal run_idx + run_idx += 1 + + if mode_ncu and run_idx == 20: + torch.cuda.cudart().cudaProfilerStart() + + _ = torch_rope_quantize(q_in, k_in, pos_ids) + + if mode_ncu and run_idx == 20: + torch.cuda.cudart().cudaProfilerStop() + + else: + raise ValueError(f"Unknown provider: {provider}") + + if mode_ncu: + measurements = bench_gpu_time(execute) + else: + measurements = bench_gpu_time_with_cudagraph(execute) + + # Calculate statistics + ms = np.median(measurements) + min_ms = np.percentile(measurements, 20) + max_ms = np.percentile(measurements, 80) + + return ms, min_ms, max_ms + + +# Create separate benchmark functions for each architecture +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens"], + x_vals=[768] if mode_ncu else [1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 768], + line_arg="provider", + line_vals=["flashinfer", "torch"], + line_names=["FlashInfer", "PyTorch Compiled"], + styles=[("blue", "-"), ("blue", "--")], + ylabel="Latency (ms)", + plot_name="mla-rope-benchmark", + args={}, + ) +) +def benchmark_mla(provider, num_tokens): + return benchmark_config("mla", num_tokens, provider) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens"], + x_vals=[768] if mode_ncu else [1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 768], + line_arg="provider", + line_vals=["flashinfer", "torch"], + line_names=["FlashInfer", "PyTorch Compiled"], + styles=[("red", "-"), ("red", "--")], + ylabel="Latency (ms)", + plot_name="gqa-rope-benchmark", + args={}, + ) +) +def benchmark_gqa(provider, num_tokens): + return benchmark_config("gqa", num_tokens, provider) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens"], + x_vals=[768] if mode_ncu else [1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 768], + line_arg="provider", + line_vals=["flashinfer", "torch"], + line_names=["FlashInfer", "PyTorch Compiled"], + styles=[("green", "-"), ("green", "--")], + ylabel="Latency (ms)", + plot_name="mha-rope-benchmark", + args={}, + ) +) +def benchmark_mha(provider, num_tokens): + return benchmark_config("mha", num_tokens, provider) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens"], + x_vals=[768] if mode_ncu else [1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 768], + line_arg="enable_pdl", + line_vals=[False, True], + line_names=["enable_pdl=False", "enable_pdl=True"], + styles=[("blue", "-"), ("red", "-")], + ylabel="Latency (ms)", + plot_name="rope-pdl-benchmark", + args={}, + ) +) +def benchmark_pdl(enable_pdl, num_tokens): + return benchmark_config("mla", num_tokens, "flashinfer", enable_pdl=enable_pdl) + + +if __name__ == "__main__": + # Run all benchmarks and generate individual plots + print("Running MLA benchmark...") + benchmark_mla.run(print_data=False, show_plots=True, save_path=".") + + print("Running GQA benchmark...") + benchmark_gqa.run(print_data=False, show_plots=True, save_path=".") + + print("Running MHA benchmark...") + benchmark_mha.run(print_data=False, show_plots=True, save_path=".") + + print("Running PDL benchmark...") + benchmark_pdl.run(print_data=False, show_plots=True, save_path=".") + + # Collect results for summary table + token_counts = ( + [1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 768] if not mode_ncu else [768] + ) + + print("\n=== Summary Table ===") + print( + f"{'Tokens':<8} {'MLA-FI (ms)':<12} {'MLA-Torch (ms)':<14} {'GQA-FI (ms)':<12} {'GQA-Torch (ms)':<14} {'MHA-FI (ms)':<12} {'MHA-Torch (ms)':<14}" + ) + print("-" * 90) + for num_tokens in token_counts: + mla_fi_ms, _, _ = benchmark_config("mla", num_tokens, "flashinfer") + mla_torch_ms, _, _ = benchmark_config("mla", num_tokens, "torch") + gqa_fi_ms, _, _ = benchmark_config("gqa", num_tokens, "flashinfer") + gqa_torch_ms, _, _ = benchmark_config("gqa", num_tokens, "torch") + mha_fi_ms, _, _ = benchmark_config("mha", num_tokens, "flashinfer") + mha_torch_ms, _, _ = benchmark_config("mha", num_tokens, "torch") + print( + f"{num_tokens:<8} {mla_fi_ms:<12.5f} {mla_torch_ms:<14.5f} {gqa_fi_ms:<12.5f} {gqa_torch_ms:<14.5f} {mha_fi_ms:<12.5f} {mha_torch_ms:<14.5f}" + ) + + print("\nConfiguration details:") + print(" MLA: 128 Q heads, 1 K head, 64+512 dims") + print(" GQA: 32 Q heads, 8 K heads, 64+64 dims") + print(" MHA: 32 Q heads, 32 K heads, 64+64 dims") + + print("\nPlot files saved to current directory:") + print(" mla-rope-benchmark.png (FlashInfer vs PyTorch)") + print(" gqa-rope-benchmark.png (FlashInfer vs PyTorch)") + print(" mha-rope-benchmark.png (FlashInfer vs PyTorch)") + print(" rope-pdl-benchmark.png (enable_pdl=False vs enable_pdl=True)") diff --git a/benchmarks/bench_rope_quantize_fp8_append_cache.py b/benchmarks/bench_rope_quantize_fp8_append_cache.py new file mode 100644 index 0000000000..3119b9fef8 --- /dev/null +++ b/benchmarks/bench_rope_quantize_fp8_append_cache.py @@ -0,0 +1,342 @@ +""" +Copyright (c) 2024 by FlashInfer team. + +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 os +import sys +import argparse +import flashinfer +import numpy as np +import torch +from flashinfer.testing.utils import bench_gpu_time_with_cudagraph +from flashinfer.utils import get_gpu_memory_bandwidth + +# Add the project root to Python path to import test helpers +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) +from tests.test_helpers.rope_reference import RotaryEmbedding + + +def benchmark_config( + config_name, + num_tokens, + batch_size=4, + page_size=16, + enable_pdl=False, + single_run=False, +): + """Benchmark a specific attention configuration with paged KV cache append.""" + input_dtype = torch.bfloat16 + device = "cuda" + quant_dtype = torch.float8_e4m3fn + + # Configuration-specific parameters + if config_name == "mla": + # MLA: DeepSeek-style multi-latent attention + num_qo_heads, num_kv_heads = 128, 1 + rope_dim, no_rope_dim = 64, 512 + elif config_name == "gqa": + # GQA: Grouped-query attention (e.g., Llama-style) + num_qo_heads, num_kv_heads = 32, 8 + rope_dim, no_rope_dim = 64, 64 + elif config_name == "mha": + # MHA: Standard multi-head attention + num_qo_heads, num_kv_heads = 32, 32 + rope_dim, no_rope_dim = 64, 64 + else: + raise ValueError(f"Unknown config: {config_name}") + + head_dim = rope_dim + no_rope_dim + + # Create input tensors + if config_name == "mla": + # MLA: 2D K tensors (shared) + q_rope = torch.randn( + num_tokens, num_qo_heads, rope_dim, dtype=input_dtype, device=device + ) + q_nope = torch.randn( + num_tokens, num_qo_heads, no_rope_dim, dtype=input_dtype, device=device + ) + k_rope = torch.randn(num_tokens, rope_dim, dtype=input_dtype, device=device) + k_nope = torch.randn(num_tokens, no_rope_dim, dtype=input_dtype, device=device) + v = None + else: + # GQA/MHA: 3D K/V tensors + q_rope = torch.randn( + num_tokens, num_qo_heads, rope_dim, dtype=input_dtype, device=device + ) + q_nope = torch.randn( + num_tokens, num_qo_heads, no_rope_dim, dtype=input_dtype, device=device + ) + k_rope = torch.randn( + num_tokens, num_kv_heads, rope_dim, dtype=input_dtype, device=device + ) + k_nope = torch.randn( + num_tokens, num_kv_heads, no_rope_dim, dtype=input_dtype, device=device + ) + v = torch.randn( + num_tokens, num_kv_heads, head_dim, dtype=input_dtype, device=device + ) + + # Create RoPE reference for cos/sin cache (ensure it covers this run) + max_seq_len = int(num_tokens) + rope_ref = RotaryEmbedding( + head_size=head_dim, + rotary_dim=rope_dim, + max_position_embeddings=max_seq_len, + base=10000, + is_neox_style=False, + dtype=input_dtype, + device=device, + ) + pos_ids = torch.arange(num_tokens, device=device, dtype=torch.int32) + + # Build paged metadata (single request with all tokens) + kv_append_length = torch.tensor( + [num_tokens] + [0] * (batch_size - 1), dtype=torch.int32, device=device + ) + kv_append_indptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + torch.cumsum(kv_append_length, dim=0), + ] + ) + num_pages_per_req = torch.tensor( + [(num_tokens + page_size - 1) // page_size] + [0] * (batch_size - 1), + dtype=torch.int32, + device=device, + ) + kv_page_indptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + torch.cumsum(num_pages_per_req, dim=0), + ] + ) + kv_page_indices = torch.arange( + kv_page_indptr[-1].item(), dtype=torch.int32, device=device + ) + kv_last_page_len = torch.tensor( + [num_tokens % page_size if num_tokens % page_size != 0 else page_size] + + [0] * (batch_size - 1), + dtype=torch.int32, + device=device, + ) + + # Get batch_indices and positions + seq_lens = flashinfer.get_seq_lens(kv_page_indptr, kv_last_page_len, page_size) + batch_indices, positions = flashinfer.get_batch_indices_positions( + kv_append_indptr, seq_lens, num_tokens + ) + + # Allocate caches + max_pages = kv_page_indptr[-1].item() + + if config_name == "mla": + ckv_cache = torch.zeros( + max_pages, page_size, no_rope_dim, dtype=quant_dtype, device=device + ) + kpe_cache = torch.zeros( + max_pages, page_size, rope_dim, dtype=quant_dtype, device=device + ) + paged_kv_cache = (ckv_cache, kpe_cache) + else: + # GQA/MHA: use NHD layout + k_cache = torch.zeros( + max_pages, + page_size, + num_kv_heads, + head_dim, + dtype=quant_dtype, + device=device, + ) + v_cache = torch.zeros( + max_pages, + page_size, + num_kv_heads, + head_dim, + dtype=quant_dtype, + device=device, + ) + paged_kv_cache = (k_cache, v_cache) + + run_idx = 0 + + def execute(): + if single_run: + import torch.cuda.nvtx as nvtx + + nvtx.range_push("rope_append") + nonlocal run_idx + run_idx += 1 + + flashinfer.rope.rope_quantize_fp8_append_paged_kv_cache( + q_rope=q_rope, + k_rope=k_rope, + q_nope=q_nope, + k_nope=k_nope, + v=v, + cos_sin_cache=rope_ref.cos_sin_cache, + pos_ids=pos_ids, + paged_kv_cache=paged_kv_cache, + kv_indices=kv_page_indices, + kv_indptr=kv_page_indptr, + batch_indices=batch_indices, + positions=positions, + page_size=page_size, + kv_layout="NHD" if config_name != "mla" else "NHD", + quantize_dtype=quant_dtype, + quant_scale_q=1.0, + quant_scale_kv=1.0, + is_neox=False, + enable_pdl=enable_pdl, + ) + if single_run: + # Ensure kernels complete inside the NVTX range for ncu filtering + torch.cuda.synchronize() + nvtx.range_pop() + + if single_run: + execute() + return None, None, None, None, None + measurements = bench_gpu_time_with_cudagraph(execute) + + # Calculate I/O bytes + # Inputs: q_rope, k_rope, q_nope, k_nope, v (if not MLA), cos_sin_cache, pos_ids + io_bytes = ( + q_rope.numel() * q_rope.element_size() + + k_rope.numel() * k_rope.element_size() + + q_nope.numel() * q_nope.element_size() + + k_nope.numel() * k_nope.element_size() + + rope_ref.cos_sin_cache.numel() * rope_ref.cos_sin_cache.element_size() + + pos_ids.numel() * pos_ids.element_size() + ) + + if v is not None: + io_bytes += v.numel() * v.element_size() + + # Outputs: q_rope_out, q_nope_out (FP8), cache writes (FP8) + io_bytes += ( + q_rope.numel() * torch.finfo(quant_dtype).bits // 8 + + q_nope.numel() * torch.finfo(quant_dtype).bits // 8 + ) + + if config_name == "mla": + # MLA writes to ckv_cache and kpe_cache + io_bytes += ( + num_tokens * no_rope_dim * torch.finfo(quant_dtype).bits // 8 + + num_tokens * rope_dim * torch.finfo(quant_dtype).bits // 8 + ) + else: + # GQA/MHA writes to k_cache and v_cache + io_bytes += ( + num_tokens * num_kv_heads * head_dim * torch.finfo(quant_dtype).bits // 8 + + num_tokens * num_kv_heads * head_dim * torch.finfo(quant_dtype).bits // 8 + ) + + # Calculate statistics + ms = np.median(measurements) + min_ms = np.percentile(measurements, 20) + max_ms = np.percentile(measurements, 80) + + # Calculate bandwidth in GB/s + bandwidth_gb_s = io_bytes / ms / 1e6 + + # Calculate TFLOPs (FP operations) + # RoPE: 6 FLOPs per dimension pair (2 muls + 1 sub for real, 2 muls + 1 add for imag) + # For Q: num_tokens * num_qo_heads * (rope_dim/2) pairs * 6 FLOPs + # For K: depends on architecture + q_flops = num_tokens * num_qo_heads * (rope_dim / 2) * 6 + + if config_name == "mla": + # MLA: K is 2D (no head dimension) + k_flops = num_tokens * (rope_dim / 2) * 6 + else: + # GQA/MHA: K is 3D (has head dimension) + k_flops = num_tokens * num_kv_heads * (rope_dim / 2) * 6 + + total_flops = q_flops + k_flops + tflops = ( + total_flops / ms / 1e9 + ) # TFLOPs (operations per ms = operations per second / 1e12) + + return ms, min_ms, max_ms, bandwidth_gb_s, tflops + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--ncu-single", action="store_true", help="Run a single execute() for ncu" + ) + parser.add_argument( + "--config", type=str, default="", help="Config name: mla/gqa/mha" + ) + parser.add_argument("--num-tokens", type=int, default=0) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument("--enable-pdl", type=int, default=0) + args, unknown = parser.parse_known_args() + + if args.ncu_single: + # Minimal single-run for ncu profiling + cfg = args.config or "mla" + ntok = int(args.num_tokens) + pgsz = int(args.page_size) + en_pdl = bool(int(args.enable_pdl)) + # Force a single execution path + benchmark_config(cfg, ntok, page_size=pgsz, enable_pdl=en_pdl, single_run=True) + sys.exit(0) + + # Get GPU information (for display only) + device = torch.device("cuda:0") + gpu_name = torch.cuda.get_device_name(0) + gpu_peak_bandwidth = get_gpu_memory_bandwidth(device) + print(f"\nDetected GPU: {gpu_name}") + print(f"Theoretical Peak Memory Bandwidth: {gpu_peak_bandwidth:.2f} GB/s") + print() + + # Token counts to benchmark + token_counts = [1, 32, 128, 384, 768, 1024, 2048, 4096, 8192] + + # Helper function to print a table for a specific configuration + def print_config_table(config_name, config_desc): + page_size_to_benchmark = 32 + print(f"\n{'=' * 100}") + print(f" {config_name.upper()}: {config_desc}") + print(f"{'=' * 100}") + + print( + f"{'Tokens':<10} {'Time (ms)':<12} {'BW (GB/s)':<12} {'BW% (Peak)':<14} {'TFLOPs':<12}" + ) + print("-" * 70) + for num_tokens in token_counts: + ms, _, _, bw, tflops = benchmark_config( + config_name, num_tokens, page_size=page_size_to_benchmark + ) + bw_pct = (bw / gpu_peak_bandwidth) * 100 + print( + f"{num_tokens:<10} {ms:<12.5f} {bw:<12.2f} {bw_pct:<14.1f} {tflops:<12.3f}" + ) + + # Print tables for each configuration + print_config_table("mla", "128 Q heads, 1 K head, 64+512 dims (DeepSeek-style)") + print_config_table("gqa", "32 Q heads, 8 K heads, 64+64 dims (Llama-style)") + print_config_table("mha", "32 Q heads, 32 K heads, 64+64 dims (Standard)") + + print("\n" + "=" * 100) + print("Configuration details:") + print(" Page size: 32, Batch size: 4") + print(" Token range: 1 (single decode) → 8192 (large prefill)") + print(f" GPU: {gpu_name}") + print(f" Theoretical Peak Memory Bandwidth: {gpu_peak_bandwidth:.2f} GB/s") + print(" BW% calculated as: (achieved_bandwidth / peak_bandwidth) * 100") + print("=" * 100) diff --git a/benchmarks/bench_sampling.py b/benchmarks/bench_sampling.py index 3eba949c95..cc2406e43f 100644 --- a/benchmarks/bench_sampling.py +++ b/benchmarks/bench_sampling.py @@ -1,7 +1,8 @@ +import numpy as np import torch -from triton.testing import do_bench import flashinfer +from flashinfer.testing.utils import bench_gpu_time def normal_distribution(std): @@ -27,6 +28,18 @@ def init_seed_sampling(*args, **kwargs): return flashinfer.sampling.sampling_from_probs(*args, **kwargs) +def init_seed_sampling_from_logits(*args, **kwargs): + torch.manual_seed(42) + return flashinfer.sampling.sampling_from_logits(*args, **kwargs) + + +def init_seed_sampling_from_softmax_logits(logits, *args, **kwargs): + torch.manual_seed(42) + return flashinfer.sampling.sampling_from_probs( + torch.softmax(logits, dim=-1), *args, **kwargs + ) + + def init_seed_top_k_sampling(*args, **kwargs): torch.manual_seed(42) return flashinfer.sampling.top_k_sampling_from_probs(*args, **kwargs) @@ -55,11 +68,12 @@ def main(): samples = torch.zeros( batch_size, dtype=torch.int32, device=probs.device ) - ms = do_bench( + measurements = bench_gpu_time( lambda: init_seed_sampling(probs, deterministic=deterministic), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = ( probs.numel() * probs.element_size() @@ -67,7 +81,7 @@ def main(): ) bandwidth = io * 1e-6 / ms print( - f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" ) print("---") @@ -87,13 +101,14 @@ def main(): samples = torch.zeros( batch_size, dtype=torch.int32, device=probs.device ) - ms = do_bench( + measurements = bench_gpu_time( lambda: init_seed_top_k_sampling( probs, k, deterministic=deterministic ), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = ( probs.numel() * probs.element_size() @@ -101,7 +116,7 @@ def main(): ) bandwidth = io * 1e-6 / ms print( - f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, k: {k}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" ) print("---") @@ -122,13 +137,14 @@ def main(): samples = torch.zeros( batch_size, dtype=torch.int32, device=probs.device ) - ms = do_bench( + measurements = bench_gpu_time( lambda: init_seed_top_p_sampling( probs, p, deterministic=deterministic ), - warmup=100, - rep=1000, + dry_run_time_ms=100, + repeat_time_ms=1000, ) + ms = np.median(measurements) io = ( probs.numel() * probs.element_size() @@ -136,9 +152,154 @@ def main(): ) bandwidth = io * 1e-6 / ms print( - f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, p: {p}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, p: {p}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" ) + print("---") + print("sampling from softmax(logits)") + for vocab_size in [128512]: + for batch_size in [1, 16, 32, 64, 128, 256, 512]: + for distrib in [ + normal_distribution(1), + normal_distribution(5), + gumbel_distribution(0.1), + gumbel_distribution(1), + ]: + for deterministic in [True, False]: + logits = distrib((batch_size, vocab_size), device="cuda") + samples = torch.zeros( + batch_size, dtype=torch.int32, device=logits.device + ) + measurements = bench_gpu_time( + lambda: init_seed_sampling_from_softmax_logits( + logits, samples, deterministic=deterministic + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + io = ( + logits.numel() * logits.element_size() + + samples.numel() * samples.element_size() + ) + bandwidth = io * 1e-6 / ms + print( + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + ) + + print("---") + print("sampling from logits") + for vocab_size in [128512]: + for batch_size in [1, 16, 32, 64, 128, 256, 512]: + for distrib in [ + normal_distribution(1), + normal_distribution(5), + gumbel_distribution(0.1), + gumbel_distribution(1), + ]: + for deterministic in [True, False]: + logits = distrib((batch_size, vocab_size), device="cuda") + samples = torch.zeros( + batch_size, dtype=torch.int32, device=logits.device + ) + measurements = bench_gpu_time( + lambda: init_seed_sampling_from_logits( + logits, samples, deterministic=deterministic + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + + io = ( + logits.numel() * logits.element_size() + + samples.numel() * samples.element_size() + ) + bandwidth = io * 1e-6 / ms + print( + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, deterministic: {deterministic}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + ) + + print("---") + print("top-p renorm probs") + for vocab_size in [128512]: + for batch_size in [1, 16, 32, 64, 128, 256, 512]: + torch.manual_seed(42) + for distrib in [ + normal_distribution(1), + normal_distribution(5), + gumbel_distribution(0.1), + gumbel_distribution(1), + ]: + for p in [0.1, 0.5, 0.9]: + logits = distrib((batch_size, vocab_size), device="cuda") + probs = torch.softmax(logits, dim=-1) + measurements = bench_gpu_time( + lambda: flashinfer.sampling.top_p_renorm_probs(probs, p), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + + io = probs.numel() * probs.element_size() * 2 + bandwidth = io * 1e-6 / ms + print( + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, p: {p}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + ) + + print("---") + print("top-k renorm probs") + for vocab_size in [128512]: + for batch_size in [1, 16, 32, 64, 128, 256, 512]: + torch.manual_seed(42) + for distrib in [ + normal_distribution(1), + normal_distribution(5), + gumbel_distribution(0.1), + gumbel_distribution(1), + ]: + for k in [10, 100, 1000, 5000]: + logits = distrib((batch_size, vocab_size), device="cuda") + probs = torch.softmax(logits, dim=-1) + measurements = bench_gpu_time( + lambda: flashinfer.sampling.top_k_renorm_probs(probs, k), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + + io = probs.numel() * probs.element_size() * 2 + bandwidth = io * 1e-6 / ms + print( + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + ) + + print("---") + print("top-k mask logits") + for vocab_size in [128512]: + for batch_size in [1, 16, 32, 64, 128, 256, 512]: + torch.manual_seed(42) + for distrib in [ + normal_distribution(1), + normal_distribution(5), + gumbel_distribution(0.1), + gumbel_distribution(1), + ]: + for k in [10, 100, 1000, 5000]: + logits = distrib((batch_size, vocab_size), device="cuda") + measurements = bench_gpu_time( + lambda: flashinfer.sampling.top_k_mask_logits(logits, k), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + + io = logits.numel() * logits.element_size() * 2 + bandwidth = io * 1e-6 / ms + print( + f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s" + ) + if __name__ == "__main__": main() diff --git a/benchmarks/bench_sliding_window.py b/benchmarks/bench_sliding_window.py new file mode 100644 index 0000000000..3454ae79d8 --- /dev/null +++ b/benchmarks/bench_sliding_window.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +import itertools +from dataclasses import dataclass + +import torch +from triton.testing import do_bench + +# Optional: pin to a device via env CUDA_VISIBLE_DEVICES +DEVICE = torch.device("cuda:0") + +import flashinfer + + +@dataclass +class Case: + batch_size: int + kv_len: int + qo_len: int + window_left: int + num_kv_heads: int + num_qo_heads: int + head_dim: int + page_size: int + + +# Same grids as your pytest params +BATCH_SIZES = [16] +KV_LENS = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384] +QO_LENS = [1, 16] +WINDOW_LEFTS = [128] +NUM_KV_HEADS = [8] +NUM_QO_HEADS = [32] +HEAD_DIMS = [64] +PAGE_SIZES = [16] + +DTYPE = torch.float16 +INDEX_DTYPE = torch.int32 + + +# --- FLOPs & bytes estimation helpers ---------------------------------------- +def total_attn_pairs(kv_len: int, qo_len: int, window_left: int) -> int: + total = 0 + for i in range(1, qo_len + 1): + total += min(kv_len - i + 1, window_left) + return total + + +def estimate_flops( + batch: int, + num_qo_heads: int, + head_dim: int, + kv_len: int, + qo_len: int, + window_left: int, +) -> float: + """ + Matmul-only FLOPs (ignoring softmax) ~ QK^T + softmax@V ≈ 4 * head_dim * total_pairs + """ + pairs = total_attn_pairs(kv_len, qo_len, window_left) + return 4.0 * head_dim * pairs * batch * num_qo_heads + + +def estimate_min_io_bytes( + batch: int, + num_kv_heads: int, + num_qo_heads: int, + head_dim: int, + kv_len: int, + qo_len: int, + window_left: int, + dtype_size: int = 2, +) -> int: + """ + A *rough* lower-bound I/O estimate (no re-reads, no intermediate buffers): + - Read Q: B * qo * Hq * D + - Read unique K, V from KV cache: B * min(kv_len, window_left) * Hk * D each + - Write O: B * qo * Hq * D + NOTE: Real kernels may read more/less due to paging, softmax stats, and reuse. + """ + uniq_kv = min(kv_len, window_left) + q_bytes = batch * qo_len * num_qo_heads * head_dim * dtype_size + k_bytes = batch * uniq_kv * num_kv_heads * head_dim * dtype_size + v_bytes = batch * uniq_kv * num_kv_heads * head_dim * dtype_size + o_bytes = batch * qo_len * num_qo_heads * head_dim * dtype_size + return q_bytes + k_bytes + v_bytes + o_bytes + + +# --- Benchmark runner --------------------------------------------------------- +def run_one(case: Case, warmup=25, rep=100): + torch.cuda.synchronize() + # Inputs (mirror your test) + q = torch.randn( + case.batch_size * case.qo_len, + case.num_qo_heads, + case.head_dim, + dtype=DTYPE, + device=DEVICE, + ) + q_indptr = ( + torch.arange(0, case.batch_size + 1, device=DEVICE, dtype=INDEX_DTYPE) + * case.qo_len + ) + + num_pages_per_seq = (case.kv_len + case.page_size - 1) // case.page_size + total_num_pages = num_pages_per_seq * case.batch_size + + k_data = torch.randn( + total_num_pages, + case.page_size, + case.num_kv_heads, + case.head_dim, + dtype=DTYPE, + device=DEVICE, + ) + v_data = torch.randn( + total_num_pages, + case.page_size, + case.num_kv_heads, + case.head_dim, + dtype=DTYPE, + device=DEVICE, + ) + + kv_indptr = ( + torch.arange(0, case.batch_size + 1, device=DEVICE, dtype=INDEX_DTYPE) + * num_pages_per_seq + ) + kv_indices = torch.arange(0, total_num_pages, device=DEVICE, dtype=INDEX_DTYPE) + kv_last_page_len = torch.full( + (case.batch_size,), + (case.kv_len - 1) % case.page_size + 1, + dtype=INDEX_DTYPE, + device=DEVICE, + ) + + # Workspace & wrapper + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=DEVICE) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, "NHD", backend="fa2" + ) + + # Plan + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + case.num_qo_heads, + case.num_kv_heads, + case.head_dim, + case.page_size, + window_left=case.window_left, + causal=False, + ) + + # Warmup + timing + def _run(): + o = wrapper.run(q, (k_data, v_data)) + return o + + # Warmup (helps stabilize clocks/caches) + for _ in range(warmup): + _ = _run() + torch.cuda.synchronize() + + ms = do_bench(_run, rep=rep) # returns milliseconds per run + torch.cuda.synchronize() + + # Metrics + flops = estimate_flops( + case.batch_size, + case.num_qo_heads, + case.head_dim, + case.kv_len, + case.qo_len, + case.window_left, + ) + bytes_min = estimate_min_io_bytes( + case.batch_size, + case.num_kv_heads, + case.num_qo_heads, + case.head_dim, + case.kv_len, + case.qo_len, + case.window_left, + dtype_size=2, + ) + + s = ms / 1e3 + tflops = (flops / s) / 1e12 + gbps = (bytes_min / s) / 1e9 + + return ms, tflops, gbps + + +def main(): + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + + configs = [ + Case(b, kv, qo, w, hk, hq, d, ps) + for b, qo, w, hk, hq, d, ps, kv in itertools.product( + BATCH_SIZES, + QO_LENS, + WINDOW_LEFTS, + NUM_KV_HEADS, + NUM_QO_HEADS, + HEAD_DIMS, + PAGE_SIZES, + KV_LENS, + ) + ] + + # Header + print( + "batch kv_len qo_len win_left kv_heads qo_heads head_dim page_size | " + "latency_ms est_TFLOPs est_GB/s(min-IO)" + ) + print("-" * 110) + + for c in configs: + try: + ms, tflops, gbps = run_one(c) + print( + f"{c.batch_size:5d} {c.kv_len:6d} {c.qo_len:6d} {c.window_left:8d} " + f"{c.num_kv_heads:8d} {c.num_qo_heads:8d} {c.head_dim:8d} {c.page_size:9d} | " + f"{ms:10.3f} {tflops:10.3f} {gbps:15.3f}" + ) + except Exception as e: + # Keep going if some configs are unsupported by the current build + print( + f"{c.batch_size:5d} {c.kv_len:6d} {c.qo_len:6d} {c.window_left:8d} " + f"{c.num_kv_heads:8d} {c.num_qo_heads:8d} {c.head_dim:8d} {c.page_size:9d} | " + f"ERROR: {type(e).__name__}: {e}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_softmax.py b/benchmarks/bench_softmax.py new file mode 100755 index 0000000000..6da8dc9fcb --- /dev/null +++ b/benchmarks/bench_softmax.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +""" +Benchmark script comparing torch.softmax vs flashinfer.softmax performance. +Creates a heatmap showing speedup across different batch sizes and hidden dimensions. +""" + +import numpy as np +import torch +import matplotlib.pyplot as plt +import seaborn as sns +from typing import List, Tuple +import flashinfer +from flashinfer.testing.utils import bench_gpu_time + + +@torch.inference_mode() +def benchmark_torch_softmax(logits: torch.Tensor) -> float: + """Benchmark torch's native softmax.""" + measurements = bench_gpu_time( + lambda: torch.softmax(logits, dim=-1), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + return np.median(measurements) + + +@torch.inference_mode() +def benchmark_flashinfer_softmax(logits: torch.Tensor) -> float: + """Benchmark flashinfer's softmax.""" + measurements = bench_gpu_time( + lambda: flashinfer.sampling.softmax(logits, temperature=None, enable_pdl=False), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + return np.median(measurements) + + +def run_benchmark( + batch_sizes: List[int], hidden_sizes: List[int] +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Run benchmarks for all combinations of batch_size and hidden_size. + + Returns: + torch_times: 2D array of torch softmax times (ms) + flashinfer_times: 2D array of flashinfer softmax times (ms) + speedups: 2D array of speedup ratios (torch_time / flashinfer_time) + """ + n_batch = len(batch_sizes) + n_hidden = len(hidden_sizes) + + torch_times = np.zeros((n_batch, n_hidden)) + flashinfer_times = np.zeros((n_batch, n_hidden)) + speedups = np.zeros((n_batch, n_hidden)) + + print("Running benchmarks...") + print("=" * 100) + print( + f"{'Batch Size':<12} {'Hidden Size':<12} {'Torch (ms)':<15} " + f"{'FlashInfer (ms)':<18} {'Speedup':<10} {'Bandwidth (GB/s)':<18}" + ) + print("=" * 100) + + for i, batch_size in enumerate(batch_sizes): + for j, hidden_size in enumerate(hidden_sizes): + # Generate random logits + torch.manual_seed(42) + logits = torch.randn( + batch_size, hidden_size, device="cuda", dtype=torch.float32 + ) + + # Benchmark torch softmax + torch_time_ms = benchmark_torch_softmax(logits) + torch_times[i, j] = torch_time_ms + + # Benchmark flashinfer softmax + flashinfer_time_ms = benchmark_flashinfer_softmax(logits) + flashinfer_times[i, j] = flashinfer_time_ms + + # Calculate speedup + speedup = torch_time_ms / flashinfer_time_ms + speedups[i, j] = speedup + + # Calculate effective bandwidth (read + write) + io_bytes = logits.numel() * logits.element_size() * 2 + bandwidth_gb_s = io_bytes * 1e-6 / flashinfer_time_ms + + print( + f"{batch_size:<12} {hidden_size:<12} {torch_time_ms:<15.4f} " + f"{flashinfer_time_ms:<18.4f} {speedup:<10.2f}x {bandwidth_gb_s:<18.2f}" + ) + + print("=" * 100) + return torch_times, flashinfer_times, speedups + + +def plot_heatmap( + speedups: np.ndarray, + batch_sizes: List[int], + hidden_sizes: List[int], + save_path: str = "softmax_speedup_heatmap.png", +): + """Create and save a heatmap of speedup values.""" + # Create figure + fig, ax = plt.subplots(figsize=(12, 8)) + + # Create heatmap + sns.heatmap( + speedups, + annot=True, + fmt=".2f", + cmap="RdYlGn", + center=1.0, + cbar_kws={"label": "Speedup (x)"}, + xticklabels=[f"{h // 1000}K" for h in hidden_sizes], + yticklabels=batch_sizes, + ax=ax, + vmin=0.5, # Adjust color scale + vmax=max(3.0, speedups.max()), # Dynamic upper bound + ) + + ax.set_xlabel("Hidden Size", fontsize=12, fontweight="bold") + ax.set_ylabel("Batch Size", fontsize=12, fontweight="bold") + ax.set_title( + "FlashInfer Softmax Speedup vs PyTorch (Higher is Better)", + fontsize=14, + fontweight="bold", + pad=20, + ) + + plt.tight_layout() + plt.savefig(save_path, dpi=300, bbox_inches="tight") + print(f"\nHeatmap saved to: {save_path}") + + # Also create a performance comparison plot + _, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6)) + + # Plot 2: Speedup trends across batch sizes + for j, hidden_size in enumerate(hidden_sizes): + ax2.plot( + batch_sizes, + speedups[:, j], + marker="o", + label=f"Hidden={hidden_size // 1000}K", + linewidth=2, + ) + + ax2.set_xlabel("Batch Size", fontsize=12, fontweight="bold") + ax2.set_ylabel("Speedup (x)", fontsize=12, fontweight="bold") + ax2.set_title("Speedup vs Batch Size", fontsize=13, fontweight="bold") + ax2.set_xscale("log", base=2) + ax2.grid(True, alpha=0.3) + ax2.legend(fontsize=9) + ax2.axhline(y=1.0, color="red", linestyle="--", alpha=0.5, label="No speedup") + + # Plot 1: Speedup trends across hidden sizes + for i, batch_size in enumerate(batch_sizes[::2]): # Sample every other batch size + idx = i * 2 + ax1.plot( + [h // 1000 for h in hidden_sizes], + speedups[idx, :], + marker="s", + label=f"Batch={batch_size}", + linewidth=2, + ) + + ax1.set_xlabel("Hidden Size (K)", fontsize=12, fontweight="bold") + ax1.set_ylabel("Speedup (x)", fontsize=12, fontweight="bold") + ax1.set_title("Speedup vs Hidden Size", fontsize=13, fontweight="bold") + ax1.grid(True, alpha=0.3) + ax1.legend(fontsize=9) + ax1.axhline(y=1.0, color="red", linestyle="--", alpha=0.5) + + plt.tight_layout() + comparison_path = save_path.replace(".png", "_trends.png") + plt.savefig(comparison_path, dpi=300, bbox_inches="tight") + print(f"Trend plots saved to: {comparison_path}") + + +def main(): + """Main benchmark execution.""" + # Configuration + batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + hidden_sizes = [32000, 64000, 128000, 256000] + + print("=" * 100) + print("FlashInfer vs PyTorch Softmax Benchmark") + print("=" * 100) + print(f"Batch sizes: {batch_sizes}") + print(f"Hidden sizes: {hidden_sizes}") + print(f"Device: {torch.cuda.get_device_name()}") + print("=" * 100) + print() + + # Run benchmarks + _, _, speedups = run_benchmark(batch_sizes, hidden_sizes) + + # Print summary statistics + print("\nSummary Statistics:") + print("=" * 100) + print(f"Average speedup: {np.mean(speedups):.2f}x") + print(f"Median speedup: {np.median(speedups):.2f}x") + print(f"Min speedup: {np.min(speedups):.2f}x") + print(f"Max speedup: {np.max(speedups):.2f}x") + print("=" * 100) + + # Generate heatmap + plot_heatmap(speedups, batch_sizes, hidden_sizes) + + print("\nBenchmark complete!") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_tgv_gemm.py b/benchmarks/bench_tgv_gemm.py new file mode 100755 index 0000000000..b5574aa71c --- /dev/null +++ b/benchmarks/bench_tgv_gemm.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +""" +Test bench for tgv_gemm_bf16_sm100() function. +Tests the autotuner integration with TGV BF16 GEMM kernels. +""" + +import time +import csv +import torch +import torch.nn.functional as F + +from flashinfer import tgv_gemm_sm100, autotune + + +def test_tgv_gemm_bf16_sm100_perf(): + """Test tgv_gemm_bf16_sm100 with different problem sizes.""" + print("\n=== Testing tgv_gemm_bf16_sm100 with different sizes ===") + + # Test different problem sizes + test_cases = [ + (1, 7168, 2048, False, "deepseekv3, o_proj, tp=8"), + (4, 7168, 2048, False, "deepseekv3, o_proj, tp=8"), + (8, 7168, 2048, False, "deepseekv3, o_proj, tp=8"), + (16, 7168, 2048, False, "deepseekv3, o_proj, tp=8"), + (32, 7168, 2048, False, "deepseekv3, o_proj, tp=8"), + (64, 7168, 2048, False, "deepseekv3, o_proj, tp=8"), + (1, 3072, 1536, False, "deepseekv3, q_b_proj, tp=8"), + (4, 3072, 1536, False, "deepseekv3, q_b_proj, tp=8"), + (8, 3072, 1536, False, "deepseekv3, q_b_proj, tp=8"), + (16, 3072, 1536, False, "deepseekv3, q_b_proj, tp=8"), + (32, 3072, 1536, False, "deepseekv3, q_b_proj, tp=8"), + (64, 3072, 1536, False, "deepseekv3, q_b_proj, tp=8"), + (1, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (4, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (8, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (16, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (32, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (64, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (128, 1280, 2880, True, "gpt-oss-120b, qkv_proj, tp=4"), + (1, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + (4, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + (8, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + (16, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + (32, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + (64, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + (128, 2880, 1024, True, "gpt-oss-120b, o_proj, tp=4"), + ] + + # Prepare CSV output + csv_filename = "bf16_tgv_gemm_benchmark_results.csv" + csv_headers = [ + "M", + "N", + "K", + "has_bias", + "description", + "cublas_time_ms", + "tgv_time_ms", + "pdl_time_ms", + "tgv_speedup", + "pdl_speedup", + ] + + results = [] + + for m, n, k, has_bias, description in test_cases: + print(f"\n--- {description}: M={m}, N={n}, K={k}, has_bias={has_bias} ---") + flops = m * n * k * 2 / 1e12 + # Create tensors + A = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + B = torch.randn(n, k, device="cuda", dtype=torch.bfloat16).t() + if has_bias: + bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) + else: + bias = None + + # Warmup cublas + for _ in range(3): + _ = F.linear(A, B.T, bias) + + torch.cuda.synchronize() + + cublas_graph = torch.cuda.CUDAGraph() + + # Start graph capture + with torch.cuda.graph(cublas_graph): + for _ in range(100): + _ = F.linear(A, B.T, bias) + + # Warmup the graph + for _ in range(3): + cublas_graph.replay() + + torch.cuda.synchronize() + + # Benchmark using CUDA graph + start_time = time.time() + cublas_graph.replay() + torch.cuda.synchronize() + end_time = time.time() + cublas_avg_time = (end_time - start_time) / 100 + print( + f"CUBLAS average time: {cublas_avg_time * 1000:.6f} ms, {flops / cublas_avg_time:.3f} TFLOPS" + ) + + # Warmup + with autotune(tune_mode=True): + for _ in range(3): + _ = tgv_gemm_sm100(A, B, bias) + + torch.cuda.synchronize() + + tgv_graph = torch.cuda.CUDAGraph() + + # Start graph capture + with torch.cuda.graph(tgv_graph): + for _ in range(100): + _ = tgv_gemm_sm100(A, B, bias) + + # Warmup the graph + tgv_graph.replay() + + torch.cuda.synchronize() + + # Benchmark using CUDA graph + start_time = time.time() + tgv_graph.replay() + torch.cuda.synchronize() + end_time = time.time() + + tgv_avg_time = (end_time - start_time) / 100 + print( + f"TGV average time: {tgv_avg_time * 1000:.6f} ms, {flops / tgv_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / tgv_avg_time:.2f}x" + ) + + # Test with PDL + print("\nTesting with PDL...") + pdl_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(pdl_graph): + for _ in range(100): + _ = tgv_gemm_sm100(A, B, bias, pdl=True) + + # Warmup the graph + pdl_graph.replay() + + torch.cuda.synchronize() + + # Benchmark using CUDA graph + start_time = time.time() + pdl_graph.replay() + torch.cuda.synchronize() + end_time = time.time() + + pdl_avg_time = (end_time - start_time) / 100 + print( + f"PDL average time: {pdl_avg_time * 1000:.6f} ms, {flops / pdl_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / pdl_avg_time:.2f}x" + ) + + # Store results for CSV + results.append( + { + "M": m, + "N": n, + "K": k, + "has_bias": has_bias, + "description": description, + "cublas_time_ms": cublas_avg_time * 1000, + "tgv_time_ms": tgv_avg_time * 1000, + "pdl_time_ms": pdl_avg_time * 1000, + "tgv_speedup": cublas_avg_time / tgv_avg_time, + "pdl_speedup": cublas_avg_time / pdl_avg_time, + } + ) + + # Write results to CSV + print(f"\n=== Writing results to {csv_filename} ===") + with open(csv_filename, "w", newline="") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=csv_headers) + writer.writeheader() + writer.writerows(results) + + print(f"Benchmark results saved to {csv_filename}") + print(f"Total test cases: {len(results)}") + + +def test_tgv_gemm_bf16_sm100_correctness(): + """Test correctness of tgv_gemm_bf16_sm100 against reference implementation.""" + print("\n=== Testing correctness ===") + + # Create tensors + m, n, k = 64, 2048, 1024 + A = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + B = torch.randn(n, k, device="cuda", dtype=torch.bfloat16).t() + bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) + + # Reference computation + reference = torch.matmul(A, B) + bias.unsqueeze(0) + + # Test with TGV runner + out = tgv_gemm_sm100(A, B, bias) + + # Check correctness + cos_sim = F.cosine_similarity(reference.reshape(-1), out.reshape(-1), dim=0) + max_diff = torch.max(torch.abs(reference - out)).item() + mean_diff = torch.mean(torch.abs(reference - out)).item() + + print(f"Cosine similarity: {cos_sim:.6f}") + print(f"Max difference: {max_diff:.6f}") + print(f"Mean difference: {mean_diff:.6f}") + + # Check if results are close enough + if cos_sim > 0.99: + print("✓ Correctness test PASSED") + else: + print("✗ Correctness test FAILED") + + +def main(): + """Run all tests for tgv_gemm_bf16_sm100.""" + print("Starting BF16 TGV GEMM SM100 Tests") + print("=" * 50) + + try: + # Run correctness test first + test_tgv_gemm_bf16_sm100_correctness() + + # Test different problem sizes + test_tgv_gemm_bf16_sm100_perf() + + print("\n" + "=" * 50) + print("All BF16 TGV GEMM SM100 tests completed successfully!") + + except Exception as e: + print(f"\nError during testing: {e}") + import traceback + + traceback.print_exc() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_trtllm_fmha.py b/benchmarks/bench_trtllm_fmha.py new file mode 100644 index 0000000000..1615503b6e --- /dev/null +++ b/benchmarks/bench_trtllm_fmha.py @@ -0,0 +1,232 @@ +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import bench_gpu_time, bench_gpu_time_with_cudagraph + +page_size = 16 +num_kv_heads = 4 +num_qo_heads = 32 +head_dim = 128 + +workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.uint8, device="cuda:0") + + +def bench_trtllm_fmha(batch_size, seq_len, kv_cache_dtype): + torch.manual_seed(42) + seq_lens = torch.full((batch_size,), seq_len, device="cuda:0", dtype=torch.int32) + seq_lens_blocks = torch.ceil(seq_lens / page_size).int() + kv_indptr = torch.zeros(batch_size + 1, dtype=torch.int, device="cuda:0") + kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) + last_page_len = (seq_lens - (seq_lens_blocks - 1) * page_size).int() + last_page_len[last_page_len == 0] = page_size + num_blocks = kv_indptr[-1].item() + kv_indices = torch.arange(num_blocks, dtype=torch.int32, device="cuda:0") + + q = torch.rand(batch_size, num_qo_heads, head_dim, device="cuda:0").to( + torch.bfloat16 + ) + kv_data = torch.randn( + num_blocks, 2, num_kv_heads, page_size, head_dim, device="cuda:0" + ).to(torch.float8_e4m3fn if kv_cache_dtype == "fp8" else torch.float16) + + wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, "HND", backend="trtllm-gen" + ) + wrapper.plan( + kv_indptr, + kv_indices, + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + pos_encoding_mode="NONE", + q_data_type=q.dtype, + kv_data_type=kv_data.dtype, + ) + # add one warmup here + wrapper.run(q, kv_data) + torch.cuda.synchronize() + + measurements = bench_gpu_time(lambda: wrapper.run(q, kv_data)) + ms = np.median(measurements) + io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + print( + f"batch_size={batch_size}, seq_len={seq_len}, num_qo_heads={num_qo_heads}, num_kv_heads={num_kv_heads}, head_dim={head_dim}, page_size={page_size}" + ) + print(f"execution time: {ms}ms") + print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s") + + +def to_float8(x, dtype=torch.float8_e4m3fn): + finfo = torch.finfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + scale = finfo.max / amax * 0.1 + x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) + return x_scl_sat.to(dtype), scale.float().reciprocal() + + +def bench_trtllm_fmha_wrapper( + kv_layout, + batch_size, + max_seq_len, + page_size, + num_kv_heads, + head_dim, + q_dtype, + head_grp_size, + kv_cache_dtype, + window_left, + bench_with_sink, +): + torch.manual_seed(42) + device = "cuda:0" + num_qo_heads = num_kv_heads * head_grp_size + batch_size = batch_size + + # Initialize tensors + num_tokens = max_seq_len * batch_size + num_blocks = (num_tokens + page_size - 1) // page_size + + dtype_map = { + "half": torch.float16, + "bf16": torch.bfloat16, + "fp8": torch.float8_e4m3fn, + } + + q = torch.randn(batch_size, num_qo_heads, head_dim, device=device).to( + dtype_map[q_dtype] + ) + + # Sequence lengths and block tables + seq_lens = torch.full((batch_size,), max_seq_len) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) + blocks_per_seq = [(seq_len + page_size - 1) // page_size for seq_len in seq_lens] + + # Generate random but unique block IDs for all sequences + total_blocks_needed = sum(blocks_per_seq) + all_block_ids = torch.randperm( + total_blocks_needed, device=device + ) # Random permutation + + kv_cache_shape = (num_blocks, 2, num_kv_heads, page_size, head_dim) + kv_cache = torch.randn(size=kv_cache_shape).to(q.dtype) + + if kv_cache_dtype.startswith("fp8") and q_dtype != "fp8": + kv_cache, _ = to_float8(kv_cache) + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + blocks_per_seq = (seq_lens_tensor + page_size - 1) // page_size + + sinks = ( + torch.randn(num_qo_heads, device=device, dtype=torch.float32) + if bench_with_sink + else None + ) + + # Compute kv_indptr as cumulative sum of blocks per sequence + kv_indptr = ( + torch.cat( + [torch.tensor([0], device=device), torch.cumsum(blocks_per_seq, dim=0)] + ) + .int() + .to(device) + ) + + kv_indices = all_block_ids.int() + + # Calculate last page lengths + kv_last_page_len = seq_lens_tensor % page_size + kv_last_page_len[kv_last_page_len == 0] = page_size + + # trtllm-gen + wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, "HND", backend="trtllm-gen" + ) + wrapper.plan( + kv_indptr, + kv_indices, + kv_last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + pos_encoding_mode="NONE", + data_type=kv_cache.dtype, + q_data_type=q.dtype, + window_left=window_left, + ) + + # add one warmup here + wrapper.run(q, kv_cache, sinks=sinks) + torch.cuda.synchronize() + + measurements = bench_gpu_time_with_cudagraph( + lambda: wrapper.run(q, kv_cache, sinks=sinks), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + ms = np.median(measurements) + io = q.numel() * q.element_size() + kv_cache.numel() * kv_cache.element_size() + print( + f"batch_size={batch_size}, seq_len={max_seq_len}, num_qo_heads={num_qo_heads}, num_kv_heads={num_kv_heads}, head_dim={head_dim}, page_size={page_size}" + ) + print(f"execution time: {ms}ms") + print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Benchmark TRTLLM FMHA") + parser.add_argument( + "--head_dim", type=int, default=64, help="Dimension of each head" + ) + parser.add_argument( + "--num_kv_heads", type=int, default=8, help="Number of key/value heads" + ) + parser.add_argument( + "--page_size", type=int, default=16, help="Size of each page [16, 32, 64]" + ) + parser.add_argument( + "--head_grp_size", + type=int, + default=8, + help="Number of query heads per key-value head (group size)", + ) + parser.add_argument("--sink", action="store_true", help="Whether to test with sink") + parser.add_argument( + "--batch_sizes", + type=int, + nargs="+", + default=[4, 128, 256], + help="List of batch sizes to test", + ) + parser.add_argument( + "--seq_lens", + type=int, + nargs="+", + default=[1024, 4096, 8192, 16384], + help="List of sequence lengths to test", + ) + + args = parser.parse_args() + + for batch_size in args.batch_sizes: + for seq_len in args.seq_lens: + bench_trtllm_fmha_wrapper( + kv_layout="HND", + batch_size=batch_size, + max_seq_len=seq_len, + page_size=args.page_size, + num_kv_heads=args.num_kv_heads, + head_dim=args.head_dim, + q_dtype="bf16", + head_grp_size=args.head_grp_size, + kv_cache_dtype="auto", + window_left=-1, + bench_with_sink=args.sink, + ) diff --git a/benchmarks/bench_trtllm_gen_fused_moe_autotuner.py b/benchmarks/bench_trtllm_gen_fused_moe_autotuner.py new file mode 100644 index 0000000000..0aff25860e --- /dev/null +++ b/benchmarks/bench_trtllm_gen_fused_moe_autotuner.py @@ -0,0 +1,378 @@ +import argparse +from typing import Optional, Literal +import torch +import numpy as np +from flashinfer import ( + RoutingMethodType, + GatedActType, + fp4_quantize, + mxfp8_quantize, +) +from flashinfer.fused_moe import ( + trtllm_fp4_block_scale_moe, + trtllm_fp8_per_tensor_scale_moe, + trtllm_fp8_block_scale_moe, + WeightLayout, +) +from flashinfer.autotuner import autotune +from flashinfer.testing.utils import bench_gpu_time +from flashinfer.utils import device_support_pdl + +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max +FLOAT4_E2M1_MAX = 6.0 + + +def fp8_quantize(x): + max = x.abs().max().float() + scale = FLOAT8_E4M3_MAX / max + x = (x * scale).to(torch.float8_e4m3fn) + return x, 1.0 / scale + + +def bench_trtllm_gen_fused_moe_autotuner_fp8( + tune_max_num_tokens: Optional[int], + quant_mode: Literal["Fp8-Per-Tensor", "Fp8-Block"], + num_tokens: int, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + warmups: int, + iterations: int, +): + device = torch.device("cuda:0") + enable_pdl = device_support_pdl(device) + routing_logits = torch.rand(num_tokens, num_experts, device=device).to( + torch.float32 + ) + hidden_states = torch.randn(num_tokens, hidden_size, device=device).to( + torch.bfloat16 + ) + routing_bias = torch.randn(num_experts, device="cuda", dtype=torch.bfloat16) + w13 = torch.randn( + num_experts, intermediate_size * 2, hidden_size, device=device + ).to(torch.bfloat16) + w2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to( + torch.bfloat16 + ) + + is_block_scale = quant_mode == "Fp8-Block" + if not is_block_scale: + hidden_states, hidden_states_scale = fp8_quantize(hidden_states) + w13, w13_scale = fp8_quantize(w13) + w2, w2_scale = fp8_quantize(w2) + else: + # block scale quantization is too slow, so we use per-tensor quantization for now + hidden_states, hidden_states_scale = fp8_quantize(hidden_states) + w13, w13_scale = fp8_quantize(w13) + w2, w2_scale = fp8_quantize(w2) + hidden_states_scale = torch.full( + (hidden_size // 128, num_tokens), hidden_states_scale.item(), device=device + ) + w13_scale = torch.full( + (num_experts, intermediate_size * 2 // 128, hidden_size // 128), + w13_scale.item(), + device=device, + ) + w2_scale = torch.full( + (num_experts, hidden_size // 128, intermediate_size // 128), + w2_scale.item(), + device=device, + ) + + output1_scale_scalar = ( + torch.tensor([hidden_states_scale * w13_scale] * num_experts, device=device) + if not is_block_scale + else None + ) + output1_scales_gate_scalar = ( + torch.ones(num_experts, device=device, dtype=torch.float32) + if not is_block_scale + else None + ) + output2_scale_scalar = ( + torch.tensor([hidden_states_scale * w2_scale] * num_experts, device=device) + if not is_block_scale + else None + ) + + if is_block_scale: + fn = lambda: trtllm_fp8_block_scale_moe( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + w13, + w13_scale, + w2, + w2_scale, + num_experts, + top_k, + 8, # n_group + 4, # topk_group + intermediate_size, + 0, # local_expert_offset + num_experts, + 2.5, # routed_scaling_factor + None, # tile_tokens_dim + RoutingMethodType.DeepSeekV3.value, + True, # use_shuffled_weight + WeightLayout.BlockMajorK.value, # weight_layout + enable_pdl=enable_pdl, + tune_max_num_tokens=num_tokens + if tune_max_num_tokens is None + else tune_max_num_tokens, + ) + else: + fn = lambda: trtllm_fp8_per_tensor_scale_moe( + routing_logits, + None, # routing_bias + hidden_states, + w13, + output1_scale_scalar, + output1_scales_gate_scalar, + w2, + output2_scale_scalar, + num_experts, + top_k, + None, # n_group + None, # topk_group + intermediate_size, + 0, # local_expert_offset + num_experts, + 1.0, # routed_scaling_factor + False, # use_routing_scales_on_input + None, # tile_tokens_dim + RoutingMethodType.TopK.value, + enable_pdl, + num_tokens if tune_max_num_tokens is None else tune_max_num_tokens, + ) + + def bench(do_autotune): + with autotune(do_autotune): + fn() + ms_list = bench_gpu_time( + fn, + dry_run_iters=warmups, + repeat_iters=iterations, + ) + median_ms = np.median(ms_list) + return median_ms + + ms = bench(do_autotune=False) + ms_tuned = bench(do_autotune=True) + print( + f"num tokens: {num_tokens}, num experts: {num_experts}, hidden size: {hidden_size}, intermediate size: {intermediate_size}, top k: {top_k}" + ) + print(f"No autotune: {ms:.3f} ms; with autotune: {ms_tuned:.3f} ms") + + +def bench_trtllm_gen_fused_moe_autotuner_fp4( + tune_max_num_tokens: Optional[int], + quant_mode: Literal["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16"], + num_tokens: int, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + warmups: int, + iterations: int, +): + device = torch.device("cuda:0") + enable_pdl = device_support_pdl(device) + routing_logits = torch.rand(num_tokens, num_experts, device=device).to( + torch.bfloat16 + ) + hidden_states = torch.randn(num_tokens, hidden_size, device=device).to( + torch.bfloat16 + ) + if quant_mode == "NvFP4xNvFP4": + hidden_states, hidden_states_scale = fp4_quantize( + hidden_states, + torch.tensor([448.0 * 6.0], device=device), + sf_vec_size=16, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + ) + hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape( + num_tokens, -1 + ) + hidden_states_global_scale = 1.0 / 448.0 / 6.0 + elif quant_mode == "MxFP4xMxFP8": + hidden_states, hidden_states_scale = mxfp8_quantize(hidden_states, False) + hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape( + num_tokens, -1 + ) + hidden_states_global_scale = 1.0 + else: # MxFP4xBf16 + hidden_states_scale = None + hidden_states_global_scale = 1.0 + + w13 = torch.randn( + num_experts, intermediate_size * 2, hidden_size, device=device + ).to(torch.bfloat16) + w2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to( + torch.bfloat16 + ) + if quant_mode == "NvFP4xNvFP4": + w13, w13_scale = fp4_quantize( + w13, + torch.tensor([448.0 * 6.0], device=device), + sf_vec_size=16, + sf_use_ue8m0=False, + ) + w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape( + num_experts, intermediate_size * 2, -1 + ) + w2, w2_scale = fp4_quantize( + w2, + torch.tensor([448.0 * 6.0], device=device), + sf_vec_size=16, + sf_use_ue8m0=False, + ) + w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, -1 + ) + w13_global_scale = 1.0 / 448.0 / 6.0 + w2_global_scale = 1.0 / 448.0 / 6.0 + else: + w13, w13_scale = fp4_quantize( + w13, torch.tensor([1.0], device=device), sf_vec_size=32, sf_use_ue8m0=True + ) + w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape( + num_experts, intermediate_size * 2, -1 + ) + w2, w2_scale = fp4_quantize( + w2, torch.tensor([1.0], device=device), sf_vec_size=32, sf_use_ue8m0=True + ) + w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, -1 + ) + w13_global_scale = 1.0 + w2_global_scale = 1.0 + bias13 = torch.randn(num_experts, intermediate_size * 2, device=device) * 10 + bias2 = torch.randn(num_experts, intermediate_size * 2, device=device) * 10 + + output1_scale_scalar = torch.tensor( + [hidden_states_global_scale * w13_global_scale] * num_experts, device=device + ) + output1_scale_gate_scalar = torch.tensor( + [hidden_states_global_scale * w13_global_scale] * num_experts, device=device + ) + output2_scale_scalar = torch.tensor( + [hidden_states_global_scale * w2_global_scale] * num_experts, device=device + ) + fn = lambda: trtllm_fp4_block_scale_moe( + routing_logits, + None, # routing_bias + hidden_states, + hidden_states_scale, + w13, + w13_scale, + bias13, + None, # gemm1_alpha + None, # gemm1_beta + None, # gemm1_clamp_limit + w2, + w2_scale, + bias2, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + None, # n_group + None, # topk_group + intermediate_size, + 0, # local_expert_offset + num_experts, + None, # routed_scaling_factor + None, # tile_tokens_dim + RoutingMethodType.Renormalize.value, + True, + enable_pdl, + GatedActType.SwiGlu.value, # gated_act_type + None, + num_tokens if tune_max_num_tokens is None else tune_max_num_tokens, + ) + + def bench(do_autotune): + with autotune(do_autotune): + fn() + ms_list = bench_gpu_time( + fn, + dry_run_iters=warmups, + repeat_iters=iterations, + ) + median_ms = np.median(ms_list) + return median_ms + + ms = bench(do_autotune=False) + ms_tuned = bench(do_autotune=True) + print( + f"num tokens: {num_tokens}, num experts: {num_experts}, hidden size: {hidden_size}, intermediate size: {intermediate_size}, top k: {top_k}" + ) + print(f"No autotune: {ms:.3f} ms; with autotune: {ms_tuned:.3f} ms") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--quant-mode", + type=str, + default="MxFP4xMxFP8", + choices=[ + "NvFP4xNvFP4", + "MxFP4xMxFP8", + "MxFP4xBf16", + "Fp8-Per-Tensor", + "Fp8-Block", + ], + help="Quantization mode", + ) + parser.add_argument("--num-tokens", type=int, default=512, help="Number of tokens") + parser.add_argument( + "--tune-max-num-tokens", + type=int, + default=None, + help="Maximum number of tokens for tunning", + ) + parser.add_argument( + "--num-experts", type=int, default=128, help="Number of experts" + ) + parser.add_argument("--hidden-size", type=int, default=3072, help="Hidden size") + parser.add_argument( + "--intermediate-size", type=int, default=3072, help="Intermediate size" + ) + parser.add_argument("--top-k", type=int, default=4, help="Top-k experts per token") + parser.add_argument( + "--warmups", type=int, default=100, help="Number of warmup iterations" + ) + parser.add_argument( + "--iterations", type=int, default=100, help="Number of benchmark iterations" + ) + args = parser.parse_args() + if args.quant_mode in ["Fp8-Per-Tensor", "Fp8-Block"]: + bench_trtllm_gen_fused_moe_autotuner_fp8( + args.tune_max_num_tokens, + args.quant_mode, + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.warmups, + args.iterations, + ) + else: + bench_trtllm_gen_fused_moe_autotuner_fp4( + args.tune_max_num_tokens, + args.quant_mode, + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.warmups, + args.iterations, + ) diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py new file mode 100644 index 0000000000..b9ac3967dd --- /dev/null +++ b/benchmarks/bench_trtllm_gen_mla.py @@ -0,0 +1,132 @@ +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import bench_gpu_time_with_cudagraph + +num_q_heads = 128 +num_kv_heads = 1 +qk_nope_head_dim = 128 +qk_rope_head_dim = 64 +kv_lora_rank = 512 + + +def bench_trtllm_mla(batch_size, q_len_per_request, seq_len, page_size, dtype): + torch.manual_seed(42) + device = "cuda:0" + + # Initialize tensors + query = torch.randn( + batch_size, + q_len_per_request, + num_q_heads, + kv_lora_rank + qk_rope_head_dim, + device=device, + ).to(dtype) + + num_tokens = seq_len * batch_size + num_blocks = (num_tokens + page_size - 1) // page_size + + # Sequence lengths and block tables + seq_lens = [torch.randint(1, seq_len, (1,)).item() for _ in range(batch_size)] + seq_lens[-1] = seq_len + max_seq_len = max(seq_lens) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) + + blocks_per_seq = (seq_lens_tensor + page_size - 1) // page_size + max_num_blocks_per_seq = blocks_per_seq.max().item() + + # Generate random but unique block IDs for all sequences + total_blocks_needed = sum(blocks_per_seq) + all_block_ids = torch.randperm( + total_blocks_needed, device=device + ) # Random permutation + + # Generate unique block IDs for all sequences + block_id = 0 + block_tables = torch.zeros( + (batch_size, max_num_blocks_per_seq), dtype=torch.int, device=device + ) + + # Populate block tables and track block assignments + block_id = 0 + for i in range(batch_size): + num_blocks_needed = blocks_per_seq[i] + block_tables[i, :num_blocks_needed] = all_block_ids[ + block_id : block_id + num_blocks_needed + ] + block_id += num_blocks_needed + + # Create interleaved KV cache + # Allocate more than needed blocks, block_id is just enough, to mimick real-world cases + kv_cache = torch.randn( + size=(num_blocks, page_size, kv_lora_rank + qk_rope_head_dim), device=device + ).to(dtype) + # (num_blocks, 1, page_size, kv_lora_rank + qk_rope_head_dim) + + # Allocate workspace buffer + # todo(Yingyi): calculate the actual size of workspace buffer + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + # Run decode-MLA + # warmup + flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=kv_cache.unsqueeze(1), + workspace_buffer=workspace_buffer, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + block_tables=block_tables, + seq_lens=seq_lens_tensor, + max_seq_len=max_seq_len, + bmm1_scale=1.0 / ((128 + 64) ** 0.5), + bmm2_scale=1.0, + ) + # benchmark + measurements = bench_gpu_time_with_cudagraph( + lambda: flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=kv_cache.unsqueeze(1), + workspace_buffer=workspace_buffer, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + block_tables=block_tables, + seq_lens=seq_lens_tensor, + max_seq_len=max_seq_len, + bmm1_scale=1.0 / ((128 + 64) ** 0.5), + bmm2_scale=1.0, + ), + dry_run_time_ms=100, + repeat_time_ms=1000, + ) + io = ( + query.numel() * query.element_size() + + kv_cache.numel() * kv_cache.element_size() + ) + ms = np.median(measurements) + flops = ( + 2 + * num_q_heads + * (2 * kv_lora_rank + qk_rope_head_dim) + * sum(seq_lens) + * q_len_per_request + ) + print( + f"batch_size={batch_size}, q_len_per_request={q_len_per_request}, seq_len={seq_len}, num_q_heads={num_q_heads}, num_kv_heads={num_kv_heads}, qk_nope_head_dim={qk_nope_head_dim}, qk_rope_head_dim={qk_rope_head_dim}, kv_lora_rank={kv_lora_rank}, page_size={page_size}" + ) + print(f"execution time: {ms} ms") + print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s") + print(f"FLOPs: {flops * 1e-9 / ms:.2f} TFLOPs/s") + + +if __name__ == "__main__": + for dtype in [torch.bfloat16, torch.float8_e4m3fn]: + for page_size in [32, 64]: + for batch_size in [1, 2, 4, 16, 32, 64, 128, 256, 512, 768, 1024]: + for seq_len in [1024, 4096, 8192]: + for q_len_per_request in [1, 2, 4, 8, 16]: + bench_trtllm_mla( + batch_size, q_len_per_request, seq_len, page_size, dtype + ) diff --git a/benchmarks/flashinfer_benchmark.py b/benchmarks/flashinfer_benchmark.py new file mode 100644 index 0000000000..bd02172eb2 --- /dev/null +++ b/benchmarks/flashinfer_benchmark.py @@ -0,0 +1,207 @@ +import argparse +import sys + +from routines.attention import parse_attention_args, run_attention_test +from routines.flashinfer_benchmark_utils import ( + benchmark_apis, + full_output_columns, + output_column_dict, +) +from routines.gemm import parse_gemm_args, run_gemm_test +from routines.moe import parse_moe_args, run_moe_test + + +def run_test(args): + """ + Route & run a single FlashInfer test case with test routine. + + Args: + args: Parsed command line arguments containing test configuration + """ + + ## Depending on routine type, route to corresponding test routine + if args.routine in benchmark_apis["attention"]: + res = run_attention_test(args) + elif args.routine in benchmark_apis["gemm"]: + res = run_gemm_test(args) + elif args.routine in benchmark_apis["moe"]: + res = run_moe_test(args) + else: + raise ValueError(f"Unsupported routine: {args.routine}") + + # Write results to output file if specified + if args.output_path is not None: + with open(args.output_path, "a") as fout: + for cur_res in res: + for key in output_column_dict["general"]: + cur_res[key] = getattr(args, key) + + output_line = ",".join( + [str(cur_res[col]) for col in full_output_columns] + ) + fout.write(output_line + "\n") + fout.flush() + return + + +def parse_args(line=sys.argv[1:]): + """ + Parse command line arguments for test configuration. + First parse shared arguments, then parse routine-specific arguments. + + Args: + line: Command line arguments (default: sys.argv[1:]) + + Returns: + Parsed argument namespace + """ + + ## Shared arguments + parser = argparse.ArgumentParser() + parser.add_argument( + "--routine", + "-R", + type=str, + required=True, + choices=list(benchmark_apis["attention"]) + + list(benchmark_apis["gemm"]) + + list(benchmark_apis["moe"]), + ) + args, _ = parser.parse_known_args(line[:]) + + parser.add_argument( + "--no_cuda_graph", + action="store_true", + default=False, + help="Disable CUDA graph to execute kernels outside of the graph.", + ) + parser.add_argument( + "--use_cupti", + action="store_true", + default=False, + help="Use CUPTI for timing GPU kernels when available.", + ) + parser.add_argument( + "--refcheck", + action="store_true", + default=False, + help="Run reference check that ensures outputs correct.", + ) + parser.add_argument( + "--allow_output_mismatch", + action="store_true", + default=False, + help="Allow output mismatch between backends during reference checks. Error message will be printed but test will continue.", + ) + parser.add_argument( + "--random_seed", type=int, default=42, help="Random seed for reproducibility." + ) + parser.add_argument( + "--verbose", "-v", action="count", help="Set verbosity level.", default=0 + ) + parser.add_argument( + "--output_path", + type=str, + required=False, + default=None, + help="Output path for results. If not specified, results will not be written to a file.", + ) + parser.add_argument( + "--num_iters", + "-n", + type=int, + required=False, + default=30, + help="Number of iterations to run for measurement.", + ) + parser.add_argument( + "--dry_run_iters", + "-d", + type=int, + required=False, + default=5, + help="Number of dry runs.", + ) + parser.add_argument( + "--case_tag", + type=str, + required=False, + default=None, + help="Optional tag for the test case for annotating output.", + ) + parser.add_argument( + "--generate_repro_command", + action="store_true", + default=False, + help="If set, will print reproducer command and store it to output csv.", + ) + parser.add_argument( + "--repro_command", + type=str, + required=False, + default="", + help="Placeholder for generated reproducer command for the test case. Not to be used directly.", + ) + + ## Check routine and pass on to routine-specific argument parser + if args.routine in benchmark_apis["attention"]: + args = parse_attention_args(line, parser) + elif args.routine in benchmark_apis["gemm"]: + args = parse_gemm_args(line, parser) + elif args.routine in benchmark_apis["moe"]: + args = parse_moe_args(line, parser) + else: + raise ValueError(f"Unsupported routine: {args.routine}") + + if args.generate_repro_command: + args.repro_command = "python3 flashinfer_benchmark.py " + " ".join(line) + return args + + +if __name__ == "__main__": + # Parse testlist argument first + testlist_parser = argparse.ArgumentParser(add_help=False) + testlist_parser.add_argument( + "--testlist", + type=str, + required=False, + default=None, + help="Optional testlist file to run multiple cases.", + ) + testlist_parser.add_argument( + "--output_path", + type=str, + required=False, + default=None, + help="Output path for results csv.", + ) + testlist_args, _ = testlist_parser.parse_known_args() + + # Setup output file if specified + if testlist_args.output_path is not None: + with open(testlist_args.output_path, "w") as fout: + fout.write(",".join(full_output_columns) + "\n") + + # Process tests either from testlist file or command line arguments + if testlist_args.testlist is not None: + # If testlist, run each test in the testlist + with open(testlist_args.testlist, "r") as f: + import shlex + + for line in f.readlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + try: + line_args = parse_args(shlex.split(line)) + line_args.output_path = testlist_args.output_path + run_test(line_args) + except Exception as e: + print(f"[ERROR] Error running test: {line}") + print(f"[ERROR] Error: {e}") + continue + else: + # If no testlist, just run the command + args = parse_args() + args.output_path = testlist_args.output_path + run_test(args) diff --git a/benchmarks/routines/__init__.py b/benchmarks/routines/__init__.py new file mode 100644 index 0000000000..a0cddcb0c0 --- /dev/null +++ b/benchmarks/routines/__init__.py @@ -0,0 +1,15 @@ +""" +Copyright (c) 2023 by FlashInfer team. + +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. +""" diff --git a/benchmarks/routines/attention.py b/benchmarks/routines/attention.py new file mode 100644 index 0000000000..e88b176f13 --- /dev/null +++ b/benchmarks/routines/attention.py @@ -0,0 +1,1961 @@ +from collections import defaultdict + +import numpy as np +import torch + +import flashinfer +from flashinfer.testing.utils import ( + attention_tb_per_sec_with_actual_seq_lens, + attention_tflops_per_sec_with_actual_seq_lens, + bench_gpu_time, +) + +from .flashinfer_benchmark_utils import ( + dtype_str_to_torch_dtype, + get_device, + print_perf_metrics, + is_close_stats, + filter_backends_by_compute_capability, +) + + +def normalize_backends(backends): + """ + Normalize backend names planned for deprecation and print warnings. + Currently: + - Replaces deprecated 'trtllm-gen-native' with 'trtllm-native'. + + Args: + backends: List of backend names + + Returns: + List of normalized backend names + """ + normalized = [] + for backend in backends: + if backend == "trtllm-gen-native": + print( + "[WARNING] Backend name 'trtllm-gen-native' has been renamed to 'trtllm-native' and will be removed in a future release. " + ) + normalized.append("trtllm-native") + else: + normalized.append(backend) + return normalized + + +def run_attention_test(args): + """ + Run an attention test. + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.routine == "BatchDecodeWithPagedKVCacheWrapper": + return testBatchDecodeWithPagedKVCacheWrapper(args) + elif args.routine == "BatchPrefillWithPagedKVCacheWrapper": + return testBatchPrefillWithPagedKVCacheWrapper(args) + elif args.routine == "BatchPrefillWithRaggedKVCacheWrapper": + return testBatchPrefillWithRaggedKVCacheWrapper(args) + elif args.routine == "BatchMLAPagedAttentionWrapper": + return testBatchMLAPagedAttentionWrapper(args) + else: + print(f"[ERROR] Unsupported routine: {args.routine}") + return [] + + +def parse_attention_args(line, parser): + """ + Parse command line arguments for attention test configuration. + + Args: + line: Command line arguments + parser: ArgumentParser object already populated with shared arguments + + Returns: + Parsed argument namespace + """ + parser.add_argument( + "--backends", + type=str, + required=False, + nargs="+", + default=["fa2"], + choices=[ + "fa2", + "fa2_tc", + "fa3", + "cudnn", + "cutlass", + "trtllm-gen", + "trtllm-native", + "trtllm-gen-native", # Deprecated, will be removed in future + ], + help="Kernel backends to test. Default: fa2", + ) + parser.add_argument( + "--page_size", + type=int, + required=False, + default=0, + help="Page size for paged attention. Required for paged attention. Ignored for non-paged attention.", + ) + parser.add_argument( + "--batch_size", type=int, required=True, help="Batch size of test case." + ) + parser.add_argument( + "--s_qo", + type=int, + required=False, + default=1, + help="Max sequence length of the query. Should be 1 for decode.", + ) + parser.add_argument( + "--s_kv", + type=int, + required=True, + help="Max sequence length of the key and value.", + ) + parser.add_argument( + "--num_qo_heads", type=int, required=True, help="Number of query heads." + ) + parser.add_argument( + "--num_kv_heads", type=int, required=True, help="Number of key and value heads." + ) + parser.add_argument( + "--head_dim_qk", + type=int, + required=False, + help="Head dimension of the query and key for prefill and decode MHA/GQA/MQA.", + ) + parser.add_argument( + "--head_dim_vo", + type=int, + required=False, + help="Head dimension of the value and output for prefill and decode MHA/GQA/MQ.", + ) + parser.add_argument( + "--head_dim_ckv", + type=int, + required=False, + help="Head dimension of compressed kv-cache tensor (without rope).", + ) + parser.add_argument( + "--head_dim_kpe", + type=int, + required=False, + help="Head dimension of the rope part of the kv-cache tensor.", + ) + parser.add_argument( + "--q_dtype", + type=str, + required=False, + default="bfloat16", + help="Data type of the query. Currently only bfloat16 is supported.", + ) + parser.add_argument( + "--kv_dtype", + type=str, + required=False, + default="bfloat16", + help="Data type of the key and value. Currently only bfloat16 is supported.", + ) + parser.add_argument( + "--causal", + action="store_true", + default=False, + help="Causal masking. Note: not padding masking. Only used for prefill tests.", + ) + parser.add_argument( + "--random_actual_seq_len", + action="store_true", + default=False, + help="Use random actual sequence lengths for the query and key and value. Random values are generated between 1 and maximum sequence length. If False, use maximum sequence length.", + ) + + args = parser.parse_args(line) + + # Normalize backend names (handle deprecated names) + args.backends = normalize_backends(args.backends) + + if args.verbose >= 1: + print(f"[INFO] {args = }") + return args + + +def sample_actual_seq_lens(max_seqlen, batch_size, device, random_actual_seq_len): + """ + Get an array of actual sequence lengths for given batch size and max sequence length. + If random_actual_seq_len is True, sample actual sequence lengths randomly. + Otherwise, set all actual sequence lengths to max_seqlen. + + Args: + max_seqlen: Maximum sequence length. + batch_size: Batch size. + device: Device to sample on. + random_actual_seq_len: Whether to sample actual sequence lengths randomly. + + Returns: + actual_seq_lens: Actual sequence lengths for each batch. + """ + if random_actual_seq_len: + actual_seq_lens = torch.randint( + 1, max_seqlen + 1, (batch_size, 1, 1, 1), device=device, dtype=torch.int32 + ) + else: + actual_seq_lens = torch.full( + (batch_size, 1, 1, 1), max_seqlen, device=device, dtype=torch.int32 + ) + return actual_seq_lens + + +def testBatchDecodeWithPagedKVCacheWrapper(args): + """ + Test BatchDecodeWithPagedKVCacheWrapper API and equivalent cuDNN API. + Supports fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native backends. + + This test: + 1. Creates paged KV cache and query tensors + 2. Runs decode attention with different backends + 3. Verifies outputs match between backends + 4. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testBatchDecodeWithPagedKVCacheWrapper") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + # Basic setup + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + q_init_dtype = torch.bfloat16 + kv_init_dtype = torch.bfloat16 + rtol = 2e-1 + atol = 1e-2 + res = [] + + # Handle different query data types. + q_dtype = dtype_str_to_torch_dtype(args.q_dtype) + if q_dtype not in [torch.bfloat16, torch.float8_e4m3fn]: + print(f"[ERROR] Unsupported q_dtype: {args.q_dtype}") + return res + + # Handle different KV cache data types. + kv_dtype = dtype_str_to_torch_dtype(args.kv_dtype) + if kv_dtype not in [torch.bfloat16, torch.float8_e4m3fn]: + print(f"[ERROR] Unsupported kv_dtype: {args.kv_dtype}") + return res + + # Parse and validate backend configurations + backends = args.backends + page_size = args.page_size + batch_size = args.batch_size + s_qo = args.s_qo + s_kv = args.s_kv + num_qo_heads = args.num_qo_heads + num_kv_heads = args.num_kv_heads + head_dim_qk = args.head_dim_qk + head_dim_vo = args.head_dim_vo + is_cuda_graph_compatible = not args.no_cuda_graph + # return_lse = not args.no_lse # TO-DO: Add support for this + run_refcheck = args.refcheck + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + # Check for backend-specific constraints + if "fa2" in backends: + remove_fa2 = False + head_grp_size = ( + num_qo_heads // num_kv_heads + ) # If 5, FA2 backend is not supported. + if head_grp_size == 5: + print( + "[INFO] FA2 backend is not supported for this configuration. Skipping." + ) + remove_fa2 = True + if remove_fa2: + backends.remove("fa2") + + if "fa2_tc" in backends: + remove_fa2_tc = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] FA2_TC backend does not support FP8. Skipping.") + remove_fa2_tc = True + if remove_fa2_tc: + backends.remove("fa2_tc") + + if "cudnn" in backends: + remove_cudnn = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] cuDNN backend does not support FP8. Skipping.") + remove_cudnn = True + if remove_cudnn: + backends.remove("cudnn") + + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + + # Sample sequence lengths and create tensors + actual_seq_lens_kv = sample_actual_seq_lens( + s_kv, batch_size, device, args.random_actual_seq_len + ) + sum_seq_kv = torch.sum(actual_seq_lens_kv).item() + avg_seq_len_kv = sum_seq_kv // batch_size + + if args.verbose >= 1: + print(f"[VERBOSE] Average actual seq len: {avg_seq_len_kv}") + if args.verbose >= 2: + print(f"[VVERBOSE] {actual_seq_lens_kv.flatten() = }") + + # Create query tensor + q = torch.rand( + batch_size, num_qo_heads, head_dim_qk, device=device, dtype=q_init_dtype + ) + if args.verbose >= 2: + print(f"[VVERBOSE] {q.shape = }") + + # Create KV cache + num_pages_per_seq = (s_kv + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + + if args.verbose >= 2: + print(f"[VVERBOSE] {num_pages_per_seq = }") + print(f"[VVERBOSE] {total_num_pages = }") + + # Initialize KV cache with appropriate shape and stride + kv_cache_shape = ( + total_num_pages, + 2, # 2 for key and value + num_kv_heads, + page_size, + head_dim_qk, + ) + kv_cache = torch.randn(size=kv_cache_shape, dtype=kv_init_dtype).to(device) + + # Keep a copy for TRT-LLM which uses different strides + if "trtllm-gen" in backends: + kv_cache_for_trt = kv_cache.detach().clone() + + kv_cache = kv_cache.as_strided( + kv_cache.shape, + ( + 2 * page_size * num_kv_heads * head_dim_qk, + page_size * num_kv_heads * head_dim_qk, + head_dim_qk, + num_kv_heads * head_dim_qk, + 1, + ), + ) + k_cache_view, v_cache_view = kv_cache[:, 0, :, :, :], kv_cache[:, 1, :, :, :] + + if "trtllm-gen" in backends: + # kv_cache now has different tensor stride and logical values. Copy over values to kv_cache_for_trt. + # Result is kv_cache and kv_cache_for_trt have the same logical values but different tensor strides. + kv_cache_for_trt.copy_(kv_cache) + + v_cache = v_cache_view.as_strided( + v_cache_view.shape, + ( + 2 * page_size * num_kv_heads * head_dim_qk, + head_dim_qk, + num_kv_heads * head_dim_qk, + 1, + ), + ) + k_cache = k_cache_view.as_strided( + k_cache_view.shape, + ( + 2 * page_size * num_kv_heads * head_dim_qk, + head_dim_qk, + num_kv_heads * head_dim_qk, + 1, + ), + ) + + # Now initialize the page tables + block_tables = torch.tensor( + [ + [k + i * num_pages_per_seq for k in range(num_pages_per_seq)] + for i in range(batch_size) + ], + dtype=torch.int, + device=device, + ) + + kv_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum( + (actual_seq_lens_kv.flatten() + page_size - 1) // page_size, dim=0 + ), + ] + ) + .int() + .to(device) + ) + + # kv_indices[-1] is the total number of actual pages + kv_indices = torch.zeros(kv_indptr[-1], device=device, dtype=torch.int32) + for i in range(len(kv_indptr) - 1): + start_idx = kv_indptr[i] + end_idx = kv_indptr[i + 1] + kv_indices[start_idx:end_idx] = torch.arange( + i * num_pages_per_seq, + i * num_pages_per_seq + (end_idx - start_idx), + device=device, + ) + + kv_last_page_len = ( + torch.where( + actual_seq_lens_kv.flatten() % page_size == 0, + torch.full((batch_size,), page_size, device=device), + actual_seq_lens_kv.flatten() % page_size, + ) + .int() + .to(device) + ) + + ragged_q = ( + torch.arange(0, batch_size + 1, device=device) * (num_qo_heads * head_dim_qk) + ).long() # For cuDNN + + scale = float(1.0 / (head_dim_qk**0.5)) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + if args.verbose >= 2: + print(f"[VVERBOSE] {kv_cache.shape = }") + print(f"[VVERBOSE] {kv_cache.stride() = }") + print(f"[VVERBOSE] {block_tables.shape = }") + print(f"[VVERBOSE] {kv_indptr.shape = }") + print(f"[VVERBOSE] {kv_indices.shape = }") + print(f"[VVERBOSE] {kv_last_page_len.shape = }") + print(f"[VVERBOSE] {scale = }") + + # Prepare wrappers + backend_wrappers = {} + for backend in backends: + if backend in ["fa2", "fa2_tc", "trtllm-gen"]: + plan_kv_indptr = ( + kv_indptr.clone().detach() if backend == "trtllm-gen" else kv_indptr + ) + backend_wrappers[backend] = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, + "HND", + use_cuda_graph=is_cuda_graph_compatible, + use_tensor_cores=(backend != "fa2"), + paged_kv_indptr_buffer=plan_kv_indptr, + paged_kv_indices_buffer=kv_indices, + paged_kv_last_page_len_buffer=kv_last_page_len, + backend=backend, + ) + backend_wrappers[backend].plan( + plan_kv_indptr, + kv_indices, + kv_last_page_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + page_size, + q_data_type=q_dtype, + data_type=kv_dtype, + block_tables=block_tables, + ) + + ## If FP8, prepare + k_scale, v_scale = None, None + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + q = q.to(q_dtype) + if kv_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + k_data, v_data = torch.chunk(kv_cache, 2, dim=1) + k_scale = k_data.amax().item() / 256 + v_scale = v_data.amax().item() / 256 + k_fp8 = (k_data / k_scale).to(kv_dtype) + v_fp8 = (v_data / v_scale).to(kv_dtype) + kv_cache = torch.cat([k_fp8, v_fp8], dim=1) + if "trtllm-gen" in backends: + k_data, v_data = torch.chunk(kv_cache_for_trt, 2, dim=1) + k_fp8 = (k_data / k_scale).to(kv_dtype) + v_fp8 = (v_data / v_scale).to(kv_dtype) + kv_cache_for_trt = torch.cat([k_fp8, v_fp8], dim=1) + + def run_backend_wrapper(backend): + if backend in ["fa2", "fa2_tc", "trtllm-gen"]: + return backend_wrappers[backend].run( + q, kv_cache, k_scale=k_scale, v_scale=v_scale + ) + elif backend == "cudnn": + return flashinfer.decode.cudnn_batch_decode_with_kv_cache( + q, + k_cache, + v_cache, + scale, + workspace_buffer, + max_sequence_kv=s_kv, + actual_seq_lens_kv=actual_seq_lens_kv, + block_tables=block_tables, + is_cuda_graph_compatible=is_cuda_graph_compatible, + batch_offsets_q=ragged_q, + batch_offsets_o=ragged_q, + ) + elif backend == "trtllm-native": + return flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q.contiguous(), + kv_cache=kv_cache, + workspace_buffer=workspace_buffer, + block_tables=block_tables, + seq_lens=actual_seq_lens_kv, + max_seq_len=s_kv, + bmm1_scale=scale if k_scale is None else k_scale * scale, + bmm2_scale=1.0 if v_scale is None else v_scale, + ) + else: + print(f"[ERROR] Backend {backend} not supported") + return res + + has_reference_output = False + # Iterate over each backend: + for cur_backend in backends: + # Clear workspace buffer to prevent unexpected interactions between backends. + workspace_buffer.zero_() + if run_refcheck: + outputs[cur_backend] = run_backend_wrapper(cur_backend).detach().clone() + if cur_backend == "fa2": + has_reference_output = True + reference_output = outputs[cur_backend] + # Unified benchmark entry: prefer graph if compatible and not using CUPTI + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend_wrapper(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=(is_cuda_graph_compatible and cur_backend != "fa2"), + ) + + # Perform reference check + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 1: + if run_refcheck and has_reference_output: + if reference_output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + if args.verbose >= 2: + print( + "[VVERBOSE] Reference output is FP8. Converting to float32 for reference check." + ) + reference_output = reference_output.to(torch.float32) + tested_outputs = [output.to(torch.float32) for output in tested_outputs] + for i in range(len(tested_outputs)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats(reference_output, tested_outputs[i], rtol, atol) + if num_different_elements > 0: + print( + f"[ERROR] Output tensor mismatch between backends fa2 and {tested_backends[i]}: " + f"{num_different_elements} / {num_elements} ({num_different_elements_percentage:.2f}%) elements are different" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch" + ) + # Compute perf metrics + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + actual_seq_lens_kv_flat = actual_seq_lens_kv.flatten().to("cpu") + actual_seq_lens_q_flat = torch.ones_like(actual_seq_lens_kv_flat) + tflops = attention_tflops_per_sec_with_actual_seq_lens( + actual_seq_lens_q_flat, + actual_seq_lens_kv_flat, + head_dim_qk, + head_dim_vo, + num_qo_heads, + False, + median_time, + ) + tb_per_sec = attention_tb_per_sec_with_actual_seq_lens( + actual_seq_lens_q_flat, + actual_seq_lens_kv_flat, + head_dim_qk, + head_dim_vo, + num_qo_heads, + num_kv_heads, + median_time, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + o_dtype=q_dtype, + ) + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["page_size"] = page_size + cur_res["batch_size"] = batch_size + cur_res["s_qo"] = s_qo + cur_res["s_kv"] = s_kv + cur_res["num_qo_heads"] = num_qo_heads + cur_res["num_kv_heads"] = num_kv_heads + cur_res["head_dim_qk"] = head_dim_qk + cur_res["head_dim_vo"] = head_dim_vo + cur_res["causal"] = False + cur_res["q_dtype"] = q_dtype + cur_res["kv_dtype"] = kv_dtype + cur_res["avg_actual_seq_len"] = avg_seq_len_kv + cur_res["random_actual_seq_len"] = args.random_actual_seq_len + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +def testBatchPrefillWithPagedKVCacheWrapper(args): + """ + Test BatchPrefillWithPagedKVCacheWrapper API and equivalent cuDNN API. + Supports fa2, fa3, trtllm-gen, trtllm-native, and cudnn backends. + + This test: + 1. Creates paged KV cache and query tensors for prefill + 2. Runs prefill attention with different backends + 3. Verifies outputs match between backends (if refcheck enabled) + 4. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: Dictionary containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testBatchPrefillWithPagedKVCacheWrapper") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + # Basic setup + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + q_init_dtype = torch.bfloat16 + kv_init_dtype = torch.bfloat16 + rtol = 2e-1 + atol = 1e-2 + res = [] + + q_dtype = dtype_str_to_torch_dtype(args.q_dtype) + if q_dtype not in [torch.bfloat16, torch.float8_e4m3fn]: + print(f"[ERROR] Unsupported q_dtype: {args.q_dtype}") + return res + + kv_dtype = dtype_str_to_torch_dtype(args.kv_dtype) + if kv_dtype not in [torch.bfloat16, torch.float8_e4m3fn]: + print(f"[ERROR] Unsupported kv_dtype: {args.kv_dtype}") + return res + + # Parse and validate backend configurations + backends = args.backends + page_size = args.page_size + batch_size = args.batch_size + s_qo = args.s_qo + s_kv = args.s_kv + num_qo_heads = args.num_qo_heads + num_kv_heads = args.num_kv_heads + head_dim_qk = args.head_dim_qk + head_dim_vo = args.head_dim_vo + causal = args.causal + is_cuda_graph_compatible = not args.no_cuda_graph + # return_lse = not args.no_lse # TO-DO: Add support for this + run_refcheck = args.refcheck + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + # Check for backend-specific constraints + if "fa2" in backends: + remove_fa2 = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + print("[INFO] FA2 backend does not support FP8. Skipping.") + remove_fa2 = True + if remove_fa2: + backends.remove("fa2") + if "cudnn" in backends: + remove_cudnn = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] cuDNN backend does not support FP8. Skipping.") + remove_cudnn = True + if remove_cudnn: + backends.remove("cudnn") + + if "trtllm-gen" in backends: + remove_trtllm = False + if not causal: + print("[INFO] trtllm-gen backend currently requires causal = True") + remove_trtllm = True + if remove_trtllm: + backends.remove("trtllm-gen") + if "trtllm-native" in backends: + remove_trtllm_native = False + if not causal: + print("[INFO] trtllm-native backend currently requires causal = True") + remove_trtllm_native = True + if remove_trtllm_native: + backends.remove("trtllm-native") + + if "cutlass" in backends: + print("[INFO] CUTLASS backend does not support prefill. Skipping.") + remove_cutlass = True + if remove_cutlass: + backends.remove("cutlass") + + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + # Check for layer-specific constraints + layer_not_supported = False + if s_qo > s_kv: + print("[ERROR] s_qo > s_kv is not supported. Exiting.") + layer_not_supported = True + if layer_not_supported: + print("[ERROR] Layer not supported. Exiting.") + return res + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + + # Sample sequence lengths. + # If s_qo == s_kv, then make sampled actual_seq_lens_kv the same as actual_seq_lens_q. + # IF s_qo < s_kv, then sample actual_seq_lens_kv separately. Then ensure actual_seq_lens_kv is at least as long as actual_seq_lens_q. + actual_seq_lens_q = sample_actual_seq_lens( + s_qo, batch_size, None, args.random_actual_seq_len + ) + if s_qo == s_kv: + if args.verbose >= 2: + print( + "[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q" + ) + actual_seq_lens_kv = actual_seq_lens_q.clone() + else: # s_qo < s_kv + if args.verbose >= 2: + print("[VVERBOSE] s_qo < s_kv, sampling actual_seq_lens_kv") + actual_seq_lens_kv = sample_actual_seq_lens( + s_kv, batch_size, None, args.random_actual_seq_len + ) + actual_seq_lens_kv = torch.maximum(actual_seq_lens_kv, actual_seq_lens_q) + + avg_seq_len_q = actual_seq_lens_q.sum().item() // batch_size + avg_seq_len_kv = actual_seq_lens_kv.sum().item() // batch_size + if args.verbose >= 1: + print(f"[VERBOSE] Average actual qo seq len: {avg_seq_len_q}") + print(f"[VERBOSE] Average actual kv seq len: {avg_seq_len_kv}") + if args.verbose >= 2: + print(f"[VVERBOSE] {actual_seq_lens_q.flatten() = }") + print(f"[VVERBOSE] {actual_seq_lens_kv.flatten() = }") + + cumsum_s_qo = torch.sum(actual_seq_lens_q) + q = torch.randn( + cumsum_s_qo, num_qo_heads, head_dim_qk, device=device, dtype=q_init_dtype + ) + if args.verbose >= 2: + print(f"[VVERBOSE] {q.shape = }") + + # Create KV cache + num_pages_per_seq = (s_kv + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + + if args.verbose >= 2: + print(f"[VVERBOSE] {num_pages_per_seq = }") + print(f"[VVERBOSE] {total_num_pages = }") + + kv_cache_shape = (total_num_pages, 2, num_kv_heads, page_size, head_dim_qk) + kv_cache = torch.randn(size=kv_cache_shape, dtype=kv_init_dtype).to(device) + kv_cache = kv_cache.as_strided( + kv_cache.shape, + ( + 2 * page_size * num_kv_heads * head_dim_qk, + page_size * num_kv_heads * head_dim_qk, + head_dim_qk, + num_kv_heads * head_dim_qk, + 1, + ), + ) + k_cache_view, v_cache_view = kv_cache[:, 0, :, :, :], kv_cache[:, 1, :, :, :] + + v_cache = v_cache_view.as_strided( + v_cache_view.shape, + ( + 2 * page_size * num_kv_heads * head_dim_qk, + head_dim_qk, + num_kv_heads * head_dim_qk, + 1, + ), + ) + k_cache = k_cache_view.as_strided( + k_cache_view.shape, + ( + 2 * page_size * num_kv_heads * head_dim_qk, + head_dim_qk, + num_kv_heads * head_dim_qk, + 1, + ), + ) + + # Now initialize the page tables + block_tables = torch.tensor( + [ + [k + i * num_pages_per_seq for k in range(num_pages_per_seq)] + for i in range(batch_size) + ], + dtype=torch.int, + device=device, + ) + + actual_seq_lens_q_device = actual_seq_lens_q.to(device) + actual_seq_lens_kv_device = actual_seq_lens_kv.to(device) + q_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_q_device.view(-1), dim=0) + * head_dim_qk + * num_qo_heads, + ] + ) + .long() + .to(device) + ) # For cuDNN + qo_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_q_device.view(-1), dim=0), + ] + ) + .int() + .to(device) + ) + + # Because actual_seq_lens_kv is the same as actual_seq_lens_q, kv_indptr will become the same as qo_indptr + kv_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum( + (actual_seq_lens_kv_device.flatten() + page_size - 1) // page_size, + dim=0, + ), + ] + ) + .int() + .to(device) + ) + kv_indices = torch.zeros(kv_indptr[-1], device=device, dtype=torch.int32) + for i in range(len(kv_indptr) - 1): + start_idx = kv_indptr[i] + end_idx = kv_indptr[i + 1] + kv_indices[start_idx:end_idx] = torch.arange( + i * num_pages_per_seq, + i * num_pages_per_seq + (end_idx - start_idx), + device=device, + ) + kv_last_page_len = ( + torch.where( + actual_seq_lens_kv_device.flatten() % page_size == 0, + torch.full((batch_size,), page_size, device=device), + actual_seq_lens_kv_device.flatten() % page_size, + ) + .int() + .to(device) + ) + + scale = float(1.0 / (head_dim_qk**0.5)) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + if args.verbose >= 2: + print(f"[VVERBOSE] {kv_cache.shape = }") + print(f"[VVERBOSE] {kv_cache.stride() = }") + print(f"[VVERBOSE] {block_tables.shape = }") + print(f"[VVERBOSE] {qo_indptr.shape = }") + print(f"[VVERBOSE] {qo_indptr.dtype = }") + print(f"[VVERBOSE] {kv_indptr.shape = }") + print(f"[VVERBOSE] {kv_indices.shape = }") + print(f"[VVERBOSE] {kv_last_page_len.shape = }") + print(f"[VVERBOSE] {scale = }") + + # Prepare wrappers + backend_wrappers = {} + for backend in backends: + if backend in ["fa2", "fa3", "trtllm-gen"]: + backend_wrappers[backend] = ( + flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + "HND", + use_cuda_graph=is_cuda_graph_compatible + if backend != "fa2" + else False, + qo_indptr_buf=qo_indptr, + paged_kv_indptr_buf=kv_indptr, + paged_kv_indices_buf=kv_indices, + paged_kv_last_page_len_buf=kv_last_page_len, + backend=backend, + ) + ) + backend_wrappers[backend].plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + page_size, + pos_encoding_mode="NONE", + causal=causal, + q_data_type=q_dtype, + kv_data_type=kv_dtype, + block_tables=block_tables, + ) + + k_scale, v_scale = None, None + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + q = q.to(q_dtype) + if kv_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + k_data, v_data = torch.chunk(kv_cache, 2, dim=1) + k_scale = k_data.amax().item() / 256 + v_scale = v_data.amax().item() / 256 + k_fp8 = (k_data / k_scale).to(kv_dtype) + v_fp8 = (v_data / v_scale).to(kv_dtype) + kv_cache = torch.cat([k_fp8, v_fp8], dim=1) + + def run_backend_wrapper(backend): + if backend in ["fa2", "fa3", "trtllm-gen"]: + return backend_wrappers[backend].run( + q, kv_cache, k_scale=k_scale, v_scale=v_scale + ) + elif backend == "cudnn": + return flashinfer.prefill.cudnn_batch_prefill_with_kv_cache( + q, + k_cache, + v_cache, + scale, + workspace_buffer, + max_token_per_sequence=s_qo, + max_sequence_kv=s_kv, + actual_seq_lens_q=actual_seq_lens_q_device, + actual_seq_lens_kv=actual_seq_lens_kv_device, + block_tables=block_tables, + causal=causal, + return_lse=True, + is_cuda_graph_compatible=is_cuda_graph_compatible, + batch_offsets_q=q_indptr, + batch_offsets_o=q_indptr, + )[0] + elif backend == "trtllm-native": + return flashinfer.prefill.trtllm_batch_context_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=workspace_buffer, + block_tables=block_tables, + seq_lens=actual_seq_lens_kv_device, + max_q_len=s_qo, + max_kv_len=s_kv, + bmm1_scale=scale if k_scale is None else k_scale * scale, + bmm2_scale=1.0 if v_scale is None else v_scale, + batch_size=batch_size, + cum_seq_lens_q=qo_indptr, + cum_seq_lens_kv=kv_indptr, + ) + else: + print(f"[ERROR] Backend {backend} not supported") + return res + + has_reference_output = False + # Iterate over each backend: + for cur_backend in backends: + # Clear workspace buffer to prevent unexpected interactions between backends. + workspace_buffer.zero_() + if run_refcheck: + outputs[cur_backend] = run_backend_wrapper(cur_backend).detach().clone() + if cur_backend == "fa2": + has_reference_output = True + reference_output = outputs[cur_backend] + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend_wrapper(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=(is_cuda_graph_compatible and cur_backend != "fa2"), + ) + + # Perform reference check + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 1: + if run_refcheck and has_reference_output: + if reference_output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + if args.verbose >= 2: + print( + "[VVERBOSE] Reference output is FP8. Converting to float32 for reference check." + ) + reference_output = reference_output.to(torch.float32) + tested_outputs = [output.to(torch.float32) for output in tested_outputs] + for i in range(len(tested_backends)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats(reference_output, tested_outputs[i], rtol, atol) + if num_different_elements > 0: + print( + f"[ERROR] Output tensor mismatch between backends fa2 and {tested_backends[i]}: " + f"{num_different_elements} / {num_elements} ({num_different_elements_percentage:.2f}%) elements are different" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch" + ) + + # Compute perf metrics + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + actual_seq_lens_q_flat = actual_seq_lens_q.flatten().to("cpu") + actual_seq_lens_kv_flat = actual_seq_lens_kv.flatten().to("cpu") + tflops = attention_tflops_per_sec_with_actual_seq_lens( + actual_seq_lens_q_flat, + actual_seq_lens_kv_flat, + head_dim_qk, + head_dim_vo, + num_qo_heads, + causal, + median_time, + ) + tb_per_sec = attention_tb_per_sec_with_actual_seq_lens( + actual_seq_lens_q_flat, + actual_seq_lens_kv_flat, + head_dim_qk, + head_dim_vo, + num_qo_heads, + num_kv_heads, + median_time, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + o_dtype=q_dtype, + ) + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["page_size"] = page_size + cur_res["batch_size"] = batch_size + cur_res["s_qo"] = s_qo + cur_res["s_kv"] = s_kv + cur_res["num_qo_heads"] = num_qo_heads + cur_res["num_kv_heads"] = num_kv_heads + cur_res["head_dim_qk"] = head_dim_qk + cur_res["head_dim_vo"] = head_dim_vo + cur_res["causal"] = causal + cur_res["q_dtype"] = q_dtype + cur_res["kv_dtype"] = kv_dtype + cur_res["avg_actual_seq_len"] = avg_seq_len_q + cur_res["random_actual_seq_len"] = args.random_actual_seq_len + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +def testBatchPrefillWithRaggedKVCacheWrapper(args): + """ + Test BatchPrefillWithRaggedKVCacheWrapper API and equivalent cuDNN API. + Supports fa2, fa3, cutlass, and cudnn backends. + + This test: + 1. Creates ragged KV cache and query tensors for prefill + 2. Runs prefill attention with different backends + 3. Verifies outputs match between backends (if refcheck enabled) + 4. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: Dictionary containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testBatchPrefillWithRaggedKVCacheWrapper") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + # Basic setup + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + q_init_dtype = torch.bfloat16 + kv_init_dtype = torch.bfloat16 + rtol = 2e-1 + atol = 1e-2 + res = [] + + q_dtype = dtype_str_to_torch_dtype(args.q_dtype) + if q_dtype not in [torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2]: + print(f"[ERROR] Unsupported q_dtype: {args.q_dtype}") + return res + kv_dtype = dtype_str_to_torch_dtype(args.kv_dtype) + if kv_dtype not in [torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2]: + print(f"[ERROR] Unsupported kv_dtype: {args.kv_dtype}") + return res + + # Parse and validate backend configurations + backends = args.backends + batch_size = args.batch_size + s_qo = args.s_qo + s_kv = args.s_kv + num_qo_heads = args.num_qo_heads + num_kv_heads = args.num_kv_heads + head_dim_qk = args.head_dim_qk + head_dim_vo = args.head_dim_vo + causal = args.causal + is_cuda_graph_compatible = not args.no_cuda_graph + # return_lse = not args.no_lse # TO-DO: Add support for this + run_refcheck = args.refcheck + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + # Check for backend-specific constraints + if "fa2" in backends: + remove_fa2 = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + print("[INFO] FA2 backend does not support FP8. Skipping.") + remove_fa2 = True + if remove_fa2: + backends.remove("fa2") + if "cudnn" in backends: + remove_cudnn = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] CUDNN backend does not support FP8. Skipping.") + remove_cudnn = True + if remove_cudnn: + backends.remove("cudnn") + + if "cutlass" in backends: + remove_cutlass = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] CUTLASS backend does not support FP8. Skipping.") + remove_cutlass = True + if not ( + (head_dim_qk == 128 and head_dim_qk == head_dim_vo) or head_dim_qk == 192 + ): + print("[INFO] CUTLASS backend requires head dimension to be 128 or 192") + remove_cutlass = True + if remove_cutlass: + backends.remove("cutlass") + + if "trtllm-gen" in backends: + print("[INFO] trtllm-gen backend does not support ragged prefill. Skipping.") + remove_trtllm = True + if remove_trtllm: + backends.remove("trtllm-gen") + if "trtllm-native" in backends: + remove_trtllm_native = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] trtllm-native backend does not support FP8. Skipping.") + remove_trtllm_native = True + if not (head_dim_qk == 192 and head_dim_vo == 128): + print( + "[INFO] trtllm-native backend requires head_dim_qk == 192 and head_dim_vo == 128" + ) + remove_trtllm_native = True + if remove_trtllm_native: + backends.remove("trtllm-native") + + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + # Check for layer-specific constraints + layer_not_supported = False + if s_qo > s_kv: + print("[ERROR] s_qo > s_kv is not supported. Exiting.") + layer_not_supported = True + if layer_not_supported: + print("[ERROR] Layer not supported. Exiting.") + return res + + backend_times = {backend: [] for backend in backends} + outputs = {} + + # Sample sequence lengths. + # If s_qo == s_kv, then make sampled actual_seq_lens_kv the same as actual_seq_lens_q. + # IF s_qo < s_kv, then sample actual_seq_lens_kv separately. Then ensure actual_seq_lens_kv is at least as long as actual_seq_lens_q. + actual_seq_lens_q = sample_actual_seq_lens( + s_qo, batch_size, None, args.random_actual_seq_len + ) + if s_qo == s_kv: + if args.verbose >= 2: + print( + "[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q" + ) + actual_seq_lens_kv = actual_seq_lens_q.clone() + else: # s_qo < s_kv + if args.verbose >= 2: + print("[VVERBOSE] s_qo < s_kv, sampling actual_seq_lens_kv") + actual_seq_lens_kv = sample_actual_seq_lens( + s_kv, batch_size, None, args.random_actual_seq_len + ) + actual_seq_lens_kv = torch.maximum(actual_seq_lens_kv, actual_seq_lens_q) + + avg_seq_len_q = actual_seq_lens_q.sum().item() // batch_size + avg_seq_len_kv = actual_seq_lens_kv.sum().item() // batch_size + if args.verbose >= 1: + print(f"[VERBOSE] Average actual qo seq len: {avg_seq_len_q}") + print(f"[VERBOSE] Average actual kv seq len: {avg_seq_len_kv}") + if args.verbose >= 2: + print(f"[VVERBOSE] {actual_seq_lens_q.flatten() = }") + print(f"[VVERBOSE] {actual_seq_lens_kv.flatten() = }") + + cumsum_s_qo = torch.sum(actual_seq_lens_q) + cumsum_s_kv = torch.sum(actual_seq_lens_kv) + q = torch.randn( + cumsum_s_qo, num_qo_heads, head_dim_qk, device=device, dtype=q_init_dtype + ) + if args.verbose >= 2: + print(f"[VVERBOSE] {q.shape = }") + + k = torch.randn( + cumsum_s_kv, num_kv_heads, head_dim_qk, device=device, dtype=kv_init_dtype + ) + v = torch.randn( + cumsum_s_kv, num_kv_heads, head_dim_vo, device=device, dtype=kv_init_dtype + ) + + block_tables = None + + ## The following are for BatchPrefillWithRaggedKVCacheWrapper + actual_seq_lens_q_device = actual_seq_lens_q.to(device) + actual_seq_lens_kv_device = actual_seq_lens_kv.to(device) + + q_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_q_device.view(-1), dim=0) + * head_dim_qk + * num_qo_heads, + ] + ) + .long() + .to(device) + ) # For cuDNN + + k_indptr = torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_kv_device.view(-1), dim=0) + * head_dim_qk + * num_kv_heads, + ] + ).long() + + v_indptr = torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_kv_device.view(-1), dim=0) + * head_dim_vo + * num_kv_heads, + ] + ).long() + + o_indptr = torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_q_device.view(-1), dim=0) + * head_dim_vo + * num_qo_heads, + ] + ).long() + + batch_offsets_stats = torch.cat( + [ + torch.zeros( + 1, + device=actual_seq_lens_q_device.device, + dtype=actual_seq_lens_q_device.dtype, + ), + torch.cumsum(actual_seq_lens_q_device.flatten(), dim=0) * num_qo_heads, + ] + ).cuda() + + qo_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_q_device.view(-1), dim=0), + ] + ) + .int() + .to(device) + ) + # Because actual_seq_lens_kv is the same as actual_seq_lens_q, kv_indptr will become the same as qo_indptr + kv_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum(actual_seq_lens_kv_device.view(-1), dim=0), + ] + ) + .int() + .to(device) + ) + + scale = float(1.0 / (head_dim_qk**0.5)) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + if args.verbose >= 2: + print(f"[VVERBOSE] {k.shape = }") + print(f"[VVERBOSE] {v.shape = }") + print(f"[VVERBOSE] {qo_indptr.shape = }") + print(f"[VVERBOSE] {kv_indptr.shape = }") + print(f"[VVERBOSE] {scale = }") + + # Prepare wrappers + backend_wrappers = {} + for backend in backends: + if backend in ["cutlass", "fa2", "fa3", "trtllm-gen"]: + backend_wrappers[backend] = ( + flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, + "NHD", + use_cuda_graph=is_cuda_graph_compatible + if backend != "fa2" + else False, + qo_indptr_buf=qo_indptr, + kv_indptr_buf=kv_indptr, + backend=backend, + ) + ) + backend_wrappers[backend].plan( + qo_indptr, + kv_indptr, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo=head_dim_vo, + causal=causal, + q_data_type=q_dtype, + kv_data_type=kv_dtype, + ) + + k_scale, v_scale = None, None + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + q = q.to(q_dtype) + if kv_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + k_scale = k.amax().item() / 256 + v_scale = v.amax().item() / 256 + k = (k / k_scale).to(kv_dtype) + v = (v / v_scale).to(kv_dtype) + + def run_backend_wrapper(backend): + if backend in ["cutlass", "fa2", "fa3", "trtllm-gen"]: + return backend_wrappers[backend].run_return_lse(q, k, v)[0] + elif backend == "cudnn": + return flashinfer.prefill.cudnn_batch_prefill_with_kv_cache( + q, + k, + v, + scale, + workspace_buffer, + max_token_per_sequence=s_qo, + max_sequence_kv=s_kv, + actual_seq_lens_q=actual_seq_lens_q_device, + actual_seq_lens_kv=actual_seq_lens_kv_device, + block_tables=block_tables, + causal=causal, + return_lse=True, + batch_offsets_q=q_indptr, + batch_offsets_k=k_indptr, + batch_offsets_v=v_indptr, + batch_offsets_o=o_indptr, + batch_offsets_stats=batch_offsets_stats, + is_cuda_graph_compatible=True, + )[0] + elif backend == "trtllm-native": + return flashinfer.prefill.trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=workspace_buffer, + seq_lens=actual_seq_lens_kv_device, + max_q_len=s_qo, + max_kv_len=s_kv, + bmm1_scale=scale, + bmm2_scale=1.0, + o_sf_scale=-1, + batch_size=batch_size, + window_left=-1, + cum_seq_lens_q=qo_indptr, + cum_seq_lens_kv=kv_indptr, + enable_pdl=False, + is_causal=causal, + return_lse=True, + )[0] + else: + print(f"[ERROR] Backend {backend} not supported") + return res + + has_reference_output = False + # Iterate over each backend: + for cur_backend in backends: + # Clear workspace buffer to prevent unexpected interactions between backends. + workspace_buffer.zero_() + if run_refcheck: + outputs[cur_backend] = run_backend_wrapper(cur_backend).detach().clone() + if cur_backend == "fa2": + has_reference_output = True + reference_output = outputs[cur_backend] + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend_wrapper(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=True, + enable_cupti=args.use_cupti, + use_cuda_graph=(is_cuda_graph_compatible and cur_backend != "fa2"), + ) + + # Perform reference check + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 1: + if run_refcheck and has_reference_output: + if reference_output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + if args.verbose >= 2: + print( + "[VVERBOSE] Reference output is FP8. Converting to float32 for reference check." + ) + reference_output = reference_output.to(torch.float32) + tested_outputs = [output.to(torch.float32) for output in tested_outputs] + for i in range(len(tested_backends)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats(reference_output, tested_outputs[i], rtol, atol) + if num_different_elements > 0: + print( + f"[ERROR] Output tensor mismatch between backends fa2 and {tested_backends[i]}: " + f"{num_different_elements} / {num_elements} ({num_different_elements_percentage:.2f}%) elements are different" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch" + ) + + # Compute perf metrics + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + actual_seq_lens_q_flat = actual_seq_lens_q.flatten().to("cpu") + actual_seq_lens_kv_flat = actual_seq_lens_kv.flatten().to("cpu") + tflops = attention_tflops_per_sec_with_actual_seq_lens( + actual_seq_lens_q_flat, + actual_seq_lens_kv_flat, + head_dim_qk, + head_dim_vo, + num_qo_heads, + causal, + median_time, + ) + tb_per_sec = attention_tb_per_sec_with_actual_seq_lens( + actual_seq_lens_q_flat, + actual_seq_lens_kv_flat, + head_dim_qk, + head_dim_vo, + num_qo_heads, + num_kv_heads, + median_time, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + o_dtype=q_dtype, + ) + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["page_size"] = 0 # No page size for ragged + cur_res["batch_size"] = batch_size + cur_res["s_qo"] = s_qo + cur_res["s_kv"] = s_kv + cur_res["num_qo_heads"] = num_qo_heads + cur_res["num_kv_heads"] = num_kv_heads + cur_res["head_dim_qk"] = head_dim_qk + cur_res["head_dim_vo"] = head_dim_vo + cur_res["causal"] = causal + cur_res["q_dtype"] = q_dtype + cur_res["kv_dtype"] = kv_dtype + cur_res["avg_actual_seq_len"] = avg_seq_len_q + cur_res["random_actual_seq_len"] = args.random_actual_seq_len + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +def testBatchMLAPagedAttentionWrapper(args): + """ + Test BatchMLAPagedAttentionWrapper and equivalent APIs. + Supports fa2, fa3, cutlass, and trtllm-native. + + This test: + 1. Creates paged query and key-value cache tensors + 2. Runs MLA with different backends + 3. Verifies outputs match between backends + 4. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testBatchMLAPagedAttentionWrapper") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + # Basic setup + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + q_init_dtype = torch.bfloat16 + kv_init_dtype = torch.bfloat16 + rtol = 2e-1 + atol = 1e-2 + res = [] + + # Handle different query data types. + q_dtype = dtype_str_to_torch_dtype(args.q_dtype) + if q_dtype not in [torch.bfloat16, torch.float8_e4m3fn]: + print(f"[ERROR] Unsupported q_dtype: {args.q_dtype}") + return res + + # Handle different KV cache data types. + kv_dtype = dtype_str_to_torch_dtype(args.kv_dtype) + if kv_dtype not in [torch.bfloat16, torch.float8_e4m3fn]: + print(f"[ERROR] Unsupported kv_dtype: {args.kv_dtype}") + return res + + backends = args.backends + page_size = args.page_size + batch_size = args.batch_size + s_qo = args.s_qo + s_kv = args.s_kv + num_qo_heads = args.num_qo_heads + # num_kv_heads not used in MLA + # head_dim_qk = args.head_dim_qk + assert args.head_dim_ckv is not None, "head_dim_ckv must be provided for MLA" + assert args.head_dim_kpe is not None, "head_dim_kpe must be provided for MLA" + head_dim_ckv = args.head_dim_ckv + head_dim_kpe = args.head_dim_kpe + is_cuda_graph_compatible = not args.no_cuda_graph + causal = False # False for MLA + run_refcheck = args.refcheck + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + # Check for backend-specific constraints + if "fa2" in backends: + remove_fa2 = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] FA2 backend does not support FP8. Skipping.") + remove_fa2 = True + if remove_fa2: + backends.remove("fa2") + if "fa3" in backends: + remove_fa3 = False + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] FA3 backend does not support FP8. Skipping.") + remove_fa3 = True + if remove_fa3: + backends.remove("fa3") + if "cutlass" in backends: + remove_cutlass = False + if page_size not in [32, 64]: + print( + "[INFO] Cutlass MLA backend only supports page size 32 or 64. Skipping." + ) + remove_cutlass = True + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] or kv_dtype in [ + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + print("[INFO] Cutlass MLA backend does not support FP8. Skipping.") + remove_cutlass = True + if remove_cutlass: + backends.remove("cutlass") + if "trtllm-native" in backends: + remove_trtllm_native = False + if page_size not in [32, 64]: + print( + "[INFO] trtllm-native backend only supports page size 32 or 64. Skipping." + ) + remove_trtllm_native = True + if remove_trtllm_native: + backends.remove("trtllm-native") + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + + actual_seq_lens_kv = sample_actual_seq_lens( + s_kv, batch_size, device, args.random_actual_seq_len + ) + sum_seq_kv = torch.sum(actual_seq_lens_kv).item() + avg_seq_len_kv = sum_seq_kv // batch_size + + if args.verbose >= 1: + print(f"[VERBOSE] Average actual seq len: {avg_seq_len_kv}") + if args.verbose >= 2: + print(f"[VVERBOSE] {actual_seq_lens_kv.flatten() = }") + + q_nope = torch.rand( + batch_size, num_qo_heads, head_dim_ckv, dtype=q_init_dtype, device="cuda" + ) + q_pe = torch.zeros( + batch_size, num_qo_heads, head_dim_kpe, dtype=q_init_dtype, device="cuda" + ) + q = torch.cat([q_nope, q_pe], dim=2) + + if args.verbose >= 2: + print(f"[VVERBOSE] {q_nope.shape = }") + print(f"[VVERBOSE] {q_pe.shape = }") + print(f"[VVERBOSE] {q.shape = }") + + # Create KV cache + num_pages_per_seq = (s_kv + page_size - 1) // page_size + total_num_pages = num_pages_per_seq * batch_size + + # Now initialize the page tables + block_tables = torch.tensor( + [ + [k + i * num_pages_per_seq for k in range(num_pages_per_seq)] + for i in range(batch_size) + ], + dtype=torch.int, + device=device, + ) + + if args.verbose >= 2: + print(f"[VVERBOSE] {num_pages_per_seq = }") + print(f"[VVERBOSE] {total_num_pages = }") + print(f"[VVERBOSE] {block_tables.shape = }") + + # Initialize KV cache with appropriate shape and stride + ckv_cache_shape = ( + total_num_pages, + page_size, + head_dim_ckv, + ) + ckv_cache = torch.randn(size=ckv_cache_shape, dtype=kv_init_dtype, device=device) + + kpe_cache_shape = ( + total_num_pages, + page_size, + head_dim_kpe, + ) + kpe_cache = torch.randn(size=kpe_cache_shape, dtype=kv_init_dtype, device=device) + kv_cache = torch.cat([ckv_cache, kpe_cache], dim=2) + + qo_indptr = torch.arange(0, batch_size + 1, device=device).int() + kv_indptr = ( + torch.cat( + [ + torch.tensor([0], device=device), + torch.cumsum( + (actual_seq_lens_kv.flatten() + page_size - 1) // page_size, dim=0 + ), + ] + ) + .int() + .to(device) + ) + + # kv_indices[-1] is the total number of actual pages + kv_indices = torch.zeros(kv_indptr[-1], device=device, dtype=torch.int32) + for i in range(len(kv_indptr) - 1): + start_idx = kv_indptr[i] + end_idx = kv_indptr[i + 1] + kv_indices[start_idx:end_idx] = torch.arange( + i * num_pages_per_seq, + i * num_pages_per_seq + (end_idx - start_idx), + device=device, + ) + + sm_scale = 1.0 / ((128 + 64) ** 0.5) # For DeepSeek-R1 + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + + if args.verbose >= 2: + print(f"[VVERBOSE] {ckv_cache.shape = }") + print(f"[VVERBOSE] {kpe_cache.shape = }") + print(f"[VVERBOSE] {kv_cache.shape = }") + print(f"[VVERBOSE] {qo_indptr.shape = }") + print(f"[VVERBOSE] {kv_indptr.shape = }") + print(f"[VVERBOSE] {kv_indices.shape = }") + print(f"[VVERBOSE] {actual_seq_lens_kv.shape = }") + print(f"[VVERBOSE] {sm_scale = }") + print(f"[VVERBOSE] {workspace_buffer.shape = }") + + # Create wrapper + backend_wrappers = {} + for backend in backends: + if backend in ["fa2", "fa3", "cutlass"]: + backend_wrappers[backend] = flashinfer.mla.BatchMLAPagedAttentionWrapper( + float_workspace_buffer=workspace_buffer, + use_cuda_graph=is_cuda_graph_compatible, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=actual_seq_lens_kv, + backend=backend, + ) + if backend != "cutlass": + backend_wrappers[backend].plan( + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=actual_seq_lens_kv, + num_heads=num_qo_heads, + head_dim_ckv=head_dim_ckv, + head_dim_kpe=head_dim_kpe, + page_size=page_size, + causal=causal, + sm_scale=sm_scale, + q_data_type=q_dtype, + kv_data_type=kv_dtype, + ) + + if q_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + q = q.to(q_dtype) + q_pe = q_pe.to(q_dtype) + q_nope = q_nope.to(q_dtype) + if kv_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + ckv_cache = ckv_cache.to(kv_dtype) + kpe_cache = kpe_cache.to(kv_dtype) + kv_cache = kv_cache.to(kv_dtype) + + def run_backend_wrapper(backend): + if backend in ["fa2", "fa3"]: + return backend_wrappers[backend].run( + q_nope, + q_pe, + ckv_cache, + kpe_cache, + page_table=block_tables, + return_lse=False, + ) + elif backend == "cutlass": + return backend_wrappers[backend].run( + q_nope, + q_pe, + ckv_cache, + kpe_cache, + kv_len=actual_seq_lens_kv.flatten(), + page_table=block_tables, + return_lse=False, + ) + elif backend == "trtllm-native": + return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=q.unsqueeze(1), + kv_cache=kv_cache.unsqueeze(1), + workspace_buffer=workspace_buffer, + qk_nope_head_dim=128, # To-do: Why?? + kv_lora_rank=head_dim_ckv, + qk_rope_head_dim=head_dim_kpe, + block_tables=block_tables, + seq_lens=actual_seq_lens_kv, + max_seq_len=s_kv, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + ).squeeze(1) + else: + print(f"[ERROR] Unsupported backend: {backend}") + return res + + has_reference_output = False + # Iterate over each backend: + for cur_backend in backends: + # Clear workspace buffer to prevent unexpected interactions between backends. + workspace_buffer.zero_() + if run_refcheck: + outputs[cur_backend] = run_backend_wrapper(cur_backend).detach().clone() + if cur_backend == "fa2": + has_reference_output = True + reference_output = outputs[cur_backend] + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend_wrapper(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=(is_cuda_graph_compatible and cur_backend != "fa2"), + ) + + # Perform reference check + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 1: + if run_refcheck and has_reference_output: + if reference_output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + reference_output = reference_output.to(torch.float32) + tested_outputs = [output.to(torch.float32) for output in tested_outputs] + for i in range(len(tested_outputs)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats(reference_output, tested_outputs[i], rtol, atol) + if num_different_elements > 0: + print( + f"[ERROR] Output tensor mismatch between backends fa2 and {tested_backends[i]}: " + f"{num_different_elements} / {num_elements} ({num_different_elements_percentage:.2f}%) elements are different" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch" + ) + # Compute perf metrics + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + actual_seq_lens_kv_flat = actual_seq_lens_kv.flatten().to("cpu") + actual_seq_lens_q_flat = torch.ones_like( + actual_seq_lens_kv.flatten().to("cpu") + ) + o_mem_bytes = ( + actual_seq_lens_q_flat.numel() + * num_qo_heads + * head_dim_ckv + * q_dtype.itemsize + ) + qkv_mem_bytes = sum( + [ + _.numel() * _.element_size() + for _ in [q_nope, q_pe, ckv_cache, kpe_cache] + ] + ) + total_mem_bytes = o_mem_bytes + qkv_mem_bytes + tb_per_sec = (total_mem_bytes / (median_time * 1e9)).item() + tflops_total = ( + 2 + * torch.dot( + actual_seq_lens_q_flat.to(torch.float32), + actual_seq_lens_kv_flat.to(torch.float32), + ) + * num_qo_heads + * (2 * head_dim_ckv + head_dim_kpe) + ) + tflops = (tflops_total / (median_time * 1e9)).item() + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + # TO-Do: + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["page_size"] = page_size + cur_res["batch_size"] = batch_size + cur_res["s_qo"] = s_qo + cur_res["s_kv"] = s_kv + cur_res["num_qo_heads"] = num_qo_heads + cur_res["head_dim_ckv"] = head_dim_ckv + cur_res["head_dim_kpe"] = head_dim_kpe + cur_res["causal"] = False + cur_res["q_dtype"] = q_dtype + cur_res["kv_dtype"] = kv_dtype + cur_res["avg_actual_seq_len"] = avg_seq_len_kv + cur_res["random_actual_seq_len"] = args.random_actual_seq_len + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res diff --git a/benchmarks/routines/flashinfer_benchmark_utils.py b/benchmarks/routines/flashinfer_benchmark_utils.py new file mode 100644 index 0000000000..8798f8340f --- /dev/null +++ b/benchmarks/routines/flashinfer_benchmark_utils.py @@ -0,0 +1,312 @@ +import torch + +from flashinfer.testing.utils import set_seed +from flashinfer.utils import get_compute_capability + +# Output columns for the test results. +output_column_dict = { + "perf": [ + "routine", + "median_time", + "std_time", + "tflops", + "tb_per_sec", + "backend", + ], + "attention": [ + "page_size", + "batch_size", + "s_qo", + "s_kv", + "num_qo_heads", + "num_kv_heads", + "head_dim_qk", + "head_dim_vo", + "head_dim_ckv", + "head_dim_kpe", + "causal", + "q_dtype", + "kv_dtype", + "avg_actual_seq_len", + "random_actual_seq_len", + ], + "gemm": [ + "m", + "n", + "k", + "group_size", + "tile_size", + "scale_major_mode", + "out_dtype", + "mma_sm", + "use_128x4_sf_layout", + "use_nvfp4", + ], + "moe": [ + "num_tokens", + "hidden_size", + "intermediate_size", + "num_experts", + "top_k", + "n_group", + "topk_group", + "routed_scaling_factor", + "local_expert_offset", + "local_num_experts", + "tile_tokens_dim", + "routing_method", + "use_shuffled_weight", + "weight_layout", + "use_routing_bias", + "use_routing_scales_on_input", + "input_dtype", + "weight_dtype", + "gated_act", + # CUTLASS fused MoE specific + "cutlass_variant", + "quantized_input", + "tp_size", + "tp_rank", + "ep_size", + "ep_rank", + ], + "general": [ + "refcheck", + "no_cuda_graph", + "use_cupti", + "allow_output_mismatch", + "random_seed", + "case_tag", + "generate_repro_command", + "repro_command", + ], +} + +full_output_columns = ( + output_column_dict["perf"] + + output_column_dict["attention"] + + output_column_dict["gemm"] + + output_column_dict["moe"] + + output_column_dict["general"] +) + +benchmark_apis = { + "attention": [ + "BatchDecodeWithPagedKVCacheWrapper", + "BatchPrefillWithPagedKVCacheWrapper", + "BatchPrefillWithRaggedKVCacheWrapper", + "BatchMLAPagedAttentionWrapper", + ], + "gemm": [ + "gemm_fp8_nt_groupwise", + "group_gemm_fp8_nt_groupwise", + "bmm_fp8", + "mm_fp4", + ], + "moe": [ + "trtllm_fp4_block_scale_moe", + "trtllm_fp8_block_scale_moe", + "trtllm_fp8_per_tensor_scale_moe", + "cutlass_fused_moe", + ], +} + + +def print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec): + output_backend_width = 15 + print( + f"[PERF] {backend.ljust(output_backend_width)[:output_backend_width]}:: median time {median_time:.3f} ms; std {std_time:.3f} ms; achieved tflops {tflops:.3f} TFLOPs/sec; achieved tb_per_sec {tb_per_sec:.3f} TB/sec" + ) + + +def get_device(args): + # Synchronize to ensure that the device is ready after previous tests + torch.cuda.empty_cache() + torch.cuda.synchronize() + set_seed(args.random_seed) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + gpu_name = torch.cuda.get_device_name(torch.cuda.current_device()).replace(" ", "_") + if args.verbose >= 2: + print(f"[VVERBOSE] {gpu_name = }") + return device + + +def is_close_stats(input, other, rtol=1e-5, atol=1e-8): + close_tensor = torch.isclose(input, other, rtol=rtol, atol=atol) + num_elements = close_tensor.numel() + num_different_elements = num_elements - close_tensor.sum().item() + return ( + num_different_elements, # number of different elements + num_elements, # total number of elements in tensor + num_different_elements / num_elements * 100.0, + ) + + +def dtype_str_to_torch_dtype(dtype_str): + if dtype_str == "bfloat16": + return torch.bfloat16 + elif dtype_str == "float16": + return torch.float16 + elif dtype_str == "float32": + return torch.float32 + elif dtype_str == "float64": + return torch.float64 + elif dtype_str == "fp8_e4m3": + return torch.float8_e4m3fn + elif dtype_str == "fp8_e5m2": + return torch.float8_e5m2 + else: + raise ValueError(f"Unsupported dtype: {dtype_str}") + + +routine_cc_to_supported_backends = { + # ATTENTION + "BatchDecodeWithPagedKVCacheWrapper": { + # NOTE: trtllm-native calls trtllm_batch_decode_with_kv_cache + "7.5": ["fa2"], + "8.0": ["fa2", "fa2_tc", "cudnn"], + "8.6": ["fa2", "fa2_tc", "cudnn"], + "8.9": ["fa2", "fa2_tc", "cudnn"], + "9.0": ["fa2", "fa2_tc", "cudnn", "trtllm-native"], + "10.0": ["fa2", "fa2_tc", "cudnn", "trtllm-gen", "trtllm-native"], + "10.3": ["fa2", "fa2_tc", "cudnn", "trtllm-gen", "trtllm-native"], + "12.0": ["fa2", "fa2_tc", "cudnn", "trtllm-native"], + }, + "BatchPrefillWithPagedKVCacheWrapper": { + # NOTE: trtllm-native calls trtllm_batch_context_with_kv_cache + "7.5": [], + "8.0": ["fa2", "cudnn"], + "8.6": ["fa2", "cudnn"], + "8.9": ["fa2", "cudnn"], + "9.0": ["fa2", "fa3", "cudnn"], + "10.0": ["fa2", "cudnn", "trtllm-gen", "trtllm-native"], + "10.3": ["fa2", "cudnn", "trtllm-gen", "trtllm-native"], + "12.0": ["fa2", "cudnn"], + }, + "BatchPrefillWithRaggedKVCacheWrapper": { + # NOTE: trtllm-native calls trtllm_ragged_attention_deepseek + "7.5": [], + "8.0": ["fa2", "cudnn"], + "8.6": ["fa2", "cudnn"], + "8.9": ["fa2", "cudnn"], + "9.0": ["fa2", "fa3", "cudnn"], + "10.0": ["fa2", "cudnn", "cutlass", "trtllm-native"], + "10.3": ["fa2", "cudnn", "cutlass", "trtllm-native"], + "12.0": ["fa2", "cudnn"], + }, + "BatchMLAPagedAttentionWrapper": { + # NOTE: trtllm-native calls trtllm_batch_decode_with_kv_cache_mla + "7.5": [], + "8.0": ["fa2"], + "8.6": ["fa2"], + "8.9": ["fa2"], + "9.0": ["fa2", "fa3"], + "10.0": ["fa2", "cutlass", "trtllm-native"], + "10.3": ["fa2", "cutlass", "trtllm-native"], + "12.0": ["fa2"], + }, + # GEMM + "gemm_fp8_nt_groupwise": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["cutlass"], + "10.3": ["cutlass"], + "12.0": [], + }, + "group_gemm_fp8_nt_groupwise": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["cutlass"], + "10.3": ["cutlass"], + "12.0": [], + }, + "bmm_fp8": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": ["cudnn", "cublas"], + "9.0": ["cudnn", "cublas"], + "10.0": ["cudnn", "cublas", "cutlass"], + "10.3": ["cudnn", "cublas", "cutlass"], + "12.0": ["cudnn", "cublas"], + }, + "mm_fp4": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["cudnn", "trtllm", "cutlass"], + "10.3": ["cudnn", "trtllm", "cutlass"], + "12.0": ["cudnn", "cutlass"], + "12.1": ["cudnn", "cutlass"], + }, + # MOE + "trtllm_fp4_block_scale_moe": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["trtllm"], + "10.3": ["trtllm"], + "12.0": [], + }, + "trtllm_fp8_block_scale_moe": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["trtllm"], + "10.3": ["trtllm"], + "12.0": [], + }, + "trtllm_fp8_per_tensor_scale_moe": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["trtllm"], + "10.3": ["trtllm"], + "12.0": [], + }, + "cutlass_fused_moe": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": [], + "10.0": ["cutlass"], + "10.3": ["cutlass"], + "12.0": [], + }, +} + + +def filter_backends_by_compute_capability(backends, routine, device): + # FlashInfer currently does not have an isSupported() function that checks support. + # WAR: Use helper function to check support. + major, minor = get_compute_capability(device) + compute_capability = f"{major}.{minor}" + + # If the compute capability is not supported, return an empty list. + cc_to_supported_backends = routine_cc_to_supported_backends[routine] + supported_backends = cc_to_supported_backends.get(compute_capability, []) + backends_to_remove = [] + for backend in backends: + if backend not in supported_backends: + backends_to_remove.append(backend) + for backend in backends_to_remove: + backends.remove(backend) + print( + f"[WARNING] {backend} for routine {routine} is not supported on compute capability {compute_capability}. Skipping." + ) + return backends diff --git a/benchmarks/routines/gemm.py b/benchmarks/routines/gemm.py new file mode 100644 index 0000000000..17336189d0 --- /dev/null +++ b/benchmarks/routines/gemm.py @@ -0,0 +1,1000 @@ +from collections import defaultdict + +import numpy as np +import torch +import torch.nn.functional as F +from einops import einsum + +import flashinfer +from flashinfer.autotuner import autotune +from flashinfer.testing.utils import ( + bench_gpu_time, + dequantize_fp8, + quantize_fp8, +) + +from .flashinfer_benchmark_utils import ( + dtype_str_to_torch_dtype, + get_device, + print_perf_metrics, + is_close_stats, + filter_backends_by_compute_capability, +) + + +def run_gemm_test(args): + """ + Run a gemm test. + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.routine == "gemm_fp8_nt_groupwise": + return testGemmFp8NtGroupwise(args) + elif args.routine == "group_gemm_fp8_nt_groupwise": + return testGroupGemmFp8NtGroupwise(args) + elif args.routine == "bmm_fp8": + return testBmmFp8(args) + elif args.routine == "mm_fp4": + return testMmFp4(args) + else: + raise ValueError(f"Unsupported routine: {args.routine}") + + +def parse_gemm_args(line, parser): + """ + Parse command line arguments for gemm test configuration. + + Args: + line: Command line arguments + parser: ArgumentParser object already populated with shared arguments + + Returns: + Parsed argument namespace + """ + parser.add_argument( + "--batch_size", + type=int, + required=False, + default=1, + help="Batch size of test case.", + ) + parser.add_argument( + "--m", type=int, required=True, help="Number of rows in the first matrix." + ) + parser.add_argument( + "--n", type=int, required=True, help="Number of columns in the second matrix." + ) + parser.add_argument( + "--k", + type=int, + required=True, + help="Number of columns in the first matrix and number of rows in the second matrix.", + ) + parser.add_argument( + "--tile_size", + type=int, + required=False, + default=128, + help="Tile size for the gemm operation.", + ) + parser.add_argument( + "--group_size", + type=int, + required=False, + default=1, + help="Group size for the group gemm operation.", + ) + parser.add_argument( + "--scale_major_mode", + type=str, + required=False, + default="MN", + choices=["MN", "K"], + help="Scale major mode.", + ) + parser.add_argument( + "--input_dtype", + type=str, + required=False, + default="fp8_e4m3", + help="Data type of the input.", + ) + parser.add_argument( + "--mat2_dtype", + type=str, + required=False, + default="fp8_e4m3", + help="Data type of the mat2.", + ) + parser.add_argument( + "--out_dtype", + type=str, + required=False, + default="bfloat16", + help="Data type of the output.", + ) + parser.add_argument( + "--mma_sm", + type=int, + required=False, + default=1, + choices=[1, 2], + help="How many SMs to use for the MMA operation, must be 1 or 2", + ) + parser.add_argument( + "--backends", + type=str, + required=False, + nargs="+", + default=["cudnn"], + choices=["cudnn", "cublas", "trtllm", "cutlass"], + help="Kernel backends to test. Default: cudnn", + ) + parser.add_argument( + "--use_128x4_sf_layout", + action="store_true", + help="Use 128x4 SF layout for the input and mat2.", + ) + parser.add_argument( + "--use_nvfp4", + action="store_true", + help="In mm_fp4, whether to use nvfp4 quantization or mxfp4 quantization, defaults to False.", + ) + parser.add_argument( + "--autotune", + action="store_true", + default=False, + help=("Enable autotuner warmup for supported routines (mm_fp4 and bmm_fp8)."), + ) + + args = parser.parse_args(line) + if args.verbose >= 1: + print(f"[INFO] {args = }") + return args + + +def to_float8(x, dtype=torch.float8_e4m3fn): + finfo = torch.finfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + scale = finfo.max / amax + x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) + return x_scl_sat.to(dtype), scale.float().reciprocal() + + +def testGemmFp8NtGroupwise(args): + """ + Test gemm_fp8_nt_groupwise API. + + This test: + 1. Generates random input tensors + 2. Quantizes input tensors to FP8 + 3. Runs gemm_fp8_nt_groupwise + 4. Runs reference check + 5. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testGemmFp8NtGroupwise") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + ## Parse input arguments + backends = args.backends + m = args.m + n = args.n + k = args.k + tile_size = args.tile_size + scale_major_mode = args.scale_major_mode + mma_sm = args.mma_sm + is_cuda_graph_compatible = not args.no_cuda_graph + run_refcheck = args.refcheck + res = [] + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + out_dtype = dtype_str_to_torch_dtype(args.out_dtype) + if out_dtype not in [torch.bfloat16, torch.float16]: + raise ValueError(f"Unsupported output dtype: {args.out_dtype}") + ## Done parsing input arguments + + if "trtllm" in backends: + remove_trtllm = False + if scale_major_mode != "MN": + print( + "[INFO] trtllm only supports MN scale_major_mode, removing trtllm from backends" + ) + remove_trtllm = True + if k < 256: + print("[INFO] trtllm only supports k >= 256, removing trtllm from backends") + remove_trtllm = True + if remove_trtllm: + backends.remove("trtllm") + + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return + + ## Prepare input tensors + a_val = torch.randn((m, k), dtype=torch.float, device=device) + b_val = torch.randn((n, k), dtype=torch.float, device=device) / np.sqrt(k) + + if args.verbose >= 2: + print(f"[VVERBOSE] {a_val.shape = }") + print(f"[VVERBOSE] {b_val.shape = }") + + if scale_major_mode == "K": + a_scale_shape = (m, k // tile_size) + b_scale_shape = (n // tile_size, k // tile_size) + else: + a_scale_shape = (k // tile_size, m) + b_scale_shape = (k // tile_size, n // tile_size) + + a_tile_shape = (1, tile_size) + b_tile_shape = (tile_size, tile_size) + + a_fp8, a_scale = quantize_fp8(a_val, a_scale_shape, a_tile_shape, scale_major_mode) + b_fp8, b_scale = quantize_fp8(b_val, b_scale_shape, b_tile_shape, scale_major_mode) + + if args.verbose >= 2: + print(f"[VVERBOSE] {a_fp8.shape = }") + print(f"[VVERBOSE] {b_fp8.shape = }") + print(f"[VVERBOSE] {a_scale.shape = }") + print(f"[VVERBOSE] {b_scale.shape = }") + + a_dequant = dequantize_fp8(a_fp8, a_scale, scale_major_mode) + b_dequant = dequantize_fp8(b_fp8, b_scale, scale_major_mode) + + def run_backend(backend): + if backend in ["cutlass", "trtllm"]: + return flashinfer.gemm.gemm_fp8_nt_groupwise( + a=a_fp8, + b=b_fp8, + a_scale=a_scale, + b_scale=b_scale, + scale_major_mode=scale_major_mode, + out_dtype=out_dtype, + mma_sm=mma_sm, + backend=backend, + ) + else: + raise ValueError(f"Unsupported backend: {backend}") + + has_reference_output = False + if run_refcheck: + reference_output = einsum(a_dequant, b_dequant, "m k, n k -> m n").to(out_dtype) + has_reference_output = True + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + for cur_backend in backends: + if run_refcheck: + outputs[cur_backend] = run_backend(cur_backend).detach() + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=True, # GEMMs are very MMA-heavy, so prefer sleep to reduce throttling. + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 0: + if run_refcheck and has_reference_output: + for i in range(len(tested_backends)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats( + reference_output, tested_outputs[i], rtol=1e-2, atol=1e-2 + ) + if num_different_elements > 0: + print( + f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements" + ) + + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + + problem_flops = 2 * m * n * k + problem_bytes = (m * k + n * k) * torch.float8_e4m3fn.itemsize + ( + m * n + ) * out_dtype.itemsize + tflops = problem_flops / (10**9 * median_time) # in TFLOPs/sec + tb_per_sec = problem_bytes / (10**9 * median_time) # in TB/sec + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["m"] = m + cur_res["n"] = n + cur_res["k"] = k + cur_res["tile_size"] = tile_size + cur_res["scale_major_mode"] = scale_major_mode + cur_res["out_dtype"] = out_dtype + cur_res["mma_sm"] = mma_sm + cur_res["backend"] = backend + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +def testGroupGemmFp8NtGroupwise(args): + """ + Test group_gemm_fp8_nt_groupwise API. + + This test: + 1. Generates random input tensors + 2. Quantizes input tensors to FP8 + 3. Runs group_gemm_fp8_nt_groupwise + 4. Runs reference check + 5. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testGroupGemmFp8NtGroupwise") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + ## Parse input arguments + backends = ["cutlass"] # Cutlass is currently the only supported backend + m = args.m + n = args.n + k = args.k + group_size = args.group_size + tile_size = args.tile_size + scale_major_mode = args.scale_major_mode + mma_sm = args.mma_sm + is_cuda_graph_compatible = not args.no_cuda_graph + run_refcheck = args.refcheck + res = [] + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + out_dtype = dtype_str_to_torch_dtype(args.out_dtype) + if out_dtype not in [torch.bfloat16, torch.float16]: + raise ValueError(f"Unsupported output dtype: {args.out_dtype}") + ## Done parsing input arguments + ## Prepare input tensors + a_val = torch.randn((group_size * m, k), dtype=torch.float, device="cuda") + b_val = torch.randn((group_size, n, k), dtype=torch.float, device="cuda") / np.sqrt( + k + ) + + if args.verbose >= 2: + print(f"[VVERBOSE] {a_val.shape = }") + print(f"[VVERBOSE] {b_val.shape = }") + + if scale_major_mode == "K": + a_scale_shape = (group_size * m, k // tile_size) + b_scale_shape = (group_size, n // tile_size, k // tile_size) + else: + a_scale_shape = (k // tile_size, m * group_size) + b_scale_shape = (group_size, k // tile_size, n // tile_size) + + a_tile_shape = (1, tile_size) + b_tile_shape = (1, tile_size, tile_size) + + a_fp8, a_scale = quantize_fp8(a_val, a_scale_shape, a_tile_shape, scale_major_mode) + b_fp8, b_scale = quantize_fp8(b_val, b_scale_shape, b_tile_shape, scale_major_mode) + + a_dequant = dequantize_fp8(a_fp8, a_scale, scale_major_mode) + b_dequant = dequantize_fp8(b_fp8, b_scale, scale_major_mode) + + m_indptr = torch.arange(0, group_size + 1, dtype=torch.int32, device="cuda") * m + + if args.verbose >= 2: + print(f"[VVERBOSE] {a_fp8.shape = }") + print(f"[VVERBOSE] {b_fp8.shape = }") + print(f"[VVERBOSE] {a_scale.shape = }") + print(f"[VVERBOSE] {b_scale.shape = }") + print(f"[VVERBOSE] {m_indptr.shape = }") + + def run_backend(backend): + if backend == "cutlass": + return flashinfer.gemm.group_gemm_fp8_nt_groupwise( + a=a_fp8, + b=b_fp8, + a_scale=a_scale, + b_scale=b_scale, + m_indptr=m_indptr, + scale_major_mode=scale_major_mode, + out_dtype=out_dtype, + mma_sm=mma_sm, + ) + else: + raise ValueError(f"Unsupported backend: {backend}") + + has_reference_output = False + if run_refcheck: + reference_output = ( + einsum( + a_dequant.view((group_size, m, k)), b_dequant, "b m k, b n k -> b m n" + ) + .view((group_size * m, n)) + .to(out_dtype) + ) + has_reference_output = True + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + for cur_backend in backends: + if run_refcheck: + outputs[cur_backend] = run_backend(cur_backend).detach() + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=True, # GEMMs are very MMA-heavy, so prefer sleep to reduce throttling. + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 0: + if run_refcheck and has_reference_output: + for i in range(len(tested_backends)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats( + reference_output, tested_outputs[i], rtol=1e-2, atol=1e-2 + ) + if num_different_elements > 0: + print( + f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements" + ) + + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + problem_flops = 2 * m * n * k * group_size + problem_bytes = ( + group_size * m * k + group_size * n * k + ) * torch.float8_e4m3fn.itemsize + (group_size * m * n) * out_dtype.itemsize + tflops = problem_flops / (10**9 * median_time) # in TFLOPs/sec + tb_per_sec = problem_bytes / (10**9 * median_time) # in TB/sec + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["m"] = m + cur_res["n"] = n + cur_res["k"] = k + cur_res["group_size"] = group_size + cur_res["tile_size"] = tile_size + cur_res["scale_major_mode"] = scale_major_mode + cur_res["out_dtype"] = out_dtype + cur_res["mma_sm"] = mma_sm + cur_res["backend"] = backend + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +def testBmmFp8(args): + """ + Test bmm_fp8 API. + + This test: + 1. Generates random input tensors + 2. Quantizes input tensors to FP8 + 3. Runs bmm_fp8 + 4. Runs reference check + 5. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testBmmFp8") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + ## Parse input arguments + backends = args.backends + batch_size = args.batch_size + m = args.m + n = args.n + k = args.k + input_dtype = args.input_dtype + mat2_dtype = args.mat2_dtype + res_dtype = args.out_dtype + backends = args.backends + is_cuda_graph_compatible = not args.no_cuda_graph + run_refcheck = args.refcheck + autotune_supported_backends = [ + "cutlass", + ] + res = [] + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + if input_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2]: + raise ValueError( + f"Unsupported input dtype: {input_dtype}. Supported dtypes are fp8_e4m3 and fp8_e5m2." + ) + + mat2_dtype = dtype_str_to_torch_dtype(args.mat2_dtype) + if mat2_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2]: + raise ValueError( + f"Unsupported mat2 dtype: {mat2_dtype}. Supported dtypes are fp8_e4m3 and fp8_e5m2." + ) + + res_dtype = dtype_str_to_torch_dtype(args.out_dtype) + if res_dtype not in [torch.bfloat16, torch.float16]: + raise ValueError( + f"Unsupported res dtype: {res_dtype}. Supported dtypes are bfloat16 and float16." + ) + ## Done parsing input arguments + + if getattr(args, "autotune", False): + backends_to_remove = [] + for cur_backend in backends: + if cur_backend not in autotune_supported_backends: + print(f"[INFO] {cur_backend} backend does not support autotune") + backends_to_remove.append(cur_backend) + for cur_backend in backends_to_remove: + backends.remove(cur_backend) + + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return + + ## Prepare input tensors + input = torch.randn([batch_size, m, k], device=device, dtype=torch.bfloat16) + input_fp8, input_inv_s = to_float8(input, dtype=input_dtype) + + mat2 = torch.randn( + [batch_size, n, k], device=device, dtype=torch.bfloat16 + ).transpose(-2, -1) + mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=mat2_dtype) + + if args.verbose >= 2: + print(f"[VVERBOSE] {input_fp8.shape = }") + print(f"[VVERBOSE] {input_fp8.dtype = }") + print(f"[VVERBOSE] {mat2_fp8.shape = }") + print(f"[VVERBOSE] {mat2_fp8.dtype = }") + print(f"[VVERBOSE] {input_inv_s = }") + print(f"[VVERBOSE] {input_inv_s.dtype = }") + print(f"[VVERBOSE] {mat2_inv_s = }") + print(f"[VVERBOSE] {mat2_inv_s.dtype = }") + + def run_backend(backend): + if backend in ["cudnn", "cublas", "cutlass"]: + return flashinfer.gemm.bmm_fp8( + A=input_fp8, + B=mat2_fp8, + A_scale=input_inv_s, + B_scale=mat2_inv_s, + dtype=res_dtype, + backend=backend, + ) + else: + raise ValueError(f"Unsupported backend: {backend}") + + has_reference_output = False + if run_refcheck: + reference_output = torch.bmm(input, mat2) + has_reference_output = True + + if getattr(args, "autotune", False): + warmup_iters = ( + args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10 + ) + for cur_backend in backends: + if cur_backend in autotune_supported_backends: + if args.verbose >= 1: + print(f"[INFO] Autotune warmup for bmm_fp8: {warmup_iters} iters") + with autotune(True): + for _ in range(warmup_iters): + run_backend(cur_backend) + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + for cur_backend in backends: + if run_refcheck: + outputs[cur_backend] = run_backend(cur_backend).detach() + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=True, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 0: + if run_refcheck and has_reference_output: + if reference_output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: + print( + "[INFO] Reference output is FP8. Converting to float32 for reference check." + ) + reference_output = reference_output.to(torch.float32) + tested_outputs = [output.to(torch.float32) for output in tested_outputs] + for i in range(len(tested_backends)): + cos_sim = F.cosine_similarity( + reference_output.reshape(-1), + tested_outputs[i].reshape(-1), + dim=0, + ) + if cos_sim < 0.99: + print( + f"[ERROR] Output tensor mismatch between backends {tested_backends[0]} and {tested_backends[i]}" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}" + ) + + for backend in backends: + backend_name = backend + ( + "_autotune" + if ( + getattr(args, "autotune", False) + and backend in autotune_supported_backends + ) + else "" + ) + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + problem_flops = 2 * m * n * k * batch_size + problem_bytes = ( + m * k * input_dtype.itemsize + + n * k * mat2_dtype.itemsize + + m * n * res_dtype.itemsize + ) + tflops = problem_flops / (10**9 * median_time) # in TFLOPs/sec + tb_per_sec = problem_bytes / (10**9 * median_time) # in TB/sec + print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["batch_size"] = batch_size + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["m"] = m + cur_res["n"] = n + cur_res["k"] = k + cur_res["input_dtype"] = input_dtype + cur_res["mat2_dtype"] = mat2_dtype + cur_res["out_dtype"] = res_dtype + cur_res["backend"] = backend_name + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +def testMmFp4(args): + """ + Test mm_fp4 API. + + This test: + 1. Generates random input tensors + 2. Quantizes input tensors to FP4 + 3. Runs mm_fp4 + 4. Runs reference check + 5. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testMmFp4") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + ## Parse input arguments + backends = args.backends + m = args.m + n = args.n + k = args.k + res_dtype = args.out_dtype + backends = args.backends + is_cuda_graph_compatible = not args.no_cuda_graph + run_refcheck = args.refcheck + use_128x4_sf_layout = args.use_128x4_sf_layout + use_nvfp4 = args.use_nvfp4 + autotune_supported_backends = ["cutlass", "trtllm"] + res = [] + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + + res_dtype = dtype_str_to_torch_dtype(args.out_dtype) + if res_dtype not in [torch.bfloat16, torch.float16]: + raise ValueError( + f"Unsupported res dtype: {res_dtype}. Supported dtypes are bfloat16 and float16." + ) + ## Done parsing input arguments + + if "trtllm" in backends: + remove_trtllm = False + if res_dtype == torch.float16: + print("[INFO] trtllm backend does not support float16 output") + remove_trtllm = True + if remove_trtllm: + backends.remove("trtllm") + if not use_nvfp4: + print( + "[INFO] trtllm backend does not support mxfp4 quantization (use_nvfp4=False)" + ) + backends.remove("trtllm") + if "cutlass" in backends: + remove_cutlass = False + if not use_128x4_sf_layout: + print("[INFO] cutlass backend does not support use_128x4_sf_layout=False") + remove_cutlass = True + if not use_nvfp4: + print( + "[INFO] cutlass backend does not support mxfp4 quantization (use_nvfp4=False)" + ) + backends.remove("cutlass") + if remove_cutlass: + backends.remove("cutlass") + if "cudnn" in backends: + remove_cudnn = False + if not use_128x4_sf_layout: + print("[INFO] cudnn backend does not support use_128x4_sf_layout=False") + remove_cudnn = True + if remove_cudnn: + backends.remove("cudnn") + if getattr(args, "autotune", False): + backends_to_remove = [] + for cur_backend in backends: + if cur_backend not in autotune_supported_backends: + print(f"[INFO] {cur_backend} backend does not support autotune") + backends_to_remove.append(cur_backend) + for cur_backend in backends_to_remove: + backends.remove(cur_backend) + + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return + + input = torch.randn([m, k], device=device, dtype=torch.bfloat16) + mat2 = torch.randn([n, k], device=device, dtype=torch.bfloat16) + a_sf_layout = ( + flashinfer.SfLayout.layout_128x4 + if use_128x4_sf_layout + else flashinfer.SfLayout.layout_8x4 + ) + + global_sf_input = (448 * 6) / input.float().abs().nan_to_num().max() + global_sf_mat2 = (448 * 6) / mat2.float().abs().nan_to_num().max() + + if use_nvfp4: + input_fp4, input_inv_s = flashinfer.nvfp4_quantize( + input, global_sf_input, sfLayout=a_sf_layout, do_shuffle=False + ) + mat2_fp4, mat2_inv_s = flashinfer.nvfp4_quantize( + mat2, + global_sf_mat2, + sfLayout=flashinfer.SfLayout.layout_128x4, + do_shuffle=False, + ) + else: # mxfp4 + input_fp4, input_inv_s = flashinfer.mxfp4_quantize(input) + mat2_fp4, mat2_inv_s = flashinfer.mxfp4_quantize(mat2) + + if "trtllm" in backends: + mat2_fp4_trtllm, mat2_inv_s_trtllm = flashinfer.nvfp4_quantize( + mat2, + global_sf_mat2, + sfLayout=flashinfer.SfLayout.layout_128x4, + do_shuffle=True, + ) + + if args.verbose >= 2: + print(f"[VVERBOSE] {input_fp4.shape = }") + print(f"[VVERBOSE] {input_fp4.dtype = }") + print(f"[VVERBOSE] {mat2_fp4.shape = }") + print(f"[VVERBOSE] {mat2_fp4.dtype = }") + + alpha = 1.0 / (global_sf_input * global_sf_mat2) if use_nvfp4 else None + # res = torch.empty([m, n], device="cuda", dtype=res_dtype) + + def run_backend(backend): + if backend in ["cudnn", "trtllm", "cutlass"]: + return flashinfer.gemm.mm_fp4( + a=input_fp4, + b=mat2_fp4.T if backend != "trtllm" else mat2_fp4_trtllm.T, + a_descale=input_inv_s, + b_descale=mat2_inv_s.T if backend != "trtllm" else mat2_inv_s_trtllm.T, + alpha=alpha, + out_dtype=res_dtype, + block_size=16 + if use_nvfp4 + else 32, # nvfp4 only supports 16; mxfp4 only supports 32. + use_8x4_sf_layout=not use_128x4_sf_layout, + backend=backend, + use_nvfp4=use_nvfp4, + ) + else: + raise ValueError(f"Unsupported backend: {backend}") + + has_reference_output = False + if run_refcheck: + reference_output = torch.mm(input, mat2.T) + has_reference_output = True + + if getattr(args, "autotune", False): + warmup_iters = ( + args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10 + ) + for cur_backend in backends: + if cur_backend in autotune_supported_backends: + if args.verbose >= 1: + print(f"[INFO] Autotune warmup for mm_fp4: {warmup_iters} iters") + with autotune(True): + for _ in range(warmup_iters): + run_backend(cur_backend) + + # Storage for timing results and outputs + backend_times = {backend: [] for backend in backends} + outputs = {} + for cur_backend in backends: + if run_refcheck: + outputs[cur_backend] = run_backend(cur_backend).detach() + backend_times[cur_backend] = bench_gpu_time( + fn=lambda: run_backend(cur_backend), + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=True, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 0: + if run_refcheck and has_reference_output: + for i in range(len(tested_backends)): + cos_sim = F.cosine_similarity( + reference_output.reshape(-1), + tested_outputs[i].reshape(-1), + dim=0, + ) + if cos_sim < 0.97: + print( + f"[ERROR] Output tensor mismatch between backends {tested_backends[0]} and {tested_backends[i]}" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch with cos_sim={cos_sim}" + ) + + for backend in backends: + backend_name = backend + ( + "_autotune" + if ( + getattr(args, "autotune", False) + and backend in autotune_supported_backends + ) + else "" + ) + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + problem_flops = 2 * m * n * k + problem_bytes = ( + m * k * 0.5 + n * k * 0.5 + m * n * res_dtype.itemsize + ) # 0.5 for fp4 + tflops = problem_flops / (10**9 * median_time) # in TFLOPs/sec + tb_per_sec = problem_bytes / (10**9 * median_time) # in TB/sec + print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["m"] = m + cur_res["n"] = n + cur_res["k"] = k + cur_res["out_dtype"] = res_dtype + cur_res["use_128x4_sf_layout"] = use_128x4_sf_layout + cur_res["backend"] = backend_name + cur_res["use_nvfp4"] = use_nvfp4 + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res diff --git a/benchmarks/routines/moe.py b/benchmarks/routines/moe.py new file mode 100644 index 0000000000..6af3425c73 --- /dev/null +++ b/benchmarks/routines/moe.py @@ -0,0 +1,1596 @@ +from collections import defaultdict +from typing import Optional + +import numpy as np +import torch + +import flashinfer +from flashinfer.autotuner import autotune +from flashinfer.fused_moe import ( + WeightLayout, + trtllm_fp4_block_scale_moe, + trtllm_fp8_block_scale_moe, + trtllm_fp8_per_tensor_scale_moe, + cutlass_fused_moe, + convert_to_block_layout, +) +from flashinfer import fp4_quantize, shuffle_matrix_a +from flashinfer.testing.utils import ( + bench_gpu_time, +) + +from .flashinfer_benchmark_utils import ( + dtype_str_to_torch_dtype, + get_device, + print_perf_metrics, + filter_backends_by_compute_capability, +) + + +def run_moe_test(args): + """ + Run a MOE test. + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.routine == "trtllm_fp4_block_scale_moe": + return testTrtllmFp4BlockScaleMoe(args) + elif args.routine == "trtllm_fp8_block_scale_moe": + return testTrtllmFp8BlockScaleMoe(args) + elif args.routine == "trtllm_fp8_per_tensor_scale_moe": + return testTrtllmFp8PerTensorScaleMoe(args) + elif args.routine == "cutlass_fused_moe": + return testCutlassFusedMoe(args) + else: + raise ValueError(f"Unsupported routine: {args.routine}") + + +def parse_moe_args(line, parser): + """ + Parse command line arguments for MOE test configuration. + + Args: + line: Command line arguments + parser: ArgumentParser object already populated with shared arguments + + Returns: + Parsed argument namespace + """ + parser.add_argument( + "--num_tokens", type=int, required=True, help="Number of input tokens." + ) + parser.add_argument( + "--hidden_size", type=int, required=True, help="Hidden dimension size." + ) + parser.add_argument( + "--intermediate_size", + type=int, + required=True, + help="Intermediate dimension size.", + ) + parser.add_argument( + "--num_experts", type=int, required=True, help="Total number of experts." + ) + parser.add_argument( + "--top_k", + type=int, + required=True, + help="Number of experts to route to per token.", + ) + parser.add_argument( + "--n_group", + type=int, + required=False, + default=None, + help="Number of expert groups (for DeepSeek routing). Only used with DeepSeekV3 routing method.", + ) + parser.add_argument( + "--topk_group", + type=int, + required=False, + default=None, + help="Number of groups to consider for top-k routing. Only used with DeepSeekV3 routing method.", + ) + parser.add_argument( + "--routed_scaling_factor", + type=float, + required=False, + default=2.5, + help="Scaling factor for routing.", + ) + parser.add_argument( + "--local_expert_offset", + type=int, + required=False, + default=0, + help="Offset of local experts in global expert space.", + ) + parser.add_argument( + "--local_num_experts", + type=int, + required=False, + default=None, + help="Number of experts handled by this device. Defaults to num_experts.", + ) + parser.add_argument( + "--tile_tokens_dim", + type=int, + required=False, + default=8, + help="Tile dimension for tokens.", + ) + parser.add_argument( + "--routing_method", + type=str, + required=False, + default="deepseek_v3", + choices=[ + "renormalize", + "deepseek_v3", + "llama4", + "renormalize_naive", + "topk", + ], + help=( + "Routing method: renormalize | deepseek_v3 | llama4 | renormalize_naive | topk." + ), + ) + parser.add_argument( + "--use_shuffled_weight", + action="store_true", + default=False, + help="Whether to use shuffled weight layout.", + ) + parser.add_argument( + "--weight_layout", + type=int, + required=False, + default=0, + choices=[0, 1, 2], + help="Weight layout: 0=MajorK, 1=MajorMn, 2=BlockMajorK.", + ) + parser.add_argument( + "--use_routing_bias", + action="store_true", + default=False, + help="Whether to use routing bias.", + ) + parser.add_argument( + "--use_routing_scales_on_input", + action="store_true", + default=False, + help="Whether to use routing scales on input (for Llama4 routing).", + ) + parser.add_argument( + "--input_dtype", + type=str, + required=False, + default="bfloat16", + help="Data type of the input hidden states.", + ) + parser.add_argument( + "--weight_dtype", + type=str, + required=False, + default="bfloat16", + help="Data type of the weights (before quantization).", + ) + parser.add_argument( + "--gated_act", + type=str, + required=False, + default="swiglu", + choices=["swiglu", "geglu"], + help="Type of gated activation function: swiglu | geglu.", + ) + parser.add_argument( + "--autotune", + action="store_true", + default=False, + help=( + "Enable autotuner warmup for supported routines (trtllm_fp4_block_scale_moe and cutlass_fused_moe)." + ), + ) + + # CUTLASS fused MoE specific + parser.add_argument( + "--cutlass_variant", + type=str, + required=False, + default="base", + choices=["base", "fp8", "nvfp4"], + help="Variant for cutlass_fused_moe benchmark: base (no quant), fp8 (per-tensor), nvfp4 (fp4 blockscale)", + ) + parser.add_argument( + "--quantized_input", + action="store_true", + default=False, + help="Quantize input activations (only used for nvfp4).", + ) + parser.add_argument( + "--tp_size", + type=int, + required=False, + default=1, + help="Tensor parallel size for cutlass_fused_moe.", + ) + parser.add_argument( + "--tp_rank", + type=int, + required=False, + default=0, + help="Tensor parallel rank for cutlass_fused_moe.", + ) + parser.add_argument( + "--ep_size", + type=int, + required=False, + default=1, + help="Expert parallel size for cutlass_fused_moe.", + ) + parser.add_argument( + "--ep_rank", + type=int, + required=False, + default=0, + help="Expert parallel rank for cutlass_fused_moe.", + ) + + args = parser.parse_args(line) + + # Normalize routing method (map string to internal int expected by kernels) + routing_method_name_to_type = { + "renormalize": 1, + "deepseek_v3": 2, + "llama4": 3, + "renormalize_naive": 4, + "topk": 5, + } + args.routing_method_type = routing_method_name_to_type[args.routing_method] + + # Normalize gated act type (map string to internal int expected by kernels) + gated_act_name_to_type = { + "swiglu": 0, + "geglu": 1, + } + args.gated_act_type = gated_act_name_to_type[args.gated_act] + + if args.verbose >= 1: + print(f"[INFO] {args = }") + return args + + +def create_trtllm_moe_test_data( + num_tokens: int, + hidden_size: int, + intermediate_size: int, + num_experts: int, + routing_method_type: int, + use_routing_bias: bool, + input_dtype: torch.dtype, + weight_dtype: torch.dtype, + device: torch.device, + moe_kernel_type: str = "fp8_per_tensor", +): + """ + Create test data for TensorRT-LLM fused MoE benchmarking (trtllm_*_moe APIs). + + This helper prepares inputs for the TensorRT-LLM fused MoE kernels exposed via + flashinfer.fused_moe (e.g., trtllm_fp4_block_scale_moe, trtllm_fp8_block_scale_moe, + trtllm_fp8_per_tensor_scale_moe). It is NOT used for CUTLASS MoE benchmarks, + which construct their own inputs specific to the CUTLASS path. + + Returns: + Tuple of tensors needed for trtllm fused MoE computation + """ + # Create routing logits - dtype depends on both routing method AND MOE kernel type + # Different MOE kernels have different routing_logits dtype requirements: + + if moe_kernel_type == "fp8_block_scale": + # FP8 block scale MOE always expects float32 routing logits (line 333 in kernel_launcher.cu) + routing_logits = torch.randn( + (num_tokens, num_experts), device=device, dtype=torch.float32 + ) + elif moe_kernel_type == "fp8_per_tensor": + # FP8 per-tensor MOE dtype depends on use_routing_scales_on_input parameter + # For Llama4: use_routing_scales_on_input=True -> bfloat16 + # For others: use_routing_scales_on_input=False -> float32 + if routing_method_type == 3: # Llama4 uses routing scales on input + routing_logits = torch.randn( + (num_tokens, num_experts), device=device, dtype=torch.bfloat16 + ) + else: + routing_logits = torch.randn( + (num_tokens, num_experts), device=device, dtype=torch.float32 + ) + elif moe_kernel_type == "fp4_block_scale": + # FP4 block scale MOE follows the test pattern: float32 for DeepSeekV3, bfloat16 for others + if routing_method_type == 2: # DeepSeekV3 - uses float32 + routing_logits = torch.randn( + (num_tokens, num_experts), device=device, dtype=torch.float32 + ) + else: # All other routing methods (Renormalize, RenormalizeNaive, Llama4) - use bfloat16 + routing_logits = torch.randn( + (num_tokens, num_experts), device=device, dtype=torch.bfloat16 + ) + else: + raise ValueError(f"Unknown MOE kernel type: {moe_kernel_type}") + + # Create routing bias if needed - always bfloat16 + routing_bias = None + if use_routing_bias: + routing_bias = torch.randn(num_experts, device=device, dtype=torch.bfloat16) + + # Create hidden states - always start with bfloat16 for proper quantization + hidden_states = 2 * torch.randn( + (num_tokens, hidden_size), device=device, dtype=torch.bfloat16 + ) + + # Create weights - always start with bfloat16 for proper quantization + gemm1_weights = torch.randn( + (num_experts, 2 * intermediate_size, hidden_size), + device=device, + dtype=torch.bfloat16, + ) + gemm2_weights = torch.randn( + (num_experts, hidden_size, intermediate_size), + device=device, + dtype=torch.bfloat16, + ) + + return routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights + + +def calculate_fp4_global_scale_factor(tensor): + """Calculate global scale factor for FP4 quantization.""" + # Calculate as a tensor on the same device + # Using the same formula as in test files: FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax + tensor_amax = tensor.abs().max().to(torch.float32) + # FLOAT8_E4M3_MAX = 448, FLOAT4_E2M1_MAX = 6 + global_scale = (448.0 * 6.0) / tensor_amax + return global_scale + + +def quant_fp4_simple(a, a_global_sf, use_ue8m0=False, is_sf_swizzled_layout=True): + """ + Simplified FP4 quantization for benchmarking. + In production, use the actual fp4_quantize function. + """ + sf_vec_size = 16 + + # Use the actual fp4_quantize function from flashinfer + a_fp4, a_sf = fp4_quantize( + a, a_global_sf, sf_vec_size, use_ue8m0, is_sf_swizzled_layout + ) + + return a_fp4, a_sf, a_global_sf + + +def quant_fp4_batches_simple( + a, num_experts, use_ue8m0=False, is_sf_swizzled_layout=True +): + """Simplified FP4 batch quantization for benchmarking.""" + quant_a = [] + sfs = [] + global_sfs = [] + for i in range(num_experts): + # Calculate global scale factor (returns tensor) + a_global_sf = calculate_fp4_global_scale_factor(a[i]) + a_fp4, a_sf, _ = quant_fp4_simple( + a[i], a_global_sf, use_ue8m0, is_sf_swizzled_layout + ) + quant_a.append(a_fp4) + sfs.append(a_sf) + global_sfs.append(a_global_sf) + + result_quant_a = torch.stack(quant_a) + result_sfs = torch.stack(sfs) + result_global_sfs = torch.stack(global_sfs) + + return result_quant_a, result_sfs, result_global_sfs + + +def calculate_moe_tflops( + num_tokens: int, + hidden_size: int, + intermediate_size: int, + num_experts: int, + top_k: int, + time_ms: float, +) -> float: + """ + Calculate TFLOPS for MOE operation. + + MOE computation involves: + 1. First GEMM: [num_tokens, hidden_size] x [num_experts, hidden_size, 2*intermediate_size] + 2. Activation function (SwiGLU gate) + 3. Second GEMM: [num_tokens, intermediate_size] x [num_experts, intermediate_size, hidden_size] + + For each token, we only compute for top_k experts. + + """ + # FLOPS per token per expert (base calculation) + flops_per_token_per_expert = ( + 2 * hidden_size * 2 * intermediate_size # First GEMM + + 2 * intermediate_size * hidden_size # Second GEMM + ) + + total_flops = num_tokens * top_k * flops_per_token_per_expert + tflops = total_flops / (time_ms * 1e-3) / 1e12 # Convert to TFLOPS + return tflops + + +def calculate_moe_bandwidth( + num_tokens: int, + hidden_size: int, + intermediate_size: int, + num_experts: int, + top_k: int, + time_ms: float, + input_dtype: torch.dtype, + weight_dtype: torch.dtype, + input_format: Optional[str] = None, + weight_format: Optional[str] = None, + routing_logits_dtype: Optional[torch.dtype] = torch.float32, + active_experts: Optional[int] = None, +) -> float: + """ + Calculate memory bandwidth for MOE operation in TB/sec. + + Args: + input_format: Override for input representation ("fp8" or "fp4"); None uses dtype.itemsize + weight_format: Override for weight representation ("fp8" or "fp4"); None uses dtype.itemsize + routing_logits_dtype: Dtype for routing logits memory accounting (default float32) + """ + + # Get effective byte sizes + def get_effective_bytes(dtype: torch.dtype, fmt: Optional[str]) -> float: + if fmt == "fp4": + return 0.5 + if fmt == "fp8": + return 1.0 + return dtype.itemsize + + input_bytes_per_element = get_effective_bytes(input_dtype, input_format) + weight_bytes_per_element = get_effective_bytes(weight_dtype, weight_format) + + # Input memory: hidden states + routing logits + # Note: routing logits dtype depends on kernel; pass in when known, default float32; None means excluded + routing_logits_bytes = ( + 0 if routing_logits_dtype is None else routing_logits_dtype.itemsize + ) + input_bytes = ( + # Count hidden states once; kernels typically reuse inputs for multiple experts + num_tokens * hidden_size * input_bytes_per_element + + num_tokens * num_experts * routing_logits_bytes + ) + + # Weight memory (reuse weights across tokens by grouping tokens per expert) + # Assume each active expert's weights are read once per run. + weight_bytes_per_expert = ( + 2 * intermediate_size * hidden_size * weight_bytes_per_element # gemm1 + + hidden_size * intermediate_size * weight_bytes_per_element # gemm2 + ) + if active_experts is not None: + num_active_experts = active_experts + else: + num_active_experts = min(num_experts, top_k * num_tokens) + weight_bytes = num_active_experts * weight_bytes_per_expert + + # Output memory (typically full precision) + output_bytes = num_tokens * hidden_size * input_dtype.itemsize + + total_bytes = input_bytes + weight_bytes + output_bytes + tb_per_sec = total_bytes / (time_ms * 1e-3) / 1e12 # Convert to TB/sec + return tb_per_sec + + +def _compute_routing(router_logits: torch.Tensor, top_k: int): + routing_weights = torch.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.float() + return routing_weights, selected_experts + + +def _dynamic_per_tensor_fp8_quant(x: torch.Tensor): + fp8_max = torch.finfo(torch.float8_e4m3fn).max + x_max = x.abs().max().float().clamp(min=1e-6) + scale = x_max / fp8_max + inv_scale = 1.0 / scale + out = (x.float() * inv_scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + return out, scale.view((1,)) + + +def testTrtllmFp4BlockScaleMoe(args): + """ + Test trtllm_fp4_block_scale_moe API (TensorRT-LLM fused MoE). + + This test: + 1. Creates quantized FP4 weights and scales + 2. Runs FP4 block scale MOE + 3. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testTrtllmFp4BlockScaleMoe") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype) + + # Parse configuration + num_tokens = args.num_tokens + hidden_size = args.hidden_size + intermediate_size = args.intermediate_size + num_experts = args.num_experts + top_k = args.top_k + n_group = ( + args.n_group + if hasattr(args, "n_group") and args.n_group is not None and args.n_group > 0 + else None + ) + topk_group = ( + args.topk_group + if hasattr(args, "topk_group") + and args.topk_group is not None + and args.topk_group > 0 + else None + ) + routed_scaling_factor = ( + args.routed_scaling_factor + if hasattr(args, "routed_scaling_factor") + and args.routed_scaling_factor is not None + else None + ) + local_expert_offset = args.local_expert_offset + local_num_experts = args.local_num_experts or num_experts + tile_tokens_dim = args.tile_tokens_dim + routing_method_type = args.routing_method_type + use_shuffled_weight = args.use_shuffled_weight + weight_layout = args.weight_layout + is_cuda_graph_compatible = not args.no_cuda_graph + gated_act_type = args.gated_act_type + res = [] + + backends = ["trtllm"] + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + if args.verbose >= 1: + print( + f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, " + f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}" + ) + + # Create test data + routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights = ( + create_trtllm_moe_test_data( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + routing_method_type, + args.use_routing_bias, + input_dtype, + weight_dtype, + device, + moe_kernel_type="fp4_block_scale", + ) + ) + + # For FP4, we need to properly quantize weights and create scales + use_ue8m0 = False + + # Calculate global scale factor for hidden states + hidden_states_scale_global = calculate_fp4_global_scale_factor(hidden_states) + + # Quantize weights using proper FP4 quantization + gemm1_weights_fp4_bytes, gemm1_scales_fp4_bytes, gemm1_scales_global = ( + quant_fp4_batches_simple(gemm1_weights, num_experts, use_ue8m0, True) + ) + gemm2_weights_fp4_bytes, gemm2_scales_fp4_bytes, gemm2_scales_global = ( + quant_fp4_batches_simple(gemm2_weights, num_experts, use_ue8m0, True) + ) + + # Quantize hidden states + hidden_states_fp4_bytes, hidden_states_scale_fp4_bytes, _ = quant_fp4_simple( + hidden_states, hidden_states_scale_global, use_ue8m0, True + ) + + # Reshape hidden states for the kernel (pack 2 FP4 values into 1 byte) + # Keep as uint8 format for FP4 packed data + hidden_states_fp4 = hidden_states_fp4_bytes.view(torch.uint8).reshape( + hidden_states.shape[0], hidden_states.shape[1] // 2 + ) + # Hidden-states scale for FP4 must be 2D: [num_tokens, hidden_size // 16] + hidden_states_scale_linear_fp4 = hidden_states_scale_fp4_bytes.view( + torch.float8_e4m3fn + ) + # Ensure expected shape (16 elements per hidden value for NvFP4) + expected_scale_elems = (num_tokens * hidden_size) // 16 + if hidden_states_scale_linear_fp4.numel() != expected_scale_elems: + if args.verbose >= 1: + print( + f"[INFO] Adjusting FP4 hidden_states_scale from {hidden_states_scale_linear_fp4.numel()} to {expected_scale_elems} elements" + ) + hidden_states_scale_linear_fp4 = torch.ones( + expected_scale_elems, device=device, dtype=torch.float8_e4m3fn + ) + hidden_states_scale_linear_fp4 = hidden_states_scale_linear_fp4.reshape( + num_tokens, hidden_size // 16 + ) + + # Prepare weights for kernel + # For FP4 weights, keep them as uint8 (packed format) - don't convert to float8_e4m3fn + gemm1_weights_fp4 = gemm1_weights_fp4_bytes.view(torch.uint8).reshape( + num_experts, 2 * intermediate_size, hidden_size // 2 + ) + # Scale factors should be viewed as float8_e4m3fn + gemm1_weights_scale = gemm1_scales_fp4_bytes.view(torch.float8_e4m3fn).reshape( + num_experts, 2 * intermediate_size, hidden_size // 16 + ) + + gemm2_weights_fp4 = gemm2_weights_fp4_bytes.view(torch.uint8).reshape( + num_experts, hidden_size, intermediate_size // 2 + ) + gemm2_weights_scale = gemm2_scales_fp4_bytes.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // 16 + ) + + # Optional parameters for FP4 (using None for simplicity in benchmarking) + gemm1_bias = None + gemm1_alpha = None + gemm1_beta = None + gemm1_clamp_limit = None + gemm2_bias = None + + # Create scale scalars (simplified - in practice these would be computed) + output1_scale_scalar = torch.ones( + local_num_experts, device=device, dtype=torch.float32 + ) + output1_scale_gate_scalar = torch.ones( + local_num_experts, device=device, dtype=torch.float32 + ) + output2_scale_scalar = torch.ones( + local_num_experts, device=device, dtype=torch.float32 + ) + + if args.verbose >= 2: + print(f"[VVERBOSE] routing_logits.shape = {routing_logits.shape}") + print(f"[VVERBOSE] hidden_states.shape = {hidden_states.shape}") + print(f"[VVERBOSE] gemm1_weights_fp4.shape = {gemm1_weights_fp4.shape}") + print(f"[VVERBOSE] gemm2_weights_fp4.shape = {gemm2_weights_fp4.shape}") + + def run_fp4_moe(): + return trtllm_fp4_block_scale_moe( + routing_logits=routing_logits, + routing_bias=routing_bias, + hidden_states=hidden_states_fp4, + hidden_states_scale=hidden_states_scale_linear_fp4, + gemm1_weights=gemm1_weights_fp4, + gemm1_weights_scale=gemm1_weights_scale, + gemm1_bias=gemm1_bias, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm2_weights=gemm2_weights_fp4, + gemm2_weights_scale=gemm2_weights_scale, + gemm2_bias=gemm2_bias, + output1_scale_scalar=output1_scale_scalar, + output1_scale_gate_scalar=output1_scale_gate_scalar, + output2_scale_scalar=output2_scale_scalar, + num_experts=num_experts, + top_k=top_k, + n_group=n_group, + topk_group=topk_group, + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=routed_scaling_factor, + tile_tokens_dim=tile_tokens_dim, + routing_method_type=routing_method_type, + gated_act_type=gated_act_type, + do_finalize=True, + ) + + backend = "trtllm" + + # Optional autotune warmup (supported for FP4 TRTLlm fused MoE) + if getattr(args, "autotune", False): + warmup_iters = ( + args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10 + ) + backend = "trtllm_autotune" + if args.verbose >= 1: + print( + f"[INFO] Autotune warmup for FP4 block scale MoE: {warmup_iters} iters" + ) + with autotune(True): + for _ in range(warmup_iters): + run_fp4_moe() + + # Benchmark timing + times = bench_gpu_time( + fn=run_fp4_moe, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + # Compute performance metrics + median_time = np.median(times) + std_time = np.std(times) + tflops = calculate_moe_tflops( + num_tokens, hidden_size, intermediate_size, num_experts, top_k, median_time + ) + tb_per_sec = calculate_moe_bandwidth( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + median_time, + input_dtype, + weight_dtype, + input_format="fp4", + weight_format="fp4", + routing_logits_dtype=routing_logits.dtype, + ) + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["num_tokens"] = num_tokens + cur_res["hidden_size"] = hidden_size + cur_res["intermediate_size"] = intermediate_size + cur_res["num_experts"] = num_experts + cur_res["top_k"] = top_k + cur_res["n_group"] = n_group + cur_res["topk_group"] = topk_group + cur_res["routed_scaling_factor"] = routed_scaling_factor + cur_res["local_expert_offset"] = local_expert_offset + cur_res["local_num_experts"] = local_num_experts + cur_res["tile_tokens_dim"] = tile_tokens_dim + cur_res["routing_method"] = args.routing_method + cur_res["use_shuffled_weight"] = use_shuffled_weight + cur_res["weight_layout"] = weight_layout + cur_res["use_routing_bias"] = args.use_routing_bias + cur_res["use_routing_scales_on_input"] = args.use_routing_scales_on_input + cur_res["input_dtype"] = input_dtype + cur_res["weight_dtype"] = weight_dtype + cur_res["gated_act"] = args.gated_act + res.append(cur_res) + + return res + + +def testCutlassFusedMoe(args): + """ + Benchmark cutlass_fused_moe (CUTLASS MoE) with variants mirroring tests in tests/moe/test_trtllm_cutlass_fused_moe.py + Variants: + - base: no quantization + - fp8: per-tensor fp8 for weights and activation scale + - nvfp4: FP4 block-scale weights, optional quantized input + Supports TP/EP via tp_size/tp_rank and ep_size/ep_rank. + """ + if args.verbose >= 1: + print("[INFO] Running testCutlassFusedMoe") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + + # Shapes + num_tokens = args.num_tokens + hidden_size = args.hidden_size + intermediate_size = args.intermediate_size + num_experts = args.num_experts + top_k = args.top_k + tp_size = getattr(args, "tp_size", 1) + tp_rank = getattr(args, "tp_rank", 0) + ep_size = getattr(args, "ep_size", 1) + ep_rank = getattr(args, "ep_rank", 0) + is_cuda_graph_compatible = not args.no_cuda_graph + res = [] + backends = ["cutlass"] + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + # Create base tensors + torch.manual_seed(args.random_seed) + x = torch.randn(num_tokens, hidden_size, dtype=input_dtype, device=device) + w31_weight = ( + torch.randn( + num_experts, + 2 * intermediate_size, + hidden_size, + dtype=input_dtype, + device=device, + ) + / 10 + ) + w2_weight = ( + torch.randn( + num_experts, + hidden_size, + intermediate_size, + dtype=input_dtype, + device=device, + ) + / 10 + ) + + # Routing + router_logits = torch.randn( + num_tokens, num_experts, dtype=input_dtype, device=device + ) + routing_weights, selected_experts = _compute_routing(router_logits, top_k) + + if args.verbose >= 2: + print(f"[VVERBOSE] x.shape = {x.shape}") + print(f"[VVERBOSE] w31_weight.shape = {w31_weight.shape}") + print(f"[VVERBOSE] w2_weight.shape = {w2_weight.shape}") + + # Build local weights per EP/TP like tests do + experts_per_rank = num_experts // max(ep_size, 1) + expert_start = ep_rank * experts_per_rank + expert_end = expert_start + experts_per_rank + w31_ep = w31_weight[expert_start:expert_end, :] + w2_ep = w2_weight[expert_start:expert_end, :] + + def build_tp_shards(w31_ep_tensor: torch.Tensor, w2_ep_tensor: torch.Tensor): + if tp_size <= 1: + return w31_ep_tensor, w2_ep_tensor + # Split w31 into w3 and w1 along intermediate dim + w3_weight, w1_weight = torch.chunk(w31_ep_tensor, 2, dim=1) + shard = intermediate_size // tp_size + start = tp_rank * shard + end = start + shard + w3_local = w3_weight[:, start:end, :] + w1_local = w1_weight[:, start:end, :] + w31_local = torch.cat([w3_local, w1_local], dim=1) + w2_local = w2_ep_tensor[:, :, start:end] + return w31_local.contiguous(), w2_local.contiguous() + + w31_local, w2_local = build_tp_shards(w31_ep, w2_ep) + + # Prepare variant-specific inputs (outside of the timed/captured region) + variant = getattr(args, "cutlass_variant", "base") + out = torch.empty_like(x) + + if variant == "base": + + def run_cutlass(): + return cutlass_fused_moe( + x, + selected_experts.to(torch.int), + routing_weights, + w31_local, + w2_local, + input_dtype, + tp_size=tp_size, + tp_rank=tp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + quant_scales=None, + output=out, + ) + + elif variant == "fp8": + # Per-tensor FP8 for weights and activation scale + w31_weight_fp8 = torch.empty_like(w31_local, dtype=torch.float8_e4m3fn) + w2_weight_fp8 = torch.empty_like(w2_local, dtype=torch.float8_e4m3fn) + local_num_experts = w31_local.shape[0] + w31_scales = torch.empty(local_num_experts, 2, dtype=input_dtype, device=device) + w2_scales = torch.empty(local_num_experts, 1, dtype=input_dtype, device=device) + + # Quantize weights per expert + for expert_id in range(local_num_experts): + w31_expert = w31_local[expert_id] + w2_expert = w2_local[expert_id] + w31_q, s31 = _dynamic_per_tensor_fp8_quant(w31_expert) + w2_q, s2 = _dynamic_per_tensor_fp8_quant(w2_expert) + w31_weight_fp8[expert_id].copy_(w31_q) + w2_weight_fp8[expert_id].copy_(w2_q) + # Store the same scalar twice to mimic test layout (avoid torch.tensor()) + w31_scales[expert_id, 0] = s31.to(dtype=input_dtype, device=device) + w31_scales[expert_id, 1] = s31.to(dtype=input_dtype, device=device) + w2_scales[expert_id, 0] = s2.to(dtype=input_dtype, device=device) + + x_quant, hidden_states_scale = _dynamic_per_tensor_fp8_quant(x) + hidden_states_scale_scalar = hidden_states_scale[0].to(device) + + # Note: follow tests quant_scales format + # [w1_scales * hidden_states_scale, 1.0, 1.0 * w2_scales, hidden_states_scale] + w1_scales = w31_scales[:, 1] + one_const = torch.ones((), device=device) + quant_scales = [ + (w1_scales * hidden_states_scale_scalar).float().squeeze(), + one_const, + w2_scales.squeeze().float(), + hidden_states_scale_scalar, + ] + + def run_cutlass(): + return cutlass_fused_moe( + x_quant, + selected_experts.to(torch.int), + routing_weights, + w31_weight_fp8, + w2_weight_fp8, + input_dtype, + tp_size=tp_size, + tp_rank=tp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + quant_scales=quant_scales, + output=out, + ) + + elif variant == "nvfp4": + # NVFP4: FP4 block-scale weights, optional quantized input + FLOAT4_E2M1_MAX = 6.0 + FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + + def round_up(x_val, y): + return (x_val + y - 1) // y * y + + e = w31_local.shape[0] + n = w2_local.shape[2] # local intermediate size after TP + k = hidden_size + quant_blocksize = 16 + + # Weight quantization buffers + w1_q = torch.empty((e, 2 * n, k // 2), device=device, dtype=torch.uint8) + w2_q = torch.empty((e, k, n // 2), device=device, dtype=torch.uint8) + w1_blockscale = torch.empty( + (e, round_up(2 * n, 128), round_up(k // quant_blocksize, 4)), + device=device, + dtype=torch.float8_e4m3fn, + ) + w2_blockscale = torch.empty( + (e, round_up(k, 128), round_up(n // quant_blocksize, 4)), + device=device, + dtype=torch.float8_e4m3fn, + ) + w1_gs = torch.empty((e,), device=device, dtype=torch.float32) + w2_gs = torch.empty((e,), device=device, dtype=torch.float32) + + # Quantize from local shards + for expert in range(e): + w1_src = w31_local[expert] + # w31 layout is [2n, k]; w2 layout is [k, n] + w2_src = w2_local[expert].contiguous() # [hidden_size, n] + w1_amax = torch.abs(w1_src).max().to(torch.float32) + w2_amax = torch.abs(w2_src).max().to(torch.float32) + w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax + w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax + w1_q[expert], w1_blockscale[expert] = fp4_quantize(w1_src, w1_gs[expert]) + w2_q[expert], w2_blockscale[expert] = fp4_quantize(w2_src, w2_gs[expert]) + + a1_gs = torch.ones((), device=device, dtype=torch.float32) + a2_gs = torch.ones((), device=device, dtype=torch.float32) + + hidden_states = x + input_sf = None + if getattr(args, "quantized_input", False): + hidden_states, input_sf = fp4_quantize(x, a1_gs) + + quant_scales = [ + a1_gs, + w1_blockscale.view(torch.int32), + 1.0 / (a1_gs * w1_gs), + a2_gs, + w2_blockscale.view(torch.int32), + 1.0 / (a2_gs * w2_gs), + ] + + def run_cutlass(): + return cutlass_fused_moe( + hidden_states, + selected_experts.to(torch.int), + routing_weights, + w1_q.contiguous().view(torch.long), + w2_q.contiguous().view(torch.long), + input_dtype, + tp_size=tp_size, + tp_rank=tp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + quant_scales=quant_scales, + input_sf=input_sf, + output=out, + ) + else: + raise ValueError(f"Unknown cutlass_variant: {variant}") + + backend = "cutlass" + + # Optional autotune warmup (supported for CUTLASS fused MoE) + if getattr(args, "autotune", False): + warmup_iters = ( + args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10 + ) + backend = "cutlass_autotune" + if args.verbose >= 1: + print(f"[INFO] Autotune warmup for CUTLASS fused MoE: {warmup_iters} iters") + with autotune(True): + for _ in range(warmup_iters): + run_cutlass() + + # Measure + times = bench_gpu_time( + fn=run_cutlass, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + median_time = np.median(times) + std_time = np.std(times) + tflops = calculate_moe_tflops( + num_tokens, hidden_size, intermediate_size, num_experts, top_k, median_time + ) + tb_per_sec = calculate_moe_bandwidth( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + median_time, + input_dtype, + input_dtype, + input_format=( + "fp8" + if variant == "fp8" + else ( + "fp4" + if (variant == "nvfp4" and getattr(args, "quantized_input", False)) + else None + ) + ), + weight_format=( + "fp8" if variant == "fp8" else ("fp4" if variant == "nvfp4" else None) + ), + routing_logits_dtype=router_logits.dtype, + active_experts=int(selected_experts.unique().numel()), + ) + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["num_tokens"] = num_tokens + cur_res["hidden_size"] = hidden_size + cur_res["intermediate_size"] = intermediate_size + cur_res["num_experts"] = num_experts + cur_res["top_k"] = top_k + # Routing method/weight layout not applicable; leave defaults + cur_res["use_shuffled_weight"] = False + cur_res["weight_layout"] = 0 + cur_res["use_routing_scales_on_input"] = False + cur_res["input_dtype"] = input_dtype + cur_res["weight_dtype"] = input_dtype + # CUTLASS fused MoE specific + cur_res["cutlass_variant"] = variant + cur_res["quantized_input"] = args.quantized_input + cur_res["tp_size"] = tp_size + cur_res["tp_rank"] = tp_rank + cur_res["ep_size"] = ep_size + cur_res["ep_rank"] = ep_rank + res.append(cur_res) + + return res + + +def testTrtllmFp8BlockScaleMoe(args): + """ + Test trtllm_fp8_block_scale_moe API (TensorRT-LLM fused MoE). + + This test: + 1. Creates quantized FP8 weights and block scales + 2. Runs FP8 block scale MOE + 3. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testTrtllmFp8BlockScaleMoe") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype) + + # Parse configuration + num_tokens = args.num_tokens + hidden_size = args.hidden_size + intermediate_size = args.intermediate_size + num_experts = args.num_experts + top_k = args.top_k + n_group = ( + args.n_group + if hasattr(args, "n_group") and args.n_group is not None and args.n_group > 0 + else None + ) + topk_group = ( + args.topk_group + if hasattr(args, "topk_group") + and args.topk_group is not None + and args.topk_group > 0 + else None + ) + routed_scaling_factor = ( + args.routed_scaling_factor + if hasattr(args, "routed_scaling_factor") + and args.routed_scaling_factor is not None + else None + ) + local_expert_offset = args.local_expert_offset + local_num_experts = args.local_num_experts or num_experts + tile_tokens_dim = args.tile_tokens_dim + routing_method_type = args.routing_method_type + use_shuffled_weight = args.use_shuffled_weight + weight_layout = args.weight_layout + is_cuda_graph_compatible = not args.no_cuda_graph + res = [] + backends = ["trtllm"] + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + if args.verbose >= 1: + print( + f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, " + f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}" + ) + + # Create test data + routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights = ( + create_trtllm_moe_test_data( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + routing_method_type, + args.use_routing_bias, + input_dtype, + weight_dtype, + device, + moe_kernel_type="fp8_block_scale", + ) + ) + + # For FP8 block scale, create quantized weights and block scales + # Quantize to FP8 + gemm1_weights_fp8 = gemm1_weights.to(torch.float8_e4m3fn) + gemm2_weights_fp8 = gemm2_weights.to(torch.float8_e4m3fn) + + # Optionally shuffle weights and convert to BlockMajorK layout to match kernel expectation + if use_shuffled_weight: + # This tile size follows test implementations + epilogue_tile_m = 64 + + gemm1_weights_fp8_shuffled = [] + gemm2_weights_fp8_shuffled = [] + for i in range(num_experts): + tmp_w1 = shuffle_matrix_a( + gemm1_weights_fp8[i].view(torch.uint8), epilogue_tile_m + ) + tmp_w2 = shuffle_matrix_a( + gemm2_weights_fp8[i].view(torch.uint8), epilogue_tile_m + ) + if weight_layout == WeightLayout.BlockMajorK: + block_k = 128 + tmp_w1 = convert_to_block_layout(tmp_w1, block_k) + tmp_w2 = convert_to_block_layout(tmp_w2, block_k) + gemm1_weights_fp8_shuffled.append(tmp_w1) + gemm2_weights_fp8_shuffled.append(tmp_w2) + + kernel_gemm1_weights = torch.stack(gemm1_weights_fp8_shuffled).view( + torch.float8_e4m3fn + ) + kernel_gemm2_weights = torch.stack(gemm2_weights_fp8_shuffled).view( + torch.float8_e4m3fn + ) + else: + kernel_gemm1_weights = gemm1_weights_fp8 + kernel_gemm2_weights = gemm2_weights_fp8 + + # Create block scale tensors for hidden states and weights (use float32 for scales) + # TensorRT-LLM FP8 block-scale expects hidden_states_scale shape [hidden_size // 128, num_tokens] + hidden_states_scale = 2.0 * torch.ones( + (hidden_size // 128, num_tokens), device=device, dtype=torch.float32 + ) + gemm1_weights_scale = 2.0 * torch.ones( + (num_experts, 2 * intermediate_size // 128, hidden_size // 128), + device=device, + dtype=torch.float32, + ) + gemm2_weights_scale = 2.0 * torch.ones( + (num_experts, hidden_size // 128, intermediate_size // 128), + device=device, + dtype=torch.float32, + ) + + if args.verbose >= 2: + print(f"[VVERBOSE] routing_logits.shape = {routing_logits.shape}") + print(f"[VVERBOSE] hidden_states.shape = {hidden_states.shape}") + print(f"[VVERBOSE] gemm1_weights_fp8.shape = {gemm1_weights_fp8.shape}") + print(f"[VVERBOSE] gemm2_weights_fp8.shape = {gemm2_weights_fp8.shape}") + + # Match test heuristic for tile_tokens_dim when using BlockMajorK + if use_shuffled_weight and weight_layout == WeightLayout.BlockMajorK: + + def _next_pow2(x: int) -> int: + x = max(1, x) + x -= 1 + x |= x >> 1 + x |= x >> 2 + x |= x >> 4 + x |= x >> 8 + x |= x >> 16 + return x + 1 + + tokens_per_expert = max(1, (num_tokens * top_k) // max(local_num_experts, 1)) + suggested_tile = min(max(_next_pow2(tokens_per_expert), 8), 64) + if suggested_tile != tile_tokens_dim and args.verbose >= 1: + print( + f"[INFO] Overriding tile_tokens_dim {tile_tokens_dim} -> {suggested_tile} for BlockMajorK" + ) + tile_tokens_dim = suggested_tile + + def run_fp8_block_moe(): + # Quantize hidden states to FP8 for block scale MOE + hidden_states_fp8 = hidden_states.to(torch.float8_e4m3fn) + # Note: FP8 block scale MOE expects int64_t for n_group/topk_group, not Optional[int64_t] + # So we convert None to 0 to indicate "no groups" mode + return trtllm_fp8_block_scale_moe( + routing_logits=routing_logits, + routing_bias=routing_bias, + hidden_states=hidden_states_fp8, + hidden_states_scale=hidden_states_scale, + gemm1_weights=kernel_gemm1_weights, + gemm1_weights_scale=gemm1_weights_scale, + gemm2_weights=kernel_gemm2_weights, + gemm2_weights_scale=gemm2_weights_scale, + num_experts=num_experts, + top_k=top_k, + n_group=n_group if n_group is not None else 0, + topk_group=topk_group if topk_group is not None else 0, + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=routed_scaling_factor, + tile_tokens_dim=tile_tokens_dim, + routing_method_type=routing_method_type, + use_shuffled_weight=use_shuffled_weight, + weight_layout=weight_layout, + enable_pdl=True, + ) + + # Benchmark timing + times = bench_gpu_time( + fn=run_fp8_block_moe, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + # Compute performance metrics + median_time = np.median(times) + std_time = np.std(times) + tflops = calculate_moe_tflops( + num_tokens, hidden_size, intermediate_size, num_experts, top_k, median_time + ) + tb_per_sec = calculate_moe_bandwidth( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + median_time, + input_dtype, + weight_dtype, + input_format="fp8", + weight_format="fp8", + routing_logits_dtype=routing_logits.dtype, + ) + + backend = "trtllm" + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["num_tokens"] = num_tokens + cur_res["hidden_size"] = hidden_size + cur_res["intermediate_size"] = intermediate_size + cur_res["num_experts"] = num_experts + cur_res["top_k"] = top_k + cur_res["n_group"] = n_group + cur_res["topk_group"] = topk_group + cur_res["routed_scaling_factor"] = routed_scaling_factor + cur_res["local_expert_offset"] = local_expert_offset + cur_res["local_num_experts"] = local_num_experts + cur_res["tile_tokens_dim"] = tile_tokens_dim + cur_res["routing_method"] = args.routing_method + cur_res["use_shuffled_weight"] = use_shuffled_weight + cur_res["weight_layout"] = weight_layout + cur_res["use_routing_bias"] = args.use_routing_bias + cur_res["use_routing_scales_on_input"] = args.use_routing_scales_on_input + cur_res["input_dtype"] = input_dtype + cur_res["weight_dtype"] = weight_dtype + res.append(cur_res) + + return res + + +def testTrtllmFp8PerTensorScaleMoe(args): + """ + Test trtllm_fp8_per_tensor_scale_moe API (TensorRT-LLM fused MoE). + + This test: + 1. Creates quantized FP8 weights and per-tensor scales + 2. Runs FP8 per-tensor scale MOE + 3. Measures performance metrics (TFLOPS, TB/sec) + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.verbose >= 1: + print("[INFO] Running testTrtllmFp8PerTensorScaleMoe") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype) + + # Parse configuration + num_tokens = args.num_tokens + hidden_size = args.hidden_size + intermediate_size = args.intermediate_size + num_experts = args.num_experts + top_k = args.top_k + n_group = ( + args.n_group + if hasattr(args, "n_group") and args.n_group is not None and args.n_group > 0 + else None + ) + topk_group = ( + args.topk_group + if hasattr(args, "topk_group") + and args.topk_group is not None + and args.topk_group > 0 + else None + ) + routed_scaling_factor = ( + args.routed_scaling_factor + if hasattr(args, "routed_scaling_factor") + and args.routed_scaling_factor is not None + else None + ) + local_expert_offset = args.local_expert_offset + local_num_experts = args.local_num_experts or num_experts + tile_tokens_dim = args.tile_tokens_dim + routing_method_type = args.routing_method_type + use_routing_scales_on_input = args.use_routing_scales_on_input + is_cuda_graph_compatible = not args.no_cuda_graph + res = [] + backends = ["trtllm"] + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + if args.verbose >= 1: + print( + f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, " + f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}" + ) + + # Create test data + routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights = ( + create_trtllm_moe_test_data( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + routing_method_type, + args.use_routing_bias, + input_dtype, + weight_dtype, + device, + moe_kernel_type="fp8_per_tensor", + ) + ) + + # For FP8 per-tensor scale, create quantized weights and per-tensor scales + # Quantize to FP8 + gemm1_weights_fp8 = gemm1_weights.to(torch.float8_e4m3fn) + gemm2_weights_fp8 = gemm2_weights.to(torch.float8_e4m3fn) + + # Quantize hidden states to FP8 for per-tensor scale + hidden_states_fp8 = hidden_states.to(torch.float8_e4m3fn) + + # Create per-tensor scale scalars + output1_scales_scalar = torch.ones( + local_num_experts, device=device, dtype=torch.float32 + ) + output1_scales_gate_scalar = torch.ones( + local_num_experts, device=device, dtype=torch.float32 + ) + output2_scales_scalar = torch.ones( + local_num_experts, device=device, dtype=torch.float32 + ) + + if args.verbose >= 2: + print(f"[VVERBOSE] routing_logits.shape = {routing_logits.shape}") + print(f"[VVERBOSE] hidden_states.shape = {hidden_states.shape}") + print(f"[VVERBOSE] gemm1_weights_fp8.shape = {gemm1_weights_fp8.shape}") + print(f"[VVERBOSE] gemm2_weights_fp8.shape = {gemm2_weights_fp8.shape}") + + def run_fp8_per_tensor_moe(): + # Note: FP8 per-tensor MOE expects int64_t for n_group/topk_group, not Optional[int64_t] + # So we convert None to 0 to indicate "no groups" mode + return trtllm_fp8_per_tensor_scale_moe( + routing_logits=routing_logits, + routing_bias=routing_bias, + hidden_states=hidden_states_fp8, + gemm1_weights=gemm1_weights_fp8, + output1_scales_scalar=output1_scales_scalar, + output1_scales_gate_scalar=output1_scales_gate_scalar, + gemm2_weights=gemm2_weights_fp8, + output2_scales_scalar=output2_scales_scalar, + num_experts=num_experts, + top_k=top_k, + n_group=n_group if n_group is not None else 0, + topk_group=topk_group if topk_group is not None else 0, + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=routed_scaling_factor, + use_routing_scales_on_input=use_routing_scales_on_input, + tile_tokens_dim=tile_tokens_dim, + routing_method_type=routing_method_type, + ) + + # Benchmark timing + times = bench_gpu_time( + fn=run_fp8_per_tensor_moe, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + l2_flush=True, + l2_flush_size_mb=256, + l2_flush_device=device, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + ) + + # Compute performance metrics + median_time = np.median(times) + std_time = np.std(times) + tflops = calculate_moe_tflops( + num_tokens, hidden_size, intermediate_size, num_experts, top_k, median_time + ) + tb_per_sec = calculate_moe_bandwidth( + num_tokens, + hidden_size, + intermediate_size, + num_experts, + top_k, + median_time, + input_dtype, + weight_dtype, + input_format="fp8", + weight_format="fp8", + routing_logits_dtype=routing_logits.dtype, + ) + + backend = "trtllm" + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["num_tokens"] = num_tokens + cur_res["hidden_size"] = hidden_size + cur_res["intermediate_size"] = intermediate_size + cur_res["num_experts"] = num_experts + cur_res["top_k"] = top_k + cur_res["n_group"] = n_group + cur_res["topk_group"] = topk_group + cur_res["routed_scaling_factor"] = routed_scaling_factor + cur_res["local_expert_offset"] = local_expert_offset + cur_res["local_num_experts"] = local_num_experts + cur_res["tile_tokens_dim"] = tile_tokens_dim + cur_res["routing_method"] = args.routing_method + cur_res["use_routing_bias"] = args.use_routing_bias + cur_res["use_routing_scales_on_input"] = use_routing_scales_on_input + cur_res["input_dtype"] = input_dtype + cur_res["weight_dtype"] = weight_dtype + res.append(cur_res) + + return res diff --git a/benchmarks/samples/sample_testlist.txt b/benchmarks/samples/sample_testlist.txt new file mode 100644 index 0000000000..051b793e57 --- /dev/null +++ b/benchmarks/samples/sample_testlist.txt @@ -0,0 +1,50 @@ +## Attention Prefill +# Paged prefill for Llama 3.1 70B +--routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "Llama-3.1-70B" +--routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 32 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "Llama-3.1-70B" + +# Ragged prefill for DeepSeep-R1 +--routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "DeepSeek-R1" +--routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "DeepSeek-R1" + +## Attention Decode +--routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "Llama-3.1-70B" +--routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "Llama-3.1-70B" + +## Attention MLA +# DeepSeek-R1 +--routine BatchMLAPagedAttentionWrapper --backends trtllm-gen-native fa2 fa3 --page_size 32 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_ckv 512 --head_dim_kpe 64 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --generate_repro_command --case_tag "DeepSeek-R1" + +## FP8 bmm +--routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +--routine bmm_fp8 --batch_size 64 --m 4 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command + +## FP8 GEMM with groupwise scaling +--routine gemm_fp8_nt_groupwise --m 4 --n 1024 --k 7168 --mma_sm 1 --scale_major_mode MN --backends cutlass --refcheck -vv --generate_repro_command +--routine gemm_fp8_nt_groupwise --m 16 --n 1024 --k 7168 --mma_sm 1 --scale_major_mode MN --backends cutlass --refcheck -vv --generate_repro_command + +## FP8 group GEMM with groupwise scaling +--routine group_gemm_fp8_nt_groupwise --m 4 --n 1024 --k 7168 --mma_sm 1 --group_size 2 --scale_major_mode MN --refcheck -vv --generate_repro_command +--routine group_gemm_fp8_nt_groupwise --m 16 --n 1024 --k 7168 --mma_sm 1 --group_size 2 --scale_major_mode MN --refcheck -vv --generate_repro_command + +## FP4 GEMM +# non-autotuned +--routine mm_fp4 --m 1 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +--routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command + +# autotuned +--routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --autotune --refcheck -vv --generate_repro_command + +## MoE +--routine trtllm_fp4_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 256 --top_k 8 --n_group 8 --topk_group 4 --routed_scaling_factor 2.5 --use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight -vv --generate_repro_command --case_tag "trtllm_moe_sample" +--routine trtllm_fp4_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 8 --routing_method renormalize_naive --use_shuffled_weight -vv --generate_repro_command --case_tag "trtllm_moe_sample" +--routine trtllm_fp8_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 256 --top_k 8 --n_group 8 --topk_group 4 --routed_scaling_factor 2.5 --use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight -vv --generate_repro_command --case_tag "trtllm_moe_sample" +--routine trtllm_fp8_per_tensor_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 1 --routed_scaling_factor 2.5 --use_routing_bias --routing_method llama4 --use_routing_scales_on_input -vv --generate_repro_command --case_tag "trtllm_moe_sample" +--routine trtllm_fp8_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 1 --routing_method renormalize --use_shuffled_weight -vv --generate_repro_command --case_tag "trtllm_moe_sample" + +# CUTLASS MoE API +--routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant base --input_dtype float16 -vv --generate_repro_command --case_tag "cutlass_moe_base" +--routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant fp8 --input_dtype float16 -vv --generate_repro_command --case_tag "cutlass_moe_fp8_scale" +--routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant nvfp4 --input_dtype float16 -vv --generate_repro_command --case_tag "cutlass_moe_nvfp4_weights" +--routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 -vv --generate_repro_command --case_tag "cutlass_moe_nvfp4_weights_quantized" +--routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 8 --top_k 2 --cutlass_variant base --input_dtype float16 --tp_size 2 --tp_rank 0 --ep_size 4 --ep_rank 0 -vv --generate_repro_command --case_tag "cutlass_moe_nvfp4_ep_tp" diff --git a/benchmarks/samples/sample_testlist_output.csv b/benchmarks/samples/sample_testlist_output.csv new file mode 100644 index 0000000000..d856d37ab0 --- /dev/null +++ b/benchmarks/samples/sample_testlist_output.csv @@ -0,0 +1,53 @@ +routine,median_time,std_time,tflops,tb_per_sec,backend,page_size,batch_size,s_qo,s_kv,num_qo_heads,num_kv_heads,head_dim_qk,head_dim_vo,head_dim_ckv,head_dim_kpe,causal,q_dtype,kv_dtype,avg_actual_seq_len,random_actual_seq_len,m,n,k,group_size,tile_size,scale_major_mode,out_dtype,mma_sm,use_128x4_sf_layout,use_nvfp4,num_tokens,hidden_size,intermediate_size,num_experts,top_k,n_group,topk_group,routed_scaling_factor,local_expert_offset,local_num_experts,tile_tokens_dim,routing_method,use_shuffled_weight,weight_layout,use_routing_bias,use_routing_scales_on_input,input_dtype,weight_dtype,gated_act,cutlass_variant,quantized_input,tp_size,tp_rank,ep_size,ep_rank,refcheck,no_cuda_graph,use_cupti,allow_output_mismatch,random_seed,case_tag,generate_repro_command,repro_command +BatchPrefillWithPagedKVCacheWrapper,0.01244799979031086,0.0009464459008260536,13.963516944729905,0.3050282827732261,fa2,16,1,1024,1024,64,8,128,128,,,True,torch.bfloat16,torch.bfloat16,103,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchPrefillWithPagedKVCacheWrapper,0.01839040070772171,0.00021363710731210026,9.45155349045863,0.20646597430613514,cudnn,16,1,1024,1024,64,8,128,128,,,True,torch.bfloat16,torch.bfloat16,103,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchPrefillWithPagedKVCacheWrapper,0.008396799862384795,5.550615129103214e-05,20.70048814413847,0.45219512936224815,trtllm-gen,16,1,1024,1024,64,8,128,128,,,True,torch.bfloat16,torch.bfloat16,103,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchPrefillWithPagedKVCacheWrapper,0.4833280146121979,0.003473954933671819,250.42114497152383,0.9745931908746264,fa2,16,32,1024,1024,64,8,128,128,,,True,torch.bfloat16,torch.bfloat16,399,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 32 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchPrefillWithPagedKVCacheWrapper,0.3817088007926941,0.0008139816712432105,317.08871937101173,1.2340511694301386,cudnn,16,32,1024,1024,64,8,128,128,,,True,torch.bfloat16,torch.bfloat16,399,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 32 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchPrefillWithPagedKVCacheWrapper,0.7442896127700807,0.00045553586925676576,162.6188955741738,0.6328829314799427,trtllm-gen,16,32,1024,1024,64,8,128,128,,,True,torch.bfloat16,torch.bfloat16,399,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 32 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchPrefillWithRaggedKVCacheWrapper,0.016127999871969223,0.00017067009388203107,26.943492277380717,1.0463492146555617,fa2,0,1,1024,1024,128,128,192,128,,,True,torch.bfloat16,torch.bfloat16,103,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +BatchPrefillWithRaggedKVCacheWrapper,0.012083200365304947,9.971927609146905e-05,35.962710777165306,1.3966101272685556,cutlass,0,1,1024,1024,128,128,192,128,,,True,torch.bfloat16,torch.bfloat16,103,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +BatchPrefillWithRaggedKVCacheWrapper,0.018636800348758698,0.00019618687934467522,23.316483080151837,0.9054944885495858,cudnn,0,1,1024,1024,128,128,192,128,,,True,torch.bfloat16,torch.bfloat16,103,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +BatchPrefillWithRaggedKVCacheWrapper,0.49769599735736847,0.0053031528422255855,217.96787873723878,1.7256503660070768,fa2,0,16,1024,1024,128,128,192,128,,,True,torch.bfloat16,torch.bfloat16,327,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +BatchPrefillWithRaggedKVCacheWrapper,0.5328896045684814,0.0014588070313195434,203.57263468827725,1.6116833067056566,cutlass,0,16,1024,1024,128,128,192,128,,,True,torch.bfloat16,torch.bfloat16,327,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +BatchPrefillWithRaggedKVCacheWrapper,0.3123199939727783,0.0005266243249355331,347.3416460473396,2.7499016924124846,cudnn,0,16,1024,1024,128,128,192,128,,,True,torch.bfloat16,torch.bfloat16,327,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +BatchDecodeWithPagedKVCacheWrapper,0.03481600061058998,0.00022415261777036224,0.07905882214290773,0.0108235292219457,fa2,16,1,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,84,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.010452799871563912,4.634408684090133e-05,0.2633277240376533,0.03605081936229778,fa2_tc,16,1,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,84,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.010683200135827065,8.551724418800796e-05,0.25764864132510307,0.035273325895698635,cudnn,16,1,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,84,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.005734400078654289,3.59118315932623e-05,0.4799999934162147,0.06571428481293416,trtllm-gen,16,1,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,84,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.006188799999654293,9.38182546863553e-05,0.444756980376447,0.06088934850391835,trtllm-gen-native,16,1,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,84,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.055296000093221664,7.235326722628814e-05,4.756740732721512,0.604074073055686,fa2,16,16,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.017208000272512437,9.39948911411913e-05,15.28525870726272,1.941125027372111,fa2_tc,16,16,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.013764800131320953,9.368704015807337e-05,19.108794424228098,2.426688341372557,cudnn,16,16,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.00963199995458126,3.5251792291711315e-05,27.30780079321905,3.467906993096757,trtllm-gen,16,16,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchDecodeWithPagedKVCacheWrapper,0.009657599776983262,6.567264399555296e-05,27.2354148105071,3.458714460254226,trtllm-gen-native,16,16,1,1024,64,8,128,128,,,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,True,42,Llama-3.1-70B,True,python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +BatchMLAPagedAttentionWrapper,0.024420800060033797,0.00010761519902284579,91.55081939697266,0.9553665704090659,trtllm-gen-native,32,16,1,1024,128,,,,512,64,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchMLAPagedAttentionWrapper --backends trtllm-gen-native fa2 fa3 --page_size 32 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_ckv 512 --head_dim_kpe 64 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --generate_repro_command --case_tag DeepSeek-R1 +BatchMLAPagedAttentionWrapper,0.04095999896526337,0.0004555166043636676,54.58359909057617,0.5696000143893066,fa2,32,16,1,1024,128,,,,512,64,False,torch.bfloat16,torch.bfloat16,501,True,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,DeepSeek-R1,True,python3 flashinfer_benchmark.py --routine BatchMLAPagedAttentionWrapper --backends trtllm-gen-native fa2 fa3 --page_size 32 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_ckv 512 --head_dim_kpe 64 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --generate_repro_command --case_tag DeepSeek-R1 +bmm_fp8,0.2860383987426758,0.0004968334033247884,13.13843316323707,0.02569322172234466,cudnn,,256,,,,,,,,,,,,,,1,1024,7168,,,,torch.bfloat16,,,,,,,,,,,,,,,,,,,,torch.float8_e4m3fn,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +bmm_fp8,0.28600159883499143,0.00047219507506758893,13.14012368919739,0.025696527676546826,cublas,,256,,,,,,,,,,,,,,1,1024,7168,,,,torch.bfloat16,,,,,,,,,,,,,,,,,,,,torch.float8_e4m3fn,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +bmm_fp8,0.2657311916351318,0.00010204157533447904,14.142473681298727,0.027656700573153076,cutlass,,256,,,,,,,,,,,,,,1,1024,7168,,,,torch.bfloat16,,,,,,,,,,,,,,,,,,,,torch.float8_e4m3fn,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +bmm_fp8,0.07516320049762726,0.000206792690088806,49.99915329734576,0.09814504905539344,cudnn,,64,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,,,,,,,,,,,,,,,,,,,torch.float8_e4m3fn,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 64 --m 4 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +bmm_fp8,0.07495999932289124,0.0002273436657742046,50.13469074101705,0.09841110014187592,cublas,,64,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,,,,,,,,,,,,,,,,,,,torch.float8_e4m3fn,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 64 --m 4 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +bmm_fp8,0.072297602891922,9.86900944781081e-05,51.980926526955464,0.10203513954712654,cutlass,,64,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,,,,,,,,,,,,,,,,,,,torch.float8_e4m3fn,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 64 --m 4 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +gemm_fp8_nt_groupwise,0.01966080069541931,1.0062279270694632e-06,2.986666561025716,0.37520832006189414,cutlass,,,,,,,,,,,,,,,,4,1024,7168,,128,MN,torch.bfloat16,1,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine gemm_fp8_nt_groupwise --m 4 --n 1024 --k 7168 --mma_sm 1 --scale_major_mode MN --backends cutlass --refcheck -vv --generate_repro_command +gemm_fp8_nt_groupwise,0.019865599274635316,1.2664492914798042e-06,11.823505586357996,0.37690723025708733,cutlass,,,,,,,,,,,,,,,,16,1024,7168,,128,MN,torch.bfloat16,1,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine gemm_fp8_nt_groupwise --m 16 --n 1024 --k 7168 --mma_sm 1 --scale_major_mode MN --backends cutlass --refcheck -vv --generate_repro_command +group_gemm_fp8_nt_groupwise,0.02232320010662079,3.682025865219179e-05,5.260917406065297,0.6609174280359654,cutlass,,,,,,,,,,,,,,,,4,1024,7168,2,128,MN,torch.bfloat16,1,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine group_gemm_fp8_nt_groupwise --m 4 --n 1024 --k 7168 --mma_sm 1 --group_size 2 --scale_major_mode MN --refcheck -vv --generate_repro_command +group_gemm_fp8_nt_groupwise,0.02252800017595291,9.332681200911203e-05,20.852363473498134,0.6647272675354804,cutlass,,,,,,,,,,,,,,,,16,1024,7168,2,128,MN,torch.bfloat16,1,,,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine group_gemm_fp8_nt_groupwise --m 16 --n 1024 --k 7168 --mma_sm 1 --group_size 2 --scale_major_mode MN --refcheck -vv --generate_repro_command +mm_fp4,0.012697599828243256,1.7496054581842973e-06,1.1561290478966861,0.28947581036726805,cudnn,,,,,,,,,,,,,,,,1,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 1 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +mm_fp4,0.009216000139713288,3.619193125084151e-05,1.5928888647409136,0.39883332728707516,cutlass,,,,,,,,,,,,,,,,1,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 1 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +mm_fp4,0.01085439994931221,3.680557773573178e-05,1.3524528365043527,0.3386320770530395,trtllm,,,,,,,,,,,,,,,,1,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 1 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +mm_fp4,0.012697599828243256,1.7233627387457594e-06,4.6245161915867445,0.2908064555465576,cudnn,,,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +mm_fp4,0.009216000139713288,5.744219533265305e-07,6.3715554589636545,0.4006666605926154,cutlass,,,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +mm_fp4,0.011059200018644333,0.00010158394987753556,5.309629620678304,0.33388888832599684,trtllm,,,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +mm_fp4,0.009216000139713288,3.6526144810097353e-05,6.3715554589636545,0.4006666605926154,cutlass_autotune,,,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --autotune --refcheck -vv --generate_repro_command +mm_fp4,0.01085439994931221,0.00010171082787792661,5.409811346017411,0.34018868083389336,trtllm_autotune,,,,,,,,,,,,,,,,4,1024,7168,,,,torch.bfloat16,,True,True,,,,,,,,,,,,,,,,,,,,,,,,,,True,False,False,False,42,None,True,python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --autotune --refcheck -vv --generate_repro_command +trtllm_fp4_block_scale_moe,0.22354559898376464,0.0001550481673529622,230.55523251765356,1.817630057791967,trtllm,,,,,,,,,,,,,,,,,,,,,,,,,,1024,1024,1024,256,8,8,4,2.5,0,256,8,deepseek_v3,True,0,True,False,torch.bfloat16,torch.bfloat16,swiglu,,,,,,,False,False,False,False,42,trtllm_moe_sample,True,python3 flashinfer_benchmark.py --routine trtllm_fp4_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 256 --top_k 8 --n_group 8 --topk_group 4 --routed_scaling_factor 2.5 --use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +trtllm_fp4_block_scale_moe,0.22620320320129395,0.00015420901713100778,227.84649740850875,0.9027731398581356,trtllm,,,,,,,,,,,,,,,,,,,,,,,,,,1024,1024,1024,128,8,None,None,2.5,0,128,8,renormalize_naive,True,0,False,False,torch.bfloat16,torch.bfloat16,swiglu,,,,,,,False,False,False,False,42,trtllm_moe_sample,True,python3 flashinfer_benchmark.py --routine trtllm_fp4_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 8 --routing_method renormalize_naive --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +trtllm_fp8_block_scale_moe,0.556544017791748,0.00016468714317887162,92.60652509840739,1.45451329296815,trtllm,,,,,,,,,,,,,,,,,,,,,,,,,,1024,1024,1024,256,8,8,4,2.5,0,256,8,deepseek_v3,True,0,True,False,torch.bfloat16,torch.bfloat16,,,,,,,,False,False,False,False,42,trtllm_moe_sample,True,python3 flashinfer_benchmark.py --routine trtllm_fp8_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 256 --top_k 8 --n_group 8 --topk_group 4 --routed_scaling_factor 2.5 --use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +trtllm_fp8_per_tensor_scale_moe,0.12308800220489502,0.00022594891363694267,52.34020236412443,3.2989491154796857,trtllm,,,,,,,,,,,,,,,,,,,,,,,,,,1024,1024,1024,128,1,None,None,2.5,0,128,8,llama4,,,True,True,torch.bfloat16,torch.bfloat16,,,,,,,,False,False,False,False,42,trtllm_moe_sample,True,python3 flashinfer_benchmark.py --routine trtllm_fp8_per_tensor_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 1 --routed_scaling_factor 2.5 --use_routing_bias --routing_method llama4 --use_routing_scales_on_input -vv --generate_repro_command --case_tag trtllm_moe_sample +trtllm_fp8_block_scale_moe,0.10864640474319458,0.00013707306207295686,59.29741494187403,3.7398678857383114,trtllm,,,,,,,,,,,,,,,,,,,,,,,,,,1024,1024,1024,128,1,None,None,2.5,0,128,8,renormalize,True,0,False,False,torch.bfloat16,torch.bfloat16,,,,,,,,False,False,False,False,42,trtllm_moe_sample,True,python3 flashinfer_benchmark.py --routine trtllm_fp8_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 1 --routing_method renormalize --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +cutlass_fused_moe,0.026214399933815004,0.00010491445634120407,0.24000000060594173,0.00812988283302598,cutlass,,,,,,,,,,,,,,,,,,,,,,,,,,32,128,128,2,2,,,,,,,,False,0,,False,torch.float16,torch.float16,,base,False,1,0,1,0,False,False,False,False,42,cutlass_moe_base,True,python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant base --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_base +cutlass_fused_moe,0.025804799795150758,0.0001189026818394618,0.2438095257449853,0.004290674637235764,cutlass,,,,,,,,,,,,,,,,,,,,,,,,,,32,128,128,2,2,,,,,,,,False,0,,False,torch.float16,torch.float16,,fp8,False,1,0,1,0,False,False,False,False,42,cutlass_moe_fp8_scale,True,python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant fp8 --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_fp8_scale +cutlass_fused_moe,0.02990399897098541,0.00010532265873466246,0.21038845025724934,0.002195826720824563,cutlass,,,,,,,,,,,,,,,,,,,,,,,,,,32,128,128,2,2,,,,,,,,False,0,,False,torch.float16,torch.float16,,nvfp4,False,1,0,1,0,False,False,False,False,42,cutlass_moe_nvfp4_weights,True,python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant nvfp4 --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_nvfp4_weights +cutlass_fused_moe,0.02949439883232117,0.00010231710901612176,0.21331019614156588,0.0020180102784388863,cutlass,,,,,,,,,,,,,,,,,,,,,,,,,,32,128,128,2,2,,,,,,,,False,0,,False,torch.float16,torch.float16,,nvfp4,True,1,0,1,0,False,False,False,False,42,cutlass_moe_nvfp4_weights_quantized,True,python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_nvfp4_weights_quantized +cutlass_fused_moe,0.025190401077270507,0.00010808926975429793,0.24975608688012632,0.031890242538648944,cutlass,,,,,,,,,,,,,,,,,,,,,,,,,,32,128,128,8,2,,,,,,,,False,0,,False,torch.float16,torch.float16,,base,False,2,0,4,0,False,False,False,False,42,cutlass_moe_nvfp4_ep_tp,True,python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 8 --top_k 2 --cutlass_variant base --input_dtype float16 --tp_size 2 --tp_rank 0 --ep_size 4 --ep_rank 0 -vv --generate_repro_command --case_tag cutlass_moe_nvfp4_ep_tp diff --git a/benchmarks/samples/sample_testlist_output.txt b/benchmarks/samples/sample_testlist_output.txt new file mode 100644 index 0000000000..69a3961f87 --- /dev/null +++ b/benchmarks/samples/sample_testlist_output.txt @@ -0,0 +1,394 @@ +[INFO] args = Namespace(routine='BatchPrefillWithPagedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='Llama-3.1-70B', generate_repro_command=True, repro_command='', backends=['fa2', 'fa3', 'cudnn', 'trtllm-gen'], page_size=16, batch_size=1, s_qo=1024, s_kv=1024, num_qo_heads=64, num_kv_heads=8, head_dim_qk=128, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=True, random_actual_seq_len=True) +[INFO] Running testBatchPrefillWithPagedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +[WARNING] fa3 for routine BatchPrefillWithPagedKVCacheWrapper is not supported on compute capability 10.0. Skipping. +[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q +[VERBOSE] Average actual qo seq len: 103 +[VERBOSE] Average actual kv seq len: 103 +[VVERBOSE] actual_seq_lens_q.flatten() = tensor([103], dtype=torch.int32) +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([103], dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([103, 64, 128]) +[VVERBOSE] num_pages_per_seq = 64 +[VVERBOSE] total_num_pages = 64 +[VVERBOSE] kv_cache.shape = torch.Size([64, 2, 8, 16, 128]) +[VVERBOSE] kv_cache.stride() = (32768, 16384, 128, 1024, 1) +[VVERBOSE] block_tables.shape = torch.Size([1, 64]) +[VVERBOSE] qo_indptr.shape = torch.Size([2]) +[VVERBOSE] qo_indptr.dtype = torch.int32 +[VVERBOSE] kv_indptr.shape = torch.Size([2]) +[VVERBOSE] kv_indices.shape = torch.Size([7]) +[VVERBOSE] kv_last_page_len.shape = torch.Size([1]) +[VVERBOSE] scale = 0.08838834764831843 +[PERF] fa2 :: median time 0.012 ms; std 0.001 ms; achieved tflops 13.964 TFLOPs/sec; achieved tb_per_sec 0.305 TB/sec +[PERF] cudnn :: median time 0.018 ms; std 0.000 ms; achieved tflops 9.452 TFLOPs/sec; achieved tb_per_sec 0.206 TB/sec +[PERF] trtllm-gen :: median time 0.008 ms; std 0.000 ms; achieved tflops 20.700 TFLOPs/sec; achieved tb_per_sec 0.452 TB/sec +[INFO] args = Namespace(routine='BatchPrefillWithPagedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='Llama-3.1-70B', generate_repro_command=True, repro_command='', backends=['fa2', 'fa3', 'cudnn', 'trtllm-gen'], page_size=16, batch_size=32, s_qo=1024, s_kv=1024, num_qo_heads=64, num_kv_heads=8, head_dim_qk=128, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=True, random_actual_seq_len=True) +[INFO] Running testBatchPrefillWithPagedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 fa3 cudnn trtllm-gen --page_size 16 --batch_size 32 --s_qo 1024 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +[WARNING] fa3 for routine BatchPrefillWithPagedKVCacheWrapper is not supported on compute capability 10.0. Skipping. +[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q +[VERBOSE] Average actual qo seq len: 399 +[VERBOSE] Average actual kv seq len: 399 +[VVERBOSE] actual_seq_lens_q.flatten() = tensor([103, 436, 861, 271, 107, 72, 701, 21, 615, 122, 467, 215, 331, 459, + 88, 373, 100, 872, 664, 131, 662, 309, 770, 344, 492, 414, 806, 386, + 192, 956, 277, 161], dtype=torch.int32) +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([103, 436, 861, 271, 107, 72, 701, 21, 615, 122, 467, 215, 331, 459, + 88, 373, 100, 872, 664, 131, 662, 309, 770, 344, 492, 414, 806, 386, + 192, 956, 277, 161], dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([12778, 64, 128]) +[VVERBOSE] num_pages_per_seq = 64 +[VVERBOSE] total_num_pages = 2048 +[VVERBOSE] kv_cache.shape = torch.Size([2048, 2, 8, 16, 128]) +[VVERBOSE] kv_cache.stride() = (32768, 16384, 128, 1024, 1) +[VVERBOSE] block_tables.shape = torch.Size([32, 64]) +[VVERBOSE] qo_indptr.shape = torch.Size([33]) +[VVERBOSE] qo_indptr.dtype = torch.int32 +[VVERBOSE] kv_indptr.shape = torch.Size([33]) +[VVERBOSE] kv_indices.shape = torch.Size([815]) +[VVERBOSE] kv_last_page_len.shape = torch.Size([32]) +[VVERBOSE] scale = 0.08838834764831843 +[PERF] fa2 :: median time 0.483 ms; std 0.003 ms; achieved tflops 250.421 TFLOPs/sec; achieved tb_per_sec 0.975 TB/sec +[PERF] cudnn :: median time 0.382 ms; std 0.001 ms; achieved tflops 317.089 TFLOPs/sec; achieved tb_per_sec 1.234 TB/sec +[PERF] trtllm-gen :: median time 0.744 ms; std 0.000 ms; achieved tflops 162.619 TFLOPs/sec; achieved tb_per_sec 0.633 TB/sec +[INFO] args = Namespace(routine='BatchPrefillWithRaggedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='DeepSeek-R1', generate_repro_command=True, repro_command='', backends=['fa2', 'fa3', 'cutlass', 'cudnn'], page_size=0, batch_size=1, s_qo=1024, s_kv=1024, num_qo_heads=128, num_kv_heads=128, head_dim_qk=192, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=True, random_actual_seq_len=True) +[INFO] Running testBatchPrefillWithRaggedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 1 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +[WARNING] fa3 for routine BatchPrefillWithRaggedKVCacheWrapper is not supported on compute capability 10.0. Skipping. +[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q +[VERBOSE] Average actual qo seq len: 103 +[VERBOSE] Average actual kv seq len: 103 +[VVERBOSE] actual_seq_lens_q.flatten() = tensor([103], dtype=torch.int32) +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([103], dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([103, 128, 192]) +[VVERBOSE] k.shape = torch.Size([103, 128, 192]) +[VVERBOSE] v.shape = torch.Size([103, 128, 128]) +[VVERBOSE] qo_indptr.shape = torch.Size([2]) +[VVERBOSE] kv_indptr.shape = torch.Size([2]) +[VVERBOSE] scale = 0.07216878364870323 +[PERF] fa2 :: median time 0.016 ms; std 0.000 ms; achieved tflops 26.943 TFLOPs/sec; achieved tb_per_sec 1.046 TB/sec +[PERF] cutlass :: median time 0.012 ms; std 0.000 ms; achieved tflops 35.963 TFLOPs/sec; achieved tb_per_sec 1.397 TB/sec +[PERF] cudnn :: median time 0.019 ms; std 0.000 ms; achieved tflops 23.316 TFLOPs/sec; achieved tb_per_sec 0.905 TB/sec +[INFO] args = Namespace(routine='BatchPrefillWithRaggedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='DeepSeek-R1', generate_repro_command=True, repro_command='', backends=['fa2', 'fa3', 'cutlass', 'cudnn'], page_size=0, batch_size=16, s_qo=1024, s_kv=1024, num_qo_heads=128, num_kv_heads=128, head_dim_qk=192, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=True, random_actual_seq_len=True) +[INFO] Running testBatchPrefillWithRaggedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1 +[WARNING] fa3 for routine BatchPrefillWithRaggedKVCacheWrapper is not supported on compute capability 10.0. Skipping. +[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q +[VERBOSE] Average actual qo seq len: 327 +[VERBOSE] Average actual kv seq len: 327 +[VVERBOSE] actual_seq_lens_q.flatten() = tensor([103, 436, 861, 271, 107, 72, 701, 21, 615, 122, 467, 215, 331, 459, + 88, 373], dtype=torch.int32) +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([103, 436, 861, 271, 107, 72, 701, 21, 615, 122, 467, 215, 331, 459, + 88, 373], dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([5242, 128, 192]) +[VVERBOSE] k.shape = torch.Size([5242, 128, 192]) +[VVERBOSE] v.shape = torch.Size([5242, 128, 128]) +[VVERBOSE] qo_indptr.shape = torch.Size([17]) +[VVERBOSE] kv_indptr.shape = torch.Size([17]) +[VVERBOSE] scale = 0.07216878364870323 +[PERF] fa2 :: median time 0.498 ms; std 0.005 ms; achieved tflops 217.968 TFLOPs/sec; achieved tb_per_sec 1.726 TB/sec +[PERF] cutlass :: median time 0.533 ms; std 0.001 ms; achieved tflops 203.573 TFLOPs/sec; achieved tb_per_sec 1.612 TB/sec +[PERF] cudnn :: median time 0.312 ms; std 0.001 ms; achieved tflops 347.342 TFLOPs/sec; achieved tb_per_sec 2.750 TB/sec +[INFO] args = Namespace(routine='BatchDecodeWithPagedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='Llama-3.1-70B', generate_repro_command=True, repro_command='', backends=['fa2', 'fa2_tc', 'cudnn', 'trtllm-gen', 'trtllm-gen-native'], page_size=16, batch_size=1, s_qo=1, s_kv=1024, num_qo_heads=64, num_kv_heads=8, head_dim_qk=128, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=False, random_actual_seq_len=True) +[INFO] Running testBatchDecodeWithPagedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 1 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +[VERBOSE] Average actual seq len: 84 +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([84], device='cuda:0', dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([1, 64, 128]) +[VVERBOSE] num_pages_per_seq = 64 +[VVERBOSE] total_num_pages = 64 +[VVERBOSE] kv_cache.shape = torch.Size([64, 2, 8, 16, 128]) +[VVERBOSE] kv_cache.stride() = (32768, 16384, 128, 1024, 1) +[VVERBOSE] block_tables.shape = torch.Size([1, 64]) +[VVERBOSE] kv_indptr.shape = torch.Size([2]) +[VVERBOSE] kv_indices.shape = torch.Size([6]) +[VVERBOSE] kv_last_page_len.shape = torch.Size([1]) +[VVERBOSE] scale = 0.08838834764831843 +[ERROR] Output tensor mismatch between backends fa2 and cudnn: 5063 / 8192 (61.80%) elements are different +[PERF] fa2 :: median time 0.035 ms; std 0.000 ms; achieved tflops 0.079 TFLOPs/sec; achieved tb_per_sec 0.011 TB/sec +[PERF] fa2_tc :: median time 0.010 ms; std 0.000 ms; achieved tflops 0.263 TFLOPs/sec; achieved tb_per_sec 0.036 TB/sec +[PERF] cudnn :: median time 0.011 ms; std 0.000 ms; achieved tflops 0.258 TFLOPs/sec; achieved tb_per_sec 0.035 TB/sec +[PERF] trtllm-gen :: median time 0.006 ms; std 0.000 ms; achieved tflops 0.480 TFLOPs/sec; achieved tb_per_sec 0.066 TB/sec +[PERF] trtllm-gen-nati:: median time 0.006 ms; std 0.000 ms; achieved tflops 0.445 TFLOPs/sec; achieved tb_per_sec 0.061 TB/sec +[INFO] args = Namespace(routine='BatchDecodeWithPagedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='Llama-3.1-70B', generate_repro_command=True, repro_command='', backends=['fa2', 'fa2_tc', 'cudnn', 'trtllm-gen', 'trtllm-gen-native'], page_size=16, batch_size=16, s_qo=1, s_kv=1024, num_qo_heads=64, num_kv_heads=8, head_dim_qk=128, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=False, random_actual_seq_len=True) +[INFO] Running testBatchDecodeWithPagedKVCacheWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc cudnn trtllm-gen trtllm-gen-native --page_size 16 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag Llama-3.1-70B +[VERBOSE] Average actual seq len: 501 +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([ 84, 874, 167, 691, 274, 736, 63, 813, 781, 450, 794, 226, 510, 499, + 524, 541], device='cuda:0', dtype=torch.int32) +[VVERBOSE] q.shape = torch.Size([16, 64, 128]) +[VVERBOSE] num_pages_per_seq = 64 +[VVERBOSE] total_num_pages = 1024 +[VVERBOSE] kv_cache.shape = torch.Size([1024, 2, 8, 16, 128]) +[VVERBOSE] kv_cache.stride() = (32768, 16384, 128, 1024, 1) +[VVERBOSE] block_tables.shape = torch.Size([16, 64]) +[VVERBOSE] kv_indptr.shape = torch.Size([17]) +[VVERBOSE] kv_indices.shape = torch.Size([509]) +[VVERBOSE] kv_last_page_len.shape = torch.Size([16]) +[VVERBOSE] scale = 0.08838834764831843 +[PERF] fa2 :: median time 0.055 ms; std 0.000 ms; achieved tflops 4.757 TFLOPs/sec; achieved tb_per_sec 0.604 TB/sec +[PERF] fa2_tc :: median time 0.017 ms; std 0.000 ms; achieved tflops 15.285 TFLOPs/sec; achieved tb_per_sec 1.941 TB/sec +[PERF] cudnn :: median time 0.014 ms; std 0.000 ms; achieved tflops 19.109 TFLOPs/sec; achieved tb_per_sec 2.427 TB/sec +[PERF] trtllm-gen :: median time 0.010 ms; std 0.000 ms; achieved tflops 27.308 TFLOPs/sec; achieved tb_per_sec 3.468 TB/sec +[PERF] trtllm-gen-nati:: median time 0.010 ms; std 0.000 ms; achieved tflops 27.235 TFLOPs/sec; achieved tb_per_sec 3.459 TB/sec +[INFO] args = Namespace(routine='BatchMLAPagedAttentionWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='DeepSeek-R1', generate_repro_command=True, repro_command='', backends=['trtllm-gen-native', 'fa2', 'fa3'], page_size=32, batch_size=16, s_qo=1, s_kv=1024, num_qo_heads=128, num_kv_heads=128, head_dim_qk=None, head_dim_vo=None, head_dim_ckv=512, head_dim_kpe=64, q_dtype='bfloat16', kv_dtype='bfloat16', causal=False, random_actual_seq_len=True) +[INFO] Running testBatchMLAPagedAttentionWrapper +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchMLAPagedAttentionWrapper --backends trtllm-gen-native fa2 fa3 --page_size 32 --batch_size 16 --s_qo 1 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_ckv 512 --head_dim_kpe 64 --random_actual_seq_len -vv --refcheck --q_dtype bfloat16 --kv_dtype bfloat16 --generate_repro_command --case_tag DeepSeek-R1 +[WARNING] fa3 for routine BatchMLAPagedAttentionWrapper is not supported on compute capability 10.0. Skipping. +[VERBOSE] Average actual seq len: 501 +[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([ 84, 874, 167, 691, 274, 736, 63, 813, 781, 450, 794, 226, 510, 499, + 524, 541], device='cuda:0', dtype=torch.int32) +[VVERBOSE] q_nope.shape = torch.Size([16, 128, 512]) +[VVERBOSE] q_pe.shape = torch.Size([16, 128, 64]) +[VVERBOSE] q.shape = torch.Size([16, 128, 576]) +[VVERBOSE] num_pages_per_seq = 32 +[VVERBOSE] total_num_pages = 512 +[VVERBOSE] block_tables.shape = torch.Size([16, 32]) +[VVERBOSE] ckv_cache.shape = torch.Size([512, 32, 512]) +[VVERBOSE] kpe_cache.shape = torch.Size([512, 32, 64]) +[VVERBOSE] kv_cache.shape = torch.Size([512, 32, 576]) +[VVERBOSE] qo_indptr.shape = torch.Size([17]) +[VVERBOSE] kv_indptr.shape = torch.Size([17]) +[VVERBOSE] kv_indices.shape = torch.Size([258]) +[VVERBOSE] actual_seq_lens_kv.shape = torch.Size([16, 1, 1, 1]) +[VVERBOSE] sm_scale = 0.041666666666666664 +[VVERBOSE] workspace_buffer.shape = torch.Size([134217728]) +[PERF] trtllm-gen-nati:: median time 0.024 ms; std 0.000 ms; achieved tflops 91.551 TFLOPs/sec; achieved tb_per_sec 0.955 TB/sec +[PERF] fa2 :: median time 0.041 ms; std 0.000 ms; achieved tflops 54.584 TFLOPs/sec; achieved tb_per_sec 0.570 TB/sec +[INFO] args = Namespace(routine='bmm_fp8', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=256, m=1, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cublas', 'cutlass'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testBmmFp8 +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +[VVERBOSE] input_fp8.shape = torch.Size([256, 1, 7168]) +[VVERBOSE] input_fp8.dtype = torch.float8_e4m3fn +[VVERBOSE] mat2_fp8.shape = torch.Size([256, 7168, 1024]) +[VVERBOSE] mat2_fp8.dtype = torch.float8_e4m3fn +[VVERBOSE] input_inv_s = tensor(0.0109, device='cuda:0') +[VVERBOSE] input_inv_s.dtype = torch.float32 +[VVERBOSE] mat2_inv_s = tensor(0.0135, device='cuda:0') +[VVERBOSE] mat2_inv_s.dtype = torch.float32 +[PERF] cudnn :: median time 0.286 ms; std 0.000 ms; achieved tflops 13.138 TFLOPs/sec; achieved tb_per_sec 0.026 TB/sec +[PERF] cublas :: median time 0.286 ms; std 0.000 ms; achieved tflops 13.140 TFLOPs/sec; achieved tb_per_sec 0.026 TB/sec +[PERF] cutlass :: median time 0.266 ms; std 0.000 ms; achieved tflops 14.142 TFLOPs/sec; achieved tb_per_sec 0.028 TB/sec +[INFO] args = Namespace(routine='bmm_fp8', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=64, m=4, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cublas', 'cutlass'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testBmmFp8 +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 64 --m 4 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command +[VVERBOSE] input_fp8.shape = torch.Size([64, 4, 7168]) +[VVERBOSE] input_fp8.dtype = torch.float8_e4m3fn +[VVERBOSE] mat2_fp8.shape = torch.Size([64, 7168, 1024]) +[VVERBOSE] mat2_fp8.dtype = torch.float8_e4m3fn +[VVERBOSE] input_inv_s = tensor(0.0109, device='cuda:0') +[VVERBOSE] input_inv_s.dtype = torch.float32 +[VVERBOSE] mat2_inv_s = tensor(0.0131, device='cuda:0') +[VVERBOSE] mat2_inv_s.dtype = torch.float32 +[PERF] cudnn :: median time 0.075 ms; std 0.000 ms; achieved tflops 49.999 TFLOPs/sec; achieved tb_per_sec 0.098 TB/sec +[PERF] cublas :: median time 0.075 ms; std 0.000 ms; achieved tflops 50.135 TFLOPs/sec; achieved tb_per_sec 0.098 TB/sec +[PERF] cutlass :: median time 0.072 ms; std 0.000 ms; achieved tflops 51.981 TFLOPs/sec; achieved tb_per_sec 0.102 TB/sec +[INFO] args = Namespace(routine='gemm_fp8_nt_groupwise', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=4, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cutlass'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testGemmFp8NtGroupwise +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine gemm_fp8_nt_groupwise --m 4 --n 1024 --k 7168 --mma_sm 1 --scale_major_mode MN --backends cutlass --refcheck -vv --generate_repro_command +[VVERBOSE] a_val.shape = torch.Size([4, 7168]) +[VVERBOSE] b_val.shape = torch.Size([1024, 7168]) +[VVERBOSE] a_fp8.shape = torch.Size([4, 7168]) +[VVERBOSE] b_fp8.shape = torch.Size([1024, 7168]) +[VVERBOSE] a_scale.shape = torch.Size([56, 4]) +[VVERBOSE] b_scale.shape = torch.Size([56, 8]) +[PERF] cutlass :: median time 0.020 ms; std 0.000 ms; achieved tflops 2.987 TFLOPs/sec; achieved tb_per_sec 0.375 TB/sec +[INFO] args = Namespace(routine='gemm_fp8_nt_groupwise', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=16, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cutlass'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testGemmFp8NtGroupwise +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine gemm_fp8_nt_groupwise --m 16 --n 1024 --k 7168 --mma_sm 1 --scale_major_mode MN --backends cutlass --refcheck -vv --generate_repro_command +[VVERBOSE] a_val.shape = torch.Size([16, 7168]) +[VVERBOSE] b_val.shape = torch.Size([1024, 7168]) +[VVERBOSE] a_fp8.shape = torch.Size([16, 7168]) +[VVERBOSE] b_fp8.shape = torch.Size([1024, 7168]) +[VVERBOSE] a_scale.shape = torch.Size([56, 16]) +[VVERBOSE] b_scale.shape = torch.Size([56, 8]) +[PERF] cutlass :: median time 0.020 ms; std 0.000 ms; achieved tflops 11.824 TFLOPs/sec; achieved tb_per_sec 0.377 TB/sec +[INFO] args = Namespace(routine='group_gemm_fp8_nt_groupwise', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=4, n=1024, k=7168, tile_size=128, group_size=2, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testGroupGemmFp8NtGroupwise +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine group_gemm_fp8_nt_groupwise --m 4 --n 1024 --k 7168 --mma_sm 1 --group_size 2 --scale_major_mode MN --refcheck -vv --generate_repro_command +[VVERBOSE] a_val.shape = torch.Size([8, 7168]) +[VVERBOSE] b_val.shape = torch.Size([2, 1024, 7168]) +[VVERBOSE] a_fp8.shape = torch.Size([8, 7168]) +[VVERBOSE] b_fp8.shape = torch.Size([2, 1024, 7168]) +[VVERBOSE] a_scale.shape = torch.Size([56, 8]) +[VVERBOSE] b_scale.shape = torch.Size([2, 56, 8]) +[VVERBOSE] m_indptr.shape = torch.Size([3]) +[PERF] cutlass :: median time 0.022 ms; std 0.000 ms; achieved tflops 5.261 TFLOPs/sec; achieved tb_per_sec 0.661 TB/sec +[INFO] args = Namespace(routine='group_gemm_fp8_nt_groupwise', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=16, n=1024, k=7168, tile_size=128, group_size=2, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False) +[INFO] Running testGroupGemmFp8NtGroupwise +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine group_gemm_fp8_nt_groupwise --m 16 --n 1024 --k 7168 --mma_sm 1 --group_size 2 --scale_major_mode MN --refcheck -vv --generate_repro_command +[VVERBOSE] a_val.shape = torch.Size([32, 7168]) +[VVERBOSE] b_val.shape = torch.Size([2, 1024, 7168]) +[VVERBOSE] a_fp8.shape = torch.Size([32, 7168]) +[VVERBOSE] b_fp8.shape = torch.Size([2, 1024, 7168]) +[VVERBOSE] a_scale.shape = torch.Size([56, 32]) +[VVERBOSE] b_scale.shape = torch.Size([2, 56, 8]) +[VVERBOSE] m_indptr.shape = torch.Size([3]) +[PERF] cutlass :: median time 0.023 ms; std 0.000 ms; achieved tflops 20.852 TFLOPs/sec; achieved tb_per_sec 0.665 TB/sec +[INFO] args = Namespace(routine='mm_fp4', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=1, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cutlass', 'trtllm'], use_128x4_sf_layout=True, use_nvfp4=True, autotune=False) +[INFO] Running testMmFp4 +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine mm_fp4 --m 1 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +[VVERBOSE] input_fp4.shape = torch.Size([1, 3584]) +[VVERBOSE] input_fp4.dtype = torch.uint8 +[VVERBOSE] mat2_fp4.shape = torch.Size([1024, 3584]) +[VVERBOSE] mat2_fp4.dtype = torch.uint8 +[PERF] cudnn :: median time 0.013 ms; std 0.000 ms; achieved tflops 1.156 TFLOPs/sec; achieved tb_per_sec 0.289 TB/sec +[PERF] cutlass :: median time 0.009 ms; std 0.000 ms; achieved tflops 1.593 TFLOPs/sec; achieved tb_per_sec 0.399 TB/sec +[PERF] trtllm :: median time 0.011 ms; std 0.000 ms; achieved tflops 1.352 TFLOPs/sec; achieved tb_per_sec 0.339 TB/sec +[INFO] args = Namespace(routine='mm_fp4', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=4, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cutlass', 'trtllm'], use_128x4_sf_layout=True, use_nvfp4=True, autotune=False) +[INFO] Running testMmFp4 +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --refcheck -vv --generate_repro_command +[VVERBOSE] input_fp4.shape = torch.Size([4, 3584]) +[VVERBOSE] input_fp4.dtype = torch.uint8 +[VVERBOSE] mat2_fp4.shape = torch.Size([1024, 3584]) +[VVERBOSE] mat2_fp4.dtype = torch.uint8 +[PERF] cudnn :: median time 0.013 ms; std 0.000 ms; achieved tflops 4.625 TFLOPs/sec; achieved tb_per_sec 0.291 TB/sec +[PERF] cutlass :: median time 0.009 ms; std 0.000 ms; achieved tflops 6.372 TFLOPs/sec; achieved tb_per_sec 0.401 TB/sec +[PERF] trtllm :: median time 0.011 ms; std 0.000 ms; achieved tflops 5.310 TFLOPs/sec; achieved tb_per_sec 0.334 TB/sec +[INFO] args = Namespace(routine='mm_fp4', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=1, m=4, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cutlass', 'trtllm'], use_128x4_sf_layout=True, use_nvfp4=True, autotune=True) +[INFO] Running testMmFp4 +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine mm_fp4 --m 4 --n 1024 --k 7168 --out_dtype bfloat16 --backends cudnn cutlass trtllm --use_128x4_sf_layout --use_nvfp4 --autotune --refcheck -vv --generate_repro_command +[INFO] cudnn backend does not support autotune +[VVERBOSE] input_fp4.shape = torch.Size([4, 3584]) +[VVERBOSE] input_fp4.dtype = torch.uint8 +[VVERBOSE] mat2_fp4.shape = torch.Size([1024, 3584]) +[VVERBOSE] mat2_fp4.dtype = torch.uint8 +[INFO] Autotune warmup for mm_fp4: 5 iters +2025-09-23 00:32:18,077 - INFO - autotuner.py:256 - flashinfer.jit: [Autotuner]: Autotuning process starts ... +2025-09-23 00:32:18,224 - INFO - autotuner.py:262 - flashinfer.jit: [Autotuner]: Autotuning process ends +[INFO] Autotune warmup for mm_fp4: 5 iters +2025-09-23 00:32:18,225 - INFO - autotuner.py:256 - flashinfer.jit: [Autotuner]: Autotuning process starts ... +2025-09-23 00:32:18,247 - INFO - autotuner.py:262 - flashinfer.jit: [Autotuner]: Autotuning process ends +[PERF] cutlass_autotun:: median time 0.009 ms; std 0.000 ms; achieved tflops 6.372 TFLOPs/sec; achieved tb_per_sec 0.401 TB/sec +[PERF] trtllm_autotune:: median time 0.011 ms; std 0.000 ms; achieved tflops 5.410 TFLOPs/sec; achieved tb_per_sec 0.340 TB/sec +[INFO] args = Namespace(routine='trtllm_fp4_block_scale_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='trtllm_moe_sample', generate_repro_command=True, repro_command='', num_tokens=1024, hidden_size=1024, intermediate_size=1024, num_experts=256, top_k=8, n_group=8, topk_group=4, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=True, weight_layout=0, use_routing_bias=True, use_routing_scales_on_input=False, input_dtype='bfloat16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testTrtllmFp4BlockScaleMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine trtllm_fp4_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 256 --top_k 8 --n_group 8 --topk_group 4 --routed_scaling_factor 2.5 --use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +[INFO] Configuration: tokens=1024, hidden=1024, intermediate=1024, experts=256, top_k=8 +[VVERBOSE] routing_logits.shape = torch.Size([1024, 256]) +[VVERBOSE] hidden_states.shape = torch.Size([1024, 1024]) +[VVERBOSE] gemm1_weights_fp4.shape = torch.Size([256, 2048, 512]) +[VVERBOSE] gemm2_weights_fp4.shape = torch.Size([256, 1024, 512]) +[PERF] trtllm :: median time 0.224 ms; std 0.000 ms; achieved tflops 230.555 TFLOPs/sec; achieved tb_per_sec 1.818 TB/sec +[INFO] args = Namespace(routine='trtllm_fp4_block_scale_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='trtllm_moe_sample', generate_repro_command=True, repro_command='', num_tokens=1024, hidden_size=1024, intermediate_size=1024, num_experts=128, top_k=8, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='renormalize_naive', use_shuffled_weight=True, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='bfloat16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=4, gated_act_type=0) +[INFO] Running testTrtllmFp4BlockScaleMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine trtllm_fp4_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 8 --routing_method renormalize_naive --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +[INFO] Configuration: tokens=1024, hidden=1024, intermediate=1024, experts=128, top_k=8 +[VVERBOSE] routing_logits.shape = torch.Size([1024, 128]) +[VVERBOSE] hidden_states.shape = torch.Size([1024, 1024]) +[VVERBOSE] gemm1_weights_fp4.shape = torch.Size([128, 2048, 512]) +[VVERBOSE] gemm2_weights_fp4.shape = torch.Size([128, 1024, 512]) +[PERF] trtllm :: median time 0.226 ms; std 0.000 ms; achieved tflops 227.846 TFLOPs/sec; achieved tb_per_sec 0.903 TB/sec +[INFO] args = Namespace(routine='trtllm_fp8_block_scale_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='trtllm_moe_sample', generate_repro_command=True, repro_command='', num_tokens=1024, hidden_size=1024, intermediate_size=1024, num_experts=256, top_k=8, n_group=8, topk_group=4, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=True, weight_layout=0, use_routing_bias=True, use_routing_scales_on_input=False, input_dtype='bfloat16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testTrtllmFp8BlockScaleMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine trtllm_fp8_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 256 --top_k 8 --n_group 8 --topk_group 4 --routed_scaling_factor 2.5 --use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +[INFO] Configuration: tokens=1024, hidden=1024, intermediate=1024, experts=256, top_k=8 +[VVERBOSE] routing_logits.shape = torch.Size([1024, 256]) +[VVERBOSE] hidden_states.shape = torch.Size([1024, 1024]) +[VVERBOSE] gemm1_weights_fp8.shape = torch.Size([256, 2048, 1024]) +[VVERBOSE] gemm2_weights_fp8.shape = torch.Size([256, 1024, 1024]) +[PERF] trtllm :: median time 0.557 ms; std 0.000 ms; achieved tflops 92.607 TFLOPs/sec; achieved tb_per_sec 1.455 TB/sec +[INFO] args = Namespace(routine='trtllm_fp8_per_tensor_scale_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='trtllm_moe_sample', generate_repro_command=True, repro_command='', num_tokens=1024, hidden_size=1024, intermediate_size=1024, num_experts=128, top_k=1, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='llama4', use_shuffled_weight=False, weight_layout=0, use_routing_bias=True, use_routing_scales_on_input=True, input_dtype='bfloat16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=3, gated_act_type=0) +[INFO] Running testTrtllmFp8PerTensorScaleMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine trtllm_fp8_per_tensor_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 1 --routed_scaling_factor 2.5 --use_routing_bias --routing_method llama4 --use_routing_scales_on_input -vv --generate_repro_command --case_tag trtllm_moe_sample +[INFO] Configuration: tokens=1024, hidden=1024, intermediate=1024, experts=128, top_k=1 +[VVERBOSE] routing_logits.shape = torch.Size([1024, 128]) +[VVERBOSE] hidden_states.shape = torch.Size([1024, 1024]) +[VVERBOSE] gemm1_weights_fp8.shape = torch.Size([128, 2048, 1024]) +[VVERBOSE] gemm2_weights_fp8.shape = torch.Size([128, 1024, 1024]) +[PERF] trtllm :: median time 0.123 ms; std 0.000 ms; achieved tflops 52.340 TFLOPs/sec; achieved tb_per_sec 3.299 TB/sec +[INFO] args = Namespace(routine='trtllm_fp8_block_scale_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='trtllm_moe_sample', generate_repro_command=True, repro_command='', num_tokens=1024, hidden_size=1024, intermediate_size=1024, num_experts=128, top_k=1, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='renormalize', use_shuffled_weight=True, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='bfloat16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=1, gated_act_type=0) +[INFO] Running testTrtllmFp8BlockScaleMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine trtllm_fp8_block_scale_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 1024 --num_experts 128 --top_k 1 --routing_method renormalize --use_shuffled_weight -vv --generate_repro_command --case_tag trtllm_moe_sample +[INFO] Configuration: tokens=1024, hidden=1024, intermediate=1024, experts=128, top_k=1 +[VVERBOSE] routing_logits.shape = torch.Size([1024, 128]) +[VVERBOSE] hidden_states.shape = torch.Size([1024, 1024]) +[VVERBOSE] gemm1_weights_fp8.shape = torch.Size([128, 2048, 1024]) +[VVERBOSE] gemm2_weights_fp8.shape = torch.Size([128, 1024, 1024]) +[PERF] trtllm :: median time 0.109 ms; std 0.000 ms; achieved tflops 59.297 TFLOPs/sec; achieved tb_per_sec 3.740 TB/sec +[INFO] args = Namespace(routine='cutlass_fused_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='cutlass_moe_base', generate_repro_command=True, repro_command='', num_tokens=32, hidden_size=128, intermediate_size=128, num_experts=2, top_k=2, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=False, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='float16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testCutlassFusedMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant base --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_base +[VVERBOSE] x.shape = torch.Size([32, 128]) +[VVERBOSE] w31_weight.shape = torch.Size([2, 256, 128]) +[VVERBOSE] w2_weight.shape = torch.Size([2, 128, 128]) +[PERF] cutlass :: median time 0.026 ms; std 0.000 ms; achieved tflops 0.240 TFLOPs/sec; achieved tb_per_sec 0.008 TB/sec +[INFO] args = Namespace(routine='cutlass_fused_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='cutlass_moe_fp8_scale', generate_repro_command=True, repro_command='', num_tokens=32, hidden_size=128, intermediate_size=128, num_experts=2, top_k=2, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=False, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='float16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='fp8', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testCutlassFusedMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant fp8 --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_fp8_scale +[VVERBOSE] x.shape = torch.Size([32, 128]) +[VVERBOSE] w31_weight.shape = torch.Size([2, 256, 128]) +[VVERBOSE] w2_weight.shape = torch.Size([2, 128, 128]) +[PERF] cutlass :: median time 0.026 ms; std 0.000 ms; achieved tflops 0.244 TFLOPs/sec; achieved tb_per_sec 0.004 TB/sec +[INFO] args = Namespace(routine='cutlass_fused_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='cutlass_moe_nvfp4_weights', generate_repro_command=True, repro_command='', num_tokens=32, hidden_size=128, intermediate_size=128, num_experts=2, top_k=2, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=False, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='float16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='nvfp4', quantized_input=False, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testCutlassFusedMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant nvfp4 --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_nvfp4_weights +[VVERBOSE] x.shape = torch.Size([32, 128]) +[VVERBOSE] w31_weight.shape = torch.Size([2, 256, 128]) +[VVERBOSE] w2_weight.shape = torch.Size([2, 128, 128]) +[PERF] cutlass :: median time 0.030 ms; std 0.000 ms; achieved tflops 0.210 TFLOPs/sec; achieved tb_per_sec 0.002 TB/sec +[INFO] args = Namespace(routine='cutlass_fused_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='cutlass_moe_nvfp4_weights_quantized', generate_repro_command=True, repro_command='', num_tokens=32, hidden_size=128, intermediate_size=128, num_experts=2, top_k=2, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=False, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='float16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='nvfp4', quantized_input=True, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testCutlassFusedMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 2 --top_k 2 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 -vv --generate_repro_command --case_tag cutlass_moe_nvfp4_weights_quantized +[VVERBOSE] x.shape = torch.Size([32, 128]) +[VVERBOSE] w31_weight.shape = torch.Size([2, 256, 128]) +[VVERBOSE] w2_weight.shape = torch.Size([2, 128, 128]) +[PERF] cutlass :: median time 0.029 ms; std 0.000 ms; achieved tflops 0.213 TFLOPs/sec; achieved tb_per_sec 0.002 TB/sec +[INFO] args = Namespace(routine='cutlass_fused_moe', no_cuda_graph=False, use_cupti=False, refcheck=False, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='cutlass_moe_nvfp4_ep_tp', generate_repro_command=True, repro_command='', num_tokens=32, hidden_size=128, intermediate_size=128, num_experts=8, top_k=2, n_group=None, topk_group=None, routed_scaling_factor=2.5, local_expert_offset=0, local_num_experts=None, tile_tokens_dim=8, routing_method='deepseek_v3', use_shuffled_weight=False, weight_layout=0, use_routing_bias=False, use_routing_scales_on_input=False, input_dtype='float16', weight_dtype='bfloat16', gated_act='swiglu', autotune=False, cutlass_variant='base', quantized_input=False, tp_size=2, tp_rank=0, ep_size=4, ep_rank=0, routing_method_type=2, gated_act_type=0) +[INFO] Running testCutlassFusedMoe +[INFO] FlashInfer version: 0.3.1 +[VVERBOSE] gpu_name = 'NVIDIA_B200' +[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine cutlass_fused_moe --num_tokens 32 --hidden_size 128 --intermediate_size 128 --num_experts 8 --top_k 2 --cutlass_variant base --input_dtype float16 --tp_size 2 --tp_rank 0 --ep_size 4 --ep_rank 0 -vv --generate_repro_command --case_tag cutlass_moe_nvfp4_ep_tp +[VVERBOSE] x.shape = torch.Size([32, 128]) +[VVERBOSE] w31_weight.shape = torch.Size([8, 256, 128]) +[VVERBOSE] w2_weight.shape = torch.Size([8, 128, 128]) +[PERF] cutlass :: median time 0.025 ms; std 0.000 ms; achieved tflops 0.250 TFLOPs/sec; achieved tb_per_sec 0.032 TB/sec diff --git a/benchmarks/test_flashinfer_benchmark.py b/benchmarks/test_flashinfer_benchmark.py new file mode 100644 index 0000000000..4d9236cae0 --- /dev/null +++ b/benchmarks/test_flashinfer_benchmark.py @@ -0,0 +1,63 @@ +import flashinfer_benchmark +import pytest + + +@pytest.mark.parametrize("batch_size", [16, 32]) +@pytest.mark.parametrize("s_kv", [1024, 2048]) +@pytest.mark.parametrize("page_size", [8, 16]) +@pytest.mark.parametrize("is_cuda_graph_compatible", [False, True]) +def test_BatchDecodeWithPagedKVCacheWrapper_routine( + batch_size, s_kv, page_size, is_cuda_graph_compatible +): + args = flashinfer_benchmark.parse_args( + f"--routine BatchDecodeWithPagedKVCacheWrapper --backends fa2 fa2_tc --page_size {page_size} --batch_size {batch_size} --s_qo 1 --s_kv {s_kv} --num_qo_heads 64 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck {'--no_cuda_graph' if not is_cuda_graph_compatible else ''}".split() + ) + flashinfer_benchmark.run_test(args) + + +@pytest.mark.parametrize("batch_size", [16, 32]) +@pytest.mark.parametrize("s_kv", [1024, 2048]) +@pytest.mark.parametrize("page_size", [8, 16]) +@pytest.mark.parametrize("is_cuda_graph_compatible", [False]) +def test_BatchPrefillWithPagedKVCacheWrapper_routine( + batch_size, s_kv, page_size, is_cuda_graph_compatible +): + args = flashinfer_benchmark.parse_args( + f"--routine BatchPrefillWithPagedKVCacheWrapper --backends fa2 --page_size {page_size} --batch_size {batch_size} --s_qo {s_kv} --s_kv {s_kv} --num_qo_heads 8 --num_kv_heads 8 --head_dim_qk 128 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal {'--no_cuda_graph' if not is_cuda_graph_compatible else ''}".split() + ) + flashinfer_benchmark.run_test(args) + + +@pytest.mark.parametrize("batch_size", [16, 32]) +@pytest.mark.parametrize("s_kv", [1024, 2048]) +@pytest.mark.parametrize("is_cuda_graph_compatible", [False]) +def test_BatchPrefillWithRaggedKVCacheWrapper_routine( + batch_size, s_kv, is_cuda_graph_compatible +): + args = flashinfer_benchmark.parse_args( + f"--routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 --batch_size {batch_size} --s_qo {s_kv} --s_kv {s_kv} --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 -vv --refcheck --causal {'--no_cuda_graph' if not is_cuda_graph_compatible else ''}".split() + ) + flashinfer_benchmark.run_test(args) + + +@pytest.mark.parametrize("m", [1024, 4096]) +@pytest.mark.parametrize("n", [1024]) +@pytest.mark.parametrize("k", [1024, 2048]) +@pytest.mark.parametrize("mma_sm", [1, 2]) +def test_gemm_fp8_nt_groupwise(m, n, k, mma_sm): + args = flashinfer_benchmark.parse_args( + f"--routine gemm_fp8_nt_groupwise --m {m} --n {n} --k {k} --mma_sm {mma_sm} --no_cuda_graph --refcheck -vv".split() + ) + flashinfer_benchmark.run_test(args) + + +@pytest.mark.parametrize("m", [1024, 4096]) +@pytest.mark.parametrize("n", [1024]) +@pytest.mark.parametrize("k", [1024, 2048]) +@pytest.mark.parametrize("mma_sm", [1, 2]) +@pytest.mark.parametrize("group_size", [1, 2]) +def test_group_gemm_fp8_nt_groupwise(m, n, k, mma_sm, group_size): + args = flashinfer_benchmark.parse_args( + f"--routine group_gemm_fp8_nt_groupwise --m {m} --n {n} --k {k} --mma_sm {mma_sm} --group_size {group_size} --no_cuda_graph --refcheck -vv".split() + ) + flashinfer_benchmark.run_test(args) diff --git a/build_backend.py b/build_backend.py new file mode 100644 index 0000000000..d14f5787f7 --- /dev/null +++ b/build_backend.py @@ -0,0 +1,179 @@ +""" +Copyright (c) 2023 by FlashInfer team. + +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 os +import shutil +from pathlib import Path + +from setuptools import build_meta as orig +from build_utils import get_git_version + +_root = Path(__file__).parent.resolve() +_data_dir = _root / "flashinfer" / "data" + + +def _create_build_metadata(): + """Create build metadata file with version information.""" + version_file = _root / "version.txt" + if version_file.exists(): + with open(version_file, "r") as f: + version = f.read().strip() + else: + version = "0.0.0+unknown" + + # Add dev suffix if specified + dev_suffix = os.environ.get("FLASHINFER_DEV_RELEASE_SUFFIX", "") + if dev_suffix: + version = f"{version}.dev{dev_suffix}" + + # Get git version + git_version = get_git_version(cwd=_root) + + # Append local version suffix if available + local_version = os.environ.get("FLASHINFER_LOCAL_VERSION") + if local_version: + # Use + to create a local version identifier that will appear in wheel name + version = f"{version}+{local_version}" + + # Create build metadata in the source tree + package_dir = Path(__file__).parent / "flashinfer" + build_meta_file = package_dir / "_build_meta.py" + + # Check if we're in a git repository + git_dir = Path(__file__).parent / ".git" + in_git_repo = git_dir.exists() + + # If file exists and not in git repo (installing from sdist), keep existing file + if build_meta_file.exists() and not in_git_repo: + print("Build metadata file already exists (not in git repo), keeping it") + return version + + # In git repo (editable) or file doesn't exist, create/update it + with open(build_meta_file, "w") as f: + f.write('"""Build metadata for flashinfer package."""\n') + f.write(f'__version__ = "{version}"\n') + f.write(f'__git_version__ = "{git_version}"\n') + + print(f"Created build metadata file with version {version}") + return version + + +# Create build metadata as soon as this module is imported +_create_build_metadata() + + +def write_if_different(path: Path, content: str) -> None: + if path.exists() and path.read_text() == content: + return + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content) + + +def _create_data_dir(use_symlinks=True): + _data_dir.mkdir(parents=True, exist_ok=True) + + def ln(source: str, target: str) -> None: + src = _root / source + dst = _data_dir / target + if dst.exists(): + if dst.is_symlink(): + dst.unlink() + elif dst.is_dir(): + shutil.rmtree(dst) + else: + dst.unlink() + + if use_symlinks: + dst.symlink_to(src, target_is_directory=True) + else: + # For wheel/sdist, copy actual files instead of symlinks + if src.exists(): + shutil.copytree(src, dst, symlinks=False, dirs_exist_ok=True) + + ln("3rdparty/cutlass", "cutlass") + ln("3rdparty/spdlog", "spdlog") + ln("csrc", "csrc") + ln("include", "include") + + +def _prepare_for_wheel(): + # For wheel, copy actual files instead of symlinks so they are included in the wheel + if _data_dir.exists(): + shutil.rmtree(_data_dir) + _create_data_dir(use_symlinks=False) + + # Copy license files from licenses/ to root to avoid nested path in wheel + licenses_dir = _root / "licenses" + if licenses_dir.exists(): + for license_file in licenses_dir.glob("*.txt"): + shutil.copy2( + license_file, + _root / f"LICENSE.{license_file.stem.removeprefix('LICENSE.')}.txt", + ) + + +def _prepare_for_editable(): + # For editable install, use symlinks so changes are reflected immediately + if _data_dir.exists(): + shutil.rmtree(_data_dir) + _create_data_dir(use_symlinks=True) + + +def _prepare_for_sdist(): + # For sdist, copy actual files instead of symlinks so they are included in the tarball + if _data_dir.exists(): + shutil.rmtree(_data_dir) + _create_data_dir(use_symlinks=False) + + +def get_requires_for_build_wheel(config_settings=None): + _prepare_for_wheel() + return [] + + +def get_requires_for_build_sdist(config_settings=None): + _prepare_for_sdist() + return [] + + +def get_requires_for_build_editable(config_settings=None): + _prepare_for_editable() + return [] + + +def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None): + _prepare_for_wheel() + return orig.prepare_metadata_for_build_wheel(metadata_directory, config_settings) + + +def prepare_metadata_for_build_editable(metadata_directory, config_settings=None): + _prepare_for_editable() + return orig.prepare_metadata_for_build_editable(metadata_directory, config_settings) + + +def build_editable(wheel_directory, config_settings=None, metadata_directory=None): + _prepare_for_editable() + return orig.build_editable(wheel_directory, config_settings, metadata_directory) + + +def build_sdist(sdist_directory, config_settings=None): + _prepare_for_sdist() + return orig.build_sdist(sdist_directory, config_settings) + + +def build_wheel(wheel_directory, config_settings=None, metadata_directory=None): + _prepare_for_wheel() + return orig.build_wheel(wheel_directory, config_settings, metadata_directory) diff --git a/build_utils.py b/build_utils.py new file mode 100644 index 0000000000..726a628204 --- /dev/null +++ b/build_utils.py @@ -0,0 +1,46 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +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. +""" + +"""Shared build utilities for flashinfer packages.""" + +import subprocess +from pathlib import Path +from typing import Optional + + +def get_git_version(cwd: Optional[Path] = None) -> str: + """ + Get git commit hash. + + Args: + cwd: Working directory for git command. If None, uses current directory. + + Returns: + Git commit hash or "unknown" if git is not available. + """ + try: + git_version = ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=cwd, + stderr=subprocess.DEVNULL, + ) + .decode("ascii") + .strip() + ) + return git_version + except Exception: + return "unknown" diff --git a/ci/bash.sh b/ci/bash.sh index 437628a5ed..5ced84e367 100644 --- a/ci/bash.sh +++ b/ci/bash.sh @@ -70,6 +70,10 @@ echo "ENV VARIABLES: ${DOCKER_ENV}" echo "VOLUMES: ${DOCKER_VOLUMNS}" echo "COMMANDS: '${COMMAND[@]}'" +# Pull the latest docker image +echo "Pulling latest image: ${DOCKER_IMAGE_NAME}" +${DOCKER_BINARY} pull ${DOCKER_IMAGE_NAME} + # By default we cleanup - remove the container once it finish running (--rm) # and share the PID namespace (--pid=host) so the process inside does not have # pid 1 and SIGKILL is propagated to the process inside (jenkins can kill it). diff --git a/ci/docker-tags.yml b/ci/docker-tags.yml new file mode 100644 index 0000000000..3619e1bf5b --- /dev/null +++ b/ci/docker-tags.yml @@ -0,0 +1,4 @@ +flashinfer/flashinfer-ci-cu126: 20251104-d528f0c +flashinfer/flashinfer-ci-cu128: 20251104-d528f0c +flashinfer/flashinfer-ci-cu129: 20251104-d528f0c +flashinfer/flashinfer-ci-cu130: 20251104-d528f0c diff --git a/ci/scripts/jenkins/git_utils.py b/ci/scripts/jenkins/git_utils.py index 114375d571..221da61b5d 100644 --- a/ci/scripts/jenkins/git_utils.py +++ b/ci/scripts/jenkins/git_utils.py @@ -30,13 +30,13 @@ def compress_query(query: str) -> str: query = query.replace("\n", "") - query = re.sub("\s+", " ", query) + query = re.sub(r"\s+", " ", query) return query def post(url: str, body: Optional[Any] = None, auth: Optional[Tuple[str, str]] = None): logging.info(f"Requesting POST to {url} with {body}") - headers = {} + headers: Dict[Any, Any] = {} req = request.Request(url, headers=headers, method="POST") if auth is not None: auth_str = base64.b64encode(f"{auth[0]}:{auth[1]}".encode()) @@ -46,9 +46,8 @@ def post(url: str, body: Optional[Any] = None, auth: Optional[Tuple[str, str]] = body = "" req.add_header("Content-Type", "application/json; charset=utf-8") - data = json.dumps(body) - data = data.encode("utf-8") - req.add_header("Content-Length", len(data)) + data = json.dumps(body).encode("utf-8") + req.add_header("Content-Length", str(len(data))) with request.urlopen(req, data) as response: return response.read() @@ -119,9 +118,8 @@ def _request( logging.info(f"Requesting {method} to {full_url} with {body}") req = request.Request(full_url, headers=self.headers(), method=method.upper()) req.add_header("Content-Type", "application/json; charset=utf-8") - data = json.dumps(body) - data = data.encode("utf-8") - req.add_header("Content-Length", len(data)) + data = json.dumps(body).encode("utf-8") + req.add_header("Content-Length", str(len(data))) try: with request.urlopen(req, data) as response: @@ -129,12 +127,12 @@ def _request( except error.HTTPError as e: msg = str(e) error_data = e.read().decode() - raise RuntimeError(f"Error response: {msg}\n{error_data}") + raise RuntimeError(f"Error response: {msg}\n{error_data}") from e logging.info(f"Got response from {full_url}: {content}") try: response = json.loads(content) - except json.decoder.JSONDecodeError as e: + except json.decoder.JSONDecodeError: return content return response @@ -206,12 +204,11 @@ def find_ccs(body: str) -> List[str]: matches = re.findall(r"(cc( @[-A-Za-z0-9]+)+)", body, flags=re.MULTILINE) matches = [full for full, last in matches] - reviewers = [] + reviewers = set() for match in matches: if match.startswith("cc "): match = match.replace("cc ", "") users = [x.strip() for x in match.split("@")] - reviewers += users + reviewers.update(users) - reviewers = set(x for x in reviewers if x != "") - return list(reviewers) + return [x for x in reviewers if x != ""] diff --git a/csrc/activation.cu b/csrc/activation.cu deleted file mode 100644 index 0e9cef7c41..0000000000 --- a/csrc/activation.cu +++ /dev/null @@ -1,128 +0,0 @@ -/* - * Copyright (c) 2024 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include - -#include "pytorch_extension_utils.h" - -using namespace flashinfer; - -__device__ __forceinline__ float silu(const float& val) { return val / (1.0f + __expf(-val)); } - -__device__ __forceinline__ float gelu(const float& val) { - constexpr float kAlpha = M_SQRT1_2; - return val * 0.5f * (1.0f + ::erf(val * kAlpha)); -} - -__device__ __forceinline__ float gelu_tanh(const float& val) { - const float cdf = - 0.5f * (1.0f + math::tanh((0.7978845608028654f * (val + 0.044715f * val * val * val)))); - return val * cdf; -} - -void silu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { - int d = input.size(-1) / 2; - int64_t num_tokens = input.numel() / input.size(-1); - - const c10::cuda::OptionalCUDAGuard device_guard(out.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - - DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { - uint32_t vec_size = 16 / sizeof(c_type); - cudaLaunchConfig_t config; - config.gridDim = num_tokens; - config.blockDim = std::min(d / vec_size, 1024U); - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; - config.numAttrs = 1; - config.attrs = attrs; - - auto kernel = flashinfer::activation::act_and_mul_kernel; - - cudaLaunchKernelEx(&config, kernel, static_cast(out.data_ptr()), - static_cast(input.data_ptr()), d); - - cudaError_t err = cudaGetLastError(); - TORCH_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); - - return true; - }); -} - -void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { - int d = input.size(-1) / 2; - int64_t num_tokens = input.numel() / input.size(-1); - - const c10::cuda::OptionalCUDAGuard device_guard(out.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - - DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { - uint32_t vec_size = 16 / sizeof(c_type); - cudaLaunchConfig_t config; - config.gridDim = num_tokens; - config.blockDim = std::min(d / vec_size, 1024U); - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; - config.numAttrs = 1; - config.attrs = attrs; - - auto kernel = flashinfer::activation::act_and_mul_kernel; - - cudaLaunchKernelEx(&config, kernel, static_cast(out.data_ptr()), - static_cast(input.data_ptr()), d); - - cudaError_t err = cudaGetLastError(); - TORCH_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); - - return true; - }); -} - -void gelu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { - int d = input.size(-1) / 2; - int64_t num_tokens = input.numel() / input.size(-1); - const c10::cuda::OptionalCUDAGuard device_guard(out.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - - DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { - uint32_t vec_size = 16 / sizeof(c_type); - cudaLaunchConfig_t config; - config.gridDim = num_tokens; - config.blockDim = std::min(d / vec_size, 1024U); - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; - config.numAttrs = 1; - config.attrs = attrs; - - auto kernel = flashinfer::activation::act_and_mul_kernel; - - cudaLaunchKernelEx(&config, kernel, static_cast(out.data_ptr()), - static_cast(input.data_ptr()), d); - - cudaError_t err = cudaGetLastError(); - TORCH_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); - - return true; - }); -} diff --git a/csrc/aot_extension_utils.h b/csrc/aot_extension_utils.h deleted file mode 100644 index acc9ddc5f0..0000000000 --- a/csrc/aot_extension_utils.h +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright (c) 2024 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "generated/dispatch.inc" -#include "pytorch_extension_utils.h" - -#define DISPATCH_head_dim(expr, const_expr, ...) \ - _DISPATCH_SWITCH("head_dim", expr, _DISPATCH_CASES_head_dim(const_expr, __VA_ARGS__)) - -#define DISPATCH_head_dim_sm90(expr1, expr2, const_expr1, const_expr2, ...) \ - _DISPATCH_SWITCH_U16x2("head_dim_qk", "head_dim_vo", expr1, expr2, \ - _DISPATCH_CASES_head_dim_sm90(const_expr1, const_expr2, __VA_ARGS__)) - -#define DISPATCH_pos_encoding_mode(expr, const_expr, ...) \ - _DISPATCH_SWITCH("positional encoding mode", expr, \ - _DISPATCH_CASES_pos_encoding_mode(const_expr, __VA_ARGS__)) - -#define DISPATCH_use_fp16_qk_reduction(expr, const_expr, ...) \ - _DISPATCH_SWITCH("use_fp16_qk_reduction", expr, \ - _DISPATCH_CASES_use_fp16_qk_reduction(const_expr, __VA_ARGS__)) - -#define DISPATCH_mask_mode(expr, const_expr, ...) \ - _DISPATCH_SWITCH("mask_mode", expr, _DISPATCH_CASES_mask_mode(const_expr, __VA_ARGS__)) - -#define DISPATCH_PYTORCH_QKV_DTYPE_TO_CTYPE(q_dtype, kv_dtype, c_type_q, c_type_kv, ...) \ - [&]() -> bool { \ - if (kv_dtype == q_dtype) { \ - return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(q_dtype, c_type_q, [&] { \ - using c_type_kv = c_type_q; \ - return __VA_ARGS__(); \ - }); \ - } else { \ - return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(q_dtype, c_type_q, [&] { \ - return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8(kv_dtype, c_type_kv, \ - [&] { return __VA_ARGS__(); }); \ - }); \ - } \ - }() diff --git a/csrc/batch_attention.cu b/csrc/batch_attention.cu new file mode 100644 index 0000000000..a3d36b7981 --- /dev/null +++ b/csrc/batch_attention.cu @@ -0,0 +1,187 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_attention_config.inc" +#include "tvm_ffi_utils.h" + +namespace flashinfer { + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +template +cudaError_t BatchPagedAttentionPersistent(const Params params_1, const Params params_2, + const uint32_t num_blks_x, const uint32_t num_blks_y, + const cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +Array BatchPagedAttentionPlan(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t head_dim_o, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); + + HolisticPlanInfo<2> plan_info; + + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); + + cudaError_t status = TwoStageHolisticPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), + batch_size, num_qo_heads, num_kv_heads, head_dim_o, causal, stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to plan persistent paged attention, error: " << cudaGetErrorString(status); + + return Array(plan_info.ToVector()); +} + +void BatchPagedAttentionRun(TensorView float_workspace_buffer, TensorView int_workspace_buffer, + Array plan_info_vec, TensorView q, TensorView k_cache, + TensorView v_cache, TensorView kv_indices, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t layout_code, int64_t num_qo_heads, int64_t num_kv_heads, + int64_t page_size, + double v_scale, // must use double due to pytorch binding + double sm_scale, + double logits_soft_cap ADDITIONAL_FUNC_PARAMS PROFILER_FUNC_PARAMS) { + HolisticPlanInfo<2> plan_info; + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + // NOTE (Yilong): assume both q and o are NHD + unsigned int q_stride_n = q.stride(0); + unsigned int q_stride_h = q.stride(1); + + // layout only constraint paged KV + const QKVLayout kv_layout = static_cast(layout_code); + unsigned int k_stride_page = k_cache.stride(0); + unsigned int v_stride_page = v_cache.stride(0); + unsigned int k_stride_n, k_stride_h, v_stride_n, v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_h = k_cache.stride(2); + k_stride_n = k_cache.stride(1); + v_stride_h = v_cache.stride(2); + v_stride_n = v_cache.stride(1); + } else { + k_stride_h = k_cache.stride(1); + k_stride_n = k_cache.stride(2); + v_stride_h = v_cache.stride(1); + v_stride_n = v_cache.stride(2); + } + + cudaSetDevice(q.device().device_id); + const cudaStream_t stream = get_stream(q.device()); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + AttentionVariant, PersistentParams, [&] { + PersistentParams params[2]; + IdType* len_kv_chunk = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.len_kv_chunk_offset); + for (int i = 0; i < 2; i++) { + params[i].q = static_cast(q.data_ptr()); + params[i].k = static_cast(k_cache.data_ptr()); + params[i].v = static_cast(v_cache.data_ptr()); + + params[i].q_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].q_indptr_offset); + params[i].kv_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].kv_indptr_offset); + params[i].partial_indptr = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.tasks[i].partial_indptr_offset); + params[i].kv_indices = static_cast(kv_indices.data_ptr()); + params[i].q_len = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].q_len_offset); + params[i].kv_len = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].kv_len_offset); + params[i].q_start = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].q_start_offset); + params[i].kv_start = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].kv_start_offset); + params[i].kv_end = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].kv_end_offset); + params[i].kv_head_idx_arr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].kv_head_idx_offset); + params[i].work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.tasks[i].work_indptr_offset); + params[i].len_kv_chunk = len_kv_chunk + i; + + params[i].final_o = static_cast(o.data_ptr()); + params[i].final_lse = + maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr; + params[i].partial_o = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); + params[i].partial_lse = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_lse_offset); + + // for state reduction + params[i].merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + params[i].merge_o_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_o_indices_offset); + params[i].num_packed_qo_len = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.num_qo_len_offset); + + params[i].num_kv_heads = num_kv_heads; + params[i].gqa_group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params[i].page_size = uint_fastdiv(page_size); + + params[i].q_stride_n = q_stride_n; + params[i].q_stride_h = q_stride_h; + params[i].k_stride_page = k_stride_page; + params[i].k_stride_h = k_stride_h; + params[i].k_stride_n = k_stride_n; + params[i].v_stride_page = v_stride_page; + params[i].v_stride_h = v_stride_h; + params[i].v_stride_n = v_stride_n; + + params[i].sm_scale = sm_scale; + params[i].v_scale = v_scale; + params[i].logits_soft_cap = logits_soft_cap; + // NOTE(Wenxuan) directly using the additional_params_decl from generate_additional_params + // will be problematic because of the params[i] + ADDITIONAL_PARAMS_SETTER + PROFILER_PARAMS_SETTER + } + + cudaError_t status = BatchPagedAttentionPersistent<128, 16, HEAD_DIM_QK, HEAD_DIM_VO, + MASK_MODE, AttentionVariant>( + params[0], params[1], plan_info.num_blks_x, plan_info.num_blks_y, stream); + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to run persistent paged attention, error: " << cudaGetErrorString(status); + return true; + }); +} diff --git a/csrc/batch_attention_customize_config.jinja b/csrc/batch_attention_customize_config.jinja new file mode 100644 index 0000000000..3cf9312748 --- /dev/null +++ b/csrc/batch_attention_customize_config.jinja @@ -0,0 +1,117 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace flashinfer; + +#define ADDITIONAL_FUNC_PARAMS {{ additional_func_params }} +#define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }} + +#ifdef FLASHINFER_ENABLE_PROFILER +#define PROFILER_PARAMS_SETTER \ + params[i].profiler_buffer = static_cast(profiler_buffer.data_ptr()); +#else +#define PROFILER_PARAMS_SETTER +#endif + +{{ variant_decl }} + +template +struct StandardAttention : AttentionVariantBase { + float sm_scale_log2; + float soft_cap_pre_tanh_scale; + static constexpr bool use_logits_soft_cap = UseLogitsSoftCap; + PROFILER_CLOSURE_PARAMS_DECL + + template + __device__ __host__ StandardAttention(const Params& params, uint32_t batch_idx, + uint8_t* smem_ptr) { + if constexpr (UseLogitsSoftCap) { + soft_cap_pre_tanh_scale = params.sm_scale * math::ptx_rcp(params.logits_soft_cap); + sm_scale_log2 = math::log2e * params.logits_soft_cap; + }else{ + sm_scale_log2 = params.sm_scale * math::log2e; + } + } + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, { + if constexpr (UseLogitsSoftCap) { + logits = float(math::tanh(logits * soft_cap_pre_tanh_scale)); + } + return logits; + }) +}; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + using AttentionVariant = {{ variant_name }}; \ + __VA_ARGS__(); \ + }) + +using DTypeQ = {{ dtype_q }}; +using DTypeKV = {{ dtype_kv }}; +using DTypeO = {{ dtype_o }}; +using IdType = {{ idtype }}; + +constexpr int HEAD_DIM_QK = {{ head_dim_qk }}; +constexpr int HEAD_DIM_VO = {{ head_dim_vo }}; +constexpr auto POS_ENCODING_MODE = {{ pos_encoding_mode }}; + +struct PersistentParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + DTypeO* partial_o; + float* partial_lse; + DTypeO* final_o; + float* final_lse; + + IdType* q_indptr; + IdType* kv_indptr; + IdType* partial_indptr; + IdType* kv_indices; + IdType* q_len; + IdType* kv_len; + IdType* q_start; + IdType* kv_start; + IdType* kv_end; + IdType* kv_head_idx_arr; + IdType* work_indptr; + IdType* len_kv_chunk; + + // for state reduction + IdType* merge_indptr; + IdType* merge_o_indices; + IdType* num_packed_qo_len; + + uint32_t num_kv_heads; + uint_fastdiv gqa_group_size; + uint_fastdiv page_size; + + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_page; + uint32_t k_stride_h; + uint32_t k_stride_n; + uint32_t v_stride_page; + uint32_t v_stride_h; + uint32_t v_stride_n; + + float sm_scale; + float logits_soft_cap; + float v_scale; + {{ additional_params_decl }} + + PROFILER_PARAMS_DECL +}; diff --git a/csrc/batch_attention_jit_binding.cu b/csrc/batch_attention_jit_binding.cu new file mode 100644 index 0000000000..1b25eb0a48 --- /dev/null +++ b/csrc/batch_attention_jit_binding.cu @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_attention_config.inc" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchPagedAttentionPlan(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t head_dim_o, bool causal); + +void BatchPagedAttentionRun(TensorView float_workspace_buffer, TensorView int_workspace_buffer, + Array plan_info_vec, TensorView q, TensorView k_cache, + TensorView v_cache, TensorView kv_indices, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t layout_code, int64_t num_qo_heads, int64_t num_kv_heads, + int64_t page_size, double v_scale, double sm_scale, + double logits_soft_cap ADDITIONAL_FUNC_PARAMS PROFILER_FUNC_PARAMS); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, &BatchPagedAttentionPlan); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, &BatchPagedAttentionRun); diff --git a/csrc/batch_attention_paged_kernel_inst.jinja b/csrc/batch_attention_paged_kernel_inst.jinja new file mode 100644 index 0000000000..b6a915feca --- /dev/null +++ b/csrc/batch_attention_paged_kernel_inst.jinja @@ -0,0 +1,9 @@ +#include +#include "batch_attention_config.inc" + +namespace flashinfer { +template cudaError_t BatchPagedAttentionPersistent< + /*CTA_TILE_Q_1=*/128, /*CTA_TILE_Q_2=*/16, {{head_dim_qk}}, {{head_dim_vo}}, {{mask_mode}}, + {{ variant_name }}, PersistentParams>(const PersistentParams params_1, const PersistentParams params_2, + const uint32_t num_blks_x, const uint32_t num_blks_y, const cudaStream_t stream); +}; // namespace flashinfer diff --git a/csrc/batch_decode.cu b/csrc/batch_decode.cu index bf258b2739..c3ce1e2ecf 100644 --- a/csrc/batch_decode.cu +++ b/csrc/batch_decode.cu @@ -16,45 +16,45 @@ #include #include #include -#include #include "batch_decode_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" namespace flashinfer { template cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, - float* tmp_s, cudaStream_t stream); + float* tmp_s, bool enable_pdl, + cudaStream_t stream); } // namespace flashinfer using namespace flashinfer; -at::Tensor BatchDecodeWithPagedKVCachePlan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchDecodeWithPagedKVCachePlan( + TensorView float_workspace_buffer, TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, TensorView indptr, int64_t batch_size, int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, - at::Tensor empty_q_data, at::Tensor empty_kv_data) { + TensorView empty_q_data, TensorView empty_kv_data) { size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); DecodePlanInfo plan_info; - auto q_scalar_type = empty_q_data.scalar_type(); - auto kv_scalar_type = empty_kv_data.scalar_type(); - - TORCH_CHECK(head_dim_qk == head_dim_vo, - "CUDA cores template only supports equal head dim for QK and VO, please use tensor " - "cores template for different head dim"); + TVM_FFI_ICHECK_EQ(head_dim_qk, head_dim_vo) + << "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"; - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); DISPATCH_context( DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { @@ -69,27 +69,26 @@ at::Tensor BatchDecodeWithPagedKVCachePlan( batch_size, num_qo_heads, page_size, enable_cuda_graph, /*stream=*/stream, work_estimation_func); - TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", - cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchDecodeWithPagedKVCache failed with error " << cudaGetErrorString(status); return true; }); }); - return vec_to_tensor(plan_info.ToVector()); + return Array(plan_info.ToVector()); } -void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, - at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, - at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, - int64_t kv_layout_code, - int64_t window_left ADDITIONAL_FUNC_PARAMS) { +void BatchDecodeWithPagedKVCacheRun(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q, TensorView paged_k_cache, + TensorView paged_v_cache, TensorView paged_kv_indptr, + TensorView paged_kv_indices, TensorView paged_kv_last_page_len, + TensorView o, Optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { DecodePlanInfo plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); QKVLayout kv_layout = static_cast(kv_layout_code); - auto device = q.device(); int64_t batch_size = q.size(0); int64_t num_qo_heads = q.size(1); int64_t num_kv_heads, page_size; @@ -104,23 +103,19 @@ void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, uint32_t head_dim_qk = q.size(2); uint32_t head_dim_vo = paged_v_cache.size(3); - TORCH_CHECK(head_dim_qk == head_dim_vo, - "CUDA cores template only supports equal head dim for QK and VO, please use tensor " - "cores template for different head dim"); + TVM_FFI_ICHECK_EQ(head_dim_qk, head_dim_vo) + << "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"; - if (maybe_lse) { - const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); - TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + if (maybe_lse.has_value()) { + const auto& lse = maybe_lse.value(); + TVM_FFI_ICHECK_EQ(lse.size(0), batch_size); + TVM_FFI_ICHECK_EQ(lse.size(1), num_qo_heads); } void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); - // get q_scalar_type and kv_scalar_type - auto q_scalar_type = q.scalar_type(); - auto kv_scalar_type = paged_k_cache.scalar_type(); - // get q_stride_n and q_stride_h const auto q_stride_n = q.stride(0); const auto q_stride_h = q.stride(1); @@ -129,11 +124,14 @@ void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, const int64_t* kv_cache_strides = nullptr; auto k_strides = paged_k_cache.strides(); auto v_strides = paged_v_cache.strides(); - TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + TVM_FFI_ICHECK_EQ(k_strides.size(), v_strides.size()); + for (int i = 0; i < k_strides.size(); ++i) { + TVM_FFI_ICHECK_EQ(k_strides[i], v_strides[i]); + } kv_cache_strides = k_strides.data(); - const c10::cuda::OptionalCUDAGuard device_guard(device); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(q.device().device_id); + const cudaStream_t stream = get_stream(q.device()); DISPATCH_context( DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, @@ -150,7 +148,8 @@ void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, params.q = static_cast(q.data_ptr()); params.paged_kv = paged_kv; params.o = static_cast(o.data_ptr()); - params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.lse = + maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.padded_batch_size = 0; params.num_qo_heads = num_qo_heads; params.q_stride_n = q_stride_n; @@ -187,10 +186,10 @@ void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, cudaError_t status = flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, - tmp_s, + tmp_s, enable_pdl, /*stream=*/stream); - TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", - cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchDecodeWithPagedKVCache failed with error " << cudaGetErrorString(status); return true; }); } diff --git a/csrc/batch_decode_config.inc b/csrc/batch_decode_config.inc deleted file mode 100644 index f0e89fd560..0000000000 --- a/csrc/batch_decode_config.inc +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -// NOTE(Zihao): this is the include file for AOT mode -#pragma once -#include -#include -#include -#include - -#include "aot_default_additional_params.h" -#include "aot_extension_utils.h" - -using IdType = int32_t; - -#define ADDITIONAL_FUNC_PARAMS BATCH_DECODE_ADDITIONAL_FUNC_PARAMS -#define ADDITIONAL_PARAMS_SETTER BATCH_DECODE_ADDITIONAL_PARAMS_SETTER - -#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, \ - POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, \ - AttentionVariant, Params, ...) \ - { \ - DISPATCH_PYTORCH_QKV_DTYPE_TO_CTYPE(q_scalar_type, kv_scalar_type, DTypeQ, DTypeKV, [&] { \ - using DTypeO = DTypeQ; \ - using Params = BatchDecodeParams; \ - constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; \ - return DISPATCH_head_dim(head_dim_qk, HEAD_DIM_QK, [&] { \ - [[maybe_unused]] constexpr int HEAD_DIM_VO = HEAD_DIM_QK; \ - return DISPATCH_BOOL(window_left > -1, USE_SLIDING_WINDOW, [&] { \ - return DISPATCH_BOOL(logits_soft_cap > 0.f, USE_LOGITS_SOFT_CAP, [&] { \ - using AttentionVariant = \ - DefaultAttention; \ - __VA_ARGS__(); \ - return true; \ - }); \ - }); \ - }); \ - }); \ - } diff --git a/csrc/batch_decode_customize_config.jinja b/csrc/batch_decode_customize_config.jinja index da80370220..68006492ac 100644 --- a/csrc/batch_decode_customize_config.jinja +++ b/csrc/batch_decode_customize_config.jinja @@ -43,6 +43,7 @@ struct Params { IdType q_stride_n; IdType q_stride_h; int32_t window_left; + bool enable_pdl; IdType* request_indices; IdType* kv_tile_indices; diff --git a/csrc/batch_decode_jit_binding.cu b/csrc/batch_decode_jit_binding.cu new file mode 100644 index 0000000000..0ce644fbc2 --- /dev/null +++ b/csrc/batch_decode_jit_binding.cu @@ -0,0 +1,42 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchDecodeWithPagedKVCachePlan( + TensorView float_workspace_buffer, TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, TensorView indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + TensorView empty_q_data, TensorView empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q, TensorView paged_k_cache, + TensorView paged_v_cache, TensorView paged_kv_indptr, + TensorView paged_kv_indices, TensorView paged_kv_last_page_len, + TensorView o, Optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +// Batched decode with paged KV-Cache plan +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchDecodeWithPagedKVCachePlan); +// Batched decode with paged KV-Cache run +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, BatchDecodeWithPagedKVCacheRun); diff --git a/csrc/batch_decode_jit_pybind.cu b/csrc/batch_decode_jit_pybind.cu deleted file mode 100644 index 0816ba13ab..0000000000 --- a/csrc/batch_decode_jit_pybind.cu +++ /dev/null @@ -1,40 +0,0 @@ -/* - * Copyright (c) 2023-2025 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "batch_decode_config.inc" -#include "pytorch_extension_utils.h" - -at::Tensor BatchDecodeWithPagedKVCachePlan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, - int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, - int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, - at::Tensor empty_q_data, at::Tensor empty_kv_data); - -void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, - at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, - at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, - int64_t kv_layout_code, - int64_t window_left ADDITIONAL_FUNC_PARAMS); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // Batched decode with paged KV-Cache plan - m.def("plan", BatchDecodeWithPagedKVCachePlan); - // Batched decode with paged KV-Cache run - m.def("run", BatchDecodeWithPagedKVCacheRun); -} diff --git a/csrc/batch_decode_kernel_inst.jinja b/csrc/batch_decode_kernel_inst.jinja index df0c57bdb2..c35362d834 100644 --- a/csrc/batch_decode_kernel_inst.jinja +++ b/csrc/batch_decode_kernel_inst.jinja @@ -8,6 +8,6 @@ namespace flashinfer { template cudaError_t BatchDecodeWithPagedKVCacheDispatched<{{ head_dim_qk }}, {{ pos_encoding_mode }}, {{ variant_name }}, Params>( Params params, {{ dtype_o }}* tmp_v, - float* tmp_s, cudaStream_t stream); + float* tmp_s, bool enable_pdl, cudaStream_t stream); }; diff --git a/csrc/batch_decode_mla_binding.cu b/csrc/batch_decode_mla_binding.cu new file mode 100644 index 0000000000..3bb416d971 --- /dev/null +++ b/csrc/batch_decode_mla_binding.cu @@ -0,0 +1,24 @@ +#include "mla_config.inc" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchDecodeWithPagedKVCachePlanMLA(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t page_size, + bool enable_cuda_graph); + +void BatchDecodeWithPagedKVCacheRunMLA( + TensorView float_workspace_buffer, TensorView int_workspace_buffer, + Array plan_info_vec, TensorView q_nope, TensorView q_pe, TensorView paged_ckv_cache, + TensorView paged_kpe_cache, TensorView paged_kv_indptr, TensorView paged_kv_indices, + TensorView paged_kv_last_page_len, TensorView o, double sm_scale, int64_t window_left, + double logits_soft_cap, double rope_scale, double rope_theta, Optional maybe_lse, + bool enable_pdl); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchDecodeWithPagedKVCachePlanMLA); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, BatchDecodeWithPagedKVCacheRunMLA); diff --git a/csrc/batch_decode_mla_config.jinja b/csrc/batch_decode_mla_config.jinja index dba7d5b322..e6f126254d 100644 --- a/csrc/batch_decode_mla_config.jinja +++ b/csrc/batch_decode_mla_config.jinja @@ -2,8 +2,8 @@ #include #include -using namespace flashinfer; - +// using namespace flashinfer; +// avoid "at::Layout" is ambiguous error using DTypeQ = {{ dtype_q }}; using DTypeKV = {{ dtype_kv }}; using DTypeO = {{ dtype_o }}; @@ -16,6 +16,6 @@ constexpr int HEAD_DIM_KPE = {{ head_dim_kpe }}; constexpr int QO_TILE_LEN = {{ qo_tile_len }}; -using Params = BatchDecodeParamsMLA; +using Params = flashinfer::BatchDecodeParamsMLA; using AttentionVariant = - DefaultAttention; + flashinfer::DefaultAttention; diff --git a/csrc/batch_decode_mla_cute_sm80.cu b/csrc/batch_decode_mla_cute_sm80.cu index 2d442b6faf..5679076438 100644 --- a/csrc/batch_decode_mla_cute_sm80.cu +++ b/csrc/batch_decode_mla_cute_sm80.cu @@ -1,24 +1,30 @@ + #include #include -#include #include "mla_config.inc" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; -std::vector BatchDecodeWithPagedKVCachePlanMLA( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, unsigned int batch_size, - unsigned int num_qo_heads, unsigned int page_size, bool enable_cuda_graph, - int64_t cuda_stream) { +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchDecodeWithPagedKVCachePlanMLA(ffi::TensorView float_workspace_buffer, + ffi::TensorView int_workspace_buffer, + ffi::TensorView page_locked_int_workspace_buffer, + ffi::TensorView indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t page_size, + bool enable_cuda_graph) { size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); DecodePlanInfo plan_info; - cudaStream_t stream = reinterpret_cast(cuda_stream); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatchedMlaCuteSM80< HEAD_DIM_CKV, HEAD_DIM_KPE, QO_TILE_LEN, AttentionVariant, Params>; @@ -31,34 +37,36 @@ std::vector BatchDecodeWithPagedKVCachePlanMLA( batch_size, num_qo_heads, page_size, enable_cuda_graph, /*stream=*/stream, work_estimation_func); - TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCachePlanMLA failed with error ", - cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchDecodeWithPagedKVCachePlanMLA failed with error " << cudaGetErrorString(status); - return plan_info.ToVector(); + return Array(plan_info.ToVector()); } void BatchDecodeWithPagedKVCacheRunMLA( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - std::vector plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, - at::Tensor paged_ckv_cache, at::Tensor paged_kpe_cache, at::Tensor paged_kv_indptr, - at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, at::Tensor o, float sm_scale, - int window_left, float logits_soft_cap, float rope_scale, float rope_theta, - std::optional maybe_lse, int64_t cuda_stream) { + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q_nope, ffi::TensorView q_pe, + ffi::TensorView paged_ckv_cache, ffi::TensorView paged_kpe_cache, + ffi::TensorView paged_kv_indptr, ffi::TensorView paged_kv_indices, + ffi::TensorView paged_kv_last_page_len, ffi::TensorView o, double sm_scale, int64_t window_left, + double logits_soft_cap, double rope_scale, double rope_theta, + Optional maybe_lse, + bool enable_pdl // fake placeholder, sm80 does not support pdl +) { DecodePlanInfo plan_info; - plan_info.FromVector(plan_info_vec); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); - auto device = q_nope.device(); int64_t batch_size = q_nope.size(0); int64_t num_qo_heads = q_nope.size(1); int64_t page_size = paged_ckv_cache.size(1); - if (maybe_lse) { - const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q_nope.size(0)); - TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q_nope.size(1)); + if (maybe_lse.has_value()) { + const auto& lse = maybe_lse.value(); + TVM_FFI_ICHECK_EQ(lse.size(0), batch_size); + TVM_FFI_ICHECK_EQ(lse.size(1), num_qo_heads); } - TORCH_CHECK(logits_soft_cap >= 0.f, "logits_soft_cap must be non-negative"); + TVM_FFI_ICHECK_GE(logits_soft_cap, 0.f) << "logits_soft_cap must be non-negative"; void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); @@ -70,10 +78,11 @@ void BatchDecodeWithPagedKVCacheRunMLA( static_cast(paged_kv_indices.data_ptr()), static_cast(paged_kv_indptr.data_ptr()), static_cast(paged_kv_last_page_len.data_ptr())); - Params params(static_cast(q_nope.data_ptr()), static_cast(q_pe.data_ptr()), - /*q_offset=*/nullptr, paged_kv, static_cast(o.data_ptr()), - /*lse=*/(maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr), - num_qo_heads, window_left, logits_soft_cap, sm_scale, rope_scale, rope_theta); + Params params( + static_cast(q_nope.data_ptr()), static_cast(q_pe.data_ptr()), + /*q_offset=*/nullptr, paged_kv, static_cast(o.data_ptr()), + /*lse=*/(maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr), + num_qo_heads, window_left, logits_soft_cap, sm_scale, rope_scale, rope_theta); DTypeO* tmp_v = nullptr; float* tmp_s = nullptr; @@ -94,10 +103,11 @@ void BatchDecodeWithPagedKVCacheRunMLA( } params.padded_batch_size = plan_info.padded_batch_size; - cudaStream_t stream = reinterpret_cast(cuda_stream); + cudaSetDevice(paged_ckv_cache.device().device_id); + const cudaStream_t stream = get_stream(paged_ckv_cache.device()); cudaError_t status = BatchDecodeWithPagedKVCacheDispatchedMlaCuteSM80( params, tmp_v, tmp_s, /*stream=*/stream); - TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", - cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchDecodeWithPagedKVCache failed with error " << cudaGetErrorString(status); } diff --git a/csrc/batch_decode_mla_plan.cu b/csrc/batch_decode_mla_plan.cu index 25b3313bc5..7925a14f27 100644 --- a/csrc/batch_decode_mla_plan.cu +++ b/csrc/batch_decode_mla_plan.cu @@ -1,26 +1,29 @@ #include #include -#include #include "mla_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; -at::Tensor BatchDecodeWithPagedKVCachePlanMLA(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, - at::Tensor indptr, int64_t batch_size, - int64_t num_qo_heads, int64_t page_size, - bool enable_cuda_graph, int64_t cuda_stream) { +using tvm::ffi::Array; + +Array BatchDecodeWithPagedKVCachePlanMLA(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t page_size, + bool enable_cuda_graph) { + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); + size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); DecodePlanInfo plan_info; - cudaStream_t stream = reinterpret_cast(cuda_stream); auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatchedMLA maybe_lse, int64_t cuda_stream); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - m.def("plan", BatchDecodeWithPagedKVCachePlanMLA); - m.def("run", BatchDecodeWithPagedKVCacheRunMLA); -} diff --git a/csrc/batch_decode_mla_run.cu b/csrc/batch_decode_mla_run.cu index e39fb0de8b..35d533b536 100644 --- a/csrc/batch_decode_mla_run.cu +++ b/csrc/batch_decode_mla_run.cu @@ -1,38 +1,43 @@ #include #include -#include #include "mla_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; +using tvm::ffi::Array; +using tvm::ffi::Optional; + void BatchDecodeWithPagedKVCacheRunMLA( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q_nope, at::Tensor q_pe, at::Tensor paged_ckv_cache, at::Tensor paged_kpe_cache, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, double sm_scale, int64_t window_left, double logits_soft_cap, double rope_scale, - double rope_theta, std::optional maybe_lse, int64_t cuda_stream) { + TensorView float_workspace_buffer, TensorView int_workspace_buffer, + Array plan_info_vec, TensorView q_nope, TensorView q_pe, TensorView paged_ckv_cache, + TensorView paged_kpe_cache, TensorView paged_kv_indptr, TensorView paged_kv_indices, + TensorView paged_kv_last_page_len, TensorView o, double sm_scale, int64_t window_left, + double logits_soft_cap, double rope_scale, double rope_theta, Optional maybe_lse, + bool enable_pdl) { DecodePlanInfo plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); - auto device = q_nope.device(); int64_t batch_size = q_nope.size(0); int64_t num_qo_heads = q_nope.size(1); int64_t page_size = paged_ckv_cache.size(1); - if (maybe_lse) { - const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q_nope.size(0)); - TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q_nope.size(1)); + if (maybe_lse.has_value()) { + const auto& lse = maybe_lse.value(); + TVM_FFI_ICHECK_EQ(lse.size(0), batch_size); + TVM_FFI_ICHECK_EQ(lse.size(1), num_qo_heads); } - TORCH_CHECK(logits_soft_cap >= 0.f, "logits_soft_cap must be non-negative"); + TVM_FFI_ICHECK_GE(logits_soft_cap, 0.f) << "logits_soft_cap must be non-negative"; void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + cudaSetDevice(q_nope.device().device_id); + const cudaStream_t stream = get_stream(q_nope.device()); + paged_kv_mla_t paged_kv( page_size, HEAD_DIM_CKV, HEAD_DIM_KPE, batch_size, static_cast(paged_ckv_cache.data_ptr()), paged_ckv_cache.strides().data(), @@ -42,7 +47,7 @@ void BatchDecodeWithPagedKVCacheRunMLA( static_cast(paged_kv_last_page_len.data_ptr())); Params params(static_cast(q_nope.data_ptr()), static_cast(q_pe.data_ptr()), /*q_offset=*/nullptr, paged_kv, static_cast(o.data_ptr()), - /*lse=*/(maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr), + /*lse=*/(maybe_lse ? static_cast(maybe_lse.value().data_ptr()) : nullptr), num_qo_heads, window_left, logits_soft_cap, sm_scale, rope_scale, rope_theta); DTypeO* tmp_v = nullptr; @@ -64,10 +69,11 @@ void BatchDecodeWithPagedKVCacheRunMLA( } params.padded_batch_size = plan_info.padded_batch_size; - cudaStream_t stream = reinterpret_cast(cuda_stream); cudaError_t status = BatchDecodeWithPagedKVCacheDispatchedMLA(params, tmp_v, tmp_s, /*stream=*/stream); - TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", - cudaGetErrorString(status)); + Params>(params, tmp_v, tmp_s, enable_pdl, + /*stream=*/stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchDecodeWithPagedKVCache failed with error: " << cudaGetErrorString(status); } diff --git a/csrc/batch_mla_binding.cu b/csrc/batch_mla_binding.cu new file mode 100644 index 0000000000..b39192de6a --- /dev/null +++ b/csrc/batch_mla_binding.cu @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_mla_config.inc" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchMLAPagedAttentionPlan(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len, int64_t num_heads, int64_t head_dim_o, + bool causal); + +void BatchMLAPagedAttentionRun(TensorView float_workspace_buffer, TensorView int_workspace_buffer, + Array plan_info_vec, TensorView q_nope, TensorView q_pe, + TensorView ckv_cache, TensorView kpe_cache, TensorView kv_indices, + TensorView o, Optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, double sm_scale, + bool return_lse_base_on_e); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchMLAPagedAttentionPlan); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, BatchMLAPagedAttentionRun); diff --git a/csrc/batch_mla_config.jinja b/csrc/batch_mla_config.jinja index e29a2b50f6..2525d34273 100644 --- a/csrc/batch_mla_config.jinja +++ b/csrc/batch_mla_config.jinja @@ -11,7 +11,7 @@ using namespace flashinfer; #ifdef FLASHINFER_ENABLE_PROFILER -#define ADDITIONAL_FUNC_PARAMS , at::Tensor profiler_buffer +#define ADDITIONAL_FUNC_PARAMS , Tensor profiler_buffer #define ADDITIONAL_PARAMS_SETTER \ params.profiler_buffer = static_cast(profiler_buffer.data_ptr()); #else diff --git a/csrc/batch_mla_plan.cu b/csrc/batch_mla_plan.cu index 0276102cab..1f7176e452 100644 --- a/csrc/batch_mla_plan.cu +++ b/csrc/batch_mla_plan.cu @@ -14,30 +14,32 @@ * limitations under the License. */ #include -#include #include "batch_mla_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; -at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len, - int64_t num_heads, int64_t head_dim_o, bool causal) { +using tvm::ffi::Array; + +Array BatchMLAPagedAttentionPlan(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len, int64_t num_heads, int64_t head_dim_o, + bool causal) { size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); MLAPlanInfo plan_info; int batch_size = kv_len.size(0); - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); cudaError_t status = MLAPlan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, @@ -45,8 +47,8 @@ at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), batch_size, num_heads, head_dim_o, causal, stream); + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to plan MLA, error: " << cudaGetErrorString(status); - TORCH_CHECK(status == cudaSuccess, "Failed to plan MLA, error: ", cudaGetErrorString(status)); - - return vec_to_tensor(plan_info.ToVector()); + return Array(plan_info.ToVector()); } diff --git a/csrc/batch_mla_pybind.cu b/csrc/batch_mla_pybind.cu deleted file mode 100644 index 438fc39519..0000000000 --- a/csrc/batch_mla_pybind.cu +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "batch_mla_config.inc" -#include "pytorch_extension_utils.h" - -at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len, - int64_t num_heads, int64_t head_dim_o, bool causal); - -void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, - at::Tensor ckv_cache, at::Tensor kpe_cache, at::Tensor kv_indices, - at::Tensor o, std::optional maybe_lse, - int64_t mask_mode_code, int64_t num_heads, int64_t page_size, - double sm_scale); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - m.def("plan", &BatchMLAPagedAttentionPlan); - m.def("run", &BatchMLAPagedAttentionRun); -} diff --git a/csrc/batch_mla_run.cu b/csrc/batch_mla_run.cu index ef4c4848ea..9d950787ad 100644 --- a/csrc/batch_mla_run.cu +++ b/csrc/batch_mla_run.cu @@ -13,40 +13,38 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include #include #include #include -#include #include "batch_mla_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; -void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, - at::Tensor ckv_cache, at::Tensor kpe_cache, at::Tensor kv_indices, - at::Tensor o, std::optional maybe_lse, - int64_t mask_mode_code, int64_t num_heads, int64_t page_size, - double sm_scale) { +using tvm::ffi::Array; +using tvm::ffi::Optional; + +void BatchMLAPagedAttentionRun(TensorView float_workspace_buffer, TensorView int_workspace_buffer, + Array plan_info_vec, TensorView q_nope, TensorView q_pe, + TensorView ckv_cache, TensorView kpe_cache, TensorView kv_indices, + TensorView o, Optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, double sm_scale, + bool return_lse_base_on_e) { // q_nope: [n, num_heads, head_dim_ckv] // q_pe: [n, num_heads, head_dim_kpe] // ckv_cache: [num_pages, page_size, head_dim_ckv] // kpe_cache: [num_pages, page_size, head_dim_kpe] MLAPlanInfo plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); - - auto device = q_nope.device(); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); void* float_buffer_ptr = float_workspace_buffer.data_ptr(); void* int_buffer_ptr = int_workspace_buffer.data_ptr(); const MaskMode mask_mode = static_cast(mask_mode_code); - auto q_scalar_type = q_nope.scalar_type(); - auto kv_scalar_type = ckv_cache.scalar_type(); - unsigned int q_nope_stride_n = q_nope.stride(0); unsigned int q_nope_stride_h = q_nope.stride(1); unsigned int q_pe_stride_n = q_pe.stride(0); @@ -58,8 +56,8 @@ void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int unsigned int o_stride_n = o.stride(0); unsigned int o_stride_h = o.stride(1); - const c10::cuda::OptionalCUDAGuard device_guard(device); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(q_nope.device().device_id); + const cudaStream_t stream = get_stream(q_nope.device()); DISPATCH_context( DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, [&] { @@ -94,7 +92,7 @@ void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_partial_stride_offset); params.final_o = static_cast(o.data_ptr()); params.final_lse = - maybe_lse.has_value() ? static_cast(maybe_lse->data_ptr()) : nullptr; + maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.partial_o = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); params.partial_lse = @@ -115,11 +113,12 @@ void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int params.o_stride_h = o_stride_h; params.sm_scale = sm_scale; + params.return_lse_base_on_e = return_lse_base_on_e; cudaError_t status = mla::BatchMLAPagedAttention( params, plan_info.num_blks_x, plan_info.num_blks_y, stream); - TORCH_CHECK(status == cudaSuccess, - "Failed to run MLA, error: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to run MLA, error: " << cudaGetErrorString(status); }); } diff --git a/csrc/batch_mla_sm90_binding.cu b/csrc/batch_mla_sm90_binding.cu new file mode 100644 index 0000000000..f2af49766a --- /dev/null +++ b/csrc/batch_mla_sm90_binding.cu @@ -0,0 +1,39 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_mla_sm90_config.inc" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchMLAPagedAttentionSM90Plan(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len, int64_t num_heads, + int64_t head_dim_o, bool causal); + +void BatchMLAPagedAttentionSM90Run(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q_nope, TensorView q_pe, TensorView ckv_cache, + TensorView kpe_cache, TensorView kv_indices, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, double sm_scale, + bool return_lse_base_on_e ADDITIONAL_FUNC_PARAMS); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchMLAPagedAttentionSM90Plan); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, BatchMLAPagedAttentionSM90Run); diff --git a/csrc/batch_mla_sm90_plan.cu b/csrc/batch_mla_sm90_plan.cu index 35f6d0fa82..d297ebab90 100644 --- a/csrc/batch_mla_sm90_plan.cu +++ b/csrc/batch_mla_sm90_plan.cu @@ -14,31 +14,32 @@ * limitations under the License. */ #include -#include #include "batch_mla_sm90_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; -at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, - at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len, int64_t num_heads, int64_t head_dim_o, - bool causal) { +using tvm::ffi::Array; + +Array BatchMLAPagedAttentionSM90Plan(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, + TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len, int64_t num_heads, + int64_t head_dim_o, bool causal) { size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); MLAPlanInfo plan_info; int batch_size = kv_len.size(0); - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); cudaError_t status = MLAPlan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, @@ -47,7 +48,8 @@ at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), batch_size, num_heads, head_dim_o, causal, stream); - TORCH_CHECK(status == cudaSuccess, "Failed to plan MLA, error: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to plan MLA, error: " << cudaGetErrorString(status); - return vec_to_tensor(plan_info.ToVector()); + return Array(plan_info.ToVector()); } diff --git a/csrc/batch_mla_sm90_pybind.cu b/csrc/batch_mla_sm90_pybind.cu deleted file mode 100644 index d089ad0c5c..0000000000 --- a/csrc/batch_mla_sm90_pybind.cu +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "batch_mla_sm90_config.inc" -#include "pytorch_extension_utils.h" - -at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, - at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len, int64_t num_heads, int64_t head_dim_o, - bool causal); - -void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q_nope, at::Tensor q_pe, at::Tensor ckv_cache, - at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t num_heads, int64_t page_size, - double sm_scale ADDITIONAL_FUNC_PARAMS); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - m.def("plan", &BatchMLAPagedAttentionSM90Plan); - m.def("run", &BatchMLAPagedAttentionSM90Run); -} diff --git a/csrc/batch_mla_sm90_run.cu b/csrc/batch_mla_sm90_run.cu index 3da072eb03..c6670ca342 100644 --- a/csrc/batch_mla_sm90_run.cu +++ b/csrc/batch_mla_sm90_run.cu @@ -16,38 +16,35 @@ #include #include #include -#include #include "batch_mla_sm90_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; -void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q_nope, at::Tensor q_pe, at::Tensor ckv_cache, - at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t num_heads, int64_t page_size, - double sm_scale ADDITIONAL_FUNC_PARAMS) { +using tvm::ffi::Array; +using tvm::ffi::Optional; + +void BatchMLAPagedAttentionSM90Run(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q_nope, TensorView q_pe, TensorView ckv_cache, + TensorView kpe_cache, TensorView kv_indices, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, double sm_scale, + bool return_lse_base_on_e ADDITIONAL_FUNC_PARAMS) { // q_nope: [n, num_heads, head_dim_ckv] // q_pe: [n, num_heads, head_dim_kpe] // ckv_cache: [num_pages, page_size, head_dim_ckv] // kpe_cache: [num_pages, page_size, head_dim_kpe] MLAPlanInfo plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); - - auto device = q_nope.device(); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); void* float_buffer_ptr = float_workspace_buffer.data_ptr(); void* int_buffer_ptr = int_workspace_buffer.data_ptr(); const MaskMode mask_mode = static_cast(mask_mode_code); - auto q_scalar_type = q_nope.scalar_type(); - auto kv_scalar_type = ckv_cache.scalar_type(); - unsigned int q_nope_stride_n = q_nope.stride(0); unsigned int q_nope_stride_h = q_nope.stride(1); unsigned int q_pe_stride_n = q_pe.stride(0); @@ -59,8 +56,8 @@ void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, unsigned int o_stride_n = o.stride(0); unsigned int o_stride_h = o.stride(1); - const c10::cuda::OptionalCUDAGuard device_guard(device); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(q_nope.device().device_id); + const cudaStream_t stream = get_stream(q_nope.device()); DISPATCH_context( DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, [&] { @@ -95,7 +92,7 @@ void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_partial_stride_offset); params.final_o = static_cast(o.data_ptr()); params.final_lse = - maybe_lse.has_value() ? static_cast(maybe_lse->data_ptr()) : nullptr; + maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.partial_o = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); params.partial_lse = @@ -114,6 +111,7 @@ void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, params.kpe_stride_n = kpe_stride_n; params.o_stride_n = o_stride_n; params.o_stride_h = o_stride_h; + params.return_lse_base_on_e = return_lse_base_on_e; ADDITIONAL_PARAMS_SETTER @@ -123,7 +121,7 @@ void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, mla::BatchMLAPageAttentionHopper( params, plan_info.num_blks_x, plan_info.num_blks_y, stream); - TORCH_CHECK(status == cudaSuccess, - "Failed to run MLA, error: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to run MLA, error: " << cudaGetErrorString(status); }); } diff --git a/csrc/batch_pod.cu b/csrc/batch_pod.cu new file mode 100644 index 0000000000..33aa5e753f --- /dev/null +++ b/csrc/batch_pod.cu @@ -0,0 +1,350 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "batch_pod_config.inc" +#include "tvm_ffi_utils.h" + +namespace flashinfer { +template +cudaError_t BatchPODWithKVCacheTensorDispatched(PrefillParams prefill_params, + typename PrefillParams::DTypeO* tmp_v_p, + float* tmp_s_p, DecodeParams decode_params, + typename DecodeParams::DTypeO* tmp_v_d, + float* tmp_s_d, bool enable_pdl, + cudaStream_t stream, int* sm_aware_sched); + +} // namespace flashinfer + +using namespace flashinfer; + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +void batch_pod_with_kv_cache_tensor( + // Prefill params + TensorView float_workspace_buffer_p, TensorView int_workspace_buffer_p, + Array plan_info_vec_p, TensorView q_p, TensorView paged_k_cache_p, + TensorView paged_v_cache_p, TensorView qo_indptr_p, TensorView paged_kv_indptr_p, + TensorView paged_kv_indices_p, TensorView paged_kv_last_page_len_p, TensorView o_p, + Optional maybe_lse_p, int64_t mask_mode_code_p, int64_t layout_p, + int64_t window_left_p, Optional maybe_custom_mask_p, + Optional maybe_mask_indptr_p, Optional maybe_alibi_slopes_p, + double logits_soft_cap_p, double sm_scale_p, double rope_rcp_scale_p, double rope_rcp_theta_p, + // Decode params + TensorView float_workspace_buffer_d, TensorView int_workspace_buffer_d, + Array plan_info_vec_d, TensorView q_d, TensorView paged_k_cache_d, + TensorView paged_v_cache_d, TensorView qo_indptr_d, TensorView paged_kv_indptr_d, + TensorView paged_kv_indices_d, TensorView paged_kv_last_page_len_d, TensorView o_d, + Optional maybe_lse_d, int64_t mask_mode_code_d, int64_t layout_d, + int64_t window_left_d, Optional maybe_custom_mask_d, + Optional maybe_mask_indptr_d, Optional maybe_alibi_slopes_d, + double logits_soft_cap_d, double sm_scale_d, double rope_rcp_scale_d, double rope_rcp_theta_d, + bool enable_pdl, TensorView sm_aware_sched) { + // Prefill setup + PrefillPlanInfo plan_info_p; + plan_info_p.FromVector(std::vector(plan_info_vec_p.begin(), plan_info_vec_p.end())); + QKVLayout kv_layout_p = static_cast(layout_p); + int64_t batch_size_p = paged_kv_indptr_p.size(0) - 1; + int64_t num_qo_heads = q_p.size(1); + + int64_t num_kv_heads_p, page_size_p; + uint32_t head_dim_qk_p = q_p.size(2); + if (kv_layout_p == QKVLayout::kHND) { + num_kv_heads_p = paged_k_cache_p.size(1); + page_size_p = paged_k_cache_p.size(2); + } else { + page_size_p = paged_k_cache_p.size(1); + num_kv_heads_p = paged_k_cache_p.size(2); + } + + if (maybe_lse_p.has_value()) { + const auto& lse = maybe_lse_p.value(); + TVM_FFI_ICHECK_EQ(lse.size(0), q_p.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q_p.size(1)); + } + + void* float_buffer_ptr_p = static_cast(float_workspace_buffer_p.data_ptr()); + void* int_buffer_ptr_p = static_cast(int_workspace_buffer_p.data_ptr()); + + const MaskMode mask_mode_p = static_cast(mask_mode_code_p); + + // get q_stride_n and q_stride_h + const auto q_stride_n_p = q_p.stride(0); + const auto q_stride_h_p = q_p.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides_p = nullptr; + auto k_strides_p = paged_k_cache_p.strides(); + auto v_strides_p = paged_v_cache_p.strides(); + TVM_FFI_ICHECK_EQ(k_strides_p.size(), v_strides_p.size()); + for (int i = 0; i < k_strides_p.size(); ++i) { + TVM_FFI_ICHECK_EQ(k_strides_p[i], v_strides_p[i]); + } + kv_cache_strides_p = k_strides_p.data(); + + cudaSetDevice(float_workspace_buffer_p.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer_p.device()); + + // Decode setup (TensorView decode = batched prefill) + PrefillPlanInfo plan_info_d; + plan_info_d.FromVector(std::vector(plan_info_vec_d.begin(), plan_info_vec_d.end())); + QKVLayout kv_layout_d = static_cast(layout_d); + int64_t batch_size_d = paged_kv_indptr_d.size(0) - 1; + int64_t num_qo_heads_d = q_d.size(1); + + TVM_FFI_ICHECK_EQ(num_qo_heads, num_qo_heads_d) + << "POD currently requires same # Query heads for prefill and decode"; + + int64_t num_kv_heads_d, page_size_d; + uint32_t head_dim_qk_d = q_d.size(2); + if (kv_layout_d == QKVLayout::kHND) { + num_kv_heads_d = paged_k_cache_d.size(1); + page_size_d = paged_k_cache_d.size(2); + } else { + page_size_d = paged_k_cache_d.size(1); + num_kv_heads_d = paged_k_cache_d.size(2); + } + TVM_FFI_ICHECK_EQ(num_kv_heads_p, num_kv_heads_d) + << "POD currently requires same # KV heads for prefill and decode; Prefill: " + << num_kv_heads_p << ", Decode: " << num_kv_heads_d; + + if (maybe_lse_d.has_value()) { + const auto& lse = maybe_lse_d.value(); + TVM_FFI_ICHECK_EQ(lse.size(0), q_d.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q_d.size(1)); + } + + void* float_buffer_ptr_d = static_cast(float_workspace_buffer_d.data_ptr()); + void* int_buffer_ptr_d = static_cast(int_workspace_buffer_d.data_ptr()); + + const MaskMode mask_mode_d = static_cast(mask_mode_code_d); + + // get q_stride_n and q_stride_h + const auto q_stride_n_d = q_d.stride(0); + const auto q_stride_h_d = q_d.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides_d = nullptr; + auto k_strides_d = paged_k_cache_d.strides(); + auto v_strides_d = paged_v_cache_d.strides(); + TVM_FFI_ICHECK_EQ(k_strides_d.size(), v_strides_d.size()); + for (int i = 0; i < k_strides_d.size(); ++i) { + TVM_FFI_ICHECK_EQ(k_strides_d[i], v_strides_d[i]); + } + kv_cache_strides_d = k_strides_d.data(); + + // Already handled by prefill + // cudaSetDevice(float_workspace_buffer_d.device().device_id); + // const cudaStream_t stream = get_stream(float_workspace_buffer_d.device()); + + DISPATCH_context( + MASK_MODE_P, MASK_MODE_D, DTypeQ, DTypeKV, HEAD_DIM_QK, USE_SLIDING_WINDOW_P, + USE_SLIDING_WINDOW_D, USE_LOGITS_SOFT_CAP, [&] { + PrefillParams prefill_params; + DTypeO* tmp_v_p = nullptr; + float* tmp_s_p = nullptr; + { + PrefillParams& params = prefill_params; + params.q = static_cast(q_p.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads_p, page_size_p, HEAD_DIM_VO, batch_size_p, kv_layout_p, + static_cast(paged_k_cache_p.data_ptr()), + static_cast(paged_v_cache_p.data_ptr()), kv_cache_strides_p, + static_cast(paged_kv_indices_p.data_ptr()), + static_cast(paged_kv_indptr_p.data_ptr()), + static_cast(paged_kv_last_page_len_p.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr_p.data_ptr()); + params.o = static_cast(o_p.data_ptr()); + + params.lse = maybe_lse_p.has_value() ? static_cast(maybe_lse_p.value().data_ptr()) + : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n_p; + params.q_stride_h = q_stride_h_p; + params.window_left = window_left_p; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + params.maybe_mask_indptr = + maybe_mask_indptr_p.has_value() + ? static_cast(maybe_mask_indptr_p.value().data_ptr()) + : nullptr; + params.maybe_alibi_slopes = + maybe_alibi_slopes_p.has_value() + ? static_cast(maybe_alibi_slopes_p.value().data_ptr()) + : nullptr; + params.logits_soft_cap = logits_soft_cap_p; + params.sm_scale = sm_scale_p; + params.rope_rcp_scale = rope_rcp_scale_p; + params.rope_rcp_theta = rope_rcp_theta_p; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.kv_tile_indices_offset); + params.o_indptr = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.kv_chunk_size_ptr_offset); + if (plan_info_p.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.merge_indptr_offset); + tmp_v_p = GetPtrFromBaseOffset(float_buffer_ptr_p, plan_info_p.v_offset); + tmp_s_p = GetPtrFromBaseOffset(float_buffer_ptr_p, plan_info_p.s_offset); + if (plan_info_p.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info_p.padded_batch_size; + params.max_total_num_rows = plan_info_p.total_num_rows; + if (plan_info_p.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr_p, plan_info_p.total_num_rows_offset); + } + } + + DecodeParams decode_params; + DTypeO* tmp_v_d = nullptr; + float* tmp_s_d = nullptr; + { + DecodeParams& params = decode_params; + params.q = static_cast(q_d.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads_d, page_size_d, HEAD_DIM_VO, batch_size_d, kv_layout_d, + static_cast(paged_k_cache_d.data_ptr()), + static_cast(paged_v_cache_d.data_ptr()), kv_cache_strides_d, + static_cast(paged_kv_indices_d.data_ptr()), + static_cast(paged_kv_indptr_d.data_ptr()), + static_cast(paged_kv_last_page_len_d.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr_d.data_ptr()); + params.o = static_cast(o_d.data_ptr()); + + params.lse = maybe_lse_d.has_value() ? static_cast(maybe_lse_d.value().data_ptr()) + : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n_d; + params.q_stride_h = q_stride_h_d; + params.window_left = window_left_d; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + params.maybe_mask_indptr = + maybe_mask_indptr_d.has_value() + ? static_cast(maybe_mask_indptr_d.value().data_ptr()) + : nullptr; + params.maybe_alibi_slopes = + maybe_alibi_slopes_d.has_value() + ? static_cast(maybe_alibi_slopes_d.value().data_ptr()) + : nullptr; + params.logits_soft_cap = logits_soft_cap_d; + params.sm_scale = sm_scale_d; + params.rope_rcp_scale = rope_rcp_scale_d; + params.rope_rcp_theta = rope_rcp_theta_d; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.kv_tile_indices_offset); + params.o_indptr = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.kv_chunk_size_ptr_offset); + if (plan_info_d.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.merge_indptr_offset); + tmp_v_d = GetPtrFromBaseOffset(float_buffer_ptr_d, plan_info_d.v_offset); + tmp_s_d = GetPtrFromBaseOffset(float_buffer_ptr_d, plan_info_d.s_offset); + if (plan_info_d.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info_d.padded_batch_size; + params.max_total_num_rows = plan_info_d.total_num_rows; + if (plan_info_d.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr_d, plan_info_d.total_num_rows_offset); + } + } + + constexpr bool use_custom_mask_p = MASK_MODE_P == MaskMode::kCustom; + using PrefillAttentionVariant = + DefaultAttention; + constexpr bool use_custom_mask_d = MASK_MODE_D == MaskMode::kCustom; + using DecodeAttentionVariant = + DefaultAttention; + + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + int num_sm = 0; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + // SM-aware scheduling buffer uses num_sm + 2 entries + // num_sm entries for counters for each SM, and + // 2 entries for keeping track of blockIds for prefill and decode + assert( + sm_aware_sched.ndim() == 1 && sm_aware_sched.size(0) == num_sm + 2 && + "sm_aware_sched tensor has incorrect shape or type, should be (num_sm + 2,) of int32"); + DISPATCH_CTA_TILE_Q(plan_info_p.cta_tile_q, CTA_TILE_Q_P, { + constexpr size_t CTA_TILE_Q_D = 16; + cudaError_t status = flashinfer::BatchPODWithKVCacheTensorDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_FP16_QK_REDUCTION, CTA_TILE_Q_P, + MASK_MODE_P, CTA_TILE_Q_D, MASK_MODE_D, PrefillAttentionVariant, + DecodeAttentionVariant>(prefill_params, tmp_v_p, tmp_s_p, decode_params, tmp_v_d, + tmp_s_d, enable_pdl, stream, + static_cast(sm_aware_sched.data_ptr())); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchPODWithKVCache kernel launch failed, error: " << cudaGetErrorString(status); + return status; + }); + }); +} diff --git a/csrc/batch_pod_customize_config.jinja b/csrc/batch_pod_customize_config.jinja new file mode 100644 index 0000000000..9f27b42953 --- /dev/null +++ b/csrc/batch_pod_customize_config.jinja @@ -0,0 +1,43 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace flashinfer; + +using DTypeQ = {{ dtype_q }}; +using DTypeKV = {{ dtype_kv }}; +using DTypeO = {{ dtype_o }}; +using IdType = {{ idtype }}; +constexpr int HEAD_DIM_QK = {{ head_dim_qk }}; +constexpr int HEAD_DIM_VO = {{ head_dim_vo }}; +constexpr bool USE_FP16_QK_REDUCTION = {{ use_fp16_qk_reduction }}; + +constexpr auto USE_LOGITS_SOFT_CAP_P = {{ use_logits_soft_cap_p }}; +constexpr auto POS_ENCODING_MODE_P = {{ pos_encoding_mode_p }}; +constexpr auto USE_SLIDING_WINDOW_P = {{ use_sliding_window_p }}; + +constexpr auto USE_LOGITS_SOFT_CAP_D = {{ use_logits_soft_cap_d }}; +constexpr auto POS_ENCODING_MODE_D = {{ pos_encoding_mode_d }}; +constexpr auto USE_SLIDING_WINDOW_D = {{ use_sliding_window_d }}; + +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr bool USE_LOGITS_SOFT_CAP = false; + +using PrefillParams = BatchPrefillPagedParams; +using DecodeParams = BatchPrefillPagedParams; + +#define DISPATCH_context(MASK_MODE_P, MASK_MODE_D, DTypeQ, DTypeKV, HEAD_DIM_QK, \ + USE_SLIDING_WINDOW_P, USE_SLIDING_WINDOW_D, USE_LOGITS_SOFT_CAP, ...) \ + DISPATCH_MASK_MODE(mask_mode_p, MASK_MODE_P, { \ + DISPATCH_MASK_MODE(mask_mode_d, MASK_MODE_D, { \ + __VA_ARGS__(); \ + }); \ +}); diff --git a/csrc/batch_pod_jit_binding.cu b/csrc/batch_pod_jit_binding.cu new file mode 100644 index 0000000000..c7a8a5ea6b --- /dev/null +++ b/csrc/batch_pod_jit_binding.cu @@ -0,0 +1,44 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_pod_config.inc" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +void batch_pod_with_kv_cache_tensor( + // Prefill params + TensorView float_workspace_buffer_p, TensorView int_workspace_buffer_p, + Array plan_info_vec_p, TensorView q_p, TensorView paged_k_cache_p, + TensorView paged_v_cache_p, TensorView qo_indptr_p, TensorView paged_kv_indptr_p, + TensorView paged_kv_indices_p, TensorView paged_kv_last_page_len_p, TensorView o_p, + Optional maybe_lse_p, int64_t mask_mode_code_p, int64_t layout_p, + int64_t window_left_p, Optional maybe_custom_mask_p, + Optional maybe_mask_indptr_p, Optional maybe_alibi_slopes_p, + double logits_soft_cap_p, double sm_scale_p, double rope_rcp_scale_p, double rope_rcp_theta_p, + // Decode params + TensorView float_workspace_buffer_d, TensorView int_workspace_buffer_d, + Array plan_info_vec_d, TensorView q_d, TensorView paged_k_cache_d, + TensorView paged_v_cache_d, TensorView qo_indptr_d, TensorView paged_kv_indptr_d, + TensorView paged_kv_indices_d, TensorView paged_kv_last_page_len_d, TensorView o_d, + Optional maybe_lse_d, int64_t mask_mode_code_d, int64_t layout_d, + int64_t window_left_d, Optional maybe_custom_mask_d, + Optional maybe_mask_indptr_d, Optional maybe_alibi_slopes_d, + double logits_soft_cap_d, double sm_scale_d, double rope_rcp_scale_d, double rope_rcp_theta_d, + bool enable_pdl, TensorView sm_aware_sched); + +// Batch-request prefill attention with KV-Cache operator +TVM_FFI_DLL_EXPORT_TYPED_FUNC(batch_pod_with_kv_cache_tensor, batch_pod_with_kv_cache_tensor); diff --git a/csrc/batch_pod_kernel_inst.jinja b/csrc/batch_pod_kernel_inst.jinja new file mode 100644 index 0000000000..cb2c39d32b --- /dev/null +++ b/csrc/batch_pod_kernel_inst.jinja @@ -0,0 +1,31 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "batch_pod_config.inc" + +using namespace flashinfer; + +namespace flashinfer { +constexpr auto use_custom_mask_p = {{ mask_mode_p }} == MaskMode::kCustom; +constexpr auto use_custom_mask_d = {{ mask_mode_d }} == MaskMode::kCustom; +// Not sure about the below declaration +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; + +{% for cta_tile_q in [16, 64, 128] %} +template cudaError_t BatchPODWithKVCacheTensorDispatched< + {{ head_dim_qk }}, {{ head_dim_vo }}, POS_ENCODING_MODE, + {{ use_fp16_qk_reduction }}, /*CTA_TILE_Q_P=*/{{cta_tile_q}}, {{ mask_mode_p }}, + /*CTA_TILE_Q_D=*/16, {{ mask_mode_d }}, {{ variant_name_p }}, + {{ variant_name_d }}, PrefillParams, DecodeParams>( + PrefillParams prefill_params, {{ dtype_o }}* tmp_v_p, float *tmp_s_p, + DecodeParams decode_params, {{ dtype_o }}* tmp_v_d, float *tmp_s_d, + bool enable_pdl, cudaStream_t stream, int* sm_aware_sched); +{% endfor %} +} diff --git a/csrc/batch_prefill.cu b/csrc/batch_prefill.cu index c5b5cefd64..9e0d77582f 100644 --- a/csrc/batch_prefill.cu +++ b/csrc/batch_prefill.cu @@ -16,11 +16,10 @@ #include #include #include -#include #include "batch_prefill_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" namespace flashinfer { @@ -28,54 +27,63 @@ template cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, - float* tmp_s, cudaStream_t stream); + float* tmp_s, bool enable_pdl, + cudaStream_t stream); template cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, - float* tmp_s, cudaStream_t stream); + float* tmp_s, bool enable_pdl, + cudaStream_t stream); } // namespace flashinfer using namespace flashinfer; -at::Tensor BatchPrefillWithKVCachePlan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchPrefillWithKVCachePlan( + TensorView float_workspace_buffer, TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, - int64_t head_dim_vo, bool causal) { + int64_t head_dim_vo, bool causal, int64_t window_left, int64_t fixed_split_size, + bool disable_split_kv, int64_t num_colocated_ctas = 0) { size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); PrefillPlanInfo plan_info; - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); cudaError_t status = PrefillPlan( float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), - int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), - kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, - head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), total_num_rows, batch_size, num_qo_heads, + num_kv_heads, head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, + /*sizeof_dtype_o=*/2, window_left, fixed_split_size, disable_split_kv, num_colocated_ctas, + stream); - TORCH_CHECK(status == cudaSuccess, - "Failed to plan prefill with error: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to plan prefill with error: " << cudaGetErrorString(status); - return vec_to_tensor(plan_info.ToVector()); + return Array(plan_info.ToVector()); } -void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS) { +void BatchPrefillWithRaggedKVCacheRun(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q, TensorView k, TensorView v, + TensorView qo_indptr, TensorView kv_indptr, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { PrefillPlanInfo plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); QKVLayout kv_layout = static_cast(layout); int64_t num_qo_heads = q.size(1); @@ -95,10 +103,10 @@ void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, v_stride_n = v.stride(1); } - if (maybe_lse) { + if (maybe_lse.has_value()) { const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); - TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + TVM_FFI_ICHECK_EQ(lse.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q.size(1)); } void* float_buffer_ptr = float_workspace_buffer.data_ptr(); @@ -106,11 +114,8 @@ void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, const MaskMode mask_mode = static_cast(mask_mode_code); - auto q_scalar_type = q.scalar_type(); - auto kv_scalar_type = k.scalar_type(); - - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); DISPATCH_context( DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, @@ -122,7 +127,8 @@ void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, params.k = static_cast(k.data_ptr()); params.v = static_cast(v.data_ptr()); params.o = static_cast(o.data_ptr()); - params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.lse = + maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.q_indptr = static_cast(qo_indptr.data_ptr()); params.kv_indptr = static_cast(kv_indptr.data_ptr()); params.num_qo_heads = num_qo_heads; @@ -185,27 +191,27 @@ void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, - RaggedParams>(params, tmp_v, tmp_s, stream); + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); }); - TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", - cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchPrefillWithRaggedKVCache failed with error " << cudaGetErrorString(status); return true; }); } -void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, - at::Tensor paged_v_cache, at::Tensor qo_indptr, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, - at::Tensor paged_kv_last_page_len, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS) { +void BatchPrefillWithPagedKVCacheRun(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q, TensorView paged_k_cache, + TensorView paged_v_cache, TensorView qo_indptr, + TensorView paged_kv_indptr, TensorView paged_kv_indices, + TensorView paged_kv_last_page_len, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { PrefillPlanInfo plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); QKVLayout kv_layout = static_cast(layout); - auto device = q.device(); int64_t batch_size = paged_kv_indptr.size(0) - 1; int64_t num_qo_heads = q.size(1); int64_t num_kv_heads, page_size; @@ -220,30 +226,29 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, if (maybe_lse) { const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); - TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + TVM_FFI_ICHECK_EQ(lse.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q.size(1)); } void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); const MaskMode mask_mode = static_cast(mask_mode_code); - auto q_scalar_type = q.scalar_type(); - auto kv_scalar_type = paged_k_cache.scalar_type(); // get q_stride_n and q_stride_h const auto q_stride_n = q.stride(0); const auto q_stride_h = q.stride(1); // get kv_cache_strides - const int64_t* kv_cache_strides = nullptr; - auto k_strides = paged_k_cache.strides(); - auto v_strides = paged_v_cache.strides(); - TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); - kv_cache_strides = k_strides.data(); + const int64_t* kv_cache_strides = paged_k_cache.strides().data(); + TVM_FFI_ICHECK_EQ(paged_k_cache.ndim(), paged_v_cache.ndim()); + for (int i = 0; i < paged_k_cache.ndim(); ++i) { + TVM_FFI_ICHECK_EQ(paged_k_cache.stride(i), paged_v_cache.stride(i)) + << "k/v strides differs at " << i; + } - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); DISPATCH_context( DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, @@ -263,7 +268,7 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, params.q_indptr = static_cast(qo_indptr.data_ptr()); params.o = static_cast(o.data_ptr()); - params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.lse = maybe_lse ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.num_qo_heads = num_qo_heads; params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); params.q_stride_n = q_stride_n; @@ -319,11 +324,11 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, - PagedParams>(params, tmp_v, tmp_s, stream); + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); }); - TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", - cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchPrefillWithPagedKVCache failed with error " << cudaGetErrorString(status); return true; }); } diff --git a/csrc/batch_prefill_config.inc b/csrc/batch_prefill_config.inc deleted file mode 100644 index 5dcd730b86..0000000000 --- a/csrc/batch_prefill_config.inc +++ /dev/null @@ -1,60 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -// NOTE(Zihao): this is the include file for AOT mode -#pragma once -#include -#include -#include -#include -#include -#include - -#include "aot_default_additional_params.h" -#include "aot_extension_utils.h" - -using IdType = int32_t; - -#define ADDITIONAL_FUNC_PARAMS BATCH_PREFILL_ADDITIONAL_FUNC_PARAMS -#define ADDITIONAL_PARAMS_SETTER BATCH_PREFILL_ADDITIONAL_PARAMS_SETTER - -#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, \ - POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, \ - USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ - { \ - DISPATCH_mask_mode(mask_mode, MASK_MODE, [&] { \ - return DISPATCH_PYTORCH_QKV_DTYPE_TO_CTYPE( \ - q_scalar_type, kv_scalar_type, DTypeQ, DTypeKV, [&] { \ - using DTypeO = DTypeQ; \ - using RaggedParams = BatchPrefillRaggedParams; \ - using PagedParams = BatchPrefillPagedParams; \ - constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; \ - constexpr bool USE_FP16_QK_REDUCTION = false; \ - constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ - return DISPATCH_head_dim(head_dim_qk, HEAD_DIM_QK, [&] { \ - [[maybe_unused]] constexpr int HEAD_DIM_VO = HEAD_DIM_QK; \ - return DISPATCH_BOOL(window_left > -1, USE_SLIDING_WINDOW, [&] { \ - return DISPATCH_BOOL(logits_soft_cap > 0.f, USE_LOGITS_SOFT_CAP, [&] { \ - using AttentionVariant = \ - DefaultAttention; \ - __VA_ARGS__(); \ - return true; \ - }); \ - }); \ - }); \ - }); \ - }); \ - } diff --git a/csrc/batch_prefill_fp8_paged_sm90_kernel_inst.jinja b/csrc/batch_prefill_fp8_paged_sm90_kernel_inst.jinja new file mode 100644 index 0000000000..f719a3313a --- /dev/null +++ b/csrc/batch_prefill_fp8_paged_sm90_kernel_inst.jinja @@ -0,0 +1,15 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + +{% for same_scheduler_for_all_heads in ["true", "false"] %} +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <{{ head_dim_qk }}, + {{ mask_mode }}, + /*USE_SLIDING_WINDOW=*/{{ use_sliding_window }}, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/{{ same_scheduler_for_all_heads }}, + {{ variant_name }}, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); +{% endfor %} + +}; // namespace flashinfer diff --git a/csrc/batch_prefill_fp8_ragged_sm90_kernel_inst.jinja b/csrc/batch_prefill_fp8_ragged_sm90_kernel_inst.jinja new file mode 100644 index 0000000000..8225edbb00 --- /dev/null +++ b/csrc/batch_prefill_fp8_ragged_sm90_kernel_inst.jinja @@ -0,0 +1 @@ +// TODO: Not implemented yet diff --git a/csrc/batch_prefill_fp8_sm90.cu b/csrc/batch_prefill_fp8_sm90.cu new file mode 100644 index 0000000000..7c8680dc0b --- /dev/null +++ b/csrc/batch_prefill_fp8_sm90.cu @@ -0,0 +1,188 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "tvm_ffi_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchPrefillWithKVCacheSM90Plan( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + ffi::TensorView page_locked_int_workspace_buffer, ffi::TensorView qo_indptr, + ffi::TensorView kv_indptr, ffi::TensorView kv_len_arr, int64_t total_num_rows, + int64_t batch_size, int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, + bool enable_cuda_graph, int64_t head_dim_qk, int64_t head_dim_vo, bool causal, + int64_t window_left) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); + + flashinfer::PrefillPlanSM90Info plan_info; + + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); + + cudaError_t status = PrefillSM90Plan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len_arr.data_ptr()), + total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, + /*sizeof_dtype_o=*/2, stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "PrefillSM90Plan failed with error: " << cudaGetErrorString(status); + + return Array(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run(ffi::TensorView float_workspace_buffer, + ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q, + ffi::TensorView k, ffi::TensorView v, + ffi::TensorView qo_indptr, ffi::TensorView kv_indptr, + ffi::TensorView o, Optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left, + bool enable_pdl // placeholder + ADDITIONAL_FUNC_PARAMS) { + return; // TODO: Implement this function +} + +void BatchPrefillWithPagedKVCacheSM90Run( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q, ffi::TensorView paged_k_cache, + ffi::TensorView paged_v_cache, ffi::TensorView qo_indptr, ffi::TensorView paged_kv_indptr, + ffi::TensorView paged_kv_indices, ffi::TensorView paged_kv_last_page_len, ffi::TensorView o, + Optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); + + if (maybe_lse.has_value()) { + const auto& lse = maybe_lse.value(); + TVM_FFI_ICHECK_EQ(lse.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse.value().data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = + BatchFP8PrefillWithPagedKVCacheDispatched(params, enable_pdl, + stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchPrefillWithPagedKVCacheSM90Run failed with error: " + << cudaGetErrorString(status); + return true; + }); + }); +} diff --git a/csrc/batch_prefill_jit_binding.cu b/csrc/batch_prefill_jit_binding.cu new file mode 100644 index 0000000000..3dda0f115a --- /dev/null +++ b/csrc/batch_prefill_jit_binding.cu @@ -0,0 +1,50 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchPrefillWithKVCachePlan( + TensorView float_workspace_buffer, TensorView int_workspace_buffer, + TensorView page_locked_int_workspace_buffer, TensorView qo_indptr, TensorView kv_indptr, + TensorView kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal, int64_t window_left, int64_t fixed_split_size, + bool disable_split_kv, int64_t num_colocated_ctas); + +void BatchPrefillWithRaggedKVCacheRun(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q, TensorView k, TensorView v, + TensorView qo_indptr, TensorView kv_indptr, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun(TensorView float_workspace_buffer, + TensorView int_workspace_buffer, Array plan_info_vec, + TensorView q, TensorView paged_k_cache, + TensorView paged_v_cache, TensorView qo_indptr, + TensorView paged_kv_indptr, TensorView paged_kv_indices, + TensorView paged_kv_last_page_len, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchPrefillWithKVCachePlan); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(ragged_run, BatchPrefillWithRaggedKVCacheRun); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(paged_run, BatchPrefillWithPagedKVCacheRun); diff --git a/csrc/batch_prefill_jit_pybind.cu b/csrc/batch_prefill_jit_pybind.cu deleted file mode 100644 index 024fcde70e..0000000000 --- a/csrc/batch_prefill_jit_pybind.cu +++ /dev/null @@ -1,49 +0,0 @@ -/* - * Copyright (c) 2023-2025 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "batch_prefill_config.inc" -#include "pytorch_extension_utils.h" - -at::Tensor BatchPrefillWithKVCachePlan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, - int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, - int64_t head_dim_vo, bool causal); - -void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS); - -void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, - at::Tensor paged_v_cache, at::Tensor qo_indptr, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, - at::Tensor paged_kv_last_page_len, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // Batch-request prefill attention with KV-Cache plan - m.def("plan", BatchPrefillWithKVCachePlan); - // Batch-request prefill attention with KV-Cache operator - m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); - // Batch-request prefill attention with KV-Cache operator - m.def("paged_run", BatchPrefillWithPagedKVCacheRun); -} diff --git a/csrc/batch_prefill_paged_kernel_inst.jinja b/csrc/batch_prefill_paged_kernel_inst.jinja index 6678877035..f62bfecf68 100644 --- a/csrc/batch_prefill_paged_kernel_inst.jinja +++ b/csrc/batch_prefill_paged_kernel_inst.jinja @@ -8,7 +8,7 @@ constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom; {% for cta_tile_q in [16, 64, 128] %} template cudaError_t BatchPrefillWithPagedKVCacheDispatched< /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}}, - {{ variant_name }}, PagedParams>(PagedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, cudaStream_t stream); + {{ variant_name }}, PagedParams>(PagedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); {% endfor %} }; // namespace flashinfer diff --git a/csrc/batch_prefill_paged_sm90_kernel_inst.jinja b/csrc/batch_prefill_paged_sm90_kernel_inst.jinja index 9e159e0064..ec937a94b3 100644 --- a/csrc/batch_prefill_paged_sm90_kernel_inst.jinja +++ b/csrc/batch_prefill_paged_sm90_kernel_inst.jinja @@ -10,7 +10,7 @@ template cudaError_t BatchPrefillWithPagedKVCacheDispatched {{ mask_mode }}, /*USE_SLIDING_WINDOW=*/{{ use_sliding_window }}, /*SAME_SCHEDULER_FOR_ALL_HEADS=*/{{ same_scheduler_for_all_heads }}, - {{ variant_name }}, PagedParams>(PagedParams& params, cudaStream_t stream); + {{ variant_name }}, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); {% endfor %} }; // namespace flashinfer diff --git a/csrc/batch_prefill_ragged_kernel_inst.jinja b/csrc/batch_prefill_ragged_kernel_inst.jinja index 39c49227cd..bdd1408184 100644 --- a/csrc/batch_prefill_ragged_kernel_inst.jinja +++ b/csrc/batch_prefill_ragged_kernel_inst.jinja @@ -8,7 +8,7 @@ constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom; {% for cta_tile_q in [16, 64, 128] %} template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}}, - {{ variant_name }}, RaggedParams>(RaggedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, cudaStream_t stream); + {{ variant_name }}, RaggedParams>(RaggedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); {% endfor %} }; // namespace flashinfer diff --git a/csrc/batch_prefill_ragged_sm90_kernel_inst.jinja b/csrc/batch_prefill_ragged_sm90_kernel_inst.jinja index 90a44b4d79..d13d700606 100644 --- a/csrc/batch_prefill_ragged_sm90_kernel_inst.jinja +++ b/csrc/batch_prefill_ragged_sm90_kernel_inst.jinja @@ -10,7 +10,7 @@ template cudaError_t BatchPrefillWithRaggedKVCacheDispatched {{ mask_mode }}, /*USE_SLIDING_WINDOW=*/{{ use_sliding_window }}, /*SAME_SCHEDULER_FOR_ALL_HEADS=*/{{ same_scheduler_for_all_heads }}, - {{ variant_name }}>(RaggedParams& params, cudaStream_t stream); + {{ variant_name }}>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); {% endfor %} }; // namespace flashinfer diff --git a/csrc/batch_prefill_sm90.cu b/csrc/batch_prefill_sm90.cu index 66a68b73f7..1cf78bab59 100644 --- a/csrc/batch_prefill_sm90.cu +++ b/csrc/batch_prefill_sm90.cu @@ -18,70 +18,75 @@ #include #include #include -#include #include "batch_prefill_sm90_config.inc" -#include "pytorch_conversion_utils.h" -#include "pytorch_extension_utils.h" +#include "tvm/ffi/container/array.h" +#include "tvm_ffi_utils.h" namespace flashinfer { template -cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, cudaStream_t stream); +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); template -cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, cudaStream_t stream); +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); } // namespace flashinfer using namespace flashinfer; -at::Tensor BatchPrefillWithKVCacheSM90Plan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, - int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, - int64_t head_dim_vo, bool causal) { +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchPrefillWithKVCacheSM90Plan( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + ffi::TensorView page_locked_int_workspace_buffer, ffi::TensorView qo_indptr, + ffi::TensorView kv_indptr, ffi::TensorView kv_len_arr, int64_t total_num_rows, + int64_t batch_size, int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, + bool enable_cuda_graph, int64_t head_dim_qk, int64_t head_dim_vo, bool causal, + int64_t window_left) { size_t float_workspace_size_in_bytes = - float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer); size_t int_workspace_size_in_bytes = - int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + int_workspace_buffer.size(0) * get_element_size(int_workspace_buffer); flashinfer::PrefillPlanSM90Info plan_info; - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); - cudaError_t status = - PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, - int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), - int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), - kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, - batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, - causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + cudaError_t status = PrefillSM90Plan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len_arr.data_ptr()), + total_num_rows, batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, + /*sizeof_dtype_o=*/2, stream); - TORCH_CHECK(status == cudaSuccess, - "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "PrefillSM90Plan failed with error: " << cudaGetErrorString(status); - return vec_to_tensor(plan_info.ToVector()); + return Array(plan_info.ToVector()); } -void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, - std::optional maybe_lse, - int64_t mask_mode_code, int64_t layout, - int64_t window_left ADDITIONAL_FUNC_PARAMS) { +void BatchPrefillWithRaggedKVCacheSM90Run( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView qo_indptr, ffi::TensorView kv_indptr, ffi::TensorView o, + Optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { PrefillPlanSM90Info plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); if (maybe_lse) { const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); - TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + TVM_FFI_ICHECK_EQ(lse.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q.size(1)); } void* float_buffer_ptr = float_workspace_buffer.data_ptr(); @@ -90,12 +95,10 @@ void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, int64_t head_dim_qk = q.size(2); int64_t head_dim_vo = v.size(2); - auto q_scalar_type = q.scalar_type(); - auto kv_scalar_type = k.scalar_type(); - QKVLayout kv_layout = static_cast(layout); - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); const MaskMode mask_mode = static_cast(mask_mode_code); bool use_swa = window_left != -1; @@ -108,7 +111,7 @@ void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, params.k_ptr = static_cast(k.data_ptr()); params.v_ptr = static_cast(v.data_ptr()); params.o_ptr = static_cast(o.data_ptr()); - params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.lse_ptr = maybe_lse ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.q_stride_n = q.stride(0); params.q_stride_h = q.stride(1); params.o_stride_n = o.stride(0); @@ -141,6 +144,8 @@ void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); params.work_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); ADDITIONAL_PARAMS_SETTER @@ -148,28 +153,30 @@ void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, - AttentionVariant>(params, stream); - TORCH_CHECK(status == cudaSuccess, - "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", - cudaGetErrorString(status)); + AttentionVariant>(params, enable_pdl, stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchPrefillWithRaggedKVCacheSM90Run failed with error: " + << cudaGetErrorString(status); return true; }); }); } void BatchPrefillWithPagedKVCacheSM90Run( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, - int64_t window_left ADDITIONAL_FUNC_PARAMS) { + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q, ffi::TensorView paged_k_cache, + ffi::TensorView paged_v_cache, ffi::TensorView qo_indptr, ffi::TensorView paged_kv_indptr, + ffi::TensorView paged_kv_indices, ffi::TensorView paged_kv_last_page_len, ffi::TensorView o, + Optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { PrefillPlanSM90Info plan_info; - plan_info.FromVector(tensor_to_vec(plan_info_vec)); + plan_info.FromVector(std::vector(plan_info_vec.begin(), plan_info_vec.end())); if (maybe_lse) { const auto& lse = *maybe_lse; - TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); - TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + TVM_FFI_ICHECK_EQ(lse.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(lse.size(1), q.size(1)); } QKVLayout kv_layout = static_cast(layout); int64_t num_kv_heads, page_size; @@ -186,11 +193,8 @@ void BatchPrefillWithPagedKVCacheSM90Run( void* float_buffer_ptr = float_workspace_buffer.data_ptr(); void* int_buffer_ptr = int_workspace_buffer.data_ptr(); - auto q_scalar_type = q.scalar_type(); - auto kv_scalar_type = paged_k_cache.scalar_type(); - - const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); - cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaSetDevice(float_workspace_buffer.device().device_id); + const cudaStream_t stream = get_stream(float_workspace_buffer.device()); const MaskMode mask_mode = static_cast(mask_mode_code); bool use_swa = window_left != -1; @@ -203,7 +207,7 @@ void BatchPrefillWithPagedKVCacheSM90Run( params.k_ptr = static_cast(paged_k_cache.data_ptr()); params.v_ptr = static_cast(paged_v_cache.data_ptr()); params.o_ptr = static_cast(o.data_ptr()); - params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.lse_ptr = maybe_lse ? static_cast(maybe_lse.value().data_ptr()) : nullptr; params.q_stride_n = q.stride(0); params.q_stride_h = q.stride(1); params.o_stride_n = o.stride(0); @@ -238,6 +242,8 @@ void BatchPrefillWithPagedKVCacheSM90Run( GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); params.work_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); params.kv_indices = static_cast(paged_kv_indices.data_ptr()); ADDITIONAL_PARAMS_SETTER @@ -246,10 +252,11 @@ void BatchPrefillWithPagedKVCacheSM90Run( DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, - AttentionVariant>(params, stream); - TORCH_CHECK(status == cudaSuccess, - "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", - cudaGetErrorString(status)); + AttentionVariant>(params, enable_pdl, stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "BatchPrefillWithPagedKVCacheSM90Run failed with error: " + << cudaGetErrorString(status); return true; }); }); diff --git a/csrc/batch_prefill_sm90_config.inc b/csrc/batch_prefill_sm90_config.inc deleted file mode 100644 index e74f1a40b0..0000000000 --- a/csrc/batch_prefill_sm90_config.inc +++ /dev/null @@ -1,55 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -// NOTE(Zihao): this is the include file for AOT mode -#pragma once -#include -#include -#include - -#include "aot_default_additional_params.h" -#include "aot_extension_utils.h" - -using IdType = int32_t; - -#define ADDITIONAL_FUNC_PARAMS BATCH_PREFILL_SM90_ADDITIONAL_FUNC_PARAMS -#define ADDITIONAL_PARAMS_SETTER BATCH_PREFILL_SM90_ADDITIONAL_PARAMS_SETTER - -#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, \ - USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, \ - PagedParams, ...) \ - { \ - DISPATCH_mask_mode(mask_mode, MASK_MODE, [&] { \ - if (q_scalar_type != kv_scalar_type) { \ - return false; \ - } \ - return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(q_scalar_type, dtype_q, [&] { \ - using DTypeQ = cutlass_dtype_t; \ - using DTypeKV = DTypeQ; \ - using DTypeO = DTypeQ; \ - using RaggedParams = BatchPrefillRaggedParams; \ - using PagedParams = BatchPrefillPagedParams; \ - return DISPATCH_head_dim_sm90(head_dim_qk, head_dim_vo, HEAD_DIM_QK, HEAD_DIM_VO, [&] { \ - return DISPATCH_BOOL(window_left > -1, USE_SLIDING_WINDOW, [&] { \ - return DISPATCH_BOOL(logits_soft_cap > 0.f, USE_LOGITS_SOFT_CAP, [&] { \ - using AttentionVariant = DefaultAttention; \ - __VA_ARGS__(); \ - return true; \ - }); \ - }); \ - }); \ - }); \ - }); \ - } diff --git a/csrc/batch_prefill_sm90_customize_config.jinja b/csrc/batch_prefill_sm90_customize_config.jinja index 5b10355fc8..b37ecac60d 100644 --- a/csrc/batch_prefill_sm90_customize_config.jinja +++ b/csrc/batch_prefill_sm90_customize_config.jinja @@ -4,7 +4,6 @@ #include #include #include -#include "pytorch_extension_utils.h" #define ADDITIONAL_FUNC_PARAMS {{ additional_func_params }} #define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }} @@ -43,6 +42,7 @@ struct RaggedParams { IdType* kv_lens; IdType* head_indices; IdType* work_indptr; + IdType* batch_indices; struct AdditionalParams { {{ additional_params_decl }} @@ -88,6 +88,7 @@ struct PagedParams { IdType* kv_lens; IdType* head_indices; IdType* work_indptr; + IdType* batch_indices; struct AdditionalParams { {{ additional_params_decl }} diff --git a/csrc/batch_prefill_sm90_jit_binding.cu b/csrc/batch_prefill_sm90_jit_binding.cu new file mode 100644 index 0000000000..259a1e22d9 --- /dev/null +++ b/csrc/batch_prefill_sm90_jit_binding.cu @@ -0,0 +1,49 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "batch_prefill_sm90_config.inc" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Array; +using tvm::ffi::Optional; + +Array BatchPrefillWithKVCacheSM90Plan( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + ffi::TensorView page_locked_int_workspace_buffer, ffi::TensorView qo_indptr, + ffi::TensorView kv_indptr, ffi::TensorView kv_len_arr, int64_t total_num_rows, + int64_t batch_size, int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, + bool enable_cuda_graph, int64_t head_dim_qk, int64_t head_dim_vo, bool causal, + int64_t window_left); + +void BatchPrefillWithRaggedKVCacheSM90Run( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView qo_indptr, ffi::TensorView kv_indptr, ffi::TensorView o, + Optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + ffi::TensorView float_workspace_buffer, ffi::TensorView int_workspace_buffer, + Array plan_info_vec, ffi::TensorView q, ffi::TensorView paged_k_cache, + ffi::TensorView paged_v_cache, ffi::TensorView qo_indptr, ffi::TensorView paged_kv_indptr, + ffi::TensorView paged_kv_indices, ffi::TensorView paged_kv_last_page_len, ffi::TensorView o, + Optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, BatchPrefillWithKVCacheSM90Plan); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(ragged_run, BatchPrefillWithRaggedKVCacheSM90Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(paged_run, BatchPrefillWithPagedKVCacheSM90Run); diff --git a/csrc/batch_prefill_sm90_jit_pybind.cu b/csrc/batch_prefill_sm90_jit_pybind.cu deleted file mode 100644 index 0b7e964c85..0000000000 --- a/csrc/batch_prefill_sm90_jit_pybind.cu +++ /dev/null @@ -1,48 +0,0 @@ -/* - * Copyright (c) 2023-2025 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "batch_prefill_sm90_config.inc" -#include "pytorch_extension_utils.h" - -at::Tensor BatchPrefillWithKVCacheSM90Plan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, - int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, - int64_t head_dim_vo, bool causal); - -void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, - std::optional maybe_lse, - int64_t mask_mode_code, int64_t layout, - int64_t window_left ADDITIONAL_FUNC_PARAMS); - -void BatchPrefillWithPagedKVCacheSM90Run( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, - int64_t window_left ADDITIONAL_FUNC_PARAMS); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // Batch-request prefill attention with KV-Cache plan - m.def("plan", BatchPrefillWithKVCacheSM90Plan); - // Batch-request prefill attention with KV-Cache operator - m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); - // Batch-request prefill attention with KV-Cache operator - m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); -} diff --git a/csrc/blackwell_fmha_plan.cu b/csrc/blackwell_fmha_plan.cu new file mode 100644 index 0000000000..ef9b1475ea --- /dev/null +++ b/csrc/blackwell_fmha_plan.cu @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "flashinfer/attention/blackwell/plan.cuh" +#include "tvm_ffi_utils.h" + +void blackwell_fmha_plan(TensorView qo_segment_offsets, TensorView kv_segment_offsets, + TensorView work_indptr, TensorView qo_tile_indices, + TensorView head_indices, TensorView batch_indices, int64_t qo_tile_size, + int64_t num_heads, int64_t num_buckets, bool causal) { + cudaSetDevice(qo_segment_offsets.device().device_id); + const cudaStream_t stream = get_stream(qo_tile_indices.device()); + int batch_size = qo_segment_offsets.size(0) - 1; + + auto status = flashinfer::plan_kernel_wrapper( + static_cast(qo_segment_offsets.data_ptr()), + static_cast(kv_segment_offsets.data_ptr()), + /*qo_lens=*/nullptr, + /*kv_lens=*/nullptr, static_cast(work_indptr.data_ptr()), + static_cast(qo_tile_indices.data_ptr()), static_cast(head_indices.data_ptr()), + static_cast(batch_indices.data_ptr()), qo_tile_size, batch_size, num_heads, num_buckets, + causal, /*enable_pdl=*/true, stream); + TVM_FFI_ICHECK_EQ(status, cudaSuccess) + << "Failed to plan blackwell fmha" << cudaGetErrorString(status); +} diff --git a/csrc/bmm_fp8.cu b/csrc/bmm_fp8.cu index f640b645c7..ea8417b617 100644 --- a/csrc/bmm_fp8.cu +++ b/csrc/bmm_fp8.cu @@ -18,35 +18,35 @@ #include -#include "pytorch_extension_utils.h" +#include "tvm_ffi_utils.h" -void bmm_fp8(at::Tensor A, at::Tensor B, at::Tensor D, at::Tensor A_scale, at::Tensor B_scale, - at::Tensor workspace_buffer, int64_t cublas_handle) { - TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor"); - TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor"); - TORCH_CHECK(D.is_cuda(), "D must be a CUDA tensor"); - TORCH_CHECK(A.dim() == 3, "Expected 3D tensor for A"); - TORCH_CHECK(B.dim() == 3, "Expected 3D tensor for B"); - TORCH_CHECK(D.dim() == 3, "Expected 3D tensor for D"); - TORCH_CHECK(A.size(0) == B.size(0) && A.size(0) == D.size(0), "Batch sizes must match"); - TORCH_CHECK(A.size(2) == B.size(1), "Incompatible matrix sizes"); - TORCH_CHECK(A.size(1) == D.size(1) && B.size(2) == D.size(2), - "Result tensor has incorrect shape"); +void bmm_fp8(TensorView A, TensorView B, TensorView D, TensorView A_scale, TensorView B_scale, + TensorView workspace_buffer, int64_t cublas_handle) { + CHECK_CUDA(A); + CHECK_CUDA(B); + CHECK_CUDA(D); + CHECK_DIM(3, A); + CHECK_DIM(3, B); + CHECK_DIM(3, D); + TVM_FFI_ICHECK(A.size(0) == B.size(0) && A.size(0) == D.size(0)) << "Batch sizes must match"; + TVM_FFI_ICHECK(A.size(2) == B.size(1)) << "Incompatible matrix sizes"; + TVM_FFI_ICHECK(A.size(1) == D.size(1) && B.size(2) == D.size(2)) + << "Result tensor has incorrect shape"; // PyTorch is row major by default. cuBLASLt is column major by default. // We need row major D as expected. // A ^ T * B = D, so D ^ T = B ^ T * A - DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8(B.scalar_type(), b_type, [&] { - return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8(A.scalar_type(), a_type, [&] { - return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(D.scalar_type(), d_type, [&] { + DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(B.dtype(), b_type, [&] { + return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(A.dtype(), a_type, [&] { + return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(D.dtype(), d_type, [&] { auto batch_size = A.size(0); auto m = A.size(1); auto k = A.size(2); auto n = B.size(2); auto lt_handle = reinterpret_cast(cublas_handle); - const c10::cuda::OptionalCUDAGuard device_guard(A.device()); - auto stream = at::cuda::getCurrentCUDAStream(); + cudaSetDevice(A.device().device_id); + auto stream = get_stream(A.device()); auto status = flashinfer::bmm_fp8::bmm_fp8_internal_cublaslt( workspace_buffer.data_ptr(), workspace_buffer.numel(), @@ -54,8 +54,8 @@ void bmm_fp8(at::Tensor A, at::Tensor B, at::Tensor D, at::Tensor A_scale, at::T static_cast(D.data_ptr()), batch_size, n, m, k, static_cast(B_scale.data_ptr()), static_cast(A_scale.data_ptr()), lt_handle, stream); - TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, - "bmm_fp8_internal_cublaslt failed: ", cublasGetStatusString(status)); + TVM_FFI_ICHECK(status == CUBLAS_STATUS_SUCCESS) + << "bmm_fp8_internal_cublaslt failed: " << cublasGetStatusString(status); return true; }); }); diff --git a/csrc/cascade.cu b/csrc/cascade.cu index eee2bcd07e..98e4a590dc 100644 --- a/csrc/cascade.cu +++ b/csrc/cascade.cu @@ -15,119 +15,116 @@ */ #include -#include "pytorch_extension_utils.h" +#include "tvm_ffi_utils.h" using namespace flashinfer; +using tvm::ffi::Optional; -void merge_state(at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, - at::Tensor v_merged, at::Tensor s_merged) { +void merge_state(TensorView v_a, TensorView s_a, TensorView v_b, TensorView s_b, + TensorView v_merged, TensorView s_merged) { CHECK_INPUT(v_a); CHECK_INPUT(s_a); CHECK_INPUT(v_b); CHECK_INPUT(s_b); - auto device = v_a.device(); - CHECK_EQ(s_a.device(), device); - CHECK_EQ(v_b.device(), device); - CHECK_EQ(s_b.device(), device); + CHECK_DEVICE(s_a, v_a); + CHECK_DEVICE(v_b, v_a); + CHECK_DEVICE(s_b, v_a); CHECK_DIM(3, v_a); CHECK_DIM(2, s_a); CHECK_DIM(3, v_b); CHECK_DIM(2, s_b); CHECK_SHAPE(v_a, v_b); CHECK_SHAPE(s_a, s_b); - CHECK_EQ(v_a.size(0), s_a.size(0)); - CHECK_EQ(v_a.size(1), s_b.size(1)); + TVM_FFI_ICHECK_EQ(v_a.size(0), s_a.size(0)); + TVM_FFI_ICHECK_EQ(v_a.size(1), s_b.size(1)); unsigned int seq_len = v_a.size(0); unsigned int num_heads = v_a.size(1); unsigned int head_dim = v_a.size(2); - const c10::cuda::OptionalCUDAGuard device_guard(v_a.device()); - auto stream = at::cuda::getCurrentCUDAStream(); + cudaSetDevice(v_a.device().device_id); + auto stream = get_stream(v_a.device()); - bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v_a.scalar_type(), c_type, [&] { + bool success = DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(v_a.dtype(), c_type, [&] { cudaError_t status = MergeState(static_cast(v_a.data_ptr()), static_cast(s_a.data_ptr()), static_cast(v_b.data_ptr()), static_cast(s_b.data_ptr()), static_cast(v_merged.data_ptr()), static_cast(s_merged.data_ptr()), seq_len, num_heads, head_dim, stream); - TORCH_CHECK(status == cudaSuccess, - "MergeState kernel launch failed: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "MergeState kernel launch failed: " << cudaGetErrorString(status); return true; }); - - TORCH_CHECK(success, "MergeState kernel launch failed: unsupported data type"); + TVM_FFI_ICHECK(success) << "MergeState kernel launch failed: unsupported data type."; } -void merge_state_in_place(at::Tensor v, at::Tensor s, at::Tensor v_other, at::Tensor s_other, - std::optional mask) { +void merge_state_in_place(TensorView v, TensorView s, TensorView v_other, TensorView s_other, + Optional mask) { CHECK_INPUT(v); CHECK_INPUT(s); CHECK_INPUT(v_other); CHECK_INPUT(s_other); - auto device = v.device(); - CHECK_EQ(s.device(), device); - CHECK_EQ(v_other.device(), device); - CHECK_EQ(s_other.device(), device); + CHECK_DEVICE(s, v); + CHECK_DEVICE(v_other, v); + CHECK_DEVICE(s_other, v); CHECK_DIM(3, v); CHECK_DIM(2, s); CHECK_DIM(3, v_other); CHECK_DIM(2, s_other); CHECK_SHAPE(v, v_other); CHECK_SHAPE(s, s_other); - CHECK_EQ(v.size(0), s.size(0)); - CHECK_EQ(v.size(1), s.size(1)); + TVM_FFI_ICHECK_EQ(v.size(0), s.size(0)); + TVM_FFI_ICHECK_EQ(v.size(1), s.size(1)); uint8_t* mask_ptr = nullptr; if (mask.has_value()) { CHECK_DIM(1, mask.value()); - CHECK_EQ(v.size(0), mask.value().size(0)); - CHECK_EQ(mask.value().device(), device); + TVM_FFI_ICHECK_EQ(v.size(0), mask.value().size(0)); + CHECK_DEVICE(mask.value(), v); mask_ptr = static_cast(mask.value().data_ptr()); } unsigned int seq_len = v.size(0); unsigned int num_heads = v.size(1); unsigned int head_dim = v.size(2); - const c10::cuda::OptionalCUDAGuard device_guard(v.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v.scalar_type(), c_type, [&] { + cudaSetDevice(v.device().device_id); + auto stream = get_stream(v.device()); + bool success = DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(v.dtype(), c_type, [&] { cudaError_t status = MergeStateInPlace( static_cast(v.data_ptr()), static_cast(s.data_ptr()), static_cast(v_other.data_ptr()), static_cast(s_other.data_ptr()), seq_len, num_heads, head_dim, mask_ptr, stream); - TORCH_CHECK(status == cudaSuccess, - "MergeStateInPlace kernel launch failed: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "MergeStateInPlace kernel launch failed: " << cudaGetErrorString(status); return true; }); - TORCH_CHECK(success, "MergeStateInPlace kernel launch failed: unsupported data type"); + TVM_FFI_ICHECK(success) << "MergeStateInPlace kernel launch failed: unsupported data type."; } -void merge_states(at::Tensor v, at::Tensor s, at::Tensor v_merged, at::Tensor s_merged) { +void merge_states(TensorView v, TensorView s, TensorView v_merged, TensorView s_merged) { CHECK_INPUT(v); CHECK_INPUT(s); - auto device = v.device(); - CHECK_EQ(s.device(), device); + CHECK_DEVICE(s, v); CHECK_DIM(4, v); CHECK_DIM(3, s); - CHECK_EQ(v.size(0), s.size(0)); - CHECK_EQ(v.size(1), s.size(1)); - CHECK_EQ(v.size(2), s.size(2)); + TVM_FFI_ICHECK_EQ(v.size(0), s.size(0)); + TVM_FFI_ICHECK_EQ(v.size(1), s.size(1)); + TVM_FFI_ICHECK_EQ(v.size(2), s.size(2)); unsigned int seq_len = v.size(0); unsigned int num_index_sets = v.size(1); unsigned int num_heads = v.size(2); unsigned int head_dim = v.size(3); - const c10::cuda::OptionalCUDAGuard device_guard(v.device()); - auto stream = at::cuda::getCurrentCUDAStream(); - bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v.scalar_type(), c_type, [&] { + cudaSetDevice(v.device().device_id); + auto stream = get_stream(v.device()); + bool success = DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(v.dtype(), c_type, [&] { cudaError_t status = MergeStates( static_cast(v.data_ptr()), static_cast(s.data_ptr()), static_cast(v_merged.data_ptr()), static_cast(s_merged.data_ptr()), num_index_sets, seq_len, num_heads, head_dim, stream); - TORCH_CHECK(status == cudaSuccess, - "MergeStates kernel launch failed: ", cudaGetErrorString(status)); + TVM_FFI_ICHECK(status == cudaSuccess) + << "MergeStates kernel launch failed: " << cudaGetErrorString(status); return true; }); - TORCH_CHECK(success, "MergeStates kernel launch failed: unsupported data type"); + TVM_FFI_ICHECK(success) << "MergeStates kernel launch failed: unsupported data type."; } diff --git a/csrc/cudnn_sdpa_kernel_launcher.cu b/csrc/cudnn_sdpa_kernel_launcher.cu new file mode 100644 index 0000000000..12f7ac1dad --- /dev/null +++ b/csrc/cudnn_sdpa_kernel_launcher.cu @@ -0,0 +1,1202 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cudnn_sdpa_utils.h" +#include "tvm_ffi_utils.h" + +#ifdef CUDNN_SDPA_CUBIN_PATH +static const std::string cudnn_sdpa_cubin_path = std::string(CUDNN_SDPA_CUBIN_PATH); +#else +static_assert(false, "CUDNN_SDPA_CUBIN_PATH macro is not defined when compiling"); +#endif + +namespace flashinfer { + +namespace cudnn_sdpa_kernel_launcher { + +#include + +using tvm::ffi::Optional; + +inline __host__ int clz(int x) { + for (int i = 31; i >= 0; --i) { + if ((1 << i) & x) { + return 31 - i; + } + } + return 32; +} + +inline __host__ int find_log_2(int x, bool round_up = false) { + int a = 31 - clz(x); + if (round_up) { + a += (x & (x - 1)) ? 1 : 0; + } + return a; +} + +inline __host__ void setFastDivisor(cudnn_sdpa::FastDivisor_t& d, uint32_t val) { + uint32_t p = 31 + find_log_2(2 * val, true); + uint32_t m = (uint32_t)(((1ull << p) + (uint32_t)(2 * val) - 1) / (uint32_t)(2 * val)); + + d.val = val; + d.mul = m; + d.shr = p - 32; +} + +static std::once_flag init_cudnn_cubin_flag; + +constexpr size_t DIMS_QKV = 4; +constexpr int32_t BYTES_PER_ELEMENT = 2; + +enum KernelType { PREFILL, PREFILL_DEEPSEEK, DECODE }; + +enum PrefillType { + KERNEL_PREFILL, + KERNEL_PREFILL_DEEPSEEK, + KERNEL_PREFILL_CAUSAL, + KERNEL_PREFILL_DEEPSEEK_CAUSAL, + KERNEL_NUM_PREFILL_TYPES +}; + +void init_cudnn_cubin(std::map& cubin_map) { + cubin_map[PREFILL] = + getCubin(cudnn_sdpa_cubin_path + "/" + "cudnn_sm100_fprop_sdpa_prefill_d128_bf16.cubin", + "ff14e8dcfc04d9b3a912dd44056be37d9aa8a85976e0070494ca0cce0524f2a1"); + + cubin_map[DECODE] = + getCubin(cudnn_sdpa_cubin_path + "/" + "cudnn_sm100_fprop_sdpa_decode_d128_bf16.cubin", + "e7ce0408b4c3a36c42616498228534ee64cab785ef570af5741deaf9dd1b475c"); + + cubin_map[PREFILL_DEEPSEEK] = + getCubin(cudnn_sdpa_cubin_path + "/" + "cudnn_sm100_fprop_sdpa_prefill_d192_bf16.cubin", + "2190967b8733e193cdcecc054eeb7c2907080a158a33fe7ba2004523a4aff6f9"); +} + +auto get_cudnn_cubin(KernelType kernel_type) -> std::string { + static std::map cubin_map; + std::call_once(init_cudnn_cubin_flag, init_cudnn_cubin, std::ref(cubin_map)); + return cubin_map[kernel_type]; +} + +__global__ static void __launch_bounds__(128) + qkv_tma_setup_decode(const unsigned int b, const unsigned int h_qo, const unsigned int h_kv, + const unsigned int d, const unsigned int total_num_pages, + const unsigned int page_size, const unsigned int split_factor, + const unsigned int tile_m_1, const unsigned int tile_n_1, + const unsigned int kv_strides_2, const unsigned int kv_strides_1, + const unsigned int kv_strides_0, void* q_ptr, const void* k_ptr, + const void* v_ptr, void* o_ptr, void* partial_o_ptr, + tma::cudaTmaDesc* tma_desc_q_array, tma::cudaTmaDesc* tma_desc_k, + tma::cudaTmaDesc* tma_desc_v, tma::cudaTmaDesc* tma_desc_o_array, + tma::cudaTmaDesc* tma_desc_partial_o_array, int64_t* batch_strides_dev) { + const int tid = threadIdx.x; + + constexpr unsigned int DIMS_QKV = 4; + constexpr unsigned int BYTES_PER_ELEMENT = 2; + + std::array tensor_traversal_stride_qkv = {1, 1, 1, 1}; + std::array tensor_box_size_qo = {64, 1, 1, 1}; + std::array tensor_box_size_kv = {64, std::min(tile_n_1, page_size), 1, 1}; + std::array tensor_box_size_partial_o = {32, 1, 1, 1}; + + std::array tensor_size_qo = {d, 1 /* s_qo */, h_qo, b}; + std::array tensor_size_kv = {d, page_size, h_kv, total_num_pages}; + + std::array tensor_stride_qo = {h_qo * d * BYTES_PER_ELEMENT, + d * BYTES_PER_ELEMENT, 0}; + std::array tensor_stride_kv = {kv_strides_2 * (BYTES_PER_ELEMENT), + kv_strides_1 * (BYTES_PER_ELEMENT), + kv_strides_0 * (BYTES_PER_ELEMENT)}; + + std::array tensor_size_partial_o = {d, split_factor, h_qo, b}; + std::array tensor_stride_partial_o = { + h_qo * d * b * sizeof(float), d * b * sizeof(float), d * h_qo * sizeof(float)}; + + tma::cudaSetTmaTileDescriptor( + reinterpret_cast(tma_desc_k), k_ptr, DIMS_QKV, tensor_size_kv.data(), + tensor_stride_kv.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_kv.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + reinterpret_cast(tma_desc_v), v_ptr, DIMS_QKV, tensor_size_kv.data(), + tensor_stride_kv.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_kv.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + int64_t batch_offset_qo = 0; + int64_t batch_offset_partial_o = 0; +#pragma unroll 1 + for (int i = 0; i < b; ++i) { + batch_strides_dev[i] = batch_offset_qo; + uint16_t* per_batch_q_ptr = + reinterpret_cast(static_cast(q_ptr) + batch_offset_qo); + uint16_t* per_batch_out_ptr = + reinterpret_cast(static_cast(o_ptr) + batch_offset_qo); + // The two below comes from half + float* per_batch_partial_o_ptr = + reinterpret_cast(static_cast(partial_o_ptr) + (batch_offset_partial_o)); + + tma::cudaTmaDesc desc_q; + tma::cudaTmaDesc desc_o; + tma::cudaTmaDesc desc_partial_o; + + tma::cudaSetTmaTileDescriptor(&desc_q, (void*)per_batch_q_ptr, DIMS_QKV, tensor_size_qo.data(), + tensor_stride_qo.data(), tensor_traversal_stride_qkv.data(), + tensor_box_size_qo.data(), tma::cudaTmaDescFormat::BF16_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + &desc_o, (void*)per_batch_out_ptr, DIMS_QKV, tensor_size_qo.data(), tensor_stride_qo.data(), + tensor_traversal_stride_qkv.data(), tensor_box_size_qo.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor(&desc_partial_o, (void*)per_batch_partial_o_ptr, DIMS_QKV, + tensor_size_partial_o.data(), tensor_stride_partial_o.data(), + tensor_traversal_stride_qkv.data(), + tensor_box_size_partial_o.data(), tma::cudaTmaDescFormat::F32_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + reinterpret_cast(tma_desc_q_array)[i] = desc_q; + reinterpret_cast(tma_desc_o_array)[i] = desc_o; + reinterpret_cast(tma_desc_partial_o_array)[i] = desc_partial_o; + + batch_offset_qo += d * h_qo * BYTES_PER_ELEMENT; + batch_offset_partial_o += d * h_qo * sizeof(float); + } +} + +__global__ static void __launch_bounds__(128) + qkv_tma_setup_prefill(const unsigned int b, const unsigned int h_qo, const unsigned int h_kv, + const unsigned int d_qk, const unsigned int d_vo, const bool is_kv_ragged, + const unsigned int page_size, const unsigned int total_num_pages, + + const int64_t k_strides_2, const int64_t k_strides_1, + const int64_t k_strides_0, const int64_t v_strides_2, + const int64_t v_strides_1, const int64_t v_strides_0, + + int32_t* actual_seq_lens_q_data, int32_t* actual_seq_lens_kv_data, + + void* q_ptr, void* k_ptr, void* v_ptr, void* o_ptr, + + tma::cudaTmaDesc* tma_desc_q_array, tma::cudaTmaDesc* tma_desc_k, + tma::cudaTmaDesc* tma_desc_v, tma::cudaTmaDesc* tma_desc_o_array + /* const int64_t *batch_offset_array */) { + const int tid = threadIdx.x; + + constexpr unsigned int DIMS_QKV = 4; + constexpr unsigned int TILE_M_1 = 128; + constexpr unsigned int TILE_N_1 = 128; + constexpr unsigned int BYTES_PER_ELEMENT = 2; + std::array tensor_traversal_stride_qkv = {1, 1, 1, 1}; + + if (is_kv_ragged) { + int64_t batch_offset_k = 0; + int64_t batch_offset_v = 0; + std::array tensor_box_size_kv = {64, TILE_N_1, 1, 1}; + +#pragma unroll 1 + for (int i = 0; i < b; ++i) { + const uint32_t actual_s_kv = static_cast(actual_seq_lens_kv_data[i]); + std::array packed_tensor_size_k = {d_qk, actual_s_kv, h_kv, 1}; + std::array packed_tensor_stride_k = {h_kv * d_qk * BYTES_PER_ELEMENT, + d_qk * BYTES_PER_ELEMENT, 0}; + std::array packed_tensor_size_v = {d_vo, actual_s_kv, h_kv, 1}; + std::array packed_tensor_stride_v = {h_kv * d_vo * BYTES_PER_ELEMENT, + d_vo * BYTES_PER_ELEMENT, 0}; + + uint16_t* k_batch_ptr = + reinterpret_cast(reinterpret_cast(k_ptr) + batch_offset_k); + uint16_t* v_batch_ptr = + reinterpret_cast(reinterpret_cast(v_ptr) + batch_offset_v); + + tma::cudaSetTmaTileDescriptor(&tma_desc_k[i], (void*)k_batch_ptr, DIMS_QKV, + packed_tensor_size_k.data(), packed_tensor_stride_k.data(), + tensor_traversal_stride_qkv.data(), tensor_box_size_kv.data(), + tma::cudaTmaDescFormat::BF16_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor(&tma_desc_v[i], (void*)v_batch_ptr, DIMS_QKV, + packed_tensor_size_v.data(), packed_tensor_stride_v.data(), + tensor_traversal_stride_qkv.data(), tensor_box_size_kv.data(), + tma::cudaTmaDescFormat::BF16_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + batch_offset_k += static_cast(actual_s_kv) * d_qk * h_kv * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + batch_offset_v += static_cast(actual_s_kv) * d_vo * h_kv * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + } + } else { + bool kv_cache_enabled = d_qk == 192 ? false : true; + + std::array tensor_size_k = {d_qk, page_size, h_kv, total_num_pages}; + std::array tensor_stride_k = {k_strides_2 * (BYTES_PER_ELEMENT), + k_strides_1 * (BYTES_PER_ELEMENT), + k_strides_0 * (BYTES_PER_ELEMENT)}; + std::array tensor_size_v = {d_vo, page_size, h_kv, total_num_pages}; + std::array tensor_stride_v = {v_strides_2 * (BYTES_PER_ELEMENT), + v_strides_1 * (BYTES_PER_ELEMENT), + v_strides_0 * (BYTES_PER_ELEMENT)}; + std::array tensor_box_size_k = { + 64, kv_cache_enabled ? std::min(TILE_N_1, page_size) : TILE_N_1, 1, 1}; + + tma::cudaSetTmaTileDescriptor( + reinterpret_cast(tma_desc_k), k_ptr, DIMS_QKV, tensor_size_k.data(), + tensor_stride_k.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_k.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + reinterpret_cast(tma_desc_v), v_ptr, DIMS_QKV, tensor_size_v.data(), + tensor_stride_v.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_k.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + } + + int64_t batch_offset_q = 0; + int64_t batch_offset_k = 0; + int64_t batch_offset_v = 0; + int64_t batch_offset_o = 0; + std::array tensor_box_size_q = {64, TILE_M_1, 1, 1}; + +#pragma unroll 1 + for (int i = 0; i < b; ++i) { + const uint32_t actual_s_q = static_cast(actual_seq_lens_q_data[i]); + + // batch_offset_qo = batch_offset_array ? batch_offset_array[i] : batch_offset_qo; + std::array packed_tensor_size_q = {d_qk, actual_s_q, h_qo, 1}; + std::array packed_tensor_stride_q = {h_qo * d_qk * BYTES_PER_ELEMENT, + d_qk * BYTES_PER_ELEMENT, 0}; + std::array packed_tensor_size_o = {d_vo, actual_s_q, h_qo, 1}; + std::array packed_tensor_stride_o = {h_qo * d_vo * BYTES_PER_ELEMENT, + d_vo * BYTES_PER_ELEMENT, 0}; + + uint16_t* per_batch_q_ptr = + reinterpret_cast(reinterpret_cast(q_ptr) + batch_offset_q); + uint16_t* per_batch_out_ptr = + reinterpret_cast(reinterpret_cast(o_ptr) + batch_offset_o); + + tma::cudaTmaDesc desc_q; + tma::cudaTmaDesc desc_o; + + tma::cudaSetTmaTileDescriptor( + &desc_q, (void*)per_batch_q_ptr, DIMS_QKV, packed_tensor_size_q.data(), + packed_tensor_stride_q.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_q.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + &desc_o, (void*)per_batch_out_ptr, DIMS_QKV, packed_tensor_size_o.data(), + packed_tensor_stride_o.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_q.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + reinterpret_cast(tma_desc_q_array)[i] = desc_q; + reinterpret_cast(tma_desc_o_array)[i] = desc_o; + + batch_offset_q += static_cast(actual_s_q) * d_qk * h_qo * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + batch_offset_o += static_cast(actual_s_q) * d_vo * h_qo * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + } +} + +static void create_packed_tma_desc_kv_prefill(int b, int32_t* actual_seq_lens_kv_data, int64_t d_qk, + int64_t d_vo, int64_t h_kv, + uint32_t* tensor_traversal_stride_qkv, + uint32_t* tensor_box_size_kv, + tma::cudaTmaDesc* packed_tma_desc_k, + tma::cudaTmaDesc* packed_tma_desc_v, TensorView k, + TensorView v) { + int64_t batch_offset_k = 0; + int64_t batch_offset_v = 0; + // tma descriptors for packed q and o + for (int i = 0; i < b; ++i) { + const uint32_t actual_s_kv = static_cast(actual_seq_lens_kv_data[i]); + std::array packed_tensor_size_k = {d_qk, actual_s_kv, h_kv, 1}; + std::array packed_tensor_stride_k = {h_kv * d_qk * BYTES_PER_ELEMENT, + d_qk * BYTES_PER_ELEMENT, 0}; + std::array packed_tensor_size_v = {d_vo, actual_s_kv, h_kv, 1}; + std::array packed_tensor_stride_v = {h_kv * d_vo * BYTES_PER_ELEMENT, + d_vo * BYTES_PER_ELEMENT, 0}; + + uint16_t* k_ptr = reinterpret_cast(k.data_ptr() + batch_offset_k); + uint16_t* v_ptr = reinterpret_cast(v.data_ptr() + batch_offset_v); + + tma::cudaSetTmaTileDescriptor( + &packed_tma_desc_k[i], (void*)k_ptr, DIMS_QKV, packed_tensor_size_k.data(), + packed_tensor_stride_k.data(), tensor_traversal_stride_qkv, tensor_box_size_kv, + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + &packed_tma_desc_v[i], (void*)v_ptr, DIMS_QKV, packed_tensor_size_v.data(), + packed_tensor_stride_v.data(), tensor_traversal_stride_qkv, tensor_box_size_kv, + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + batch_offset_k += static_cast(actual_s_kv) * d_qk * h_kv * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + batch_offset_v += static_cast(actual_s_kv) * d_vo * h_kv * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + } +} + +static void create_packed_tma_desc_qo_prefill(int b, int32_t* actual_seq_lens_q_data, int64_t d_qk, + int64_t d_vo, int64_t h_qo, + uint32_t* tensor_traversal_stride_qkv, + uint32_t* tensor_box_size_q, + tma::cudaTmaDesc* packed_tma_desc_q, + tma::cudaTmaDesc* packed_tma_desc_o, TensorView q, + TensorView out, int64_t* batch_offset_array) { + int64_t batch_offset_q = 0; + int64_t batch_offset_o = 0; + // tma descriptors for packed q and o + for (int i = 0; i < b; ++i) { + const uint32_t actual_s_q = static_cast(actual_seq_lens_q_data[i]); + + batch_offset_q = batch_offset_array ? batch_offset_array[i] : batch_offset_q; + batch_offset_o = batch_offset_array ? batch_offset_array[i] : batch_offset_o; + std::array packed_tensor_size_q = {d_qk, actual_s_q, h_qo, 1}; + std::array packed_tensor_stride_q = {h_qo * d_qk * BYTES_PER_ELEMENT, + d_qk * BYTES_PER_ELEMENT, 0}; + std::array packed_tensor_size_o = {d_vo, actual_s_q, h_qo, 1}; + std::array packed_tensor_stride_o = {h_qo * d_vo * BYTES_PER_ELEMENT, + d_vo * BYTES_PER_ELEMENT, 0}; + + uint16_t* q_ptr = reinterpret_cast(q.data_ptr() + batch_offset_q); + uint16_t* out_ptr = reinterpret_cast(out.data_ptr() + batch_offset_o); + + tma::cudaSetTmaTileDescriptor( + &packed_tma_desc_q[i], (void*)q_ptr, DIMS_QKV, packed_tensor_size_q.data(), + packed_tensor_stride_q.data(), tensor_traversal_stride_qkv, tensor_box_size_q, + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + &packed_tma_desc_o[i], (void*)out_ptr, DIMS_QKV, packed_tensor_size_o.data(), + packed_tensor_stride_o.data(), tensor_traversal_stride_qkv, tensor_box_size_q, + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + batch_offset_q += static_cast(actual_s_q) * d_qk * h_qo * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + batch_offset_o += static_cast(actual_s_q) * d_vo * h_qo * + BYTES_PER_ELEMENT; // Becomes a no-op if batch_offset_array is provided + } +} + +void setup_prefill(CUfunction* prefill_func) { + // Use cu++filt to get the kernel name + std::string kernel_name_deepseek_causal = + "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_" + "128x128x192ILb1ELb0EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE"; + + std::string kernel_name_causal = + "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_" + "128x128x128ILb1ELb1EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE"; + + std::string kernel_name_deepseek = + "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_" + "128x128x192ILb0ELb0EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE"; + + std::string kernel_name = + "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_" + "128x128x128ILb0ELb1EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE"; + + std::string cubin = get_cudnn_cubin(PREFILL); + std::string cubin_deepseek = get_cudnn_cubin(PREFILL_DEEPSEEK); + + if (cubin.empty()) { + throw std::runtime_error("Failed to load cubin for prefill"); + } + if (cubin_deepseek.empty()) { + throw std::runtime_error("Failed to load cubin for prefill_deepseek"); + } + + CUmodule hmod{0}; + CUmodule hmod_deepseek{0}; + if (cuModuleLoadData(&hmod_deepseek, cubin_deepseek.data()) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleLoadData for prefill_deepseek"); + } + + if (cuModuleLoadData(&hmod, cubin.data()) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleLoadData for prefill"); + } + + if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL], hmod, kernel_name.c_str()) != + CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleGetFunction for prefill"); + } + + if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL_DEEPSEEK], hmod_deepseek, + kernel_name_deepseek.c_str()) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleGetFunction for prefill_deepseek"); + } + + if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL_CAUSAL], hmod, kernel_name_causal.c_str()) != + CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleGetFunction for prefill"); + } + + if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL_DEEPSEEK_CAUSAL], hmod_deepseek, + kernel_name_deepseek_causal.c_str()) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleGetFunction for prefill_deepseek"); + } +}; + +void setup_decode(CUfunction* hfunc_decode, CUfunction* lean_attn_reduction) { + constexpr int NUM_DECODE_KERNELS = 5; + + std::string decode_kernel_name[NUM_DECODE_KERNELS] = { + "_Z44cudnn_sm100_fprop_sdpa_decode_bf16_" + "Mx128x128ILb1ELi1EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_ifPfNS0_7stridesES5_" + "S5_PKjS9_S9_jjNS0_11FastDivisorE", + "_Z44cudnn_sm100_fprop_sdpa_decode_bf16_" + "Mx128x128ILb1ELi8EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_ifPfNS0_7stridesES5_" + "S5_PKjS9_S9_jjNS0_11FastDivisorE", + "_Z44cudnn_sm100_fprop_sdpa_decode_bf16_" + "Mx128x128ILb1ELi16EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_ifPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE", + "_Z44cudnn_sm100_fprop_sdpa_decode_bf16_" + "Mx128x128ILb1ELi32EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_ifPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE", + "_Z44cudnn_sm100_fprop_sdpa_decode_bf16_" + "Mx128x128ILb1ELi64EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_ifPfNS0_" + "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE", + }; + + std::string lean_attn_reduction_kernel_name = + "_Z19lean_attn_reductionN4fmha19AttentionDescriptorEiP13__nv_bfloat16PfS3_S3_NS_7stridesES4_" + "S4_S4_Pl"; + + std::string cubin = get_cudnn_cubin(DECODE); + if (cubin.empty()) { + throw std::runtime_error("Failed to load cubin for decode"); + } + + CUmodule hmod{0}; + if (cuModuleLoadData(&hmod, cubin.data()) != CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleLoadData for decode"); + } + + for (int i = 0; i < NUM_DECODE_KERNELS; i++) { + if (cuModuleGetFunction(&hfunc_decode[i], hmod, decode_kernel_name[i].c_str()) != + CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleGetFunction for decode at location " + + std::to_string(i) + " " + decode_kernel_name[i]); + } + } + if (cuModuleGetFunction(lean_attn_reduction, hmod, lean_attn_reduction_kernel_name.c_str()) != + CUDA_SUCCESS) { + throw std::runtime_error("Failed to cuModuleGetFunction for lean_attn_reduction decode"); + } +}; + +void prefill(int64_t b, int64_t s_qo, int64_t max_s_kv, TensorView q, TensorView k_cache, + TensorView v_cache, double scale, TensorView workspace_buffer, + TensorView actual_seq_lens_q, TensorView actual_seq_lens_kv, + TensorView actual_seq_lens_q_gpu, TensorView actual_seq_lens_kv_gpu, + TensorView block_tables, bool causal, bool return_lse, TensorView out, TensorView lse, + Optional batch_offset_q_array, Optional batch_offset_o_array, + Optional batch_offset_k_array, Optional batch_offset_v_array, + bool is_cuda_graph_compatible) { + constexpr size_t SMEM_SIZE = 227 * 1024; // All smem + constexpr int64_t TILE_M_1 = 128; + constexpr int64_t TILE_N_1 = 128; + + constexpr int32_t NUM_THREADS = 512; + + const CUstream stream = get_stream(q.device()); + + int64_t* batch_offset_q_array_data = nullptr; + int64_t* batch_offset_o_array_data = nullptr; + int64_t* batch_offset_k_array_data = nullptr; + int64_t* batch_offset_v_array_data = nullptr; + int64_t* batch_offset_array_data = nullptr; + if (batch_offset_q_array.has_value()) { + batch_offset_array_data = static_cast( + batch_offset_q_array.value().data_ptr()); // Fix this to make it operational later + } + + // Step 1: Setup the kernel pointer + + static CUfunction prefill_func[KERNEL_NUM_PREFILL_TYPES] = {nullptr, nullptr, nullptr, nullptr}; + + int64_t d_qk = q.size(2); + + int64_t d_vo = v_cache.ndim() == 3 ? v_cache.size(2) : v_cache.size(3); + + if (prefill_func[0] == nullptr) { + setup_prefill(prefill_func); + + for (int i = 0; i < KERNEL_NUM_PREFILL_TYPES; i++) { + if (prefill_func[i] != nullptr) { + cuErrCheck(cuFuncSetAttribute(prefill_func[i], + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, SMEM_SIZE)); + cuErrCheck(cuFuncSetAttribute(prefill_func[i], + CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, 100)); + cuErrCheck(cuFuncSetAttribute(prefill_func[i], + CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)); + } + } + } + + // Step 2: Extract attention descriptor + + int64_t h_qo = q.size(1); + + int64_t h_kv = k_cache.size(1); + + int64_t page_size = k_cache.ndim() == 4 ? k_cache.size(2) : 1; + + int64_t s_kv = max_s_kv; + + int64_t num_pages_per_seq = static_cast(std::ceil(1.0 * s_kv / page_size)); + + int64_t total_num_pages = k_cache.ndim() == 4 ? k_cache.size(0) : 1; + + bool kv_cache_enabled = d_qk == 192 ? false : true; + + // Step 3: Setup the launch configuration + + CUlaunchConfig config; + + constexpr int NUM_ATTRS = 1; + CUlaunchAttribute attrs[NUM_ATTRS]; + config.numAttrs = NUM_ATTRS; + attrs[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[0].value.clusterDim.x = 1; + attrs[0].value.clusterDim.y = 1; + attrs[0].value.clusterDim.z = 1; + config.attrs = attrs; + + config.sharedMemBytes = SMEM_SIZE; + config.hStream = stream; + + if (is_cuda_graph_compatible == false) { + CHECK_CPU(actual_seq_lens_q); + CHECK_CPU(actual_seq_lens_kv); + auto actual_seq_lens_q_data = static_cast(actual_seq_lens_q.data_ptr()); + auto actual_seq_lens_kv_data = static_cast(actual_seq_lens_kv.data_ptr()); + + uint32_t actual_num_tiles_per_head = std::transform_reduce( + actual_seq_lens_q_data, actual_seq_lens_q_data + b, 0U, std::plus<>(), [](int32_t seq_len) { + return static_cast(std::ceil(seq_len / (TILE_M_1 * 2.0f))); + }); + config.gridDimX = actual_num_tiles_per_head; + + } else { + config.gridDimX = static_cast(std::ceil(s_qo / (TILE_M_1 * 2.0f))) * b; + } + + config.gridDimY = h_qo; + config.gridDimZ = 1; + + config.blockDimX = NUM_THREADS; + config.blockDimY = 1; + config.blockDimZ = 1; + + // Step 4: Set up the launch arguments + + auto k_strides = k_cache.strides(); + auto v_strides = v_cache.strides(); + + bool is_kv_ragged = k_cache.ndim() == 3; + + std::array tensor_traversal_stride_qkv = {1, 1, 1, 1}; + std::array tensor_size_k = {d_qk, page_size, h_kv, total_num_pages}; + std::array tensor_stride_k = {k_strides[2] * (BYTES_PER_ELEMENT), + k_strides[1] * (BYTES_PER_ELEMENT), + k_strides[0] * (BYTES_PER_ELEMENT)}; + std::array tensor_size_v = {d_vo, page_size, h_kv, total_num_pages}; + std::array tensor_stride_v = {v_strides[2] * (BYTES_PER_ELEMENT), + v_strides[1] * (BYTES_PER_ELEMENT), + v_strides[0] * (BYTES_PER_ELEMENT)}; + + std::array tensor_box_size_q = {64, TILE_M_1, 1, 1}; + std::array tensor_box_size_k = { + 64, kv_cache_enabled ? std::min(TILE_N_1, page_size) : TILE_N_1, 1, 1}; + std::array tensor_box_size_v = { + 64, kv_cache_enabled ? std::min(TILE_N_1, page_size) : TILE_N_1, 1, 1}; + + uint64_t batch_offset_qo = 0; + int8_t* workspace_start = static_cast(workspace_buffer.data_ptr()); + + // These tensors are allocated in the workspace buffer + // Using 2 * b for q and o + std::unique_ptr packed_tma_desc(new tma::cudaTmaDesc[(4 * b)]); + auto packed_tma_desc_q = packed_tma_desc.get(); + auto packed_tma_desc_o = packed_tma_desc.get() + b; + auto tma_desc_k_host = packed_tma_desc.get() + (2 * b); + auto tma_desc_v_host = packed_tma_desc.get() + (3 * b); + + tma::cudaTmaDesc* packed_tma_desc_q_dev = reinterpret_cast(workspace_start); + tma::cudaTmaDesc* packed_tma_desc_o_dev = + reinterpret_cast(workspace_start + sizeof(tma::cudaTmaDesc) * b); + + // These TMA descriptors are allocated in the host and passed by value + tma::cudaTmaDesc* tma_desc_k = + reinterpret_cast(workspace_start + sizeof(tma::cudaTmaDesc) * (2 * b)); + tma::cudaTmaDesc* tma_desc_v = + reinterpret_cast(workspace_start + sizeof(tma::cudaTmaDesc) * (3 * b)); + + if (is_cuda_graph_compatible == false) { + if (is_kv_ragged) { + auto actual_seq_lens_kv_data = static_cast(actual_seq_lens_kv.data_ptr()); + create_packed_tma_desc_kv_prefill( + b, actual_seq_lens_kv_data, d_qk, d_vo, h_kv, tensor_traversal_stride_qkv.data(), + tensor_box_size_k.data(), tma_desc_k_host, tma_desc_v_host, k_cache, v_cache); + } else { + // tma descriptors for k and v + tma::cudaSetTmaTileDescriptor( + tma_desc_k_host, k_cache.data_ptr(), DIMS_QKV, tensor_size_k.data(), + tensor_stride_k.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_k.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor( + tma_desc_v_host, v_cache.data_ptr(), DIMS_QKV, tensor_size_v.data(), + tensor_stride_v.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_v.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + } + auto actual_seq_lens_q_data = static_cast(actual_seq_lens_q.data_ptr()); + create_packed_tma_desc_qo_prefill(b, actual_seq_lens_q_data, d_qk, d_vo, h_qo, + tensor_traversal_stride_qkv.data(), tensor_box_size_q.data(), + packed_tma_desc_q, packed_tma_desc_o, q, out, + batch_offset_array_data); + + cudaMemcpyAsync(workspace_start, packed_tma_desc.get(), sizeof(tma::cudaTmaDesc) * (4 * b), + cudaMemcpyHostToDevice, stream); + } else { + dim3 grid(1, 1, 1); + dim3 block(128, 1, 1); + + cudaStream_t raw_stream = get_stream(q.device()); + + cudaError_t err = cudaStreamQuery(raw_stream); + if (!(err == cudaSuccess || err == cudaErrorNotReady)) { + throw std::runtime_error("CUDA cudnn stream error" + std::string(cudaGetErrorString(err))); + } + + qkv_tma_setup_prefill<<>>( + b, h_qo, h_kv, d_qk, d_vo, is_kv_ragged, page_size, total_num_pages, k_cache.stride(2), + k_cache.stride(1), k_cache.stride(0), v_cache.stride(2), v_cache.stride(1), + v_cache.stride(0), static_cast(actual_seq_lens_q_gpu.data_ptr()), + static_cast(actual_seq_lens_kv_gpu.data_ptr()), q.data_ptr(), k_cache.data_ptr(), + v_cache.data_ptr(), out.data_ptr(), packed_tma_desc_q_dev, tma_desc_k, tma_desc_v, + packed_tma_desc_o_dev); + } + + cudnn_sdpa::AttentionDescriptor_t attn_desc{ + static_cast(b), static_cast(h_qo), static_cast(h_kv), + static_cast(h_kv), static_cast(s_qo), static_cast(s_kv), + static_cast(d_qk), static_cast(h_qo / h_kv), is_kv_ragged}; + + float attn_scale = scale; + + cudnn_sdpa::strides_t lse_strides = {h_qo * s_qo, 1, h_qo, 1}; + + cudnn_sdpa::FastDivisor_t page_size_div; + setFastDivisor(page_size_div, page_size); + + uint32_t page_size32 = static_cast(page_size); + uint32_t num_pages_per_seq32 = static_cast(num_pages_per_seq); + + void* lse_tensor_pointer = return_lse ? lse.data_ptr() : NULL; + + void* actual_seq_lens_q_gpu_pointer = static_cast(actual_seq_lens_q_gpu.data_ptr()); + void* actual_seq_lens_kv_gpu_pointer = static_cast(actual_seq_lens_kv_gpu.data_ptr()); + void* block_tables_pointer = d_qk == 192 ? NULL : static_cast(block_tables.data_ptr()); + + auto print_cudaTmaDescTiled = [](tma::cudaTmaDescTiled* desc) { + printf("addr %p", desc->tensor_common0); + printf(" common1 %x", desc->tensor_common1); + printf(" stride %x", (desc->tensor_stride_lower[0] << 4)); + printf(" stride %x", (desc->tensor_stride_lower[1] << 4)); + printf(" stride %x", (desc->tensor_stride_lower[2] << 4)); + printf(" stride %x", (desc->tensor_stride_lower[3] << 4)); + printf(" stride %x", desc->tensor_stride_upper); + printf(" size0 %x", desc->tensor_size[0]); + printf(" size1 %x", desc->tensor_size[1]); + printf(" size2 %x", desc->tensor_size[2]); + printf(" size3 %x", desc->tensor_size[3]); + printf(" size4 %x", desc->tensor_size[4]); + printf(" stride %x", desc->traversal_stride_box_0); + printf(" box_size_end %d", desc->box_size_end); + printf("\n"); + }; + + // for (int i = 0; i < b; i++) { + // print_cudaTmaDescTiled(reinterpret_cast(&packed_tma_desc_q[i])); + // print_cudaTmaDescTiled(reinterpret_cast(&packed_tma_desc_o[i])); + // } + // print_cudaTmaDescTiled(reinterpret_cast(tma_desc_v_host)); + + void* args[14]; + args[0] = (void*)&attn_desc; + args[1] = (void*)&packed_tma_desc_q_dev; + args[2] = (void*)&tma_desc_k; + args[3] = (void*)&attn_scale; + args[4] = &lse_tensor_pointer; + args[5] = (void*)&lse_strides; + args[6] = (void*)&tma_desc_v; + args[7] = (void*)&packed_tma_desc_o_dev; + args[8] = &actual_seq_lens_q_gpu_pointer; + args[9] = &actual_seq_lens_kv_gpu_pointer; + args[10] = &block_tables_pointer; + args[11] = &page_size32; + args[12] = &num_pages_per_seq32; + args[13] = &page_size_div; + + auto err_launch = CUDA_SUCCESS; + + auto choice = KERNEL_PREFILL; + if (causal) { + choice = d_qk == 192 ? KERNEL_PREFILL_DEEPSEEK_CAUSAL : KERNEL_PREFILL_CAUSAL; + } else { + choice = d_qk == 192 ? KERNEL_PREFILL_DEEPSEEK : KERNEL_PREFILL; + } + + err_launch = cuLaunchKernelEx(&config, prefill_func[choice], (void**)args, nullptr); + + if (err_launch != CUDA_SUCCESS) { + const char* errstr = NULL; + cuGetErrorString(err_launch, &errstr); + throw std::runtime_error("Failed to cuLaunchKernelEx for prefill"); + } +} + +static int32_t compute_split_factor(int32_t b, int32_t h_kv, int32_t h_qo, int32_t s_kv, + uint32_t sm_count) { + uint32_t split_factor = 1; + if ((b * h_kv <= (sm_count / 2))) { + split_factor = std::ceil(1.f * sm_count / (b * h_kv)); + int i = 2; + for (; i < 128; i *= 2) { + if (split_factor <= (i + (i / 2) + (i / 4))) { + split_factor = i; + break; + } + } + if (i == 128) { + split_factor = 64; + } + if ((h_qo / h_kv) <= 8) { + while (std::ceil(1.f * s_kv / split_factor) < (h_qo / h_kv)) { + split_factor /= 2; + } + if (s_kv <= 512) { + split_factor = 1; + } + } else { + if (s_kv <= 1024) { + split_factor = 1; + } + } + if (split_factor == 0) { + split_factor = 1; + } + } + return split_factor; +} + +int32_t get_kernel_id(int32_t q_heads_per_kv) { + auto kernel_id = 0; + if (q_heads_per_kv == 1) { + kernel_id = 0; + } else if (q_heads_per_kv <= 8) { + kernel_id = 1; + } else if (q_heads_per_kv <= 16) { + kernel_id = 2; + } else if (q_heads_per_kv <= 32) { + kernel_id = 3; + } else { + kernel_id = 4; + } + return kernel_id; +} + +void setup_tma_desc_decode(int64_t b, int64_t s_kv, int64_t h_qo, int64_t h_kv, int64_t d, + int64_t total_num_pages, TensorView q, TensorView out, + TensorView k_cache, TensorView v_cache, int32_t split_factor, + int64_t page_size, int8_t* partial_o_dev, tma::cudaTmaDesc* tma_desc_q, + tma::cudaTmaDesc* tma_desc_o, tma::cudaTmaDesc* tma_desc_partial_o, + tma::cudaTmaDesc* tma_desc_k, tma::cudaTmaDesc* tma_desc_v) { + auto kid = get_kernel_id(h_qo / h_kv); + int64_t TILE_M_1 = 1; + int64_t TILE_N_1 = 128; + switch (kid) { + case 0: + TILE_M_1 = 1; + break; + case 1: + TILE_M_1 = 8; + break; + case 2: + TILE_M_1 = 16; + break; + case 3: + TILE_M_1 = 32; + break; + case 4: + TILE_M_1 = 64; + break; + } + + constexpr int64_t DIMS_QKV = 4; + + std::array tensor_traversal_stride_qkv = {1, 1, 1, 1}; + std::array tensor_box_size_qo = {64, 1, 1, 1}; + std::array tensor_box_size_kv = {64, std::min(TILE_N_1, page_size), 1, 1}; + std::array tensor_box_size_partial_o = {32, 1, 1, 1}; + + std::array tensor_size_qo = {d, 1 /* s_qo */, h_qo, b}; + std::array tensor_size_kv = {d, page_size, h_kv, total_num_pages}; + + auto kv_strides = k_cache.strides(); + + std::array tensor_stride_qo = {h_qo * d * BYTES_PER_ELEMENT, + d * BYTES_PER_ELEMENT, 0}; + std::array tensor_stride_kv = {kv_strides[2] * (BYTES_PER_ELEMENT), + kv_strides[1] * (BYTES_PER_ELEMENT), + kv_strides[0] * (BYTES_PER_ELEMENT)}; + + std::array tensor_size_partial_o = {d, split_factor, h_qo, b}; + std::array tensor_stride_partial_o = { + h_qo * d * b * sizeof(float), d * b * sizeof(float), d * h_qo * sizeof(float)}; + uint16_t* q_ptr = reinterpret_cast(q.data_ptr()); + uint16_t* out_ptr = reinterpret_cast(out.data_ptr()); + float* partial_o_ptr = reinterpret_cast(partial_o_dev); + + int64_t batch_offset_qo = 0; + + for (int64_t i = 0; i < b; i++) { + tma::cudaSetTmaTileDescriptor( + &tma_desc_q[i], q_ptr + batch_offset_qo, DIMS_QKV, tensor_size_qo.data(), + tensor_stride_qo.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_qo.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + tma::cudaSetTmaTileDescriptor( + &tma_desc_o[i], out_ptr + batch_offset_qo, DIMS_QKV, tensor_size_qo.data(), + tensor_stride_qo.data(), tensor_traversal_stride_qkv.data(), tensor_box_size_qo.data(), + tma::cudaTmaDescFormat::BF16_RN, tma::cudaTmaDescSwizzle::SWIZZLE_128B); + tma::cudaSetTmaTileDescriptor(&tma_desc_partial_o[i], partial_o_ptr + batch_offset_qo, DIMS_QKV, + tensor_size_partial_o.data(), tensor_stride_partial_o.data(), + tensor_traversal_stride_qkv.data(), + tensor_box_size_partial_o.data(), tma::cudaTmaDescFormat::F32_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); + batch_offset_qo += h_qo * d; + } + + tma::cudaSetTmaTileDescriptor(tma_desc_k, k_cache.data_ptr(), DIMS_QKV, tensor_size_kv.data(), + tensor_stride_kv.data(), tensor_traversal_stride_qkv.data(), + tensor_box_size_kv.data(), tma::cudaTmaDescFormat::BF16_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); + + tma::cudaSetTmaTileDescriptor(tma_desc_v, v_cache.data_ptr(), DIMS_QKV, tensor_size_kv.data(), + tensor_stride_kv.data(), tensor_traversal_stride_qkv.data(), + tensor_box_size_kv.data(), tma::cudaTmaDescFormat::BF16_RN, + tma::cudaTmaDescSwizzle::SWIZZLE_128B); +} + +void decode(int64_t max_s_kv, TensorView q, TensorView k_cache, TensorView v_cache, double scale, + TensorView workspace_buffer, TensorView actual_seq_lens_kv, + TensorView actual_seq_lens_kv_gpu, TensorView block_tables, TensorView out, + Optional batch_offset_q_array, Optional batch_offset_o_array, + bool is_cuda_graph_compatible) { + constexpr size_t SMEM_SIZE = 227 * 1024; // All smem + constexpr size_t REDUCTION_MEM_SIZE = 128 * 1024; + constexpr int64_t TILE_N_1 = 128; + + constexpr int32_t NUM_THREADS = 384; + + int64_t* batch_offset_q_array_data = nullptr; + if (batch_offset_q_array.has_value()) { + batch_offset_q_array_data = static_cast(batch_offset_q_array.value().data_ptr()); + } + + const CUstream stream = get_stream(q.device()); + + constexpr int NUM_DECODE_KERNELS = 5; + static CUfunction hfunc_decode[NUM_DECODE_KERNELS] = {nullptr, nullptr, nullptr, nullptr, + nullptr}; + static CUfunction lean_attn_reduction{nullptr}; + + static int sm_count = 0; + + // Setup decode kernels + if (hfunc_decode[0] == nullptr) { + setup_decode(hfunc_decode, &lean_attn_reduction); + + for (int i = 0; i < NUM_DECODE_KERNELS; i++) { + if (hfunc_decode[i] != nullptr) { + cuErrCheck(cuFuncSetAttribute(hfunc_decode[i], + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, SMEM_SIZE)); + cuErrCheck(cuFuncSetAttribute(hfunc_decode[i], + CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, 100)); + cuErrCheck(cuFuncSetAttribute(hfunc_decode[i], + CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)); + } + } + if (lean_attn_reduction != nullptr) { + cuErrCheck(cuFuncSetAttribute(lean_attn_reduction, + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + REDUCTION_MEM_SIZE)); + cuErrCheck(cuFuncSetAttribute(lean_attn_reduction, + CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, 100)); + cuErrCheck(cuFuncSetAttribute(lean_attn_reduction, + CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)); + } + + // Get number of SMs per GPU + int device_id; + cudaGetDevice(&device_id); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id); + } + + int64_t b = q.size(0); + int64_t h_qo = q.size(1); + int64_t d = q.size(2); + + int64_t h_kv = k_cache.size(1); + + int64_t page_size = k_cache.ndim() == 4 ? k_cache.size(2) : 1; + + int64_t total_num_pages = k_cache.ndim() == 4 ? k_cache.size(0) : 1; + + int64_t s_kv = max_s_kv; + + int64_t s_qo = 1; + + int32_t split_factor = compute_split_factor(b, h_kv, h_qo, s_kv, sm_count); + + split_factor = 1; // Fix split factor. Setting it to 1 for now + + // Set up TMA descriptors for Q, K, V, O + auto qo_strides = q.strides(); + auto kv_strides = v_cache.strides(); + + // Launch config for main kernel + CUlaunchConfig config; + CUlaunchAttribute attrs[1]; + attrs[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[0].value.clusterDim.x = 1; + attrs[0].value.clusterDim.y = 1; + attrs[0].value.clusterDim.z = 1; + + const unsigned int CTAs_y = h_kv * std::ceil(1.f * (h_qo / h_kv) / 64); + + config.gridDimX = split_factor; // Number of CTAs per row + config.gridDimY = CTAs_y; + config.gridDimZ = b; + + config.blockDimX = NUM_THREADS; + config.blockDimY = 1; + config.blockDimZ = 1; + + config.attrs = attrs; + config.sharedMemBytes = SMEM_SIZE; + + config.hStream = stream; + config.numAttrs = 1; + + int8_t* workspace_start = static_cast(workspace_buffer.data_ptr()); + int8_t* partial_o_dev = workspace_start; + int8_t* tma_descriptor_start = + partial_o_dev + (b * s_qo * h_qo * d * sizeof(float) * split_factor); + + int8_t* batch_strides_dev = tma_descriptor_start + ((5 * b) * sizeof(tma::cudaTmaDesc)); + + tma::cudaTmaDesc* packed_tma_desc_q_dev = + reinterpret_cast(tma_descriptor_start); + tma::cudaTmaDesc* packed_tma_desc_o_dev = + reinterpret_cast(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc)); + tma::cudaTmaDesc* packed_tma_desc_partial_o_dev = + reinterpret_cast(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc) * 2); + tma::cudaTmaDesc* tma_desc_k_dev = + reinterpret_cast(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc) * 3); + tma::cudaTmaDesc* tma_desc_v_dev = + reinterpret_cast(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc) * 4); + + int8_t* lse_dev = batch_strides_dev + (b * sizeof(int64_t)); + + if (is_cuda_graph_compatible) { + dim3 grid(1, 1, 1); + dim3 block(128, 1, 1); + auto kid = get_kernel_id(h_qo / h_kv); + int64_t TILE_M_1 = 1; + switch (kid) { + case 0: + TILE_M_1 = 1; + break; + case 1: + TILE_M_1 = 8; + break; + case 2: + TILE_M_1 = 16; + break; + case 3: + TILE_M_1 = 32; + break; + case 4: + TILE_M_1 = 64; + break; + } + + qkv_tma_setup_decode<<>>( + b, h_qo, h_kv, d, total_num_pages, page_size, split_factor, TILE_M_1, TILE_N_1, + kv_strides[2], kv_strides[1], kv_strides[0], q.data_ptr(), k_cache.data_ptr(), + v_cache.data_ptr(), out.data_ptr(), partial_o_dev, packed_tma_desc_q_dev, tma_desc_k_dev, + tma_desc_v_dev, packed_tma_desc_o_dev, packed_tma_desc_partial_o_dev, + reinterpret_cast(batch_strides_dev)); + } else { + std::unique_ptr tma_desc_host(new tma::cudaTmaDesc[5 * b]); + + tma::cudaTmaDesc* tma_desc_q = tma_desc_host.get(); + tma::cudaTmaDesc* tma_desc_o = tma_desc_host.get() + b; + tma::cudaTmaDesc* tma_desc_partial_o = tma_desc_host.get() + b * 2; + tma::cudaTmaDesc* tma_desc_k = tma_desc_host.get() + b * 3; + tma::cudaTmaDesc* tma_desc_v = tma_desc_host.get() + b * 4; + + setup_tma_desc_decode(b, max_s_kv, h_qo, h_kv, d, total_num_pages, q, out, k_cache, v_cache, + split_factor, page_size, partial_o_dev, tma_desc_q, tma_desc_o, + tma_desc_partial_o, tma_desc_k, tma_desc_v); + + std::unique_ptr batch_strides(new int64_t[b]); + for (int i = 0; i < b; i++) { + batch_strides[i] = (i)*d * h_qo; + } + cudaMemcpyAsync(batch_strides_dev, batch_strides.get(), sizeof(int64_t) * b, + cudaMemcpyHostToDevice, stream); + + cudaMemcpyAsync(tma_descriptor_start, tma_desc_host.get(), sizeof(tma::cudaTmaDesc) * (5 * b), + cudaMemcpyHostToDevice, stream); + } + + cudnn_sdpa::AttentionDescriptor_t attnDesc{b, h_qo, h_kv, h_kv, s_qo, + max_s_kv, d, h_qo / h_kv, 0}; + + cudnn_sdpa::FastDivisor_t page_size_div; + setFastDivisor(page_size_div, page_size); + + uint32_t page_size32 = static_cast(page_size); + uint32_t num_pages_per_seq32 = static_cast(max_s_kv / page_size); + + void* args[15]; + + float attn_scale = scale; + void* actual_seq_lens_q_gpu_pointer = nullptr; + void* actual_seq_lens_kv_gpu_pointer = static_cast(actual_seq_lens_kv_gpu.data_ptr()); + void* block_tables_pointer = static_cast(block_tables.data_ptr()); + + cudnn_sdpa::strides_t lse_strides = {h_qo, 1, h_qo, 1}; + cudnn_sdpa::strides_t partial_lse_strides = {h_qo, 1, h_qo * b, 1}; + + cudnn_sdpa::strides_t partial_o_strides = {split_factor * h_qo * d, h_qo * d, d, 1}; + + args[0] = (void*)&attnDesc; + args[1] = (void*)&packed_tma_desc_q_dev; + args[2] = (void*)&tma_desc_k_dev; + args[3] = (void*)&split_factor; + args[4] = (void*)&attn_scale; + args[5] = (void*)&lse_dev; + args[6] = split_factor == 1 ? (void*)&lse_strides : (void*)&partial_lse_strides; + args[7] = (void*)&tma_desc_v_dev; + args[8] = + split_factor == 1 ? (void*)&packed_tma_desc_o_dev : (void*)&packed_tma_desc_partial_o_dev; + args[9] = (void*)&actual_seq_lens_q_gpu_pointer; + args[10] = (void*)&actual_seq_lens_kv_gpu_pointer; + args[11] = (void*)&block_tables_pointer; + args[12] = (void*)&page_size32; + args[13] = (void*)&num_pages_per_seq32; + args[14] = (void*)&page_size_div; + + auto kernel_id = get_kernel_id(attnDesc.q_heads_per_kv); + + auto err_launch = cuLaunchKernelEx(&config, hfunc_decode[kernel_id], (void**)args, nullptr); + if (err_launch != CUDA_SUCCESS) { + std::cerr << "cuLaunchKernelEx failed with error code " << err_launch << std::endl; + throw std::runtime_error("cuLaunchKernelEx failed for decode"); + } + + // Now setting up the reduction kernel + if (split_factor > 1) { + // TODO: Add support for split_factor > 1 + void* args_lean_attn_reduction[11]; + void* o_dev = out.data_ptr(); + + void* lse_final_dev = nullptr; + + cudnn_sdpa::strides_t o_strides = {h_qo * d, d, 1}; + + args_lean_attn_reduction[0] = (void*)&attnDesc; + args_lean_attn_reduction[1] = (void*)&split_factor; + args_lean_attn_reduction[2] = (void*)&o_dev; + args_lean_attn_reduction[3] = (void*)&partial_o_dev; + args_lean_attn_reduction[4] = (void*)&lse_final_dev; + args_lean_attn_reduction[5] = (void*)&lse_dev; + args_lean_attn_reduction[6] = (void*)&o_strides; + args_lean_attn_reduction[7] = (void*)&partial_o_strides; + args_lean_attn_reduction[8] = (void*)&lse_strides; + args_lean_attn_reduction[9] = (void*)&partial_lse_strides; + args_lean_attn_reduction[10] = (void*)&batch_strides_dev; + + // Launch config for reduction kernel + + CUlaunchConfig reduction_config; + + reduction_config.gridDimX = h_qo; + reduction_config.gridDimY = b; // Same as CTAs_z of main kernel + reduction_config.gridDimZ = 1; + + reduction_config.blockDimX = 128; // 128 threads per block + reduction_config.blockDimY = 1; + reduction_config.blockDimZ = 1; + + reduction_config.sharedMemBytes = REDUCTION_MEM_SIZE; + + CUlaunchAttribute reduction_attrs[1]; + reduction_attrs[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + reduction_attrs[0].value.clusterDim.x = 1; + reduction_attrs[0].value.clusterDim.y = 1; + reduction_attrs[0].value.clusterDim.z = 1; + + reduction_config.hStream = stream; + reduction_config.numAttrs = 1; + reduction_config.attrs = reduction_attrs; + + auto err_launch = cuLaunchKernelEx(&reduction_config, lean_attn_reduction, + (void**)args_lean_attn_reduction, nullptr); + if (err_launch != CUDA_SUCCESS) { + std::cerr << "cuLaunchKernelEx failed with error code " << err_launch << std::endl; + throw std::runtime_error("cuLaunchKernelEx failed for decode"); + } + } +} + +} // namespace cudnn_sdpa_kernel_launcher + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(prefill, cudnn_sdpa_kernel_launcher::prefill); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(decode, cudnn_sdpa_kernel_launcher::decode); + +} // namespace flashinfer diff --git a/csrc/cudnn_sdpa_utils.h b/csrc/cudnn_sdpa_utils.h new file mode 100644 index 0000000000..665898e0b5 --- /dev/null +++ b/csrc/cudnn_sdpa_utils.h @@ -0,0 +1,373 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: LicenseRef-NvidiaProprietary + * + * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual + * property and proprietary rights in and to this material, related + * documentation and any modifications thereto. Any use, reproduction, + * disclosure or distribution of this material and related documentation + * without an express license agreement from NVIDIA CORPORATION or + * its affiliates is strictly prohibited. + */ + +#pragma once + +namespace flashinfer { +namespace cudnn_sdpa_kernel_launcher { + +namespace tma { + +typedef enum { TILED = 0, IM2COL } cudaTmaDescType; + +typedef enum { TENSOR_ZFILL = 0, TENSOR_CFILL } cudaTmaDescOobFillMode; + +typedef enum { FP32_TO_TF32_DISABLED = 0, FP32_TO_TF32_ENABLED } cudaTmaDescFp32toTf32Mode; + +typedef enum { + SWIZZLE_DISABLED = 0, + SWIZZLE_32B, + SWIZZLE_64B, + SWIZZLE_128B, + SWIZZLE_MAX +} cudaTmaDescSwizzle; + +typedef enum { + INTERLEAVE_DISABLED, + INTERLEAVE_16B, + INTERLEAVE_32B, + INTERLEAVE_MAX +} cudaTmaDescInterleave; + +typedef enum { + PROMOTION_DISABLED = 0, + PROMOTION_64B, + PROMOTION_128B, + PROMOTION_256B +} cudaTmaDescPromotion; + +typedef enum { + U8 = 0, + U16, + U32, + S32, + U64, + S64, + F16_RN, + F32_RN, + F32_FTZ_RN, + F64_RN, + BF16_RN, + FORMAT_MAX +} cudaTmaDescFormat; + +typedef struct { + uint64_t tensor_common0; + uint32_t tensor_common1; + + uint32_t tensor_stride_lower[4]; //< 36b of 64b with 4B aligned + uint32_t tensor_stride_upper; + uint32_t tensor_size[5]; //< value -1 + uint32_t traversal_stride_box_0; //< packed 3b (-1) + + uint32_t box_size_end; +} cudaTmaDescTiled; + +typedef struct { + uint64_t tensor_common0; + uint32_t tensor_common1; + + uint32_t tensor_stride_lower[4]; + uint32_t tensor_stride_upper; + uint32_t tensor_size[5]; + uint32_t traversal_stride_range_c; + + uint32_t box_corner_dhw; + uint32_t range_ndhw; +} cudaTmaDescIm2Col; + +typedef struct alignas(64) { + uint64_t data[8]; +} cudaTmaDesc; + +static inline __host__ __device__ void set_tensor_common_0(cudaTmaDesc* p_desc, uint64_t addr) { + cudaTmaDescTiled* desc = reinterpret_cast(p_desc); + desc->tensor_common0 = 0; + desc->tensor_common0 |= addr; +} + +static inline __host__ __device__ void set_tensor_common_1( + cudaTmaDesc* p_desc, cudaTmaDescType desc_type, uint32_t dims, cudaTmaDescFormat format, + cudaTmaDescInterleave interleave, cudaTmaDescSwizzle swizzle, cudaTmaDescOobFillMode fill, + cudaTmaDescFp32toTf32Mode f32_to_tf32, cudaTmaDescPromotion promotion) { + cudaTmaDescTiled* desc = reinterpret_cast(p_desc); + + desc->tensor_common1 = 0; + desc->tensor_common1 |= desc_type == TILED ? 0x0 : 0x1; + + constexpr uint32_t VERSION_SHIFT = 1; + constexpr uint32_t VERSION_BITS = 3; + desc->tensor_common1 |= (1u << VERSION_SHIFT); + + constexpr uint32_t DIM_BITS = 3; + constexpr uint32_t DIM_SHIFT = VERSION_SHIFT + VERSION_BITS; + constexpr uint32_t DIM_MASK = (1u << DIM_BITS) - 1; + desc->tensor_common1 |= ((dims - 1) & DIM_MASK) << DIM_SHIFT; + + constexpr uint32_t FORMAT_BITS = 4; + constexpr uint32_t FORMAT_SHIFT = DIM_SHIFT + DIM_BITS; + constexpr uint32_t FORMAT_MASK = (1u << FORMAT_BITS) - 1; + desc->tensor_common1 |= (static_cast(format) & FORMAT_MASK) << FORMAT_SHIFT; + + constexpr uint32_t INTERLEAVE_BITS = 2; + constexpr uint32_t INTERLEAVE_SHIFT = FORMAT_SHIFT + FORMAT_BITS; + constexpr uint32_t INTERLEAVE_MASK = (1u << INTERLEAVE_BITS) - 1; + desc->tensor_common1 |= (static_cast(interleave) & INTERLEAVE_MASK) << INTERLEAVE_SHIFT; + + constexpr uint32_t SWIZZLE_BITS = 2; + constexpr uint32_t SWIZZLE_SHIFT = INTERLEAVE_SHIFT + INTERLEAVE_BITS; + constexpr uint32_t SWIZZLE_MASK = (1u << SWIZZLE_BITS) - 1; + desc->tensor_common1 |= (static_cast(swizzle) & SWIZZLE_MASK) << SWIZZLE_SHIFT; + + constexpr uint32_t FILL_BITS = 1; + constexpr uint32_t FILL_SHIFT = SWIZZLE_SHIFT + SWIZZLE_BITS; + constexpr uint32_t FILL_MASK = (1u << FILL_BITS) - 1; + desc->tensor_common1 |= (static_cast(fill) & FILL_MASK) << FILL_SHIFT; + + constexpr uint32_t F32_TO_TF32_BITS = 1; + constexpr uint32_t F32_TO_TF32_SHIFT = FILL_SHIFT + FILL_BITS; + constexpr uint32_t F32_TO_TF32_MASK = (1u << F32_TO_TF32_BITS) - 1; + desc->tensor_common1 |= (static_cast(f32_to_tf32) & F32_TO_TF32_MASK) + << F32_TO_TF32_SHIFT; + + constexpr uint32_t PROMOTION_BITS = 2; + constexpr uint32_t PROMOTION_SHIFT = F32_TO_TF32_SHIFT + F32_TO_TF32_BITS; + constexpr uint32_t PROMOTION_MASK = (1u << PROMOTION_BITS) - 1; + desc->tensor_common1 |= (static_cast(promotion) & PROMOTION_MASK) << PROMOTION_SHIFT; +} + +static inline __host__ __device__ void set_tensor_stride(cudaTmaDesc* p_desc, + uint64_t* p_tensor_stride, uint32_t dims) { + cudaTmaDescTiled* desc = reinterpret_cast(p_desc); + + constexpr uint32_t TENSOR_STRIDE_UPPER_BITS = 4; + constexpr uint32_t TENSOR_STRIDE_UPPER_MASK = (1u << TENSOR_STRIDE_UPPER_BITS) - 1; + + for (uint32_t i = 0; i < dims - 1; i++) { + desc->tensor_stride_lower[i] = 0u; + uint64_t tensor_stride_lower_64b = (p_tensor_stride[i] >> 4) & 0xFFFFFFFFlu; + desc->tensor_stride_lower[i] = static_cast(tensor_stride_lower_64b); + } + desc->tensor_stride_upper = 0u; + + for (uint32_t i = 0; i < dims - 1; i++) { + uint64_t tensor_stride = p_tensor_stride[i]; + tensor_stride = tensor_stride >> 4; + uint64_t tensor_stride_upper = tensor_stride >> 32; + uint32_t tensor_stride_upper_32b = static_cast(tensor_stride_upper); + desc->tensor_stride_upper |= + ((tensor_stride_upper_32b & TENSOR_STRIDE_UPPER_MASK) << (i * TENSOR_STRIDE_UPPER_BITS)); + } +} + +static inline __host__ __device__ void set_tensor_size(cudaTmaDesc* p_desc, uint32_t* p_tensor_size, + uint32_t dims) { + cudaTmaDescTiled* desc = reinterpret_cast(p_desc); + for (uint32_t dim = 0; dim < dims; dim++) { + desc->tensor_size[dim] = p_tensor_size[dim] - 1; + } +} + +static inline __host__ __device__ void set_traversal_stride_tiled(cudaTmaDesc* p_desc, + uint32_t* p_traversal_stride, + uint32_t dims) { + cudaTmaDescTiled* desc = reinterpret_cast(p_desc); + + desc->traversal_stride_box_0 = 0; + + constexpr uint32_t TRAVERSAL_STRIDE_BITS = 3; + constexpr uint32_t TRAVERSAL_STRIDE_MASK = (1u << TRAVERSAL_STRIDE_BITS) - 1; + + for (uint32_t dim = 0; dim < dims; dim++) { + uint32_t traversal_stride = p_traversal_stride[dim] - 1; + traversal_stride = (traversal_stride & TRAVERSAL_STRIDE_MASK) << (dim * TRAVERSAL_STRIDE_BITS); + desc->traversal_stride_box_0 |= traversal_stride; + } +} + +static inline __host__ __device__ void set_box_size(cudaTmaDesc* p_desc, uint32_t* p_box_size, + uint32_t dims) { + cudaTmaDescTiled* desc = reinterpret_cast(p_desc); + + desc->box_size_end = 0; + + constexpr uint32_t BOX_SIZE_BITS = 8; + constexpr uint32_t BOX_SIZE_MASK = (1 << BOX_SIZE_BITS) - 1; + + if (dims > 1) { + uint32_t box_size_0 = p_box_size[0] - 1; + box_size_0 = box_size_0 & BOX_SIZE_MASK; + box_size_0 = box_size_0 << 24; + desc->traversal_stride_box_0 |= box_size_0; + } + + for (uint32_t dim = 1; dim < dims; dim++) { + uint32_t box_size = p_box_size[dim] - 1; + box_size = box_size & BOX_SIZE_MASK; + box_size = box_size << ((dim - 1) * BOX_SIZE_BITS); + desc->box_size_end |= box_size; + } +} + +static inline __host__ __device__ void set_traversal_stride_im2col(cudaTmaDesc* p_desc, + uint32_t* p_traversal_stride, + uint32_t dims) { + cudaTmaDescIm2Col* desc = reinterpret_cast(p_desc); + + constexpr uint32_t TRAVERSAL_STRIDE_BITS = 3; + constexpr uint32_t TRAVERSAL_STRIDE_MASK = (1u << (TRAVERSAL_STRIDE_BITS + 1)) - 1; + + desc->traversal_stride_range_c = 0; + for (uint32_t dim = 0; dim < dims; dim++) { + uint32_t traversal_stride = p_traversal_stride[dim] - 1; + traversal_stride = (traversal_stride & TRAVERSAL_STRIDE_MASK) << (dim * TRAVERSAL_STRIDE_BITS); + desc->traversal_stride_range_c |= traversal_stride; + } +} + +static inline __host__ __device__ void set_range_c(cudaTmaDesc* p_desc, uint32_t range_c) { + cudaTmaDescIm2Col* desc = reinterpret_cast(p_desc); + + constexpr uint32_t RANGE_C_BITS = 8; + constexpr uint32_t RANGE_C_MASK = (1u << RANGE_C_BITS) - 1; + + range_c = range_c & RANGE_C_MASK; + desc->traversal_stride_range_c |= ((range_c - 1) << 24); +} + +static inline __host__ __device__ void set_box_corner_dhw(cudaTmaDesc* p_desc, + uint32_t* p_base_corner, + uint32_t* p_far_corner, uint32_t dims) { + cudaTmaDescIm2Col* desc = reinterpret_cast(p_desc); + + desc->box_corner_dhw = 0; + + uint32_t box_base_corner = 0, box_far_corner = 0; + uint32_t box_corner_dhw = 0; + + if (dims == 3) { + constexpr uint32_t BOX_CORNER_BITS = 16; + constexpr uint32_t BOX_CORNER_MASK = (1u << BOX_CORNER_BITS) - 1; + + box_base_corner = p_base_corner[0] & BOX_CORNER_MASK; + box_far_corner = p_far_corner[0] & BOX_CORNER_MASK; + } + + if (dims == 4) { + constexpr uint32_t BOX_CORNER_BITS = 8; + constexpr uint32_t BOX_CORNER_MASK = (1u << BOX_CORNER_BITS) - 1; + + box_base_corner = p_base_corner[0] & BOX_CORNER_MASK; + box_base_corner |= ((p_base_corner[1] & BOX_CORNER_MASK) << BOX_CORNER_BITS); + + box_far_corner = p_far_corner[0] & BOX_CORNER_MASK; + box_far_corner |= ((p_far_corner[1] & BOX_CORNER_MASK) << BOX_CORNER_BITS); + } + + if (dims == 5) { + constexpr uint32_t BOX_CORNER_BITS = 5; + constexpr uint32_t BOX_CORNER_MASK = (1u << BOX_CORNER_BITS) - 1; + + box_base_corner = p_base_corner[0] & BOX_CORNER_MASK; + box_base_corner |= ((p_base_corner[1] & BOX_CORNER_MASK) << BOX_CORNER_BITS); + box_base_corner |= ((p_base_corner[2] & BOX_CORNER_MASK) << (2 * BOX_CORNER_BITS)); + + box_far_corner = p_far_corner[0] & BOX_CORNER_MASK; + box_far_corner |= ((p_far_corner[1] & BOX_CORNER_MASK) << BOX_CORNER_BITS); + box_far_corner |= ((p_far_corner[2] & BOX_CORNER_MASK) << (2 * BOX_CORNER_BITS)); + } + + box_corner_dhw = box_base_corner; + box_corner_dhw |= (box_far_corner << 16); + + desc->box_corner_dhw = box_corner_dhw; +} + +static inline __host__ __device__ void set_range_ndhw(cudaTmaDesc* p_desc, uint32_t ndhw) { + cudaTmaDescIm2Col* desc = reinterpret_cast(p_desc); + + constexpr uint32_t RANGE_NDHW_BITS = 10; + constexpr uint32_t RANGE_NDHW_MASK = (1u << RANGE_NDHW_BITS) - 1; + + desc->range_ndhw = 0; + desc->range_ndhw = ((ndhw - 1) & RANGE_NDHW_MASK); +} + +static inline __host__ __device__ void cudaSetTmaTileDescriptor( + cudaTmaDesc* p_desc, const void* p_addr, uint32_t dims, uint32_t* p_tensor_size, + uint64_t* p_tensor_stride, uint32_t* p_traversal_stride, uint32_t* p_box_size, + cudaTmaDescFormat format, cudaTmaDescSwizzle swizzle, + cudaTmaDescPromotion promotion = cudaTmaDescPromotion::PROMOTION_DISABLED, + cudaTmaDescInterleave interleave = cudaTmaDescInterleave::INTERLEAVE_DISABLED, + cudaTmaDescOobFillMode fill_oob = cudaTmaDescOobFillMode::TENSOR_ZFILL, + cudaTmaDescFp32toTf32Mode round_to_tf32 = cudaTmaDescFp32toTf32Mode::FP32_TO_TF32_DISABLED) { + set_tensor_common_0(p_desc, reinterpret_cast(p_addr)); + set_tensor_common_1(p_desc, TILED, dims, format, interleave, swizzle, fill_oob, round_to_tf32, + promotion); + set_tensor_stride(p_desc, p_tensor_stride, dims); + set_tensor_size(p_desc, p_tensor_size, dims); + set_traversal_stride_tiled(p_desc, p_traversal_stride, dims); + set_box_size(p_desc, p_box_size, dims); +} + +static inline __host__ __device__ void cudaSetTmaIm2ColDescriptor( + cudaTmaDesc* p_desc, const void* p_addr, uint32_t dims, uint32_t* p_tensor_size, + uint64_t* p_tensor_stride, uint32_t* p_traversal_stride, uint32_t range_c, uint32_t range_ndhw, + uint32_t* p_box_base_corner_dhw, uint32_t* p_box_far_corner_dhw, cudaTmaDescFormat format, + cudaTmaDescSwizzle swizzle, + cudaTmaDescPromotion promotion = cudaTmaDescPromotion::PROMOTION_DISABLED, + cudaTmaDescInterleave interleave = cudaTmaDescInterleave::INTERLEAVE_DISABLED, + cudaTmaDescOobFillMode fill_oob = cudaTmaDescOobFillMode::TENSOR_ZFILL, + cudaTmaDescFp32toTf32Mode round_to_tf32 = cudaTmaDescFp32toTf32Mode::FP32_TO_TF32_DISABLED) { + set_tensor_common_0(p_desc, reinterpret_cast(p_addr)); + set_tensor_common_1(p_desc, IM2COL, dims, format, interleave, swizzle, fill_oob, round_to_tf32, + promotion); + + set_tensor_stride(p_desc, p_tensor_stride, dims); + set_tensor_size(p_desc, p_tensor_size, dims); + set_traversal_stride_im2col(p_desc, p_traversal_stride, dims); + set_range_c(p_desc, range_c); + set_box_corner_dhw(p_desc, p_box_base_corner_dhw, p_box_far_corner_dhw, dims); + set_range_ndhw(p_desc, range_ndhw); +} + +} // namespace tma + +namespace cudnn_sdpa { + +typedef struct AttentionDescriptor { + // Input parameters + // b - batch + // q_h - num heads of q/dq/o/do + // k_h - num heads of k/dk + // v_h - num heads of v/dv + // s_q - max sequence length of q + // s_kv - max sequence length of kv + // d - hidden dim (head dim) + // is_kv_ragged - whether kv is ragged + uint32_t b, q_h, k_h, v_h, s_q, s_kv, d, q_heads_per_kv, is_kv_ragged; +} AttentionDescriptor_t; + +typedef struct FastDivisor { + uint32_t val, shr, mul; +} FastDivisor_t; + +typedef struct strides { + int64_t strides[4]; + + int __attribute__((host)) __attribute__((device)) operator[](int i) const { return strides[i]; } +} strides_t; + +} // namespace cudnn_sdpa + +} // namespace cudnn_sdpa_kernel_launcher +} // namespace flashinfer diff --git a/csrc/cutlass_mla.cu b/csrc/cutlass_mla.cu new file mode 100644 index 0000000000..f68df30bea --- /dev/null +++ b/csrc/cutlass_mla.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "tvm_ffi_utils.h" + +using namespace flashinfer; +using namespace flashinfer::attention; + +void CutlassMLAPagedAttention(ffi::TensorView workspace, ffi::TensorView out, ffi::TensorView lse, + ffi::TensorView q_nope_pe, ffi::TensorView ckv_kpe_cache, + ffi::TensorView kv_lens, ffi::TensorView page_table) { + cudaSetDevice(q_nope_pe.device().device_id); + const cudaStream_t stream = get_stream(q_nope_pe.device()); + + int device_index = q_nope_pe.device().device_id; + int batches = q_nope_pe.size(0); + int page_count_per_seq = page_table.size(1); + int page_count_total = ckv_kpe_cache.size(0); + int page_size = ckv_kpe_cache.size(1); + + DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(q_nope_pe.dtype(), c_type, [&] { + using cutlass_t = cutlass_dtype_t; + auto status = runMla( + workspace.data_ptr(), out.data_ptr(), lse.data_ptr(), q_nope_pe.data_ptr(), + ckv_kpe_cache.data_ptr(), kv_lens.data_ptr(), page_table.data_ptr(), batches, + page_count_per_seq, page_count_total, page_size, device_index, stream); + + TVM_FFI_ICHECK(status == cudaSuccess) + << "Failed to run CutlassMLAPagedAttention: " << cudaGetErrorString(status); + return true; + }); +} diff --git a/csrc/dsv3_router_gemm.cu b/csrc/dsv3_router_gemm.cu new file mode 100644 index 0000000000..2d44147d97 --- /dev/null +++ b/csrc/dsv3_router_gemm.cu @@ -0,0 +1,152 @@ +#include "flashinfer/gemm/dsv3_router_gemm.cuh" +#include "tvm_ffi_utils.h" + +namespace flashinfer::trtllm_dsv3_router_gemm { +template +void invokeRouterGemm(float* output, T const* mat_a, T const* mat_b, cudaStream_t stream, + bool use_pdl = false) { + constexpr int VPT = 16 / sizeof(T); + constexpr int kBlockSize = 128; + cudaLaunchConfig_t config; + config.gridDim = kNumExperts; + config.blockDim = kBlockSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = use_pdl; + config.numAttrs = 1; + config.attrs = attrs; + auto status = cudaLaunchKernelEx( + &config, router_gemm_kernel, output, + mat_a, mat_b); + TVM_FFI_ICHECK(status == cudaSuccess) + << "cudaLaunchKernelEx failed with error code " << cudaGetErrorString(status); +} + +template void invokeRouterGemm<__nv_bfloat16, 1, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 2, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 3, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 4, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 5, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 6, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 7, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 8, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 9, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 10, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 11, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 12, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 13, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 14, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 15, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template void invokeRouterGemm<__nv_bfloat16, 16, 256, 7168>(float*, __nv_bfloat16 const*, + __nv_bfloat16 const*, cudaStream_t, + bool); + +template +struct LoopUnroller { + static void unroll(int num_tokens, float* output, __nv_bfloat16 const* input, + __nv_bfloat16 const* weights, cudaStream_t stream, bool launch_with_pdl) { + if (num_tokens == kBegin) { + invokeRouterGemm<__nv_bfloat16, kBegin, kNumExperts, kHiddenDim>(output, input, weights, + stream, launch_with_pdl); + } else { + LoopUnroller::unroll( + num_tokens, output, input, weights, stream, launch_with_pdl); + } + } +}; + +template +struct LoopUnroller { + static void unroll(int num_tokens, float* output, __nv_bfloat16 const* input, + __nv_bfloat16 const* weights, cudaStream_t stream, bool launch_with_pdl) { + if (num_tokens == kEnd) { + invokeRouterGemm<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(output, input, weights, stream, + launch_with_pdl); + } else { + throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); + } + } +}; + +void dsv3_router_gemm_op(TensorView mat_a, TensorView mat_b, TensorView out, bool launch_with_pdl) { + int const num_tokens = mat_a.sizes()[0]; + int const num_experts = mat_b.sizes()[1]; + int const hidden_dim = mat_a.sizes()[1]; + auto const out_dtype_ = out.dtype(); + auto const data_type = mat_a.dtype(); + constexpr int kNumExperts = 256; + constexpr int kHiddenDim = 7168; + std::vector output_size = {mat_a.sizes()[0], mat_b.sizes()[1]}; + TVM_FFI_ICHECK(mat_a.dim() == 2 && mat_b.dim() == 2) << "mat_a and mat_b must be 2D tensors"; + TVM_FFI_ICHECK(mat_a.strides()[1] == 1 && out.strides()[1] == 1) + << "mat_a and out must be row-major"; + TVM_FFI_ICHECK(mat_b.strides()[0] == 1) << "mat_b must be column-major"; + auto stream = get_stream(mat_a.device()); + bool use_custom_kernel = false; + if (num_tokens >= 1 && num_tokens <= 16 && num_experts == kNumExperts && + hidden_dim == kHiddenDim && encode_dlpack_dtype(data_type) == bfloat16_code && + encode_dlpack_dtype(out_dtype_) == float32_code) { + use_custom_kernel = true; + } + + if (use_custom_kernel) { + LoopUnroller<1, 16, kNumExperts, kHiddenDim>::unroll( + num_tokens, reinterpret_cast(out.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream, launch_with_pdl); + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "Unsupported input tensor size"; + } +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(dsv3_router_gemm_op, + flashinfer::trtllm_dsv3_router_gemm::dsv3_router_gemm_op); + +} // namespace flashinfer::trtllm_dsv3_router_gemm diff --git a/csrc/flashinfer_cascade_binding.cu b/csrc/flashinfer_cascade_binding.cu new file mode 100644 index 0000000000..0fe7e7c06c --- /dev/null +++ b/csrc/flashinfer_cascade_binding.cu @@ -0,0 +1,34 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tvm_ffi_utils.h" + +using tvm::ffi::Optional; + +void merge_state(TensorView v_a, TensorView s_a, TensorView v_b, TensorView s_b, + TensorView v_merged, TensorView s_merged); + +void merge_state_in_place(TensorView v, TensorView s, TensorView v_other, TensorView s_other, + Optional mask); + +void merge_states(TensorView v, TensorView s, TensorView v_merged, TensorView s_merged); + +// Merge two self-attention states +TVM_FFI_DLL_EXPORT_TYPED_FUNC(merge_state, merge_state); +// Merge another self-attention state in-place. +TVM_FFI_DLL_EXPORT_TYPED_FUNC(merge_state_in_place, merge_state_in_place); +// "Merge multiple self-attention states" +TVM_FFI_DLL_EXPORT_TYPED_FUNC(merge_states, merge_states); diff --git a/csrc/flashinfer_cascade_ops.cu b/csrc/flashinfer_cascade_ops.cu deleted file mode 100644 index de14192f6d..0000000000 --- a/csrc/flashinfer_cascade_ops.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "pytorch_extension_utils.h" - -void merge_state(at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, - at::Tensor v_merged, at::Tensor s_merged); - -void merge_state_in_place(at::Tensor v, at::Tensor s, at::Tensor v_other, at::Tensor s_other, - std::optional mask); - -void merge_states(at::Tensor v, at::Tensor s, at::Tensor v_merged, at::Tensor s_merged); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // Merge two self-attention states - m.def("merge_state", merge_state); - // Merge another self-attention state in-place. - m.def("merge_state_in_place", merge_state_in_place); - // "Merge multiple self-attention states" - m.def("merge_states", merge_states); -} diff --git a/csrc/flashinfer_comm_ops.cu b/csrc/flashinfer_comm_ops.cu deleted file mode 100644 index 38864688cd..0000000000 --- a/csrc/flashinfer_comm_ops.cu +++ /dev/null @@ -1,37 +0,0 @@ -// flashinfer: adapted from sglang + vllm code -// refer to: https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/common_extension.cc -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "flashinfer/distributed/comm_ops.h" -#include "pytorch_extension_utils.h" - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - m.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta); - m.def("register_graph_buffers", ®ister_graph_buffers); - m.def("dispose", &dispose); - m.def("meta_size", &meta_size); - m.def("register_buffer", ®ister_buffer); - - m.def( - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " - "int rank, bool full_nvlink) -> int"); - m.impl("init_custom_ar", torch::kCUDA, &init_custom_ar); - - m.def( - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " - "int reg_buffer_sz_bytes, int num_ctas) -> ()"); - m.impl("all_reduce", torch::kCUDA, &all_reduce); -} diff --git a/csrc/flashinfer_gemm_ops.cu b/csrc/flashinfer_gemm_binding.cu similarity index 50% rename from csrc/flashinfer_gemm_ops.cu rename to csrc/flashinfer_gemm_binding.cu index 3ca23f40cb..52d0551413 100644 --- a/csrc/flashinfer_gemm_ops.cu +++ b/csrc/flashinfer_gemm_binding.cu @@ -13,18 +13,15 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "pytorch_extension_utils.h" -void bmm_fp8(at::Tensor A, at::Tensor B, at::Tensor D, at::Tensor A_scale, at::Tensor B_scale, - at::Tensor workspace_buffer, int64_t cublas_handle); +#include "tvm_ffi_utils.h" -void CutlassSegmentGEMM(at::Tensor workspace_buffer, at::Tensor all_problems, at::Tensor x_ptr, - at::Tensor w_ptr, at::Tensor y_ptr, at::Tensor x_ld, at::Tensor w_ld, - at::Tensor y_ld, at::Tensor empty_x_data, bool weight_column_major); +void bmm_fp8(TensorView A, TensorView B, TensorView D, TensorView A_scale, TensorView B_scale, + TensorView workspace_buffer, int64_t cublas_handle); -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // "Cutlass Segment GEMM" - m.def("cutlass_segment_gemm", CutlassSegmentGEMM); - // "BMM FP8" - m.def("bmm_fp8", bmm_fp8); -} +void CutlassSegmentGEMM(TensorView workspace_buffer, TensorView all_problems, TensorView x_ptr, + TensorView w_ptr, TensorView y_ptr, TensorView x_ld, TensorView w_ld, + TensorView y_ld, TensorView empty_x_data, bool weight_column_major); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cutlass_segment_gemm, CutlassSegmentGEMM); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(bmm_fp8, bmm_fp8); diff --git a/csrc/flashinfer_gemm_sm90_ops.cu b/csrc/flashinfer_gemm_sm90_binding.cu similarity index 53% rename from csrc/flashinfer_gemm_sm90_ops.cu rename to csrc/flashinfer_gemm_sm90_binding.cu index 38a9eb9eb7..d5d86abd69 100644 --- a/csrc/flashinfer_gemm_sm90_ops.cu +++ b/csrc/flashinfer_gemm_sm90_binding.cu @@ -13,15 +13,13 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "pytorch_extension_utils.h" +#include "tvm_ffi_utils.h" -void CutlassSegmentGEMMSM90(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor all_problems, at::Tensor x_ptr, at::Tensor w_ptr, - at::Tensor y_ptr, at::Tensor x_stride, at::Tensor weight_stride, - at::Tensor y_stride, at::Tensor empty_x_data, at::Tensor empty_y_data, +void CutlassSegmentGEMMSM90(TensorView float_workspace_buffer, TensorView int_workspace_buffer, + TensorView all_problems, TensorView x_ptr, TensorView w_ptr, + TensorView y_ptr, TensorView x_stride, TensorView weight_stride, + TensorView y_stride, TensorView empty_x_data, TensorView empty_y_data, bool weight_column_major); -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // "Cutlass Segment GEMM operator for SM90" - m.def("cutlass_segment_gemm_sm90", CutlassSegmentGEMMSM90); -} +// "Cutlass Segment GEMM operator for SM90" +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cutlass_segment_gemm_sm90, CutlassSegmentGEMMSM90); diff --git a/csrc/flashinfer_mla_binding.cu b/csrc/flashinfer_mla_binding.cu new file mode 100644 index 0000000000..c8fc1a61e3 --- /dev/null +++ b/csrc/flashinfer_mla_binding.cu @@ -0,0 +1,22 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +void CutlassMLAPagedAttention(TensorView workspace, TensorView out, TensorView lse, + TensorView q_nope_pe, TensorView ckv_kpe_cache, TensorView kv_lens, + TensorView page_table); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cutlass_mla_paged_attention, CutlassMLAPagedAttention); diff --git a/csrc/flashinfer_norm_binding.cu b/csrc/flashinfer_norm_binding.cu new file mode 100644 index 0000000000..09fd2b585a --- /dev/null +++ b/csrc/flashinfer_norm_binding.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +void rmsnorm(TensorView out, TensorView input, TensorView weight, double eps, bool enable_pdl); + +void fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight, double eps, + bool enable_pdl); + +void gemma_rmsnorm(TensorView out, TensorView input, TensorView weight, double eps, + bool enable_pdl); + +void gemma_fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight, double eps, + bool enable_pdl); + +void layernorm(Tensor out, Tensor input, Tensor gamma, Tensor beta, double eps); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm, rmsnorm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm, fused_add_rmsnorm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_rmsnorm, gemma_rmsnorm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(layernorm, layernorm); diff --git a/csrc/flashinfer_norm_ops.cu b/csrc/flashinfer_norm_ops.cu deleted file mode 100644 index 100e4657a2..0000000000 --- a/csrc/flashinfer_norm_ops.cu +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "pytorch_extension_utils.h" - -void rmsnorm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double eps, bool enable_pdl); - -void fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps, - bool enable_pdl); - -void gemma_rmsnorm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double eps, - bool enable_pdl); - -void gemma_fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, - double eps, bool enable_pdl); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // Root mean square normalization - m.def("rmsnorm", rmsnorm); - // Fused add root mean square normalization - m.def("fused_add_rmsnorm", fused_add_rmsnorm); - // Gemma Root mean square normalization - m.def("gemma_rmsnorm", gemma_rmsnorm); - // Gemma Fused add root mean square normalization - m.def("gemma_fused_add_rmsnorm", gemma_fused_add_rmsnorm); -} diff --git a/csrc/flashinfer_ops.cu b/csrc/flashinfer_ops.cu deleted file mode 100644 index 2fe10c6e70..0000000000 --- a/csrc/flashinfer_ops.cu +++ /dev/null @@ -1,313 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "aot_default_additional_params.h" -#include "pytorch_extension_utils.h" - -//========== activation ========== - -void silu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl); -void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl); -void gelu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl); - -//========== cascade ========== - -void merge_state(at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, - at::Tensor v_merged, at::Tensor s_merged); - -void merge_state_in_place(at::Tensor v, at::Tensor s, at::Tensor v_other, at::Tensor s_other, - std::optional mask); - -void merge_states(at::Tensor v, at::Tensor s, at::Tensor v_merged, at::Tensor s_merged); - -//========== decode ========== - -void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, - at::Tensor o, std::optional maybe_lse, int64_t layout, - int64_t window_left SINGLE_DECODE_ADDITIONAL_FUNC_PARAMS); - -at::Tensor BatchDecodeWithPagedKVCachePlan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, - int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, - int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, - at::Tensor empty_q_data, at::Tensor empty_kv_data); - -void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, - at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, - at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, - int64_t kv_layout_code, - int64_t window_left BATCH_DECODE_ADDITIONAL_FUNC_PARAMS); - -//========== gemm ========== - -void bmm_fp8(at::Tensor A, at::Tensor B, at::Tensor D, at::Tensor A_scale, at::Tensor B_scale, - at::Tensor workspace_buffer, int64_t cublas_handle); - -void CutlassSegmentGEMM(at::Tensor workspace_buffer, at::Tensor all_problems, at::Tensor x_ptr, - at::Tensor w_ptr, at::Tensor y_ptr, at::Tensor x_ld, at::Tensor w_ld, - at::Tensor y_ld, at::Tensor empty_x_data, bool weight_column_major); - -//========== norm ========== - -void rmsnorm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double eps, bool enable_pdl); - -void fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps, - bool enable_pdl); - -void gemma_rmsnorm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double eps, - bool enable_pdl); - -void gemma_fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, - double eps, bool enable_pdl); - -//========== page ========== - -void append_paged_kv_cache(at::Tensor append_key, at::Tensor append_value, at::Tensor batch_indices, - at::Tensor positions, at::Tensor paged_k_cache, at::Tensor paged_v_cache, - at::Tensor kv_indices, at::Tensor kv_indptr, at::Tensor kv_last_page_len, - int64_t layout); - -void append_paged_mla_kv_cache(at::Tensor append_ckv, at::Tensor append_kpe, - at::Tensor batch_indices, at::Tensor positions, at::Tensor ckv_cache, - at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor kv_indptr, - at::Tensor kv_last_page_len); - -void block_sparse_indices_to_vector_sparse_offsets( - at::Tensor block_sparse_indices, at::Tensor block_sparse_indptr, - at::Tensor vector_sparse_offsets, at::Tensor vector_sparse_indptr, at::Tensor kv_len_arr, - int64_t stride_block, int64_t stride_n, int64_t batch_size, int64_t block_size); - -//========== prefill ========== - -void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, - at::Tensor o, std::optional maybe_lse, - int64_t mask_mode_code, int64_t layout, - int64_t window_left SINGLE_PREFILL_ADDITIONAL_FUNC_PARAMS); - -at::Tensor BatchPrefillWithKVCachePlan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, - int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, - int64_t head_dim_vo, bool causal); - -void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, - at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor k, at::Tensor v, - at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, - int64_t layout, - int64_t window_left BATCH_PREFILL_ADDITIONAL_FUNC_PARAMS); - -void BatchPrefillWithPagedKVCacheRun( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, - int64_t window_left BATCH_PREFILL_ADDITIONAL_FUNC_PARAMS); - -//========== pod-attention ========= -void pod_with_kv_cache_tensor( - // Prefill params - at::Tensor q_p, at::Tensor k_p, at::Tensor v_p, at::Tensor tmp_p, at::Tensor o_p, - std::optional maybe_lse_p, int64_t mask_mode_code_p, int64_t layout_p, - int64_t window_left_p, std::optional maybe_custom_mask_p, - std::optional maybe_alibi_slopes_p, double logits_soft_cap_p, double sm_scale_p, - double rope_rcp_scale_p, double rope_rcp_theta_p, - // Decode params - at::Tensor float_workspace_buffer_d, at::Tensor int_workspace_buffer_d, - at::Tensor plan_info_vec, at::Tensor q_d, at::Tensor paged_k_cache_d, - at::Tensor paged_v_cache_d, at::Tensor qo_indptr_d, at::Tensor paged_kv_indptr_d, - at::Tensor paged_kv_indices_d, at::Tensor paged_kv_last_page_len_d, at::Tensor o_d, - std::optional maybe_lse_d, int64_t mask_mode_code_d, int64_t layout_d, - int64_t window_left, std::optional maybe_custom_mask_d, - std::optional maybe_mask_indptr_d, std::optional maybe_alibi_slopes_d, - double logits_soft_cap_d, double sm_scale_d, double rope_rcp_scale_d, double rope_rcp_theta_d); -//========== quantization ========== - -void packbits(at::Tensor x, const std::string& bitorder, at::Tensor y); - -void segment_packbits(at::Tensor x, at::Tensor input_indptr, at::Tensor output_indptr, - const std::string& bitorder, at::Tensor y); - -//========== rope ========== - -void apply_rope(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, at::Tensor indptr, - at::Tensor offsets, int64_t rotary_dim, bool interleave, double rope_scale, - double rope_theta); - -void apply_llama31_rope(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, - at::Tensor indptr, at::Tensor offsets, int64_t rotary_dim, bool interleave, - double rope_scale, double rope_theta, double low_freq_factor, - double high_freq_factor, double old_context_length); - -void apply_rope_pos_ids(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, - at::Tensor pos_ids, int64_t rotary_dim, bool interleave, double rope_scale, - double rope_theta); - -void apply_llama31_rope_pos_ids(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, - at::Tensor pos_ids, int64_t rotary_dim, bool interleave, - double rope_scale, double rope_theta, double low_freq_factor, - double high_freq_factor, double old_context_length); - -void apply_rope_pos_ids_cos_sin_cache(at::Tensor q, at::Tensor k, at::Tensor q_rope, - at::Tensor k_rope, at::Tensor cos_sin_cache, - at::Tensor pos_ids, bool interleave); - -//========== sampling ========== - -void sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, bool deterministic, - std::optional gen); - -void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_p_arr, double top_p_val, - bool deterministic, std::optional gen); - -void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_k_arr, int64_t top_k_val, - bool deterministic, std::optional gen); - -void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_min_p_arr, double min_p_val, - bool deterministic, std::optional gen); - -void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_k_arr, double top_k_val, - std::optional maybe_top_p_arr, double top_p_val, - bool deterministic, std::optional gen); - -void top_p_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, - std::optional maybe_top_p_arr, double top_p_val); - -void top_k_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, - std::optional maybe_top_k_arr, int64_t top_k_val); - -void top_k_mask_logits(at::Tensor logits, at::Tensor mask_logits, - std::optional maybe_top_k_arr, int64_t top_k_val); - -void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_ids, - at::Tensor target_probs, at::Tensor output_token_ids, - at::Tensor output_accepted_token_num, - at::Tensor output_emitted_draft_token_num, bool deterministic, - std::optional gen); - -//========== Torch Library ========== - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // activation - // Fused SiLU and Mul - m.def("silu_and_mul", silu_and_mul); - // Fused GeLU Tanh and Mul - m.def("gelu_tanh_and_mul", gelu_tanh_and_mul); - // Fused GeLU and Mul - m.def("gelu_and_mul", gelu_and_mul); - - // cascade - // Merge two self-attention states - m.def("merge_state", merge_state); - // Merge another self-attention state in-place. - m.def("merge_state_in_place", merge_state_in_place); - // "Merge multiple self-attention states" - m.def("merge_states", merge_states); - - // decode - // "Single-request decode with KV-Cache operator" - m.def("single_decode_with_kv_cache", single_decode_with_kv_cache); - m.def("batch_decode_with_paged_kv_cache_plan", BatchDecodeWithPagedKVCachePlan); - m.def("batch_decode_with_paged_kv_cache_run", BatchDecodeWithPagedKVCacheRun); - - // gemm - // BMM FP8 - m.def("bmm_fp8", bmm_fp8); - // Cutlass Segment GEMM operator - m.def("cutlass_segment_gemm", CutlassSegmentGEMM); - - // norm - // Root mean square normalization - m.def("rmsnorm", rmsnorm); - // Fused add root mean square normalization - m.def("fused_add_rmsnorm", fused_add_rmsnorm); - // Gemma Root mean square normalization - m.def("gemma_rmsnorm", gemma_rmsnorm); - // Gemma Fused add root mean square normalization - m.def("gemma_fused_add_rmsnorm", gemma_fused_add_rmsnorm); - - // page - // Append paged KV-Cache operator - m.def("append_paged_kv_cache", append_paged_kv_cache); - // Append paged MLA KV-Cache operator - m.def("append_paged_mla_kv_cache", append_paged_mla_kv_cache); - // Precompute block sparse offsets - m.def("block_sparse_indices_to_vector_sparse_offsets", - block_sparse_indices_to_vector_sparse_offsets); - - // prefill - // Single-request prefill attention with KV-Cache operator - m.def("single_prefill_with_kv_cache", single_prefill_with_kv_cache); - m.def("batch_prefill_with_kv_cache_plan", BatchPrefillWithKVCachePlan); - m.def("batch_prefill_with_ragged_kv_cache_run", BatchPrefillWithRaggedKVCacheRun); - m.def("batch_prefill_with_paged_kv_cache_run", BatchPrefillWithPagedKVCacheRun); - - // pod-attention - // Temporarily disabled because we don't generate the implementation yet. - // m.def("pod_with_kv_cache_tensor", pod_with_kv_cache_tensor); - - // quantization - // GPU packbits operator - m.def("packbits", packbits); - // GPU segment packbits operator - m.def("segment_packbits", segment_packbits); - - // rope - // "Apply RoPE" - m.def("apply_rope", apply_rope); - // "Apply Llama 3.1 style RoPE" - m.def("apply_llama31_rope", apply_llama31_rope); - // "Apply RoPE with positional ids" - m.def("apply_rope_pos_ids", apply_rope_pos_ids); - // "Apply Llama 3.1 style RoPE with positional ids" - m.def("apply_llama31_rope_pos_ids", apply_llama31_rope_pos_ids); - // "Apply RoPE with positional ids and cosine/sine cache" - m.def("apply_rope_pos_ids_cos_sin_cache", apply_rope_pos_ids_cos_sin_cache); - - // sampling - // Sample from probabilities - m.def("sampling_from_probs", sampling_from_probs); - // Top-k sampling from probabilities - m.def("top_k_sampling_from_probs", top_k_sampling_from_probs); - // Min-p sampling from probabilities - m.def("min_p_sampling_from_probs", min_p_sampling_from_probs); - // Top-p sampling from probabilities - m.def("top_p_sampling_from_probs", top_p_sampling_from_probs); - // Top-k and top-p sampling from probabilities - m.def("top_k_top_p_sampling_from_probs", top_k_top_p_sampling_from_probs); - // Renormalize probabilities by top-k mask - m.def("top_k_renorm_probs", top_k_renorm_probs); - // Renormalize probabilities by top-p mask - m.def("top_p_renorm_probs", top_p_renorm_probs); - // Mask logits by top-k mask - m.def("top_k_mask_logits", top_k_mask_logits); - // Speculative sampling from sequence of probabilities - m.def("chain_speculative_sampling", chain_speculative_sampling); -} diff --git a/csrc/flashinfer_ops_sm90.cu b/csrc/flashinfer_ops_sm90.cu deleted file mode 100644 index 855df95487..0000000000 --- a/csrc/flashinfer_ops_sm90.cu +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "aot_default_additional_params.h" -#include "pytorch_extension_utils.h" - -void CutlassSegmentGEMMSM90(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor all_problems, at::Tensor x_ptr, at::Tensor w_ptr, - at::Tensor y_ptr, at::Tensor x_stride, at::Tensor weight_stride, - at::Tensor y_stride, at::Tensor empty_x_data, at::Tensor empty_y_data, - bool weight_column_major); - -void single_prefill_with_kv_cache_sm90( - at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, at::Tensor o, - std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, - int64_t window_left SINGLE_PREFILL_SM90_ADDITIONAL_FUNC_PARAMS); - -at::Tensor BatchPrefillWithKVCacheSM90Plan( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, - at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, - int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, - int64_t head_dim_vo, bool causal); - -void BatchPrefillWithRaggedKVCacheSM90Run( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, - at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, - int64_t window_left BATCH_PREFILL_SM90_ADDITIONAL_FUNC_PARAMS); - -void BatchPrefillWithPagedKVCacheSM90Run( - at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, - at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, - at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, - at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, - int64_t window_left BATCH_PREFILL_SM90_ADDITIONAL_FUNC_PARAMS); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // "Cutlass Segment GEMM operator for SM90" - m.def("cutlass_segment_gemm_sm90", CutlassSegmentGEMMSM90); - m.def("single_prefill_with_kv_cache_sm90", single_prefill_with_kv_cache_sm90); - m.def("batch_prefill_with_kv_cache_sm90_plan", BatchPrefillWithKVCacheSM90Plan); - m.def("batch_prefill_with_ragged_kv_cache_sm90_run", BatchPrefillWithRaggedKVCacheSM90Run); - m.def("batch_prefill_with_paged_kv_cache_sm90_run", BatchPrefillWithPagedKVCacheSM90Run); -} diff --git a/csrc/flashinfer_page_binding.cu b/csrc/flashinfer_page_binding.cu new file mode 100644 index 0000000000..dbab4f5cb8 --- /dev/null +++ b/csrc/flashinfer_page_binding.cu @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +using tvm::ffi::Tensor; + +void append_paged_kv_cache(TensorView append_key, TensorView append_value, TensorView batch_indices, + TensorView positions, TensorView paged_k_cache, TensorView paged_v_cache, + TensorView kv_indices, TensorView kv_indptr, TensorView kv_last_page_len, + int64_t layout); + +void append_paged_mla_kv_cache(TensorView append_ckv, TensorView append_kpe, + TensorView batch_indices, TensorView positions, TensorView ckv_cache, + TensorView kpe_cache, TensorView kv_indices, TensorView kv_indptr, + TensorView kv_last_page_len); + +void block_sparse_indices_to_vector_sparse_offsets( + TensorView block_sparse_indices, TensorView block_sparse_indptr, + TensorView vector_sparse_offsets, TensorView vector_sparse_indptr, TensorView kv_len_arr, + int64_t stride_block, int64_t stride_n, int64_t batch_size, int64_t block_size); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(append_paged_kv_cache, append_paged_kv_cache); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(append_paged_mla_kv_cache, append_paged_mla_kv_cache); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(block_sparse_indices_to_vector_sparse_offsets, + block_sparse_indices_to_vector_sparse_offsets); diff --git a/csrc/flashinfer_page_ops.cu b/csrc/flashinfer_page_ops.cu deleted file mode 100644 index 0ba8c295b0..0000000000 --- a/csrc/flashinfer_page_ops.cu +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "pytorch_extension_utils.h" - -void append_paged_kv_cache(at::Tensor append_key, at::Tensor append_value, at::Tensor batch_indices, - at::Tensor positions, at::Tensor paged_k_cache, at::Tensor paged_v_cache, - at::Tensor kv_indices, at::Tensor kv_indptr, at::Tensor kv_last_page_len, - int64_t layout); - -void append_paged_mla_kv_cache(at::Tensor append_ckv, at::Tensor append_kpe, - at::Tensor batch_indices, at::Tensor positions, at::Tensor ckv_cache, - at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor kv_indptr, - at::Tensor kv_last_page_len); - -void block_sparse_indices_to_vector_sparse_offsets( - at::Tensor block_sparse_indices, at::Tensor block_sparse_indptr, - at::Tensor vector_sparse_offsets, at::Tensor vector_sparse_indptr, at::Tensor kv_len_arr, - int64_t stride_block, int64_t stride_n, int64_t batch_size, int64_t block_size); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // "Append paged KV-Cache operator" - m.def("append_paged_kv_cache", append_paged_kv_cache); - // "Append paged MLA KV-Cache operator" - m.def("append_paged_mla_kv_cache", append_paged_mla_kv_cache); - // "Precompute block sparse offsets" - m.def("block_sparse_indices_to_vector_sparse_offsets", - block_sparse_indices_to_vector_sparse_offsets); -} diff --git a/csrc/flashinfer_quantization_ops.cu b/csrc/flashinfer_quantization_binding.cu similarity index 57% rename from csrc/flashinfer_quantization_ops.cu rename to csrc/flashinfer_quantization_binding.cu index 9ff295b49d..9e22fccc23 100644 --- a/csrc/flashinfer_quantization_ops.cu +++ b/csrc/flashinfer_quantization_binding.cu @@ -13,16 +13,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "pytorch_extension_utils.h" +#include "tvm_ffi_utils.h" -void packbits(at::Tensor x, const std::string& bitorder, at::Tensor y); +void packbits(TensorView x, const std::string& bitorder, TensorView y); -void segment_packbits(at::Tensor x, at::Tensor input_indptr, at::Tensor output_indptr, - const std::string& bitorder, at::Tensor y); +void segment_packbits(TensorView x, TensorView input_indptr, TensorView output_indptr, + const std::string& bitorder, TensorView y); -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // GPU packbits operator - m.def("packbits", packbits); - // GPU segment packbits operator - m.def("segment_packbits", segment_packbits); -} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(packbits, packbits); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(segment_packbits, segment_packbits); diff --git a/csrc/flashinfer_rope_binding.cu b/csrc/flashinfer_rope_binding.cu new file mode 100644 index 0000000000..94809da735 --- /dev/null +++ b/csrc/flashinfer_rope_binding.cu @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +using tvm::ffi::Tensor; + +void apply_rope(TensorView q, TensorView k, TensorView q_rope, TensorView k_rope, TensorView indptr, + TensorView offsets, int64_t rotary_dim, bool interleave, double rope_scale, + double rope_theta); + +void apply_llama31_rope(TensorView q, TensorView k, TensorView q_rope, TensorView k_rope, + TensorView indptr, TensorView offsets, int64_t rotary_dim, bool interleave, + double rope_scale, double rope_theta, double low_freq_factor, + double high_freq_factor, double old_context_length); + +void apply_rope_pos_ids(TensorView q, TensorView k, TensorView q_rope, TensorView k_rope, + TensorView pos_ids, int64_t rotary_dim, bool interleave, double rope_scale, + double rope_theta); + +void apply_llama31_rope_pos_ids(TensorView q, TensorView k, TensorView q_rope, TensorView k_rope, + TensorView pos_ids, int64_t rotary_dim, bool interleave, + double rope_scale, double rope_theta, double low_freq_factor, + double high_freq_factor, double old_context_length); + +void apply_rope_pos_ids_cos_sin_cache(TensorView q, TensorView k, TensorView q_rope, + TensorView k_rope, TensorView cos_sin_cache, + TensorView pos_ids, bool interleave); + +void rope_quantize(TensorView q_rope_in, TensorView k_rope_in, TensorView q_nope_in, + TensorView k_nope_in, TensorView q_rope_out, TensorView k_rope_out, + TensorView q_nope_out, TensorView k_nope_out, TensorView cos_sin_cache, + TensorView pos_ids, double quant_scale_q, double quant_scale_kv, bool interleave, + bool enable_pdl); + +void rope_quantize_append_paged_kv_cache( + TensorView q_rope_in, TensorView k_rope_in, TensorView q_nope_in, TensorView k_nope_in, + TensorView v_in, TensorView q_rope_out, TensorView q_nope_out, TensorView cos_sin_cache, + TensorView pos_ids, TensorView k_cache, TensorView v_cache, TensorView ckv_cache, + TensorView kpe_cache, TensorView kv_indices, TensorView kv_indptr, TensorView batch_indices, + TensorView positions, int64_t kv_layout_code, int64_t page_size, double quant_scale_q, + double quant_scale_kv, bool interleave, bool enable_pdl); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(apply_rope, apply_rope); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(apply_llama31_rope, apply_llama31_rope); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(apply_rope_pos_ids, apply_rope_pos_ids); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(apply_llama31_rope_pos_ids, apply_llama31_rope_pos_ids); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(apply_rope_pos_ids_cos_sin_cache, apply_rope_pos_ids_cos_sin_cache); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(rope_quantize, rope_quantize); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(rope_quantize_append_paged_kv_cache, + rope_quantize_append_paged_kv_cache); diff --git a/csrc/flashinfer_rope_ops.cu b/csrc/flashinfer_rope_ops.cu deleted file mode 100644 index 84c5efe6fa..0000000000 --- a/csrc/flashinfer_rope_ops.cu +++ /dev/null @@ -1,53 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include - -#include "pytorch_extension_utils.h" - -void apply_rope(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, at::Tensor indptr, - at::Tensor offsets, int64_t rotary_dim, bool interleave, double rope_scale, - double rope_theta); - -void apply_llama31_rope(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, - at::Tensor indptr, at::Tensor offsets, int64_t rotary_dim, bool interleave, - double rope_scale, double rope_theta, double low_freq_factor, - double high_freq_factor, double old_context_length); - -void apply_rope_pos_ids(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, - at::Tensor pos_ids, int64_t rotary_dim, bool interleave, double rope_scale, - double rope_theta); - -void apply_llama31_rope_pos_ids(at::Tensor q, at::Tensor k, at::Tensor q_rope, at::Tensor k_rope, - at::Tensor pos_ids, int64_t rotary_dim, bool interleave, - double rope_scale, double rope_theta, double low_freq_factor, - double high_freq_factor, double old_context_length); - -void apply_rope_pos_ids_cos_sin_cache(at::Tensor q, at::Tensor k, at::Tensor q_rope, - at::Tensor k_rope, at::Tensor cos_sin_cache, - at::Tensor pos_ids, bool interleave); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // "Apply RoPE" - m.def("apply_rope", apply_rope); - // "Apply Llama 3.1 style RoPE" - m.def("apply_llama31_rope", apply_llama31_rope); - // "Apply RoPE with positional ids" - m.def("apply_rope_pos_ids", apply_rope_pos_ids); - // "Apply Llama 3.1 style RoPE with positional ids" - m.def("apply_llama31_rope_pos_ids", apply_llama31_rope_pos_ids); - // "Apply RoPE with positional ids and cosine/sine cache" - m.def("apply_rope_pos_ids_cos_sin_cache", apply_rope_pos_ids_cos_sin_cache); -} diff --git a/csrc/flashinfer_sampling_binding.cu b/csrc/flashinfer_sampling_binding.cu new file mode 100644 index 0000000000..8e4bbb98b8 --- /dev/null +++ b/csrc/flashinfer_sampling_binding.cu @@ -0,0 +1,87 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +using tvm::ffi::Optional; + +void softmax(TensorView workspace_buffer, TensorView logits, TensorView output, + Optional maybe_temperature_arr, double temperature_val, bool enable_pdl); + +void sampling_from_probs(TensorView probs, TensorView output, Optional maybe_indices, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset); + +void sampling_from_logits(TensorView logits, TensorView output, Optional maybe_indices, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset); + +void top_p_sampling_from_probs(TensorView probs, TensorView output, + Optional maybe_indices, + Optional maybe_top_p_arr, double top_p_val, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset); + +void top_k_sampling_from_probs(TensorView probs, TensorView output, + Optional maybe_indices, + Optional maybe_top_k_arr, int64_t top_k_val, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset); + +void min_p_sampling_from_probs(TensorView probs, TensorView output, + Optional maybe_indices, + Optional maybe_min_p_arr, double min_p_val, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset); + +void top_k_top_p_sampling_from_probs(TensorView probs, TensorView output, + Optional maybe_indices, + Optional maybe_top_k_arr, double top_k_val, + Optional maybe_top_p_arr, double top_p_val, + bool deterministic, uint64_t philox_seed, + uint64_t philox_offset); + +void top_p_renorm_probs(TensorView probs, TensorView renorm_probs, + Optional maybe_top_p_arr, double top_p_val); + +void top_k_renorm_probs(TensorView probs, TensorView renorm_probs, + Optional maybe_top_k_arr, int64_t top_k_val); + +void top_k_mask_logits(TensorView logits, TensorView mask_logits, + Optional maybe_top_k_arr, int64_t top_k_val); + +void chain_speculative_sampling(TensorView draft_probs, TensorView draft_token_ids, + TensorView target_probs, TensorView output_token_ids, + TensorView output_accepted_token_num, + TensorView output_emitted_draft_token_num, bool deterministic, + uint64_t philox_seed, uint64_t philox_offset); + +// Softmax +TVM_FFI_DLL_EXPORT_TYPED_FUNC(softmax, softmax); +// Sample from probabilities +TVM_FFI_DLL_EXPORT_TYPED_FUNC(sampling_from_probs, sampling_from_probs); +// Sample from logits +TVM_FFI_DLL_EXPORT_TYPED_FUNC(sampling_from_logits, sampling_from_logits); +// Top-k sampling from probabilities +TVM_FFI_DLL_EXPORT_TYPED_FUNC(top_k_sampling_from_probs, top_k_sampling_from_probs); +// Min-p sampling from probabilities +TVM_FFI_DLL_EXPORT_TYPED_FUNC(min_p_sampling_from_probs, min_p_sampling_from_probs); +// Top-p sampling from probabilities +TVM_FFI_DLL_EXPORT_TYPED_FUNC(top_p_sampling_from_probs, top_p_sampling_from_probs); +// Top-k and top-p sampling from probabilities +TVM_FFI_DLL_EXPORT_TYPED_FUNC(top_k_top_p_sampling_from_probs, top_k_top_p_sampling_from_probs); +// Renormalize probabilities by top-k mask +TVM_FFI_DLL_EXPORT_TYPED_FUNC(top_k_renorm_probs, top_k_renorm_probs); +// Renormalize probabilities by top-p mask +TVM_FFI_DLL_EXPORT_TYPED_FUNC(top_p_renorm_probs, top_p_renorm_probs); +// Mask logits by top-k mask +TVM_FFI_DLL_EXPORT_TYPED_FUNC(top_k_mask_logits, top_k_mask_logits); +// Speculative sampling from sequence of probabilities +TVM_FFI_DLL_EXPORT_TYPED_FUNC(chain_speculative_sampling, chain_speculative_sampling); diff --git a/csrc/flashinfer_sampling_ops.cu b/csrc/flashinfer_sampling_ops.cu deleted file mode 100644 index 7a0547f07b..0000000000 --- a/csrc/flashinfer_sampling_ops.cu +++ /dev/null @@ -1,77 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "pytorch_extension_utils.h" - -void sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, bool deterministic, - std::optional gen); - -void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_p_arr, double top_p_val, - bool deterministic, std::optional gen); - -void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_k_arr, int64_t top_k_val, - bool deterministic, std::optional gen); - -void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_min_p_arr, double min_p_val, - bool deterministic, std::optional gen); - -void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_k_arr, double top_k_val, - std::optional maybe_top_p_arr, double top_p_val, - bool deterministic, std::optional gen); - -void top_p_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, - std::optional maybe_top_p_arr, double top_p_val); - -void top_k_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, - std::optional maybe_top_k_arr, int64_t top_k_val); - -void top_k_mask_logits(at::Tensor logits, at::Tensor mask_logits, - std::optional maybe_top_k_arr, int64_t top_k_val); - -void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_ids, - at::Tensor target_probs, at::Tensor output_token_ids, - at::Tensor output_accepted_token_num, - at::Tensor output_emitted_draft_token_num, bool deterministic, - std::optional gen); - -TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { - // Sample from probabilities - m.def("sampling_from_probs", sampling_from_probs); - // Top-k sampling from probabilities - m.def("top_k_sampling_from_probs", top_k_sampling_from_probs); - // Min-p sampling from probabilities - m.def("min_p_sampling_from_probs", min_p_sampling_from_probs); - // Top-p sampling from probabilities - m.def("top_p_sampling_from_probs", top_p_sampling_from_probs); - // Top-k and top-p sampling from probabilities - m.def("top_k_top_p_sampling_from_probs", top_k_top_p_sampling_from_probs); - // Renormalize probabilities by top-k mask - m.def("top_k_renorm_probs", top_k_renorm_probs); - // Renormalize probabilities by top-p mask - m.def("top_p_renorm_probs", top_p_renorm_probs); - // Mask logits by top-k mask - m.def("top_k_mask_logits", top_k_mask_logits); - // Speculative sampling from sequence of probabilities - m.def("chain_speculative_sampling", chain_speculative_sampling); -} diff --git a/csrc/flashinfer_xqa_binding.cu b/csrc/flashinfer_xqa_binding.cu new file mode 100644 index 0000000000..dc06614763 --- /dev/null +++ b/csrc/flashinfer_xqa_binding.cu @@ -0,0 +1,45 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tvm_ffi_utils.h" + +#if MLA_WRAPPER +void xqa_wrapper_mla(int64_t multiProcessorCount, double qScale, + tvm::ffi::Optional qScaleTensor, TensorView output, TensorView q, + TensorView kCacheVLLM, TensorView vCacheVLLM, TensorView kvCachePageList, + int64_t maxSeqLen, TensorView seqLen, int64_t batchSize, double kvCacheScale, + tvm::ffi::Optional kvScaleTensor, TensorView semaphores, + TensorView scratch, bool enable_pdl); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(xqa_wrapper_mla, xqa_wrapper_mla); + +#else + +void xqa_wrapper(bool run_sm90_fp8_mha, int64_t multiProcessorCount, int64_t nbKHeads, + int64_t slidingWinSize, double qScale, tvm::ffi::Optional qScaleTensor, + TensorView output, double rcpOutScale, TensorView q, + tvm::ffi::Optional attentionSinks, TensorView kCacheVLLM, + TensorView vCacheVLLM, TensorView kvCachePageList, int64_t maxSeqLen, + TensorView seqLen, int64_t batchSize, double kvCacheScale, + tvm::ffi::Optional kvScaleTensor, +#if SPEC_DEC + int64_t qSeqLen, TensorView qCuSeqLens, TensorView mask, +#endif + TensorView semaphores, TensorView scratch, bool enable_pdl); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(xqa_wrapper, xqa_wrapper); + +#endif diff --git a/csrc/fmhaReduction.cu b/csrc/fmhaReduction.cu new file mode 100644 index 0000000000..1f1ca8c755 --- /dev/null +++ b/csrc/fmhaReduction.cu @@ -0,0 +1,358 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "flashinfer/exception.h" +#include "flashinfer/trtllm/common/cudaTypeUtils.cuh" +#include "flashinfer/trtllm/common/cudaUtils.h" +#include "flashinfer/trtllm/fmha/fmhaReduction.h" +#include "flashinfer/trtllm/fmha/kernelUtils.h" +#include "flashinfer/utils.cuh" + +namespace tensorrt_llm { +namespace kernels { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define NumThreadsPerCta 512 + +template +__global__ void __launch_bounds__(NumThreadsPerCta, 2) + fmhaReductionKernel(KernelParams const params, int32_t numCtasForReduction, + int32_t numCtasForAllHeads, int32_t numHeadDimCtasV) { + // clang-format off + // The shape of partialO buffer: [batchSize, numHeadCtas, numCtasQ, numCtasKv, TileSizePerCtaQ, headDimPerCta]. + // The shape of final O buffer: [batchSize, numCtasQ, numHeadsQ, headDim]. + // The shape of attentionSinks buffer: [numHeadsQ]. + // The shape of partialStats buffer: [batchSize, numHeadCtas, numCtasQ, numCtasKv, TileSizePerCtaQ], where each element is a float2 (max/sum). + // The shape of softmaxStats buffer: [batchSize, numCtasQ, numHeadsQ], where each element is a float2 (max/sum). + // Note that numValidRows includes both numValidTokens and numHeadsQPerKv if grouping headsQ. + // clang-format on + + // The batchIdx. + int32_t const batchIdx{static_cast(blockIdx.z)}; + // The headCtaIdxO. + int32_t const headCtaIdxO{static_cast(blockIdx.y)}; + // The headDimCtaIdxV. + int32_t const headDimCtaIdxV{static_cast(blockIdx.y % numHeadDimCtasV)}; + // The headGrpIdxO. + int32_t const headGrpIdxO{static_cast(blockIdx.y / numHeadDimCtasV)}; + // The ctaIdxQ. + int32_t const ctaIdxQ{static_cast(blockIdx.x % params.mMaxNumCtasQ)}; + // The ctaIdx for the reduction work. + int32_t const ctaIdxForReduction{static_cast(blockIdx.x / params.mMaxNumCtasQ)}; + // The headIdxO. + int32_t const headIdxO{headGrpIdxO * TileSizePerCtaQ}; + // The warpGrpThreadIdx. + int32_t const warpGrpThreadIdx{static_cast(threadIdx.x)}; + + // The number of validRows. + int32_t const numValidRows{TileSizePerCtaQ}; + // The actual number of seqLenKv. + int32_t seqLenKv{params.ptrSeqLensKv[batchIdx]}; + // Consider the causal-mask speculative decoding. + seqLenKv = seqLenKv - ((params.mMaxSeqLenQ - 1) - ctaIdxQ); + // The actual number of CtasKv (TileSizeKv is always 128 for now). + int32_t numCtasKv{min((seqLenKv + 127) / 128, params.mMaxNumCtasKv)}; + + // The tileIdx in the batch/head dimension. + int64_t const batchHeadTileIdx{ + ((batchIdx * static_cast(gridDim.y) + headCtaIdxO) * params.mMaxNumCtasQ + ctaIdxQ)}; + + // The offset of the partialStats buffer. + int64_t const partialStatsOffset{batchHeadTileIdx * params.mMaxNumCtasKv * TileSizePerCtaQ}; + // The offset of the partialO buffer. + int64_t const partialOOffset{partialStatsOffset * HeadDimPerCta}; + // The offset of the softmaxStats buffer. + int64_t const softmaxStatsOffset{ + ((batchIdx * params.mMaxNumCtasQ + ctaIdxQ) * numCtasForAllHeads + headGrpIdxO) * + TileSizePerCtaQ}; + // The offset of the O buffer. + int64_t const oOffset{softmaxStatsOffset * HeadDim + headDimCtaIdxV * HeadDimPerCta}; + + // The partialStats pointer. + float2* partialStatsPtr = reinterpret_cast(params.ptrPartialStats) + partialStatsOffset; + // The partialO pointer. + DtypePartialO* partialOPtr = + reinterpret_cast(params.ptrPartialO) + partialOOffset; + // The softmaxStats pointer. + float2* softmaxStatsPtr = reinterpret_cast(params.ptrSoftmaxStats) + softmaxStatsOffset; + // The O pointer. + DtypeO* oPtr = reinterpret_cast(params.ptrO) + oOffset; + // The attentionSinks pointer. + float const* attentionSinksPtr = params.ptrAttentionSinks + headIdxO; + + // Whether to store the softmax stats. + bool const storesSoftmaxStats{params.ptrSoftmaxStats != nullptr}; + + // The softmaxScaleLog2. + float const softmaxScaleLog2 = params.mScaleSoftmaxLog2; + + int32_t constexpr NumBytesPerPartialElt{sizeof(DtypePartialO)}; + static_assert(NumBytesPerPartialElt == 2, + "The data type of partialO should be either fp16 or bf16."); + + // The threads in the warp-group should load different values from one partial output + // [numValidRows, headDim], and then iterate over partial outputs from different CTAs. + int32_t constexpr NumEltsPer16BVec{16 / NumBytesPerPartialElt}; + static_assert((HeadDimPerCta * NumBytesPerPartialElt) % 16 == 0, "Not implemented"); + + // The number of unrolled iterations to issue multiple LDGs. + int32_t constexpr UnrollSize{4}; + + // The number of processed rows in one slice where each CTA will process one slice. + int32_t constexpr NumBytesPerHeadDim{HeadDimPerCta * NumBytesPerPartialElt}; + int32_t constexpr NumBytePerSlice{NumThreadsPerCta * 16}; + static_assert(NumBytePerSlice % NumBytesPerHeadDim == 0, "Not implemented"); + int32_t constexpr NumRowsPerSlice{NumBytePerSlice / NumBytesPerHeadDim}; + // The actual number of tensor slices for the reduction. + int32_t numSlices{(numValidRows + NumRowsPerSlice - 1) / NumRowsPerSlice}; + + // The number of slices that each CTA will process. + int32_t numSlicesPerCta{(numSlices + numCtasForReduction - 1) / numCtasForReduction}; + // The start slice index for the current CTA. + int32_t startSliceIdx{ctaIdxForReduction * numSlicesPerCta}; + // The end slice index for the current CTA. + int32_t endSliceIdx{min(startSliceIdx + numSlicesPerCta, numSlices)}; + + // The total number of rows in the partial buffers. + int32_t numRowsInPartialBuffers{TileSizePerCtaQ}; + + // Iterate over different slices. + // Split the reduction work across multiple CtasKv to reduce the latency. + for (int32_t sliceIdx = startSliceIdx; sliceIdx < endSliceIdx; ++sliceIdx) { + // The base offset that each thread points to. + int32_t const baseOffset{warpGrpThreadIdx * NumEltsPer16BVec}; + // The index in the row dimension. + int32_t const rowIdx{sliceIdx * NumRowsPerSlice + (baseOffset / HeadDimPerCta)}; + // Does this thread point to a valid row ? + bool const isValidRow{rowIdx < numValidRows}; + int32_t validRowIdx{min(rowIdx, numValidRows - 1)}; + int32_t loadRowIdx{validRowIdx}; + // The index in the headDim dimension. + int32_t const headDimIdx{baseOffset % HeadDimPerCta}; + // The memory load offset. + int64_t const destMemOffset{loadRowIdx * HeadDimPerCta + headDimIdx}; + // The memory store offset. + int64_t gmemStoreOffset{validRowIdx * HeadDim + headDimIdx}; + // The local headIdxO. + int32_t localHeadIdxO{validRowIdx}; + +// Wait for the primary kernel to complete. +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + // Add offset to the pointers. + float2* localPartialStatsPtr = partialStatsPtr + loadRowIdx; + DtypePartialO* localPartialOPtr = partialOPtr + destMemOffset; + + // Reduce max, sum and partialO vectors from different CtasKv. + float sumVal{0.f}; + float oldMaxVal{-FLT_MAX}, maxVal{-FLT_MAX}; + float outputVals[NumEltsPer16BVec] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + for (int32_t ii = 0; ii < numCtasKv; ii += UnrollSize) { + // The partialStats array and partialO array. + float2 partialStatsArray[UnrollSize]; + uint4 partialOArray[UnrollSize]; +#pragma unroll + for (int32_t jj = 0; jj < UnrollSize; ++jj) { + int32_t ctaIdxKv = min(ii + jj, numCtasKv - 1); + partialStatsArray[jj] = localPartialStatsPtr[ctaIdxKv * numRowsInPartialBuffers]; + partialOArray[jj] = *reinterpret_cast( + localPartialOPtr + ctaIdxKv * numRowsInPartialBuffers * HeadDimPerCta); + } +#pragma unroll + for (int32_t jj = 0; jj < UnrollSize; ++jj) { + // Whether the ctaIdxKv is valid. + bool const isValidCtaIdxKv = (ii + jj) < numCtasKv; + // The local max and sum values. + auto partialStats = partialStatsArray[jj]; + float localMax = partialStats.x; + float localSum = partialStats.y; + // Update the max value. + maxVal = fmaxf(maxVal, localMax); + // Compute the correction scales. + float corrScale0 = isValidCtaIdxKv ? exp2f(softmaxScaleLog2 * (oldMaxVal - maxVal)) : 1.f; + float corrScale1 = isValidCtaIdxKv ? exp2f(softmaxScaleLog2 * (localMax - maxVal)) : 0.f; + // Update the old max value. + oldMaxVal = maxVal; + // The partialO value. + uint4 vec = partialOArray[jj]; + // Reduce sum and finalO. + sumVal = sumVal * corrScale0 + localSum * corrScale1; + convertToFloatAndAccumulate(outputVals, vec, corrScale0, corrScale1); + } + } + + // Update the sums with the attention sink value. + if (attentionSinksPtr != nullptr) { + float attentionSinkVal = + exp2f(attentionSinksPtr[localHeadIdxO] * M_LOG2E - maxVal * softmaxScaleLog2); + // Multiply the attention sink value by 448.f if the MMA data type is e4m3 as the sum value + // has also included the 448.f quantization scale. + sumVal += IsE4m3Bmm ? attentionSinkVal * 448.f : attentionSinkVal; + } + + // Stores the final softmax stats values to global memory if needed (Helix attention, which + // splits seqLenKv across GPUs). + if (storesSoftmaxStats && isValidRow && headDimIdx == 0) { + // The softmaxScale. + float softmaxScale = (softmaxScaleLog2 * (1.f / M_LOG2E)); + // The sumScale to unscale the 448.f quantization scale from P. + float sumScale = IsE4m3Bmm ? (1.f / 448.f) : 1.f; + // The final max and sum values. + float2 stats{maxVal * softmaxScale, sumVal * sumScale}; + // Store the final max and sum values to global memory. + reinterpret_cast(softmaxStatsPtr)[validRowIdx] = stats; + } + + // The final normalized scale. + // If the output data type is e4m3, make sure that sumVal is divided by the quantization scale + // (448.f), so 1.0f / (sumVal / 448.f) = 448.f / sumVal. + float normalizedScale{IsE4m3Bmm ? (448.f / sumVal) : (1.0f / sumVal)}; + float2 normalizedScale2{normalizedScale, normalizedScale}; + + // Apply the normalized scale to the reduced O values. + for (int ii = 0; ii < NumEltsPer16BVec / 2; ++ii) { + float2& f2 = reinterpret_cast(outputVals)[ii]; + mul(f2, f2, normalizedScale2); + } + + // Convert the float values to DtypeO, and Store it to global memory. + if (isValidRow) { + convertAndStoreToGmem(reinterpret_cast(oPtr + gmemStoreOffset), outputVals); + } + } + +// Trigger the secondary kernel. +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define SELECT_FMHA_REDUCTION_KERNEL(HeadDimPerCta) \ + if (kernelMeta.mDataTypeQ == DATA_TYPE_E4M3) { \ + if (kernelMeta.mDataTypeO == DATA_TYPE_E4M3) { \ + kernel = &fmhaReductionKernel<64, 512, HeadDimPerCta, true, __nv_fp8_e4m3, half>; \ + } else if (kernelMeta.mDataTypeO == DATA_TYPE_FP16) { \ + kernel = &fmhaReductionKernel<64, 512, HeadDimPerCta, true, half, half>; \ + } else if (kernelMeta.mDataTypeO == DATA_TYPE_BF16) { \ + kernel = &fmhaReductionKernel<64, 512, HeadDimPerCta, true, __nv_bfloat16, __nv_bfloat16>; \ + } else { \ + FLASHINFER_CHECK(false, "Not implemented"); \ + } \ + } else { \ + FLASHINFER_CHECK(kernelMeta.mDataTypeQ == kernelMeta.mDataTypeO, "Not implemented"); \ + if (kernelMeta.mDataTypeQ == DATA_TYPE_FP16) { \ + kernel = &fmhaReductionKernel<64, 512, HeadDimPerCta, false, half, half>; \ + } else if (kernelMeta.mDataTypeQ == DATA_TYPE_BF16) { \ + kernel = &fmhaReductionKernel<64, 512, HeadDimPerCta, false, __nv_bfloat16, __nv_bfloat16>; \ + } else { \ + FLASHINFER_CHECK(false, "Not implemented"); \ + } \ + } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void runFmhaReduction(TllmGenFmhaKernelMetaInfo const& kernelMeta, KernelParams const& params, + int32_t multiProcessorCount, bool enable_pdl, cudaStream_t stream) { + // Skip the kernel if not using the separate reduction kernel. + if (!isGmemReductionWithSeparateKernel( + static_cast(kernelMeta.mMultiCtasKvMode))) { + return; + } + + // This should only be enabled when the keepsMmaAbForGeneration MLA kernel (either 1-CTA or 2-CTA) + // is used. + FLASHINFER_CHECK( + kernelMeta.mHeadDimQk == 576 && kernelMeta.mHeadDimV == 512 && + isKeepsMmaAbForGenerationKernel(static_cast(kernelMeta.mKernelType)), + "Not implemented"); + // The tileSizeQ and tileSizeKv should be 64 and 128 for those kernels. + FLASHINFER_CHECK(kernelMeta.mTileSizeQ == 64 && kernelMeta.mTileSizeKv == 128, "Not implemented"); + + // The headDimPerCtaV. + int32_t const headDimPerCtaV = + kernelMeta.m2CtaMma ? kernelMeta.mHeadDimPerCtaV * 2 : kernelMeta.mHeadDimPerCtaV; + FLASHINFER_CHECK(headDimPerCtaV == 128 || headDimPerCtaV == 256 || headDimPerCtaV == 512, + "Not implemented"); + + // The number of slices for the reduction work. + int32_t const numSlices = (headDimPerCtaV * /* bytesPerPartialElt */ 2 * kernelMeta.mTileSizeQ) / + (NumThreadsPerCta * 16); + // The number of Ctas for all heads. + int32_t const numCtasForAllHeads{params.mNumHeadsQ / kernelMeta.mTileSizeQ}; + // The number of Ctas for headDim. + int32_t const numHeadDimCtasV{kernelMeta.mHeadDimV / headDimPerCtaV}; + + // The 512 threads will split the reduction work of TileSizePerCtaQ * HeadDimPerCta. + dim3 blockDim(NumThreadsPerCta); + dim3 gridDim; + // Each CTA processes one tokenQ. + gridDim.x = params.mMaxNumCtasQ; + // The head dimension. + gridDim.y = numCtasForAllHeads * numHeadDimCtasV; + // The batch dimension. + gridDim.z = params.mBatchSize; + + // The maximum number of Ctas for the reduction work. + // This avoids having too many waves of CTAs which can have obvious launching overheads. + int32_t const maxNumCtasForReduction{(multiProcessorCount * 2) / + static_cast(gridDim.x * gridDim.y * gridDim.z)}; + // The number of Ctas for the reduction work. + int32_t const numCtasForReduction{std::min(maxNumCtasForReduction, numSlices)}; + // Launch more CTAs to split the reduction work if needed. + gridDim.x *= numCtasForReduction; + + // The PDL attribute. + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = enable_pdl ? 1 : 0; + cudaLaunchConfig_t config; + config.gridDim = gridDim; + config.blockDim = blockDim; + config.stream = stream; + config.dynamicSmemBytes = 0; + config.attrs = attribute; + config.numAttrs = 1; + + // Select the kernel function pointer. + void (*kernel)(KernelParams const, int32_t, int32_t, int32_t) = nullptr; + if (headDimPerCtaV == 128) { + SELECT_FMHA_REDUCTION_KERNEL(128); + } else if (headDimPerCtaV == 256) { + SELECT_FMHA_REDUCTION_KERNEL(256); + } else if (headDimPerCtaV == 512) { + SELECT_FMHA_REDUCTION_KERNEL(512); + } + + // Launch the kernel. + cudaLaunchKernelEx(&config, kernel, params, numCtasForReduction, numCtasForAllHeads, + numHeadDimCtasV); + cudaError_t err = cudaGetLastError(); + FLASHINFER_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernels +} // namespace tensorrt_llm diff --git a/csrc/fmha_cutlass_sm100.cu b/csrc/fmha_cutlass_sm100.cu new file mode 100644 index 0000000000..c50116fa7f --- /dev/null +++ b/csrc/fmha_cutlass_sm100.cu @@ -0,0 +1,131 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "tvm_ffi_utils.h" + +using tvm::ffi::Optional; + +#define DISPATCH_mask_mode(mask_mode, MASK_MODE, ...) \ + [&]() -> bool { \ + if (mask_mode == MaskMode::kNone) { \ + constexpr MaskMode MASK_MODE = MaskMode::kNone; \ + return __VA_ARGS__(); \ + } else if (mask_mode == MaskMode::kCausal) { \ + constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ + return __VA_ARGS__(); \ + } \ + return false; \ + }() + +#define DISPATCH_head_dim(head_dim_qk, head_dim_vo, HEAD_DIM_QK, HEAD_DIM_VO, ...) \ + [&]() -> bool { \ + if (head_dim_qk == 192 && head_dim_vo == 128) { \ + constexpr int HEAD_DIM_QK = 192; \ + constexpr int HEAD_DIM_VO = 128; \ + return __VA_ARGS__(); \ + } else if (head_dim_qk == 128 && head_dim_vo == 128) { \ + constexpr int HEAD_DIM_QK = 128; \ + constexpr int HEAD_DIM_VO = 128; \ + return __VA_ARGS__(); \ + } else if (head_dim_qk == 64 && head_dim_vo == 64) { \ + constexpr int HEAD_DIM_QK = 64; \ + constexpr int HEAD_DIM_VO = 64; \ + return __VA_ARGS__(); \ + } \ + return false; \ + }() + +#define DISPATCH_DTYPE_IN_OUT(in_dtype, out_dtype, c_type_in, c_type_out, ...) \ + [&]() -> bool { \ + if (in_dtype == out_dtype) { \ + return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(in_dtype, c_type_in, [&] { \ + using c_type_out = c_type_in; \ + return __VA_ARGS__(); \ + }); \ + } \ + return false; \ + }() + +#define DISPATCH_context(DTypeIn, DTypeOut, HEAD_DIM_QK, HEAD_DIM_VO, MaskMode, ...) \ + { \ + DISPATCH_mask_mode(mask_mode, MaskMode, [&] { \ + return DISPATCH_DTYPE_IN_OUT(scalar_type_in, scalar_type_out, DTypeIn, DTypeOut, [&] { \ + return DISPATCH_head_dim(head_dim_qk, head_dim_vo, HEAD_DIM_QK, HEAD_DIM_VO, \ + [&] { return __VA_ARGS__(); }); \ + }); \ + }); \ + } + +using namespace flashinfer; + +void FMHACutlassSM100Run(ffi::TensorView workspace_buffer, ffi::TensorView q, ffi::TensorView k, + ffi::TensorView v, ffi::TensorView qo_segment_offsets, + ffi::TensorView kv_segment_offsets, ffi::TensorView work_indptr, + ffi::TensorView qo_tile_indices, ffi::TensorView qo_head_indices, + ffi::TensorView batch_indices, ffi::TensorView o, + Optional maybe_lse, int64_t mask_mode_code, + double sm_scale, int64_t num_qo_heads, int64_t num_kv_heads, + int64_t head_dim_qk, int64_t head_dim_vo, int64_t max_qo_len) { + TVM_FFI_ICHECK_EQ(q.dtype(), k.dtype()); + auto scalar_type_in = q.dtype(); + auto scalar_type_out = o.dtype(); + MaskMode mask_mode = static_cast(mask_mode_code); + int total_qo_len = q.size(0); + int total_kv_len = k.size(0); + int batch_size = qo_segment_offsets.size(0) - 1; + int q_stride_n = q.stride(0); + int q_stride_h = q.stride(1); + int k_stride_n = k.stride(0); + int k_stride_h = k.stride(1); + int v_stride_n = v.stride(0); + int v_stride_h = v.stride(1); + + cudaSetDevice(qo_segment_offsets.device().device_id); + const cudaStream_t stream = get_stream(o.device()); + + DISPATCH_context(DTypeIn, DTypeOut, HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, [&] { + using cutlass_type_in = cutlass_dtype_t; + using cutlass_type_out = cutlass_dtype_t; + using TILE_Q = _256; + using TILE_KV = _128; + using D_QK = cute::Int; + using D_VO = cute::Int; + using TileShapeQK = Shape; + using TileShapePV = Shape; + using CutlassMaskMode = + typename std::conditional::type; + auto status = run_fmha_fwd( + workspace_buffer.data_ptr(), static_cast(q.data_ptr()), + static_cast(k.data_ptr()), static_cast(v.data_ptr()), + static_cast(qo_segment_offsets.data_ptr()), + static_cast(kv_segment_offsets.data_ptr()), static_cast(work_indptr.data_ptr()), + static_cast(qo_tile_indices.data_ptr()), + static_cast(qo_head_indices.data_ptr()), static_cast(batch_indices.data_ptr()), + static_cast(o.data_ptr()), + maybe_lse.has_value() ? static_cast(maybe_lse.value().data_ptr()) : nullptr, + mask_mode_code, sm_scale, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, q_stride_n, + q_stride_h, k_stride_n, k_stride_h, v_stride_n, v_stride_h, batch_size, total_qo_len, + total_kv_len, max_qo_len, stream); + TVM_FFI_ICHECK_EQ(status, cudaSuccess) + << "Cutlass FMHA forward pass failed" << cudaGetErrorString(status); + + return true; + }); +} diff --git a/csrc/fmha_cutlass_sm100_binding.cu b/csrc/fmha_cutlass_sm100_binding.cu new file mode 100644 index 0000000000..ddb3b8d9cd --- /dev/null +++ b/csrc/fmha_cutlass_sm100_binding.cu @@ -0,0 +1,34 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +using tvm::ffi::Optional; + +void FMHACutlassSM100Run(TensorView workspace_buffer, TensorView q, TensorView k, TensorView v, + TensorView qo_segment_offsets, TensorView kv_segment_offsets, + TensorView work_indptr, TensorView qo_tile_indices, + TensorView qo_head_indices, TensorView batch_indices, TensorView o, + Optional maybe_lse, int64_t mask_mode_code, double sm_scale, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t head_dim_qk, + int64_t head_dim_vo, int64_t max_qo_len); + +void blackwell_fmha_plan(TensorView qo_segment_offsets, TensorView kv_segment_offsets, + TensorView work_indptr, TensorView qo_tile_indices, + TensorView head_indices, TensorView batch_indices, int64_t qo_tile_size, + int64_t num_heads, int64_t num_buckets, bool causal); + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, FMHACutlassSM100Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(plan, blackwell_fmha_plan); diff --git a/csrc/fp4_gemm_cutlass.cu b/csrc/fp4_gemm_cutlass.cu new file mode 100644 index 0000000000..ae9b0aa658 --- /dev/null +++ b/csrc/fp4_gemm_cutlass.cu @@ -0,0 +1,193 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include +#include +#include +#include +#include + +#include "flashinfer/gemm/cutlass_gemm_configs.h" +#include "flashinfer/gemm/fp4_gemm_cutlass.h" +#include "flashinfer/gemm/fp4_gemm_cutlass_template.h" +#include "tvm_ffi_utils.h" + +using flashinfer::gemm::ClusterShape; +using flashinfer::gemm::CutlassFp4GemmRunner; +using flashinfer::gemm::CutlassFp4GemmRunnerInterface; +using flashinfer::gemm::CutlassGemmConfig; +using flashinfer::gemm::CutlassTileConfigSM100; +using flashinfer::gemm::EpilogueScheduleType; +using flashinfer::gemm::FP4GemmType; +using flashinfer::gemm::MainloopScheduleType; + +namespace flashinfer { +namespace gemm { +template class CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4>; +template class CutlassFp4GemmRunner; +} // namespace gemm +} // namespace flashinfer + +namespace torch_ext { + +namespace { + +CutlassGemmConfig getFp4GemmConfig(int64_t m, int64_t n, int64_t k, int64_t tactic) { + auto getCutlassFp4GemmConfigs = []() { + CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4> gemmRunner; + return gemmRunner.getConfigs(); + }; + static std::vector globalConfigs = getCutlassFp4GemmConfigs(); + TVM_FFI_ICHECK(tactic >= 0 && tactic < globalConfigs.size()) + << "tactic must be between 0 and " << globalConfigs.size(); + return globalConfigs[tactic]; +} + +template +void runGemm(TensorView out, TensorView mat1, TensorView mat2, TensorView mat1Scale, + TensorView mat2Scale, TensorView globalScale, int64_t m, int64_t n, int64_t k, + int64_t batch_count, CutlassGemmConfig const& gemmConfig, + TensorView workspace_buffer) { + CutlassFp4GemmRunner gemmRunner; + + int64_t const required_workspace_size = gemmRunner.getWorkspaceSize(m, n, k, batch_count); + int64_t const provided_workspace_size = + workspace_buffer.numel() * get_element_size(workspace_buffer); + + auto runKernel = [&](void* workspace) { + gemmRunner.gemm(out.data_ptr(), mat1.data_ptr(), mat2.data_ptr(), mat1Scale.data_ptr(), + mat2Scale.data_ptr(), static_cast(globalScale.data_ptr()), m, n, k, + batch_count, gemmConfig, reinterpret_cast(workspace), + required_workspace_size, get_stream(mat1.device())); + }; + + if (provided_workspace_size < required_workspace_size) { + Tensor new_workspace = + alloc_tensor({required_workspace_size}, DLDataType{kDLInt, 8, 1}, mat1.device()); + runKernel(new_workspace.data_ptr()); + } else { + runKernel(workspace_buffer.data_ptr()); + } +} + +constexpr auto FLOAT4_E2M1X2 = dl_uint8; // uint8_t +constexpr auto SF_DTYPE = dl_uint8; // uint8_t + +// mat1: [B, M, K / 2], FLOAT4_E2M1X2 or [B, M, K], FLOAT8_E4M3FN +// mat2: [B, N, K / 2], FLOAT4_E2M1X2 +// out: [B, M, N], fp16/bf16/fp32 +// mat1Scale: ceil(M / 128) * 128 * ceil(K / sfVecSize / 4) * 4, SF_DTYPE (UE4M3 or UE8M0) +// mat2Scale: ceil(N / 128) * 128 * ceil(K / sfVecSize / 4) * 4, SF_DTYPE (UE4M3 or UE8M0) +// globalScale: [1], 1 / (((448 * 6) / mat1.abs().max()) * ((448 * 6) / mat2.abs().max())) +// B = 1 for GEMM op as a special case +void fp4_bmm_impl(TensorView mat1, TensorView mat2, TensorView mat1Scale, TensorView mat2Scale, + TensorView globalScale, TensorView out, TensorView workspace_buffer, + int64_t tactic) { + CHECK_INPUT_AND_TYPE(mat1, FLOAT4_E2M1X2); + CHECK_INPUT_AND_TYPE(mat2, FLOAT4_E2M1X2); + + int mat2_k_scale = 1; + + CHECK_INPUT_AND_TYPE(mat1Scale, SF_DTYPE); + CHECK_INPUT_AND_TYPE(mat2Scale, SF_DTYPE); + + CHECK_INPUT_AND_TYPE(globalScale, dl_float32); + + int64_t m, n, k, b; + if (mat1.ndim() == 2) { + TVM_FFI_ICHECK_EQ(mat2.ndim(), 2) << "mat2 must be a matrix"; + TVM_FFI_ICHECK_EQ(mat1.size(1), mat2.size(1) * mat2_k_scale) + << "mat1 and mat2 shapes cannot be multiplied (" << mat1.size(0) << "x" << mat1.size(1) + << " and " << mat2.size(0) << "x" << mat2.size(1) << ")"; + m = mat1.size(0); + n = mat2.size(0); + k = mat2.size(1) * 2; + b = 1; + } else if (mat1.ndim() == 3) { + TVM_FFI_ICHECK_EQ(mat2.ndim(), 3) << "mat2 must be a batch of matrices"; + TVM_FFI_ICHECK_EQ(mat1.size(0), mat2.size(0)) << "mat1 and mat2 must have the same batch size (" + << mat1.size(0) << " and " << mat2.size(0) << ")"; + TVM_FFI_ICHECK_EQ(mat1.size(2), mat2.size(2) * mat2_k_scale) + << "mat1 and mat2 shapes cannot be multiplied (" << mat1.size(1) << "x" << mat1.size(2) + << " and " << mat2.size(1) << "x" << mat2.size(2) << ")"; + m = mat1.size(1); + n = mat2.size(1); + k = mat2.size(2) * 2; + b = mat1.size(0); + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "mat1 must be a matrix or a batch of matrices"; + } + + // No heuristic for now, we rely on the autotuner to select the best tactic. + if (tactic == -1) { + tactic = 0; + } + auto config = getFp4GemmConfig(m, n, k, tactic); + + constexpr int alignment = 32; + TVM_FFI_ICHECK_EQ(k % alignment, 0) + << "Expected k to be divisible by " << alignment << ", but got mat1 shape: (" << mat1.size(0) + << "x" << mat1.size(1) << "), k: " << k << "."; + TVM_FFI_ICHECK_EQ(n % alignment, 0) + << "Expected n to be divisible by " << alignment << ", but got mat2 shape: (" << mat2.size(0) + << "x" << mat2.size(1) << ")."; + + // Validate out dimensions + std::vector out_shape = + mat1.ndim() == 2 ? std::vector{m, n} : std::vector{b, m, n}; + TVM_FFI_ICHECK_EQ(out.ndim(), out_shape.size()) + << "out must have " << out_shape.size() << " dimensions, but got " << out.ndim(); + for (int i = 0; i < out_shape.size(); ++i) { + TVM_FFI_ICHECK_EQ(out.size(i), out_shape[i]) + << "out shape mismatch at dimension " << i << ": expected " << out_shape[i] << ", got " + << out.size(i); + } + + switch (encode_dlpack_dtype(out.dtype())) { + case float16_code: + runGemm(out, mat1, mat2, mat1Scale, mat2Scale, globalScale, m, n, k, b, config, + workspace_buffer); + break; + case bfloat16_code: + runGemm<__nv_bfloat16>(out, mat1, mat2, mat1Scale, mat2Scale, globalScale, m, n, k, b, config, + workspace_buffer); + break; + default: + TVM_FFI_ICHECK(false) << "out_dtype must be one of fp16/bf16."; + } +} + +} // namespace + +void fp4_gemm(TensorView mat1, TensorView mat2, TensorView mat1Scale, TensorView mat2Scale, + TensorView globalScale, TensorView out, TensorView workspace_buffer, int64_t tactic) { + fp4_bmm_impl(mat1, mat2, mat1Scale, mat2Scale, globalScale, out, workspace_buffer, tactic); +} + +int64_t fp4_gemm_tactic_num() { + auto getCutlassConfigs = []() { + CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4> gemmRunner; + return gemmRunner.getConfigs(); + }; + static int64_t totalTactics = getCutlassConfigs().size(); + return totalTactics; +} + +} // namespace torch_ext + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_gemm, torch_ext::fp4_gemm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_gemm_tactic_num, torch_ext::fp4_gemm_tactic_num); diff --git a/csrc/fp4_gemm_cutlass.jinja b/csrc/fp4_gemm_cutlass.jinja new file mode 100644 index 0000000000..27d6f3f659 --- /dev/null +++ b/csrc/fp4_gemm_cutlass.jinja @@ -0,0 +1,31 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "flashinfer/gemm/fp4_gemm_cutlass_template.h" + +namespace flashinfer { +namespace gemm { +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 1, 1, _1SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 2, 1, _1SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 4, 1, _1SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 1, 1, _2SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 2, 1, _2SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 4, 1, _2SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 2, 1, _2SM) +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 4, 1, _2SM) + +} // namespace gemm +} // namespace flashinfer diff --git a/csrc/fp4_gemm_cutlass_sm120.cu b/csrc/fp4_gemm_cutlass_sm120.cu new file mode 100644 index 0000000000..30080f0fce --- /dev/null +++ b/csrc/fp4_gemm_cutlass_sm120.cu @@ -0,0 +1,187 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include +#include +#include +#include +#include + +#include "flashinfer/gemm/cutlass_gemm_configs.h" +// Use SM120-specific dispatch template (includes fp4_gemm_cutlass.h) +#include "flashinfer/gemm/fp4_gemm_cutlass_template_sm120.h" +#include "tvm_ffi_utils.h" + +using flashinfer::gemm::ClusterShape; +using flashinfer::gemm::CutlassFp4GemmRunner; +using flashinfer::gemm::CutlassGemmConfig; +using flashinfer::gemm::CutlassTileConfigSM120; +using flashinfer::gemm::EpilogueScheduleType; +using flashinfer::gemm::FP4GemmType; +using flashinfer::gemm::MainloopScheduleType; + +namespace torch_ext { + +namespace { + +CutlassGemmConfig getFp4GemmConfig(int64_t m, int64_t n, int64_t k, int64_t tactic) { + auto getCutlassFp4GemmConfigs = []() { + CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4> gemmRunner; + return gemmRunner.getConfigs(); + }; + static std::vector globalConfigs = getCutlassFp4GemmConfigs(); + TVM_FFI_ICHECK(tactic >= 0 && tactic < globalConfigs.size()) + << "tactic must be between 0 and " << globalConfigs.size(); + return globalConfigs[tactic]; +} + +template +void runGemm(TensorView out, TensorView mat1, TensorView mat2, TensorView mat1Scale, + TensorView mat2Scale, TensorView globalScale, int64_t m, int64_t n, int64_t k, + int64_t batch_count, CutlassGemmConfig const& gemmConfig, + TensorView workspace_buffer) { + CutlassFp4GemmRunner gemmRunner; + + int64_t const required_workspace_size = gemmRunner.getWorkspaceSize(m, n, k, batch_count); + int64_t const provided_workspace_size = + workspace_buffer.numel() * get_element_size(workspace_buffer); + + auto runKernel = [&](void* workspace) { + gemmRunner.gemm(out.data_ptr(), mat1.data_ptr(), mat2.data_ptr(), mat1Scale.data_ptr(), + mat2Scale.data_ptr(), static_cast(globalScale.data_ptr()), m, n, k, + batch_count, gemmConfig, reinterpret_cast(workspace), + required_workspace_size, get_stream(mat1.device())); + }; + + if (provided_workspace_size < required_workspace_size) { + Tensor new_workspace = + alloc_tensor({required_workspace_size}, DLDataType{kDLInt, 8, 1}, mat1.device()); + runKernel(new_workspace.data_ptr()); + } else { + runKernel(workspace_buffer.data_ptr()); + } +} + +constexpr auto FLOAT4_E2M1X2 = dl_uint8; // uint8_t +constexpr auto SF_DTYPE = dl_uint8; // uint8_t + +void fp4_bmm_impl(TensorView mat1, TensorView mat2, TensorView mat1Scale, TensorView mat2Scale, + TensorView globalScale, TensorView out, TensorView workspace_buffer, + int64_t tactic) { + // Validate inputs + TVM_FFI_ICHECK_EQ(mat1.dtype(), FLOAT4_E2M1X2) << "mat1 must be FLOAT4_E2M1X2 (uint8)"; + TVM_FFI_ICHECK_EQ(mat2.dtype(), FLOAT4_E2M1X2) << "mat2 must be FLOAT4_E2M1X2 (uint8)"; + TVM_FFI_ICHECK_EQ(mat1Scale.dtype(), SF_DTYPE) << "mat1Scale must be SF_DTYPE (uint8)"; + TVM_FFI_ICHECK_EQ(mat2Scale.dtype(), SF_DTYPE) << "mat2Scale must be SF_DTYPE (uint8)"; + TVM_FFI_ICHECK_EQ(globalScale.dtype(), dl_float32) << "globalScale must be float"; + TVM_FFI_ICHECK_EQ(mat1.device().device_type, kDLCUDA) << "mat1 must be on CUDA device"; + TVM_FFI_ICHECK_EQ(mat2.device().device_type, kDLCUDA) << "mat2 must be on CUDA device"; + TVM_FFI_ICHECK_EQ(mat1Scale.device().device_type, kDLCUDA) << "mat1Scale must be on CUDA device"; + TVM_FFI_ICHECK_EQ(mat2Scale.device().device_type, kDLCUDA) << "mat2Scale must be on CUDA device"; + TVM_FFI_ICHECK_EQ(globalScale.device().device_type, kDLCUDA) + << "globalScale must be on CUDA device"; + TVM_FFI_ICHECK_EQ(out.device().device_type, kDLCUDA) << "out must be on CUDA device"; + TVM_FFI_ICHECK_EQ(workspace_buffer.device().device_type, kDLCUDA) + << "workspace_buffer must be on CUDA device"; + + // Check device consistency + CHECK_DEVICE(mat1, mat2); + CHECK_DEVICE(mat1, mat1Scale); + CHECK_DEVICE(mat1, mat2Scale); + CHECK_DEVICE(mat1, globalScale); + CHECK_DEVICE(mat1, out); + CHECK_DEVICE(mat1, workspace_buffer); + + // Get dimensions + int64_t b = 1; + int64_t m, k_packed, n; + + if (mat1.ndim() == 2) { + m = mat1.size(0); + k_packed = mat1.size(1); + } else if (mat1.ndim() == 3) { + b = mat1.size(0); + m = mat1.size(1); + k_packed = mat1.size(2); + } else { + TVM_FFI_ICHECK(false) << "mat1 must be 2D or 3D tensor"; + } + + if (mat2.ndim() == 2) { + n = mat2.size(0); + TVM_FFI_ICHECK_EQ(mat2.size(1), k_packed) << "mat2.size(1) must match mat1.size(-1)"; + } else if (mat2.ndim() == 3) { + TVM_FFI_ICHECK_EQ(mat2.size(0), b) << "Batch dimensions must match"; + n = mat2.size(1); + TVM_FFI_ICHECK_EQ(mat2.size(2), k_packed) << "mat2.size(2) must match mat1.size(-1)"; + } else { + TVM_FFI_ICHECK(false) << "mat2 must be 2D or 3D tensor"; + } + + // k_packed stores 2 FP4 values per byte + int64_t k = k_packed * 2; + + TVM_FFI_ICHECK_EQ(globalScale.numel(), 1) << "globalScale must be a scalar tensor"; + + // Configure the kernel + CutlassGemmConfig config = + (tactic >= 0) ? getFp4GemmConfig(m, n, k, tactic) + : CutlassGemmConfig(CutlassTileConfigSM120::CtaShape128x128x128B, + MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, + ClusterShape::ClusterShape_1x1x1); + + // Validate output dimensions + std::vector out_shape = + (b > 1) ? std::vector{b, m, n} : std::vector{m, n}; + TVM_FFI_ICHECK_EQ(out.ndim(), out_shape.size()) + << "out must have " << out_shape.size() << " dimensions"; + for (size_t i = 0; i < out_shape.size(); ++i) { + TVM_FFI_ICHECK_EQ(out.size(i), out_shape[i]) + << "out.size(" << i << "): expected " << out_shape[i] << ", got " << out.size(i); + } + + switch (encode_dlpack_dtype(out.dtype())) { + case float16_code: + runGemm(out, mat1, mat2, mat1Scale, mat2Scale, globalScale, m, n, k, b, config, + workspace_buffer); + break; + case bfloat16_code: + runGemm<__nv_bfloat16>(out, mat1, mat2, mat1Scale, mat2Scale, globalScale, m, n, k, b, config, + workspace_buffer); + break; + default: + TVM_FFI_ICHECK(false) << "out_dtype must be one of fp16/bf16."; + } +} + +} // namespace + +void fp4_gemm(TensorView mat1, TensorView mat2, TensorView mat1Scale, TensorView mat2Scale, + TensorView globalScale, TensorView out, TensorView workspace_buffer, int64_t tactic) { + fp4_bmm_impl(mat1, mat2, mat1Scale, mat2Scale, globalScale, out, workspace_buffer, tactic); +} + +int64_t fp4_gemm_tactic_num() { + static const int64_t totalTactics = + CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4>{}.getConfigs().size(); + return totalTactics; +} + +} // namespace torch_ext + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_gemm, torch_ext::fp4_gemm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_gemm_tactic_num, torch_ext::fp4_gemm_tactic_num); diff --git a/csrc/fp4_gemm_cutlass_sm120.jinja b/csrc/fp4_gemm_cutlass_sm120.jinja new file mode 100644 index 0000000000..0a0a9ba947 --- /dev/null +++ b/csrc/fp4_gemm_cutlass_sm120.jinja @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Use SM120-specific dispatch template +#include "flashinfer/gemm/fp4_gemm_cutlass_template_sm120.h" + +namespace flashinfer { +namespace gemm { +// SM120/121 only supports 1x1x1 cluster shape +INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 1, 1, _1SM) + +} // namespace gemm +} // namespace flashinfer diff --git a/csrc/fp8_gemm_cutlass.cu b/csrc/fp8_gemm_cutlass.cu new file mode 100644 index 0000000000..d2e6a63d82 --- /dev/null +++ b/csrc/fp8_gemm_cutlass.cu @@ -0,0 +1,170 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include +#include +#include +#include +#include + +#include "flashinfer/gemm/cutlass_gemm_configs.h" +#include "flashinfer/gemm/fp8_gemm_cutlass.h" +#include "flashinfer/gemm/fp8_gemm_cutlass_template.h" +#include "tvm_ffi_utils.h" + +using flashinfer::gemm::ClusterShape; +using flashinfer::gemm::CutlassFp8GemmRunner; +using flashinfer::gemm::CutlassFp8GemmRunnerInterface; +using flashinfer::gemm::CutlassGemmConfig; +using flashinfer::gemm::CutlassTileConfigSM100; +using flashinfer::gemm::EpilogueScheduleType; +using flashinfer::gemm::MainloopScheduleType; + +namespace flashinfer { +namespace gemm { +template class CutlassFp8GemmRunner<__nv_bfloat16>; +template class CutlassFp8GemmRunner; +} // namespace gemm +} // namespace flashinfer + +namespace torch_ext { + +namespace { + +CutlassGemmConfig getFp8GemmConfig(int64_t m, int64_t n, int64_t k, int64_t tactic) { + auto getCutlassFp8GemmConfigs = []() { + CutlassFp8GemmRunner<__nv_bfloat16> gemmRunner; + return gemmRunner.getConfigs(); + }; + static std::vector globalConfigs = getCutlassFp8GemmConfigs(); + TVM_FFI_ICHECK(tactic >= 0 && tactic < globalConfigs.size()) + << "tactic must be between 0 and " << globalConfigs.size(); + return globalConfigs[tactic]; +} + +template +void runGemm(TensorView out, TensorView mat1, TensorView mat2, TensorView scale_a, + TensorView scale_b, int64_t m, int64_t n, int64_t k, int64_t b, + CutlassGemmConfig const& gemmConfig, TensorView workspace_buffer) { + CutlassFp8GemmRunner gemmRunner; + + int64_t const required_workspace_size = gemmRunner.getWorkspaceSize(m, n, k); + int64_t const provided_workspace_size = + workspace_buffer.numel() * get_element_size(workspace_buffer); + + auto runKernel = [&](void* workspace) { + gemmRunner.gemm( + static_cast<__nv_fp8_e4m3*>(mat1.data_ptr()), static_cast<__nv_fp8_e4m3*>(mat2.data_ptr()), + static_cast(scale_a.data_ptr()), static_cast(scale_b.data_ptr()), + out.data_ptr(), m, n, k, b, gemmConfig, static_cast(workspace), + required_workspace_size, get_stream(mat1.device())); + }; + + if (provided_workspace_size < required_workspace_size) { + Tensor new_workspace = + alloc_tensor({required_workspace_size}, DLDataType{kDLInt, 8, 1}, mat1.device()); + + runKernel(new_workspace.data_ptr()); + } else { + runKernel(workspace_buffer.data_ptr()); + } +} + +void fp8_bmm_impl(TensorView mat1, TensorView mat2, TensorView scale_a, TensorView scale_b, + TensorView out, TensorView workspace_buffer, int64_t tactic) { + CHECK_INPUT(mat1); + CHECK_INPUT(mat2); + CHECK_INPUT(scale_a); + CHECK_INPUT(scale_b); + + int mat2_k_scale = 1; + + int64_t m, n, k, b; + if (mat1.ndim() == 2) { + TVM_FFI_ICHECK_EQ(mat2.ndim(), 2) << "mat2 must be a matrix"; + TVM_FFI_ICHECK_EQ(mat1.size(1), mat2.size(1) * mat2_k_scale) + << "mat1 and mat2 shapes cannot be multiplied (" << mat1.size(0) << "x" << mat1.size(1) + << " and " << mat2.size(0) << "x" << mat2.size(1) << ")"; + m = mat1.size(0); + n = mat2.size(0); + k = mat2.size(1); + b = 1; + } else if (mat1.ndim() == 3) { + TVM_FFI_ICHECK_EQ(mat2.ndim(), 3) << "mat2 must be a batch of matrices"; + TVM_FFI_ICHECK_EQ(mat1.size(0), mat2.size(0)) << "mat1 and mat2 must have the same batch size (" + << mat1.size(0) << " and " << mat2.size(0) << ")"; + TVM_FFI_ICHECK_EQ(mat1.size(2), mat2.size(2) * mat2_k_scale) + << "mat1 and mat2 shapes cannot be multiplied (" << mat1.size(1) << "x" << mat1.size(2) + << " and " << mat2.size(1) << "x" << mat2.size(2) << ")"; + m = mat1.size(1); + n = mat2.size(1); + k = mat2.size(2); + b = mat1.size(0); + } else { + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "mat1 must be a matrix or a batch of matrices"; + } + + // No heuristic for now, we rely on the autotuner to select the best tactic. + if (tactic == -1) { + tactic = 0; + } + auto config = getFp8GemmConfig(m, n, k, tactic); + + // Validate out dimensions + std::vector out_shape = + mat1.ndim() == 2 ? std::vector{m, n} : std::vector{b, m, n}; + TVM_FFI_ICHECK_EQ(out.ndim(), out_shape.size()) + << "out must have " << out_shape.size() << " dimensions, but got " << out.ndim(); + for (int i = 0; i < out_shape.size(); ++i) { + TVM_FFI_ICHECK_EQ(out.size(i), out_shape[i]) + << "out shape mismatch at dimension " << i << ": expected " << out_shape[i] << ", got " + << out.size(i); + } + + switch (encode_dlpack_dtype(out.dtype())) { + case float16_code: + runGemm(out, mat1, mat2, scale_a, scale_b, m, n, k, b, config, workspace_buffer); + break; + case bfloat16_code: + runGemm<__nv_bfloat16>(out, mat1, mat2, scale_a, scale_b, m, n, k, b, config, + workspace_buffer); + break; + default: + TVM_FFI_LOG_AND_THROW(NotImplementedError) << "out_dtype must be one of fp16/bf16."; + } +} + +} // namespace + +void fp8_gemm(TensorView mat1, TensorView mat2, TensorView scale_a, TensorView scale_b, + TensorView out, TensorView workspace_buffer, int64_t tactic) { + fp8_bmm_impl(mat1, mat2, scale_a, scale_b, out, workspace_buffer, tactic); +} + +int64_t fp8_gemm_tactic_num() { + auto getCutlassConfigs = []() { + CutlassFp8GemmRunner<__nv_bfloat16> gemmRunner; + return gemmRunner.getConfigs(); + }; + static int64_t totalTactics = getCutlassConfigs().size(); + return totalTactics; +} + +} // namespace torch_ext + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_gemm, torch_ext::fp8_gemm); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_gemm_tactic_num, torch_ext::fp8_gemm_tactic_num); diff --git a/csrc/fp8_gemm_cutlass.jinja b/csrc/fp8_gemm_cutlass.jinja new file mode 100644 index 0000000000..333d7af525 --- /dev/null +++ b/csrc/fp8_gemm_cutlass.jinja @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "flashinfer/gemm/fp8_gemm_template_sm100.h" + +namespace flashinfer { +namespace gemm { + INSTANCE_FP8_GEMM_TEMPLATE_SM100({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 1, 1, _1SM); + INSTANCE_FP8_GEMM_TEMPLATE_SM100({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 2, 1, _1SM); + INSTANCE_FP8_GEMM_TEMPLATE_SM100({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 4, 1, _1SM); + INSTANCE_FP8_GEMM_TEMPLATE_SM100({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 1, 1, _2SM); + INSTANCE_FP8_GEMM_TEMPLATE_SM100({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 2, 1, _2SM); +} // namespace gemm +} // namespace flashinfer diff --git a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu new file mode 100644 index 0000000000..6469b9a0cd --- /dev/null +++ b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_instantiation.cu @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cutlass_fused_moe_kernels.cuh" +#include "moe_kernels.h" + +namespace tensorrt_llm::kernels::cutlass_kernels { +template class CutlassMoeFCRunner; + +#ifdef ENABLE_BF16 +template class CutlassMoeFCRunner<__nv_bfloat16, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_bfloat16, uint8_t>; +template class CutlassMoeFCRunner<__nv_bfloat16, cutlass::uint4b_t>; +#endif + +template class CutlassMoeFCRunner; +template class CutlassMoeFCRunner; +template class CutlassMoeFCRunner; +#ifdef ENABLE_FP8 +// template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, half>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, half, half>; +#ifdef ENABLE_BF16 +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, __nv_fp8_e4m3>; +#endif +#endif +#ifdef ENABLE_FP4 +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, half>; +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, half, half>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half, half>; +template class CutlassMoeFCRunner; +#ifdef ENABLE_BF16 +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, __nv_bfloat16, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_bfloat16, __nv_fp4_e2m1>; +#endif +#endif + +// Explicit instantiations for finalizeMoeRoutingKernelLauncher to ensure +// symbols are emitted in the JIT library for common data types. +INSTANTIATE_FINALIZE_MOE_ROUTING(half, half, half); +INSTANTIATE_FINALIZE_MOE_ROUTING(float, float, float); +#ifdef ENABLE_BF16 +INSTANTIATE_FINALIZE_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, __nv_bfloat16); +#endif +} // namespace tensorrt_llm::kernels::cutlass_kernels diff --git a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh new file mode 100644 index 0000000000..465241546d --- /dev/null +++ b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh @@ -0,0 +1,4875 @@ +/* + * Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/workspace.h" + +// Ignore CUTLASS warnings about type punning +#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif +#include "cute/tensor.hpp" +#include "cutlass/conv/convolution.h" +// Order matters here, packed_stride.hpp is missing cute and convolution includes +#include "cutlass/array.h" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass_extensions/epilogue/thread/fused_activations.h" + +#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#pragma GCC diagnostic pop +#endif + +#include "moe_kernels.h" +#include "moe_util_kernels.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" +#include "tensorrt_llm/kernels/preQuantScaleKernel.h" +#include "tensorrt_llm/kernels/quantization.cuh" + +#ifndef CUDART_VERSION +#error CUDART_VERSION Undefined! +#elif (CUDART_VERSION >= 11050) +#include +#include + +#include +#else +#include "3rdparty/cub/cub.cuh" +#endif + +using namespace tensorrt_llm::kernels; +using namespace tensorrt_llm::common; + +namespace tensorrt_llm::kernels::cutlass_kernels { +/** + * Takes the input maps and prepares the expanded maps for min latency + * @param num_active_experts_per_node: Number of active experts on current node + * @param experts_to_token_scores: The score of each token for each activated expert. 0 if the + * expert is not chosen by the token. Only the first num_active_experts_per_ rows are valid + * @param active_expert_global_ids: The global expert id for each activated expert + * Only the first num_active_experts_per_ values are valid + * @param expert_first_token_offset: Store the first token offset for each expert + */ +template +__device__ __forceinline__ void initTensor(T* value, int const tid, int const total_num, + T const init_value) { + for (int i = tid; i < total_num; i += BLOCK_SIZE) { + value[i] = init_value; + } +} + +template +__device__ __forceinline__ void setLocalExperts(int* s_local_experts, + T const* token_selected_experts, + int const total_num_experts, int const tid, + int const start_expert, int const end_expert) { + for (int i = tid; i < total_num_experts; i += BLOCK_SIZE) { + int const expert = token_selected_experts[i]; + + // If expert is in the current node, subtract start_expert to shift the range to [0, + // num_experts_per_node) + bool is_valid_expert = expert >= start_expert && expert < end_expert; + if (is_valid_expert) { + int local_expert_id = expert - start_expert; + if (s_local_experts[local_expert_id] == 0) { + s_local_experts[local_expert_id] = + 1; // @TODO: Make sure that we allow duplicated write here + } + } + } + __syncthreads(); +} + +template +__device__ __forceinline__ void prefixSum(T* out, T* in, int const num, int const tid) { + typedef cub::BlockScan BlockScan; + __shared__ typename BlockScan::TempStorage tempStorage; + + T threadData = 0; + if (tid < num) { + threadData = in[tid]; + } + + BlockScan(tempStorage).InclusiveSum(threadData, threadData); + __syncthreads(); + + if (tid < num) { + out[tid] = threadData; + } + __syncthreads(); +} + +__device__ __forceinline__ void setActiveNum(int& num_active, int& num_active_offset_start, + int& num_active_offset_end, int const cluster_size, + int const cluster_rank) { + int num_remainder = num_active % cluster_size; + int num_active_per_node = + max(0, num_active - 1) / cluster_size; // num_active_per_node shouldn't be neg + if (cluster_rank < num_remainder) { + num_active = num_active_per_node + 1; + num_active_offset_start = cluster_rank * num_active; + } else { + num_active = num_active_per_node; + num_active_offset_start = cluster_rank * num_active_per_node + num_remainder; + } + num_active_offset_end = num_active_offset_start + num_active; +} + +template +__global__ void buildMinLatencyActiveExpertMapsKernel( + int* num_active_experts_per_node, float* experts_to_token_scores, int* active_expert_global_ids, + int64_t* expert_first_token_offset, int const* token_selected_experts, + float const* token_final_scales, int64_t const num_tokens, int const num_experts_per_token, + int const start_expert, int const end_expert, int const num_experts_per_node, + bool const smart_routing, int const cluster_rank, int const cluster_size, + int const num_experts_smem) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + // Use one block to process the min latency case + int tid = threadIdx.x; + // 0. init the global memory experts_to_token_scores [num_experts_per_node, num_token] + int const total_local_scales = num_experts_per_node * num_tokens; + initTensor(experts_to_token_scores, tid, total_local_scales, 0.0f); + initTensor(active_expert_global_ids, tid, num_experts_per_node, -1); + + __threadfence(); //@Todo: check do I need this fence for previous zero setting + + // 1. mask for the active expert: 1 stands for active + extern __shared__ int s_local_experts[]; + int* s_store_experts = s_local_experts + num_experts_smem; + initTensor(s_local_experts, tid, num_experts_smem, 0); + __syncthreads(); + + // 2. set the shared array s_local_experts[] + int const total_num_experts = num_tokens * num_experts_per_token; + setLocalExperts(s_local_experts, token_selected_experts, total_num_experts, tid, + start_expert, end_expert); + + // 3. perform prefix sum to acquire the store position and total active experts + //@TODO: Use cub first, might need to change it to self-defined api + prefixSum(s_store_experts, s_local_experts, num_experts_smem, tid); + + // 4. store the num of active experts + int num_active = s_store_experts[num_experts_smem - 1]; + int num_active_offset_start = 0; + int num_active_offset_end = 0; + + if (smart_routing) { + setActiveNum(num_active, num_active_offset_start, num_active_offset_end, cluster_size, + cluster_rank); + } + + if (tid == 0) { + *num_active_experts_per_node = num_active; + } + + // 5. store the global expert id for each expert + if (smart_routing) { + for (int i = tid; i < num_experts_smem; i += BLOCK_SIZE) { + if (s_local_experts[i]) { + int offset = s_store_experts[i] - 1; + if (offset >= num_active_offset_start && offset < num_active_offset_end) { + active_expert_global_ids[offset - num_active_offset_start] = i; + } else { + s_local_experts[i] = 0; + } + } + } + __syncthreads(); // Need sync to update the s_local_experts + } else { + for (int i = tid; i < num_experts_smem; i += BLOCK_SIZE) { + if (s_local_experts[i]) { + int offset = s_store_experts[i] - 1; + active_expert_global_ids[offset] = i + start_expert; + } + } + } + + // 6. store the scale values + __threadfence(); //@Todo: check do I need this fence for previous zero setting + for (int i = tid; i < total_num_experts; i += BLOCK_SIZE) { + int const expert = token_selected_experts[i]; + + // If expert is not in the current node, set it to num_experts_per_node + // If expert is in the current node, subtract start_expert to shift the range to [0, + // num_experts_per_node) + bool is_valid_expert = + smart_routing ? s_local_experts[expert] : (expert >= start_expert && expert < end_expert); + + if (is_valid_expert) { + int token = i / num_experts_per_token; + float const scale = token_final_scales[i]; + int offset = s_store_experts[expert - start_expert] - 1 - num_active_offset_start; + experts_to_token_scores[offset * num_tokens + token] = scale; + } + } + // 7. set default value for redundant memory + for (int i_exp = num_active + tid; i_exp < num_experts_per_node; i_exp += BLOCK_SIZE) { + active_expert_global_ids[i_exp] = -1; + } + // 8. set expert_first_token_offset + for (int i_exp = tid; i_exp < num_experts_per_node + 1; i_exp += BLOCK_SIZE) { + if (i_exp < num_active) { + expert_first_token_offset[i_exp] = i_exp * num_tokens; + } else { + expert_first_token_offset[i_exp] = num_active * num_tokens; + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +void buildMinLatencyActiveExpertMaps( + int* num_active_experts_per_node, float* experts_to_token_scores, int* active_expert_global_ids, + int64_t* expert_first_token_offset, int const* token_selected_experts, + float const* token_final_scales, int64_t const num_tokens, int const experts_per_token, + int const start_expert, int const end_expert, int const num_experts_per_node, + int const cluster_rank, int const cluster_size, int const num_experts_smem, bool enable_pdl, + cudaStream_t const stream) { + TLLM_CHECK_WITH_INFO(num_experts_per_node == (end_expert - start_expert), + "num_experts_per_node must be equal to end_expert - start_expert"); + + TLLM_CHECK_WITH_INFO(num_experts_per_node <= 256, + "don't support num_experts_per_node > 256 cases"); + + int const threads = 256; + int const blocks = 1; + bool const smart_routing = cluster_size > 1; + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = num_experts_smem * sizeof(int) * 2; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, buildMinLatencyActiveExpertMapsKernel, + num_active_experts_per_node, experts_to_token_scores, active_expert_global_ids, + expert_first_token_offset, token_selected_experts, token_final_scales, + num_tokens, experts_per_token, start_expert, end_expert, num_experts_per_node, + smart_routing, cluster_rank, cluster_size, num_experts_smem); +} + +template +__global__ void fusedBuildExpertMapsSortFirstTokenKernel( + int const* const token_selected_experts, int* const permuted_row_to_unpermuted_row, + int* const unpermuted_row_to_permuted_row, int64_t* const expert_first_token_offset, + int64_t const num_tokens, int const experts_per_token, int const start_expert, + int const end_expert, int const num_experts_per_node) { + // Only using block wise collective so we can only have one block + assert(gridDim.x == 1); + + assert(start_expert <= end_expert); + assert(num_experts_per_node == (end_expert - start_expert)); + assert(num_experts_per_node <= (1 << LOG2_NUM_EXPERTS)); + + int const token = blockIdx.x * BLOCK_SIZE + threadIdx.x; + + bool is_valid_token = token < num_tokens; + + // This is the masked expert id for this token + int local_token_selected_experts[EXPERTS_PER_TOKEN]; + // This is the final permuted rank of this token (ranked by selected expert) + int local_token_permuted_indices[EXPERTS_PER_TOKEN]; + + // Wait PDL before reading token_selected_experts +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +// build expert map +// we need to populate expert ids for all threads, even if there are +// fewer tokens +#pragma unroll + for (int i = 0; i < EXPERTS_PER_TOKEN; i++) { + int const expert = is_valid_token ? token_selected_experts[token * EXPERTS_PER_TOKEN + i] + : num_experts_per_node; + + // If the token is not valid, set the expert id to num_experts_per_node + 1 + // If expert is not in the current node, set it to num_experts_per_node + // If expert is in the current node, subtract start_expert to shift the range to [0, + // num_experts_per_node) + bool is_valid_expert = expert >= start_expert && expert < end_expert; + local_token_selected_experts[i] = !is_valid_token ? num_experts_per_node + 1 + : is_valid_expert ? (expert - start_expert) + : num_experts_per_node; + } + + // TODO: decompose cub's sort to expose the bucket starts, and just return + // that to elide the binary search + + // sort the expert map + using BlockRadixRank = cub::BlockRadixRank; + extern __shared__ unsigned char temp_storage[]; + auto& sort_temp = *reinterpret_cast(temp_storage); + + // Sanity check that the number of bins do correspond to the number of experts + static_assert(BlockRadixRank::BINS_TRACKED_PER_THREAD * BLOCK_SIZE >= (1 << LOG2_NUM_EXPERTS)); + assert(BlockRadixRank::BINS_TRACKED_PER_THREAD * BLOCK_SIZE >= num_experts_per_node); + + int local_expert_first_token_offset[BlockRadixRank::BINS_TRACKED_PER_THREAD]; + + cub::BFEDigitExtractor extractor(0, LOG2_NUM_EXPERTS); + BlockRadixRank(sort_temp).RankKeys(local_token_selected_experts, local_token_permuted_indices, + extractor, local_expert_first_token_offset); + +// We are done with compute, launch the dependent kernels while the stores are in flight +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + + // write to shared memory and global memory + if (is_valid_token) { +#pragma unroll + for (int i = 0; i < EXPERTS_PER_TOKEN; i++) { + int const unpermuted_row = i * num_tokens + token; + int const permuted_row = local_token_permuted_indices[i]; + permuted_row_to_unpermuted_row[permuted_row] = unpermuted_row; + unpermuted_row_to_permuted_row[unpermuted_row] = permuted_row; + } + } + +#pragma unroll + for (int expert_id = 0; expert_id < BlockRadixRank::BINS_TRACKED_PER_THREAD; expert_id++) { + int out_expert_id = expert_id + token * BlockRadixRank::BINS_TRACKED_PER_THREAD; + if (out_expert_id < num_experts_per_node + 1) { + expert_first_token_offset[out_expert_id] = local_expert_first_token_offset[expert_id]; + } + } +} + +template +bool fusedBuildExpertMapsSortFirstTokenDispatch( + int const* token_selected_experts, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, + int const start_expert, int const end_expert, bool enable_pdl, cudaStream_t stream) { + TLLM_CHECK_WITH_INFO(num_experts_per_node == (end_expert - start_expert), + "num_experts_per_node must be equal to end_expert - start_expert"); + int const threads = BLOCK_SIZE; + int const blocks = (num_tokens + threads - 1) / threads; + TLLM_CHECK_WITH_INFO(blocks == 1, "Current implementation requires single block"); + + using BlockRadixRank = cub::BlockRadixRank; + size_t shared_size = sizeof(typename BlockRadixRank::TempStorage); + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = shared_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + auto kernel = + &fusedBuildExpertMapsSortFirstTokenKernel; + + int device = 0; + int max_smem_per_block = 0; + check_cuda_error(cudaGetDevice(&device)); + check_cuda_error( + cudaDeviceGetAttribute(&max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + if (shared_size >= static_cast(max_smem_per_block)) { + // This should mean that + // cudaFuncSetAttribute(cutlass::Kernel, + // cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) wouldn't work. + return false; + } + + check_cuda_error( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_size)); + check_cuda_error( + cudaLaunchKernelEx(&config, kernel, token_selected_experts, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, expert_first_token_offset, num_tokens, + experts_per_token, start_expert, end_expert, num_experts_per_node)); + + return true; +} + +template +bool fusedBuildExpertMapsSortFirstTokenBlockSize( + int const* token_selected_experts, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, + int const start_expert, int const end_expert, bool enable_pdl, cudaStream_t stream) { + int const block_size = num_tokens; + if (num_tokens > 256) { + TLLM_LOG_TRACE( + "Number of tokens %d is greater than 256, which is not supported for fused moe prologues", + num_tokens); + return false; + } + + auto func = &fusedBuildExpertMapsSortFirstTokenDispatch<32, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + if (block_size > 32 && block_size <= 64) { + func = &fusedBuildExpertMapsSortFirstTokenDispatch<64, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + } else if (block_size > 64 && block_size <= 128) { + func = &fusedBuildExpertMapsSortFirstTokenDispatch<128, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + } else if (block_size > 128 && block_size <= 256) { + func = &fusedBuildExpertMapsSortFirstTokenDispatch<256, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + } + + return func(token_selected_experts, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, expert_first_token_offset, num_tokens, + num_experts_per_node, experts_per_token, start_expert, end_expert, enable_pdl, + stream); +} + +template +bool fusedBuildExpertMapsSortFirstTokenBlockSize( + int const* token_selected_experts, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, + int const start_expert, int const end_expert, bool enable_pdl, cudaStream_t stream) { + auto func = &fusedBuildExpertMapsSortFirstTokenBlockSize<1, LOG2_NUM_EXPERTS>; + switch (experts_per_token) { + case 1: { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<1, LOG2_NUM_EXPERTS>; + break; + } + case 2: { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<2, LOG2_NUM_EXPERTS>; + break; + } + case 4: { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<4, LOG2_NUM_EXPERTS>; + break; + } + case 6: { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<6, LOG2_NUM_EXPERTS>; + break; + } + case 8: { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<8, LOG2_NUM_EXPERTS>; + break; + } + default: { + TLLM_LOG_TRACE("Top-K value %d does not have supported fused moe prologues", + experts_per_token); + return false; + } + } + return func(token_selected_experts, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, expert_first_token_offset, num_tokens, + num_experts_per_node, experts_per_token, start_expert, end_expert, enable_pdl, + stream); +} + +bool fusedBuildExpertMapsSortFirstToken( + int const* token_selected_experts, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, + int const start_expert, int const end_expert, bool enable_pdl, cudaStream_t stream) { + // We need enough bits to represent [0, num_experts_per_node+1] (inclusive) i.e. + // num_experts_per_node + 2 values This is floor(log2(num_experts_per_node+1)) + 1 + int expert_log = static_cast(log2(num_experts_per_node + 1)) + 1; + if (expert_log <= 9) { + auto funcs = std::array{&fusedBuildExpertMapsSortFirstTokenBlockSize<1>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<2>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<3>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<4>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<5>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<6>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<7>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<8>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<9>}; + + return funcs[expert_log - 1](token_selected_experts, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, expert_first_token_offset, + num_tokens, num_experts_per_node, experts_per_token, start_expert, + end_expert, enable_pdl, stream); + } + TLLM_LOG_TRACE("Experts per node %d does not have supported fused moe prologues", + num_experts_per_node); + return false; +} + +int64_t computeNumTokensPerBlock(int64_t const num_tokens, int64_t const num_experts_per_node) { + for (int64_t num_tokens_per_block = 32; num_tokens_per_block <= 1024; num_tokens_per_block *= 2) { + int64_t const num_blocks_per_seq = + tensorrt_llm::common::ceilDiv(num_tokens, num_tokens_per_block); + if (num_blocks_per_seq * num_experts_per_node <= num_tokens_per_block) { + return num_tokens_per_block; + } + } + return 1024; +} + +template +__global__ void blockExpertPrefixSumKernel(int const* token_selected_experts, + int* blocked_expert_counts, + int* blocked_row_to_unpermuted_row, + int64_t const num_tokens, + int64_t const num_experts_per_token, + int const start_expert_id) { + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + + // target_expert_id and expert_id are offset by start_expert_id + int const target_expert_id = blockIdx.x; + int const block_id = blockIdx.y; + int const num_blocks_per_seq = gridDim.y; + int const token_id = block_id * kNumTokensPerBlock + threadIdx.x; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + int expanded_token_id = -1; + if (token_id < num_tokens) { + for (int i = 0; i < num_experts_per_token; i++) { + // TODO(enweiz): Fix uncoalesced access with shared memory. + int const expert_id = + token_selected_experts[token_id * num_experts_per_token + i] - start_expert_id; + if (expert_id == target_expert_id) { + expanded_token_id = i * num_tokens + token_id; + break; + } + } + } + + int const has_matched = expanded_token_id >= 0 ? 1 : 0; + int index; + BlockScan(temp_storage).ExclusiveSum(has_matched, index); + + if (has_matched) { + blocked_row_to_unpermuted_row[target_expert_id * num_tokens + block_id * kNumTokensPerBlock + + index] = expanded_token_id; + } + if (threadIdx.x == kNumTokensPerBlock - 1) { + blocked_expert_counts[target_expert_id * num_blocks_per_seq + block_id] = index + has_matched; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +void blockExpertPrefixSum(int const* token_selected_experts, int* blocked_expert_counts, + int* blocked_row_to_unpermuted_row, int64_t const num_tokens, + int64_t const num_experts_per_node, int64_t const num_experts_per_token, + int64_t const num_tokens_per_block, int64_t const num_blocks_per_seq, + int const start_expert_id, bool enable_pdl, cudaStream_t stream) { + dim3 const blocks(num_experts_per_node, num_blocks_per_seq); + dim3 const threads(num_tokens_per_block); + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + auto func = blockExpertPrefixSumKernel<1024>; + if (num_tokens_per_block <= 32) { + func = blockExpertPrefixSumKernel<32>; + } else if (num_tokens_per_block <= 64) { + func = blockExpertPrefixSumKernel<64>; + } else if (num_tokens_per_block <= 128) { + func = blockExpertPrefixSumKernel<128>; + } else if (num_tokens_per_block <= 256) { + func = blockExpertPrefixSumKernel<256>; + } else if (num_tokens_per_block <= 512) { + func = blockExpertPrefixSumKernel<512>; + } + cudaLaunchKernelEx(&config, func, token_selected_experts, blocked_expert_counts, + blocked_row_to_unpermuted_row, num_tokens, num_experts_per_token, + start_expert_id); +} + +template +__global__ void globalExpertPrefixSumLargeKernel(int const* blocked_expert_counts, + int* blocked_expert_counts_cumsum, + int64_t* expert_first_token_offset, + int64_t const num_experts_per_node, + int64_t const num_blocks_per_seq, + int64_t const num_elem_per_thread) { + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + + int offset = threadIdx.x * num_elem_per_thread; + int cnt = 0; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + // Note: Because of limited registers, cannot store thread-level prefix sum or enable #pragma + // unroll + for (int i = 0; i < num_elem_per_thread; i++) { + // TODO(enweiz): Fix uncoalesced access with shared memory. + if (offset + i < num_experts_per_node * num_blocks_per_seq) { + cnt += blocked_expert_counts[offset + i]; + } + } + + int cumsum; + BlockScan(temp_storage).ExclusiveSum(cnt, cumsum); + + for (int i = 0; i < num_elem_per_thread; i++) { + if (offset + i < num_experts_per_node * num_blocks_per_seq) { + blocked_expert_counts_cumsum[offset + i] = cumsum; + if ((offset + i) % num_blocks_per_seq == 0) { + expert_first_token_offset[(offset + i) / num_blocks_per_seq] = cumsum; + } + cumsum += blocked_expert_counts[offset + i]; + if ((offset + i) == num_experts_per_node * num_blocks_per_seq - 1) { + expert_first_token_offset[num_experts_per_node] = cumsum; + } + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +__global__ void globalExpertPrefixSumKernel(int const* blocked_expert_counts, + int* blocked_expert_counts_cumsum, + int64_t* expert_first_token_offset, + int64_t const num_experts_per_node, + int64_t const num_blocks_per_seq) { + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + int const cnt = threadIdx.x < num_experts_per_node * num_blocks_per_seq + ? blocked_expert_counts[threadIdx.x] + : 0; + int cumsum; + BlockScan(temp_storage).ExclusiveSum(cnt, cumsum); + + if (threadIdx.x < num_experts_per_node * num_blocks_per_seq) { + blocked_expert_counts_cumsum[threadIdx.x] = cumsum; + if (threadIdx.x % num_blocks_per_seq == 0) { + expert_first_token_offset[threadIdx.x / num_blocks_per_seq] = cumsum; + } + if (threadIdx.x == num_experts_per_node * num_blocks_per_seq - 1) { + expert_first_token_offset[num_experts_per_node] = cumsum + cnt; + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +void globalExpertPrefixSum(int const* blocked_expert_counts, int* blocked_expert_counts_cumsum, + int64_t* expert_first_token_offset, int64_t const num_experts_per_node, + int64_t const num_tokens_per_block, int64_t const num_blocks_per_seq, + bool enable_pdl, cudaStream_t stream) { + int64_t const num_elements = num_experts_per_node * num_blocks_per_seq; + + cudaLaunchConfig_t config; + config.gridDim = 1; + config.blockDim = 1024; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + if (num_elements <= 1024) { + auto func = globalExpertPrefixSumKernel<1024>; + if (num_elements <= 32) { + func = globalExpertPrefixSumKernel<32>; + config.blockDim = 32; + } else if (num_elements <= 64) { + func = globalExpertPrefixSumKernel<64>; + config.blockDim = 64; + } else if (num_elements <= 128) { + func = globalExpertPrefixSumKernel<128>; + config.blockDim = 128; + } else if (num_elements <= 256) { + func = globalExpertPrefixSumKernel<256>; + config.blockDim = 256; + } else if (num_elements <= 512) { + func = globalExpertPrefixSumKernel<512>; + config.blockDim = 512; + } + cudaLaunchKernelEx(&config, func, blocked_expert_counts, blocked_expert_counts_cumsum, + expert_first_token_offset, num_experts_per_node, num_blocks_per_seq); + } else { + auto func = globalExpertPrefixSumLargeKernel<1024>; + int64_t const num_elem_per_thread = tensorrt_llm::common::ceilDiv(num_elements, 1024); + cudaLaunchKernelEx(&config, func, blocked_expert_counts, blocked_expert_counts_cumsum, + expert_first_token_offset, num_experts_per_node, num_blocks_per_seq, + num_elem_per_thread); + } +} + +__global__ void mergeExpertPrefixSumKernel(int const* blocked_expert_counts, + int const* blocked_expert_counts_cumsum, + int const* blocked_row_to_unpermuted_row, + int* permuted_token_selected_experts, + int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, + int const num_tokens) { + int const target_expert_id = blockIdx.x; + int const block_id = blockIdx.y; + int const num_blocks_per_seq = gridDim.y; + int const token_id = block_id * blockDim.x + threadIdx.x; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + int const cnt = blocked_expert_counts[target_expert_id * num_blocks_per_seq + block_id]; + int const offset = blocked_expert_counts_cumsum[target_expert_id * num_blocks_per_seq + block_id]; + if (threadIdx.x < cnt) { + int const unpermuted_row = + blocked_row_to_unpermuted_row[target_expert_id * num_tokens + token_id]; + int const permuted_row = offset + threadIdx.x; + permuted_row_to_unpermuted_row[permuted_row] = unpermuted_row; + permuted_token_selected_experts[permuted_row] = target_expert_id; + unpermuted_row_to_permuted_row[unpermuted_row] = permuted_row; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +void mergeExpertPrefixSum(int const* blocked_expert_counts, int const* blocked_expert_counts_cumsum, + int const* blocked_row_to_unpermuted_row, + int* permuted_token_selected_experts, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int64_t const num_tokens, + int64_t const num_experts_per_node, int64_t const num_tokens_per_block, + int64_t const num_blocks_per_seq, bool enable_pdl, cudaStream_t stream) { + dim3 const blocks(num_experts_per_node, num_blocks_per_seq); + dim3 const threads(num_tokens_per_block); + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, mergeExpertPrefixSumKernel, blocked_expert_counts, + blocked_expert_counts_cumsum, blocked_row_to_unpermuted_row, + permuted_token_selected_experts, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, num_tokens); +} + +// threeStepBuildExpertMapsSortFirstToken uses three kernels to achieve the sort of +// token_selected_experts + +// 1. blockExpertPrefixSumKernel launches [num_experts_per_node, num_blocks_per_seq] CTAs; each CTA +// has num_tokens_per_block threads. blocked_row_to_unpermuted_row points to a 2D buffer of size +// [num_experts_per_node, num_tokens], which can be viewed as [num_experts_per_node, +// num_blocks_per_seq] blocks, and each block has num_tokens_per_block tokens. Note that each CTA +// corresponds to a block in blocked_row_to_unpermuted_row. Within each CTA, the threads leverage +// cub::BlockScan to compute the offsets of tokens that activate the target expert. If a thread's +// token activates the target expert, the thread stores its unpermuted_row to the buffer block with +// the offset. In addition, the kernel also stores the expert counts for each block to another 2D +// buffer blocked_expert_counts of size [num_experts_per_node, num_blocks_per_seq]. + +// 2. globalExpertPrefixSumKernel launches 1 CTA; that CTA has num_experts_per_node * +// num_blocks_per_seq threads. The kernel views blocked_expert_counts as a 1D buffer, and leverages +// cub::BlockScan to compute the prefix sum of the expert counts for each block. The prefix sum is +// stored to blocked_expert_counts_cumsum. + +// 3. mergeExpertPrefixSumKernel launches [num_experts_per_node, num_blocks_per_seq] CTAs; each CTA +// has num_tokens_per_block threads. Each CTA obtains the block-level offset from +// blocked_expert_counts_cumsum, and thus compacts blocked_row_to_unpermuted_row to +// permuted_row_to_unpermuted_row. In addition, with the block-level offsets, the kernel fills +// permuted_token_selected_experts. + +// computeNumTokensPerBlock decides num_tokens_per_block. Note that both blockExpertPrefixSumKernel +// and globalExpertPrefixSumKernel leverage cub::BlockScan, and their CTA sizes are +// num_tokens_per_block and num_experts_per_node * num_blocks_per_seq, respectively. +// computeNumTokensPerBlock tries to find a minimum CTA size for both kernels, so that the +// block-leval cub::BlockScan can be efficient. + +void threeStepBuildExpertMapsSortFirstToken( + int const* token_selected_experts, int* permuted_token_selected_experts, + int* permuted_row_to_unpermuted_row, int* unpermuted_row_to_permuted_row, + int64_t* expert_first_token_offset, int* blocked_expert_counts, + int* blocked_expert_counts_cumsum, int* blocked_row_to_unpermuted_row, int64_t const num_tokens, + int64_t const num_experts_per_node, int64_t const num_experts_per_token, + int const start_expert_id, bool enable_pdl, cudaStream_t stream) { + int64_t const num_tokens_per_block = computeNumTokensPerBlock(num_tokens, num_experts_per_node); + int64_t const num_blocks_per_seq = + tensorrt_llm::common::ceilDiv(num_tokens, num_tokens_per_block); + + blockExpertPrefixSum(token_selected_experts, blocked_expert_counts, blocked_row_to_unpermuted_row, + num_tokens, num_experts_per_node, num_experts_per_token, + num_tokens_per_block, num_blocks_per_seq, start_expert_id, enable_pdl, + stream); + sync_check_cuda_error(stream); + + globalExpertPrefixSum(blocked_expert_counts, blocked_expert_counts_cumsum, + expert_first_token_offset, num_experts_per_node, num_tokens_per_block, + num_blocks_per_seq, enable_pdl, stream); + sync_check_cuda_error(stream); + + mergeExpertPrefixSum(blocked_expert_counts, blocked_expert_counts_cumsum, + blocked_row_to_unpermuted_row, permuted_token_selected_experts, + permuted_row_to_unpermuted_row, unpermuted_row_to_permuted_row, num_tokens, + num_experts_per_node, num_tokens_per_block, num_blocks_per_seq, enable_pdl, + stream); +} + +// ============================== Infer GEMM sizes ================================= +// TODO Could linear search be better for small # experts +template +__device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices, + int64_t const arr_length, T const target) { + int64_t low = 0, high = arr_length - 1, target_location = -1; + while (low <= high) { + int64_t mid = (low + high) / 2; + + if (sorted_indices[mid] >= target) { + high = mid - 1; + } else { + low = mid + 1; + target_location = mid; + } + } + return target_location + 1; +} + +template +using sizeof_bits = cutlass::sizeof_bits< + typename cutlass_kernels::TllmToCutlassTypeAdapter>::type>; + +// Function to safely offset an pointer that may contain sub-byte types (FP4/INT4) +template +__host__ __device__ constexpr T* safe_inc_ptr(T* ptr, size_t offset) { + constexpr int adjustment = (sizeof_bits::value < 8) ? (8 / sizeof_bits::value) : 1; + assert(offset % adjustment == 0 && "Attempt to offset index to sub-byte"); + return ptr + offset / adjustment; +} + +__host__ __device__ constexpr int64_t getOffsetWeightSF( + int64_t expert_id, int64_t gemm_n, int64_t gemm_k, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType scaling_type) { + auto function = [=](int64_t min_n_dim_alignment, int64_t min_k_dim_alignment, + int64_t block_size) { + int64_t padded_gemm_n = + TmaWarpSpecializedGroupedGemmInput::alignToSfDim(gemm_n, min_n_dim_alignment); + int64_t padded_gemm_k = + TmaWarpSpecializedGroupedGemmInput::alignToSfDim(gemm_k, min_k_dim_alignment); + assert(gemm_k % block_size == 0); + return expert_id * padded_gemm_n * padded_gemm_k / block_size; + }; + switch (scaling_type) { + case TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX: + return function(TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX, + TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX, + TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize); + case TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4: + return function(TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4, + TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4, + TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize); + case TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE: + return 0; // No scaling factors, no offset + } + + assert(false && "Unrecognized scaling type"); + return 0; +} + +__host__ __device__ constexpr int64_t getOffsetActivationSF( + int64_t expert_id, int64_t token_offset, int64_t gemm_k, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType scaling_type) { + auto function = [=](int64_t min_n_dim_alignment, int64_t min_k_dim_alignment, + int64_t block_size) { + // This formulation ensures that: + // `sf_offset[i + 1] - sf_offset[i] >= padded(token_offset[i + 1] - token_offset[i])` + // is true for all possible token distributions. + int64_t padded_sf_start_offset = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( + token_offset + expert_id * (min_n_dim_alignment - 1), min_n_dim_alignment); + int64_t padded_gemm_k = + TmaWarpSpecializedGroupedGemmInput::alignToSfDim(gemm_k, min_k_dim_alignment); + assert(gemm_k % block_size == 0); + assert(padded_gemm_k % block_size == 0); + return padded_sf_start_offset * padded_gemm_k / block_size; + }; + switch (scaling_type) { + case TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX: + return function(TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX, + TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX, + TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize); + case TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4: + return function(TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4, + TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4, + TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize); + case TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE: + return 0; // No scaling factors, no offset + } + + assert(false && "Unrecognized scaling type"); + return 0; +} + +template +__device__ auto quantizePackedFPXValue( + ComputeElem& post_act_val, float global_scale_val, int64_t num_tokens_before_expert, + int64_t expert_id, int64_t token_id, int64_t elem_idx, int64_t num_cols, + TmaWarpSpecializedGroupedGemmInput::ElementSF* act_sf_flat, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType scaling_type) { + constexpr bool is_fp8 = std::is_same_v; + static constexpr int NumThreadsPerSF = VecSize / CVT_ELTS_PER_THREAD; + // Quantize the input to FP4 + static_assert(std::is_same_v || + std::is_same_v); + static_assert(ComputeElem::kElements == CVT_ELTS_PER_THREAD); + PackedVec packed_vec{}; + for (int i = 0; i < CVT_ELTS_PER_THREAD / 2; i++) { + packed_vec.elts[i].x = static_cast(post_act_val[i * 2 + 0]); + packed_vec.elts[i].y = static_cast(post_act_val[i * 2 + 1]); + } + + // We need to offset into the scaling factors for just this expert + auto act_sf_expert = act_sf_flat + getOffsetActivationSF(expert_id, num_tokens_before_expert, + num_cols, scaling_type); + + // Use `token - num_tokens_before_expert` because we want this to be relative to the start of this + // expert + auto sf_out = + cvt_quant_get_sf_out_offset( + std::nullopt /* batchIdx */, token_id - num_tokens_before_expert, elem_idx, + std::nullopt /* numRows */, num_cols / VecSize, act_sf_expert, + QuantizationSFLayout::SWIZZLED_128x4); + + // Do the conversion and set the output and scaling factor + auto func = [&]() { + if constexpr (is_fp8) { + return [](PackedVec& vec, float /* ignored */, uint8_t* SFout) -> uint64_t { + static_assert(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize == VecSize); + return cvt_warp_fp16_to_mxfp8(vec, SFout); + }; + } else { + return (scaling_type == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4) + ? &cvt_warp_fp16_to_fp4 + : &cvt_warp_fp16_to_fp4; + } + }(); + + return func(packed_vec, global_scale_val, sf_out); +} + +template +__device__ void writeSF(int64_t num_tokens_before_expert, int64_t expert_id, + int64_t source_token_id, int64_t token_id, int64_t elem_idx, + int64_t num_cols, + TmaWarpSpecializedGroupedGemmInput::ElementSF* act_sf_flat, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, + bool const swizzled_input_sf = true) { + static constexpr int NumThreadsPerSF = VecSize / ElementsPerThread; + + // We need to offset into the scaling factors for just this expert + auto act_sf_expert = + act_sf_flat + getOffsetActivationSF( + expert_id, num_tokens_before_expert, num_cols, + (VecSize == TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize) + ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 + : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX); + + // Use `token - num_tokens_before_expert` because we want this to be relative to the start of this + // expert + auto sf_out = + cvt_quant_get_sf_out_offset( + std::nullopt /* batchIdx */, token_id - num_tokens_before_expert, elem_idx, + std::nullopt /* numRows */, num_cols / VecSize, act_sf_expert, + QuantizationSFLayout::SWIZZLED_128x4); + if (sf_out) { + if (input_sf) { + if (swizzled_input_sf) { + auto const sf_in = + cvt_quant_get_sf_out_offset( + std::nullopt /* batchIdx */, source_token_id, elem_idx, std::nullopt /* numRows */, + num_cols / VecSize, + const_cast(input_sf), + QuantizationSFLayout::SWIZZLED_128x4); + *sf_out = *sf_in; + } else { + auto const sf_in = + cvt_quant_get_sf_out_offset( + std::nullopt /* batchIdx */, source_token_id, elem_idx, std::nullopt /* numRows */, + num_cols / VecSize, + const_cast(input_sf), + QuantizationSFLayout::LINEAR); + *sf_out = *sf_in; + } + } else { + *sf_out = 0x00; + } + } +} + +// ====================== Compute FP8 dequant scale only =============================== +__global__ void computeFP8DequantScaleKernel(float const** alpha_scale_ptr_array, + int64_t const num_experts_per_node, + float const* fp8_dequant) { + // First, compute the global tid. We only need 1 thread per expert. + int const expert = blockIdx.x * blockDim.x + threadIdx.x; + if (expert >= num_experts_per_node) { + return; + } + + assert(fp8_dequant != nullptr); + alpha_scale_ptr_array[expert] = fp8_dequant + expert; +} + +float const** computeFP8DequantScale(float const** alpha_scale_ptr_array, + int const num_experts_per_node, float const* fp8_dequant, + cudaStream_t stream) { + if (!fp8_dequant) { + return nullptr; + } + + int const threads = std::min(1024, num_experts_per_node); + int const blocks = (num_experts_per_node + threads - 1) / threads; + + computeFP8DequantScaleKernel<<>>(alpha_scale_ptr_array, + num_experts_per_node, fp8_dequant); + + return alpha_scale_ptr_array; +} + +template +__device__ void setupFP4BlockScalingFactors( + TmaWarpSpecializedGroupedGemmInput& layout_info, int expert, int gemm_m, int gemm_n, int gemm_k, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* fp4_act_flat, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* weight_block_scale, + int64_t num_tokens_before_expert) { + assert(layout_info.fpX_block_scaling_factors_stride_act); + assert(layout_info.fpX_block_scaling_factors_stride_weight); + + auto stride_act_ptr = reinterpret_cast( + layout_info.fpX_block_scaling_factors_stride_act); + auto stride_weight_ptr = reinterpret_cast( + layout_info.fpX_block_scaling_factors_stride_weight); + if (layout_info.swap_ab) { + // M & N swapped for transpose + stride_act_ptr[expert] = BSConfig::tile_atom_to_shape_SFB( + cute::make_shape((int)gemm_n, (int)gemm_m, (int)gemm_k, (int)1)); + stride_weight_ptr[expert] = BSConfig::tile_atom_to_shape_SFA( + cute::make_shape((int)gemm_n, (int)gemm_m, (int)gemm_k, (int)1)); + } else { + stride_act_ptr[expert] = BSConfig::tile_atom_to_shape_SFA( + cute::make_shape((int)gemm_m, (int)gemm_n, (int)gemm_k, (int)1)); + stride_weight_ptr[expert] = BSConfig::tile_atom_to_shape_SFB( + cute::make_shape((int)gemm_m, (int)gemm_n, (int)gemm_k, (int)1)); + } + + // This assert validates our current assumption that A&B can be safely transposed without needing + // to modify + assert( + BSConfig::tile_atom_to_shape_SFB( + cute::make_shape((int)gemm_n, (int)gemm_m, (int)gemm_k, 1)) == + BSConfig::tile_atom_to_shape_SFA(cute::make_shape((int)gemm_m, (int)gemm_n, (int)gemm_k, 1))); + + auto scaling_type = + std::is_same_v + ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 + : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; + layout_info.fpX_block_scaling_factors_act[expert] = + fp4_act_flat + getOffsetActivationSF(expert, num_tokens_before_expert, gemm_k, scaling_type); + + layout_info.fpX_block_scaling_factors_weight[expert] = + weight_block_scale + getOffsetWeightSF(expert, gemm_n, gemm_k, scaling_type); +} + +__device__ void computeTmaWarpSpecializedInputStrides( + TmaWarpSpecializedGroupedGemmInput& layout_info, int gemm_m, int gemm_n, int gemm_k, + int64_t out_idx) { + if (layout_info.swap_ab) { + reinterpret_cast( + layout_info.stride_act)[out_idx] = + cutlass::make_cute_packed_stride(TmaWarpSpecializedGroupedGemmInput::StrideB{}, + cute::make_shape(gemm_m, gemm_k, 1)); + reinterpret_cast( + layout_info.stride_weight)[out_idx] = + cutlass::make_cute_packed_stride(TmaWarpSpecializedGroupedGemmInput::StrideA{}, + cute::make_shape(gemm_n, gemm_k, 1)); + } else { + reinterpret_cast( + layout_info.stride_act)[out_idx] = + cutlass::make_cute_packed_stride(TmaWarpSpecializedGroupedGemmInput::StrideA{}, + cute::make_shape(gemm_m, gemm_k, 1)); + reinterpret_cast( + layout_info.stride_weight)[out_idx] = + cutlass::make_cute_packed_stride(TmaWarpSpecializedGroupedGemmInput::StrideB{}, + cute::make_shape(gemm_n, gemm_k, 1)); + } + if (layout_info.stride_c) { + // TODO Enable 1xN bias matrix as C + assert(false && "CUTLASS does not support a 1xN bias"); + } + if (layout_info.fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE) { + if (layout_info.swap_ab) { + reinterpret_cast( + layout_info.stride_d)[out_idx] = + cutlass::make_cute_packed_stride(TmaWarpSpecializedGroupedGemmInput::StrideD_T{}, + cute::make_shape(gemm_n, gemm_m, 1)); + } else { + reinterpret_cast( + layout_info.stride_d)[out_idx] = + cutlass::make_cute_packed_stride(TmaWarpSpecializedGroupedGemmInput::StrideD{}, + cute::make_shape(gemm_m, gemm_n, 1)); + } + } + if (layout_info.int4_groupwise_params.enabled) { + layout_info.int4_groupwise_params.stride_s_a[out_idx] = cutlass::make_cute_packed_stride( + TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::StrideSFA{}, + cute::make_shape( + gemm_n, + gemm_k / + (layout_info.int4_groupwise_params.use_wfp4a16 + ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size + : TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size), + 1)); + } +} + +template +__device__ void computeTmaWarpSpecializedInputPointers( + TmaWarpSpecializedGroupedGemmInput& layout_info, int64_t gemm_m, int64_t gemm_n, int64_t gemm_k, + int num_tokens_before_expert, int64_t expert, T const* in, WeightType const* weights, + TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::SFA const* w4a8_weight_scale, + ScaleBiasType const* bias, OutputType* output, float const* router_scales, + int const* permuted_row_to_unpermuted_row, int64_t const out_idx) { + // The input prior to this contains K elements per token, with `num_tokens_before_expert` tokens + layout_info.ptr_act[out_idx] = safe_inc_ptr(in, num_tokens_before_expert * gemm_k); + + // Each expert's weight matrix is a constant size NxK, get the matrix at index `expert` + layout_info.ptr_weight[out_idx] = safe_inc_ptr(weights, expert * (gemm_n * gemm_k)); + + if (layout_info.fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE) { + // The output prior to this contains N elements per token, with `num_tokens_before_expert` + // tokens + layout_info.ptr_d[out_idx] = safe_inc_ptr(output, num_tokens_before_expert * gemm_n); + } + if (layout_info.fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE) { + layout_info.fused_finalize_epilogue.ptr_source_token_index[expert] = + permuted_row_to_unpermuted_row + num_tokens_before_expert; + layout_info.fused_finalize_epilogue.ptr_router_scales[expert] = + router_scales + num_tokens_before_expert; + if (layout_info.fused_finalize_epilogue.ptr_bias != nullptr) { + layout_info.fused_finalize_epilogue.ptr_bias[expert] = bias + gemm_n * expert; + } + } + if (layout_info.int4_groupwise_params.enabled) { + // The group size of wfp4a16 is multiplied by 2 because each scale uses 1 byte instead of 2 + // bytes + layout_info.int4_groupwise_params.ptr_s_a[out_idx] = safe_inc_ptr( + w4a8_weight_scale, + expert * + (gemm_n * gemm_k / + (layout_info.int4_groupwise_params.use_wfp4a16 + ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size * 2 + : TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size))); + } +} + +// TODO Some of this setup could be cached +template +__global__ void computeStridesTmaWarpSpecializedKernel( + int64_t const* expert_first_token_offset, TmaWarpSpecializedGroupedGemmInput layout_info1, + TmaWarpSpecializedGroupedGemmInput layout_info2, int64_t num_tokens, + int64_t expanded_num_tokens, int64_t gemm1_n, int64_t gemm1_k, int64_t gemm2_n, int64_t gemm2_k, + int64_t const num_experts_per_node, T const* gemm1_in, T const* gemm2_in, + WeightType const* weights1, WeightType const* weights2, float const* alpha_scale_flat1, + float const* alpha_scale_flat2, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* fp4_act_flat1, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* fp4_act_flat2, QuantParams quant_params, + ScaleBiasType const* bias1, ScaleBiasType const* bias2, OutputType* gemm1_output, + OutputType* gemm2_output, float const* router_scales, + int const* permuted_row_to_unpermuted_row) { + // First, compute the global tid. We only need 1 thread per expert. + int const expert = blockIdx.x * blockDim.x + threadIdx.x; + if (expert >= num_experts_per_node) { + return; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + // Both gemms use the same token offset + auto const num_tokens_before_expert = expert_first_token_offset[expert]; + auto const num_tokens_including_expert = expert_first_token_offset[expert + 1]; + auto const num_tokens_to_expert = num_tokens_including_expert - num_tokens_before_expert; + auto const gemm_m = num_tokens_to_expert; + + // M and N transposed since we are using the #tokens as the N dimension + layout_info1.shape_info.problem_shapes[expert] = + TmaWarpSpecializedGroupedGemmInput::ProblemShape::UnderlyingProblemShape( + layout_info1.swap_ab ? gemm1_n : gemm_m, layout_info1.swap_ab ? gemm_m : gemm1_n, + gemm1_k); + layout_info2.shape_info.problem_shapes[expert] = + TmaWarpSpecializedGroupedGemmInput::ProblemShape::UnderlyingProblemShape( + layout_info2.swap_ab ? gemm2_n : gemm_m, layout_info2.swap_ab ? gemm_m : gemm2_n, + gemm2_k); + + if (layout_info1.int4_groupwise_params.enabled) { + layout_info1.int4_groupwise_params.shape.problem_shapes[expert] = + TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::ProblemShapeInt:: + UnderlyingProblemShape(layout_info1.swap_ab ? gemm1_n : gemm_m, + layout_info1.swap_ab ? gemm_m : gemm1_n, gemm1_k); + } + + if (layout_info2.int4_groupwise_params.enabled) { + layout_info2.int4_groupwise_params.shape.problem_shapes[expert] = + TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::ProblemShapeInt:: + UnderlyingProblemShape(layout_info2.swap_ab ? gemm2_n : gemm_m, + layout_info2.swap_ab ? gemm_m : gemm2_n, gemm2_k); + } + + if (alpha_scale_flat1 && alpha_scale_flat2) { + layout_info1.alpha_scale_ptr_array[expert] = alpha_scale_flat1 + expert; + layout_info2.alpha_scale_ptr_array[expert] = alpha_scale_flat2 + expert; + } + + auto setupIfSelected = [&](auto bs_config, auto quant_type) { + if (quant_type.fc1.weight_block_scale) { + setupFP4BlockScalingFactors( + layout_info1, expert, gemm_m, gemm1_n, gemm1_k, fp4_act_flat1, + quant_type.fc1.weight_block_scale, num_tokens_before_expert); + } + if (quant_type.fc2.weight_block_scale) { + setupFP4BlockScalingFactors( + layout_info2, expert, gemm_m, gemm2_n, gemm2_k, fp4_act_flat2, + quant_type.fc2.weight_block_scale, num_tokens_before_expert); + } + }; + + setupIfSelected(TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaledConfig{}, quant_params.fp4); + setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, + quant_params.fp8_mxfp4); + setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, + quant_params.mxfp8_mxfp4); + + assert(gemm_m <= INT32_MAX); + assert(gemm1_n > 0 && gemm1_n <= INT32_MAX); + assert(gemm1_k > 0 && gemm1_k <= INT32_MAX); + assert(gemm2_n > 0 && gemm2_n <= INT32_MAX); + assert(gemm2_k > 0 && gemm2_k <= INT32_MAX); + computeTmaWarpSpecializedInputStrides(layout_info1, gemm_m, gemm1_n, gemm1_k, expert); + computeTmaWarpSpecializedInputStrides(layout_info2, gemm_m, gemm2_n, gemm2_k, expert); + + computeTmaWarpSpecializedInputPointers( + layout_info1, gemm_m, gemm1_n, gemm1_k, num_tokens_before_expert, expert, gemm1_in, weights1, + reinterpret_cast( + quant_params.groupwise.fc1.weight_scales), + bias1, gemm1_output, nullptr, nullptr, expert); + computeTmaWarpSpecializedInputPointers( + layout_info2, gemm_m, gemm2_n, gemm2_k, num_tokens_before_expert, expert, gemm2_in, weights2, + reinterpret_cast( + quant_params.groupwise.fc2.weight_scales), + bias2, gemm2_output, router_scales, permuted_row_to_unpermuted_row, expert); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +// ========================== Permutation things ======================================= + +template +__host__ __device__ constexpr static U arrayConvert(T const& input) { + cutlass::NumericArrayConverter converter; + return converter(input); +} + +// Duplicated and permutes rows for MoE. In addition, reverse the permutation map to help with +// finalizing routing. + +// "expanded_x_row" simply means that the number of values is num_rows x k. It is "expanded" since +// we will have to duplicate some rows in the input matrix to match the dimensions. Duplicates will +// always get routed to separate experts in the end. + +// Note that the permuted_row_to_unpermuted_row map referred to here has indices in the range (0, +// k*rows_in_input - 1). However, it is set up so that index 0, rows_in_input, 2*rows_in_input ... +// (k-1)*rows_in_input all map to row 0 in the original matrix. Thus, to know where to read in the +// source matrix, we simply take the modulus of the expanded index. + +constexpr static int EXPAND_THREADS_PER_BLOCK = 256; + +template +__global__ void expandInputRowsKernel( + InputActivationsType const* unpermuted_input, ExpandedActivationsType* permuted_output, + float const* unpermuted_scales, float* permuted_scales, + int const* permuted_row_to_unpermuted_row, int64_t const num_tokens, int64_t const hidden_size, + int64_t const k, float const* fc1_act_global_scale, bool use_per_expert_act_scale, + int64_t const* expert_first_token_offset, + TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, + int64_t const num_experts_per_node, InputActivationsType const* prequant_scales = nullptr) { + static_assert(BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE || + !PRE_QUANT_AWQ, + "AWQ and Block Scaling are mutually exclusive"); + constexpr bool is_mxfp8 = + std::is_same_v && + BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX && + !PRE_QUANT_AWQ; + constexpr bool is_mxfp8_input = is_mxfp8 && std::is_same_v; + constexpr bool need_mxfp8_quant = is_mxfp8 && !is_mxfp8_input; + +#ifdef ENABLE_FP4 + constexpr bool is_nvfp4 = + std::is_same_v && + BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 && + !PRE_QUANT_AWQ; + constexpr bool is_nvfp4_input = is_nvfp4 && std::is_same_v; + constexpr bool need_nvfp4_quant = is_nvfp4 && !is_nvfp4_input; +#else + constexpr bool is_nvfp4 = false; + constexpr bool is_nvfp4_input = false; + constexpr bool need_nvfp4_quant = false; +#endif + + static_assert(need_nvfp4_quant || need_mxfp8_quant || PRE_QUANT_AWQ || + std::is_same_v, + "Only NVFP4, MXFP8 and WINT4_AFP8 supports outputting a different format as part " + "of the expansion"); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + constexpr int VecSize = is_nvfp4 ? TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize + : TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize; + + constexpr int64_t ELEM_PER_THREAD = (is_nvfp4 || is_mxfp8) + ? CVT_ELTS_PER_THREAD + : (128 / sizeof_bits::value); + + // This should be VecSize * 4 elements + // We assume at least VecSize alignment or the quantization will fail + constexpr int64_t min_k_dim_alignment = + is_nvfp4 ? TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4 + : TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX; + int64_t const padded_hidden_size = + TmaWarpSpecializedGroupedGemmInput::alignToSfDim(hidden_size, min_k_dim_alignment); + + int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; + for (int64_t permuted_row = blockIdx.x; permuted_row < num_valid_tokens; + permuted_row += gridDim.x) { + int64_t const unpermuted_row = permuted_row_to_unpermuted_row[permuted_row]; + + // Load 128-bits per thread + + constexpr int64_t ELEM_PER_BYTE = is_nvfp4_input ? 2 : 1; + using DataElem = std::conditional_t< + is_nvfp4_input, uint32_t, + std::conditional_t>>; + using OutputElem = std::conditional_t< + is_nvfp4, uint32_t, + std::conditional_t>>; + + // Duplicate and permute rows + int64_t const source_k_rank = unpermuted_row / num_tokens; + int64_t const source_row = unpermuted_row % num_tokens; + + auto const* source_row_ptr = reinterpret_cast( + unpermuted_input + source_row * hidden_size / ELEM_PER_BYTE); + // Cast first to handle when this is FP4 + auto* dest_row_ptr = reinterpret_cast(permuted_output) + + permuted_row * hidden_size / ELEM_PER_THREAD; + + int64_t const start_offset = threadIdx.x; + int64_t const stride = EXPAND_THREADS_PER_BLOCK; + int64_t const num_elems_in_col = hidden_size / ELEM_PER_THREAD; + assert(hidden_size % ELEM_PER_THREAD == 0); + assert(hidden_size % VecSize == 0); + + if constexpr (is_nvfp4 || is_mxfp8) { + static_assert(ELEM_PER_THREAD == 8, "Expecting 8 elements per thread for quantized types"); + int64_t expert = findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, + (int64_t)permuted_row + 1) - + 1; + + assert(!fc1_act_global_scale || is_nvfp4 && "Global scale is only supported for NVFP4"); + size_t act_scale_idx = use_per_expert_act_scale ? expert : 0; + float global_scale_val = fc1_act_global_scale ? fc1_act_global_scale[act_scale_idx] : 1.0f; + int64_t num_tokens_before_expert = expert_first_token_offset[expert]; + + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { + auto in_vec = source_row_ptr[elem_index]; + if constexpr (need_nvfp4_quant || need_mxfp8_quant) { + auto res = quantizePackedFPXValue( + in_vec, global_scale_val, num_tokens_before_expert, expert, permuted_row, elem_index, + padded_hidden_size, fc1_act_sf_flat, + is_nvfp4 ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 + : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX); + static_assert(sizeof(res) == sizeof(*dest_row_ptr), + "Quantized value must be the same size as the output"); + dest_row_ptr[elem_index] = res; + } else { + assert(act_scale_idx == 0 && + "Cannot use per-expert act scale for pre-quantized activations"); + writeSF(num_tokens_before_expert, expert, source_row, + permuted_row, elem_index, padded_hidden_size, + fc1_act_sf_flat, input_sf, swizzled_input_sf); + dest_row_ptr[elem_index] = in_vec; + } + } + + // Pad zeros in the extra SFs along the K dimension, we do this to ensure there are no nan + // values in the padded SF atom Use VecSize per thread since we are just writing out zeros so + // every thread can process a whole vector + size_t padding_start_offset = hidden_size / VecSize + start_offset; + size_t padding_elems_in_col = padded_hidden_size / VecSize; + for (int64_t elem_index = padding_start_offset; elem_index < padding_elems_in_col; + elem_index += stride) { + writeSF(num_tokens_before_expert, expert, /*source_row*/ -1, permuted_row, + elem_index, padded_hidden_size, fc1_act_sf_flat, + /* input_sf */ nullptr); // Pass nulltpr input_sf so we write 0 + } + } else if constexpr (PRE_QUANT_AWQ) { + static_assert(!is_nvfp4 && !is_mxfp8, "NVFP4 and MXFP8 are not supported for AWQ"); + static_assert(!std::is_same_v, + "Input and output types must be different for AWQ"); + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { + auto frag_elems = source_row_ptr[elem_index]; + + CUTLASS_PRAGMA_UNROLL + for (int e = 0; e < ELEM_PER_THREAD; e++) { + frag_elems[e] = frag_elems[e] * prequant_scales[elem_index * ELEM_PER_THREAD + e]; + } + + dest_row_ptr[elem_index] = arrayConvert(frag_elems); + } + } else { + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { + dest_row_ptr[elem_index] = source_row_ptr[elem_index]; + } + } + + if (permuted_scales && threadIdx.x == 0) { + int64_t const source_k_idx = source_row * k + source_k_rank; + permuted_scales[permuted_row] = unpermuted_scales ? unpermuted_scales[source_k_idx] : 1.0f; + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + + // Pad zeros in the extra SFs along the N dimension, we do this to ensure there are no nan values + // in the padded SF atom + if constexpr (is_nvfp4 || is_mxfp8) { + int64_t const start_offset = threadIdx.x; + int64_t const stride = EXPAND_THREADS_PER_BLOCK; + // Use VecSize per thread since we are just writing out zeros so every thread can process a + // whole vector + int64_t const padded_num_elems_in_col = padded_hidden_size / VecSize; + assert(padded_hidden_size % VecSize == 0); + + constexpr int min_num_tokens_alignment = + is_nvfp4 ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 + : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; + static_assert((min_num_tokens_alignment & (min_num_tokens_alignment - 1)) == 0, + "Min num tokens alignment must be a power of two"); + // Since we don't know a priori how much padding is needed we assume the max per expert + // NOTE: we don't use (min_num_tokens_alignment-1) to be able to do power of two divisions + int64_t num_padding_tokens = min_num_tokens_alignment * num_experts_per_node; + + for (int64_t padding_token = blockIdx.x; padding_token < num_padding_tokens; + padding_token += gridDim.x) { + int64_t expert = padding_token / min_num_tokens_alignment; + int64_t num_tokens_before_expert = expert_first_token_offset[expert]; + int64_t num_tokens_after_expert = expert_first_token_offset[expert + 1]; + int64_t tokens_to_expert = num_tokens_after_expert - num_tokens_before_expert; + int64_t padding_to_expert = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( + tokens_to_expert, min_num_tokens_alignment) - + tokens_to_expert; + int64_t expert_pad_idx = padding_token % min_num_tokens_alignment; + if (expert_pad_idx < padding_to_expert) { + for (int64_t elem_index = start_offset; elem_index < padded_num_elems_in_col; + elem_index += stride) { + writeSF(num_tokens_before_expert, expert, /*source_row*/ -1, + num_tokens_after_expert + expert_pad_idx, elem_index, + padded_hidden_size, fc1_act_sf_flat, + /* input_sf */ nullptr); // Pass nulltpr input_sf so we write 0 + } + } + } + } +} + +template +void expandInputRowsKernelLauncher( + InputActivationsType const* unpermuted_input, ExpandedActivationsType* permuted_output, + float const* unpermuted_scales, float* permuted_scales, + int const* permuted_row_to_unpermuted_row, int64_t const num_rows, int64_t const hidden_size, + int const k, int const num_experts_per_node, QuantParams const& quant_params, + bool use_per_expert_act_scale, int64_t* expert_first_token_offset, + TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, + TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, + void const* prequant_scales, bool enable_pdl, cudaStream_t stream) { +#ifdef ENABLE_FP4 + TLLM_CHECK_WITH_INFO( + (std::is_same_v && fc1_act_sf_flat) || + !use_per_expert_act_scale, + "Per-expert act scale for FC1 is only supported for NVFP4 activations"); + constexpr int64_t min_num_tokens_alignment = + std::is_same_v + ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 + : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; + int64_t num_padding_tokens = min_num_tokens_alignment * num_experts_per_node; +#else + int64_t num_padding_tokens = 0; +#endif + + static int64_t const smCount = tensorrt_llm::common::getMultiProcessorCount(); + // Note: Launching 8 blocks per SM can fully leverage the memory bandwidth (tested on B200). + int64_t const blocks = std::min(smCount * 8, std::max(num_rows * k, num_padding_tokens)); + int64_t const threads = EXPAND_THREADS_PER_BLOCK; + + auto func = [&]() { +#ifdef ENABLE_FP8 + // Always MXFP8 + if constexpr (std::is_same_v && + !std::is_same_v) { + TLLM_CHECK_WITH_INFO(quant_params.mxfp8_mxfp4.fc1.weight_block_scale || prequant_scales, + "MXFP8xMXFP4 block scaling or prequant_scales or prequant_scales " + "parameters not provided"); + return prequant_scales + ? &expandInputRowsKernel< + InputActivationsType, ExpandedActivationsType, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, true> + : &expandInputRowsKernel< + InputActivationsType, ExpandedActivationsType, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX, false>; + } + // Could be either regular FP8 or MXFP8 + else if constexpr (std::is_same_v && + std::is_same_v) { + TLLM_CHECK_WITH_INFO(!prequant_scales, "FP8 is not supported for AWQ"); + return quant_params.mxfp8_mxfp4.fc1.weight_block_scale + ? &expandInputRowsKernel< + InputActivationsType, ExpandedActivationsType, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX, false> + : &expandInputRowsKernel< + InputActivationsType, ExpandedActivationsType, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, false>; + } else +#endif +#ifdef ENABLE_FP4 + if constexpr (std::is_same_v) { + TLLM_CHECK_WITH_INFO(quant_params.fp4.fc1.weight_block_scale, + "NVFP4 block scaling is expected for FP4xFP4"); + TLLM_CHECK_WITH_INFO(!prequant_scales, "NVFP4 is not supported for AWQ"); + return &expandInputRowsKernel; + } else +#endif + { + TLLM_CHECK_WITH_INFO(!prequant_scales, + "w4afp8 Prequant scales provided for non-FP8 data type"); + return &expandInputRowsKernel; + } + }(); + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, func, unpermuted_input, permuted_output, unpermuted_scales, + permuted_scales, permuted_row_to_unpermuted_row, num_rows, hidden_size, k, + quant_params.fp4.fc1.act_global_scale, use_per_expert_act_scale, + expert_first_token_offset, fc1_act_sf_flat, input_sf, swizzled_input_sf, + num_experts_per_node, + reinterpret_cast(prequant_scales)); +} + +#define INSTANTIATE_EXPAND_INPUT_ROWS(InputActivationsType, ExpandedActivationsType) \ + template void expandInputRowsKernelLauncher( \ + InputActivationsType const* unpermuted_input, ExpandedActivationsType* permuted_output, \ + float const* unpermuted_scales, float* permuted_scales, \ + int const* permuted_row_to_unpermuted_row, int64_t const num_rows, \ + int64_t const hidden_size, int const k, int const num_experts_per_node, \ + QuantParams const& quant_params, bool use_per_expert_act_scale, \ + int64_t* expert_first_token_offset, \ + TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, \ + TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, bool const swizzled_input_sf, \ + void const* prequant_scales, bool enable_pdl, cudaStream_t stream) + +// Instantiate the data types that are used by the external pytorch op +// INSTANTIATE_EXPAND_INPUT_ROWS(float, float); +// INSTANTIATE_EXPAND_INPUT_ROWS(half, half); +// #ifdef ENABLE_BF16 +// INSTANTIATE_EXPAND_INPUT_ROWS(__nv_bfloat16, __nv_bfloat16); +// #endif + +enum class ScaleMode : int { + NO_SCALE = 0, + DEFAULT = 1, +}; + +constexpr static int FINALIZE_THREADS_PER_BLOCK = 256; + +// Final kernel to unpermute and scale +// This kernel unpermutes the original data, does the k-way reduction and performs the final skip +// connection. +template +__global__ void finalizeMoeRoutingKernel( + GemmOutputType const* expanded_permuted_rows, OutputType* reduced_unpermuted_output, + ScaleBiasType const* bias, float const* scales, int const* unpermuted_row_to_permuted_row, + int const* token_selected_experts, int64_t const padded_cols, int64_t const unpadded_cols, + int64_t const experts_per_token, int const num_experts_per_node, int const start_expert_id) { + assert(padded_cols % 4 == 0); + assert(unpadded_cols % 4 == 0); + assert(unpadded_cols <= padded_cols); + int64_t const original_row = blockIdx.x; + int64_t const num_rows = gridDim.x; + auto const offset = original_row * unpadded_cols; + OutputType* reduced_row_ptr = reduced_unpermuted_output + offset; + + // Load 128-bits per thread, according to the smallest data type we read/write + constexpr int64_t FINALIZE_ELEM_PER_THREAD = + 128 / std::min(sizeof_bits::value, sizeof_bits::value); + + int64_t const start_offset = threadIdx.x; + int64_t const stride = FINALIZE_THREADS_PER_BLOCK; + int64_t const num_elems_in_padded_col = padded_cols / FINALIZE_ELEM_PER_THREAD; + int64_t const num_elems_in_orig_col = unpadded_cols / FINALIZE_ELEM_PER_THREAD; + + using BiasElem = cutlass::Array; + using InputElem = cutlass::Array; + using OutputElem = cutlass::Array; + using ComputeElem = cutlass::Array; + auto const* bias_v = reinterpret_cast(bias); + auto const* expanded_permuted_rows_v = reinterpret_cast(expanded_permuted_rows); + auto* reduced_row_ptr_v = reinterpret_cast(reduced_row_ptr); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +#pragma unroll + for (int elem_index = start_offset; elem_index < num_elems_in_orig_col; elem_index += stride) { + ComputeElem thread_output; + thread_output.fill(0); + for (int k_idx = 0; k_idx < experts_per_token; ++k_idx) { + int64_t const k_offset = original_row * experts_per_token + k_idx; + int64_t const expert_id = token_selected_experts[k_offset] - start_expert_id; + if (expert_id < 0 || expert_id >= num_experts_per_node) { + continue; + } + + int64_t const expanded_original_row = original_row + k_idx * num_rows; + int64_t const expanded_permuted_row = unpermuted_row_to_permuted_row[expanded_original_row]; + + float const row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset]; + + auto const* expanded_permuted_rows_row_ptr = + expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_padded_col; + + ComputeElem expert_result = + arrayConvert(expanded_permuted_rows_row_ptr[elem_index]); + if (bias) { + auto const* bias_ptr = bias_v + expert_id * num_elems_in_padded_col; + expert_result = expert_result + arrayConvert(bias_ptr[elem_index]); + } + + thread_output = thread_output + row_scale * expert_result; + } + + OutputElem output_elem = arrayConvert(thread_output); + reduced_row_ptr_v[elem_index] = output_elem; + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +// Final kernel to unpermute and scale +// This kernel unpermutes the original data, does the k-way reduction and performs the final skip +// connection. +template +__global__ void finalizeMoeRoutingNoFillingKernel( + GemmOutputType const* expanded_permuted_rows, OutputType* reduced_unpermuted_output, + ScaleBiasType const* bias, float const* scales, int const* const unpermuted_row_to_permuted_row, + int const* permuted_row_to_unpermuted_row, int const* token_selected_experts, + int64_t const* expert_first_token_offset, int64_t const num_rows, int64_t const padded_cols, + int64_t const unpadded_cols, int64_t const experts_per_token, int const num_experts_per_node, + int const start_expert_id) { + assert(padded_cols % 4 == 0); + assert(unpadded_cols % 4 == 0); + assert(unpadded_cols <= padded_cols); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; + for (int64_t expanded_permuted_row = blockIdx.x; expanded_permuted_row < num_valid_tokens; + expanded_permuted_row += gridDim.x) { + int64_t unpermuted_row = permuted_row_to_unpermuted_row[expanded_permuted_row]; + + // Duplicate and permute rows + int64_t const source_k_rank = unpermuted_row / num_rows; + int64_t const source_row = unpermuted_row % num_rows; + + // If the expert is the first selected (valid) one of the corresponding token on the current EP + // rank, do reduction; otherwise, skip. + bool is_first_selected_expert = true; + for (int k_idx = 0; k_idx < source_k_rank; ++k_idx) { + int const expert_id = + token_selected_experts[source_row * experts_per_token + k_idx] - start_expert_id; + if (expert_id >= 0 && expert_id < num_experts_per_node) { + is_first_selected_expert = false; + break; + } + } + if (!is_first_selected_expert) { + continue; + } + + OutputType* reduced_row_ptr = reduced_unpermuted_output + source_row * unpadded_cols; + + // Load 128-bits per thread, according to the smallest data type we read/write + constexpr int64_t FINALIZE_ELEM_PER_THREAD = + 128 / std::min(sizeof_bits::value, sizeof_bits::value); + + int64_t const start_offset = threadIdx.x; + int64_t const stride = FINALIZE_THREADS_PER_BLOCK; + int64_t const num_elems_in_padded_col = padded_cols / FINALIZE_ELEM_PER_THREAD; + int64_t const num_elems_in_orig_col = unpadded_cols / FINALIZE_ELEM_PER_THREAD; + + using BiasElem = cutlass::Array; + using InputElem = cutlass::Array; + using OutputElem = cutlass::Array; + using ComputeElem = cutlass::Array; + auto const* bias_v = reinterpret_cast(bias); + auto const* expanded_permuted_rows_v = + reinterpret_cast(expanded_permuted_rows); + auto* reduced_row_ptr_v = reinterpret_cast(reduced_row_ptr); + + for (int elem_index = start_offset; elem_index < num_elems_in_padded_col; + elem_index += stride) { + if (elem_index >= num_elems_in_orig_col) continue; // Skip writing beyond original columns + + ComputeElem thread_output; + thread_output.fill(0); + for (int k_idx = 0; k_idx < experts_per_token; ++k_idx) { + int64_t const k_offset = source_row * experts_per_token + k_idx; + int64_t const expert_id = token_selected_experts[k_offset] - start_expert_id; + if (expert_id < 0 || expert_id >= num_experts_per_node) { + continue; + } + + int64_t const expanded_permuted_row_from_k_idx = + unpermuted_row_to_permuted_row[source_row + k_idx * num_rows]; + + float const row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset]; + + auto const* expanded_permuted_rows_row_ptr = + expanded_permuted_rows_v + expanded_permuted_row_from_k_idx * num_elems_in_padded_col; + + ComputeElem expert_result = + arrayConvert(expanded_permuted_rows_row_ptr[elem_index]); + + if (bias) { + auto const* bias_ptr = bias_v + expert_id * num_elems_in_padded_col; + expert_result = expert_result + arrayConvert(bias_ptr[elem_index]); + } + + thread_output = thread_output + row_scale * expert_result; + } + OutputElem output_elem = arrayConvert(thread_output); + reduced_row_ptr_v[elem_index] = output_elem; + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +void finalizeMoeRoutingKernelLauncher( + GemmOutputType const* expanded_permuted_rows, OutputType* reduced_unpermuted_output, + ScaleBiasType const* bias, float const* final_scales, int const* unpermuted_row_to_permuted_row, + int const* permuted_row_to_unpermuted_row, int const* token_selected_experts, + int64_t const* expert_first_token_offset, int64_t const num_rows, int64_t const padded_cols, + int64_t const unpadded_cols, int64_t const experts_per_token, + int64_t const num_experts_per_node, MOEParallelismConfig parallelism_config, + bool const enable_alltoall, bool enable_pdl, cudaStream_t stream) { + // Only add bias on rank 0 for tensor parallelism + bool const is_rank_0 = parallelism_config.tp_rank == 0; + ScaleBiasType const* bias_ptr = is_rank_0 ? bias : nullptr; + int const start_expert_id = num_experts_per_node * parallelism_config.ep_rank; + + cudaLaunchConfig_t config; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + if (parallelism_config.ep_size > 1 && enable_alltoall) { + // If all-to-all comm is enabled, finalizeMoeRouting doesn't need to fill the invalid output + // tokens with zeros. + static int const smCount = tensorrt_llm::common::getMultiProcessorCount(); + // Note: Launching 8 blocks per SM can fully leverage the memory bandwidth (tested on B200). + int64_t const blocks = smCount * 8; + int64_t const threads = FINALIZE_THREADS_PER_BLOCK; + config.gridDim = blocks; + config.blockDim = threads; + auto func = final_scales + ? &finalizeMoeRoutingNoFillingKernel + : &finalizeMoeRoutingNoFillingKernel; + cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, + final_scales, unpermuted_row_to_permuted_row, permuted_row_to_unpermuted_row, + token_selected_experts, expert_first_token_offset, num_rows, padded_cols, + unpadded_cols, experts_per_token, num_experts_per_node, start_expert_id); + } else { + // If all-gather reduce-scatter is used, finalizeMoeRouting must fill invalid output tokens with + // zeros. + int64_t const blocks = num_rows; + int64_t const threads = FINALIZE_THREADS_PER_BLOCK; + config.gridDim = blocks; + config.blockDim = threads; + auto func = final_scales ? &finalizeMoeRoutingKernel + : &finalizeMoeRoutingKernel; + cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, + final_scales, unpermuted_row_to_permuted_row, token_selected_experts, + padded_cols, unpadded_cols, experts_per_token, num_experts_per_node, + start_expert_id); + } +} + +#define INSTANTIATE_FINALIZE_MOE_ROUTING(OutputT, GemmOutputT, ScaleBiasT) \ + template void finalizeMoeRoutingKernelLauncher( \ + GemmOutputT const* expanded_permuted_rows, OutputT* reduced_unpermuted_output, \ + ScaleBiasT const* bias, float const* final_scales, \ + int const* unpermuted_row_to_permuted_row, int const* permuted_row_to_unpermuted_row, \ + int const* expert_for_source_row, int64_t const* expert_first_token_offset, \ + int64_t const num_rows, int64_t const padded_cols, int64_t const actual_cols, \ + int64_t const experts_per_token, int64_t const num_experts_per_node, \ + MOEParallelismConfig parallelism_config, bool const enable_alltoall, bool enable_pdl, \ + cudaStream_t stream); + +// // Instantiate the data types that are used by the external pytorch op +// INSTANTIATE_FINALIZE_MOE_ROUTING(half, half, half); +// INSTANTIATE_FINALIZE_MOE_ROUTING(float, float, float); +// #ifdef ENABLE_BF16 +// INSTANTIATE_FINALIZE_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, __nv_bfloat16); +// #endif + +// ============================== Activation Adaptors ================================= +template